Update app.py

app.py

@@ -4,65 +4,52 @@ import cv2
 import pytesseract
 import numpy as np
 from PIL import Image
-import sys
-import os
-
 from ultralytics import YOLO
 
 # Load model
-model = YOLO("/home/user/app/best.pt")  # আপনি যেই path এ best.pt রেখেছেন
-
+model = YOLO("/home/user/app/best.pt")
 
+# Label map
+label_map = {0: "Analog", 1: "Digital", 2: "Non-LP"}
 
-    
-# Frame processing function
 def process_frame(frame):
-    # Resize image to 640x640
-    frame_resized = cv2.resize(frame, (640, 640))
-    
-    img = cv2.cvtColor(frame_resized, cv2.COLOR_BGR2RGB)
-    img_tensor = torch.from_numpy(img).permute(2, 0, 1).float() / 255.0
-    img_tensor = img_tensor.unsqueeze(0)
-
-    # Run inference with the YOLO model (no need to manually apply nms)
-    results = model(img_tensor, augment=False)
-    
-    # Extract results (list of detections)
-    detections = results.xywh[0]  # YOLO's detection results
+    # Resize to YOLO input shape
+    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:
-        x_center, y_center, width, height, conf, cls = det.tolist()
-        if conf > 0.5:
-            # Convert from YOLO format to bounding box format
-            x1 = int((x_center - width / 2) * 640)
-            y1 = int((y_center - height / 2) * 640)
-            x2 = int((x_center + width / 2) * 640)
-            y2 = int((y_center + height / 2) * 640)
-            cls = int(cls)
-
-            label_map = {0: "Analog", 1: "Digital", 2: "Non-LP"}
-            label = label_map.get(cls, "Unknown")
-            percent = f"{conf * 100:.2f}%"
-
-            # Draw box & label
-            cv2.rectangle(frame, (x1, y1), (x2, y2), (255, 0, 0), 2)
-            cv2.putText(frame, f"{label}: {percent}", (x1, y1 - 10),
-                        cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)
-
-            # OCR
-            lp_crop = frame[y1:y2, x1:x2]
-            gray = cv2.cvtColor(lp_crop, cv2.COLOR_BGR2GRAY)
+        if len(det) < 6:
+            continue
+
+        x1, y1, x2, y2, conf, cls = det
+        x1, y1, x2, y2 = map(int, [x1, y1, x2, y2])
+        label = label_map.get(int(cls), "Unknown")
+        percent = f"{conf * 100:.2f}%"
+
+        # Draw box and label on image
+        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
+        cropped = frame[y1:y2, x1:x2]  # Use original frame for OCR
+        if cropped.size > 0:
+            gray = cv2.cvtColor(cropped, cv2.COLOR_BGR2GRAY)
             text = pytesseract.image_to_string(gray, config="--psm 6 -l ben")
             extracted_texts.append(text.strip())
             confidences.append(percent)
 
-    return frame, "\n".join(extracted_texts), ", ".join(confidences)
+    # 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)
 
 
-# Input handler
 def process_input(input_file):
     file_path = input_file.name
 
@@ -77,9 +64,8 @@ def process_input(input_file):
         if frame is None:
             return None, "Invalid image", ""
 
-    processed_frame, text, confidence = process_frame(frame)
-    processed_pil = Image.fromarray(cv2.cvtColor(processed_frame, cv2.COLOR_BGR2RGB))
-    return processed_pil, text, confidence
+    return process_frame(frame)
+
 
 interface = gr.Interface(
     fn=process_input,
@@ -90,7 +76,7 @@ interface = gr.Interface(
         gr.Textbox(label="Confidence (%)")
     ],
     title="YOLOv10n License Plate Detector (Bangla)",
-    description="Upload an image or video. Detects plates and extracts Bangla text using OCR (CPU)."
+    description="Upload an image or video. Detects license plates and extracts Bangla text using OCR."
 )
 
 interface.launch()