Add filtering step for display of bounding box

app.py

@@ -35,6 +35,18 @@ def parse_rois(roi_json):
     return roi_polygons
 
 rois = []
+names= [ 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light',
+         'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow',
+         'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee',
+         'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard',
+         'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',
+         'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch',
+         'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone',
+         'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear',
+         'hair drier', 'toothbrush' ]
+desired_classes = ['person', 'bicycle', 'car', 'motorcycle', 'bus', 'train', 'truck']
+desired_indices = [names.index(cls) for cls in desired_classes if cls in names]
+print(desired_indices)
 
 def detect(img, model, rois):
     if img is None:
@@ -81,27 +93,33 @@ def detect(img, model, rois):
             if len(det):
                 det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()
 
-                # Filter detections based on ROIs
-                filtered_detections = []
-                for *xyxy, conf, cls in reversed(det):
-                    x_center = float((xyxy[0] + xyxy[2]) / 2)
-                    y_center = float((xyxy[1] + xyxy[3]) / 2)
+            # Filter detections based on ROIs and desired classes
+            filtered_detections = []
+            for *xyxy, conf, cls in reversed(det):
+                if int(cls) not in desired_indices:
+                    continue
+                x_center = float((xyxy[0] + xyxy[2]) / 2)
+                y_center = float((xyxy[1] + xyxy[3]) / 2)
 
-                    # Check if the center of the detection is inside any of the ROI polygons
-                    for roi_points in roi_points_list:
-                        if cv2.pointPolygonTest(np.array(roi_points, dtype=np.float32), (x_center, y_center), False) >= 0:
-                            filtered_detections.append((*xyxy, conf, cls))
-                            break
+                inside_roi = False
+                for roi_points in roi_points_list:
+                    if cv2.pointPolygonTest(np.array(roi_points, dtype=np.float32), (x_center, y_center), False) >= 0:
+                        inside_roi = True
+                        break
 
-                for *xyxy, conf, cls in filtered_detections:
-                    label = f'{names[int(cls)]} {conf:.2f}'
-                    plot_one_box(xyxy, im0, label=label, color=colors[int(cls)], line_thickness=3)
+                filtered_detections.append((*xyxy, conf, cls, inside_roi))
+
+            # Plot the detections with the desired color
+            for *xyxy, conf, cls, inside_roi in filtered_detections:
+                color = (0, 255, 0) if inside_roi else (0, 0, 255)  # green for inside ROI, red for outside
+                label = f'{names[int(cls)]} {conf:.2f}'
+                plot_one_box(xyxy, im0, label=label, color=color, line_thickness=1)
 
         return Image.fromarray(im0[:,:,::-1])
 # Modify the Gradio interface to accept ROIs
 roi_example = [
     ["8-2.jpg", "yolov7", json.dumps([
-        {"coordinates": [{"x": 0.035, "y": 0.631}, {"x": 0.083, "y": 0.584}, {"x": 0.935, "y": 0.832}, {"x": 0.874, "y": 0.907}]},
+        {"coordinates": [{ "x": 0.005, "y": 0.644 },{ "x": 0.047, "y": 0.572 },{ "x": 0.961, "y": 0.834 },{ "x": 0.695, "y": 0.919 }]},
         {"coordinates": [{"x": 0.237, "y": 0.505}, {"x": 0.283, "y": 0.460}, {"x": 0.921, "y": 0.578}, {"x": 0.912, "y": 0.654}]}
     ])],
     ["9-1.jpg", "yolov7", json.dumps([
@@ -162,7 +180,7 @@ roi_example = [
         {"coordinates": [{"x": 0.629, "y": 0.255}, {"x": 0.661, "y": 0.281}, {"x": 0.006, "y": 0.602}, {"x": 0.005, "y": 0.408}]}
     ])]
 ]
-
+print(len(roi_example))
 gr_examples = [[example[0], example[1], example[2]] for example in roi_example]
 description_html = """<b>Demo for YOLOv7 Object Detection</b>: This interface is specifically tailored for <b>detecting vehicles</b> in images. The primary focus is on <b>accident-prone regions</b> on public roads. By leveraging state-of-the-art object detection techniques, this tool aims to provide insights into areas where vehicles are most at risk, helping in <b>road safety analysis</b> and <b>preventive measures</b>. Users can also define <b>Regions of Interest (ROIs)</b> to narrow down the detection area, ensuring that the analysis is focused on the most critical parts of the image."""
 
@@ -173,5 +191,6 @@ gr.Interface(
     title="YOLOv7 Object Detection for Accident-Prone Regions", 
     examples=gr_examples, 
     description=description_html,
-    live=False  # This ensures that the model doesn't run until the 'Submit' button is clicked
+    live=False,# This ensures that the model doesn't run until the 'Submit' button is clicked
+    example_ceiling=14 
 ).launch()