Update app.py

app.py

@@ -5,17 +5,14 @@ import onnxruntime as ort
 from PIL import Image
 import tempfile
 
-# Class labels for both vehicles and license plates
-CLASSES = {
-    0: "Vehicle",
-    1: "License_Plate"
-}
-
-# Different colors for each class
-COLORS = {
-    0: (255, 0, 0),  # Red for vehicles
-    1: (0, 255, 0)   # Green for license plates
-}
+# 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 = {}
 
 # Load the ONNX model
 @st.cache_resource
@@ -44,6 +41,11 @@ 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:
@@ -54,6 +56,15 @@ 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]
@@ -102,7 +113,8 @@ def process_image(image):
     
     # Draw bounding boxes on the image
     for x1, y1, x2, y2, score, class_id in results:
-        color = COLORS[class_id]
+        # Get color dynamically
+        color = get_color(class_id)
         cv2.rectangle(orig_image, (x1, y1), (x2, y2), color, 2)
         
         label = f"{CLASSES[class_id]}: {score:.2f}"
@@ -146,7 +158,6 @@ def process_video(video_path):
         (width, height)
     )
     
-    # Add progress bar for video processing
     total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
     progress_bar = st.progress(0)
     frame_count = 0
@@ -159,7 +170,6 @@ def process_video(video_path):
         processed_frame = process_image(frame)
         out.write(cv2.cvtColor(processed_frame, cv2.COLOR_RGB2BGR))
         
-        # Update progress bar
         frame_count += 1
         progress_bar.progress(frame_count / total_frames)
     
@@ -170,7 +180,7 @@ def process_video(video_path):
     return temp_file.name
 
 # Streamlit UI
-st.title("Vehicle and License Plate Detection")
+st.title("Object Detection")
 
 # Add confidence threshold slider
 confidence_threshold = st.slider(
@@ -209,15 +219,16 @@ if uploaded_file is not None:
                 processed_video = process_video(tfile.name)
                 st.video(processed_video)
 
-# Add legend
-st.markdown("### Detection Legend")
-for class_id, class_name in CLASSES.items():
-    color = COLORS[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
-    )
+# 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 vehicles and license plates.")
+st.write("Upload an image or video to detect objects.")