Upload app.py

app.py

@@ -39,11 +39,9 @@ def read_yolo_boxes(file_path):
     with open(file_path, 'r') as f:
         for line in f:
             parts = line.strip().split()
-            class_id = int(parts[0])
-            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))
+            class_name = COCO_CLASSES[int(parts[0])]
+            x, y, w, h = map(float, parts[1:5])
+            boxes.append((class_name, x, y, w, h))
     return boxes
 
 def yolo_to_pixel_coord(x, y, img_width, img_height):
@@ -88,20 +86,18 @@ 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_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))
+        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))
     
     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',
@@ -181,16 +177,14 @@ def detect_objects(image, rail_segment):
     
     results = []
     for class_name, x, y, w, h in yolo_boxes:
-        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.")
+        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():
@@ -220,30 +214,33 @@ 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
-            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)
+            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)