Upload 2 files

app.py

@@ -33,16 +33,16 @@ def initialize_models():
         # Set device
         device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         
-        # Initialize ShuffleNet model
-        model = models.shufflenet_v2_x0_5(weights=None)
-        model.fc = nn.Sequential(
-            nn.Linear(model.fc.in_features, 2),
+        # Initialize MobileNetV3 model
+        model = models.mobilenet_v3_small(weights=None)
+        model.classifier = nn.Sequential(
+            nn.Linear(576, 2),  # Direct mapping to output classes
             nn.Softmax(dim=1)
         )
         model = model.to(device)
         
         # Load model weights
-        best_model_path = "best_model_ShuffleNetV2.pth"
+        best_model_path = "best_model_mobilenet_v3_v2.pth"
         if not os.path.exists(best_model_path):
             st.error(f"Model file not found: {best_model_path}")
             return None, None, None
@@ -54,7 +54,7 @@ def initialize_models():
         model.eval()
         
         # Load YOLO model
-        yolo_model_path = "yolo11s.onnx"
+        yolo_model_path = "../yolo11s.onnx"  # Going up one directory since the app.py is in API22_FEB
         if not os.path.exists(yolo_model_path):
             st.error(f"YOLO model file not found: {yolo_model_path}")
             return device, model, None
@@ -80,7 +80,8 @@ def process_image(image, model, device):
     # Perform inference
     with torch.no_grad():
         output = model(input_tensor)
-        probabilities = output[0]
+        probabilities = output[0]  # Get probabilities for both classes
+        # Class 0 is "No Red Light", Class 1 is "Red Light"
         no_red_light_prob = probabilities[0].item()
         red_light_prob = probabilities[1].item()
         is_red_light = red_light_prob > no_red_light_prob
@@ -128,6 +129,44 @@ def put_text_with_background(img, text, position, font_scale=0.8, thickness=2, f
     # Put text
     cv2.putText(img, text, (position[0] + padding, position[1]), font, font_scale, (255, 255, 255), thickness)
 
+def calculate_iou(box1, box2):
+    """Calculate Intersection over Union between two bounding boxes."""
+    x1 = max(box1[0], box2[0])
+    y1 = max(box1[1], box2[1])
+    x2 = min(box1[2], box2[2])
+    y2 = min(box1[3], box2[3])
+    
+    intersection = max(0, x2 - x1) * max(0, y2 - y1)
+    box1_area = (box1[2] - box1[0]) * (box1[3] - box1[1])
+    box2_area = (box2[2] - box2[0]) * (box2[3] - box2[1])
+    union = box1_area + box2_area - intersection
+    
+    return intersection / union if union > 0 else 0
+
+def merge_overlapping_detections(detections, iou_threshold=0.5):
+    """Merge overlapping detections of the same class."""
+    if not detections:
+        return []
+        
+    # Sort detections by confidence
+    detections = sorted(detections, key=lambda x: x['confidence'], reverse=True)
+    merged_detections = []
+    
+    while detections:
+        best_detection = detections.pop(0)
+        i = 0
+        while i < len(detections):
+            current_detection = detections[i]
+            if (current_detection['class'] == best_detection['class'] and 
+                calculate_iou(current_detection['bbox'], best_detection['bbox']) >= iou_threshold):
+                # Remove the lower confidence detection
+                detections.pop(i)
+            else:
+                i += 1
+        merged_detections.append(best_detection)
+    
+    return merged_detections
+
 def main():
     st.title("Traffic Light Detection with Protection Area")
     
@@ -279,17 +318,23 @@ def main():
                                             'confidence': confidence,
                                             'bbox': bbox
                                         })
-                                        
-                                        # Draw detection
-                                        cv2.rectangle(cv_image, 
-                                                    (int(bbox[0]), int(bbox[1])), 
-                                                    (int(bbox[2]), int(bbox[3])), 
-                                                    (0, 0, 255), 2)
-                                        
-                                        # Add label
-                                        text = f"{class_name}: {confidence:.2%}"
-                                        put_text_with_background(cv_image, text, 
-                                                               (int(bbox[0]), int(bbox[1]) - 10))
+                    
+                    # Merge overlapping detections
+                    detection_results = merge_overlapping_detections(detection_results, iou_threshold=0.5)
+                    
+                    # Draw detections
+                    for det in detection_results:
+                        bbox = det['bbox']
+                        # Draw detection box
+                        cv2.rectangle(cv_image, 
+                                    (int(bbox[0]), int(bbox[1])), 
+                                    (int(bbox[2]), int(bbox[3])), 
+                                    (0, 0, 255), 2)
+                        
+                        # Add label
+                        text = f"{det['class']}: {det['confidence']:.2%}"
+                        put_text_with_background(cv_image, text, 
+                                               (int(bbox[0]), int(bbox[1]) - 10))
                     
                     # Add status text
                     status_text = f"Red Light: DETECTED ({red_light_prob:.1%})"

best_model_mobilenet_v3_v2.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a9d6dbfc5f368b8dd4f06f86e2ef088c0cd88c7bfd4f686800d2ef7b256b36f7
+size 3850192