cropped bbox output

app.py

@@ -109,7 +109,7 @@ def detect(image: Image.Image, labels: List[str], threshold: float = 0.3, detect
     results = object_detector(image, candidate_labels=labels, threshold=threshold)
     return [DetectionResult.from_dict(result) for result in results]
 
-@spaces.GPU
+@spaces.GGPU
 def segment(image: Image.Image, detection_results: List[DetectionResult], polygon_refinement: bool = False, segmenter_id: Optional[str] = None) -> List[DetectionResult]:
     segmenter_id = segmenter_id if segmenter_id else "martintmv/InsectSAM"
     segmentator = AutoModelForMaskGeneration.from_pretrained(segmenter_id).to("cuda")
@@ -190,19 +190,33 @@ def detections_to_json(detections):
         detections_list.append(detection_dict)
     return detections_list
 
+def crop_bounding_boxes(image: np.ndarray, detections: List[DetectionResult]) -> List[np.ndarray]:
+    crops = []
+    for detection in detections:
+        xmin, ymin, xmax, ymax = detection.box.xyxy
+        crop = image[ymin:ymax, xmin:xmax]
+        crops.append(crop)
+    return crops
+
 def process_image(image, include_json, include_bboxes):
     labels = ["insect"]
     original_image, detections = grounded_segmentation(image, labels, threshold=0.3, polygon_refinement=True)
     yellow_background_with_insects = create_yellow_background_with_insects(np.array(original_image), detections)
     annotated_image = plot_detections(yellow_background_with_insects, detections, include_bboxes)
+
+    results = [annotated_image]
+    if include_bboxes:
+        crops = crop_bounding_boxes(np.array(original_image), detections)
+        results.extend(crops)
+
     if include_json:
         detections_json = detections_to_json(detections)
         json_output_path = "insect_detections.json"
         with open(json_output_path, 'w') as json_file:
             json.dump(detections_json, json_file, indent=4)
-        return annotated_image, json.dumps(detections_json, separators=(',', ':'))
-    else:
-        return annotated_image, None
+        results.append(json.dumps(detections_json, separators=(',', ':')))
+    
+    return tuple(results)
 
 examples = [
     ["flower-night.jpg"]
@@ -211,7 +225,7 @@ examples = [
 gr.Interface(
     fn=process_image,
     inputs=[gr.Image(type="pil"), gr.Checkbox(label="Include JSON", value=False), gr.Checkbox(label="Include Bounding Boxes", value=False)],
-    outputs=[gr.Image(type="numpy"), gr.Textbox()],
+    outputs=[gr.Image(type="numpy")] + [gr.Image(type="numpy")] * 10 + [gr.Textbox()],
     title="InsectSAM 🐞",
     examples=examples
 ).launch()