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