|
|
@@ -7,12 +7,187 @@ from typing import List, Dict, Optional
|
|
|
from pathlib import Path
|
|
|
import cv2
|
|
|
import numpy as np
|
|
|
+import math
|
|
|
+import os
|
|
|
from loguru import logger
|
|
|
|
|
|
|
|
|
class GridRecovery:
|
|
|
"""网格结构恢复工具类"""
|
|
|
|
|
|
+ # ==================== 辅助方法:几何计算 ====================
|
|
|
+
|
|
|
+ @staticmethod
|
|
|
+ def _fit_line(p):
|
|
|
+ """拟合直线方程 Ax + By + C = 0"""
|
|
|
+ x1, y1 = p[0]
|
|
|
+ x2, y2 = p[1]
|
|
|
+ A = y2 - y1
|
|
|
+ B = x1 - x2
|
|
|
+ C = x2 * y1 - x1 * y2
|
|
|
+ return A, B, C
|
|
|
+
|
|
|
+ @staticmethod
|
|
|
+ def _point_line_distance(p, A, B, C):
|
|
|
+ """计算点到直线的距离"""
|
|
|
+ x, y = p
|
|
|
+ return A * x + B * y + C
|
|
|
+
|
|
|
+ @staticmethod
|
|
|
+ def _dist_sqrt(p1, p2):
|
|
|
+ """计算两点间距离"""
|
|
|
+ return np.sqrt((p1[0] - p2[0]) ** 2 + (p1[1] - p2[1]) ** 2)
|
|
|
+
|
|
|
+ @staticmethod
|
|
|
+ def _line_to_line(points1, points2, alpha=10, angle=30, max_len=None):
|
|
|
+ """将线段延长到与另一线段相交"""
|
|
|
+ x1, y1, x2, y2 = points1
|
|
|
+ ox1, oy1, ox2, oy2 = points2
|
|
|
+
|
|
|
+ current_len = GridRecovery._dist_sqrt((x1, y1), (x2, y2))
|
|
|
+ if max_len is not None and current_len >= max_len:
|
|
|
+ return points1
|
|
|
+
|
|
|
+ step_limit = current_len
|
|
|
+ effective_alpha = min(alpha, step_limit)
|
|
|
+
|
|
|
+ xy = np.array([(x1, y1), (x2, y2)], dtype="float32")
|
|
|
+ A1, B1, C1 = GridRecovery._fit_line(xy)
|
|
|
+ oxy = np.array([(ox1, oy1), (ox2, oy2)], dtype="float32")
|
|
|
+ A2, B2, C2 = GridRecovery._fit_line(oxy)
|
|
|
+
|
|
|
+ flag1 = GridRecovery._point_line_distance(np.array([x1, y1], dtype="float32"), A2, B2, C2)
|
|
|
+ flag2 = GridRecovery._point_line_distance(np.array([x2, y2], dtype="float32"), A2, B2, C2)
|
|
|
+
|
|
|
+ if (flag1 > 0 and flag2 > 0) or (flag1 < 0 and flag2 < 0):
|
|
|
+ if (A1 * B2 - A2 * B1) != 0:
|
|
|
+ x = (B1 * C2 - B2 * C1) / (A1 * B2 - A2 * B1)
|
|
|
+ y = (A2 * C1 - A1 * C2) / (A1 * B2 - A2 * B1)
|
|
|
+ p = (x, y)
|
|
|
+ r0 = GridRecovery._dist_sqrt(p, (x1, y1))
|
|
|
+ r1 = GridRecovery._dist_sqrt(p, (x2, y2))
|
|
|
+
|
|
|
+ if min(r0, r1) < effective_alpha:
|
|
|
+ if max_len is not None:
|
|
|
+ new_len = GridRecovery._dist_sqrt(p, (x2, y2)) if r0 < r1 else GridRecovery._dist_sqrt((x1, y1), p)
|
|
|
+ if new_len > max_len:
|
|
|
+ return points1
|
|
|
+
|
|
|
+ if r0 < r1:
|
|
|
+ k = abs((y2 - p[1]) / (x2 - p[0] + 1e-10))
|
|
|
+ a = math.atan(k) * 180 / math.pi
|
|
|
+ if a < angle or abs(90 - a) < angle:
|
|
|
+ points1 = np.array([p[0], p[1], x2, y2], dtype="float32")
|
|
|
+ else:
|
|
|
+ k = abs((y1 - p[1]) / (x1 - p[0] + 1e-10))
|
|
|
+ a = math.atan(k) * 180 / math.pi
|
|
|
+ if a < angle or abs(90 - a) < angle:
|
|
|
+ points1 = np.array([x1, y1, p[0], p[1]], dtype="float32")
|
|
|
+ return points1
|
|
|
+
|
|
|
+ @staticmethod
|
|
|
+ def _extend_lines(rowboxes, colboxes, alpha=50):
|
|
|
+ """延长线段使其相交"""
|
|
|
+ extension_multiplier = 2.0
|
|
|
+ row_max_lens = [GridRecovery._dist_sqrt(b[:2], b[2:]) * extension_multiplier for b in rowboxes]
|
|
|
+ col_max_lens = [GridRecovery._dist_sqrt(b[:2], b[2:]) * extension_multiplier for b in colboxes]
|
|
|
+
|
|
|
+ for i in range(len(rowboxes)):
|
|
|
+ for j in range(len(colboxes)):
|
|
|
+ rowboxes[i] = GridRecovery._line_to_line(rowboxes[i], colboxes[j], alpha=alpha, angle=30, max_len=row_max_lens[i])
|
|
|
+ colboxes[j] = GridRecovery._line_to_line(colboxes[j], rowboxes[i], alpha=alpha, angle=30, max_len=col_max_lens[j])
|
|
|
+ return rowboxes, colboxes
|
|
|
+
|
|
|
+ @staticmethod
|
|
|
+ def _filter_edge_lines(lines, img_h, img_w, margin):
|
|
|
+ """过滤贴近图像边缘的线条(padding 区域噪声)"""
|
|
|
+ filtered = []
|
|
|
+ removed = []
|
|
|
+
|
|
|
+ for line in lines:
|
|
|
+ x1, y1, x2, y2 = line
|
|
|
+ min_x, max_x = min(x1, x2), max(x1, x2)
|
|
|
+ min_y, max_y = min(y1, y2), max(y1, y2)
|
|
|
+ is_horizontal = abs(y2 - y1) < abs(x2 - x1)
|
|
|
+
|
|
|
+ should_remove = False
|
|
|
+ reason = ""
|
|
|
+
|
|
|
+ if is_horizontal:
|
|
|
+ if min_y < margin:
|
|
|
+ should_remove = True
|
|
|
+ reason = f"贴近上边缘 (min_y={min_y:.1f} < {margin})"
|
|
|
+ elif max_y > (img_h - margin):
|
|
|
+ should_remove = True
|
|
|
+ reason = f"贴近下边缘 (max_y={max_y:.1f} > {img_h - margin:.1f})"
|
|
|
+ else:
|
|
|
+ if min_x < margin:
|
|
|
+ should_remove = True
|
|
|
+ reason = f"贴近左边缘 (min_x={min_x:.1f} < {margin})"
|
|
|
+ elif max_x > (img_w - margin):
|
|
|
+ should_remove = True
|
|
|
+ reason = f"贴近右边缘 (max_x={max_x:.1f} > {img_w - margin:.1f})"
|
|
|
+
|
|
|
+ if should_remove:
|
|
|
+ removed.append((line, reason))
|
|
|
+ else:
|
|
|
+ filtered.append(line)
|
|
|
+
|
|
|
+ return filtered, removed
|
|
|
+
|
|
|
+ @staticmethod
|
|
|
+ def _filter_short_lines(lines, threshold):
|
|
|
+ """过滤过短的线段(噪声)"""
|
|
|
+ valid_lines = []
|
|
|
+ for line in lines:
|
|
|
+ x1, y1, x2, y2 = line
|
|
|
+ length = math.sqrt((x2-x1)**2 + (y2-y1)**2)
|
|
|
+ if length > threshold:
|
|
|
+ valid_lines.append(line)
|
|
|
+ return valid_lines
|
|
|
+
|
|
|
+ @staticmethod
|
|
|
+ def _filter_lines_by_bboxes(lines, bboxes, is_horizontal, tolerance=5.0):
|
|
|
+ """过滤线条,只保留与bboxes边界对齐的线条"""
|
|
|
+ if not bboxes:
|
|
|
+ return lines
|
|
|
+
|
|
|
+ if is_horizontal:
|
|
|
+ bbox_coords = {bbox[1] for bbox in bboxes} | {bbox[3] for bbox in bboxes}
|
|
|
+ else:
|
|
|
+ bbox_coords = {bbox[0] for bbox in bboxes} | {bbox[2] for bbox in bboxes}
|
|
|
+
|
|
|
+ filtered_lines = []
|
|
|
+ for line in lines:
|
|
|
+ line_coord = (line[1] + line[3]) / 2 if is_horizontal else (line[0] + line[2]) / 2
|
|
|
+ if any(abs(line_coord - coord) < tolerance for coord in bbox_coords):
|
|
|
+ filtered_lines.append(line)
|
|
|
+
|
|
|
+ return filtered_lines
|
|
|
+
|
|
|
+ @staticmethod
|
|
|
+ def _save_debug_image(debug_dir, debug_prefix, step_name, img, is_lines=False, lines=None):
|
|
|
+ """保存调试图片"""
|
|
|
+ if not debug_dir:
|
|
|
+ return
|
|
|
+ try:
|
|
|
+ os.makedirs(debug_dir, exist_ok=True)
|
|
|
+ name = f"{debug_prefix}_{step_name}.png" if debug_prefix else f"{step_name}.png"
|
|
|
+ path = os.path.join(debug_dir, name)
|
|
|
+
|
|
|
+ if is_lines and lines:
|
|
|
+ from mineru.model.table.rec.unet_table.utils_table_line_rec import draw_lines
|
|
|
+ tmp = np.zeros(img.shape[:2], dtype=np.uint8)
|
|
|
+ tmp = draw_lines(tmp, lines, color=255, lineW=2)
|
|
|
+ cv2.imwrite(path, tmp)
|
|
|
+ else:
|
|
|
+ cv2.imwrite(path, img)
|
|
|
+ logger.debug(f"Saved debug image: {path}")
|
|
|
+ except Exception as e:
|
|
|
+ logger.warning(f"Failed to save debug image {step_name}: {e}")
|
|
|
+
|
|
|
+ # ==================== 主要方法 ====================
|
|
|
+
|
|
|
@staticmethod
|
|
|
def compute_cells_from_lines(
|
|
|
hpred_up: np.ndarray,
|
|
|
@@ -58,12 +233,6 @@ class GridRecovery:
|
|
|
Returns:
|
|
|
单元格bbox列表 [[x1, y1, x2, y2], ...] (原图坐标系)
|
|
|
"""
|
|
|
- import numpy as np
|
|
|
- import cv2
|
|
|
- import math
|
|
|
- import os
|
|
|
- from loguru import logger
|
|
|
-
|
|
|
# 尝试导入MinerU的工具函数 (仅导入基础提取函数)
|
|
|
try:
|
|
|
from mineru.model.table.rec.unet_table.utils_table_line_rec import (
|
|
|
@@ -76,122 +245,6 @@ class GridRecovery:
|
|
|
logger.error("Could not import mineru utils. Please ensure MinerU is in python path.")
|
|
|
raise
|
|
|
|
|
|
- # --- Local Helper Functions for Robust Line Adjustment ---
|
|
|
- # Ported and modified from MinerU to verify larger gaps
|
|
|
-
|
|
|
- def fit_line(p):
|
|
|
- x1, y1 = p[0]
|
|
|
- x2, y2 = p[1]
|
|
|
- A = y2 - y1
|
|
|
- B = x1 - x2
|
|
|
- C = x2 * y1 - x1 * y2
|
|
|
- return A, B, C
|
|
|
-
|
|
|
- def point_line_cor(p, A, B, C):
|
|
|
- x, y = p
|
|
|
- r = A * x + B * y + C
|
|
|
- return r
|
|
|
-
|
|
|
- def dist_sqrt(p1, p2):
|
|
|
- return np.sqrt((p1[0] - p2[0]) ** 2 + (p1[1] - p2[1]) ** 2)
|
|
|
-
|
|
|
- def line_to_line(points1, points2, alpha=10, angle=30, max_len=None):
|
|
|
- x1, y1, x2, y2 = points1
|
|
|
- ox1, oy1, ox2, oy2 = points2
|
|
|
-
|
|
|
- # Calculate current line length
|
|
|
- current_len = dist_sqrt((x1, y1), (x2, y2))
|
|
|
-
|
|
|
- # If we already exceeded max_len, don't extend further
|
|
|
- if max_len is not None and current_len >= max_len:
|
|
|
- return points1
|
|
|
-
|
|
|
- # Dynamic Alpha based on CURRENT length (or capped by max extension per step)
|
|
|
- # We maintain the "step limit" to avoid huge jumps, but rely on max_len for total size.
|
|
|
- # effective_alpha = min(alpha, current_len)
|
|
|
- # (User previous logic: limit step to 1.0x length)
|
|
|
- step_limit = current_len
|
|
|
- effective_alpha = min(alpha, step_limit)
|
|
|
-
|
|
|
- # Fit lines
|
|
|
- xy = np.array([(x1, y1), (x2, y2)], dtype="float32")
|
|
|
- A1, B1, C1 = fit_line(xy)
|
|
|
- oxy = np.array([(ox1, oy1), (ox2, oy2)], dtype="float32")
|
|
|
- A2, B2, C2 = fit_line(oxy)
|
|
|
-
|
|
|
- flag1 = point_line_cor(np.array([x1, y1], dtype="float32"), A2, B2, C2)
|
|
|
- flag2 = point_line_cor(np.array([x2, y2], dtype="float32"), A2, B2, C2)
|
|
|
-
|
|
|
- # 如果位于同一侧(没穿过),尝试延伸
|
|
|
- if (flag1 > 0 and flag2 > 0) or (flag1 < 0 and flag2 < 0):
|
|
|
- if (A1 * B2 - A2 * B1) != 0:
|
|
|
- # 计算交点
|
|
|
- x = (B1 * C2 - B2 * C1) / (A1 * B2 - A2 * B1)
|
|
|
- y = (A2 * C1 - A1 * C2) / (A1 * B2 - A2 * B1)
|
|
|
- p = (x, y)
|
|
|
- r0 = dist_sqrt(p, (x1, y1))
|
|
|
- r1 = dist_sqrt(p, (x2, y2))
|
|
|
-
|
|
|
- if min(r0, r1) < effective_alpha:
|
|
|
- # Check total length constraint
|
|
|
- if max_len is not None:
|
|
|
- # Estimate new length
|
|
|
- if r0 < r1: # Extending (x1,y1) -> p
|
|
|
- new_len = dist_sqrt(p, (x2, y2))
|
|
|
- else: # Extending (x2,y2) -> p
|
|
|
- new_len = dist_sqrt((x1, y1), p)
|
|
|
-
|
|
|
- if new_len > max_len:
|
|
|
- return points1
|
|
|
-
|
|
|
- if r0 < r1:
|
|
|
- k = abs((y2 - p[1]) / (x2 - p[0] + 1e-10))
|
|
|
- a = math.atan(k) * 180 / math.pi
|
|
|
- if a < angle or abs(90 - a) < angle:
|
|
|
- points1 = np.array([p[0], p[1], x2, y2], dtype="float32")
|
|
|
- else:
|
|
|
- k = abs((y1 - p[1]) / (x1 - p[0] + 1e-10))
|
|
|
- a = math.atan(k) * 180 / math.pi
|
|
|
- if a < angle or abs(90 - a) < angle:
|
|
|
- points1 = np.array([x1, y1, p[0], p[1]], dtype="float32")
|
|
|
- return points1
|
|
|
-
|
|
|
- def custom_final_adjust_lines(rowboxes, colboxes, alpha=50):
|
|
|
- nrow = len(rowboxes)
|
|
|
- ncol = len(colboxes)
|
|
|
-
|
|
|
- # Pre-calculate Max Allowed Lengths (Original Length * Multiplier)
|
|
|
- # Multiplier = 2.0 means we allow the line to double in size, but not more.
|
|
|
- # This effectively stops short noise from becoming page-height lines.
|
|
|
- extension_multiplier = 2.0
|
|
|
-
|
|
|
- row_max_lens = [dist_sqrt(b[:2], b[2:]) * extension_multiplier for b in rowboxes]
|
|
|
- col_max_lens = [dist_sqrt(b[:2], b[2:]) * extension_multiplier for b in colboxes]
|
|
|
-
|
|
|
- for i in range(nrow):
|
|
|
- for j in range(ncol):
|
|
|
- rowboxes[i] = line_to_line(rowboxes[i], colboxes[j], alpha=alpha, angle=30, max_len=row_max_lens[i])
|
|
|
- colboxes[j] = line_to_line(colboxes[j], rowboxes[i], alpha=alpha, angle=30, max_len=col_max_lens[j])
|
|
|
- return rowboxes, colboxes
|
|
|
-
|
|
|
- def save_debug_image(step_name, img, is_lines=False, lines=None):
|
|
|
- if debug_dir:
|
|
|
- try:
|
|
|
- os.makedirs(debug_dir, exist_ok=True)
|
|
|
- name = f"{debug_prefix}_{step_name}.png" if debug_prefix else f"{step_name}.png"
|
|
|
- path = os.path.join(debug_dir, name)
|
|
|
-
|
|
|
- if is_lines and lines:
|
|
|
- # Draw lines on black background
|
|
|
- tmp = np.zeros(img.shape[:2], dtype=np.uint8)
|
|
|
- tmp = draw_lines(tmp, lines, color=255, lineW=2)
|
|
|
- cv2.imwrite(path, tmp)
|
|
|
- else:
|
|
|
- cv2.imwrite(path, img)
|
|
|
- logger.debug(f"Saved debug image: {path}")
|
|
|
- except Exception as e:
|
|
|
- logger.warning(f"Failed to save debug image {step_name}: {e}")
|
|
|
-
|
|
|
# ---------------------------------------------------------
|
|
|
|
|
|
h, w = hpred_up.shape[:2]
|
|
|
@@ -219,7 +272,7 @@ class GridRecovery:
|
|
|
logger.debug(f"Initial lines -> Rows: {len(rowboxes)}, Cols: {len(colboxes)}")
|
|
|
|
|
|
# Step 2 Debug
|
|
|
- save_debug_image("step02_raw_vectors", h_bin, is_lines=True, lines=rowboxes + colboxes)
|
|
|
+ GridRecovery._save_debug_image(debug_dir, debug_prefix, "step02_raw_vectors", h_bin, is_lines=True, lines=rowboxes + colboxes)
|
|
|
|
|
|
# ==================== 新增:边缘线过滤 ====================
|
|
|
# 2.5 过滤边缘线条(padding 区域的噪声)
|
|
|
@@ -237,67 +290,12 @@ class GridRecovery:
|
|
|
f"edge_margin={edge_margin}px (mask坐标系)"
|
|
|
)
|
|
|
|
|
|
- def filter_edge_lines(lines, img_h, img_w, margin):
|
|
|
- """
|
|
|
- 过滤贴近图像边缘的线条(padding 区域噪声)
|
|
|
-
|
|
|
- Args:
|
|
|
- lines: 线段列表 [[x1, y1, x2, y2], ...] (mask坐标系)
|
|
|
- img_h: mask图像高度
|
|
|
- img_w: mask图像宽度
|
|
|
- margin: 边缘阈值(像素,mask坐标系)
|
|
|
-
|
|
|
- Returns:
|
|
|
- (过滤后的线段列表, 被过滤的线段详情列表)
|
|
|
- """
|
|
|
- filtered = []
|
|
|
- removed = []
|
|
|
-
|
|
|
- for line in lines:
|
|
|
- x1, y1, x2, y2 = line
|
|
|
-
|
|
|
- # 计算线段的边界框
|
|
|
- min_x = min(x1, x2)
|
|
|
- max_x = max(x1, x2)
|
|
|
- min_y = min(y1, y2)
|
|
|
- max_y = max(y1, y2)
|
|
|
-
|
|
|
- # 判断线段方向(基于长度比例)
|
|
|
- is_horizontal = abs(y2 - y1) < abs(x2 - x1)
|
|
|
-
|
|
|
- should_remove = False
|
|
|
- reason = ""
|
|
|
-
|
|
|
- if is_horizontal:
|
|
|
- # 横线:检查是否贴近上下边缘
|
|
|
- if min_y < margin:
|
|
|
- should_remove = True
|
|
|
- reason = f"贴近上边缘 (min_y={min_y:.1f} < {margin})"
|
|
|
- elif max_y > (img_h - margin):
|
|
|
- should_remove = True
|
|
|
- reason = f"贴近下边缘 (max_y={max_y:.1f} > {img_h - margin:.1f})"
|
|
|
- else:
|
|
|
- # 竖线:检查是否贴近左右边缘
|
|
|
- if min_x < margin:
|
|
|
- should_remove = True
|
|
|
- reason = f"贴近左边缘 (min_x={min_x:.1f} < {margin})"
|
|
|
- elif max_x > (img_w - margin):
|
|
|
- should_remove = True
|
|
|
- reason = f"贴近右边缘 (max_x={max_x:.1f} > {img_w - margin:.1f})"
|
|
|
-
|
|
|
- if should_remove:
|
|
|
- removed.append((line, reason))
|
|
|
- else:
|
|
|
- filtered.append(line)
|
|
|
-
|
|
|
- return filtered, removed
|
|
|
-
|
|
|
# 执行边缘过滤
|
|
|
len_row_before = len(rowboxes)
|
|
|
len_col_before = len(colboxes)
|
|
|
|
|
|
- rowboxes_filtered, rowboxes_removed = filter_edge_lines(rowboxes, h, w, edge_margin)
|
|
|
- colboxes_filtered, colboxes_removed = filter_edge_lines(colboxes, h, w, edge_margin)
|
|
|
+ rowboxes_filtered, rowboxes_removed = GridRecovery._filter_edge_lines(rowboxes, h, w, edge_margin)
|
|
|
+ colboxes_filtered, colboxes_removed = GridRecovery._filter_edge_lines(colboxes, h, w, edge_margin)
|
|
|
|
|
|
# 详细日志
|
|
|
if rowboxes_removed or colboxes_removed:
|
|
|
@@ -330,7 +328,7 @@ class GridRecovery:
|
|
|
colboxes = colboxes_filtered
|
|
|
|
|
|
# Step 2.5 Debug(过滤后的干净线条)
|
|
|
- save_debug_image("step02b_edge_filtered", h_bin, is_lines=True, lines=rowboxes + colboxes)
|
|
|
+ GridRecovery._save_debug_image(debug_dir, debug_prefix, "step02b_edge_filtered", h_bin, is_lines=True, lines=rowboxes + colboxes)
|
|
|
# ==================== 边缘线过滤结束 ====================
|
|
|
|
|
|
# 3. 线段合并 (adjust_lines)
|
|
|
@@ -343,43 +341,28 @@ class GridRecovery:
|
|
|
colboxes += rboxes_col_
|
|
|
|
|
|
# Step 3 Debug
|
|
|
- save_debug_image("step03_merged_vectors", h_bin, is_lines=True, lines=rowboxes + colboxes)
|
|
|
+ GridRecovery._save_debug_image(debug_dir, debug_prefix, "step03_merged_vectors", h_bin, is_lines=True, lines=rowboxes + colboxes)
|
|
|
|
|
|
# 3.5 过滤短线 (Noise Filtering)
|
|
|
- # 在延长线段之前,过滤掉过短的线段(往往是噪声、文字下划线等)
|
|
|
- # 阈值: min(w, h) * 0.02, 至少 20px
|
|
|
filter_threshold = max(20, min(w, h) * 0.02)
|
|
|
-
|
|
|
- def filter_short_lines(lines, thresh):
|
|
|
- valid_lines = []
|
|
|
- for line in lines:
|
|
|
- x1, y1, x2, y2 = line
|
|
|
- length = math.sqrt((x2-x1)**2 + (y2-y1)**2)
|
|
|
- if length > thresh:
|
|
|
- valid_lines.append(line)
|
|
|
- return valid_lines
|
|
|
-
|
|
|
len_row_before = len(rowboxes)
|
|
|
len_col_before = len(colboxes)
|
|
|
|
|
|
- rowboxes = filter_short_lines(rowboxes, filter_threshold)
|
|
|
- colboxes = filter_short_lines(colboxes, filter_threshold)
|
|
|
+ rowboxes = GridRecovery._filter_short_lines(rowboxes, filter_threshold)
|
|
|
+ colboxes = GridRecovery._filter_short_lines(colboxes, filter_threshold)
|
|
|
|
|
|
if len(rowboxes) < len_row_before or len(colboxes) < len_col_before:
|
|
|
logger.info(f"Filtered short lines (thresh={filter_threshold:.1f}): Rows {len_row_before}->{len(rowboxes)}, Cols {len_col_before}->{len(colboxes)}")
|
|
|
- # Optional: Save filtered state
|
|
|
- save_debug_image("step03b_filtered_vectors", h_bin, is_lines=True, lines=rowboxes + colboxes)
|
|
|
+ GridRecovery._save_debug_image(debug_dir, debug_prefix, "step03b_filtered_vectors", h_bin, is_lines=True, lines=rowboxes + colboxes)
|
|
|
|
|
|
- # 4. 几何延长线段 (使用自定义的大阈值函数)
|
|
|
- # alpha=w//20 动态阈值,或者固定给一个较大的值如 100
|
|
|
- # 假设分辨率较大,100px的断连是需要被修复的
|
|
|
- dynamic_alpha = max(50, int(min(w, h) * 0.05)) # 5% of min dimension
|
|
|
+ # 4. 几何延长线段
|
|
|
+ dynamic_alpha = max(50, int(min(w, h) * 0.05))
|
|
|
logger.info(f"Using dynamic alpha for line extension: {dynamic_alpha}")
|
|
|
|
|
|
- rowboxes, colboxes = custom_final_adjust_lines(rowboxes, colboxes, alpha=dynamic_alpha)
|
|
|
+ rowboxes, colboxes = GridRecovery._extend_lines(rowboxes, colboxes, alpha=dynamic_alpha)
|
|
|
|
|
|
# Step 4 Debug
|
|
|
- save_debug_image("step04_extended_vectors", h_bin, is_lines=True, lines=rowboxes + colboxes)
|
|
|
+ GridRecovery._save_debug_image(debug_dir, debug_prefix, "step04_extended_vectors", h_bin, is_lines=True, lines=rowboxes + colboxes)
|
|
|
|
|
|
# 5. 重绘纯净Mask
|
|
|
line_mask = np.zeros((h, w), dtype=np.uint8)
|
|
|
@@ -387,20 +370,20 @@ class GridRecovery:
|
|
|
line_mask = draw_lines(line_mask, rowboxes + colboxes, color=255, lineW=2)
|
|
|
|
|
|
# Step 5a Debug (Before Dilation)
|
|
|
- save_debug_image("step05a_rerasterized", line_mask)
|
|
|
+ GridRecovery._save_debug_image(debug_dir, debug_prefix, "step05a_rerasterized", line_mask)
|
|
|
|
|
|
# 增强: 全局微膨胀
|
|
|
kernel_dilate = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
|
|
|
line_mask = cv2.dilate(line_mask, kernel_dilate, iterations=1)
|
|
|
|
|
|
# Step 5b Debug (After Dilation)
|
|
|
- save_debug_image("step05b_dilated", line_mask)
|
|
|
+ GridRecovery._save_debug_image(debug_dir, debug_prefix, "step05b_dilated", line_mask)
|
|
|
|
|
|
# 6. 反转图像
|
|
|
inv_grid = cv2.bitwise_not(line_mask)
|
|
|
|
|
|
# Step 6 Debug (Input to ConnectedComponents)
|
|
|
- save_debug_image("step06_inverted_input", inv_grid)
|
|
|
+ GridRecovery._save_debug_image(debug_dir, debug_prefix, "step06_inverted_input", inv_grid)
|
|
|
|
|
|
# 7. 连通域
|
|
|
num_labels, labels, stats, centroids = cv2.connectedComponentsWithStats(inv_grid, connectivity=8)
|
|
|
@@ -511,42 +494,9 @@ class GridRecovery:
|
|
|
for line in colboxes
|
|
|
]
|
|
|
|
|
|
- # 🆕 过滤线条:只保留与existing_bboxes边界对齐的线条
|
|
|
- # 因为OCR补偿只针对与现有单元格相邻的空单元格
|
|
|
- def filter_lines_by_bboxes(lines, bboxes, is_horizontal, tolerance=5.0):
|
|
|
- """过滤线条,只保留与bboxes边界对齐的线条"""
|
|
|
- if not bboxes:
|
|
|
- return lines
|
|
|
-
|
|
|
- # 提取所有bbox的边界坐标
|
|
|
- if is_horizontal:
|
|
|
- # 横线:检查是否与bbox的y1或y2对齐
|
|
|
- bbox_coords = set()
|
|
|
- for bbox in bboxes:
|
|
|
- bbox_coords.add(bbox[1]) # y1
|
|
|
- bbox_coords.add(bbox[3]) # y2
|
|
|
- else:
|
|
|
- # 竖线:检查是否与bbox的x1或x2对齐
|
|
|
- bbox_coords = set()
|
|
|
- for bbox in bboxes:
|
|
|
- bbox_coords.add(bbox[0]) # x1
|
|
|
- bbox_coords.add(bbox[2]) # x2
|
|
|
-
|
|
|
- # 过滤线条
|
|
|
- filtered_lines = []
|
|
|
- for line in lines:
|
|
|
- line_coord = (line[1] + line[3]) / 2 if is_horizontal else (line[0] + line[2]) / 2
|
|
|
-
|
|
|
- # 检查是否与任意bbox边界对齐
|
|
|
- is_aligned = any(abs(line_coord - coord) < tolerance for coord in bbox_coords)
|
|
|
- if is_aligned:
|
|
|
- filtered_lines.append(line)
|
|
|
-
|
|
|
- return filtered_lines
|
|
|
-
|
|
|
- # 过滤掉与existing_bboxes不对齐的干扰线条
|
|
|
- rowboxes_filtered = filter_lines_by_bboxes(rowboxes_orig, bboxes, is_horizontal=True)
|
|
|
- colboxes_filtered = filter_lines_by_bboxes(colboxes_orig, bboxes, is_horizontal=False)
|
|
|
+ # 过滤线条:只保留与existing_bboxes边界对齐的线条
|
|
|
+ rowboxes_filtered = GridRecovery._filter_lines_by_bboxes(rowboxes_orig, bboxes, is_horizontal=True)
|
|
|
+ colboxes_filtered = GridRecovery._filter_lines_by_bboxes(colboxes_orig, bboxes, is_horizontal=False)
|
|
|
|
|
|
logger.debug(
|
|
|
f"🔍 线条过滤: 横线 {len(rowboxes_orig)}→{len(rowboxes_filtered)}, "
|