app

app.py

@@ -34,8 +34,7 @@ def load_model(model_choice):
         "FLAN-T5 Small (60M)": "google/flan-t5-small",
         "FLAN-T5 Base (220M)": "google/flan-t5-base", 
         "FLAN-T5 Large (770M)": "google/flan-t5-large",
-        "FLAN-T5 XL (3B)": "google/flan-t5-xl",
-        "FLAN-T5 XXL (11B)": "google/flan-t5-xxl"
+        "FLAN-T5 XL (3B)": "google/flan-t5-xl"
     }
     
     selected_model = model_map.get(model_choice, "google/flan-t5-base")
@@ -59,61 +58,41 @@ def load_model(model_choice):
     
     return loaded_generator
 
-def extract_dl_info(image, model_choice, manual_text, use_manual_text):
+def extract_dl_info(image, model_choice):
     """
     Extracts driver's license information from an image file using OCR and AI model processing.
     
     Args:
         image: PIL Image object from Gradio
         model_choice: Selected model from radio button
-        manual_text: Manual text input from user
-        use_manual_text: Boolean to determine if manual text should be used instead of OCR
     
     Returns:
-        tuple: (raw_ocr_text, cleaned_ocr_text, json_result, status_message)
+        tuple: (raw_ocr_text, cleaned_ocr_text, json_result)
     """
     try:
-        # Determine text source
-        if use_manual_text and manual_text and manual_text.strip():
-            # Use manual text input
-            print("Using manual text input...")
-            raw_text = manual_text.strip()
-            cleaned_text = clean_ocr_text(raw_text)
-            print(f"Manual text: {raw_text}")
-            print(f"Cleaned text: {cleaned_text}")
-            status = "Using manual text input"
-        else:
-            # Use OCR on image
-            if image is None:
-                return "No image provided", "No image provided", json.dumps({"error": "No image provided"}, indent=2), "Error: No image provided"
-            
-            # Convert Gradio image to RGB
-            if hasattr(image, 'convert'):
-                image = image.convert("RGB")
-            
-            # Use Tesseract OCR to extract text
-            print("Extracting text with Tesseract OCR...")
-            status = "Extracting text with OCR..."
-            raw_text = pytesseract.image_to_string(image)
-            print(f"Raw OCR text: {raw_text}")
-            
-            # Clean the OCR text using regex (only for cleaning, not extraction)
-            cleaned_text = clean_ocr_text(raw_text)
-            print(f"Cleaned OCR text: {cleaned_text}")
-            status = "OCR completed, processing with AI..."
+        # Convert Gradio image to RGB
+        if hasattr(image, 'convert'):
+            image = image.convert("RGB")
+        
+        # Use Tesseract OCR to extract text
+        print("Extracting text with Tesseract OCR...")
+        raw_text = pytesseract.image_to_string(image)
+        print(f"Raw OCR text: {raw_text}")
+        
+        # Clean the OCR text using regex (only for cleaning, not extraction)
+        cleaned_text = clean_ocr_text(raw_text)
+        print(f"Cleaned OCR text: {cleaned_text}")
         
         # Use AI model to process cleaned OCR text and extract structured information
         print("Processing cleaned OCR text with AI model...")
         
         # Load model only when needed
         try:
-            status = "Loading AI model..."
             generator = load_model(model_choice)
-            status = "AI model loaded, extracting information..."
         except Exception as e:
             error_msg = f"Failed to load AI model: {str(e)}"
             print(error_msg)
-            return raw_text, cleaned_text, json.dumps({"error": error_msg}, indent=2), f"Error: {error_msg}"
+            return raw_text, cleaned_text, json.dumps({"error": error_msg}, indent=2)
         
         # Create specific prompts for the AI model to extract individual fields
         name_prompt = f"From this driver's license text, what is the person's full name? Text: {cleaned_text} Answer:"
@@ -176,13 +155,12 @@ def extract_dl_info(image, model_choice, manual_text, use_manual_text):
         
         # Format JSON result
         json_result = json.dumps(extracted, indent=2)
-        status = "Extraction completed successfully!"
         
-        return raw_text, cleaned_text, json_result, status
+        return raw_text, cleaned_text, json_result
     
     except Exception as e:
         error_msg = f"Processing failed: {str(e)}"
-        return error_msg, error_msg, json.dumps({"error": error_msg}, indent=2), f"Error: {error_msg}"
+        return error_msg, error_msg, json.dumps({"error": error_msg}, indent=2)
 
 # Create Gradio interface
 with gr.Blocks(title="Driver's License Information Extractor") as demo:
@@ -202,34 +180,14 @@ with gr.Blocks(title="Driver's License Information Extractor") as demo:
                     "FLAN-T5 Small (60M)",
                     "FLAN-T5 Base (220M)", 
                     "FLAN-T5 Large (770M)",
-                    "FLAN-T5 XL (3B)",
-                    "FLAN-T5 XXL (11B)"
+                    "FLAN-T5 XL (3B)"
                 ],
                 value="FLAN-T5 Base (220M)",
                 label="Select AI Model",
                 info="Larger models are more accurate but slower to load and process"
             )
             
-            use_manual_text = gr.Checkbox(
-                label="Use Manual Text Input",
-                value=False,
-                info="Check this to type text manually instead of using OCR"
-            )
-            
-            manual_text_input = gr.Textbox(
-                label="Manual Text Input",
-                lines=6,
-                placeholder="Type or paste the driver's license text here...",
-                visible=False
-            )
-            
             submit_btn = gr.Button("Extract Information", variant="primary")
-            status_text = gr.Textbox(
-                label="Status",
-                value="Ready to process",
-                interactive=False,
-                visible=True
-            )
         
         with gr.Column():
             raw_ocr = gr.Textbox(
@@ -248,20 +206,10 @@ with gr.Blocks(title="Driver's License Information Extractor") as demo:
                 label="Extracted Information (JSON)"
             )
     
-    # Show/hide manual text input based on checkbox
-    def toggle_manual_input(use_manual):
-        return gr.update(visible=use_manual)
-    
-    use_manual_text.change(
-        fn=toggle_manual_input,
-        inputs=[use_manual_text],
-        outputs=[manual_text_input]
-    )
-    
     submit_btn.click(
         fn=extract_dl_info,
-        inputs=[image_input, model_choice, manual_text_input, use_manual_text],
-        outputs=[raw_ocr, cleaned_text, json_result, status_text]
+        inputs=[image_input, model_choice],
+        outputs=[raw_ocr, cleaned_text, json_result]
     )
 
 if __name__ == "__main__":