scripts/inference/ocr_note.py

"""Note OCR pipeline — extract text from Note regions.

Supports Vietnamese and English using PaddleOCR with dual-language detection.
Lines are sorted top-to-bottom for natural reading order.
Includes image preprocessing (upscale, denoise, binarize) for low-res inputs.

Usage:
    from scripts.inference.ocr_note import NoteOCR
    ocr = NoteOCR()
    text = ocr.extract(crop_image)
"""
from __future__ import annotations

import cv2
import numpy as np


class NoteOCR:
    """Extract text from Note crop images using PaddleOCR or EasyOCR."""

    def __init__(self, langs: list[str] | None = None, backend: str = "auto"):
        """Initialize OCR engines.

        Args:
            langs: Language codes to use. Defaults to ["vi", "en"].
            backend: "paddle", "easyocr", or "auto" (try paddle first).
        """
        self.langs = langs or ["vi", "en"]
        self.backend = backend
        self.engines: dict = {}

        if backend == "auto":
            try:
                self._init_paddle()
                self.backend = "paddle"
            except Exception:
                self._init_easyocr()
                self.backend = "easyocr"
        elif backend == "paddle":
            self._init_paddle()
        else:
            self._init_easyocr()

        print(f"  NoteOCR backend: {self.backend}")

    def _init_paddle(self):
        from paddleocr import PaddleOCR
        for lang in self.langs:
            self.engines[lang] = PaddleOCR(lang=lang)

    def _init_easyocr(self):
        import easyocr
        # EasyOCR uses different lang codes: vi, en
        self.engines["easyocr"] = easyocr.Reader(self.langs, gpu=True)

    def _run_ocr(self, image: np.ndarray, lang: str) -> list[dict]:
        """Run OCR with one language, return structured line results."""
        if self.backend == "easyocr":
            return self._run_easyocr(image)

        engine = self.engines[lang]
        try:
            result = engine.ocr(image, cls=True)
        except TypeError:
            result = engine.ocr(image)

        if not result or not result[0]:
            return []

        lines = []
        for detection in result[0]:
            bbox = detection[0]  # [[x1,y1],[x2,y2],[x3,y3],[x4,y4]]
            text = detection[1][0]
            confidence = detection[1][1]

            y_center = sum(pt[1] for pt in bbox) / 4
            x_center = sum(pt[0] for pt in bbox) / 4

            lines.append({
                "text": text,
                "confidence": float(confidence),
                "y_center": y_center,
                "x_center": x_center,
                "bbox": bbox,
            })

        return lines

    def _run_easyocr(self, image: np.ndarray) -> list[dict]:
        """Run EasyOCR, return structured line results."""
        reader = self.engines["easyocr"]
        result = reader.readtext(image)

        lines = []
        for bbox, text, confidence in result:
            # EasyOCR bbox: [[x1,y1],[x2,y1],[x2,y2],[x1,y2]]
            y_center = sum(pt[1] for pt in bbox) / 4
            x_center = sum(pt[0] for pt in bbox) / 4

            lines.append({
                "text": text,
                "confidence": float(confidence),
                "y_center": y_center,
                "x_center": x_center,
                "bbox": bbox,
            })

        return lines

    def _sort_lines(self, lines: list[dict], line_height_threshold: float = 15.0) -> list[dict]:
        """Sort lines in reading order: top-to-bottom, left-to-right.

        Groups lines that are on the same vertical level (within threshold)
        and sorts them left-to-right within each group.
        """
        if not lines:
            return lines

        # Sort by y first
        lines.sort(key=lambda l: l["y_center"])

        # Group into rows (lines within threshold are same row)
        rows: list[list[dict]] = []
        current_row = [lines[0]]

        for line in lines[1:]:
            if abs(line["y_center"] - current_row[0]["y_center"]) < line_height_threshold:
                current_row.append(line)
            else:
                rows.append(current_row)
                current_row = [line]
        rows.append(current_row)

        # Sort each row left-to-right
        sorted_lines = []
        for row in rows:
            row.sort(key=lambda l: l["x_center"])
            sorted_lines.extend(row)

        return sorted_lines

    def _preprocess(self, image: np.ndarray, target_height: int = 1000) -> np.ndarray:
        """Preprocess image for better OCR: upscale, denoise, sharpen.

        Engineering drawings are often low-res scans. Upscaling + sharpening
        significantly improves character recognition, especially for Vietnamese diacritics.
        """
        h, w = image.shape[:2]

        # Upscale small images (most impactful improvement)
        if h < target_height:
            scale = target_height / h
            image = cv2.resize(image, None, fx=scale, fy=scale, interpolation=cv2.INTER_CUBIC)

        # Convert to grayscale for processing
        if len(image.shape) == 3:
            gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        else:
            gray = image

        # Denoise
        gray = cv2.fastNlMeansDenoising(gray, h=10)

        # Adaptive threshold binarization (handles uneven lighting in scans)
        binary = cv2.adaptiveThreshold(
            gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, 15, 8
        )

        # Convert back to BGR (PaddleOCR expects 3-channel)
        result = cv2.cvtColor(binary, cv2.COLOR_GRAY2BGR)
        return result

    def extract(self, image: np.ndarray, min_confidence: float = 0.1) -> dict:
        """Extract text from a Note crop image.

        Args:
            image: BGR numpy array of the cropped Note region.
            min_confidence: Minimum confidence threshold per line.

        Returns:
            dict with keys:
                - text: joined text string
                - lines: list of {text, confidence, bbox}
                - language: detected primary language
        """
        best_result = {"text": "", "lines": [], "language": "unknown", "avg_confidence": 0.0}

        # Try both raw and preprocessed image, keep best result
        images_to_try = [image, self._preprocess(image)]

        for img in images_to_try:
            for lang in self.langs:
                lines = self._run_ocr(img, lang)
                lines = [l for l in lines if l["confidence"] >= min_confidence]

                if not lines:
                    continue

                avg_conf = sum(l["confidence"] for l in lines) / len(lines)

                if avg_conf > best_result["avg_confidence"]:
                    sorted_lines = self._sort_lines(lines)
                    best_result = {
                        "text": "\n".join(l["text"] for l in sorted_lines),
                        "lines": [
                            {
                                "text": l["text"],
                                "confidence": round(l["confidence"], 4),
                                "bbox": l["bbox"],
                            }
                            for l in sorted_lines
                        ],
                        "language": lang,
                        "avg_confidence": round(avg_conf, 4),
                    }

        return best_result