Update app.py

app.py

@@ -0,0 +1,159 @@
+import torch
+import numpy as np
+import gradio as gr
+from PIL import Image
+
+from transformers import (
+    DetrImageProcessor,
+    TableTransformerForObjectDetection,
+    TrOCRProcessor,
+    VisionEncoderDecoderModel
+)
+
+# ===============================
+# Chargement des modèles
+# ===============================
+
+DEVICE = "cpu"
+
+# 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()
+
+# OCR
+ocr_processor = TrOCRProcessor.from_pretrained(
+    "microsoft/trocr-base-printed"
+)
+ocr_model = VisionEncoderDecoderModel.from_pretrained(
+    "microsoft/trocr-base-printed"
+).to(DEVICE)
+ocr_model.eval()
+
+# ===============================
+# Utils
+# ===============================
+
+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])
+    columns = []
+
+    for box in boxes:
+        placed = False
+        for col in columns:
+            if abs(col[0][0] - box[0]) < x_threshold:
+                col.append(box)
+                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.7,
+        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)
+
+    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])
+
+    # 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)
+
+
+# ===============================
+# 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."
+    )
+)
+
+demo.launch()