Upload app.py

app.py

@@ -40,7 +40,7 @@ def read_yolo_boxes(file_path):
         for line in f:
             parts = line.strip().split()
             class_id = int(parts[0])
-            if COCO_CLASSES[class_id] != 'traffic light':  # Skip traffic lights
+            if COCO_CLASSES[class_id] != 'traffic light':
                 class_name = COCO_CLASSES[class_id]
                 x, y, w, h = map(float, parts[1:5])
                 boxes.append((class_name, x, y, w, h))
@@ -88,18 +88,20 @@ def plot_boxes_and_segment(image, yolo_boxes, segment, img_width, img_height, th
     labels = {'intersecting': 'Intersecting Box', 'obstructed': 'Obstructed Box', 'not touching': 'Non-interacting Box'}
     
     for yolo_box in yolo_boxes:
-        class_id, x_center, y_center, width, height = yolo_box
-        x1, y1, x2, y2 = yolo_to_pixel_coords(x_center, y_center, width, height, img_width, img_height)
-        relationship = box_segment_relationship(yolo_box, segment, img_width, img_height, threshold)
-        color = colors[relationship]
-        label = labels[relationship]
-        ax.add_patch(plt.Rectangle((x1, y1), x2-x1, y2-y1, fill=False, edgecolor=color, linewidth=2, label=label))
+        class_name, x_center, y_center, width, height = yolo_box
+        if class_name != 'traffic light':
+            x1, y1, x2, y2 = yolo_to_pixel_coords(x_center, y_center, width, height, img_width, img_height)
+            relationship = box_segment_relationship(yolo_box, segment, img_width, img_height, threshold)
+            color = colors[relationship]
+            label = labels[relationship]
+            ax.add_patch(plt.Rectangle((x1, y1), x2-x1, y2-y1, fill=False, edgecolor=color, linewidth=2, label=label))
     
     ax.legend()
     ax.axis('off')
     plt.tight_layout()
     return fig
 
+
 # COCO classes
 COCO_CLASSES = [
     'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light',
@@ -177,16 +179,18 @@ def detect_objects(image, rail_segment):
     
     fig = plot_boxes_and_segment(image, yolo_boxes, rail_segment, img_width, img_height, threshold)
     
-    results = []
+   results = []
     for class_name, x, y, w, h in yolo_boxes:
-        result = box_segment_relationship((0, x, y, w, h), rail_segment, img_width, img_height, threshold)
-        results.append(f"{class_name} at ({x:.2f}, {y:.2f}) is {result} the segment.")
+        if class_name != 'traffic light':
+            result = box_segment_relationship((0, x, y, w, h), rail_segment, img_width, img_height, threshold)
+            results.append(f"{class_name} at ({x:.2f}, {y:.2f}) is {result} the segment.")
     
     os.remove(temp_image_path)
     os.remove(label_path)
     
     return fig, "\n".join(results), yolo_boxes
 
+
 def process_video(video_path, rail_segment, frame_skip=15):
     cap = cv2.VideoCapture(video_path)
     if not cap.isOpened():
@@ -205,7 +209,7 @@ def process_video(video_path, rail_segment, frame_skip=15):
     total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
     threshold = 10  # Set threshold (in pixels) for obstruction detection
     
-    while True:
+     while True:
         ret, frame = cap.read()
         if not ret:
             break
@@ -216,33 +220,30 @@ def process_video(video_path, rail_segment, frame_skip=15):
         
         processed_count += 1
         
-        # Convert frame to PIL Image for compatibility with detect_objects
         pil_frame = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
         
-        # Detect objects in the frame
         _, _, yolo_boxes = detect_objects(pil_frame, rail_segment)
         
-        # Draw rail segment
         pixel_segment = convert_segment_to_pixel(rail_segment, width, height)
         pts = np.array(list(zip(pixel_segment[::2], pixel_segment[1::2])), np.int32)
         pts = pts.reshape((-1, 1, 2))
         cv2.polylines(frame, [pts], True, (0, 0, 255), 2)
         
-        # Check for obstructions and draw bounding boxes
         for box in yolo_boxes:
             class_name, x, y, w, h = box
-            relationship = box_segment_relationship((0, x, y, w, h), rail_segment, width, height, threshold)
-            x1, y1, x2, y2 = yolo_to_pixel_coords(x, y, w, h, width, height)
-            
-            if relationship == "intersecting":
-                color = (0, 0, 255)  # Red for intersecting
-            elif relationship == "obstructed":
-                color = (0, 255, 255)  # Yellow for obstructed
-            else:
-                color = (0, 255, 0)  # Green for not touching
-            
-            cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)
-            cv2.putText(frame, f"{class_name} ({relationship})", (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, color, 2)
+            if class_name != 'traffic light':
+                relationship = box_segment_relationship((0, x, y, w, h), rail_segment, width, height, threshold)
+                x1, y1, x2, y2 = yolo_to_pixel_coords(x, y, w, h, width, height)
+                
+                if relationship == "intersecting":
+                    color = (0, 0, 255)  # Red for intersecting
+                elif relationship == "obstructed":
+                    color = (0, 255, 255)  # Yellow for obstructed
+                else:
+                    color = (0, 255, 0)  # Green for not touching
+                
+                cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)
+                cv2.putText(frame, f"{class_name} ({relationship})", (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, color, 2)
         
         out.write(frame)
         
@@ -312,11 +313,11 @@ class TwoStepDetection:
             
             cap.release()
             
-            #obstruction_percentage = (obstructed_frames / total_frames) * 100
+            obstruction_percentage = (obstructed_frames / total_frames) * 100
             summary = f"{processing_message}\n\nObstruction Summary:\n"
             summary += f"Total frames: {total_frames}\n"
             summary += f"Frames with obstructions: {obstructed_frames}\n"
-            #summary += f"Percentage of frames with obstructions: {obstruction_percentage:.2f}%"
+            summary += f"Percentage of frames with obstructions: {obstruction_percentage:.2f}%"
             
             return None, video_output, summary