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()