Update app.py

app.py

@@ -13,6 +13,13 @@ def load_model():
 ort_session = load_model()
 
 def preprocess_image(image, target_size=(640, 640)):
+    # Convert PIL Image to numpy array if necessary
+    if isinstance(image, Image.Image):
+        image = np.array(image)
+    
+    # Convert RGB to BGR
+    image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
+    
     # Resize image
     image = cv2.resize(image, target_size)
     # Normalize
@@ -24,10 +31,24 @@ def preprocess_image(image, target_size=(640, 640)):
     return image
 
 def postprocess_results(output, image_shape, confidence_threshold=0.25, iou_threshold=0.45):
-    # Assuming YOLO v5 output format
-    boxes = output[0]
-    scores = output[1]
-    class_ids = output[2]
+    # Handle different possible output formats
+    if isinstance(output, (list, tuple)):
+        predictions = output[0]
+    elif isinstance(output, np.ndarray):
+        predictions = output
+    else:
+        raise ValueError(f"Unexpected output type: {type(output)}")
+
+    # Reshape if necessary
+    if len(predictions.shape) == 4:
+        predictions = predictions.squeeze((0, 1))
+    elif len(predictions.shape) == 3:
+        predictions = predictions.squeeze(0)
+
+    # Extract boxes, scores, and class_ids
+    boxes = predictions[:, :4]
+    scores = predictions[:, 4]
+    class_ids = predictions[:, 5]
 
     # Filter by confidence
     mask = scores > confidence_threshold
@@ -35,6 +56,13 @@ def postprocess_results(output, image_shape, confidence_threshold=0.25, iou_thre
     scores = scores[mask]
     class_ids = class_ids[mask]
 
+    # Convert boxes from [x, y, w, h] to [x1, y1, x2, y2]
+    boxes[:, 2:] += boxes[:, :2]
+
+    # Scale boxes to image size
+    boxes[:, [0, 2]] *= image_shape[1]
+    boxes[:, [1, 3]] *= image_shape[0]
+
     # Apply NMS
     indices = cv2.dnn.NMSBoxes(boxes.tolist(), scores.tolist(), confidence_threshold, iou_threshold)
 
@@ -43,12 +71,8 @@ def postprocess_results(output, image_shape, confidence_threshold=0.25, iou_thre
         box = boxes[i]
         score = scores[i]
         class_id = class_ids[i]
-        x, y, w, h = box
-        x1 = int(x * image_shape[1])
-        y1 = int(y * image_shape[0])
-        x2 = int((x + w) * image_shape[1])
-        y2 = int((y + h) * image_shape[0])
-        results.append((x1, y1, x2, y2, score, class_id))
+        x1, y1, x2, y2 = map(int, box)
+        results.append((x1, y1, x2, y2, float(score), int(class_id)))
 
     return results
 
@@ -68,7 +92,7 @@ def process_image(image):
         label = f"License Plate: {score:.2f}"
         cv2.putText(orig_image, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
     
-    return orig_image
+    return cv2.cvtColor(orig_image, cv2.COLOR_BGR2RGB)
 
 def process_video(video_path):
     cap = cv2.VideoCapture(video_path)
@@ -88,7 +112,7 @@ def process_video(video_path):
             break
         
         processed_frame = process_image(frame)
-        out.write(processed_frame)
+        out.write(cv2.cvtColor(processed_frame, cv2.COLOR_RGB2BGR))
     
     cap.release()
     out.release()
@@ -104,12 +128,10 @@ if uploaded_file is not None:
     
     if file_type == "image":
         image = Image.open(uploaded_file)
-        image = np.array(image)
-        
         st.image(image, caption="Uploaded Image", use_column_width=True)
         
         if st.button("Detect License Plates"):
-            processed_image = process_image(image)
+            processed_image = process_image(np.array(image))
             st.image(processed_image, caption="Processed Image", use_column_width=True)
     
     elif file_type == "video":