model

app.py

@@ -1,33 +1,35 @@
+from transformers import DonutProcessor, VisionEncoderDecoderModel
+import torch
 import gradio as gr
 import fitz  # PyMuPDF
 from PIL import Image
 import io
+import json
 
+# ====== Load Donut model (no fine-tuning needed) ======
+MODEL_ID = "naver-clova-ix/donut-base-finetuned-docvqa"
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+processor = DonutProcessor.from_pretrained(MODEL_ID)
+model = VisionEncoderDecoderModel.from_pretrained(MODEL_ID).to(device)
+
+
+# ====== Convert PDF to Image ======
 def convert_pdf_to_image(pdf_file):
     """
-    Convert PDF to image(s) and return the first page as PIL Image
+    Convert first page of PDF to PIL Image and return both image and info.
     """
     if pdf_file is None:
         return "Please upload a PDF file first.", None
     
     try:
-        # Open the PDF file
         pdf_document = fitz.open(pdf_file.name)
-        
-        # Get PDF info before closing
         total_pages = len(pdf_document)
-        
-        # Get the first page
         first_page = pdf_document[0]
-        
-        # Convert page to image (pixmap)
         pix = first_page.get_pixmap(matrix=fitz.Matrix(2, 2))  # 2x zoom for better quality
-        
-        # Convert pixmap to PIL Image
         img_data = pix.tobytes("png")
         img = Image.open(io.BytesIO(img_data))
-        
-        # Get PDF info
+
         pdf_info = f"""
         **PDF Conversion Results:**
         - Total pages: {total_pages}
@@ -35,38 +37,86 @@ def convert_pdf_to_image(pdf_file):
         - Image size: {img.width} x {img.height} pixels
         - Format: PNG
         """
-        
-        # Close the PDF document
+
         pdf_document.close()
-        
         return pdf_info, img
-        
+
     except Exception as e:
         return f"Error processing PDF: {str(e)}", None
 
-# Create the Gradio interface
-with gr.Blocks(title="PDF to Image Converter") as demo:
-    gr.Markdown("# 📄 PDF to Image Converter")
-    gr.Markdown("Upload a PDF file to convert the first page to an image!")
-    
+
+# ====== Process Image with Donut ======
+def process_image(image):
+    if image is None:
+        return "No image found."
+
+    try:
+        # Preprocess image
+        pixel_values = processor(image, return_tensors="pt").pixel_values.to(device)
+        task_prompt = "<s_docvqa><s_question>Extract all driver’s license fields.<s_answer>"
+        decoder_input_ids = processor.tokenizer(
+            task_prompt, add_special_tokens=False, return_tensors="pt"
+        ).input_ids.to(device)
+
+        # Generate text output
+        outputs = model.generate(
+            pixel_values,
+            decoder_input_ids=decoder_input_ids,
+            max_length=512,
+            num_beams=4,
+        )
+        result = processor.batch_decode(outputs, skip_special_tokens=True)[0]
+
+        # Try parsing as JSON
+        try:
+            data = json.loads(result)
+            formatted = json.dumps(data, indent=2)
+        except Exception:
+            formatted = result  # fallback to raw text
+
+        return formatted
+
+    except Exception as e:
+        return f"Error during extraction: {str(e)}"
+
+
+# ====== Gradio Interface ======
+with gr.Blocks(title="Driver’s License Info Extractor") as demo:
+    gr.Markdown("# 🪪 Driver’s License Info Extractor (Prototype)")
+    gr.Markdown(
+        "Upload a **PDF of a driver’s license**. The app converts it to an image, "
+        "then uses a **Donut document understanding model** to extract key details."
+    )
+
     with gr.Row():
         with gr.Column():
             pdf_input = gr.File(
-                label="Upload PDF File", 
+                label="Upload Driver’s License PDF",
                 file_types=[".pdf"],
-                file_count="single"
+                file_count="single",
             )
-            convert_btn = gr.Button("Convert to Image", variant="primary")
+            convert_btn = gr.Button("1️⃣ Convert PDF to Image", variant="primary")
+            extract_btn = gr.Button("2️⃣ Extract Information", variant="primary")
         
         with gr.Column():
             conversion_output = gr.Markdown(label="Conversion Results")
             converted_image = gr.Image(label="Converted Image (First Page)")
-    
+            extraction_output = gr.Textbox(
+                label="Extracted License Info (JSON/Text)", lines=10
+            )
+
+    # Link buttons
     convert_btn.click(
         fn=convert_pdf_to_image,
         inputs=[pdf_input],
-        outputs=[conversion_output, converted_image]
+        outputs=[conversion_output, converted_image],
+    )
+
+    extract_btn.click(
+        fn=process_image,
+        inputs=[converted_image],
+        outputs=[extraction_output],
     )
 
 if __name__ == "__main__":
-    demo.launch(share=True)
+    demo.launch(share=True)

requirements.txt

@@ -1,3 +1,5 @@
-gradio==5.49.0
-PyMuPDF==1.23.8
-Pillow==10.0.1
+torch
+transformers
+gradio
+pymupdf
+pillow