Update app.py

app.py

@@ -1,29 +1,26 @@
 import gradio as gr
-import torch
 import cv2
+import easyocr
 import numpy as np
+import os
 from PIL import Image
-from paddleocr import PaddleOCR  # Import PaddleOCR
 from ultralytics import YOLO
+from datetime import datetime
 
-# Load model
+# Load YOLO model
 model = YOLO("/home/user/app/best.pt")
 
 # Label map
 label_map = {0: "Analog", 1: "Digital", 2: "Non-LP"}
 
-# Initialize PaddleOCR (for Bangla OCR)
-ocr = PaddleOCR(use_angle_cls=True, lang='en')  # For Bangla language
+# EasyOCR Bengali
+reader = easyocr.Reader(['bn'])
 
-def process_frame(frame):
-    # Resize to YOLO input shape
+def annotate_frame(frame):
     input_img = cv2.resize(frame, (640, 640))
     results = model(input_img)[0]
     detections = results.boxes.data.cpu().numpy()
 
-    extracted_texts = []
-    confidences = []
-
     for det in detections:
         if len(det) < 6:
             continue
@@ -33,54 +30,73 @@ def process_frame(frame):
         label = label_map.get(int(cls), "Unknown")
         percent = f"{conf * 100:.2f}%"
 
-        # Draw box and label on image
+        # Draw box and label
         cv2.rectangle(input_img, (x1, y1), (x2, y2), (255, 0, 0), 2)
         cv2.putText(input_img, f"{label}: {percent}", (x1, y1 - 10),
                     cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)
 
-        # OCR using PaddleOCR
-        cropped = frame[y1:y2, x1:x2]  # Use original frame for OCR
+        # OCR
+        cropped = frame[y1:y2, x1:x2]
         if cropped.size > 0:
-            # Convert to RGB and run OCR
-            result = ocr.ocr(cropped, cls=True)
-            for line in result[0]:
-                extracted_texts.append(line[1])  # Get the detected text
-            confidences.append(percent)
-
-    # Convert to PIL
-    annotated = cv2.cvtColor(input_img, cv2.COLOR_BGR2RGB)
-    pil_img = Image.fromarray(annotated)
-
-    return pil_img, "\n".join(extracted_texts), ", ".join(confidences)
+            ocr_result = reader.readtext(cropped)
+            for i, item in enumerate(ocr_result):
+                text = item[1].strip()
+                conf = item[2]
+                cv2.putText(input_img, text, (x1, y2 + 20 + i*25),
+                            cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 255), 2)
 
+    return cv2.cvtColor(input_img, cv2.COLOR_BGR2RGB)
 
 def process_input(input_file):
     file_path = input_file.name
+    ext = os.path.splitext(file_path)[-1].lower()
 
-    if file_path.endswith(('.mp4', '.avi', '.mov')):
+    if ext in ['.mp4', '.avi', '.mov']:
         cap = cv2.VideoCapture(file_path)
-        ret, frame = cap.read()
+        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+        fps = cap.get(cv2.CAP_PROP_FPS)
+        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+
+        # Output path
+        timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
+        output_path = f"annotated_{timestamp}.mp4"
+        out = cv2.VideoWriter(output_path, fourcc, fps, (640, 640))
+
+        while cap.isOpened():
+            ret, frame = cap.read()
+            if not ret:
+                break
+            annotated = annotate_frame(frame)
+            annotated_resized = cv2.resize(annotated, (640, 640))
+            out.write(cv2.cvtColor(annotated_resized, cv2.COLOR_RGB2BGR))
+
         cap.release()
-        if not ret:
-            return None, "Couldn't read video", ""
+        out.release()
+
+        return output_path, "", ""
+
     else:
+        # Image case
         frame = cv2.imread(file_path)
         if frame is None:
             return None, "Invalid image", ""
 
-    return process_frame(frame)
+        annotated = annotate_frame(frame)
+        pil_img = Image.fromarray(annotated)
+        return pil_img, "", ""
 
 
 interface = gr.Interface(
     fn=process_input,
     inputs=gr.File(type="filepath", label="Upload Image or Video"),
     outputs=[
-        gr.Image(type="pil", label="Detected Output"),
+        gr.Video(label="Output Video or Image") | gr.Image(type="pil", label="Output"),
         gr.Textbox(label="Detected Text (Bangla)"),
         gr.Textbox(label="Confidence (%)")
     ],
-    title="YOLOv10n License Plate Detector (Bangla)",
-    description="Upload an image or video. Detects license plates and extracts Bangla text using PaddleOCR."
+    title="YOLOv5 License Plate Detector (Bangla OCR)",
+    description="Upload an image or video. Detects license plates and extracts Bangla text using EasyOCR."
 )
 
 interface.launch()