return_rectangles

app.py

@@ -30,7 +30,7 @@ SAM_IMAGE_MODEL = load_sam_image_model(device=DEVICE)
 @spaces.GPU(duration=20)
 @torch.inference_mode()
 @torch.autocast(device_type="cuda", dtype=torch.bfloat16)
-def process_image(image_input, image_url, task_prompt, text_prompt=None, dilate=0, merge_masks=False) -> Optional[Image.Image]:
+def process_image(image_input, image_url, task_prompt, text_prompt=None, dilate=0, merge_masks=False, return_rectangles=False) -> Optional[Image.Image]:
     if not image_input:
         gr.Info("Please upload an image.")
         return None
@@ -58,28 +58,37 @@ def process_image(image_input, image_url, task_prompt, text_prompt=None, dilate=
         result=result,
         resolution_wh=image_input.size
     )
-    detections = run_sam_inference(SAM_IMAGE_MODEL, image_input, detections)
-    if len(detections) == 0:
-        gr.Info("No objects detected.")
-        return None
     images = []
-    print("mask generated:", len(detections.mask))
-    kernel_size = dilate
-    kernel = np.ones((kernel_size, kernel_size), np.uint8)
+    if return_rectangles:
+       # 创建黑色背景的图片
+        mask_image = np.zeros((image_input.size.height, image_input.size.width), dtype=np.uint8)
+        bboxes = detections.get('bboxes', [])
+        for bbox in bboxes:
+            x1, y1, x2, y2 = map(int, bbox)
+            # 在 mask_image 上绘制白色的矩形
+            cv2.rectangle(mask_image, (x1, y1), (x2, y2), 255, thickness=cv2.FILLED)
+        images = [mask_image]
+    else:
+        # sam
+        detections = run_sam_inference(SAM_IMAGE_MODEL, image_input, detections)
+        if len(detections) == 0:
+            gr.Info("No objects detected.")
+            return None
+        kernel_size = dilate
+        print("mask generated:", len(detections.mask))
+        kernel = np.ones((kernel_size, kernel_size), np.uint8)
+        for i in range(len(detections.mask)):
+            mask = detections.mask[i].astype(np.uint8) * 255
+            if dilate > 0:
+                mask = cv2.dilate(mask, kernel, iterations=1)
+            images.append(mask)
+        if merge_masks:
+        
+            merged_mask = np.zeros_like(images[0], dtype=np.uint8)
+            for mask in images:
+                merged_mask = cv2.bitwise_or(merged_mask, mask)
+            images = [merged_mask] + images
 
-    for i in range(len(detections.mask)):
-        mask = detections.mask[i].astype(np.uint8) * 255
-        if dilate > 0:
-            mask = cv2.dilate(mask, kernel, iterations=1)
-        images.append(mask)
-    
-    if merge_masks:
-        final_images = []
-        merged_mask = np.zeros_like(images[0], dtype=np.uint8)
-        for mask in images:
-            merged_mask = cv2.bitwise_or(merged_mask, mask)
-        final_images = [merged_mask]
-        return final_images
     return images
 
 
@@ -93,7 +102,7 @@ with gr.Blocks() as demo:
             )
             dilate = gr.Slider(label="dilate mask", minimum=0, maximum=50, value=10, step=1)
             merge_masks = gr.Checkbox(label="Merge masks", value=False)
-
+            return_rectangles = gr.Checkbox(label="Return rectangle masks", value=False)
             text_prompt = gr.Textbox(label='Text prompt', placeholder='Enter text prompts')
             submit_button = gr.Button(value='Submit', variant='primary')
         with gr.Column():
@@ -101,7 +110,7 @@ with gr.Blocks() as demo:
     print(image, image_url, task_prompt, text_prompt, image_gallery)
     submit_button.click(
         fn = process_image,
-        inputs = [image, image_url, task_prompt, text_prompt, dilate, merge_masks],
+        inputs = [image, image_url, task_prompt, text_prompt, dilate, merge_masks, return_rectangles],
         outputs = [image_gallery,],
         show_api=False
     )