Update app.py

app.py

@@ -1,163 +1,64 @@
+# app.py
 import gradio as gr
 import easyocr
-import numpy as np
-from PIL import Image, ImageDraw, ImageFont
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
-import torch
+from transformers import MarianMTModel, MarianTokenizer
 
-# Simplified Language Mapping
-LANG_MAP = {
-    'en': 'eng',
-    'hi': 'hin',
-    'mr': 'mar',
-    'fr': 'fra',
-    'de': 'deu',
-    'es': 'spa'
+# OCR Reader Initialization
+reader = easyocr.Reader(['en', 'hi', 'fr', 'de', 'es', 'ru'], gpu=False)  # Add more if needed
+
+# Supported Languages for Translation
+LANGUAGE_CODES = {
+    "English": "en",
+    "Hindi": "hi",
+    "French": "fr",
+    "German": "de",
+    "Spanish": "es",
+    "Russian": "ru"
 }
 
-# Initialize OCR Reader with optimized languages
-ocr_reader = easyocr.Reader(['en', 'hi'], gpu=False)
+# Function to load MarianMT model
+model_cache = {}
+def get_model(target_lang):
+    model_name = f"Helsinki-NLP/opus-mt-ROMANCE-{target_lang}" if target_lang in ['fr', 'es', 'ro', 'pt'] else f"Helsinki-NLP/opus-mt-en-{target_lang}"
+    if model_name not in model_cache:
+        tokenizer = MarianTokenizer.from_pretrained(model_name)
+        model = MarianMTModel.from_pretrained(model_name)
+        model_cache[model_name] = (tokenizer, model)
+    return model_cache[model_name]
 
-# Translation Model Cache
-class TranslationCache:
-    def __init__(self):
-        self.models = {}
-        self.tokenizers = {}
-    
-    def get_model(self, src_lang, tgt_lang):
-        model_key = f"{src_lang}-{tgt_lang}"
-        
-        if model_key not in self.models:
-            try:
-                model_name = f"Helsinki-NLP/opus-mt-{src_lang}-{tgt_lang}"
-                tokenizer = AutoTokenizer.from_pretrained(model_name)
-                model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
-                
-                self.models[model_key] = model
-                self.tokenizers[model_key] = tokenizer
-            except Exception as e:
-                print(f"Error loading translation model {model_key}: {e}")
-                return None, None
-        
-        return self.models[model_key], self.tokenizers[model_key]
+# Main function
+def translate_image_text(image, target_lang):
+    try:
+        # OCR
+        result = reader.readtext(image, detail=0, paragraph=True)
+        extracted_text = " ".join(result)
 
-# Global translation cache
-translation_cache = TranslationCache()
+        if not extracted_text.strip():
+            return "No text found in the image."
 
-def detect_language(text):
-    """Attempt to detect language more accurately"""
-    # Simple language detection based on script
-    if any('\u0900' <= char <= '\u097F' for char in text):
-        return 'hi'
-    return 'en'
+        # Get model
+        code = LANGUAGE_CODES[target_lang]
+        tokenizer, model = get_model(code)
 
-def translate_text(text, src_lang, tgt_lang):
-    """Improved translation function with better error handling"""
-    try:
-        # Ensure language codes match model requirements
-        src_lang = src_lang.lower()[:2]
-        tgt_lang = tgt_lang.lower()[:2]
-        
-        # Get model and tokenizer
-        model, tokenizer = translation_cache.get_model(src_lang, tgt_lang)
-        
-        if not model or not tokenizer:
-            return text  # Fallback to original text if model fails
-        
-        # Prepare inputs
-        inputs = tokenizer(text, return_tensors="pt", max_length=512, truncation=True)
-        
-        # Generate translation
-        with torch.no_grad():
-            outputs = model.generate(**inputs)
-        
-        # Decode translation
-        translated = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        return translated
-    except Exception as e:
-        print(f"Translation error: {e}")
-        return text
+        # Translation
+        batch = tokenizer([extracted_text], return_tensors="pt", padding=True)
+        gen = model.generate(**batch)
+        translated = tokenizer.batch_decode(gen, skip_special_tokens=True)[0]
 
-def process_image(image, target_lang):
-    """Optimized image processing with improved error handling"""
-    if image is None:
-        return "Please upload an image."
-    
-    try:
-        # Convert image to numpy
-        image_np = np.array(image)
-        
-        # Perform OCR with confidence filtering
-        results = ocr_reader.readtext(image_np, threshold=0.3, low_text=0.4)
-        
-        if not results:
-            return "No clear text detected in the image."
-        
-        # Prepare PIL image for drawing
-        pil_img = Image.fromarray(image_np)
-        draw = ImageDraw.Draw(pil_img)
-        
-        # Use a more universal font
-        try:
-            font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", 20)
-        except IOError:
-            font = ImageFont.load_default()
-        
-        # Process each detected text
-        for detection in results:
-            bbox, text, confidence = detection
-            
-            # Detect source language
-            src_lang = detect_language(text)
-            
-            # Translate text
-            translated_text = translate_text(text, src_lang, target_lang)
-            
-            # Convert bbox to integers
-            bbox = np.array(bbox).astype(int)
-            
-            # Draw bounding box
-            draw.polygon(bbox.reshape(-1, 2).tolist(), outline='red', width=2)
-            
-            # Draw translated text
-            text_bbox = bbox[0]  # Top-left corner
-            draw.text((text_bbox[0], text_bbox[1] - 25), 
-                      translated_text, 
-                      fill='yellow', 
-                      font=font)
-        
-        return np.array(pil_img)
-    
+        return translated
     except Exception as e:
-        print(f"Processing error: {e}")
-        return f"An error occurred: {str(e)}"
+        return f"Error: {str(e)}"
 
 # Gradio Interface
-def create_interface():
-    with gr.Blocks() as demo:
-        gr.Markdown("# 🌍 TravelOCR: Multilingual Signboard Translator")
-        
-        with gr.Row():
-            image_input = gr.Image(type="pil", label="Upload Signboard Image")
-            lang_dropdown = gr.Dropdown(
-                label="Target Language",
-                choices=["en", "hi", "fr", "de", "es"],
-                value="en"
-            )
-        
-        translate_btn = gr.Button("Translate & Overlay")
-        output_img = gr.Image(label="Translated Output")
-        
-        translate_btn.click(
-            fn=process_image, 
-            inputs=[image_input, lang_dropdown], 
-            outputs=output_img
-        )
-    
-    return demo
-
-# Launch the app
-demo = create_interface()
+iface = gr.Interface(
+    fn=translate_image_text,
+    inputs=[
+        gr.Image(type="filepath", label="Upload Image"),
+        gr.Dropdown(choices=list(LANGUAGE_CODES.keys()), label="Translate To")
+    ],
+    outputs=gr.Textbox(label="Translated Text"),
+    title="Image Text Translator",
+    description="Upload an image containing text, and choose a language to translate the extracted text."
+)
 
-if __name__ == "__main__":
-    demo.launch()
+iface.launch()