src/ocr_engine.py

"""
LightOnOCR engine wrapper.

Model files are cached through storage.py before Hugging Face / Transformers are
initialized. On app restarts, existing files in the persistent cache are reused.
"""

from __future__ import annotations

import base64
import io
import json
import os
import shutil
from typing import Any, Dict, Literal, Union

import cv2
import numpy as np
import requests
from PIL import Image

try:
    from .storage import ensure_dirs, get_env_overrides, is_model_cached
except ImportError:
    from storage import ensure_dirs, get_env_overrides, is_model_cached

ensure_dirs()
os.environ.update(get_env_overrides())


DEFAULT_MODEL_ID = "PetaniHandal/LightOnOCR-2-1B-1025-ft-iam-handwriting-vllm"

PROMPTS = {
    "handwriting": (
        "Extract all handwritten and printed text from this document image. "
        "Preserve reading order, line breaks, tables, totals, and important labels. "
        "Return clean Markdown only."
    ),
    "document": (
        "Perform OCR on this document. Preserve reading order, headings, tables, "
        "and key-value fields. Return clean Markdown only."
    ),
    "receipt": (
        "Extract receipt or invoice content from this image. Preserve merchant, "
        "date, item rows, prices, totals, and notes. Return clean Markdown only."
    ),
}


def _hub_cache_dir() -> str:
    return os.environ.get("HUGGINGFACE_HUB_CACHE", os.path.join(os.environ["HF_HOME"], "hub"))


class OCREngine:
    def __init__(
        self,
        preset: Literal["handwriting", "document", "receipt"] = "handwriting",
        model_id: str | None = None,
        endpoint_url: str | None = None,
        max_new_tokens: int = 4096,
        temperature: float = 0.1,
        top_p: float = 0.9,
    ):
        if preset not in PROMPTS:
            raise ValueError(f"Preset '{preset}' tidak dikenal. Pilihan: {list(PROMPTS)}")

        self.preset = preset
        self.model_id = model_id or os.getenv("LIGHTONOCR_MODEL_ID", DEFAULT_MODEL_ID)
        self.endpoint_url = endpoint_url or os.getenv("LIGHTONOCR_ENDPOINT_URL")
        self.max_new_tokens = max_new_tokens
        self.temperature = temperature
        self.top_p = top_p
        self._processor = None
        self._model = None
        self._device = None
        self._dtype = None

        if not self.endpoint_url:
            self._load_local_model()

    def _load_local_model(self) -> None:
        import torch

        try:
            from transformers import AutoProcessor, LightOnOcrForConditionalGeneration
        except ImportError:
            from transformers import AutoProcessor
            from transformers import AutoModelForSeq2SeqLM as LightOnOcrForConditionalGeneration

        device = "cuda" if torch.cuda.is_available() else "cpu"
        dtype = torch.bfloat16 if device == "cuda" else torch.float32
        adapter_config = self._load_adapter_config()
        model_to_load = adapter_config.get("base_model_name_or_path", self.model_id)
        model_local_only = is_model_cached(model_to_load)
        adapter_local_only = is_model_cached(self.model_id)

        self._processor = AutoProcessor.from_pretrained(
            self.model_id,
            cache_dir=_hub_cache_dir(),
            local_files_only=adapter_local_only,
            trust_remote_code=True,
        )
        try:
            self._model = LightOnOcrForConditionalGeneration.from_pretrained(
                model_to_load,
                cache_dir=_hub_cache_dir(),
                local_files_only=model_local_only,
                dtype=dtype,
                device_map="auto",
                attn_implementation="sdpa",
                trust_remote_code=True,
            )
        except OSError as exc:
            if model_local_only and self._looks_like_missing_weights(exc):
                self._remove_partial_model_cache(model_to_load)
                self._model = LightOnOcrForConditionalGeneration.from_pretrained(
                    model_to_load,
                    cache_dir=_hub_cache_dir(),
                    local_files_only=False,
                    dtype=dtype,
                    device_map="auto",
                    attn_implementation="sdpa",
                    trust_remote_code=True,
                )
            else:
                raise

        if adapter_config:
            from peft import PeftModel

            try:
                self._model = PeftModel.from_pretrained(
                    self._model,
                    self.model_id,
                    cache_dir=_hub_cache_dir(),
                    local_files_only=adapter_local_only,
                )
            except OSError as exc:
                if adapter_local_only and self._looks_like_missing_weights(exc):
                    self._remove_partial_model_cache(self.model_id)
                    self._model = PeftModel.from_pretrained(
                        self._model,
                        self.model_id,
                        cache_dir=_hub_cache_dir(),
                        local_files_only=False,
                    )
                else:
                    raise

        self._model.eval()
        self._device = device
        self._dtype = dtype

    def _load_adapter_config(self) -> dict[str, Any]:
        from huggingface_hub import hf_hub_download
        from huggingface_hub.utils import EntryNotFoundError

        local_only = is_model_cached(self.model_id)
        try:
            path = hf_hub_download(
                repo_id=self.model_id,
                filename="adapter_config.json",
                cache_dir=_hub_cache_dir(),
                local_files_only=local_only,
            )
        except EntryNotFoundError:
            return {}
        except OSError:
            if local_only:
                path = hf_hub_download(
                    repo_id=self.model_id,
                    filename="adapter_config.json",
                    cache_dir=_hub_cache_dir(),
                    local_files_only=False,
                )
            else:
                raise

        with open(path, "r", encoding="utf-8") as file:
            return json.load(file)

    def _looks_like_missing_weights(self, exc: OSError) -> bool:
        message = str(exc)
        return "pytorch_model.bin" in message or "model.safetensors" in message

    def _remove_partial_model_cache(self, model_id: str) -> None:
        repo_dir = f"models--{model_id.replace('/', '--')}"
        cache_dir = os.path.join(os.environ["HF_HOME"], "hub", repo_dir)
        if os.path.isdir(cache_dir):
            shutil.rmtree(cache_dir, ignore_errors=True)

    def process_image(
        self,
        image: Union[str, np.ndarray, Image.Image],
        max_size: int = 1540,
    ) -> Dict[str, Any]:
        pil_image = self._to_pil_image(image)
        pil_image = self._resize_if_needed(pil_image, max_size)

        if self.endpoint_url:
            markdown_text = self._process_with_endpoint(pil_image)
        else:
            markdown_text = self._process_local(pil_image)

        return {
            "markdown_text": markdown_text.strip(),
            "images": [],
            "model_id": self.model_id,
            "runtime": "vLLM endpoint" if self.endpoint_url else "Transformers local",
        }

    def _process_local(self, image: Image.Image) -> str:
        import torch

        assert self._processor is not None
        assert self._model is not None

        messages = [{"role": "user", "content": [{"type": "image"}]}]
        text = self._processor.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=True,
        )
        inputs = self._processor(text=[text], images=[image], return_tensors="pt").to(self._device)
        if "pixel_values" in inputs:
            inputs["pixel_values"] = inputs["pixel_values"].to(self._dtype)

        with torch.inference_mode():
            output_ids = self._model.generate(
                **inputs,
                max_new_tokens=self.max_new_tokens,
                do_sample=self.temperature > 0,
                temperature=self.temperature,
                top_p=self.top_p,
            )
        input_length = inputs["input_ids"].shape[1]
        return self._processor.tokenizer.decode(
            output_ids[0, input_length:],
            skip_special_tokens=True,
        )

    def _process_with_endpoint(self, image: Image.Image) -> str:
        buffer = io.BytesIO()
        image.save(buffer, format="PNG")
        image_base64 = base64.b64encode(buffer.getvalue()).decode("utf-8")
        payload = {
            "model": self.model_id,
            "messages": [
                {
                    "role": "user",
                    "content": [
                        {"type": "text", "text": PROMPTS[self.preset]},
                        {
                            "type": "image_url",
                            "image_url": {
                                "url": f"data:image/png;base64,{image_base64}",
                            },
                        },
                    ],
                }
            ],
            "max_tokens": self.max_new_tokens,
            "temperature": self.temperature,
            "top_p": self.top_p,
        }
        response = requests.post(self.endpoint_url, json=payload, timeout=300)
        response.raise_for_status()
        data = response.json()
        return data["choices"][0]["message"]["content"]

    def _to_pil_image(self, image: Union[str, np.ndarray, Image.Image]) -> Image.Image:
        if isinstance(image, Image.Image):
            return image.convert("RGB")
        if isinstance(image, str):
            img = cv2.imread(image)
            if img is None:
                raise ValueError(f"Gambar tidak bisa dibaca: {image}")
            return Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
        if isinstance(image, np.ndarray):
            if image.ndim == 2:
                return Image.fromarray(image).convert("RGB")
            return Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
        raise TypeError(f"Tipe gambar tidak didukung: {type(image)!r}")

    def _resize_if_needed(self, image: Image.Image, max_size: int) -> Image.Image:
        longest = max(image.size)
        if longest <= max_size:
            return image
        scale = max_size / longest
        size = (int(image.width * scale), int(image.height * scale))
        return image.resize(size, Image.Resampling.LANCZOS)