Update app.py

app.py

@@ -1,20 +1,20 @@
 import gradio as gr
 import easyocr
-import torch
 from transformers import pipeline
 import numpy as np
+from PIL import Image, ImageDraw
 
-# Load OCR model (English, Arabic, French, Chinese)
-reader = easyocr.Reader(['en', 'ar', 'fr', 'zh'])
+# Load OCR Reader (GPU-enabled)
+reader = easyocr.Reader(['en', 'ar', 'fr', 'zh'], gpu=True)  # Enable GPU if available
 
-# Load YOLOv5 model (small model for faster performance)
-model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)
+# Load Translation Model
+translator = pipeline("translation", model="Helsinki-NLP/opus-mt-mul-en", device=0)  # device=0 means use GPU
 
-# Load translation model (multilingual to English)
-translator = pipeline("translation", model="Helsinki-NLP/opus-mt-mul-en")
+# Load Object Detection Model
+detector = pipeline("object-detection", model="facebook/detr-resnet-50", device=0)
 
 def process_image(image):
-    # Convert PIL Image to numpy array if needed
+    # Ensure image is numpy
     if not isinstance(image, np.ndarray):
         image = np.array(image)
     
@@ -23,29 +23,41 @@ def process_image(image):
     extracted_texts = [res[1] for res in text_results]
     extracted_text = " ".join(extracted_texts)
 
-    # Step 2: Translate Text
+    # Step 2: Translation
     if extracted_text:
         translation = translator(extracted_text)[0]['translation_text']
     else:
         translation = "No text detected."
 
-    # Step 3: Object Detection - Bounding Boxes
-    results = model(image)
-    detected_img = results.render()[0]  # returns a list, take first image
+    # Step 3: Object Detection
+    detections = detector(Image.fromarray(image))
+    image_draw = Image.fromarray(image).convert("RGB")
+    draw = ImageDraw.Draw(image_draw)
 
-    return detected_img, extracted_text, translation
+    for det in detections:
+        box = det['box']
+        label = det['label']
+        score = det['score']
+        if score > 0.5:
+            draw.rectangle(
+                [box['xmin'], box['ymin'], box['xmax'], box['ymax']], 
+                outline="red", width=3
+            )
+            draw.text((box['xmin'], box['ymin'] - 10), f"{label} ({score:.2f})", fill="red")
 
-# Define Gradio Interface
+    return image_draw, extracted_text, translation
+
+# Gradio Interface
 iface = gr.Interface(
     fn=process_image,
-    inputs=gr.Image(type="pil", label="Upload Image (Signs, ID, License Plate)"),
+    inputs=gr.Image(type="pil", label="Upload an Image (e.g., Signs, IDs, License Plates)"),
     outputs=[
-        gr.Image(label="Detected Objects in Image"),
+        gr.Image(label="Detected Objects"),
         gr.Textbox(label="Extracted Text"),
-        gr.Textbox(label="Translated Text to English")
+        gr.Textbox(label="Translated Text (English)")
     ],
-    title="🚨 Police Smart Glasses AI Demo",
-    description="Upload an image to simulate smart glasses detecting text, translating it, and recognizing objects."
+    title="🚨 Police Smart Glasses - AI Demo (GPU Optimized)",
+    description="Upload an image to simulate smart glasses detecting text, translating it, and recognizing objects in real-time."
 )
 
 if __name__ == "__main__":