src/pageparse/ocr/printed.py

from __future__ import annotations

import shutil
import subprocess
import tempfile
from pathlib import Path

import cv2
import numpy as np


TESSERACT_CANDIDATES = [
    r"C:\Program Files\Tesseract-OCR\tesseract.exe",
    r"C:\Program Files (x86)\Tesseract-OCR\tesseract.exe",
]


class PrintedOCR:
    def __init__(self, tesseract_cmd: str | None = None) -> None:
        if tesseract_cmd is not None:
            resolved = tesseract_cmd
        else:
            resolved = shutil.which("tesseract")
            if resolved is None:
                for c in TESSERACT_CANDIDATES:
                    if Path(c).exists():
                        resolved = c
                        break
            if resolved is None:
                raise RuntimeError(
                    "Tesseract not found. Install Tesseract OCR from "
                    "https://github.com/UB-Mannheim/tesseract/wiki "
                    "or ensure it is in your PATH."
                )
        self.tesseract_cmd = resolved

    def recognize(self, image: np.ndarray) -> str:
        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
            tmp_path = tmp.name
            cv2.imwrite(tmp_path, image)

        try:
            result = subprocess.run(
                [self.tesseract_cmd, tmp_path, "stdout", "-l", "eng"],
                capture_output=True,
                text=True,
                timeout=30,
            )
            if result.returncode != 0:
                raise RuntimeError(
                    f"Tesseract failed (code {result.returncode}): {result.stderr.strip()}"
                )
            return result.stdout.strip()
        finally:
            Path(tmp_path).unlink(missing_ok=True)

    def detect_barcodes(self, image: np.ndarray) -> list[dict]:
        results = []
        try:
            from pyzbar.pyzbar import decode as pyzbar_decode
            gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) if image.ndim == 3 else image
            barcodes = pyzbar_decode(gray)
            for barcode in barcodes:
                results.append({
                    "data": barcode.data.decode("utf-8"),
                    "type": barcode.type,
                    "rect": {
                        "x": barcode.rect.left,
                        "y": barcode.rect.top,
                        "w": barcode.rect.width,
                        "h": barcode.rect.height,
                    },
                })
        except ImportError:
            pass
        except Exception as e:
            print(f"Barcode detection failed: {e}")
        return results

    def detect_tables(self, image: np.ndarray) -> list[list[str]]:
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) if image.ndim == 3 else image
        binary = cv2.adaptiveThreshold(
            ~gray, 255, cv2.ADAPTIVE_THRESH_MEAN_C, cv2.THRESH_BINARY, 15, -2,
        )

        horizontal = binary.copy()
        h_size = horizontal.shape[1] // 30
        h_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (h_size, 1))
        horizontal = cv2.erode(horizontal, h_kernel)
        horizontal = cv2.dilate(horizontal, h_kernel)

        vertical = binary.copy()
        v_size = vertical.shape[0] // 30
        v_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (1, v_size))
        vertical = cv2.erode(vertical, v_kernel)
        vertical = cv2.dilate(vertical, v_kernel)

        table_mask = cv2.add(horizontal, vertical)
        contours, _ = cv2.findContours(table_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

        tables = []
        for contour in contours:
            x, y, w, h = cv2.boundingRect(contour)
            if w > 50 and h > 50:
                cell = gray[y:y+h, x:x+w]
                try:
                    text = self.recognize(cell)
                    if text.strip():
                        tables.append(text.strip().split("\n"))
                except Exception:
                    pass

        return tables