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)