add merge mutil mask

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) -> Optional[Image.Image]:
+def process_image(image_input, image_url, task_prompt, text_prompt=None, dilate=0, merge_masks=False) -> Optional[Image.Image]:
     if not image_input:
         gr.Info("Please upload an image.")
         return None
@@ -72,6 +72,14 @@ def process_image(image_input, image_url, task_prompt, text_prompt=None, dilate=
         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
 
 
@@ -84,6 +92,8 @@ with gr.Blocks() as demo:
                 ['<OD>', '<CAPTION_TO_PHRASE_GROUNDING>', '<DENSE_REGION_CAPTION>', '<REGION_PROPOSAL>', '<OCR_WITH_REGION>', '<REFERRING_EXPRESSION_SEGMENTATION>', '<REGION_TO_SEGMENTATION>', '<OPEN_VOCABULARY_DETECTION>', '<REGION_TO_CATEGORY>', '<REGION_TO_DESCRIPTION>'], value="<CAPTION_TO_PHRASE_GROUNDING>", label="Task Prompt", info="task prompts"
             )
             dilate = gr.Slider(label="dilate mask", minimum=0, maximum=50, value=10, step=1)
+            merge_masks = gr.Checkbox(label="Merge 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():
@@ -91,7 +101,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],
+        inputs = [image, image_url, task_prompt, text_prompt, dilate, merge_masks],
         outputs = [image_gallery,],
         show_api=False
     )