core/ocr.py

"""
Paso 1: Extracción de texto mediante OCR (LightOnOCR-2-1B).
"""

import time
from typing import Dict

import torch
from PIL import Image

from config import OCR_MODEL_ID, DEVICE, DTYPE
from core.models import ModelManager


def extract_text(image: Image.Image, manager: ModelManager) -> Dict:
    """
    Extrae texto de la imagen del menú usando LightOnOCR-2.

    Returns:
        dict con 'text', 'processing_time', 'model', 'success'
    """
    start_time = time.time()

    try:
        if image.mode != "RGB":
            image = image.convert("RGB")

        # Construir conversación con la imagen (API oficial de LightOnOCR-2)
        conversation = [
            {
                "role": "user",
                "content": [{"type": "image"}],
            }
        ]

        inputs = manager.ocr_processor.apply_chat_template(
            conversation,
            images=[image],
            add_generation_prompt=True,
            tokenize=True,
            return_dict=True,
            return_tensors="pt",
        )

        # Mover tensores al dispositivo correcto
        inputs = {
            k: v.to(device=DEVICE, dtype=DTYPE) if v.is_floating_point() else v.to(DEVICE)
            for k, v in inputs.items()
        }

        print("🔍 Ejecutando OCR...")
        with torch.no_grad():
            output_ids = manager.ocr_model.generate(
                **inputs,
                max_new_tokens=2048,
                do_sample=False,
                num_beams=3,
                early_stopping=True,
            )

        # Decodificar solo los tokens nuevos (sin el prompt)
        generated_ids = output_ids[0, inputs["input_ids"].shape[1]:]
        extracted_text = manager.ocr_processor.decode(
            generated_ids, skip_special_tokens=True
        )

        return {
            "text": extracted_text.strip(),
            "processing_time": time.time() - start_time,
            "model": OCR_MODEL_ID,
            "success": True,
        }

    except Exception as e:
        return {
            "text": "",
            "processing_time": time.time() - start_time,
            "model": OCR_MODEL_ID,
            "success": False,
            "error": str(e),
        }