Add confidence threshold slider and update class labels

- Add confidence_threshold slider to UI (0.01-0.9 range)
- Update detect() function to accept dynamic confidence threshold
- Change from 43 classes to 29 aggregated traffic sign categories
- Improve NMS IOU threshold from 0.45 to 0.55
- Update config default confidence to 0.30

app.py

@@ -8,10 +8,11 @@ import io
 # Load the detector
 detector = TrafficSignDetector('config.yaml')
 
-def detect_traffic_signs(image):
+def detect_traffic_signs(image, confidence_threshold):
     """
     Process the uploaded image and return the image with detected signs.
     :param image: PIL Image or numpy array
+    :param confidence_threshold: confidence threshold from slider
     :return: tuple of (detected image, preprocessed image)
     """
     # Redirect stdout to capture all logs
@@ -25,8 +26,8 @@ def detect_traffic_signs(image):
     image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
     print(f"Converted to BGR, shape: {image.shape}")
 
-    # Perform detection (returns tuple of (detected_image, preprocessed_image))
-    result_image, preprocessed_image = detector.detect(image)
+    # Perform detection with the slider's confidence threshold
+    result_image, preprocessed_image = detector.detect(image, confidence_threshold=confidence_threshold)
 
     # Convert back to RGB for Gradio
     result_image = cv2.cvtColor(result_image, cv2.COLOR_BGR2RGB)
@@ -46,10 +47,20 @@ with gr.Blocks(title="Traffic Sign Detector") as demo:
     with gr.Row():
         preprocessed_image = gr.Image(label="Preprocessed Image (640x640, Letterboxed)")
     
+    with gr.Row():
+        confidence_threshold = gr.Slider(
+            minimum=0.01,
+            maximum=0.9,
+            value=0.30,
+            step=0.01,
+            label="Confidence Threshold",
+            info="Lower values show more detections (less confident). Adjust to find optimal balance."
+        )
+    
     detect_btn = gr.Button("Detect Traffic Signs")
     detect_btn.click(
         fn=detect_traffic_signs, 
-        inputs=input_image, 
+        inputs=[input_image, confidence_threshold], 
         outputs=[output_image, preprocessed_image],
         queue=True  # Enable queue to ensure logs are shown
     )

config.yaml

@@ -1,6 +1,6 @@
 model:
   path: 'VietCat/GTSRB-Model/models/GTSRB.pt' # Path to the YOLO model on Hugging Face Hub (will be downloaded automatically)
-  confidence_threshold: 0.001 # Minimum confidence for detections (very low for testing)
+  confidence_threshold: 0.30 # Minimum confidence for detections (0.3 filters most false positives)
 
 inference:
   box_color: (128, 0, 128) # Purple color for bounding boxes (BGR format)
@@ -8,46 +8,32 @@ inference:
   thickness: 2 # Thickness of bounding box lines
 
 classes:
-  - 'Speed limit 20'
-  - 'Speed limit 30'
-  - 'Speed limit 50'
-  - 'Speed limit 60'
-  - 'Speed limit 70'
-  - 'Speed limit 80'
-  - 'Restriction ends 80'
-  - 'Speed limit 100'
-  - 'Speed limit 120'
-  - 'No overtaking'
-  - 'No overtaking trucks'
-  - 'Priority at next intersection'
-  - 'Priority road'
-  - 'Give way'
-  - 'Stop'
-  - 'No traffic both ways'
-  - 'No trucks'
-  - 'No entry'
-  - 'Danger'
-  - 'Bend left'
-  - 'Bend right'
-  - 'Bend'
-  - 'Uneven road'
-  - 'Slippery road'
-  - 'Road narrows'
-  - 'Construction'
-  - 'Traffic signal'
-  - 'Pedestrian crossing'
-  - 'School crossing'
-  - 'Cycles crossing'
-  - 'Snow'
-  - 'Animals'
-  - 'Restriction ends'
-  - 'Go right'
-  - 'Go left'
-  - 'Go straight'
-  - 'Go right or straight'
-  - 'Go left or straight'
-  - 'Keep right'
-  - 'Keep left'
-  - 'Roundabout'
-  - 'Restriction ends overtaking'
-  - 'Restriction ends overtaking trucks'
+  - 'one_way_prohibition'
+  - 'no_parking'
+  - 'no_stopping_and_parking'
+  - 'no_turn_left'
+  - 'no_turn_right'
+  - 'no_u_turn'
+  - 'no_u_and_left_turn'
+  - 'no_u_and_right_turn'
+  - 'no_motorbike_entry_turning'
+  - 'no_car_entry_turning'
+  - 'no_truck_entry_turning'
+  - 'other_prohibition'
+  - 'indication'
+  - 'direction'
+  - 'speed_limit'
+  - 'weight_limit'
+  - 'height_limit'
+  - 'pedestrian_crossing'
+  - 'intersection_danger'
+  - 'road_danger'
+  - 'pedestrian_danger'
+  - 'construction_danger'
+  - 'slow_warning'
+  - 'other_warning'
+  - 'vehicle_permission_lane'
+  - 'vehicle_and_speed_permission_lane'
+  - 'overpass_route'
+  - 'no_more_prohibition'
+  - 'other'

model.py

@@ -149,12 +149,18 @@ class TrafficSignDetector:
         print(f"Image format: {image.dtype}, Min: {image.min()}, Max: {image.max()}, Mean: {image.mean():.1f}")
         return image
 
-    def detect(self, image):
+    def detect(self, image, confidence_threshold=None):
         """
         Perform inference on the image and draw bounding boxes.
         :param image: numpy array of the image
+        :param confidence_threshold: optional override for confidence threshold
         :return: tuple of (image with drawn bounding boxes, preprocessed image for visualization)
         """
+        # Use provided threshold or fall back to config value
+        if confidence_threshold is None:
+            confidence_threshold = self.conf_threshold
+        else:
+            confidence_threshold = float(confidence_threshold)
         print(f"\n{'='*80}")
         print(f"DETECTION PIPELINE START")
         print(f"{'='*80}")
@@ -191,11 +197,11 @@ class TrafficSignDetector:
         # Use iou_threshold for NMS (Non-Maximum Suppression) to remove overlapping boxes
         print(f"\n[STEP 4] MODEL INFERENCE")
         print(f"  - Input shape to model: {image.shape}")
-        print(f"  - Confidence threshold: {self.conf_threshold}")
-        print(f"  - IOU threshold: 0.45")
+        print(f"  - Confidence threshold: {confidence_threshold}")
+        print(f"  - IOU threshold: 0.55")
         
         # Run with conf=0.0 to get raw predictions (before filtering)
-        results_raw = self.model(image, conf=0.0, imgsz=640, iou=0.45)
+        results_raw = self.model(image, conf=0.0, imgsz=640, iou=0.55)
         raw_box_count = len(results_raw[0].boxes) if results_raw else 0
         print(f"  - Raw detections (conf=0.0): {raw_box_count}")
         
@@ -208,8 +214,8 @@ class TrafficSignDetector:
             print(f"  - Confidences > 0.0001: {sum(1 for c in all_raw_confs if c > 0.0001)}")
         
         # Now run with actual threshold
-        results = self.model(image, conf=self.conf_threshold, imgsz=640, iou=0.45)
-        print(f"  - Filtered detections (conf={self.conf_threshold}): {len(results)}")
+        results = self.model(image, conf=confidence_threshold, imgsz=640, iou=0.55)
+        print(f"  - Filtered detections (conf={confidence_threshold}): {len(results)}")
 
         # Get original dimensions for coordinate transformation
         orig_h, orig_w = original_image.shape[:2]
@@ -249,7 +255,7 @@ class TrafficSignDetector:
                 print(f"Detected: {self.classes[cls]} with conf {conf:.4f} at ({x1},{y1})-({x2},{y2})")
 
                 # Only draw if confidence meets threshold
-                if conf >= self.conf_threshold:
+                if conf >= confidence_threshold:
                     # Draw bounding box on original image
                     cv2.rectangle(original_image, (x1, y1), (x2, y2), self.box_color, self.thickness)
 
@@ -279,6 +285,7 @@ class TrafficSignDetector:
             print(f"       2) Use augmentation during training")
             print(f"       3) Check training/validation accuracy was good")
             print(f"       4) Ensure training data matches inference image types")
+            print(f"     - Try lowering the confidence threshold slider to see detections")
         
         if scale < 0.5:
             print(f"\n  ⚠️  SCALING ISSUE:")