Kaynağa Gözat

feat(添加虚线段绘制功能): 在module_debug_viz.py中新增虚线段绘制函数,支持在OCR span无文本时使用虚线框显示,提升可视化效果和调试灵活性。

zhch158_admin 4 gün önce
ebeveyn
işleme
fdef502446
1 değiştirilmiş dosya ile 76 ekleme ve 4 silme
  1. 76 4
      ocr_utils/module_debug_viz.py

+ 76 - 4
ocr_utils/module_debug_viz.py

@@ -48,6 +48,8 @@ LAYOUT_CATEGORY_COLORS_BGR = {
 # 亮蓝(BGR),在白底/浅灰流水上比黄色更易辨认;与 layout 红色框区分
 OCR_BOX_COLOR_BGR = (255, 0, 0)
 OCR_BOX_LINE_THICKNESS = 2
+OCR_BOX_DASH_LENGTH = 8
+OCR_BOX_DASH_GAP = 6
 
 
 def _to_bgr(image: Union[np.ndarray, Image.Image]) -> np.ndarray:
@@ -101,13 +103,78 @@ def draw_layout_boxes_cv2(
     return vis
 
 
+def _draw_dashed_segment(
+    vis: np.ndarray,
+    p1: np.ndarray,
+    p2: np.ndarray,
+    color: tuple,
+    thickness: int,
+    *,
+    dash_length: int = OCR_BOX_DASH_LENGTH,
+    gap_length: int = OCR_BOX_DASH_GAP,
+) -> None:
+    """在 p1→p2 上绘制虚线段。"""
+    start = p1.astype(np.float64)
+    end = p2.astype(np.float64)
+    vec = end - start
+    length = float(np.linalg.norm(vec))
+    if length < 1e-6:
+        return
+    direction = vec / length
+    pos = 0.0
+    draw = True
+    while pos < length:
+        seg = float(dash_length if draw else gap_length)
+        seg_end = min(pos + seg, length)
+        if draw:
+            s = (start + direction * pos).astype(np.int32)
+            e = (start + direction * seg_end).astype(np.int32)
+            cv2.line(
+                vis,
+                (int(s[0]), int(s[1])),
+                (int(e[0]), int(e[1])),
+                color,
+                thickness,
+                cv2.LINE_AA,
+            )
+        pos = seg_end
+        draw = not draw
+
+
+def _draw_span_outline(
+    vis: np.ndarray,
+    pts: np.ndarray,
+    color: tuple,
+    thickness: int,
+    *,
+    dashed: bool,
+) -> None:
+    n = len(pts)
+    if n < 2:
+        return
+    for i in range(n):
+        p1 = pts[i]
+        p2 = pts[(i + 1) % n]
+        if dashed:
+            _draw_dashed_segment(vis, p1, p2, color, thickness)
+        else:
+            cv2.line(
+                vis,
+                (int(p1[0]), int(p1[1])),
+                (int(p2[0]), int(p2[1])),
+                color,
+                thickness,
+                cv2.LINE_AA,
+            )
+
+
 def draw_ocr_spans_cv2(
     image: Union[np.ndarray, Image.Image],
     spans: List[Dict[str, Any]],
     *,
     max_label_chars: int = 12,
 ) -> np.ndarray:
-    """在 BGR 图像上绘制 OCR span(poly 或 bbox)。"""
+    """在 BGR 图像上绘制 OCR span(poly 或 bbox);无文字用虚线框。"""
     vis = _to_bgr(image)
     for span in spans:
         poly = span.get('poly')
@@ -121,10 +188,15 @@ def draw_ocr_spans_cv2(
                 [[x0, y0], [x1, y0], [x1, y1], [x0, y1]], dtype=np.int32
             )
         if pts is not None:
-            cv2.polylines(
-                vis, [pts], True, OCR_BOX_COLOR_BGR, OCR_BOX_LINE_THICKNESS
+            text_raw = str(span.get('text', '') or '').strip()
+            _draw_span_outline(
+                vis,
+                pts,
+                OCR_BOX_COLOR_BGR,
+                OCR_BOX_LINE_THICKNESS,
+                dashed=not text_raw,
             )
-        text = str(span.get('text', ''))[:max_label_chars]
+        text = str(span.get('text', '')).strip()[:max_label_chars]
         if text and pts is not None:
             x, y = int(pts[0][0]), int(pts[0][1])
             cv2.putText(