ocr_processor.py

import os
import uuid
from PIL import Image, ImageDraw
import numpy as np
import cv2
from gemini_ner_client import get_gemini_ner_tags

# Módulos de Doctr
from doctr.models import ocr_predictor

# Módulos locales
# Asegúrate de tener este archivo 'error_handler.py' en tu entorno
# from error_handler import ErrorHandler

# --- Configuración de Directorios ---
DATASET_BASE_DIR = "dataset"
IMAGES_DIR = os.path.join(DATASET_BASE_DIR, "imagenes")


# --- Inicialización del Modelo Doctr ---
# Se carga el modelo una sola vez al inicio del módulo
try:
    # Usaremos el modelo por defecto, que es robusto.
    OCR_MODEL = ocr_predictor(pretrained=True) 
    print("✅ Modelo Doctr cargado con éxito.")
except Exception as e:
    print(f"FATAL: Error al cargar el modelo Doctr. Asegúrate de instalar doctr[full]. Error: {e}")
    OCR_MODEL = None

# --- Función Dummy para ErrorHandler ---
# Se asume que ErrorHandler.show_error existe. Si no, usa print o logging.
class ErrorHandler:
    @staticmethod
    def show_error(title, message=""):
        print(f"[{title}] ERROR: {message}")


# --- Preprocesamiento de Imagen (Activado) ---

def preprocess_image_for_ocr(image: Image.Image) -> Image.Image:
    """
    Aplica preprocesamiento con OpenCV (conversión a escala de grises y umbral adaptativo)
    para mejorar la calidad del OCR.
    """
    """
    try:
        # Convertir PIL Image a array NumPy RGB
        img_np = np.array(image.convert('RGB'))
        
        # Convertir a escala de grises
        gray = cv2.cvtColor(img_np, cv2.COLOR_RGB2GRAY)
        
        # Aplicar umbral adaptativo (útil para imágenes con iluminación variable)
        thresh = cv2.adaptiveThreshold(
            gray, 
            255, 
            cv2.ADAPTIVE_THRESH_GAUSSIAN_C, 
            cv2.THRESH_BINARY, 
            11,  # Tamaño del bloque (debe ser impar y > 1)
            2    # Constante a restar de la media
        )
        
        # Devolver el array umbralizado como una imagen PIL
        return Image.fromarray(thresh).convert('RGB')
    
    except Exception as e:
        ErrorHandler.show_error("Error de Preprocesamiento", f"Fallo al aplicar OpenCV: {e}")
        # En caso de fallo, devolvemos la imagen original
        return image
    """
    return image


# --- Función Principal de OCR con Doctr ---

def get_ocr_data_doctr(image: Image.Image):
    """
    Ejecuta Doctr en la imagen, opcionalmente preprocesada, y devuelve 
    tokens y bboxes normalizados para LayoutXLM.
    """
    if image is None:
        raise ValueError("No se proporcionó ninguna imagen.")

    if OCR_MODEL is None:
        ErrorHandler.show_error("Error de Modelo", 
                                 "El motor Doctr no está disponible. Revise el log de inicio.")
        return None, [], None

    W, H = image.size
    tokens_data = []
    
    # 💡 PASO CLAVE: Aplicar Preprocesamiento
    # Almacenamos la imagen original ANTES del preprocesamiento
    image_orig_unprocessed = image.copy()
    processed_image = preprocess_image_for_ocr(image)
    
    try:
        # Doctr espera la imagen como un array numpy RGB
        img_np = np.array(processed_image.convert('RGB')) 
        
        # 1. Ejecutar la predicción
        result = OCR_MODEL([img_np]) # El modelo espera una lista de imágenes
        
        # 2. Parsear el resultado a nivel de 'Word'
        for page in result.pages:
            for block in page.blocks:
                for line in block.lines:
                    for word in line.words:
                        token_text = word.value.strip()
                        
                        if not token_text:
                            continue

                        # Las coordenadas de Doctr son normalizadas a 0-1 (fracciones)
                        # word.geometry: [[x_min, y_min], [x_max, y_max]]
                        bbox_frac = word.geometry 

                        # BBox original (en píxeles de la imagen original)
                        x_min, y_min, x_max, y_max = [
                            int(bbox_frac[0][0] * W), 
                            int(bbox_frac[0][1] * H), 
                            int(bbox_frac[1][0] * W), 
                            int(bbox_frac[1][1] * H)
                        ]
                        bbox_original = [x_min, y_min, x_max, y_max]

                        # BBox normalizado a 0-1000 (para LayoutXLM)
                        # Nota: Se corrigen los denominadores. LayoutXLM normaliza X e Y
                        # con respecto al tamaño completo (W o H).
                        bbox_normalized = [
                            int(x_min * 1000 / W),
                            int(y_min * 1000 / H),
                            int(x_max * 1000 / W),
                            int(y_max * 1000 / H)
                        ]
                        
                        tokens_data.append({
                            'token': token_text,
                            'bbox_norm': bbox_normalized,
                            'bbox_orig': bbox_original,
                            'ner_tag': 'O' 
                        })
                        
        # 💡 Importante: Devolvemos la imagen original sin procesar, ya que 
        # esta es la que se guarda y se usa para el dibujo en la interfaz.
        return image_orig_unprocessed, tokens_data, None

    except Exception as e:
        ErrorHandler.show_error("Error de Procesamiento OCR", e)
        return None, [], None


# --- Funciones de Visualización y Guardado ---

def draw_boxes(image: Image.Image, tokens_data: list, highlight_index: int = -1):
    """Dibuja un resaltado en la imagen para el bounding box seleccionado."""
    if image is None or not tokens_data:
        return None

    img_copy = image.copy()
    draw = ImageDraw.Draw(img_copy)
    
    if highlight_index >= 0 and highlight_index < len(tokens_data):
        # Usamos bbox_orig para dibujar en la imagen a tamaño completo
        bbox = tokens_data[highlight_index]['bbox_orig']
        x1, y1, x2, y2 = bbox
        
        outline_color = (255, 0, 0)
        draw.rectangle([x1, y1, x2, y2], outline=outline_color, width=4)
        
    return img_copy


def save_image_to_dataset(image: Image.Image) -> str:
    """
    Genera un nombre único y guarda la imagen en el directorio de dataset.
    Retorna el nombre único del archivo.
    """
    # 1. Crear el nombre único (UUID + extensión)
    unique_filename = f"{uuid.uuid4()}.jpeg"
    save_path = os.path.join(IMAGES_DIR, unique_filename)

    # 2. Asegurar que el directorio exista (es importante en un entorno efímero)
    os.makedirs(IMAGES_DIR, exist_ok=True)
    
    # 3. Guardar la imagen
    image.save(save_path, format="JPEG")
    
    print(f"Imagen guardada: {save_path}")
    return unique_filename


# --- Función de Flujo Principal ---

def process_and_setup(image_file, api_key: str): # 💡 ACEPTA LA LLAVE API
    """
    Función inicial: OCR con Doctr, **NER asistido por Gemini**, 
    configuración del estado y guardar la imagen.
    """
    if image_file is None:
        empty_df = {'token': [], 'ner_tag': []}
        return None, [], None, empty_df, None, None 
        
    # 💡 Llama a la función de OCR basada en Doctr
    image_orig, tokens_data, _ = get_ocr_data_doctr(image_file) 
    
    if image_orig is None:
        empty_df = {'token': [], 'ner_tag': []}
        return None, [], None, empty_df, "Error fatal al procesar el OCR con Doctr. Revise el log.", None
    
    # 💡 PASO NUEVO: NER Asistido por Gemini (CONDICIONAL)
    msg_ner_assist = ""
    if api_key:
        # Llama a la función que SÓLO se ejecuta si api_key no es None/vacío
        tokens_data = get_gemini_ner_tags(api_key, tokens_data)
        msg_ner_assist = " (NER asistido)"
    
    # --- Guardar la Imagen Original ---
    image_filename = save_image_to_dataset(image_orig)
    
    if not tokens_data:
        empty_df = {'token': [], 'ner_tag': []}
        msg = "OCR completado. No se detectaron tokens válidos."
        ErrorHandler.show_error(msg)
        return image_orig, [], None, empty_df, msg, image_filename
        
    # Crear el DataFrame inicial para la edición en Gradio
    df_data = {
        'token': [item['token'] for item in tokens_data],
        'ner_tag': [item['ner_tag'] for item in tokens_data]
    }
    
    # La imagen inicial no tiene resaltado
    highlighted_image = image_orig.copy() 
    msg = f"OCR de Doctr completado. Tokens detectados: {len(tokens_data)}.{msg_ner_assist}"
    print(msg)
    
    return image_orig, tokens_data, highlighted_image, df_data, msg, image_filename