scripts/inference/ocr_table.py

"""Table extraction pipeline — detect structure and OCR cells.

Primary: Microsoft Table Transformer (TATR) for structure + PaddleOCR per cell.
Fallback: img2table when TATR returns degenerate structure (<2 rows or <2 cols).

Usage:
    from scripts.inference.ocr_table import TableOCR
    table_ocr = TableOCR()
    result = table_ocr.extract(crop_image)
"""
from __future__ import annotations

from pathlib import Path

import numpy as np
import cv2


class TableOCR:
    """Extract structured table data from Table crop images."""

    def __init__(self, langs: list[str] | None = None, use_gpu: bool = True, ocr_backend: str = "auto"):
        self.langs = langs or ["vi", "en"]
        self.use_gpu = use_gpu
        self.ocr_backend = ocr_backend
        self._cell_ocr = None
        self._tatr_model = None
        self._tatr_processor = None

    @property
    def cell_ocr(self):
        """Lazy-init OCR engine for cell text extraction."""
        if self._cell_ocr is not None:
            return self._cell_ocr

        if self.ocr_backend == "auto":
            try:
                from paddleocr import PaddleOCR
                self._cell_ocr = ("paddle", PaddleOCR(lang=self.langs[0]))
                print("  TableOCR cell backend: paddle")
            except Exception:
                import easyocr
                self._cell_ocr = ("easyocr", easyocr.Reader(self.langs, gpu=self.use_gpu))
                print("  TableOCR cell backend: easyocr")
        elif self.ocr_backend == "paddle":
            from paddleocr import PaddleOCR
            self._cell_ocr = ("paddle", PaddleOCR(lang=self.langs[0]))
        else:
            import easyocr
            self._cell_ocr = ("easyocr", easyocr.Reader(self.langs, gpu=self.use_gpu))

        return self._cell_ocr

    @property
    def tatr_model(self):
        if self._tatr_model is None:
            from transformers import TableTransformerForObjectDetection
            self._tatr_model = TableTransformerForObjectDetection.from_pretrained(
                "microsoft/table-transformer-structure-recognition-v1.1-all"
            )
            if self.use_gpu:
                import torch
                if torch.cuda.is_available():
                    self._tatr_model = self._tatr_model.to("cuda")
            self._tatr_model.eval()
        return self._tatr_model

    @property
    def tatr_processor(self):
        if self._tatr_processor is None:
            from transformers import AutoImageProcessor
            self._tatr_processor = AutoImageProcessor.from_pretrained(
                "microsoft/table-transformer-structure-recognition-v1.1-all"
            )
        return self._tatr_processor

    def _preprocess_cell(self, cell_image: np.ndarray) -> np.ndarray:
        """Preprocess a cell crop for better OCR: upscale + sharpen."""
        h, w = cell_image.shape[:2]
        # Upscale very small cells
        if h < 40:
            scale = 40 / h
            cell_image = cv2.resize(cell_image, None, fx=scale, fy=scale, interpolation=cv2.INTER_CUBIC)
        # Add small white border (helps OCR avoid edge artifacts)
        cell_image = cv2.copyMakeBorder(cell_image, 5, 5, 5, 5, cv2.BORDER_CONSTANT, value=(255, 255, 255))
        return cell_image

    def _ocr_cell(self, cell_image: np.ndarray) -> str:
        """Run OCR on a single cell crop."""
        if cell_image.size == 0 or cell_image.shape[0] < 5 or cell_image.shape[1] < 5:
            return ""

        cell_image = self._preprocess_cell(cell_image)
        backend_name, engine = self.cell_ocr

        if backend_name == "paddle":
            try:
                result = engine.ocr(cell_image, cls=True)
            except TypeError:
                result = engine.ocr(cell_image)
            if not result or not result[0]:
                return ""
            texts = [det[1][0] for det in result[0] if det[1][1] > 0.3]
            return " ".join(texts)
        else:
            result = engine.readtext(cell_image)
            texts = [text for _, text, conf in result if conf > 0.1]
            return " ".join(texts)

    def _preprocess_table(self, image: np.ndarray, target_height: int = 800) -> np.ndarray:
        """Upscale table image for better structure detection and OCR."""
        h, w = image.shape[:2]
        if h < target_height:
            scale = target_height / h
            image = cv2.resize(image, None, fx=scale, fy=scale, interpolation=cv2.INTER_CUBIC)
        return image

    def _detect_structure_tatr(self, image: np.ndarray) -> dict | None:
        """Detect table structure using TATR. Returns rows/cols/cells or None."""
        import torch
        from PIL import Image

        # Upscale small table images
        image = self._preprocess_table(image)

        # Add padding around the table crop (helps TATR)
        padded = cv2.copyMakeBorder(image, 20, 20, 20, 20, cv2.BORDER_CONSTANT, value=(255, 255, 255))
        pil_img = Image.fromarray(cv2.cvtColor(padded, cv2.COLOR_BGR2RGB))

        inputs = self.tatr_processor(images=pil_img, return_tensors="pt")
        if self.use_gpu and torch.cuda.is_available():
            inputs = {k: v.to("cuda") for k, v in inputs.items()}

        with torch.no_grad():
            outputs = self.tatr_model(**inputs)

        # Post-process
        target_sizes = torch.tensor([pil_img.size[::-1]])
        if self.use_gpu and torch.cuda.is_available():
            target_sizes = target_sizes.to("cuda")

        results = self.tatr_processor.post_process_object_detection(
            outputs, threshold=0.3, target_sizes=target_sizes
        )[0]

        # TATR classes: 0=table, 1=table column, 2=table row,
        #               3=table column header, 4=table projected row header, 5=table spanning cell
        labels = results["labels"].cpu().numpy()
        boxes = results["boxes"].cpu().numpy()
        scores = results["scores"].cpu().numpy()

        pad = 20
        rows = []
        cols = []
        for label, box, score in zip(labels, boxes, scores):
            x1, y1, x2, y2 = box
            # Adjust for padding offset
            x1 = max(0, x1 - pad)
            y1 = max(0, y1 - pad)
            x2 = max(0, x2 - pad)
            y2 = max(0, y2 - pad)

            if label == 2:  # table row
                rows.append({"y1": int(y1), "y2": int(y2), "score": float(score)})
            elif label == 1:  # table column
                cols.append({"x1": int(x1), "x2": int(x2), "score": float(score)})

        if len(rows) < 2 or len(cols) < 2:
            return None

        # Sort rows top-to-bottom, cols left-to-right
        rows.sort(key=lambda r: r["y1"])
        cols.sort(key=lambda c: c["x1"])

        # Extract cells at row/col intersections
        h, w = image.shape[:2]
        cells = []
        for row_idx, row in enumerate(rows):
            for col_idx, col in enumerate(cols):
                x1 = max(0, col["x1"])
                y1 = max(0, row["y1"])
                x2 = min(w, col["x2"])
                y2 = min(h, row["y2"])

                if x2 <= x1 or y2 <= y1:
                    continue

                cell_crop = image[y1:y2, x1:x2]
                cell_text = self._ocr_cell(cell_crop)

                cells.append({
                    "row": row_idx,
                    "col": col_idx,
                    "text": cell_text,
                    "bbox_in_crop": {"x1": x1, "y1": y1, "x2": x2, "y2": y2},
                })

        return {
            "type": "table",
            "method": "tatr",
            "row_count": len(rows),
            "col_count": len(cols),
            "cells": cells,
            "as_csv": self._cells_to_csv(cells, len(rows), len(cols)),
        }

    def _extract_with_img2table(self, image: np.ndarray) -> dict | None:
        """Fallback: use img2table for structure detection."""
        try:
            from img2table.document import Image as Img2TableImage
            from img2table.ocr import PaddleOCR as Img2TablePaddle

            ocr_engine = Img2TablePaddle(lang=self.langs[0])

            # img2table expects a file path or PIL image
            import tempfile
            with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
                cv2.imwrite(tmp.name, image)
                doc = Img2TableImage(src=tmp.name)

            tables = doc.extract_tables(ocr=ocr_engine)
            if not tables:
                return None

            table = tables[0]
            df = table.df
            if df is None or df.empty:
                return None

            rows, cols = df.shape
            cells = []
            for r in range(rows):
                for c in range(cols):
                    val = str(df.iloc[r, c]) if df.iloc[r, c] is not None else ""
                    cells.append({
                        "row": r,
                        "col": c,
                        "text": val,
                        "bbox_in_crop": {"x1": 0, "y1": 0, "x2": 0, "y2": 0},
                    })

            return {
                "type": "table",
                "method": "img2table",
                "row_count": rows,
                "col_count": cols,
                "cells": cells,
                "as_csv": self._cells_to_csv(cells, rows, cols),
            }
        except Exception:
            return None

    def _extract_with_ocr_only(self, image: np.ndarray) -> dict:
        """Last resort: run OCR on full table image, return as text lines."""
        backend_name, engine = self.cell_ocr

        cells = []
        if backend_name == "paddle":
            try:
                result = engine.ocr(image, cls=True)
            except TypeError:
                result = engine.ocr(image)
            if result and result[0]:
                lines = sorted(result[0], key=lambda d: d[0][0][1])
                for idx, det in enumerate(lines):
                    text, conf = det[1][0], det[1][1]
                    if conf < 0.3:
                        continue
                    cells.append({
                        "row": idx, "col": 0, "text": text,
                        "bbox_in_crop": {
                            "x1": int(det[0][0][0]), "y1": int(det[0][0][1]),
                            "x2": int(det[0][2][0]), "y2": int(det[0][2][1]),
                        },
                    })
        else:
            result = engine.readtext(image)
            result.sort(key=lambda r: r[0][0][1])
            for idx, (bbox, text, conf) in enumerate(result):
                if conf < 0.3:
                    continue
                cells.append({
                    "row": idx, "col": 0, "text": text,
                    "bbox_in_crop": {
                        "x1": int(bbox[0][0]), "y1": int(bbox[0][1]),
                        "x2": int(bbox[2][0]), "y2": int(bbox[2][1]),
                    },
                })

        return {
            "type": "table",
            "method": "ocr_only",
            "row_count": len(cells),
            "col_count": 1,
            "cells": cells,
            "as_csv": "\n".join(c["text"] for c in cells),
        }

    def _cells_to_csv(self, cells: list[dict], n_rows: int, n_cols: int) -> str:
        """Convert cells list to CSV string."""
        grid = [[""] * n_cols for _ in range(n_rows)]
        for cell in cells:
            r, c = cell["row"], cell["col"]
            if 0 <= r < n_rows and 0 <= c < n_cols:
                grid[r][c] = cell["text"]
        return "\n".join(",".join(row) for row in grid)

    def _extract_with_line_detection(self, image: np.ndarray) -> dict | None:
        """Primary: detect table structure using classical CV line detection."""
        try:
            from table_structure import detect_table_structure
        except ImportError:
            import sys
            sys.path.insert(0, str(Path(__file__).resolve().parent))
            from table_structure import detect_table_structure

        structure = detect_table_structure(image)
        if structure is None:
            return None

        n_rows = structure["row_count"]
        n_cols = structure["col_count"]

        if n_rows < 1 or n_cols < 1:
            return None

        # Upscale image to match structure detection (which may have upscaled)
        h, w = image.shape[:2]
        struct_w = structure["image_size"]["width"]
        struct_h = structure["image_size"]["height"]
        if struct_h != h or struct_w != w:
            image = cv2.resize(image, (struct_w, struct_h), interpolation=cv2.INTER_CUBIC)

        # OCR each cell
        cells = []
        for cell_info in structure["cells"]:
            bbox = cell_info["bbox"]
            x1, y1, x2, y2 = bbox["x1"], bbox["y1"], bbox["x2"], bbox["y2"]
            cell_crop = image[y1:y2, x1:x2]
            cell_text = self._ocr_cell(cell_crop)

            cells.append({
                "row": cell_info["row"],
                "col": cell_info["col"],
                "text": cell_text,
                "bbox_in_crop": bbox,
            })

        return {
            "type": "table",
            "method": "line_detection",
            "row_count": n_rows,
            "col_count": n_cols,
            "cells": cells,
            "as_csv": self._cells_to_csv(cells, n_rows, n_cols),
        }

    def extract(self, image: np.ndarray) -> dict:
        """Extract structured table data from a Table crop image.

        Pipeline: line detection → TATR → img2table → plain OCR.

        Args:
            image: BGR numpy array of the cropped Table region.

        Returns:
            dict with type, method, row_count, col_count, cells, as_csv
        """
        # Primary: classical CV line detection (best for ruled BOM tables)
        try:
            result = self._extract_with_line_detection(image)
            if result is not None and result["row_count"] >= 2 and result["col_count"] >= 2:
                return result
        except Exception:
            pass

        # Secondary: TATR (better for PDF-style tables)
        try:
            result = self._detect_structure_tatr(image)
            if result is not None:
                return result
        except Exception:
            pass

        # Tertiary: img2table
        result = self._extract_with_img2table(image)
        if result is not None:
            return result

        # Last resort: plain OCR
        return self._extract_with_ocr_only(image)