first version

.gitignore

@@ -0,0 +1,4 @@
+.env
+/venv/
+/__pycache__/
+*.bat

README.md

@@ -1,8 +1,8 @@
 ---
 title: Layoutlmv3 Facturas Extractor
-emoji: 🚀
-colorFrom: green
-colorTo: gray
+emoji: 🏃
+colorFrom: blue
+colorTo: indigo
 sdk: gradio
 sdk_version: 5.49.1
 app_file: app.py

app.py

@@ -0,0 +1,366 @@
+import gradio as gr
+import numpy as np
+from PIL import Image, ImageDraw, ImageFont
+import torch
+from transformers import LayoutLMv3ImageProcessor, AutoProcessor, LayoutLMv3ForTokenClassification
+from doctr.models import ocr_predictor
+from doctr.io import DocumentFile
+from doctr.utils.visualization import visualize_page
+import os 
+import warnings
+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
+
+# Define el dispositivo como CPU
+device = torch.device("cpu")
+print(f"Inferencia forzada al dispositivo: {device}")
+
+# Definir las etiquetas utilizadas durante el entrenamiento
+label_list = [
+    'B-ALICUOTA', 'B-COMPROBANTE_NUMERO', 'B-CONCEPTO_GASTO', 'B-FECHA', 'B-INGRESOS_BRUTOS', 'B-IVA', 'B-JURISDICCION_GASTO', 'B-NETO', 'B-PROVEEDOR_CUIT', 'B-PROVEEDOR_RAZON_SOCIAL', 'B-TIPO', 'B-TOTAL', 'I-COMPROBANTE_NUMERO', 'I-CONCEPTO_GASTO', 'I-INGRESOS_BRUTOS', 'I-JURISDICCION_GASTO', 'I-PROVEEDOR_CUIT', 'I-PROVEEDOR_RAZON_SOCIAL', 'I-TOTAL', 'O'
+    ]
+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', 
+    'pink', 'cyan', 'lime', 'olive', 'teal', 'magenta', 
+    'navy', 'maroon', 'gold', 'silver', 'indigo', 'turquoise'
+]
+
+# 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
+    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
+    loaded_processor = AutoProcessor.from_pretrained(
+        BASE_MODEL, image_processor=image_processor
+    )
+    
+    # 2. Cargar la arquitectura base de LayoutLMv3 (sin los pesos)
+    # Se añade la configuración de las etiquetas personalizadas
+    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.")
+
+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.
+
+# Cargar el predictor OCR de DocTR
+doctr_model = ocr_predictor(det_arch='db_resnet50', reco_arch='crnn_vgg16_bn', pretrained=True)
+
+# --- 2. Función Principal de Inferencia y Visualización ---
+
+def process_invoice(image: Image.Image):
+    """
+    Realiza OCR con DocTR, NER con LayoutLMv3 y genera los resultados en tabla y imagen.
+    """
+    if image is None:
+        return None, "Por favor, carga una imagen de factura.", None, None
+
+    # 1. OCR con DocTR (obtener texto y bboxes)
+    try:
+        # 1. Asegurar el formato RGB
+        rgb_image = image.convert("RGB")
+        
+        # 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
+        img_byte_arr.seek(0)
+        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
+    
+    doctr_result = doctr_model(doctr_doc)
+    
+    if not doctr_result.pages:
+          return None, "DocTR no pudo extraer ninguna página de la imagen.", None, None
+
+    page = doctr_result.pages[0]
+
+    # Extraer texto, bboxes normalizados y fusionar a nivel de palabra
+    words_data = []
+    # La geometría de DocTR es [x_min, y_min] y [x_max, y_max] normalizada a [0, 1]
+    for block in page.blocks:
+        for line in block.lines:
+            for word in line.words:
+                text = word.value
+                # Coordenadas normalizadas a [0, 1000]
+                geom = np.array(word.geometry) * 1000
+                xmin, ymin = map(int, geom[0])
+                xmax, ymax = map(int, geom[1])
+                words_data.append({"text": text, "box": [xmin, ymin, xmax, ymax]})
+
+    words = [wd["text"] for wd in words_data]
+    boxes = [wd["box"] for wd in words_data]
+    image_width, image_height = image.size
+
+    # 2. Preprocesamiento para LayoutLMv3 (usando los resultados del OCR)
+    encoding = loaded_processor(
+        image,
+        words,
+        boxes=boxes,
+        max_length=512,
+        truncation=True,
+        padding="max_length",
+        return_tensors="pt"
+    )
+    
+    # Mover los tensores de entrada a la CPU antes de la inferencia
+    input_ids = encoding["input_ids"].to(device)
+    attention_mask = encoding["attention_mask"].to(device)
+    bbox = encoding["bbox"].to(device)
+    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(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            bbox=bbox,
+            pixel_values=pixel_values
+        )
+
+    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.
+    word_ids = encoding.word_ids()
+    predictions_final = []
+    current_word_index = None
+
+    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])
+                
+                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': [...]}, {...}]}
+    ner_candidates = {} 
+    
+    current_entity = []
+    current_label = None # Almacena la etiqueta raíz (ej. 'TOTAL')
+    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] = []
+            
+        ner_candidates[label].append({
+            'valor': " ".join(entity_list),
+            '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"]
+        tag_parts = pred_label.split('-', 1)
+        tag_type = tag_parts[0]
+        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 
+                current_entity = [word_text]
+                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)
+
+
+    # --- 5: DESDUPLICACIÓN (Seleccionar el valor más largo) ---
+    final_ner_results = []
+    
+    for label, candidates in ner_candidates.items():
+        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'],
+            'bbox_entity': best_candidate['bbox_entity']
+        })
+
+
+    # Preparar tabla de resultados (Usando final_ner_results)
+    table_data = [[res['etiqueta'], res['valor']] for res in final_ner_results]
+
+    # 6. Dibujar Bounding Boxes en la Imagen (para visualización)
+    annotated_image = image.copy()
+    draw = ImageDraw.Draw(annotated_image)
+    
+    try:
+        font = ImageFont.truetype("arial.ttf", 20)
+    except IOError:
+        font = ImageFont.load_default()
+        
+    for res in final_ner_results: # Usar final_ner_results
+        label = res['etiqueta']
+        min_x_norm, min_y_norm, max_x_norm, max_y_norm = res['bbox_entity']
+        
+        # Desnormalizar el bbox [0-1000] a píxeles
+        min_x = int(min_x_norm * image_width / 1000)
+        min_y = int(min_y_norm * image_height / 1000)
+        max_x = int(max_x_norm * image_width / 1000)
+        max_y = int(max_y_norm * image_height / 1000)
+        
+        color = label2color.get(label, 'yellow')
+        
+        draw.rectangle([min_x, min_y, max_x, max_y], outline=color, width=3)
+        draw.text((min_x, min_y - 20), label, fill=color, font=font)
+        
+    
+    # 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
+    ]
+
+# --- 3. Interfaz Gradio ---
+
+# Elementos de entrada y salida
+image_input = gr.Image(type="pil", label="Cargar Imagen de Factura", interactive=True)
+image_output = gr.Image(type="pil", label="Factura con Entidades Resaltadas")
+status_output = gr.Textbox(label="Estado", value="Carga una imagen y haz clic en 'Procesar'")
+table_output = gr.Dataframe(
+    headers=["Etiqueta", "Valor"],
+    label="Resultados de NER",
+    interactive=False,
+    col_count=(2, "fixed")
+)
+json_output = gr.JSON(label="Datos JSON Crudos (Incluye BBox Normalizados)", visible=True) 
+
+# Interfaz
+with gr.Blocks(title="NER de Facturas Argentinas con LayoutLMv3 y DocTR") as demo:
+    gr.Markdown(
+        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**.
+        """
+    )
+
+    with gr.Row():
+        with gr.Column(scale=1):
+            image_input.render()
+            process_button = gr.Button("🚀 Procesar Factura", variant="primary")
+            status_output.render()
+        with gr.Column(scale=2):
+            image_output.render()
+            table_output.render()
+            json_output.render()
+
+    process_button.click(
+        fn=process_invoice,
+        inputs=[image_input],
+        outputs=[image_output, status_output, table_output, json_output]
+    )
+
+# Lanzar la aplicación
+demo.launch()

layoutlmv3_state_dict.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:807818c88ce85767b337f03ce6ca7fd89ea14ce559c2981ea404cafc13557025
+size 503825075

requirements.txt

@@ -0,0 +1,20 @@
+# --- Requerimientos del Frameworks y Utilidades ---
+gradio>=4.0.0            # Interfaz de usuario
+pillow                  # Manipulación de imágenes (PIL)
+numpy                   # Operaciones numéricas
+
+# --- Requerimientos de OCR (DocTR) y NER (Transformers) ---
+# Usamos una versión más moderna de DocTR para asegurar compatibilidad
+python-doctr[viz,html]>=1.0.0 # Librería DocTR (incluye dependencias de CPU como Pillow)
+transformers>=4.30.0    # Librería principal para LayoutLMv3
+torch
+matplotlib # Añadir esta línea
+
+# --- Requerimientos de PyTorch ---
+# El archivo .pth requiere torch. Si lo instalas manualmente, puedes omitirlo.
+# Si quieres que pip lo instale (incluso la versión CPU), descomenta:
+# torch>=2.0.0
+
+# --- Requerimientos CRÍTICOS (ya incluidos o buena práctica) ---
+# protobuf se maneja internamente en transformers/torch.
+# Solo añadir si hay problemas específicos.