Prompting for Segmentation

Introduction

We implement a module for segmenting images using (points and bounding box) prompts, in PyTorch. As a model, we use pretrained FastSAM provided by ultralytics.

Setting Up the Environment

1. Install Conda, if not already installed.
2. Clone the repository

   git clone https://github.com/byrkbrk/prompting-for-segmentation.git
   

3. Change the directory:

   cd prompting-for-segmentation
   

4. For macos, run:

   conda env create -f prompting-for-sam_macos.yaml
   

   For linux or windows, run:

   conda env create -f prompting-for-sam_linux.yaml
   

5. Activate the environment:

   conda activate prompting-for-sam
   

Prompts

Check it out how to use:

python3 segment.py -h

Output:

Segment given image

positional arguments:
  image_name            Name of the image that be processed
  point_or_bbox_prompts
                        List of point prompts in (height, width) or bbox prompts in (x_min, y_min, x_max, y_max)

options:
  -h, --help            show this help message and exit
  --label_prompts LABEL_PROMPTS [LABEL_PROMPTS ...]
                        List of labels of prompts. Use the values 0 and 1 for negative and positive prompts, respectively
  --image_size IMAGE_SIZE [IMAGE_SIZE ...]
                        Size (height, width) to which the image be transformed
  --checkpoint_name {FastSAM-x.pt,FastSAM-s.pt}
                        Name of the pretrained model for FastSAM
  --device DEVICE       Name of the device on which the model be run

Point prompts

Using single prompt

python3 segment.py dogs.jpg 400,400 --label_prompts 1 --image_size 1024 1024

Using multiple prompts

python3 segment.py dogs.jpg 400,400 700,400 --label_prompts 0 1 --image_size 1024 1024

Bounding box prompts

Using single prompt

python3 segment.py dogs.jpg 625,625,700,700 --label_prompts 1 --image_size 1024 1024

Using multiple prompts

python3 segment.py dogs.jpg 500,200,800,900 510,210,810,610 --label_prompts 1 0 --image_size 1024 1024