feature： enable_segment_all

app.py

@@ -20,7 +20,6 @@ def load_sam_instance():
     if sam is None:
         gr.Info('Initialising SAM, hang in there...')
         if not os.path.exists(configs.model_ckpt_path):
-            gr.Info('Downloading weights from  hugging face hub')
             chkpt_path = hf_hub_download("ybelkada/segment-anything", configs.model_ckpt_path)
         else:
             chkpt_path = configs.model_ckpt_path
@@ -46,12 +45,12 @@ with block:
                 point_label_radio = gr.Radio(label='Point Label', choices=[1,0], value=1)
                 reset_btn = gr.Button('Reset')
                 run_btn = gr.Button('Run', variant = 'primary')
+            gr.Examples(examples=[['examples/cat-256.png','examples/cat-256.png']],inputs=[input_image, raw_image])
         with gr.Column():
             with gr.Tab('Cutout'):
                 cutout_gallery = gr.Gallery()
             with gr.Tab('Annotation'):
                 masks_annotated_image = gr.AnnotatedImage(label='Segments')
-            gr.Examples(examples=[['examples/cat-256.png','examples/cat-256.png']],inputs=[input_image, raw_image])
     
     # components
     components = {point_coords, point_labels, raw_image, input_image, point_label_radio, reset_btn, run_btn, cutout_gallery, masks_annotated_image}
@@ -70,7 +69,7 @@ with block:
         x, y = evt.index
         color = red if point_label_radio == 0 else blue
         img = np.array(input_image)
-        cv2.circle(img, (x, y), 10, color, -1)
+        cv2.circle(img, (x, y), 5, color, -1)
         img = Image.fromarray(img)
         point_coords.append([x,y])
         point_labels.append(point_label_radio)
@@ -78,19 +77,23 @@ with block:
     input_image.select(on_input_image_select, [input_image, point_coords, point_labels, point_label_radio], [input_image, point_coords, point_labels], queue=False)
 
     # event - inference
-    def on_run_btn_click(data):
+    def on_run_btn_click(inputs):
         sam = load_sam_instance()
-        image = data[raw_image]
-        if len(data[point_coords]) == 0:
-            masks, _ = service.predict_all(sam, image)
+        image = inputs[raw_image]
+        if len(inputs[point_coords]) == 0:
+            if configs.enable_segment_all:
+                masks, _ = service.predict_all(sam, image)
+            else:
+                gr.Warning('Segment-all disabled, set point label(s) before running')
+                return inputs[cutout_gallery], inputs[masks_annotated_image]
         else:
             masks, _ = service.predict_conditioned(sam,
                                                    image, 
-                                                   point_coords=np.array(data[point_coords]), 
-                                                   point_labels=np.array(data[point_labels]))
+                                                   point_coords=np.array(inputs[point_coords]), 
+                                                   point_labels=np.array(inputs[point_labels]))
         annotated = (image, [(masks[i], f'Mask {i}') for i in range(len(masks))])
         cutouts = [service.cutout(image, mask) for mask in masks]
-        return cutouts, annotated, masks
+        return cutouts, annotated
     run_btn.click(on_run_btn_click, components, [cutout_gallery, masks_annotated_image], queue=True)
 
 if __name__ == '__main__':

app_configs.py

@@ -2,3 +2,4 @@ model_type = r'vit_b'
 # model_ckpt_path = None
 model_ckpt_path = "checkpoints/sam_vit_b_01ec64.pth"
 device = 'cpu'
+enable_segment_all = False