selection

app.py

@@ -5,6 +5,10 @@ import pytesseract
 import re
 import json
 
+# 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"""
     # Remove extra whitespace and normalize
@@ -21,35 +25,88 @@ def clean_ocr_text(raw_text):
     
     return cleaned.strip()
 
-def extract_dl_info(image):
+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, manual_text, use_manual_text):
     """
     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)
     """
     try:
-        # 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}")
+        # 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}")
+        else:
+            # Use OCR on image
+            if image is None:
+                return "No image provided", "No image provided", json.dumps({"error": "No image provided"}, indent=2)
+            
+            # 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 a text generation pipeline for structured extraction
-        generator = pipeline("text2text-generation", model="google/flan-t5-base")
+        # 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)
         
         # 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:"
@@ -131,6 +188,32 @@ 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"
+            )
+            
+            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")
         
         with gr.Column():
@@ -150,9 +233,19 @@ 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],
+        inputs=[image_input, model_choice, manual_text_input, use_manual_text],
         outputs=[raw_ocr, cleaned_text, json_result]
     )