VanguardAI/Arabic-OCR
2
1import spaces2import gradio as gr3import torch4from PIL import Image5from qwen_vl_utils import process_vision_info6from transformers import Qwen2VLForConditionalGeneration, AutoProcessor, AutoTokenizer7from transformers import Qwen2VLProcessor, Qwen2VLImageProcessor8import traceback9import json10import os11 12# ========================================13# AIN VLM MODEL FOR OCR14# ========================================15 16# Model configuration17MODEL_ID = "MBZUAI/AIN"18 19# Image resolution settings for the processor20# The default range for the number of visual tokens per image in the model is 4-1638421# These settings balance speed and memory usage22MIN_PIXELS = 256 * 28 * 28 # Minimum resolution23MAX_PIXELS = 1280 * 28 * 28 # Maximum resolution24 25# Global model and processor26model = None27processor = None28 29# Strict OCR-focused prompt30OCR_PROMPT = """Extract all text from this image exactly as it appears. 31 32Requirements:331. Extract ONLY the text content - do not describe, analyze, or interpret the image342. Maintain the original text structure, layout, and formatting353. Preserve line breaks, paragraphs, and spacing as they appear364. Do not translate the text - keep it in its original language375. Do not add any explanations, descriptions, or additional commentary386. If there are tables, maintain their structure397. If there are headers, titles, or sections, preserve their hierarchy40 41Output only the extracted text, nothing else."""42 43 44def ensure_model_loaded():45 """Lazily load the AIN VLM model and processor."""46 global model, processor47 48 if model is not None and processor is not None:49 return50 51 print("🔄 Loading AIN VLM model...")52 53 try:54 # Determine device and dtype55 if torch.cuda.is_available():56 device_map = "auto"57 torch_dtype = "auto"58 print("✅ Using GPU (CUDA)")59 else:60 device_map = "cpu"61 torch_dtype = torch.float3262 print("✅ Using CPU")63 64 # Load model65 loaded_model = Qwen2VLForConditionalGeneration.from_pretrained(66 MODEL_ID,67 torch_dtype=torch_dtype,68 device_map=device_map,69 trust_remote_code=True,70 )71 72 # Load processor with proper configuration73 # Manual construction to avoid size parameter issues74 try:75 # First, try the standard way76 loaded_processor = AutoProcessor.from_pretrained(77 MODEL_ID,78 trust_remote_code=True,79 )80 print("✅ Processor loaded successfully (standard method)")81 except ValueError as e:82 if "size must contain 'shortest_edge' and 'longest_edge' keys" in str(e):83 print("⚠️ Standard processor loading failed, trying manual construction...")84 # Manually construct processor with correct size format85 try:86 # Load tokenizer separately87 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)88 89 # Create image processor with correct size format90 image_processor = Qwen2VLImageProcessor(91 size={"shortest_edge": 224, "longest_edge": 1120}, # Valid format92 do_resize=True,93 do_rescale=True,94 do_normalize=True,95 )96 97 # Create processor from components98 loaded_processor = Qwen2VLProcessor(99 image_processor=image_processor,100 tokenizer=tokenizer,101 )102 print("✅ Processor loaded successfully (manual construction)")103 except Exception as manual_error:104 print(f"❌ Manual construction also failed: {manual_error}")105 raise106 else:107 raise108 109 model = loaded_model110 processor = loaded_processor111 112 print("✅ Model loaded successfully!")113 114 except Exception as e:115 print(f"❌ Error loading model: {e}")116 traceback.print_exc()117 raise118 119 120@spaces.GPU(duration=100)121def extract_text_from_image(122 image: Image.Image, 123 custom_prompt: str = None, 124 max_new_tokens: int = 2048,125 min_pixels: int = None,126 max_pixels: int = None127) -> str:128 """129 Extract text from image using AIN VLM model.130 131 Args:132 image: PIL Image to process133 custom_prompt: Optional custom prompt (uses default OCR prompt if None)134 max_new_tokens: Maximum tokens to generate135 min_pixels: Minimum image resolution (optional)136 max_pixels: Maximum image resolution (optional)137 138 Returns:139 Extracted text as string140 """141 try:142 # Ensure model is loaded143 ensure_model_loaded()144 145 if model is None or processor is None:146 return "❌ Error: Model not loaded. Please refresh and try again."147 148 # Use custom prompt or default OCR prompt149 prompt_to_use = custom_prompt if custom_prompt and custom_prompt.strip() else OCR_PROMPT150 151 # Use custom resolution settings if provided, otherwise use defaults152 min_pix = min_pixels if min_pixels else MIN_PIXELS153 max_pix = max_pixels if max_pixels else MAX_PIXELS154 155 # Prepare messages in the format expected by the model156 # Include min_pixels and max_pixels in the image content for proper resizing157 messages = [158 {159 "role": "user",160 "content": [161 {162 "type": "image",163 "image": image,164 "min_pixels": min_pix,165 "max_pixels": max_pix,166 },167 {168 "type": "text",169 "text": prompt_to_use170 },171 ],172 }173 ]174 175 # Apply chat template176 text = processor.apply_chat_template(177 messages,178 tokenize=False,179 add_generation_prompt=True180 )181 182 # Process vision information183 image_inputs, video_inputs = process_vision_info(messages)184 185 # Prepare inputs186 inputs = processor(187 text=[text],188 images=image_inputs,189 videos=video_inputs,190 padding=True,191 return_tensors="pt",192 )193 194 # Move to device195 device = next(model.parameters()).device196 inputs = inputs.to(device)197 198 # Generate output199 with torch.no_grad():200 generated_ids = model.generate(201 **inputs,202 max_new_tokens=max_new_tokens,203 do_sample=False, # Greedy decoding for consistency204 )205 206 # Decode output207 generated_ids_trimmed = [208 out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)209 ]210 211 output_text = processor.batch_decode(212 generated_ids_trimmed,213 skip_special_tokens=True,214 clean_up_tokenization_spaces=False215 )216 217 result = output_text[0] if output_text else ""218 219 return result.strip() if result else "No text extracted"220 221 except Exception as e:222 error_msg = f"❌ Error during text extraction: {str(e)}"223 print(error_msg)224 traceback.print_exc()225 return error_msg226 227 228def create_gradio_interface():229 """Create the Gradio interface for AIN OCR."""230 231 # Custom CSS for better UI232 css = """233 .main-container {234 max-width: 1400px;235 margin: 0 auto;236 padding: 20px;237 }238 239 .header-text {240 text-align: center;241 color: #2c3e50;242 margin-bottom: 30px;243 }244 245 .process-button {246 background: linear-gradient(135deg, #667eea 0%, #764ba2 100%) !important;247 border: none !important;248 color: white !important;249 font-weight: bold !important;250 font-size: 1.1em !important;251 padding: 12px 24px !important;252 width: 100% !important;253 margin-top: 10px !important;254 }255 256 .process-button:hover {257 transform: translateY(-2px) !important;258 box-shadow: 0 6px 12px rgba(0,0,0,0.2) !important;259 }260 261 /* Larger font for extracted text */262 .output-textbox textarea {263 font-size: 20px !important;264 line-height: 2.0 !important;265 font-family: 'Segoe UI', 'Tahoma', 'Traditional Arabic', 'Arabic Typesetting', sans-serif !important;266 padding: 24px !important;267 direction: auto !important;268 text-align: start !important;269 }270 271 .output-textbox {272 background: #ffffff;273 border: 2px solid #e0e0e0;274 border-radius: 8px;275 box-shadow: 0 2px 8px rgba(0,0,0,0.1);276 }277 278 /* Better Arabic text support */279 .output-textbox textarea[dir="rtl"] {280 text-align: right !important;281 direction: rtl !important;282 }283 284 .info-box {285 background: #e3f2fd;286 border-left: 4px solid #2196f3;287 padding: 15px;288 margin: 10px 0;289 border-radius: 4px;290 }291 292 /* Status box styling */293 .status-box {294 background: #f0f4f8;295 border: 1px solid #d0dae6;296 border-radius: 6px;297 padding: 12px;298 margin-top: 10px;299 text-align: center;300 font-size: 14px;301 }302 303 /* Better spacing for rows and columns */304 .gradio-container {305 gap: 20px !important;306 }307 308 .contain {309 gap: 15px !important;310 }311 312 /* Image preview styling */313 .image-preview {314 border: 2px solid #e0e0e0;315 border-radius: 8px;316 box-shadow: 0 2px 8px rgba(0,0,0,0.1);317 }318 319 /* Accordion styling */320 .accordion {321 background: #f8f9fa;322 border-radius: 8px;323 margin-top: 15px;324 padding: 5px;325 }326 327 /* Clear button */328 button[variant="secondary"] {329 width: 100% !important;330 margin-top: 10px !important;331 }332 333 /* Label styling */334 label {335 font-weight: 600 !important;336 margin-bottom: 8px !important;337 }338 339 /* Better component spacing */340 .gr-form {341 gap: 12px !important;342 }343 344 /* Example images styling */345 .gr-examples {346 margin-top: 15px;347 }348 """349 350 with gr.Blocks(theme=gr.themes.Soft(), css=css, title="AIN VLM OCR") as demo:351 352 # Header353 gr.HTML("""354 <div class="header-text">355 <h1>🔍 AIN VLM - Vision Language Model OCR</h1>356 <p style="font-size: 1.1em; color: #6b7280; margin-top: 10px;">357 Advanced OCR using Vision Language Model (VLM) for accurate text extraction358 </p>359 <p style="font-size: 0.95em; color: #9ca3af; margin-top: 8px;">360 Powered by <strong>MBZUAI/AIN</strong> - Specialized for understanding and extracting text from images361 </p>362 </div>363 """)364 365 # Info box366 gr.Markdown("""367 <div class="info-box">368 <strong>ℹ️ How it works:</strong> Upload an image containing text, click "Process Image", and get the extracted text.369 The VLM model intelligently understands context and can handle handwritten text better than traditional OCR models.370 </div>371 """)372 373 # Main interface374 with gr.Row(equal_height=False):375 # Left column - Input376 with gr.Column(scale=1, min_width=400):377 # Image input378 image_input = gr.Image(379 label="📸 Upload Image",380 type="pil",381 height=400,382 elem_classes=["image-preview"]383 )384 385 # Advanced settings386 with gr.Accordion("⚙️ Advanced Settings", open=False, elem_classes=["accordion"]):387 custom_prompt = gr.Textbox(388 label="Custom Prompt (Optional)",389 placeholder="Leave empty to use default OCR prompt...",390 lines=3,391 info="Customize the prompt if you want specific extraction behavior"392 )393 394 max_tokens = gr.Slider(395 minimum=512,396 maximum=4096,397 value=2048,398 step=128,399 label="Max Tokens",400 info="Maximum length of extracted text"401 )402 403 gr.Markdown("**📐 Image Resolution Settings**")404 gr.Markdown("*Controls visual token range (4-16384) - balance quality vs speed*")405 406 with gr.Row():407 min_pixels_input = gr.Number(408 value=MIN_PIXELS,409 label="Min Pixels",410 info=f"Default: {MIN_PIXELS:,} (~{MIN_PIXELS//1000}k)",411 precision=0412 )413 max_pixels_input = gr.Number(414 value=MAX_PIXELS,415 label="Max Pixels",416 info=f"Default: {MAX_PIXELS:,} (~{MAX_PIXELS//1000}k)",417 precision=0418 )419 420 show_prompt_btn = gr.Button("👁️ Show Default Prompt", size="sm", variant="secondary")421 422 # Process button423 process_btn = gr.Button(424 "🚀 Process Image",425 variant="primary",426 elem_classes=["process-button"],427 size="lg"428 )429 430 # Clear button431 clear_btn = gr.Button("🗑️ Clear All", variant="secondary", size="lg")432 433 # Right column - Output434 with gr.Column(scale=1, min_width=500):435 # Text output with larger font436 text_output = gr.Textbox(437 label="📝 Extracted Text",438 placeholder="Extracted text will appear here...",439 lines=18,440 max_lines=22,441 show_copy_button=True,442 interactive=False,443 elem_classes=["output-textbox"],444 container=True,445 )446 447 # Status/info448 status_output = gr.Markdown(449 value="✨ *Ready to process images*",450 elem_classes=["status-box"]451 )452 453 # Examples section454 with gr.Row():455 with gr.Column():456 gr.Markdown("### 📚 Example Images")457 gr.Markdown("*Click on any example below to load it*")458 gr.Examples(459 examples=[460 ["image/app/1762329983969.png"],461 ["image/app/1762330009302.png"],462 ["image/app/1762330020168.png"],463 ],464 inputs=image_input,465 label="",466 examples_per_page=3467 )468 469 # Default prompt display470 default_prompt_display = gr.Textbox(471 label="Default OCR Prompt",472 value=OCR_PROMPT,473 lines=10,474 visible=False,475 interactive=False476 )477 478 # Event handlers479 def process_image_handler(image, custom_prompt_text, max_tokens_value, min_pix, max_pix):480 """Handle image processing."""481 if image is None:482 return "", "⚠️ Please upload an image first."483 484 try:485 status = "⏳ Processing image..."486 extracted_text = extract_text_from_image(487 image,488 custom_prompt=custom_prompt_text,489 max_new_tokens=int(max_tokens_value),490 min_pixels=int(min_pix) if min_pix else None,491 max_pixels=int(max_pix) if max_pix else None492 )493 494 if extracted_text and not extracted_text.startswith("❌"):495 status = f"✅ Text extracted successfully! ({len(extracted_text)} characters)"496 else:497 status = "⚠️ No text extracted or error occurred."498 499 return extracted_text, status500 501 except Exception as e:502 error_msg = f"❌ Error: {str(e)}"503 return error_msg, "❌ Processing failed."504 505 def clear_all_handler():506 """Clear all inputs and outputs."""507 return None, "", "", "✨ Ready to process images"508 509 def toggle_prompt_display(current_visible):510 """Toggle the visibility of the default prompt."""511 return gr.update(visible=not current_visible)512 513 # Wire up events514 process_btn.click(515 process_image_handler,516 inputs=[image_input, custom_prompt, max_tokens, min_pixels_input, max_pixels_input],517 outputs=[text_output, status_output]518 )519 520 clear_btn.click(521 clear_all_handler,522 outputs=[image_input, text_output, custom_prompt, status_output]523 )524 525 # Show/hide default prompt526 show_prompt_btn.click(527 lambda: gr.update(visible=True),528 outputs=[default_prompt_display]529 )530 531 return demo532 533 534if __name__ == "__main__":535 # Create and launch the interface536 demo = create_gradio_interface()537 demo.queue(max_size=10).launch(538 server_name="0.0.0.0",539 server_port=7860,540 share=False,541 debug=True,542 show_error=True543 )544 545