update app.py

app.py

@@ -672,7 +672,7 @@ def load_kolors_inpainting(inpaint_prompt, input_image, mask_image):
         return None
 
 def run_anything_task(input_image, text_prompt, task_type, inpaint_prompt, box_threshold, text_threshold, 
-            iou_threshold, inpaint_mode, mask_source_radio, remove_mode, remove_mask_extend, num_relation, kosmos_input, cleaner_size_limit=1080):
+            iou_threshold, inpaint_mode, mask_source_radio, remove_mode, remove_mask_extend, remove_use_segment, num_relation, kosmos_input, cleaner_size_limit=1080):
     text_prompt = getTextTrans(text_prompt, source='zh', target='en')
     inpaint_prompt = getTextTrans(inpaint_prompt, source='zh', target='en')
 
@@ -716,7 +716,7 @@ def run_anything_task(input_image, text_prompt, task_type, inpaint_prompt, box_t
             return [], gr.update(label='Please upload a image!😂😂😂😂'), time_cost_str, gr.update(visible=(time_cost_str !='')), None, None, None
 
     file_temp = int(time.time())
-    logger.info(f'run_anything_task_002/{device}_[{file_temp}]_{task_type}/{inpaint_mode}/[{mask_source_radio}]/{remove_mode}/{remove_mask_extend}_[{text_prompt}]/[{inpaint_prompt}]___1_')
+    logger.info(f'run_anything_task_002/{device}_[{file_temp}]_{task_type}/{inpaint_mode}/[{mask_source_radio}]/{remove_mode}/{remove_mask_extend}/{remove_use_segment}_[{text_prompt}]/[{inpaint_prompt}]___1_')
 
     output_images = []
 
@@ -739,6 +739,9 @@ def run_anything_task(input_image, text_prompt, task_type, inpaint_prompt, box_t
     size = image_pil.size
     H, W = size[1], size[0]
 
+    if remove_use_segment and task_type == 'remove':
+        remove_mode = 'rectangle'
+
     # run grounding dino model
     if (task_type in ['inpainting', 'outpainting'] or task_type == 'remove') and mask_source_radio == mask_source_draw:
         pass
@@ -772,7 +775,11 @@ def run_anything_task(input_image, text_prompt, task_type, inpaint_prompt, box_t
     logger.info(f'run_anything_task_[{file_temp}]_{task_type}_2_')
     if task_type == 'segment' or ((task_type in ['inpainting', 'outpainting'] or task_type == 'remove') and mask_source_radio == mask_source_segment):
         image = np.array(input_img)
-        if sam_predictor:
+        use_sam_predictor = True
+        if remove_use_segment and task_type == 'remove':
+            use_sam_predictor = False
+
+        if sam_predictor and use_sam_predictor:
             sam_predictor.set_image(image)
 
         for i in range(boxes_filt.size(0)):
@@ -780,7 +787,7 @@ def run_anything_task(input_image, text_prompt, task_type, inpaint_prompt, box_t
             boxes_filt[i][:2] -= boxes_filt[i][2:] / 2
             boxes_filt[i][2:] += boxes_filt[i][:2]
 
-        if sam_predictor:
+        if sam_predictor and use_sam_predictor:
             boxes_filt = boxes_filt.to(sam_device)
             transformed_boxes = sam_predictor.transform.apply_boxes_torch(boxes_filt, image.shape[:2])
 
@@ -1019,7 +1026,8 @@ def main_gradio(args):
                             remove_mode = gr.Radio(["segment", "rectangle"],  value="segment", label='remove mode') 
                         with gr.Column(scale=1):
                             remove_mask_extend = gr.Textbox(label="remove_mask_extend", value='10')
-
+                        with gr.Column(scale=1):
+                            remove_use_segment = gr.Checkbox(value=True, elem_id='remove_use_segment', label="use segment for removing?", info="")
             with gr.Column():
                 image_gallery = gr.Gallery(label="result images", show_label=True, elem_id="gallery", height=512, visible=True
                     ) #.style(preview=True, columns=[5], object_fit="scale-down", height="auto")   
@@ -1055,7 +1063,7 @@ def main_gradio(args):
 
             run_button.click(fn=run_anything_task, inputs=[
                             input_image, text_prompt, task_type, inpaint_prompt, box_threshold, text_threshold, 
-                            iou_threshold, inpaint_mode, mask_source_radio, remove_mode, remove_mask_extend, num_relation, kosmos_input], 
+                            iou_threshold, inpaint_mode, mask_source_radio, remove_mode, remove_mask_extend, remove_use_segment, num_relation, kosmos_input], 
                             outputs=[image_gallery, image_gallery, time_cost, time_cost, kosmos_output, kosmos_text_output, entity_output], show_progress=True, queue=True)
             
             mask_source_radio.change(fn=change_radio_display, inputs=[task_type, mask_source_radio],