ocr_platform/ocr_utils/module_debug_viz.py

"""
模块级 Debug 可视化（Layout / OCR）

用于 ``{output_dir}/debug/{subdir}/`` 下基于 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

# 各模块 debug_options 默认落盘根目录（相对 pipeline output_dir）
MODULE_DEBUG_ROOT = "debug"


def _json_default(o: Any):
    """json.dumps 的兜底序列化：处理 numpy 标量/数组（如 OCR confidence 的 float32）。"""
    if isinstance(o, np.generic):
        return o.item()
    if isinstance(o, np.ndarray):
        return o.tolist()
    if isinstance(o, (set, tuple)):
        return list(o)
    raise TypeError(f"Object of type {o.__class__.__name__} is not JSON serializable")


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 = {
    '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),
    'chart': (255, 255, 0),   # 青色（BGR 下 B=255,G=255,R=0）
    # 注意：OpenCV 为 BGR，(0,255,255) 在屏幕上呈黄色（与 title 相同），勿用于 seal
    'seal': (0, 140, 255),    # 亮橙，与红 table / 黄 title / 蓝 text 均易区分
    'abandon': (128, 128, 128),
}

# seal 常与 table 重叠：加粗线宽 + 黑色外描边
LAYOUT_HIGHLIGHT_CATEGORIES = frozenset({'seal'})
LAYOUT_HIGHLIGHT_LINE_THICKNESS = 4
LAYOUT_HIGHLIGHT_OUTLINE_BGR = (0, 0, 0)
LAYOUT_DEFAULT_LINE_THICKNESS = 2

# OCR 框线宽 (不受配色统一影响)
OCR_BOX_LINE_THICKNESS = 2
OCR_BOX_DASH_LENGTH = 8
OCR_BOX_DASH_GAP = 6


def _ocr_box_color_bgr() -> tuple:
    """亮蓝 OCR 框 (BGR)，派生自 VisualizationUtils.COLOR_MAP['ocr_box']。"""
    from ocr_utils.visualization_utils import VisualizationUtils
    return VisualizationUtils.rgb_to_bgr(VisualizationUtils.COLOR_MAP['ocr_box'])


def _seal_ocr_box_color_bgr() -> tuple:
    """印章 OCR 框 (BGR)，派生自 VisualizationUtils.COLOR_MAP['seal_ocr_box']。"""
    from ocr_utils.visualization_utils import VisualizationUtils
    return VisualizationUtils.rgb_to_bgr(VisualizationUtils.COLOR_MAP['seal_ocr_box'])


def ocr_box_color_rgb() -> tuple:
    """OCR 亮蓝 (RGB)，供 PIL / Plotly 使用。"""
    from ocr_utils.visualization_utils import VisualizationUtils
    return VisualizationUtils.COLOR_MAP['ocr_box']


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))
        thickness = (
            LAYOUT_HIGHLIGHT_LINE_THICKNESS
            if category in LAYOUT_HIGHLIGHT_CATEGORIES
            else LAYOUT_DEFAULT_LINE_THICKNESS
        )
        x1, y1, x2, y2 = int(bbox[0]), int(bbox[1]), int(bbox[2]), int(bbox[3])
        if category in LAYOUT_HIGHLIGHT_CATEGORIES:
            cv2.rectangle(
                vis, (x1, y1), (x2, y2),
                LAYOUT_HIGHLIGHT_OUTLINE_BGR,
                thickness + 2,
            )
        cv2.rectangle(vis, (x1, y1), (x2, y2), color, thickness)
        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.5 if category in LAYOUT_HIGHLIGHT_CATEGORIES else 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_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）；无文字用虚线框。
    
    span 可带 category='seal' 使用印章专用亮橙色，否则使用亮蓝。
    """
    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:
            text_raw = str(span.get('text', '') or '').strip()
            color = _seal_ocr_box_color_bgr() if span.get('category') == 'seal' else _ocr_box_color_bgr()
            _draw_span_outline(
                vis,
                pts,
                color,
                OCR_BOX_LINE_THICKNESS,
                dashed=not text_raw,
            )
        text = str(span.get('text', '')).strip()[:max_label_chars]
        if text and pts is not None:
            color = _seal_ocr_box_color_bgr() if span.get('category') == 'seal' else _ocr_box_color_bgr()
            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, color, 1, cv2.LINE_AA,
            )
    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 = resolve_module_debug_dir(output_dir, subdir)
        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, default=_json_default),
                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 = resolve_module_debug_dir(output_dir, subdir)
        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, default=_json_default),
                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