|
@@ -48,6 +48,8 @@ LAYOUT_CATEGORY_COLORS_BGR = {
|
|
|
# 亮蓝(BGR),在白底/浅灰流水上比黄色更易辨认;与 layout 红色框区分
|
|
# 亮蓝(BGR),在白底/浅灰流水上比黄色更易辨认;与 layout 红色框区分
|
|
|
OCR_BOX_COLOR_BGR = (255, 0, 0)
|
|
OCR_BOX_COLOR_BGR = (255, 0, 0)
|
|
|
OCR_BOX_LINE_THICKNESS = 2
|
|
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:
|
|
def _to_bgr(image: Union[np.ndarray, Image.Image]) -> np.ndarray:
|
|
@@ -101,13 +103,78 @@ def draw_layout_boxes_cv2(
|
|
|
return vis
|
|
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(
|
|
def draw_ocr_spans_cv2(
|
|
|
image: Union[np.ndarray, Image.Image],
|
|
image: Union[np.ndarray, Image.Image],
|
|
|
spans: List[Dict[str, Any]],
|
|
spans: List[Dict[str, Any]],
|
|
|
*,
|
|
*,
|
|
|
max_label_chars: int = 12,
|
|
max_label_chars: int = 12,
|
|
|
) -> np.ndarray:
|
|
) -> np.ndarray:
|
|
|
- """在 BGR 图像上绘制 OCR span(poly 或 bbox)。"""
|
|
|
|
|
|
|
+ """在 BGR 图像上绘制 OCR span(poly 或 bbox);无文字用虚线框。"""
|
|
|
vis = _to_bgr(image)
|
|
vis = _to_bgr(image)
|
|
|
for span in spans:
|
|
for span in spans:
|
|
|
poly = span.get('poly')
|
|
poly = span.get('poly')
|
|
@@ -121,10 +188,15 @@ def draw_ocr_spans_cv2(
|
|
|
[[x0, y0], [x1, y0], [x1, y1], [x0, y1]], dtype=np.int32
|
|
[[x0, y0], [x1, y0], [x1, y1], [x0, y1]], dtype=np.int32
|
|
|
)
|
|
)
|
|
|
if pts is not None:
|
|
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:
|
|
if text and pts is not None:
|
|
|
x, y = int(pts[0][0]), int(pts[0][1])
|
|
x, y = int(pts[0][0]), int(pts[0][1])
|
|
|
cv2.putText(
|
|
cv2.putText(
|