Update app.py

app.py

@@ -5,14 +5,11 @@ import onnxruntime as ort
 from PIL import Image
 import tempfile
 
-# Dynamically assign colors to classes
-def get_color(class_id):
-    """Generate a color for any class ID"""
-    np.random.seed(class_id)  # For consistent colors
-    return tuple(map(int, np.random.randint(0, 255, 3)))
-
-# Class labels - will be populated dynamically
-CLASSES = {}
+# Define class labels
+CLASSES = {
+    0: "Vehicle",
+    1: "License_Plate"
+}
 
 # Load the ONNX model
 @st.cache_resource
@@ -41,11 +38,6 @@ def postprocess_results(output, original_shape, confidence_threshold=0.25, iou_t
     else:
         raise ValueError(f"Unexpected output type: {type(output)}")
 
-    # Debug: Print the shape and first few entries of predictions
-    st.write(f"Debug - Predictions shape: {predictions.shape}")
-    if len(predictions) > 0:
-        st.write(f"Debug - First prediction entry: {predictions[0]}")
-
     if len(predictions.shape) == 4:
         predictions = predictions.squeeze((0, 1))
     elif len(predictions.shape) == 3:
@@ -56,15 +48,6 @@ def postprocess_results(output, original_shape, confidence_threshold=0.25, iou_t
     scores = predictions[:, 4]
     class_ids = predictions[:, 5]
 
-    # Debug: Print unique class IDs
-    unique_classes = np.unique(class_ids)
-    st.write(f"Debug - Unique class IDs found: {unique_classes}")
-    
-    # Update CLASSES dictionary with any new class IDs
-    for class_id in unique_classes:
-        if int(class_id) not in CLASSES:
-            CLASSES[int(class_id)] = f"Class_{int(class_id)}"
-
     # Filter by confidence
     mask = scores > confidence_threshold
     boxes = boxes[mask]
@@ -113,32 +96,36 @@ def process_image(image):
     
     # Draw bounding boxes on the image
     for x1, y1, x2, y2, score, class_id in results:
-        # Get color dynamically
-        color = get_color(class_id)
-        cv2.rectangle(orig_image, (x1, y1), (x2, y2), color, 2)
+        # Draw rectangle with white color
+        cv2.rectangle(orig_image, (x1, y1), (x2, y2), (255, 255, 255), 2)
+        
+        # Get class name
+        class_name = CLASSES.get(class_id, f"Class_{class_id}")
+        label = f"{class_name}: {score:.2f}"
         
-        label = f"{CLASSES[class_id]}: {score:.2f}"
-        # Calculate text size for better positioning
+        # Add label background and text
         (text_width, text_height), _ = cv2.getTextSize(
-            label, cv2.FONT_HERSHEY_SIMPLEX, 0.9, 2
+            label, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 1
         )
-        # Draw background rectangle for text
+        
+        # Draw black background for text
         cv2.rectangle(
             orig_image, 
-            (x1, y1 - text_height - 10), 
+            (x1, y1 - text_height - 4), 
             (x1 + text_width, y1), 
-            color, 
+            (0, 0, 0), 
             -1
         )
-        # Draw text
+        
+        # Draw white text
         cv2.putText(
             orig_image,
             label,
             (x1, y1 - 5),
             cv2.FONT_HERSHEY_SIMPLEX,
-            0.9,
+            0.6,
             (255, 255, 255),
-            2
+            1
         )
     
     return cv2.cvtColor(orig_image, cv2.COLOR_BGR2RGB)
@@ -180,7 +167,7 @@ def process_video(video_path):
     return temp_file.name
 
 # Streamlit UI
-st.title("Object Detection")
+st.title("Vehicle and License Plate Detection")
 
 # Add confidence threshold slider
 confidence_threshold = st.slider(
@@ -219,16 +206,4 @@ if uploaded_file is not None:
                 processed_video = process_video(tfile.name)
                 st.video(processed_video)
 
-# Add legend after processing to include all detected classes
-if CLASSES:
-    st.markdown("### Detection Legend")
-    for class_id, class_name in CLASSES.items():
-        color = get_color(class_id)
-        st.markdown(
-            f'<div style="display: flex; align-items: center;">'
-            f'<div style="width: 20px; height: 20px; background-color: rgb{color}; margin-right: 10px;"></div>'
-            f'<span>{class_name}</span></div>',
-            unsafe_allow_html=True
-        )
-
-st.write("Upload an image or video to detect objects.")
+st.write("Upload an image or video to detect vehicles and license plates.")