delete image_predictor, only use automatic mode

app.py

@@ -6,7 +6,6 @@ import numpy as np
 import supervision as sv
 import torch
 from PIL import Image
-from gradio_image_prompter import ImagePrompter
 
 from utils.models import load_models, CHECKPOINT_NAMES, MODE_NAMES, \
     MASK_GENERATION_MODE, BOX_PROMPT_MODE
@@ -46,7 +45,7 @@ if torch.cuda.get_device_properties(0).major >= 8:
     torch.backends.cudnn.allow_tf32 = True
 
 MASK_ANNOTATOR = sv.MaskAnnotator(color_lookup=sv.ColorLookup.INDEX)
-IMAGE_PREDICTORS, MASK_GENERATORS = load_models(device=DEVICE)
+MASK_GENERATORS = load_models(device=DEVICE)
 
 
 @spaces.GPU
@@ -54,39 +53,13 @@ IMAGE_PREDICTORS, MASK_GENERATORS = load_models(device=DEVICE)
 @torch.autocast(device_type="cuda", dtype=torch.bfloat16)
 def process(
     checkpoint_dropdown,
-    mode_dropdown,
     image_input,
-    image_prompter_input
 ) -> Optional[Image.Image]:
-    if mode_dropdown == BOX_PROMPT_MODE:
-        image_input = image_prompter_input["image"]
-        prompt = image_prompter_input["points"]
-        if len(prompt) == 0:
-            return image_input
-
-        model = IMAGE_PREDICTORS[checkpoint_dropdown]
-        image = np.array(image_input.convert("RGB"))
-        box = np.array([[x1, y1, x2, y2] for x1, y1, _, x2, y2, _ in prompt])
-
-        model.set_image(image)
-        masks, _, _ = model.predict(box=box, multimask_output=False)
-
-        # dirty fix; remove this later
-        if len(masks.shape) == 4:
-            masks = np.squeeze(masks)
-
-        detections = sv.Detections(
-            xyxy=sv.mask_to_xyxy(masks=masks),
-            mask=masks.astype(bool)
-        )
-        return MASK_ANNOTATOR.annotate(image_input, detections)
-
-    if mode_dropdown == MASK_GENERATION_MODE:
-        model = MASK_GENERATORS[checkpoint_dropdown]
-        image = np.array(image_input.convert("RGB"))
-        result = model.generate(image)
-        detections = sv.Detections.from_sam(result)
-        return MASK_ANNOTATOR.annotate(image_input, detections)
+    model = MASK_GENERATORS[checkpoint_dropdown]
+    image = np.array(image_input.convert("RGB"))
+    result = model.generate(image)
+    detections = sv.Detections.from_sam(result)
+    return MASK_ANNOTATOR.annotate(image_input, detections)
 
 
 with gr.Blocks() as demo:
@@ -98,21 +71,10 @@ with gr.Blocks() as demo:
             label="Checkpoint", info="Select a SAM2 checkpoint to use.",
             interactive=True
         )
-        mode_dropdown_component = gr.Dropdown(
-            choices=MODE_NAMES,
-            value=MODE_NAMES[0],
-            label="Mode",
-            info="Select a mode to use. `box prompt` if you want to generate masks for "
-                 "selected objects, `mask generation` if you want to generate masks "
-                 "for the whole image.",
-            interactive=True
-        )
     with gr.Row():
         with gr.Column():
             image_input_component = gr.Image(
                 type='pil', label='Upload image')
-            image_prompter_input_component = ImagePrompter(
-                type='pil', label='Image prompt', visible=False)
             submit_button_component = gr.Button(
                 value='Submit', variant='primary')
         with gr.Column():
@@ -123,37 +85,18 @@ with gr.Blocks() as demo:
             examples=EXAMPLES,
             inputs=[
                 checkpoint_dropdown_component,
-                mode_dropdown_component,
                 image_input_component,
-                image_prompter_input_component,
             ],
             outputs=[image_output_component],
             cache_examples=False,
             run_on_click=True
         )
 
-
-    def on_mode_dropdown_change(text):
-        return [
-            gr.Image(visible=text == MASK_GENERATION_MODE),
-            ImagePrompter(visible=text == BOX_PROMPT_MODE)
-        ]
-
-    mode_dropdown_component.change(
-        on_mode_dropdown_change,
-        inputs=[mode_dropdown_component],
-        outputs=[
-            image_input_component,
-            image_prompter_input_component
-        ]
-    )
     submit_button_component.click(
         fn=process,
         inputs=[
             checkpoint_dropdown_component,
-            mode_dropdown_component,
             image_input_component,
-            image_prompter_input_component,
         ],
         outputs=[image_output_component]
     )

utils/models.py

@@ -21,11 +21,9 @@ CHECKPOINTS = {
 def load_models(
     device: torch.device
 ) -> Tuple[Dict[str, SAM2ImagePredictor], Dict[str, SAM2AutomaticMaskGenerator]]:
-    image_predictors = {}
     mask_generators = {}
     for key, (config, checkpoint) in CHECKPOINTS.items():
         model = build_sam2(config, checkpoint, device=device)
-        image_predictors[key] = SAM2ImagePredictor(sam_model=model)
         mask_generators[key] = SAM2AutomaticMaskGenerator(
             model=model,
             points_per_side=32,
@@ -36,4 +34,4 @@ def load_models(
             crop_n_layers=1,
             box_nms_thresh=0.7,
         )
-    return image_predictors, mask_generators
+    return mask_generators