|
@@ -1,7 +1,7 @@
|
|
|
"""
|
|
"""
|
|
|
模块级 Debug 可视化(Layout / OCR)
|
|
模块级 Debug 可视化(Layout / OCR)
|
|
|
|
|
|
|
|
-用于 debug_comparison/ 下基于 inference_image 的调试图;
|
|
|
|
|
|
|
+用于 ``{output_dir}/debug/{subdir}/`` 下基于 inference_image 的调试图;
|
|
|
用户审计图由 VisualizationUtils + original_image 负责,不在此模块。
|
|
用户审计图由 VisualizationUtils + original_image 负责,不在此模块。
|
|
|
"""
|
|
"""
|
|
|
from __future__ import annotations
|
|
from __future__ import annotations
|
|
@@ -15,6 +15,22 @@ import numpy as np
|
|
|
from loguru import logger
|
|
from loguru import logger
|
|
|
from PIL import Image
|
|
from PIL import Image
|
|
|
|
|
|
|
|
|
|
+# 各模块 debug_options 默认落盘根目录(相对 pipeline output_dir)
|
|
|
|
|
+MODULE_DEBUG_ROOT = "debug"
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
|
|
+def resolve_module_debug_dir(
|
|
|
|
|
+ output_dir: Union[str, Path],
|
|
|
|
|
+ subdir: str,
|
|
|
|
|
+ *,
|
|
|
|
|
+ debug_root: str = MODULE_DEBUG_ROOT,
|
|
|
|
|
+) -> Path:
|
|
|
|
|
+ """``{output_dir}/{debug_root}/{subdir}/``,目录不存在则创建。"""
|
|
|
|
|
+ debug_dir = Path(output_dir) / debug_root / subdir
|
|
|
|
|
+ debug_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
|
+ return debug_dir
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
LAYOUT_CATEGORY_COLORS_BGR = {
|
|
LAYOUT_CATEGORY_COLORS_BGR = {
|
|
|
'table_body': (0, 0, 255),
|
|
'table_body': (0, 0, 255),
|
|
|
'table_caption': (0, 0, 200),
|
|
'table_caption': (0, 0, 200),
|
|
@@ -29,7 +45,9 @@ LAYOUT_CATEGORY_COLORS_BGR = {
|
|
|
'abandon': (128, 128, 128),
|
|
'abandon': (128, 128, 128),
|
|
|
}
|
|
}
|
|
|
|
|
|
|
|
-OCR_BOX_COLOR_BGR = (0, 255, 255)
|
|
|
|
|
|
|
+# 亮蓝(BGR),在白底/浅灰流水上比黄色更易辨认;与 layout 红色框区分
|
|
|
|
|
+OCR_BOX_COLOR_BGR = (255, 0, 0)
|
|
|
|
|
+OCR_BOX_LINE_THICKNESS = 2
|
|
|
|
|
|
|
|
|
|
|
|
|
def _to_bgr(image: Union[np.ndarray, Image.Image]) -> np.ndarray:
|
|
def _to_bgr(image: Union[np.ndarray, Image.Image]) -> np.ndarray:
|
|
@@ -103,13 +121,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, 1)
|
|
|
|
|
|
|
+ cv2.polylines(
|
|
|
|
|
+ vis, [pts], True, OCR_BOX_COLOR_BGR, OCR_BOX_LINE_THICKNESS
|
|
|
|
|
+ )
|
|
|
text = str(span.get('text', ''))[:max_label_chars]
|
|
text = str(span.get('text', ''))[: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(
|
|
|
vis, text, (x, max(y - 2, 10)),
|
|
vis, text, (x, max(y - 2, 10)),
|
|
|
- cv2.FONT_HERSHEY_SIMPLEX, 0.35, OCR_BOX_COLOR_BGR, 1,
|
|
|
|
|
|
|
+ cv2.FONT_HERSHEY_SIMPLEX, 0.35, OCR_BOX_COLOR_BGR, 1, cv2.LINE_AA,
|
|
|
)
|
|
)
|
|
|
return vis
|
|
return vis
|
|
|
|
|
|
|
@@ -130,8 +150,7 @@ def save_layout_debug(
|
|
|
return None
|
|
return None
|
|
|
try:
|
|
try:
|
|
|
fmt = (image_format or 'jpg').lstrip('.')
|
|
fmt = (image_format or 'jpg').lstrip('.')
|
|
|
- debug_dir = Path(output_dir) / 'debug_comparison' / subdir
|
|
|
|
|
- debug_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
|
|
|
+ debug_dir = resolve_module_debug_dir(output_dir, subdir)
|
|
|
vis = draw_layout_boxes_cv2(image, layout_results)
|
|
vis = draw_layout_boxes_cv2(image, layout_results)
|
|
|
img_path = debug_dir / f'{page_name}_layout_{suffix}.{fmt}'
|
|
img_path = debug_dir / f'{page_name}_layout_{suffix}.{fmt}'
|
|
|
cv2.imwrite(str(img_path), vis)
|
|
cv2.imwrite(str(img_path), vis)
|
|
@@ -179,8 +198,7 @@ def save_ocr_debug(
|
|
|
return None
|
|
return None
|
|
|
try:
|
|
try:
|
|
|
fmt = (image_format or 'png').lstrip('.')
|
|
fmt = (image_format or 'png').lstrip('.')
|
|
|
- debug_dir = Path(output_dir) / 'debug_comparison' / subdir
|
|
|
|
|
- debug_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
|
|
|
+ debug_dir = resolve_module_debug_dir(output_dir, subdir)
|
|
|
vis = draw_ocr_spans_cv2(image, spans or [])
|
|
vis = draw_ocr_spans_cv2(image, spans or [])
|
|
|
img_path = debug_dir / f'{page_name}_ocr_spans.{fmt}'
|
|
img_path = debug_dir / f'{page_name}_ocr_spans.{fmt}'
|
|
|
cv2.imwrite(str(img_path), vis)
|
|
cv2.imwrite(str(img_path), vis)
|