wechat-screenshot-vision-algorithm 0.1.0__py3-none-any.whl

wechat_screenshot_vision_algorithm/ocr/text_ocr_adapter.py

@@ -0,0 +1,627 @@
+"""Processor-side text-only OCR adapter (d2-4).
+
+Day 2 scope (this module):
+
+    Thin, engine-agnostic wrapper around PaddleOCR's text detection +
+    recognition. Feeds a single PNG (plus its ``Screenshot`` metadata)
+    through one pass and returns a flat list of **raw** ``TextBlock``
+    entries — ``(text, bbox_xyxy, confidence)`` — in reading order,
+    with bounding boxes mapped back to the **original** image
+    coordinate system.
+
+    Explicit non-goals (by design, per ADR v0.3 + §2.5.2 phase split):
+
+      - **Text normalization** (NFC / full-width ↔ half-width) is
+        deferred to the downstream text pipeline (d5
+        ``resume_text_merger`` for resume pages, or the nickname OCR
+        minimal pipeline in d3). Running it here would duplicate
+        ``lite_text_normalizer.normalize_business_text`` and cross
+        the scripts/venv ↔ backend/venv boundary for a function
+        whose output shape is a simple str. The raw OCR output must
+        survive untouched to the phase-4 / phase-5 consumers.
+      - **Nickname/speaker attribution**: ADR §2.3 phase 2
+        (NicknameBoundaryService) reads from this adapter's output;
+        it is NOT the adapter's job.
+      - **Resume page stitching / n-gram dedup**: d5
+        ``resume_text_merger`` consumes a list of ``OcrPageResult``
+        objects across one ``resume_group_id``; this module only
+        returns the per-page shape.
+      - **PP-Structure / layout analysis**: ADR v0.3 §4.3.0 permanent
+        project boundary — this adapter is **text-only** by contract;
+        PaddleOCR's ``PaddleOCR.predict()`` (or 2.x ``.ocr()``) is
+        the sole entry point used.
+      - **Head-avatar pHash / geometry**: ADR v0.3 permanently
+        discarded; the adapter does not expose any head-specific
+        API.
+
+    Scale policy (ADR §6.4.4):
+
+        The collector stamps ``screenshots[*].ocr_scale_hint`` at
+        capture time — ``1600`` for ``chat_message`` type,
+        ``1280`` for ``resume_detail`` type. The adapter treats this
+        as an **upper bound on the long-edge pre-OCR**: if the raw
+        image's long-edge exceeds the hint, it's downscaled with
+        ``cv2.INTER_AREA``; otherwise the image is passed through
+        unresized. Detection bounding boxes coming back from PaddleOCR
+        are in the *resized* coordinate space and are rescaled to
+        original coords before leaving the adapter so that downstream
+        consumers (e.g. ``side_hint_ratio`` cross-check in phase 2)
+        can reason in the same frame as the metadata.
+
+Engine indirection:
+
+    ``OcrEngine`` is a Protocol with a single
+    ``detect_and_recognize(bgr) -> list[RawOcrItem]`` method. The
+    real implementation ``PaddleOcrEngine`` lazy-imports PaddleOCR
+    so Windows collector venvs and CI containers without PaddlePaddle
+    installed can still import this module (they'd just fail to
+    construct a real engine). Unit tests inject a ``FakeOcrEngine``
+    (``dummy_engine`` helper in the tests) that returns synthetic
+    items; that keeps d2-4 test runs offline and lets d2-5's mock
+    coverage target the adapter independently of Paddle's model
+    download / CPU startup cost.
+
+Output contract:
+
+    ``OcrPageResult`` carries enough information for the scanner /
+    CLI to record session-level OCR stats (wall time, block count,
+    char count) AND for phase 2 nickname detection to operate on the
+    bboxes. It is a plain dataclass (``asdict``-safe) so JSON reports
+    stay trivial.
+
+References:
+    OCR ADR §2.5.2 phase 1 Preprocess / phase 2 NicknameBoundary
+    OCR ADR §6.4.4 ACTION_TYPE_MAP (ocr_scale_hint canonical values)
+    OCR ADR §4.3.0 project vision-layer boundary (v0.3)
+    backend/app/business/wx_match_business/paddle_ocr_subprocess.py
+        (API 冒烟 / 消费者预加载用的子进程批 OCR；与本 adapter 的 Paddle 3.x
+        ``predict`` 用法对齐，但 venv 边界独立)
+"""
+
+from __future__ import annotations
+
+import logging
+import time
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import TYPE_CHECKING, Callable, Optional, Protocol, runtime_checkable
+
+from collector_phone_android_contract import Screenshot
+
+if TYPE_CHECKING:  # avoid a hard numpy import at module-load time
+    import numpy as np
+
+LOG = logging.getLogger("wx_processor.text_ocr_adapter")
+
+#: Minimum confidence kept after PaddleOCR inference（PRD §6 块级 + 与 §11 侧车 0.7 口径一致）。
+#: 与 ``WxMatchSettings.wx_match_ocr_confidence_threshold``、``paddle_ocr_subprocess`` 对齐。
+DEFAULT_CONFIDENCE_THRESHOLD = 0.7
+
+#: Reading-order tie-break tolerance: two bboxes whose y1 differs by
+#: less than this many px are treated as "same row" and then sorted
+#: by x1. Intent: match how a human reads chat bubbles where two
+#: adjacent messages sometimes have slightly offset top pixels.
+READING_ORDER_ROW_TOLERANCE_PX = 6
+
+#: 相邻 OCR 块若垂直间距超过该值，视为新段落（PRD §6 段落号；与行号独立）。
+PARAGRAPH_BREAK_MIN_GAP_PX = 32
+
+
+# ============================================================================
+# Raw engine I/O
+# ============================================================================
+
+
+@dataclass(frozen=True)
+class RawOcrItem:
+    """One raw detection returned by an engine before filtering/mapping.
+
+    Attributes:
+        text: recognized text, unmodified.
+        bbox_quad: four (x, y) corners in the **resized** image's
+            coordinate space; PaddleOCR gives a clockwise quad
+            starting from top-left, but the adapter does NOT rely
+            on the ordering — it axis-aligns the quad via
+            ``quad_to_xyxy``.
+        confidence: scalar in ``[0, 1]``.
+    """
+
+    text: str
+    bbox_quad: tuple[tuple[float, float], ...]
+    confidence: float
+
+
+@runtime_checkable
+class OcrEngine(Protocol):
+    """Minimal engine contract: take BGR image, return raw items."""
+
+    name: str
+
+    def detect_and_recognize(self, image_bgr: "np.ndarray") -> list[RawOcrItem]:
+        """Run text detection + recognition.
+
+        Implementations MUST NOT resize the image themselves — the
+        adapter has already applied ``ocr_scale_hint`` when calling
+        this method, and internal resize would break the bbox
+        back-projection contract.
+        """
+        ...
+
+
+# ============================================================================
+# Adapter output shape
+# ============================================================================
+
+
+@dataclass
+class TextBlock:
+    """Filtered, axis-aligned block in **original** image coords.
+
+    ``bbox_xyxy`` is ``(x1, y1, x2, y2)`` with ``x1 <= x2`` and
+    ``y1 <= y2``. Coordinates are integers (rounded) because
+    downstream consumers always treat bboxes as pixel indices.
+
+    ``line_read_index``：本页阅读序下标（PRD §6 行/块序；1-based）。
+    ``paragraph_read_index``：段落序（按垂直大间断分段；1-based）。
+    """
+
+    text: str
+    bbox_xyxy: tuple[int, int, int, int]
+    confidence: float
+    line_read_index: int = 0
+    paragraph_read_index: int = 0
+
+
+@dataclass
+class OcrPageResult:
+    """Per-screenshot OCR output the scanner/CLI can stash directly.
+
+    ``text_blocks`` is already sorted in reading order; ``raw_full_text``
+    is the simple newline-join in that same order. ``raw_full_text``
+    is *raw* — callers must normalize themselves (see module docstring
+    "Explicit non-goals").
+    """
+
+    screenshot_id: str
+    scale_hint: int
+    original_size: tuple[int, int]
+    processed_size: tuple[int, int]
+    resized_applied: bool
+    text_blocks: list[TextBlock]
+    raw_full_text: str
+    engine_name: str
+    wall_ms: float
+    filtered_out_count: int = 0
+    confidence_threshold: float = DEFAULT_CONFIDENCE_THRESHOLD
+    #: If non-empty, the adapter hit a soft failure (e.g. image decoded
+    #: to ``None``) and returned an empty ``text_blocks``. The scanner
+    #: uses this to distinguish "really no text" from "could not read
+    #: the file" without making it a hard error.
+    soft_error: Optional[str] = None
+
+
+# ============================================================================
+# Errors
+# ============================================================================
+
+
+class TextOcrAdapterError(Exception):
+    """Raised when the adapter cannot produce a usable ``OcrPageResult``.
+
+    Reserved for programmer / environment errors (e.g. engine not
+    installed when expected, PNG file unreadable). Run-time OCR
+    "zero hits" or "low-confidence filtered all" scenarios return a
+    valid ``OcrPageResult`` with empty ``text_blocks`` instead, so
+    the scanner's terminal-status logic stays simple.
+    """
+
+    def __init__(self, message: str, *, error_code: str) -> None:
+        super().__init__(message)
+        self.error_code = error_code
+
+
+# ============================================================================
+# Geometry helpers (pure, easy to unit-test)
+# ============================================================================
+
+
+def resize_long_edge_to(
+    image_bgr: "np.ndarray",
+    long_edge: int,
+) -> tuple["np.ndarray", float]:
+    """Downscale to ``long_edge``; no-op if already smaller.
+
+    Returns ``(resized_image, scale_ratio)`` where ``scale_ratio`` is
+    ``new_long_edge / old_long_edge`` (``1.0`` if no resize). The
+    caller uses ``scale_ratio`` to map bboxes back to original
+    coordinates.
+    """
+    import cv2
+
+    h, w = image_bgr.shape[:2]
+    original_long = max(h, w)
+    if original_long <= long_edge:
+        return image_bgr, 1.0
+    ratio = long_edge / original_long
+    new_w = max(1, int(round(w * ratio)))
+    new_h = max(1, int(round(h * ratio)))
+    return (
+        cv2.resize(image_bgr, (new_w, new_h), interpolation=cv2.INTER_AREA),
+        ratio,
+    )
+
+
+def quad_to_xyxy(
+    quad: tuple[tuple[float, float], ...],
+) -> tuple[float, float, float, float]:
+    """Axis-align a 4-corner polygon to a ``(x1, y1, x2, y2)`` bbox.
+
+    Uses the quad's min/max envelope. This is deliberate: the
+    downstream ``NicknameBoundaryService`` reasons in axis-aligned
+    boxes (font-size ratio / y-gap / left-right alignment), so
+    preserving the rotated quad would pay a cost with no consumer.
+    """
+    xs = [p[0] for p in quad]
+    ys = [p[1] for p in quad]
+    return (min(xs), min(ys), max(xs), max(ys))
+
+
+def scale_bbox(
+    bbox_xyxy: tuple[float, float, float, float],
+    inv_ratio: float,
+) -> tuple[int, int, int, int]:
+    """Scale ``xyxy`` by ``inv_ratio`` and round to int pixels.
+
+    ``inv_ratio`` is ``1 / scale_ratio`` returned by
+    ``resize_long_edge_to`` — i.e. "resized → original" direction.
+    """
+    x1, y1, x2, y2 = bbox_xyxy
+    return (
+        max(0, int(round(x1 * inv_ratio))),
+        max(0, int(round(y1 * inv_ratio))),
+        max(0, int(round(x2 * inv_ratio))),
+        max(0, int(round(y2 * inv_ratio))),
+    )
+
+
+def assign_paragraph_read_index(blocks: list[TextBlock]) -> list[TextBlock]:
+    """为已排序块赋 ``paragraph_read_index``（垂直间隙过大则新开段落）。"""
+    if not blocks:
+        return blocks
+    out: list[TextBlock] = []
+    para = 1
+    prev_bottom: Optional[int] = None
+    for b in blocks:
+        y1 = b.bbox_xyxy[1]
+        if prev_bottom is not None and (y1 - prev_bottom) > PARAGRAPH_BREAK_MIN_GAP_PX:
+            para += 1
+        out.append(
+            TextBlock(
+                text=b.text,
+                bbox_xyxy=b.bbox_xyxy,
+                confidence=b.confidence,
+                line_read_index=b.line_read_index,
+                paragraph_read_index=para,
+            )
+        )
+        prev_bottom = b.bbox_xyxy[3]
+    return out
+
+
+def sort_reading_order(blocks: list[TextBlock]) -> list[TextBlock]:
+    """Top-to-bottom, left-to-right with a small row tolerance.
+
+    We bucket ``y1`` into row bands of
+    ``READING_ORDER_ROW_TOLERANCE_PX`` so that two bboxes sharing a
+    row but whose detected tops differ by a few pixels still sort
+    left-to-right within the row. Pure-Python, O(n log n), stable.
+    """
+
+    def key(b: TextBlock) -> tuple[int, int]:
+        y1 = b.bbox_xyxy[1]
+        row = y1 // READING_ORDER_ROW_TOLERANCE_PX
+        return (row, b.bbox_xyxy[0])
+
+    return sorted(blocks, key=key)
+
+
+# ============================================================================
+# Real PaddleOCR engine (lazy-loaded)
+# ============================================================================
+
+
+class PaddleOcrEngine:
+    """Real PaddleOCR 3.x engine, text-only.
+
+    Construction is expensive (model download on first run, ~30 s
+    startup even on cached models), so callers are expected to
+    instantiate **once** and reuse across screenshots within a
+    session.
+
+    ``use_server_model=False`` (default) matches the ADR v0.3
+    deployment: PP-OCRv5 *mobile* on CPU. The ``server`` variant is
+    reserved for GPU rollout later.
+
+    ``rec_batch`` defaults to 6 mirroring
+    ``app.core.config.Settings.paddle_ocr_rec_batch`` — changing it
+    only affects throughput, not the adapter's output shape.
+    """
+
+    name = "paddleocr_v3_text_only"
+
+    def __init__(
+        self,
+        *,
+        use_server_model: bool = False,
+        rec_batch: int = 6,
+        max_long_edge_pre_ocr: int = 2048,
+    ) -> None:
+        self._use_server_model = use_server_model
+        self._rec_batch = rec_batch
+        # Give PaddleOCR's internal det_limit a ceiling generous
+        # enough to never re-resize images the adapter has already
+        # sized to ``ocr_scale_hint``. The adapter is the single
+        # source of truth for scale policy.
+        self._max_long_edge_pre_ocr = max_long_edge_pre_ocr
+        self._ocr: object = None
+
+    def _ensure_engine(self) -> object:
+        if self._ocr is not None:
+            return self._ocr
+        from paddleocr import PaddleOCR  # noqa: PLC0415 — lazy
+        kwargs: dict[str, object] = {
+            "use_textline_orientation": True,
+            # Doc preprocessing (orientation classify + UVDoc unwarping)
+            # defaults to ON in PaddleOCR 3.x. UVDoc "rectifies" flat phone
+            # screenshots, warping bbox coordinates non-linearly along y
+            # (measured -25px top → +90px bottom on 720x1612), which breaks
+            # card/nickname spatial binding downstream. Screenshots are
+            # always flat and upright — disable both.
+            "use_doc_orientation_classify": False,
+            "use_doc_unwarping": False,
+            "lang": "ch",
+            "device": "cpu",
+            "text_recognition_batch_size": self._rec_batch,
+            "text_det_limit_type": "max",
+            "text_det_limit_side_len": self._max_long_edge_pre_ocr,
+        }
+        if not self._use_server_model:
+            kwargs["text_detection_model_name"] = "PP-OCRv5_mobile_det"
+            kwargs["text_recognition_model_name"] = "PP-OCRv5_mobile_rec"
+        self._ocr = PaddleOCR(**kwargs)
+        return self._ocr
+
+    def detect_and_recognize(self, image_bgr: "np.ndarray") -> list[RawOcrItem]:
+        ocr = self._ensure_engine()
+        results = list(ocr.predict([image_bgr]))  # type: ignore[attr-defined]
+        if not results:
+            return []
+        r = results[0]
+        texts: list[str] = list(r.get("rec_texts", []) or [])
+        scores: list[float] = list(r.get("rec_scores", []) or [])
+        polys_raw = r.get("rec_polys") or r.get("dt_polys") or []
+        items: list[RawOcrItem] = []
+        for i, text in enumerate(texts):
+            score = float(scores[i]) if i < len(scores) else 0.0
+            poly = polys_raw[i] if i < len(polys_raw) else None
+            if poly is None:
+                continue
+            quad = _normalize_poly(poly)
+            if quad is None:
+                continue
+            items.append(RawOcrItem(text=text, bbox_quad=quad, confidence=score))
+        return items
+
+
+def _normalize_poly(poly: object) -> Optional[tuple[tuple[float, float], ...]]:
+    """Coerce numpy / list-of-lists poly into a plain tuple of (x, y)."""
+    try:
+        iterable = list(poly)  # type: ignore[arg-type]
+    except TypeError:
+        return None
+    corners: list[tuple[float, float]] = []
+    for pt in iterable:
+        try:
+            x, y = float(pt[0]), float(pt[1])
+        except (TypeError, IndexError, ValueError):
+            return None
+        corners.append((x, y))
+    if len(corners) < 3:
+        return None
+    return tuple(corners)
+
+
+# ============================================================================
+# Adapter
+# ============================================================================
+
+
+class TextOcrAdapter:
+    """Process one screenshot through an injected OCR engine.
+
+    Single-session usage pattern (scanner/CLI):
+
+        engine = PaddleOcrEngine()
+        adapter = TextOcrAdapter(engine=engine)
+        for shot in metadata.screenshots:
+            png = resolver.resolve(session_dir, shot)
+            result = adapter.process_page(png, shot)
+            # hand result to phase-2 / phase-5 consumers
+    """
+
+    def __init__(
+        self,
+        engine: OcrEngine,
+        *,
+        confidence_threshold: float = DEFAULT_CONFIDENCE_THRESHOLD,
+        image_loader: Optional[Callable[[Path], "np.ndarray"]] = None,
+    ) -> None:
+        self.engine = engine
+        self.confidence_threshold = confidence_threshold
+        self._load_image = image_loader or _default_image_loader
+
+    def process_page(
+        self,
+        png_path: Path,
+        screenshot: Screenshot,
+    ) -> OcrPageResult:
+        """Run one screenshot through the OCR engine.
+
+        Raises:
+            TextOcrAdapterError: when the PNG file cannot be decoded
+                at all (propagated error_code=``ocr_image_decode_error``
+                so the scanner can mark the session row ``error``
+                with that code — this is a **non-retryable** failure
+                because redoing OCR with the same broken bytes won't
+                help).
+
+        Zero-hit / all-low-confidence cases are NOT raised — they
+        return a valid ``OcrPageResult`` with empty ``text_blocks``
+        and ``filtered_out_count`` set so Admin can distinguish.
+        """
+        t0 = time.perf_counter()
+        try:
+            image_bgr = self._load_image(png_path)
+        except FileNotFoundError:
+            raise
+        except Exception as e:  # decode failure surfaces as adapter error
+            raise TextOcrAdapterError(
+                f"failed to decode image at {png_path!s}: {e!r}",
+                error_code="ocr_image_decode_error",
+            ) from e
+        if image_bgr is None:
+            raise TextOcrAdapterError(
+                f"cv2.imdecode returned None for {png_path!s}",
+                error_code="ocr_image_decode_error",
+            )
+
+        original_h, original_w = image_bgr.shape[:2]
+        resized_bgr, scale_ratio = resize_long_edge_to(
+            image_bgr, screenshot.ocr_scale_hint
+        )
+        processed_h, processed_w = resized_bgr.shape[:2]
+        resized_applied = scale_ratio != 1.0
+
+        try:
+            raw_items = self.engine.detect_and_recognize(resized_bgr)
+        except Exception as e:
+            LOG.warning(
+                "processor.text_ocr_engine_error",
+                extra={
+                    "screenshot_id": screenshot.screenshot_id,
+                    "engine": getattr(self.engine, "name", type(self.engine).__name__),
+                    "error": str(e),
+                },
+            )
+            wall_ms = (time.perf_counter() - t0) * 1000.0
+            return OcrPageResult(
+                screenshot_id=screenshot.screenshot_id,
+                scale_hint=screenshot.ocr_scale_hint,
+                original_size=(original_w, original_h),
+                processed_size=(processed_w, processed_h),
+                resized_applied=resized_applied,
+                text_blocks=[],
+                raw_full_text="",
+                engine_name=getattr(self.engine, "name", type(self.engine).__name__),
+                wall_ms=wall_ms,
+                filtered_out_count=0,
+                confidence_threshold=self.confidence_threshold,
+                soft_error=f"engine_error:{type(e).__name__}",
+            )
+
+        inv_ratio = 1.0 / scale_ratio if scale_ratio != 0 else 1.0
+        blocks: list[TextBlock] = []
+        filtered_out = 0
+        for item in raw_items:
+            if item.confidence < self.confidence_threshold:
+                filtered_out += 1
+                continue
+            axis_aligned = quad_to_xyxy(item.bbox_quad)
+            bbox_original = scale_bbox(axis_aligned, inv_ratio)
+            # Clamp to original image bounds — shields downstream
+            # consumers from off-by-one rounding overshoot.
+            x1, y1, x2, y2 = bbox_original
+            x1 = min(max(0, x1), original_w)
+            y1 = min(max(0, y1), original_h)
+            x2 = min(max(0, x2), original_w)
+            y2 = min(max(0, y2), original_h)
+            if x2 <= x1 or y2 <= y1:
+                filtered_out += 1
+                continue
+            blocks.append(
+                TextBlock(
+                    text=item.text,
+                    bbox_xyxy=(x1, y1, x2, y2),
+                    confidence=item.confidence,
+                    line_read_index=0,
+                )
+            )
+
+        blocks = sort_reading_order(blocks)
+        blocks = [
+            TextBlock(
+                text=b.text,
+                bbox_xyxy=b.bbox_xyxy,
+                confidence=b.confidence,
+                line_read_index=idx,
+                paragraph_read_index=0,
+            )
+            for idx, b in enumerate(blocks, start=1)
+        ]
+        blocks = assign_paragraph_read_index(blocks)
+        raw_full_text = "\n".join(b.text for b in blocks)
+        wall_ms = (time.perf_counter() - t0) * 1000.0
+
+        return OcrPageResult(
+            screenshot_id=screenshot.screenshot_id,
+            scale_hint=screenshot.ocr_scale_hint,
+            original_size=(original_w, original_h),
+            processed_size=(processed_w, processed_h),
+            resized_applied=resized_applied,
+            text_blocks=blocks,
+            raw_full_text=raw_full_text,
+            engine_name=getattr(self.engine, "name", type(self.engine).__name__),
+            wall_ms=wall_ms,
+            filtered_out_count=filtered_out,
+            confidence_threshold=self.confidence_threshold,
+        )
+
+
+# ============================================================================
+# Default image loader
+# ============================================================================
+
+
+def _default_image_loader(png_path: Path) -> "np.ndarray":
+    """Load a PNG file as BGR ``np.ndarray`` via cv2.imdecode.
+
+    Uses ``imdecode`` (not ``imread``) so non-ASCII paths on Windows
+    still work — ``imread`` chokes on Unicode paths.
+    """
+    import cv2
+    import numpy as np
+
+    if not png_path.exists():
+        raise FileNotFoundError(f"PNG not found: {png_path!s}")
+    data = png_path.read_bytes()
+    arr = np.frombuffer(data, dtype=np.uint8)
+    img = cv2.imdecode(arr, cv2.IMREAD_COLOR)
+    return img
+
+
+__all__ = [
+    "DEFAULT_CONFIDENCE_THRESHOLD",
+    "READING_ORDER_ROW_TOLERANCE_PX",
+    "OcrEngine",
+    "OcrPageResult",
+    "PaddleOcrEngine",
+    "RawOcrItem",
+    "TextBlock",
+    "TextOcrAdapter",
+    "TextOcrAdapterError",
+    "quad_to_xyxy",
+    "resize_long_edge_to",
+    "scale_bbox",
+    "sort_reading_order",
+    "assign_paragraph_read_index",
+]