Martin Tomov commited on
Commit
8fc1551
Β·
verified Β·
1 Parent(s): 1b22764

gr.Checkbox(label="Include Bounding Boxes", value=True)],

Browse files
Files changed (1) hide show
  1. app.py +7 -8
app.py CHANGED
@@ -55,13 +55,9 @@ def annotate(image: Union[Image.Image, np.ndarray], detection_results: List[Dete
55
  score = detection.score
56
  box = detection.box
57
  mask = detection.mask
58
- color = (255, 255, 255) # White color for bounding boxes and text
59
-
60
- # Draw the filled rectangle (white filling)
61
- cv2.rectangle(image_cv2, (box.xmin, box.ymin), (box.xmax, box.ymax), color, -1) # Filled white rectangle
62
- # Draw the border
63
- cv2.rectangle(image_cv2, (box.xmin, box.ymin), (box.xmax, box.ymax), color, 2) # Border
64
 
 
65
  cv2.putText(image_cv2, f'{label}: {score:.2f}', (box.xmin, box.ymin - 10),
66
  cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
67
 
@@ -197,9 +193,12 @@ def detections_to_json(detections):
197
  detections_list.append(detection_dict)
198
  return detections_list
199
 
200
- def process_image(image, include_json):
201
  labels = ["insect"]
202
  original_image, detections = grounded_segmentation(image, labels, threshold=0.3, polygon_refinement=True)
 
 
 
203
  yellow_background_with_insects = create_yellow_background_with_insects(np.array(original_image), detections)
204
  if include_json:
205
  detections_json = detections_to_json(detections)
@@ -216,7 +215,7 @@ examples = [
216
 
217
  gr.Interface(
218
  fn=process_image,
219
- inputs=[gr.Image(type="pil"), gr.Checkbox(label="Include JSON", value=False)],
220
  outputs=[gr.Image(type="numpy"), gr.Textbox()],
221
  title="InsectSAM 🐞",
222
  examples=examples
 
55
  score = detection.score
56
  box = detection.box
57
  mask = detection.mask
58
+ color = np.random.randint(0, 256, size=3).tolist()
 
 
 
 
 
59
 
60
+ cv2.rectangle(image_cv2, (box.xmin, box.ymin), (box.xmax, box.ymax), color, 2)
61
  cv2.putText(image_cv2, f'{label}: {score:.2f}', (box.xmin, box.ymin - 10),
62
  cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
63
 
 
193
  detections_list.append(detection_dict)
194
  return detections_list
195
 
196
+ def process_image(image, include_json, include_bboxes):
197
  labels = ["insect"]
198
  original_image, detections = grounded_segmentation(image, labels, threshold=0.3, polygon_refinement=True)
199
+ if not include_bboxes:
200
+ for detection in detections:
201
+ detection.box = None
202
  yellow_background_with_insects = create_yellow_background_with_insects(np.array(original_image), detections)
203
  if include_json:
204
  detections_json = detections_to_json(detections)
 
215
 
216
  gr.Interface(
217
  fn=process_image,
218
+ inputs=[gr.Image(type="pil"), gr.Checkbox(label="Include JSON", value=False), gr.Checkbox(label="Include Bounding Boxes", value=True)],
219
  outputs=[gr.Image(type="numpy"), gr.Textbox()],
220
  title="InsectSAM 🐞",
221
  examples=examples