LayoutLMv3 fine-tuneado cargado directamente de Hugging Face: lucasgagneten/layoutlmv3-argentine-invoices

.gitignore

@@ -1,4 +1,6 @@
 .env
 /venv/
 /__pycache__/
-*.bat
+*.bat
+app_with_state_dict.py
+layoutlmv3_state_dict.pth

app.py

@@ -12,16 +12,16 @@ from io import BytesIO
 warnings.filterwarnings('ignore')
 
 # --- 1. Carga de Modelo y Procesador (CPU Habilitada) ---
-
-# --- CONFIGURACIÓN DE ARCHIVOS ---
-STATE_DICT_PATH = "./layoutlmv3_state_dict.pth"
-BASE_MODEL = "microsoft/layoutlmv3-base" # Usamos este para la arquitectura base
+# MODELO DE HUGGING FACE FINE-TUNEADO
+HUGGINGFACE_MODEL = "lucasgagneten/layoutlmv3-argentine-invoices" 
 
 # Define el dispositivo como CPU
 device = torch.device("cpu")
 print(f"Inferencia forzada al dispositivo: {device}")
 
 # Definir las etiquetas utilizadas durante el entrenamiento
+# Estas son necesarias para la lógica de visualización y la deduplicación,
+# aunque el modelo cargado ya contendrá esta información en su configuración.
 label_list = [
     'B-ALICUOTA',
     'B-COMPROBANTE_NUMERO',
@@ -44,6 +44,7 @@ label_list = [
     ]
 id2label = {i: label for i, label in enumerate(label_list)}
 label2id = {label: i for i, label in enumerate(label_list)}
+
 # 1. Definir una paleta de colores robusta
 color_palette = [
     'red', 'blue', 'green', 'purple', 'orange', 'brown', 
@@ -52,55 +53,43 @@ color_palette = [
 ]
 
 # 2. Extraer las etiquetas raíz únicas
-# La etiqueta 'O' (Outside) se ignora ya que no es una entidad
 root_labels = set()
 for label in label_list:
     if label != 'O':
-        # Split solo por el primer '-' para manejar etiquetas tipo 'B-ETIQUETA'
         root_label = label.split('-', 1)[-1] 
         root_labels.add(root_label)
 
 # 3. Crear el diccionario de asignación de color
 label2color = {}
-for i, root_label in enumerate(sorted(list(root_labels))): # Ordenar para consistencia
-    # Asigna un color de la paleta usando el operador módulo (%) para reciclar colores
+for i, root_label in enumerate(sorted(list(root_labels))):
     label2color[root_label] = color_palette[i % len(color_palette)]
 
 # Cargar el modelo/procesador
 try:
-    # 1. Cargar la configuración de procesamiento de imagen, FORZANDO apply_ocr=False
-    image_processor = LayoutLMv3ImageProcessor.from_pretrained(BASE_MODEL, apply_ocr=False)
-    
-    # 2. Inicializar AutoProcessor con el procesador de imagen ya configurado
+    # 1. Cargar el procesador directamente desde el modelo de HF. 
+    # El procesador de LayoutLMv3 siempre requiere que apply_ocr=False si se usa con OCR externo.
+    # AutoProcessor se encargará de cargar ImageProcessor, Tokenizer, y FeatureExtractor.
     loaded_processor = AutoProcessor.from_pretrained(
-        BASE_MODEL, image_processor=image_processor
+        HUGGINGFACE_MODEL, 
+        apply_ocr=False # Importante para usar los resultados de DocTR
     )
     
-    # 2. Cargar la arquitectura base de LayoutLMv3 (sin los pesos)
-    # Se añade la configuración de las etiquetas personalizadas
+    # 2. Cargar el modelo de Clasificación de Tokens directamente desde el repositorio de HF.
+    # Esto carga tanto la arquitectura como los pesos fine-tuneados.
     loaded_model = LayoutLMv3ForTokenClassification.from_pretrained(
-        BASE_MODEL,
-        num_labels=len(label_list),
-        id2label=id2label,
-        label2id=label2id
-    ).to(device)
-
-    # 3. Cargar los pesos fine-tuneados desde el archivo .pth
-    if os.path.exists(STATE_DICT_PATH):
-        # Mapear a la CPU para asegurar la compatibilidad
-        state_dict = torch.load(STATE_DICT_PATH, map_location=device)
-        
-        # Inyectar los pesos en el modelo
-        loaded_model.load_state_dict(state_dict)
-        print(f"Modelo fine-tuneado cargado exitosamente desde {STATE_DICT_PATH} en CPU.")
-    else:
-        print(f"Advertencia: No se encontró el archivo de pesos: {STATE_DICT_PATH}. Usando pesos iniciales del modelo base.")
+        HUGGINGFACE_MODEL
+    ).to(device) # Mover a la CPU
+    
+    # Sobrescribir id2label/label2id para consistencia, aunque ya deberían estar cargados
+    # en la configuración del modelo de HF. Esto es una precaución.
+    loaded_model.config.id2label = id2label
+    loaded_model.config.label2id = label2id
+    
+    print(f"Modelo fine-tuneado cargado exitosamente desde Hugging Face: {HUGGINGFACE_MODEL} en CPU.")
 
 except Exception as e:
-    print(f"Error fatal al cargar el modelo o procesador: {e}")
-    # En un entorno de producción, puedes optar por salir o cargar el modelo base como fallback.
-    # Por simplicidad, el código anterior se salta el fallback del modelo base,
-    # ya que la arquitectura base ya fue cargada, solo falló la inyección de pesos.
+    print(f"Error fatal al cargar el modelo o procesador desde Hugging Face: {e}")
+    # Nota: Aquí la aplicación fallaría si no puede descargar el modelo.
 
 # Cargar el predictor OCR de DocTR
 doctr_model = ocr_predictor(det_arch='db_resnet50', reco_arch='crnn_vgg16_bn', pretrained=True)
@@ -121,7 +110,6 @@ def process_invoice(image: Image.Image):
         
         # 2. Guardar la imagen en un buffer de memoria como si fuera un archivo JPG
         img_byte_arr = BytesIO()
-        # Nota: Asegúrate de que PIL pueda guardar como 'jpeg' o 'png'
         rgb_image.save(img_byte_arr, format='JPEG') 
         
         # 3. Mover el puntero al inicio del buffer y obtener los bytes
@@ -129,11 +117,9 @@ def process_invoice(image: Image.Image):
         image_bytes = img_byte_arr.read()
         
         # 4. DocTR soporta la carga de una lista de bytes de imágenes
-        # NOTA: Usamos from_images y le pasamos los bytes de UNA imagen
         doctr_doc = DocumentFile.from_images([image_bytes])
         
     except Exception as e:
-        # Imprime el error completo en tu consola para depuración
         print(f"Error detallado al cargar imagen en DocTR: {e}") 
         return None, f"Error al procesar la imagen con DocTR (conversión): {e}", None, None
     
@@ -179,7 +165,6 @@ def process_invoice(image: Image.Image):
     pixel_values = encoding["pixel_values"].to(device) # LayoutLMv3 usa 'pixel_values'
 
     # 3. Inferencia del Modelo LayoutLMv3
-    # Asegúrate de poner el modelo en modo de evaluación
     loaded_model.eval() 
     with torch.no_grad():
         outputs = loaded_model(
@@ -191,8 +176,7 @@ def process_invoice(image: Image.Image):
 
     predictions = outputs.logits.argmax(dim=-1).squeeze().tolist()
 
-    # --- Mapeo Correcto de Predicciones a Palabras del OCR (Similar al Colab) ---
-    # Esto asegura que haya una predicción limpia por cada palabra extraída por DocTR.
+    # --- Mapeo Correcto de Predicciones a Palabras del OCR ---
     word_ids = encoding.word_ids()
     predictions_final = []
     current_word_index = None
@@ -200,38 +184,32 @@ def process_invoice(image: Image.Image):
     for idx, pred_id in enumerate(predictions):
         word_idx = word_ids[idx]
         
-        # Solo procesar tokens que se mapean a palabras (no CLS, SEP, etc.)
         if word_idx is not None:
-            # Solo tomar la predicción del primer sub-token de cada palabra
             if word_idx != current_word_index:
                 if len(predictions_final) < len(words):
-                    predictions_final.append(id2label[pred_id])
+                    # Usar el id2label del modelo cargado, que ahora es la fuente de verdad
+                    predictions_final.append(loaded_model.config.id2label[pred_id]) 
                 
                 current_word_index = word_idx
             
     # --- Fin del Mapeo ---
 
 
-    # 4. Agrupación de Resultados BIO (Recolecta todos los candidatos, incluidos duplicados)
-    # ner_candidates almacenará una lista de entidades para cada etiqueta raíz.
-    # Structure: {'ETIQUETA': [{'valor': '...', 'bbox_entity': [...]}, {...}]}
+    # 4. Agrupación de Resultados BIO 
     ner_candidates = {} 
     
     current_entity = []
-    current_label = None # Almacena la etiqueta raíz (ej. 'TOTAL')
+    current_label = None 
     current_bbox_group = []
     
-    # Función auxiliar para guardar la entidad actual
     def save_current_entity(entity_list, label, bbox_list):
         if not entity_list or not label:
             return
         
-        # 1. Calcular el BBox final de la entidad (min/max de todos los bboxes de las palabras)
         all_x = [b[0] for b in bbox_list] + [b[2] for b in bbox_list]
         all_y = [b[1] for b in bbox_list] + [b[3] for b in bbox_list]
         bbox_normalized = [min(all_x), min(all_y), max(all_x), max(all_y)]
         
-        # 2. Guardar en ner_candidates (permite duplicados)
         if label not in ner_candidates:
             ner_candidates[label] = []
             
@@ -240,7 +218,6 @@ def process_invoice(image: Image.Image):
             'bbox_entity': bbox_normalized
         })
 
-    # Iterar sobre palabras y sus predicciones finales
     for word_data, pred_label in zip(words_data, predictions_final):
         word_text = word_data["text"]
         word_box = word_data["box"]
@@ -249,22 +226,17 @@ def process_invoice(image: Image.Image):
         root_label = tag_parts[1] if len(tag_parts) > 1 else None
 
         if tag_type == 'B':
-            # 1. Si hay una entidad previa, guardarla.
             save_current_entity(current_entity, current_label, current_bbox_group)
             
-            # 2. Iniciar la nueva entidad.
             current_label = root_label
             current_entity = [word_text]
             current_bbox_group = [word_box]
             
         elif tag_type == 'I':
-            # Continuar solo si el I- tag corresponde a la entidad B- tag actual
             if current_label == root_label:
                 current_entity.append(word_text)
                 current_bbox_group.append(word_box)
             else:
-                # Si no coincide (error BIO), guardar la entidad previa (si existe) y
-                # tratar el I- tag desalineado como el inicio de una nueva entidad.
                 save_current_entity(current_entity, current_label, current_bbox_group)
                 
                 current_label = root_label 
@@ -272,15 +244,12 @@ def process_invoice(image: Image.Image):
                 current_bbox_group = [word_box]
             
         elif tag_type == 'O':
-            # Si se encuentra 'O', finalizar la entidad actual si existe.
             save_current_entity(current_entity, current_label, current_bbox_group)
             
-            # Resetear
             current_entity = []
             current_label = None
             current_bbox_group = []
 
-    # Añadir la última entidad después del bucle
     save_current_entity(current_entity, current_label, current_bbox_group)
 
 
@@ -291,13 +260,9 @@ def process_invoice(image: Image.Image):
         if not candidates:
             continue
             
-        # Ordenar por longitud de la cadena de valor (mayor a menor)
         sorted_candidates = sorted(candidates, key=lambda x: len(x['valor']), reverse=True)
-        
-        # El mejor candidato es el primero (el más largo)
         best_candidate = sorted_candidates[0]
         
-        # Agregar al resultado final (ya desduplicado)
         final_ner_results.append({
             'etiqueta': label,
             'valor': best_candidate['valor'],
@@ -305,10 +270,10 @@ def process_invoice(image: Image.Image):
         })
 
 
-    # Preparar tabla de resultados (Usando final_ner_results)
+    # Preparar tabla de resultados
     table_data = [[res['etiqueta'], res['valor']] for res in final_ner_results]
 
-    # 6. Dibujar Bounding Boxes en la Imagen (para visualización)
+    # 6. Dibujar Bounding Boxes en la Imagen
     annotated_image = image.copy()
     draw = ImageDraw.Draw(annotated_image)
     
@@ -317,7 +282,7 @@ def process_invoice(image: Image.Image):
     except IOError:
         font = ImageFont.load_default()
         
-    for res in final_ner_results: # Usar final_ner_results
+    for res in final_ner_results: 
         label = res['etiqueta']
         min_x_norm, min_y_norm, max_x_norm, max_y_norm = res['bbox_entity']
         
@@ -336,7 +301,7 @@ def process_invoice(image: Image.Image):
     # 7. Devolver resultados
     return annotated_image, "Extracción de Entidades Nombradas completada.", table_data, [
         {'etiqueta': r['etiqueta'], 'valor': r['valor'], 'bbox_entity': r['bbox_entity']} 
-        for r in final_ner_results # Usar final_ner_results
+        for r in final_ner_results
     ]
 
 # --- 3. Interfaz Gradio ---
@@ -359,7 +324,7 @@ with gr.Blocks(title="NER de Facturas Argentinas con LayoutLMv3 y DocTR") as dem
         f"""
         # 🇦🇷 Extracción de Datos de Facturas Argentinas (LayoutLMv3 + DocTR)
         Carga una imagen de factura para realizar OCR (DocTR) y Reconocimiento de Entidades Nombradas (NER) 
-        con un modelo **LayoutLMv3 fine-tuneado** cargado desde **`{STATE_DICT_PATH}`**, forzando la **ejecución en CPU**.
+        con un modelo **LayoutLMv3 fine-tuneado** cargado directamente de Hugging Face: **`{HUGGINGFACE_MODEL}`**, forzando la **ejecución en CPU**.
         """
     )