Update app.py

app.py

@@ -1,127 +1,163 @@
-# File: app.py
-import os
 import gradio as gr
+import easyocr
+import numpy as np
+from PIL import Image, ImageDraw, ImageFont
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 import torch
-from PIL import Image
-import pytesseract
-from transformers import MarianMTModel, MarianTokenizer
 
-class HindiSignboardTranslator:
+# Simplified Language Mapping
+LANG_MAP = {
+    'en': 'eng',
+    'hi': 'hin',
+    'mr': 'mar',
+    'fr': 'fra',
+    'de': 'deu',
+    'es': 'spa'
+}
+
+# Initialize OCR Reader with optimized languages
+ocr_reader = easyocr.Reader(['en', 'hi'], gpu=False)
+
+# Translation Model Cache
+class TranslationCache:
     def __init__(self):
-        # OCR Configuration
-        pytesseract.pytesseract.tesseract_cmd = r'/usr/bin/tesseract'  # Adjust path as needed
+        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]
+
+# Global translation cache
+translation_cache = TranslationCache()
+
+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'
+
+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]
         
-        # Translation Model
-        model_name = 'Helsinki-NLP/opus-mt-hi-en'
-        self.model = MarianMTModel.from_pretrained(model_name)
-        self.tokenizer = MarianTokenizer.from_pretrained(model_name)
+        # 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
+
+def process_image(image, target_lang):
+    """Optimized image processing with improved error handling"""
+    if image is None:
+        return "Please upload an image."
     
-    def extract_text(self, image):
-        """
-        Extract text from Hindi signboard image
-        Args:
-            image (PIL.Image): Input image
-        Returns:
-            str: Extracted Hindi text
-        """
+    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:
-            # Ensure image is in RGB mode
-            if image.mode != 'RGB':
-                image = image.convert('RGB')
+            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
             
-            # Extract text using Tesseract
-            hindi_text = pytesseract.image_to_string(image, lang='hin')
-            return hindi_text.strip()
-        except Exception as e:
-            print(f"OCR Error: {e}")
-            return None
-    
-    def translate_text(self, hindi_text):
-        """
-        Translate Hindi text to English
-        Args:
-            hindi_text (str): Input Hindi text
-        Returns:
-            str: Translated English text
-        """
-        try:
-            # Handle empty or None input
-            if not hindi_text:
-                return "No text detected"
+            # 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)
             
-            # Tokenize and translate
-            inputs = self.tokenizer(hindi_text, return_tensors="pt", padding=True)
-            outputs = self.model.generate(**inputs)
-            english_text = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
-            return english_text
-        except Exception as e:
-            print(f"Translation Error: {e}")
-            return "Translation failed"
+            # 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)
     
-    def translate_signboard(self, image):
-        """
-        Complete pipeline for signboard translation
-        Args:
-            image (PIL.Image): Signboard image
-        Returns:
-            dict: Translation results
-        """
-        # Validate input
-        if image is None:
-            return {
-                "status": "error",
-                "message": "No image provided",
-                "original_text": "",
-                "translated_text": ""
-            }
-        
-        # Extract text via OCR
-        hindi_text = self.extract_text(image)
-        
-        if not hindi_text:
-            return {
-                "status": "error",
-                "message": "Could not extract text from image",
-                "original_text": "",
-                "translated_text": ""
-            }
-        
-        # Translate to English
-        english_text = self.translate_text(hindi_text)
-        
-        return {
-            "status": "success",
-            "original_text": hindi_text,
-            "translated_text": english_text
-        }
-
-# Initialize the translator
-translator = HindiSignboardTranslator()
+    except Exception as e:
+        print(f"Processing error: {e}")
+        return f"An error occurred: {str(e)}"
 
 # Gradio Interface
-def translate_image(image):
-    """
-    Gradio-friendly translation function
-    """
-    if image is None:
-        return "", ""
+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
+        )
     
-    result = translator.translate_signboard(image)
-    return result['original_text'], result['translated_text']
-
-# Create Gradio Interface
-iface = gr.Interface(
-    fn=translate_image,
-    inputs=gr.Image(type="pil", label="Upload Hindi Signboard"),
-    outputs=[
-        gr.Textbox(label="Original Hindi Text"),
-        gr.Textbox(label="English Translation")
-    ],
-    title="Hindi Signboard Translator",
-    description="Upload a Hindi signboard image to extract and translate its text.",
-    # Removed example images
-)
+    return demo
 
 # Launch the app
+demo = create_interface()
+
 if __name__ == "__main__":
-    iface.launch()
+    demo.launch()