Update app.py

app.py

@@ -1,159 +1,63 @@
-import torch
-import numpy as np
 import gradio as gr
+import numpy as np
 from PIL import Image
+import pytesseract
 
-from transformers import (
-    DetrImageProcessor,
-    TableTransformerForObjectDetection,
-    TrOCRProcessor,
-    VisionEncoderDecoderModel
-)
+def extract_second_column(image):
+    if image is None:
+        return "Aucune image fournie"
 
-# ===============================
-# Chargement des modèles
-# ===============================
+    image = image.convert("RGB")
+    img = np.array(image)
 
-DEVICE = "cpu"
+    # OCR avec positions
+    data = pytesseract.image_to_data(
+        img,
+        output_type=pytesseract.Output.DICT,
+        config="--psm 6"
+    )
 
-# Table detection
-table_processor = DetrImageProcessor.from_pretrained(
-    "microsoft/table-transformer-detection"
-)
-table_model = TableTransformerForObjectDetection.from_pretrained(
-    "microsoft/table-transformer-detection"
-).to(DEVICE)
-table_model.eval()
+    words = []
+    for i in range(len(data["text"])):
+        text = data["text"][i].strip()
+        if text:
+            x = data["left"][i]
+            y = data["top"][i]
+            words.append((text, x, y))
 
-# OCR
-ocr_processor = TrOCRProcessor.from_pretrained(
-    "microsoft/trocr-base-printed"
-)
-ocr_model = VisionEncoderDecoderModel.from_pretrained(
-    "microsoft/trocr-base-printed"
-).to(DEVICE)
-ocr_model.eval()
+    if not words:
+        return "Aucun texte détecté"
 
-# ===============================
-# Utils
-# ===============================
+    # Trier par X (colonnes)
+    words.sort(key=lambda w: w[1])
 
-def cluster_columns(boxes, x_threshold=25):
-    """
-    Regroupe les bounding boxes par colonnes
-    en se basant sur la position X (x_min)
-    """
-    boxes = sorted(boxes, key=lambda b: b[0])
+    # Regrouper par colonnes
     columns = []
-
-    for box in boxes:
+    for word in words:
         placed = False
         for col in columns:
-            if abs(col[0][0] - box[0]) < x_threshold:
-                col.append(box)
+            if abs(col[0][1] - word[1]) < 60:
+                col.append(word)
                 placed = True
                 break
         if not placed:
-            columns.append([box])
-
-    return columns
-
-
-def ocr_cell(image, box):
-    crop = image.crop(box)
-    pixel_values = ocr_processor(
-        crop, return_tensors="pt"
-    ).pixel_values.to(DEVICE)
-
-    with torch.no_grad():
-        generated_ids = ocr_model.generate(pixel_values)
-
-    text = ocr_processor.batch_decode(
-        generated_ids, skip_special_tokens=True
-    )[0]
-
-    return text.strip()
-
-
-# ===============================
-# Pipeline principal
-# ===============================
-
-def extract_second_column(image):
-    if image is None:
-        return "Aucune image fournie"
-
-    image = image.convert("RGB")
-
-    # 1. Détection des cellules
-    inputs = table_processor(
-        images=image, return_tensors="pt"
-    ).to(DEVICE)
-
-    with torch.no_grad():
-        outputs = table_model(**inputs)
-
-    target_sizes = torch.tensor(
-        [image.size[::-1]]
-    )
-
-    results = table_processor.post_process_object_detection(
-        outputs,
-        threshold=0.3,
-        target_sizes=target_sizes
-    )[0]
-
-    # 2. Garder uniquement les cellules
-    cells = []
-    for label, box in zip(results["labels"], results["boxes"]):
-        label_name = table_model.config.id2label[label.item()]
-        if label_name == "table cell":
-            cells.append([int(v) for v in box.tolist()])
-
-    if len(cells) == 0:
-        return "Aucune cellule détectée"
-
-    # 3. Regrouper par colonnes
-    columns = cluster_columns(cells)
+            columns.append([word])
 
     if len(columns) < 2:
         return "Moins de 2 colonnes détectées"
 
-    second_column = columns[1]
-
-    # Trier de haut en bas
-    second_column = sorted(second_column, key=lambda b: b[1])
+    # 2ᵉ colonne
+    second_column = sorted(columns[1], key=lambda w: w[2])
 
-    # 4. OCR
-    extracted_texts = []
-    for box in second_column:
-        text = ocr_cell(image, box)
-        if text:
-            extracted_texts.append(text)
-
-    if not extracted_texts:
-        return "Aucun texte OCR extrait"
-
-    return "\n".join(extracted_texts)
+    return "\n".join([w[0] for w in second_column])
 
 
-# ===============================
-# Interface Gradio
-# ===============================
-
 demo = gr.Interface(
     fn=extract_second_column,
-    inputs=gr.Image(type="pil", label="Image du tableau"),
-    outputs=gr.Textbox(
-        label="Contenu de la 2ᵉ colonne",
-        lines=20
-    ),
-    title="Extraction automatique de la 2ᵉ colonne d’un tableau",
-    description=(
-        "Upload une image de tableau (JPEG/PNG).\n"
-        "Le système détecte le tableau et extrait uniquement "
-        "les cellules de la deuxième colonne."
-    )
+    inputs=gr.Image(type="pil"),
+    outputs=gr.Textbox(lines=20),
+    title="Extraction de la 2ᵉ colonne (Facture)",
+    description="OCR + regroupement par colonnes (optimisé pour factures)"
 )
 
 demo.launch()