Add a wider result

app_gradio.py

@@ -21,6 +21,8 @@ device = torch.device(
 title = "<center><strong><font size='8'>🏃 Fast Segment Anything 🤗</font></strong></center>"
 
 news = """ # 📖 News
+        🔥 2023/07/14: Add a "wider result" button in text mode (Thanks for [gaoxinge](https://github.com/CASIA-IVA-Lab/FastSAM/pull/95)).
+
         🔥 2023/06/29: Support the text mode (Thanks for [gaoxinge](https://github.com/CASIA-IVA-Lab/FastSAM/pull/47)).
 
         🔥 2023/06/26: Support the points mode. (Better and faster interaction will come soon!)
@@ -76,6 +78,7 @@ def segment_everything(
     withContours=True,
     use_retina=True,
     text="",
+    wider=False,
     mask_random_color=True,
 ):
     input_size = int(input_size)  # 确保 imgsz 是整数
@@ -95,7 +98,7 @@ def segment_everything(
 
     if len(text) > 0:
         results = format_results(results[0], 0)
-        annotations, _ = text_prompt(results, text, input, device=device)
+        annotations, _ = text_prompt(results, text, input, device=device, wider=wider)
         annotations = np.array([annotations])
     else:
         annotations = results[0].masks.data
@@ -189,7 +192,7 @@ segm_img_t = gr.Image(label="Segmented Image with text", interactive=False, type
 global_points = []
 global_point_label = []
 
-input_size_slider_e = gr.components.Slider(minimum=512,
+input_size_slider = gr.components.Slider(minimum=512,
                                          maximum=1024,
                                          value=1024,
                                          step=64,
@@ -218,10 +221,10 @@ with gr.Blocks(css=css, title='Fast Segment Anything') as demo:
         # Submit & Clear
         with gr.Row():
             with gr.Column():
-                input_size_slider_e.render()
+                input_size_slider.render()
 
                 with gr.Row():
-                    contour_check_e = gr.Checkbox(value=True, label='withContours', info='draw the edges of the masks')
+                    contour_check = gr.Checkbox(value=True, label='withContours', info='draw the edges of the masks')
 
                     with gr.Column():
                         segment_btn_e = gr.Button("Segment Everything", variant='primary')
@@ -237,16 +240,28 @@ with gr.Blocks(css=css, title='Fast Segment Anything') as demo:
 
             with gr.Column():
                 with gr.Accordion("Advanced options", open=False):
-                    # text_box = gr.Textbox(label="text prompt")
-                    iou_threshold_e = gr.Slider(0.1, 0.9, 0.7, step=0.1, label='iou', info='iou threshold for filtering the annotations')
-                    conf_threshold_e = gr.Slider(0.1, 0.9, 0.25, step=0.05, label='conf', info='object confidence threshold')
+                    iou_threshold = gr.Slider(0.1, 0.9, 0.7, step=0.1, label='iou', info='iou threshold for filtering the annotations')
+                    conf_threshold = gr.Slider(0.1, 0.9, 0.25, step=0.05, label='conf', info='object confidence threshold')
                     with gr.Row():
-                        mor_check_e = gr.Checkbox(value=False, label='better_visual_quality', info='better quality using morphologyEx')
+                        mor_check = gr.Checkbox(value=False, label='better_visual_quality', info='better quality using morphologyEx')
                         with gr.Column():
-                            retina_check_e = gr.Checkbox(value=True, label='use_retina', info='draw high-resolution segmentation masks')
+                            retina_check = gr.Checkbox(value=True, label='use_retina', info='draw high-resolution segmentation masks')
+
                 # Description
                 gr.Markdown(description_e)
 
+    segment_btn_e.click(segment_everything,
+                        inputs=[
+                            cond_img_e,
+                            input_size_slider,
+                            iou_threshold,
+                            conf_threshold,
+                            mor_check,
+                            contour_check,
+                            retina_check,
+                        ],
+                        outputs=segm_img_e)
+
     with gr.Tab("Points mode"):
         # Images
         with gr.Row(variant="panel"):
@@ -277,7 +292,13 @@ with gr.Blocks(css=css, title='Fast Segment Anything') as demo:
             with gr.Column():
                 # Description
                 gr.Markdown(description_p)
-                
+
+    cond_img_p.select(get_points_with_draw, [cond_img_p, add_or_remove], cond_img_p)
+
+    segment_btn_p.click(segment_with_points,
+                        inputs=[cond_img_p],
+                        outputs=[segm_img_p, cond_img_p])
+
     with gr.Tab("Text mode"):
         # Images
         with gr.Row(variant="panel"):
@@ -291,14 +312,14 @@ with gr.Blocks(css=css, title='Fast Segment Anything') as demo:
         with gr.Row():
             with gr.Column():
                 input_size_slider_t = gr.components.Slider(minimum=512,
-                                         maximum=1024,
-                                         value=1024,
-                                         step=64,
-                                         label='Input_size',
-                                         info='Our model was trained on a size of 1024')
+                                                           maximum=1024,
+                                                           value=1024,
+                                                           step=64,
+                                                           label='Input_size',
+                                                           info='Our model was trained on a size of 1024')
                 with gr.Row():
                     with gr.Column():
-                        contour_check_t = gr.Checkbox(value=True, label='withContours', info='draw the edges of the masks')
+                        contour_check = gr.Checkbox(value=True, label='withContours', info='draw the edges of the masks')
                         text_box = gr.Textbox(label="text prompt", value="a black dog")
 
                     with gr.Column():
@@ -306,7 +327,7 @@ with gr.Blocks(css=css, title='Fast Segment Anything') as demo:
                         clear_btn_t = gr.Button("Clear", variant="secondary")
 
                 gr.Markdown("Try some of the examples below ⬇️")
-                gr.Examples(examples=["examples/dogs.jpg"],
+                gr.Examples(examples=[["examples/dogs.jpg"]] + examples,
                             inputs=[cond_img_e],
                             # outputs=segm_img_e,
                             # fn=segment_everything,
@@ -315,44 +336,27 @@ with gr.Blocks(css=css, title='Fast Segment Anything') as demo:
 
             with gr.Column():
                 with gr.Accordion("Advanced options", open=False):
-                    iou_threshold_t = gr.Slider(0.1, 0.9, 0.7, step=0.1, label='iou', info='iou threshold for filtering the annotations')
-                    conf_threshold_t = gr.Slider(0.1, 0.9, 0.25, step=0.05, label='conf', info='object confidence threshold')
+                    iou_threshold = gr.Slider(0.1, 0.9, 0.7, step=0.1, label='iou', info='iou threshold for filtering the annotations')
+                    conf_threshold = gr.Slider(0.1, 0.9, 0.25, step=0.05, label='conf', info='object confidence threshold')
                     with gr.Row():
-                        mor_check_t = gr.Checkbox(value=False, label='better_visual_quality', info='better quality using morphologyEx')
-                        with gr.Column():
-                            retina_check_t = gr.Checkbox(value=True, label='use_retina', info='draw high-resolution segmentation masks')
+                        mor_check = gr.Checkbox(value=False, label='better_visual_quality', info='better quality using morphologyEx')
+                        retina_check = gr.Checkbox(value=True, label='use_retina', info='draw high-resolution segmentation masks')
+                        wider_check = gr.Checkbox(value=False, label='wider', info='wider result')
 
                 # Description
                 gr.Markdown(description_e)
-        
-    cond_img_p.select(get_points_with_draw, [cond_img_p, add_or_remove], cond_img_p)
-
-    segment_btn_e.click(segment_everything,
-                        inputs=[
-                            cond_img_e,
-                            input_size_slider_e,
-                            iou_threshold_e,
-                            conf_threshold_e,
-                            mor_check_e,
-                            contour_check_e,
-                            retina_check_e,
-                        ],
-                        outputs=segm_img_e)
-
-    segment_btn_p.click(segment_with_points,
-                        inputs=[cond_img_p],
-                        outputs=[segm_img_p, cond_img_p])
     
     segment_btn_t.click(segment_everything,
                         inputs=[
                             cond_img_t,
                             input_size_slider_t,
-                            iou_threshold_t,
-                            conf_threshold_t,
-                            mor_check_t,
-                            contour_check_t,
-                            retina_check_t,
+                            iou_threshold,
+                            conf_threshold,
+                            mor_check,
+                            contour_check,
+                            retina_check,
                             text_box,
+                            wider_check,
                         ],
                         outputs=segm_img_t)
 
@@ -361,7 +365,7 @@ with gr.Blocks(css=css, title='Fast Segment Anything') as demo:
     
     def clear_text():
         return None, None, None
-    
+
     clear_btn_e.click(clear, outputs=[cond_img_e, segm_img_e])
     clear_btn_p.click(clear, outputs=[cond_img_p, segm_img_p])
     clear_btn_t.click(clear_text, outputs=[cond_img_p, segm_img_p, text_box])

utils/tools.py

@@ -9,11 +9,14 @@ import clip
 
 
 def convert_box_xywh_to_xyxy(box):
-    x1 = box[0]
-    y1 = box[1]
-    x2 = box[0] + box[2]
-    y2 = box[1] + box[3]
-    return [x1, y1, x2, y2]
+    if len(box) == 4:
+        return [box[0], box[1], box[0] + box[2], box[1] + box[3]]
+    else:
+        result = []
+        for b in box:
+            b = convert_box_xywh_to_xyxy(b)
+            result.append(b)               
+    return result
 
 
 def segment_image(image, bbox):
@@ -323,8 +326,8 @@ def fast_show_mask_gpu(
 # clip
 @torch.no_grad()
 def retriev(
-    model, preprocess, elements, search_text: str, device
-) -> int:
+    model, preprocess, elements: [Image.Image], search_text: str, device
+):
     preprocessed_images = [preprocess(image).to(device) for image in elements]
     tokenized_text = clip.tokenize([search_text]).to(device)
     stacked_images = torch.stack(preprocessed_images)
@@ -348,19 +351,16 @@ def crop_image(annotations, image_like):
     cropped_boxes = []
     cropped_images = []
     not_crop = []
-    filter_id = []
-    # annotations, _ = filter_masks(annotations)
-    # filter_id = list(_)
+    origin_id = []
     for _, mask in enumerate(annotations):
         if np.sum(mask["segmentation"]) <= 100:
-            filter_id.append(_)
             continue
+        origin_id.append(_)
         bbox = get_bbox_from_mask(mask["segmentation"])  # mask 的 bbox
         cropped_boxes.append(segment_image(image, bbox))  # 保存裁剪的图片
         # cropped_boxes.append(segment_image(image,mask["segmentation"]))
         cropped_images.append(bbox)  # 保存裁剪的图片的bbox
-
-    return cropped_boxes, cropped_images, not_crop, filter_id, annotations
+    return cropped_boxes, cropped_images, not_crop, origin_id, annotations
 
 
 def box_prompt(masks, bbox, target_height, target_width):
@@ -415,8 +415,8 @@ def point_prompt(masks, points, point_label, target_height, target_width):  # nu
     return onemask, 0
 
 
-def text_prompt(annotations, text, img_path, device):
-    cropped_boxes, cropped_images, not_crop, filter_id, annotations_ = crop_image(
+def text_prompt(annotations, text, img_path, device, wider=False, threshold=0.9):
+    cropped_boxes, cropped_images, not_crop, origin_id, annotations_ = crop_image(
         annotations, img_path
     )
     clip_model, preprocess = clip.load("./weights/CLIP_ViT_B_32.pt", device=device)
@@ -425,5 +425,18 @@ def text_prompt(annotations, text, img_path, device):
     )
     max_idx = scores.argsort()
     max_idx = max_idx[-1]
-    max_idx += sum(np.array(filter_id) <= int(max_idx))
+    max_idx = origin_id[int(max_idx)]
+
+    # find the biggest mask which contains the mask with max score
+    if wider:
+        mask0 = annotations_[max_idx]["segmentation"]
+        area0 = np.sum(mask0)
+        areas = [(i, np.sum(mask["segmentation"])) for i, mask in enumerate(annotations_) if i in origin_id]
+        areas = sorted(areas, key=lambda area: area[1], reverse=True)
+        indices = [area[0] for area in areas]
+        for index in indices:
+            if index == max_idx or np.sum(annotations_[index]["segmentation"] & mask0) / area0 > threshold:
+                max_idx = index
+                break
+
     return annotations_[max_idx]["segmentation"], max_idx

utils/tools_gradio.py

@@ -103,7 +103,7 @@ def fast_show_mask(
     annotation = annotation[sorted_indices]
 
     index = (annotation != 0).argmax(axis=0)
-    if random_color == True:
+    if random_color:
         color = np.random.random((mask_sum, 1, 1, 3))
     else:
         color = np.ones((mask_sum, 1, 1, 3)) * np.array([30 / 255, 144 / 255, 255 / 255])
@@ -121,7 +121,7 @@ def fast_show_mask(
         x1, y1, x2, y2 = bbox
         ax.add_patch(plt.Rectangle((x1, y1), x2 - x1, y2 - y1, fill=False, edgecolor='b', linewidth=1))
 
-    if retinamask == False:
+    if not retinamask:
         mask = cv2.resize(mask, (target_width, target_height), interpolation=cv2.INTER_NEAREST)
 
     return mask
@@ -145,7 +145,7 @@ def fast_show_mask_gpu(
     annotation = annotation[sorted_indices]
     # 找每个位置第一个非零值下标
     index = (annotation != 0).to(torch.long).argmax(dim=0)
-    if random_color == True:
+    if random_color:
         color = torch.rand((mask_sum, 1, 1, 3)).to(device)
     else:
         color = torch.ones((mask_sum, 1, 1, 3)).to(device) * torch.tensor(
@@ -168,7 +168,7 @@ def fast_show_mask_gpu(
                 (x1, y1), x2 - x1, y2 - y1, fill=False, edgecolor="b", linewidth=1
             )
         )
-    if retinamask == False:
+    if not retinamask:
         mask_cpu = cv2.resize(
             mask_cpu, (target_width, target_height), interpolation=cv2.INTER_NEAREST
         )