Update app.py

app.py

@@ -2,60 +2,64 @@ import re
 import gradio as gr
 import torch
 from transformers import DonutProcessor, VisionEncoderDecoderModel
-from transformers import AutoProcessor, AutoModelForVision2Seq
-
-
-processor = AutoProcessor.from_pretrained("debu-das/donut_receipt_v1.20")
-model = AutoModelForVision2Seq.from_pretrained("debu-das/donut_receipt_v1.20")
+import os
+from PIL import Image
 
+processor = DonutProcessor.from_pretrained("debu-das/donut_receipt_v1.20")
+model = VisionEncoderDecoderModel.from_pretrained("debu-das/donut_receipt_v1.20")
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 model.to(device)
 
-def process_documents(images):
+def process_document(image):
+    if isinstance(image, str):  # Si l'image est un chemin de fichier
+        image = Image.open(image).convert("RGB")
+    pixel_values = processor(image, return_tensors="pt").pixel_values
+    task_prompt = "<s_cord-v2>"
+    decoder_input_ids = processor.tokenizer(task_prompt, add_special_tokens=False, return_tensors="pt").input_ids
+    
+    outputs = model.generate(
+        pixel_values.to(device),
+        decoder_input_ids=decoder_input_ids.to(device),
+        max_length=model.decoder.config.max_position_embeddings,
+        early_stopping=True,
+        pad_token_id=processor.tokenizer.pad_token_id,
+        eos_token_id=processor.tokenizer.eos_token_id,
+        use_cache=True,
+        num_beams=1,
+        bad_words_ids=[[processor.tokenizer.unk_token_id]],
+        return_dict_in_generate=True,
+    )
+    
+    sequence = processor.batch_decode(outputs.sequences)[0]
+    sequence = sequence.replace(processor.tokenizer.eos_token, "").replace(processor.tokenizer.pad_token, "")
+    sequence = re.sub(r"<.*?>", "", sequence, count=1).strip()
+    
+    return processor.token2json(sequence)
+
+def process_batch(images):
     results = []
     for image in images:
-        # Prepare encoder inputs
-        pixel_values = processor(image, return_tensors="pt").pixel_values
-        
-        # Prepare decoder inputs
-        task_prompt = "<s_cord-v2>"
-        decoder_input_ids = processor.tokenizer(task_prompt, add_special_tokens=False, return_tensors="pt").input_ids
-          
-        # Generate answer
-        outputs = model.generate(
-            pixel_values.to(device),
-            decoder_input_ids=decoder_input_ids.to(device),
-            max_length=model.decoder.config.max_position_embeddings,
-            early_stopping=True,
-            pad_token_id=processor.tokenizer.pad_token_id,
-            eos_token_id=processor.tokenizer.eos_token_id,
-            use_cache=True,
-            num_beams=1,
-            bad_words_ids=[[processor.tokenizer.unk_token_id]],
-            return_dict_in_generate=True,
-        )
-        
-        # Postprocess
-        sequence = processor.batch_decode(outputs.sequences)[0]
-        sequence = sequence.replace(processor.tokenizer.eos_token, "").replace(processor.tokenizer.pad_token, "")
-        sequence = re.sub(r"<.*?>", "", sequence, count=1).strip()  # Remove first task start token
-        
-        results.append(processor.token2json(sequence))
-    
+        result = process_document(image)
+        results.append(result)
     return results
 
-description = "Gradio Demo for Donut, an instance of `VisionEncoderDecoderModel` fine-tuned on CORD (document parsing). To use it, simply upload multiple images and click 'submit'."
-article = "<p style='text-align: center'><a href='https://arxiv.org/abs/2111.15664' target='_blank'>Donut: OCR-free Document Understanding Transformer</a> | <a href='https://github.com/clovaai/donut' target='_blank'>Github Repo</a></p>"
+description = "Démo Gradio pour Donut, une instance du modèle `VisionEncoderDecoderModel` affiné sur CORD (analyse de documents). Pour l'utiliser, téléchargez une ou plusieurs images et cliquez sur `Submit`, ou cliquez sur l'un des exemples pour les charger."
+article = "Cette application permet maintenant de traiter plusieurs images de tickets de caisse à la fois."
 
 demo = gr.Interface(
-    fn=process_documents,
-    inputs=gr.Files(label="Upload Images"),  # Use Files to handle multiple files
+    fn=process_batch,
+    inputs=gr.File(file_count="multiple", type="filepath"),
     outputs="json",
-    title="Batch Demo: Donut 🍩 for Document Parsing",
+    title="Reconnaissance des tickets de caisse en lot 🧾",
     description=description,
     article=article,
-    examples=[["example.png"], ["example_1.png"],["example_2.png"], ["example_3.png"],["example_4.png"]],
-    cache_examples=False)
+    examples=[
+        [["example.jpg", "example_1.jpg"]],
+        [["example_2.jpg", "example_3.jpg", "example_4.jpg"]]
+    ],
+    cache_examples=False
+)
 
-demo.launch()
+port = int(os.environ.get("PORT", 7860))
+demo.launch(server_name="0.0.0.0", server_port=port)