app.py

import gradio as gr
from transformers import pipeline
from PIL import Image
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
    cleaned = re.sub(r'\s+', ' ', raw_text)
    
    # Remove common OCR artifacts
    cleaned = re.sub(r'[^\w\s\.,;:/\-@#$%&*()]', ' ', cleaned)
    
    # Normalize common OCR mistakes
    cleaned = re.sub(r'\b([A-Z])\s+([A-Z])\b', r'\1\2', cleaned)  # Fix split letters
    
    # Clean up date patterns
    cleaned = re.sub(r'(\d{2})\s*/\s*(\d{2})\s*/\s*(\d{4})', r'\1/\2/\3', cleaned)
    
    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:
        try:
            loaded_generator = pipeline("text2text-generation", model=selected_model)
            current_model = selected_model
        except Exception as e:
            # Fallback to base model
            if selected_model != "google/flan-t5-base":
                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:
        # Determine text source
        if use_manual_text and manual_text and manual_text.strip():
            # Use manual text input
            raw_text = manual_text.strip()
            cleaned_text = clean_ocr_text(raw_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
            raw_text = pytesseract.image_to_string(image)
            
            # Clean the OCR text using regex (only for cleaning, not extraction)
            cleaned_text = clean_ocr_text(raw_text)
        
        # Use AI model to process cleaned OCR text and extract structured information
        
        # 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)}"
            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:"
        address_prompt = f"From this driver's license text, what is the street address? Text: {cleaned_text} Answer:"
        license_prompt = f"From this driver's license text, what is the license number (the long number)? Text: {cleaned_text} Answer:"
        dob_prompt = f"From this driver's license text, what is the date of birth? Text: {cleaned_text} Answer:"
        exp_prompt = f"From this driver's license text, what is the expiration date? Text: {cleaned_text} Answer:"
        state_prompt = f"From this driver's license text, what is the state abbreviation? Text: {cleaned_text} Answer:"
        
        # Extract each field separately using AI
        try:
            name_result = generator(name_prompt, max_length=50, do_sample=False)
            name = name_result[0]['generated_text'].strip()
            
            address_result = generator(address_prompt, max_length=50, do_sample=False)
            address = address_result[0]['generated_text'].strip()
            
            license_result = generator(license_prompt, max_length=50, do_sample=False)
            license_number = license_result[0]['generated_text'].strip()
            
            dob_result = generator(dob_prompt, max_length=50, do_sample=False)
            date_of_birth = dob_result[0]['generated_text'].strip()
            
            exp_result = generator(exp_prompt, max_length=50, do_sample=False)
            expiration_date = exp_result[0]['generated_text'].strip()
            
            state_result = generator(state_prompt, max_length=10, do_sample=False)
            state = state_result[0]['generated_text'].strip()
            
            # Clean up the AI responses
            def clean_response(response, max_length=100):
                if not response or response == "Not found":
                    return "Not found"
                # Take only the first part if it's too long
                if len(response) > max_length:
                    response = response[:max_length].split()[0] if response.split() else "Not found"
                return response.strip()
            
            # Create the extracted data
            extracted = {
                "name": clean_response(name, 50),
                "address": clean_response(address, 100),
                "license_number": clean_response(license_number, 30),
                "date_of_birth": clean_response(date_of_birth, 20),
                "expiration_date": clean_response(expiration_date, 20),
                "state": clean_response(state, 10)
            }
            
        except Exception as e:
            extracted = {
                "name": "AI extraction failed",
                "address": "AI extraction failed", 
                "license_number": "AI extraction failed",
                "date_of_birth": "AI extraction failed",
                "expiration_date": "AI extraction failed",
                "state": "AI extraction failed",
                "error": f"AI model error: {str(e)}"
            }
        
        # Format JSON result
        json_result = json.dumps(extracted, indent=2)
        
        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)

# Create Gradio interface
with gr.Blocks(title="Driver's License Information Extractor") as demo:
    gr.Markdown("# 🆔 Driver's License Information Extractor")
    gr.Markdown("Upload a driver's license image to extract information using OCR and AI processing. The system will show raw OCR text, cleaned text, and structured JSON results.")
    
    with gr.Row():
        with gr.Column():
            image_input = gr.Image(
                label="Upload Driver's License Image", 
                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():
            raw_ocr = gr.Textbox(
                label="Raw OCR Text", 
                lines=8, 
                interactive=False,
                placeholder="Raw text extracted by Tesseract OCR will appear here..."
            )
            cleaned_text = gr.Textbox(
                label="Cleaned OCR Text", 
                lines=8, 
                interactive=False,
                placeholder="Cleaned and normalized text will appear here..."
            )
            json_result = gr.JSON(
                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]
    )

if __name__ == "__main__":
    demo.launch(share=True)