Update app.py

app.py

@@ -1,13 +1,12 @@
 import gradio as gr
 from transformers import TrOCRProcessor, VisionEncoderDecoderModel
 from PIL import Image
-import re
 import torch
 
 # ===============================
 # Charger le modèle TrOCR public
 # ===============================
-model_name = "microsoft/trocr-base-handwritten"  # modèle public
+model_name = "microsoft/trocr-base-handwritten"  # modèle OCR général
 processor = TrOCRProcessor.from_pretrained(model_name)
 model = VisionEncoderDecoderModel.from_pretrained(model_name)
 
@@ -26,31 +25,21 @@ def extract_description(image_pil):
     # Séparer le texte en lignes
     lines = [line.strip() for line in ocr_text.split("\n") if line.strip()]
 
-    # Identifier la colonne Description
-    desc_col = []
-    header_found = False
-    headers = []
+    # Détection de la colonne Description via mot-clé
+    desc_lines = []
+    found_header = False
 
-    if lines:
-        first_line = lines[0]
-        # Split en colonnes par tabulation ou espaces multiples
-        headers = re.split(r"\t+|\s{2,}", first_line)
-        try:
-            desc_index = next(i for i, h in enumerate(headers) if "description" in h.lower())
-            header_found = True
-        except StopIteration:
-            desc_index = None
+    for line in lines:
+        if found_header:
+            # toutes les lignes après le header sont considérées comme contenu de la colonne
+            desc_lines.append(line)
+        elif "description" in line.lower():
+            found_header = True
 
-    # Extraire les valeurs sous la colonne Description
-    if header_found:
-        for line in lines[1:]:
-            cols = re.split(r"\t+|\s{2,}", line)
-            if desc_index is not None and desc_index < len(cols):
-                desc_col.append(cols[desc_index])
-    else:
+    if not desc_lines:
         return "❌ Colonne 'Description' non trouvée", ocr_text
-
-    return "\n".join(desc_col), ocr_text
+    else:
+        return "\n".join(desc_lines), ocr_text
 
 # ===============================
 # Interface Gradio