Update app.py

app.py

@@ -12,20 +12,26 @@ def load_model():
 
 ort_session = load_model()
 
-# Define class names and their corresponding indices
-CLASS_NAMES = {0: 'car', 1: 'license_plate'}
-
 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
     image = image.astype(np.float32) / 255.0
+    # Transpose for ONNX input
     image = np.transpose(image, (2, 0, 1))
+    # Add batch dimension
     image = np.expand_dims(image, axis=0)
     return image
 
 def postprocess_results(output, image_shape, confidence_threshold=0.25, iou_threshold=0.45):
+    # Handle different possible output formats
     if isinstance(output, (list, tuple)):
         predictions = output[0]
     elif isinstance(output, np.ndarray):
@@ -33,33 +39,40 @@ def postprocess_results(output, image_shape, confidence_threshold=0.25, iou_thre
     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
     boxes = boxes[mask]
     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)
 
     results = []
     for i in indices:
         box = boxes[i]
         score = scores[i]
-        class_id = int(class_ids[i])
+        class_id = class_ids[i]
         x1, y1, x2, y2 = map(int, box)
-        results.append((x1, y1, x2, y2, float(score), class_id))
+        results.append((x1, y1, x2, y2, float(score), int(class_id)))
 
     return results
 
@@ -67,26 +80,29 @@ def process_image(image):
     orig_image = image.copy()
     processed_image = preprocess_image(image)
     
+    # Run inference
     inputs = {ort_session.get_inputs()[0].name: processed_image}
     outputs = ort_session.run(None, inputs)
     
     results = postprocess_results(outputs, image.shape)
     
+    # Draw bounding boxes on the image
     for x1, y1, x2, y2, score, class_id in results:
-        color = (0, 255, 0) if CLASS_NAMES.get(class_id, 'unknown') == 'car' else (255, 0, 0)
-        cv2.rectangle(orig_image, (x1, y1), (x2, y2), color, 2)
-        label = f"{CLASS_NAMES.get(class_id, 'unknown')}: {score:.2f}"
-        cv2.putText(orig_image, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, color, 2)
+        cv2.rectangle(orig_image, (x1, y1), (x2, y2), (0, 255, 0), 2)
+        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 cv2.cvtColor(orig_image, cv2.COLOR_BGR2RGB)
 
 def process_video(video_path):
     cap = cv2.VideoCapture(video_path)
     
+    # Get video properties
     width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
     height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
     fps = int(cap.get(cv2.CAP_PROP_FPS))
     
+    # Create a temporary file to store the processed video
     temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.mp4')
     out = cv2.VideoWriter(temp_file.name, cv2.VideoWriter_fourcc(*'mp4v'), fps, (width, height))
     
@@ -103,7 +119,7 @@ def process_video(video_path):
     
     return temp_file.name
 
-st.title("Vehicle and License Plate Detection")
+st.title("License Plate Detection")
 
 uploaded_file = st.file_uploader("Choose an image or video file", type=["jpg", "jpeg", "png", "mp4"])
 
@@ -114,7 +130,7 @@ if uploaded_file is not None:
         image = Image.open(uploaded_file)
         st.image(image, caption="Uploaded Image", use_column_width=True)
         
-        if st.button("Detect Objects"):
+        if st.button("Detect License Plates"):
             processed_image = process_image(np.array(image))
             st.image(processed_image, caption="Processed Image", use_column_width=True)
     
@@ -124,8 +140,8 @@ if uploaded_file is not None:
         
         st.video(tfile.name)
         
-        if st.button("Detect Objects"):
+        if st.button("Detect License Plates"):
             processed_video = process_video(tfile.name)
             st.video(processed_video)
 
-st.write("Upload an image or video to detect vehicles and license plates.")
+st.write("Upload an image or video to detect license plates.")