|
@@ -0,0 +1,213 @@
|
|
|
|
|
+"""
|
|
|
|
|
+模块级 Debug 可视化(Layout / OCR)
|
|
|
|
|
+
|
|
|
|
|
+用于 debug_comparison/ 下基于 inference_image 的调试图;
|
|
|
|
|
+用户审计图由 VisualizationUtils + original_image 负责,不在此模块。
|
|
|
|
|
+"""
|
|
|
|
|
+from __future__ import annotations
|
|
|
|
|
+
|
|
|
|
|
+import json
|
|
|
|
|
+from pathlib import Path
|
|
|
|
|
+from typing import Any, Dict, List, Optional, Union
|
|
|
|
|
+
|
|
|
|
|
+import cv2
|
|
|
|
|
+import numpy as np
|
|
|
|
|
+from loguru import logger
|
|
|
|
|
+from PIL import Image
|
|
|
|
|
+
|
|
|
|
|
+LAYOUT_CATEGORY_COLORS_BGR = {
|
|
|
|
|
+ 'table_body': (0, 0, 255),
|
|
|
|
|
+ 'table_caption': (0, 0, 200),
|
|
|
|
|
+ 'table_footnote': (0, 0, 150),
|
|
|
|
|
+ 'text': (255, 0, 0),
|
|
|
|
|
+ 'title': (0, 255, 255),
|
|
|
|
|
+ 'header': (255, 0, 255),
|
|
|
|
|
+ 'footer': (0, 165, 255),
|
|
|
|
|
+ 'image_body': (0, 255, 0),
|
|
|
|
|
+ 'image_caption': (0, 200, 0),
|
|
|
|
|
+ 'image_footnote': (0, 150, 0),
|
|
|
|
|
+ 'abandon': (128, 128, 128),
|
|
|
|
|
+}
|
|
|
|
|
+
|
|
|
|
|
+OCR_BOX_COLOR_BGR = (0, 255, 255)
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
|
|
+def _to_bgr(image: Union[np.ndarray, Image.Image]) -> np.ndarray:
|
|
|
|
|
+ if isinstance(image, Image.Image):
|
|
|
|
|
+ arr = np.array(image)
|
|
|
|
|
+ else:
|
|
|
|
|
+ arr = image.copy()
|
|
|
|
|
+ if arr.ndim == 2:
|
|
|
|
|
+ return cv2.cvtColor(arr, cv2.COLOR_GRAY2BGR)
|
|
|
|
|
+ if arr.shape[2] == 3:
|
|
|
|
|
+ return cv2.cvtColor(arr, cv2.COLOR_RGB2BGR)
|
|
|
|
|
+ return arr
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
|
|
+def draw_layout_boxes_cv2(
|
|
|
|
|
+ image: Union[np.ndarray, Image.Image],
|
|
|
|
|
+ layout_results: List[Dict[str, Any]],
|
|
|
|
|
+) -> np.ndarray:
|
|
|
|
|
+ """在 BGR 图像上绘制 layout 检测框,返回新图像。"""
|
|
|
|
|
+ vis = _to_bgr(image)
|
|
|
|
|
+ for result in layout_results:
|
|
|
|
|
+ bbox = result.get('bbox', [])
|
|
|
|
|
+ if not bbox or len(bbox) < 4:
|
|
|
|
|
+ continue
|
|
|
|
|
+ category = result.get('category', 'unknown')
|
|
|
|
|
+ color = LAYOUT_CATEGORY_COLORS_BGR.get(category, (128, 128, 128))
|
|
|
|
|
+ x1, y1, x2, y2 = int(bbox[0]), int(bbox[1]), int(bbox[2]), int(bbox[3])
|
|
|
|
|
+ cv2.rectangle(vis, (x1, y1), (x2, y2), color, 2)
|
|
|
|
|
+ label = category
|
|
|
|
|
+ confidence = result.get('confidence', result.get('score', 0))
|
|
|
|
|
+ if confidence:
|
|
|
|
|
+ label += f":{float(confidence):.2f}"
|
|
|
|
|
+ font = cv2.FONT_HERSHEY_SIMPLEX
|
|
|
|
|
+ font_scale = 0.4
|
|
|
|
|
+ text_thickness = 1
|
|
|
|
|
+ (text_width, text_height), baseline = cv2.getTextSize(
|
|
|
|
|
+ label, font, font_scale, text_thickness
|
|
|
|
|
+ )
|
|
|
|
|
+ text_y = max(y1 - baseline - 1, text_height + baseline)
|
|
|
|
|
+ cv2.rectangle(
|
|
|
|
|
+ vis,
|
|
|
|
|
+ (x1, text_y - text_height - baseline - 2),
|
|
|
|
|
+ (x1 + text_width, text_y),
|
|
|
|
|
+ color,
|
|
|
|
|
+ -1,
|
|
|
|
|
+ )
|
|
|
|
|
+ cv2.putText(
|
|
|
|
|
+ vis, label, (x1, text_y - baseline - 1),
|
|
|
|
|
+ font, font_scale, (255, 255, 255), text_thickness,
|
|
|
|
|
+ )
|
|
|
|
|
+ return vis
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
|
|
+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)。"""
|
|
|
|
|
+ vis = _to_bgr(image)
|
|
|
|
|
+ for span in spans:
|
|
|
|
|
+ poly = span.get('poly')
|
|
|
|
|
+ bbox = span.get('bbox', [])
|
|
|
|
|
+ pts = None
|
|
|
|
|
+ if poly and len(poly) >= 4:
|
|
|
|
|
+ pts = np.array(poly, dtype=np.int32).reshape(-1, 2)
|
|
|
|
|
+ elif bbox and len(bbox) >= 4:
|
|
|
|
|
+ x0, y0, x1, y1 = map(int, bbox[:4])
|
|
|
|
|
+ pts = np.array(
|
|
|
|
|
+ [[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, 1)
|
|
|
|
|
+ text = str(span.get('text', ''))[:max_label_chars]
|
|
|
|
|
+ if text and pts is not None:
|
|
|
|
|
+ x, y = int(pts[0][0]), int(pts[0][1])
|
|
|
|
|
+ cv2.putText(
|
|
|
|
|
+ vis, text, (x, max(y - 2, 10)),
|
|
|
|
|
+ cv2.FONT_HERSHEY_SIMPLEX, 0.35, OCR_BOX_COLOR_BGR, 1,
|
|
|
|
|
+ )
|
|
|
|
|
+ return vis
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
|
|
+def save_layout_debug(
|
|
|
|
|
+ image: Union[np.ndarray, Image.Image],
|
|
|
|
|
+ layout_results: List[Dict[str, Any]],
|
|
|
|
|
+ output_dir: Union[str, Path],
|
|
|
|
|
+ page_name: str,
|
|
|
|
|
+ *,
|
|
|
|
|
+ suffix: str = 'raw',
|
|
|
|
|
+ subdir: str = 'layout_detection',
|
|
|
|
|
+ image_format: str = 'jpg',
|
|
|
|
|
+ save_json: bool = True,
|
|
|
|
|
+) -> Optional[Dict[str, str]]:
|
|
|
|
|
+ """保存 layout 模块 debug 图与 JSON。"""
|
|
|
|
|
+ if not layout_results or not output_dir:
|
|
|
|
|
+ return None
|
|
|
|
|
+ try:
|
|
|
|
|
+ fmt = (image_format or 'jpg').lstrip('.')
|
|
|
|
|
+ debug_dir = Path(output_dir) / 'debug_comparison' / subdir
|
|
|
|
|
+ debug_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
|
+ vis = draw_layout_boxes_cv2(image, layout_results)
|
|
|
|
|
+ img_path = debug_dir / f'{page_name}_layout_{suffix}.{fmt}'
|
|
|
|
|
+ cv2.imwrite(str(img_path), vis)
|
|
|
|
|
+ paths: Dict[str, str] = {'image': str(img_path)}
|
|
|
|
|
+ logger.info(f"Saved layout detection image ({suffix}): {img_path}")
|
|
|
|
|
+ if save_json:
|
|
|
|
|
+ json_data = {
|
|
|
|
|
+ 'page_name': page_name,
|
|
|
|
|
+ 'suffix': suffix,
|
|
|
|
|
+ 'count': len(layout_results),
|
|
|
|
|
+ 'results': [
|
|
|
|
|
+ {
|
|
|
|
|
+ 'category': r.get('category'),
|
|
|
|
|
+ 'bbox': r.get('bbox'),
|
|
|
|
|
+ 'confidence': r.get('confidence', r.get('score', 0.0)),
|
|
|
|
|
+ }
|
|
|
|
|
+ for r in layout_results
|
|
|
|
|
+ ],
|
|
|
|
|
+ }
|
|
|
|
|
+ json_path = debug_dir / f'{page_name}_layout_{suffix}.json'
|
|
|
|
|
+ json_path.write_text(
|
|
|
|
|
+ json.dumps(json_data, ensure_ascii=False, indent=2),
|
|
|
|
|
+ encoding='utf-8',
|
|
|
|
|
+ )
|
|
|
|
|
+ paths['json'] = str(json_path)
|
|
|
|
|
+ logger.info(f"Saved layout detection JSON ({suffix}): {json_path}")
|
|
|
|
|
+ return paths
|
|
|
|
|
+ except Exception as e:
|
|
|
|
|
+ logger.warning(f"Failed to save layout debug ({suffix}): {e}")
|
|
|
|
|
+ return None
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
|
|
+def save_ocr_debug(
|
|
|
|
|
+ image: Union[np.ndarray, Image.Image],
|
|
|
|
|
+ spans: List[Dict[str, Any]],
|
|
|
|
|
+ output_dir: Union[str, Path],
|
|
|
|
|
+ page_name: str,
|
|
|
|
|
+ *,
|
|
|
|
|
+ subdir: str = 'ocr_recognition',
|
|
|
|
|
+ image_format: str = 'png',
|
|
|
|
|
+ save_json: bool = True,
|
|
|
|
|
+) -> Optional[Dict[str, str]]:
|
|
|
|
|
+ """保存 OCR 模块 debug 图与 JSON。"""
|
|
|
|
|
+ if not output_dir:
|
|
|
|
|
+ return None
|
|
|
|
|
+ try:
|
|
|
|
|
+ fmt = (image_format or 'png').lstrip('.')
|
|
|
|
|
+ debug_dir = Path(output_dir) / 'debug_comparison' / subdir
|
|
|
|
|
+ debug_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
|
+ vis = draw_ocr_spans_cv2(image, spans or [])
|
|
|
|
|
+ img_path = debug_dir / f'{page_name}_ocr_spans.{fmt}'
|
|
|
|
|
+ cv2.imwrite(str(img_path), vis)
|
|
|
|
|
+ paths: Dict[str, str] = {'image': str(img_path)}
|
|
|
|
|
+ logger.info(f"Saved OCR debug image: {img_path}")
|
|
|
|
|
+ if save_json:
|
|
|
|
|
+ json_data = {
|
|
|
|
|
+ 'page_name': page_name,
|
|
|
|
|
+ 'count': len(spans or []),
|
|
|
|
|
+ 'spans': [
|
|
|
|
|
+ {
|
|
|
|
|
+ 'bbox': s.get('bbox'),
|
|
|
|
|
+ 'poly': s.get('poly'),
|
|
|
|
|
+ 'text': s.get('text'),
|
|
|
|
|
+ 'confidence': s.get('confidence'),
|
|
|
|
|
+ }
|
|
|
|
|
+ for s in (spans or [])
|
|
|
|
|
+ ],
|
|
|
|
|
+ }
|
|
|
|
|
+ json_path = debug_dir / f'{page_name}_ocr_spans.json'
|
|
|
|
|
+ json_path.write_text(
|
|
|
|
|
+ json.dumps(json_data, ensure_ascii=False, indent=2),
|
|
|
|
|
+ encoding='utf-8',
|
|
|
|
|
+ )
|
|
|
|
|
+ paths['json'] = str(json_path)
|
|
|
|
|
+ logger.info(f"Saved OCR debug JSON: {json_path}")
|
|
|
|
|
+ return paths
|
|
|
|
|
+ except Exception as e:
|
|
|
|
|
+ logger.warning(f"Failed to save OCR debug: {e}")
|
|
|
|
|
+ return None
|