Update app.py

app.py

@@ -1,63 +1,113 @@
 import gradio as gr
 import numpy as np
-from PIL import Image
-import pytesseract
+from paddleocr import PaddleOCR
+from sklearn.cluster import KMeans
 
-def extract_second_column(image):
+ocr = PaddleOCR(
+    use_textline_orientation=True,
+    lang="fr"
+)
+
+def extract_column2_9_lines(image):
     if image is None:
-        return "Aucune image fournie"
+        return "Aucune image fournie."
 
-    image = image.convert("RGB")
     img = np.array(image)
+    result = ocr.predict(img)
+
+    if not result or len(result) == 0:
+        return "OCR exécuté mais aucun texte détecté."
+
+    data = result[0]
+    texts = data.get("rec_texts", [])
+    boxes = data.get("dt_polys", [])
+
+    if not texts:
+        return "Aucun texte exploitable détecté."
 
-    # OCR avec positions
-    data = pytesseract.image_to_data(
-        img,
-        output_type=pytesseract.Output.DICT,
-        config="--psm 6"
+    # -----------------------------
+    # 1. Collecte OCR
+    # -----------------------------
+    elements = []
+    for text, box in zip(texts, boxes):
+        text = text.strip()
+        if len(text) < 2:
+            continue
+        x_center = np.mean([p[0] for p in box])
+        y_center = np.mean([p[1] for p in box])
+        elements.append((x_center, y_center, text))
+
+    # -----------------------------
+    # 2. Clustering vertical (colonnes)
+    # -----------------------------
+    X = np.array([[e[0]] for e in elements])
+    n_cols = 6
+    kmeans = KMeans(n_clusters=n_cols, random_state=42).fit(X)
+    labels = kmeans.labels_
+
+    columns = {}
+    for (x, y, text), label in zip(elements, labels):
+        columns.setdefault(label, []).append((x, y, text))
+
+    sorted_columns = sorted(
+        columns.values(),
+        key=lambda col: np.mean([e[0] for e in col])
     )
 
-    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))
+    if len(sorted_columns) < 2:
+        return "Impossible de détecter la colonne 2."
+
+    # -----------------------------
+    # 3. Sélection colonne 2
+    # -----------------------------
+    col = sorted_columns[1]
+    col.sort(key=lambda e: e[1])  # top → bottom
+
+    # -----------------------------
+    # 4. FUSION DES LIGNES OCR
+    # -----------------------------
+    merged_lines = []
+    current_text = ""
+    last_y = None
+    Y_THRESHOLD = 18
 
-    if not words:
-        return "Aucun texte détecté"
+    for _, y, text in col:
+        if text.upper().startswith((
+            "DESIGNATION", "UNITE", "QUANT", "PRIX", "TOTAL", "LOT"
+        )):
+            continue
 
-    # Trier par X (colonnes)
-    words.sort(key=lambda w: w[1])
+        if last_y is None or abs(y - last_y) > Y_THRESHOLD:
+            if current_text:
+                merged_lines.append(current_text.strip())
+            current_text = text
+        else:
+            current_text += " " + text
 
-    # Regrouper par colonnes
-    columns = []
-    for word in words:
-        placed = False
-        for col in columns:
-            if abs(col[0][1] - word[1]) < 60:
-                col.append(word)
-                placed = True
-                break
-        if not placed:
-            columns.append([word])
+        last_y = y
 
-    if len(columns) < 2:
-        return "Moins de 2 colonnes détectées"
+    if current_text:
+        merged_lines.append(current_text.strip())
 
-    # 2ᵉ colonne
-    second_column = sorted(columns[1], key=lambda w: w[2])
+    # -----------------------------
+    # 5. 9 premières lignes
+    # -----------------------------
+    final_lines = merged_lines[:9]
 
-    return "\n".join([w[0] for w in second_column])
+    if not final_lines:
+        return "Colonne détectée mais lignes vides."
 
+    return "\n".join(final_lines)
 
+# -----------------------------
+# Interface Gradio
+# -----------------------------
 demo = gr.Interface(
-    fn=extract_second_column,
-    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)"
+    fn=extract_column2_9_lines,
+    inputs=gr.Image(type="pil", label="Image du devis"),
+    outputs=gr.Textbox(label="Colonne 2 – 9 premières lignes"),
+    title="Extraction de la colonne DESIGNATIONS",
+    description="Fusion automatique des lignes OCR (devis & tableaux)"
 )
 
-demo.launch()
+demo.launch(server_name="0.0.0.0", server_port=7860)