restored

app.py

@@ -4,11 +4,6 @@ from PIL import Image
 import pytesseract
 import re
 import json
-import os
-
-# Global variable to store the loaded model
-loaded_generator = None
-current_model = None
 
 def clean_ocr_text(raw_text):
     """Clean OCR text using regex - only for cleaning, not extraction"""
@@ -26,45 +21,12 @@ def clean_ocr_text(raw_text):
     
     return cleaned.strip()
 
-def load_model(model_choice):
-    """Load the AI model only when needed"""
-    global loaded_generator, current_model
-    
-    model_map = {
-        "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"
-    }
-    
-    selected_model = model_map.get(model_choice, "google/flan-t5-base")
-    
-    # Only load if different model is requested
-    if current_model != selected_model:
-        print(f"Loading model: {selected_model}")
-        try:
-            loaded_generator = pipeline("text2text-generation", model=selected_model)
-            current_model = selected_model
-            print(f"Model {selected_model} loaded successfully")
-        except Exception as e:
-            print(f"Error loading model {selected_model}: {e}")
-            # Fallback to base model
-            if selected_model != "google/flan-t5-base":
-                print("Falling back to base model...")
-                loaded_generator = pipeline("text2text-generation", model="google/flan-t5-base")
-                current_model = "google/flan-t5-base"
-            else:
-                raise e
-    
-    return loaded_generator
-
-def extract_dl_info(image, model_choice):
+def extract_dl_info(image):
     """
     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
     
     Returns:
         tuple: (raw_ocr_text, cleaned_ocr_text, json_result)
@@ -86,13 +48,8 @@ def extract_dl_info(image, model_choice):
         # 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:
-            generator = load_model(model_choice)
-        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)
+        # Load a text generation pipeline for structured extraction
+        generator = pipeline("text2text-generation", model="google/flan-t5-base")
         
         # 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:"
@@ -174,19 +131,6 @@ with gr.Blocks(title="Driver's License Information Extractor") as demo:
                 type="pil",
                 height=400
             )
-            
-            model_choice = gr.Radio(
-                choices=[
-                    "FLAN-T5 Small (60M)",
-                    "FLAN-T5 Base (220M)", 
-                    "FLAN-T5 Large (770M)",
-                    "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"
-            )
-            
             submit_btn = gr.Button("Extract Information", variant="primary")
         
         with gr.Column():
@@ -208,9 +152,9 @@ with gr.Blocks(title="Driver's License Information Extractor") as demo:
     
     submit_btn.click(
         fn=extract_dl_info,
-        inputs=[image_input, model_choice],
+        inputs=[image_input],
         outputs=[raw_ocr, cleaned_text, json_result]
     )
 
 if __name__ == "__main__":
-    demo.launch(server_name="127.0.0.1", server_port=7860, share=False)
+    demo.launch(share=True)