search

facebookresearch / segment-anything

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

How can i segment only one class ? #126

Open bisnoisk opened 1 year ago

FullStackSimon commented 1 year ago

Something like this perhaps?

Mouse Location or known x,y coordinates

from segment_anything import SamPredictor, sam_model_registry

points = np.array([[10, 20]])

sam = sam_model_registry["<model_type>"](checkpoint="<path/to/checkpoint>")
predictor = SamPredictor(sam)
predictor.set_image(<your_image>)
masks, _, _ = predictor.predict(point_coords=points, multimask_output=false)

Or bounding box...

from segment_anything import SamPredictor, sam_model_registry

box = np.array([10, 10, 50, 50])

sam = sam_model_registry["<model_type>"](checkpoint="<path/to/checkpoint>")
predictor = SamPredictor(sam)
predictor.set_image(<your_image>)
masks, _, _ = predictor.predict(box=box, multimask_output=false)
bisnoisk commented 1 year ago

without bbox or point cords , like i want to segment one particular coco class so what should i do ?

FullStackSimon commented 1 year ago

I'm not sure if I understand you correctly but if you don't want to provide a prompt for the predictor I think the only option you have is to generate masks for the entire image and then choose the one you want somehow.

As explained in the readme...

from segment_anything import SamAutomaticMaskGenerator, sam_model_registry
sam = sam_model_registry["<model_type>"](checkpoint="<path/to/checkpoint>")
mask_generator = SamAutomaticMaskGenerator(sam)
masks = mask_generator.generate(<your_image>)
XinyueZ commented 1 year ago

without bbox or point cords , like i want to segment one particular coco class so what should i do ?

@bisnoisk I think you should collaborate with other models or devices, ie. object detection models which can provide bboxes. Use the bboxes from there you can then feed into the SAM. Find out through the return masks.

This is what I have understood, unsure I was right or wrong, either.

klarrieu commented 1 year ago

This might be what you're looking for?

https://github.com/IDEA-Research/Grounded-Segment-Anything