Update app.py

app.py

@@ -4,13 +4,14 @@ import cv2
 import pytesseract
 import numpy as np
 from PIL import Image
-from transformers import DetrImageProcessor, TableTransformerForObjectDetection
+from transformers import DetrImageProcessor, TableTransformerForObjectDetection, TableTransformerForStructureRecognition
 
 # ===============================
 # Chargement des modèles
 # ===============================
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
+# Modèle de détection de tableau
 det_processor = DetrImageProcessor.from_pretrained(
     "microsoft/table-transformer-detection"
 )
@@ -18,10 +19,11 @@ det_model = TableTransformerForObjectDetection.from_pretrained(
     "microsoft/table-transformer-detection"
 ).to(DEVICE)
 
+# Modèle de reconnaissance de structure (cellules)
 struct_processor = DetrImageProcessor.from_pretrained(
     "microsoft/table-transformer-structure-recognition"
 )
-struct_model = TableTransformerForObjectDetection.from_pretrained(
+struct_model = TableTransformerForStructureRecognition.from_pretrained(
     "microsoft/table-transformer-structure-recognition"
 ).to(DEVICE)
 
@@ -60,26 +62,44 @@ def extract_description(image_pil):
     if not tables:
         return "❌ Aucun tableau détecté", ""
 
+    # Extraire premier tableau détecté
     table_box = tables[0].int().tolist()
     x0, y0, x1, y1 = table_box
     table_img = image[y0:y1, x0:x1]
 
+    # ---- Optionnel : vérifier visuellement le tableau ----
+    # Image.fromarray(table_img).show()
+
+    # ---- Redimensionner le tableau pour la structure ----
+    max_size = 1024
+    scale = max_size / max(table_img.shape[:2])
+    new_w, new_h = int(table_img.shape[1]*scale), int(table_img.shape[0]*scale)
+    table_resized = cv2.resize(table_img, (new_w, new_h))
+
     # ---- Structure du tableau ----
-    inputs = struct_processor(images=table_img, return_tensors="pt")
+    inputs = struct_processor(images=table_resized, return_tensors="pt")
     inputs = {k: v.to(DEVICE) for k, v in inputs.items()}
     outputs = struct_model(**inputs)
 
     results = struct_processor.post_process_object_detection(
         outputs,
-        threshold=0.7,
-        target_sizes=[table_img.shape[:2]]
+        threshold=0.5,  # seuil plus bas pour capturer plus de cellules
+        target_sizes=[table_resized.shape[:2]]
     )[0]
 
     cells = []
     for box, label in zip(results["boxes"], results["labels"]):
         label_name = struct_model.config.id2label[label.item()]
         if label_name == "table cell":
-            cells.append(box.int().tolist())
+            # Remettre les coordonnées à l'échelle originale
+            scale_x = table_img.shape[1] / table_resized.shape[1]
+            scale_y = table_img.shape[0] / table_resized.shape[0]
+            x0c, y0c, x1c, y1c = box.int().tolist()
+            x0c = int(x0c * scale_x)
+            x1c = int(x1c * scale_x)
+            y0c = int(y0c * scale_y)
+            y1c = int(y1c * scale_y)
+            cells.append([x0c, y0c, x1c, y1c])
 
     if not cells:
         return "❌ Aucune cellule détectée", ""