search

facebookresearch / segment-anything-2

The repository provides code for running inference with the Meta Segment Anything Model 2 (SAM 2), links for downloading the trained model checkpoints, and example notebooks that show how to use the model.
Apache License 2.0
10.68k stars 857 forks source link

Multiple bounding box as prompt #267

Open YangJae96 opened 2 weeks ago

YangJae96 commented 2 weeks ago

Hi. Thank you for great work!

I was wondering if it is possible to give multiple bounding boxes as a prompt.

In the demo code, it just give only 1 box as prompt and uses predictor.add_new_points_or_box to add the box.

Is it possible to give two boxes at once like below code?

boxes = np.array([[300, 0, 500, 400],
                [180, 140,290,400]], dtype=np.float32)

_, out_obj_ids, out_mask_logits = predictor.add_new_points_or_box(
    inference_state=inference_state,
    frame_idx=ann_frame_idx,
    obj_id=ann_obj_id,
    box=boxes,
)

Thanks in advance.

AlexMcClay commented 2 weeks ago

Yes it is possible.

if you do something like this

boxes = np.array([[300, 0, 500, 400],
                [180, 140,290,400]])

masks, scores, logits = predictor.predict(
                    box=boxes,
                    multimask_output=False,
                )
print("Masks: ", masks.shape)

the output will be something like this.

Masks: (2, 1, 512, 512)

basically (n, 1, image width, image height) where n is the number of boxes.

That works for me, and im going off of the example image notebook

https://github.com/facebookresearch/segment-anything-2/blob/main/notebooks/image_predictor_example.ipynb

They explain it in the "Batched prompt inputs" section