Interactive annotation with Segment Anything Model

Interactive annotation in Label Studio with Segment Anything Model

https://github.com/shondle/label-studio-ml-backend/assets/106922533/42a8a535-167c-404a-96bd-c2e2382df99a

Use Facebook’s Segment Anything Model with Label Studio!

Quickstart

Using Docker Compose (recommended)

To start the server with the lightweight mobile version of SAM, run the following command:

docker-compose up

GPU support

By default, the docker-compose file runs the model on the CPU. If you have a GPU, you can enable it by adding the following lines in docker-compose.yml:

environment:
  - NVIDIA_VISIBLE_DEVICES=all
deploy:
  reservations:
    devices:
      - driver: nvidia
        count: 1
        capabilities: [gpu]

About the models

There are two models in this repo that you can use:

Advanced Segment Anything Model
ONNX Segment Anything Model

Advanced Segment Anything Model

The Advanced Segment Anything Model introduces the ability to combine a
multitude of different prompts to achieve a prediction, and the ability to use
MobileSAM.

Mix one rectangle label with multiple positive keypoints to refine your
predictions.
Use negative keypoints to remove areas from predictions for
increased control.
Use MobileSAM, an extremely lightweight alternative to the heavy Segment
Anything Model from Facebook, to retrieve predictions. This can run inference
within a second using a laptop GPU.

ONNX Segment Anything Model

The ONNX Segment Anything Model gives you the ability to use either a single
keypoint or a single rectangle label to prompt the original SAM.

This offers a much faster prediction than using the original Segment Anything
Model.
The downside is that image size must be specified before using the ONNX model, and
cannot be generalized to other image sizes while labeling. Also, this does not yet
offer the mixed labeling and refinement that AdvancedSAM does.

Model configuration options

Each model has different pros and cons. Consider which is best for your project:

AdvancedSAM
- Mobile SAM Configuration
  - Pros: Lightweight model that can be run on laptops, and can mix many
    different combinations of input prompts to fine-tune prediction.
  - Cons: Lower accuracy than Facebook’s original SAM architecture.
- Original SAM architecture
  - Pros: Higher accuracy than MobileSAM, with ability to mix many different
    combinations of input prompts to fine-tune predictions.
  - Cons: Takes long to gather predictions (~2s to create embedding of an
    image), requires access to good GPUs.
ONNXSAM
- Original SAM Architecture
  - Pros: Much faster than when you use it in Advanced SAM.
  - Cons: Can only use one smart label per prediction. Image size must be
    defined before generating the ONNX model. Cannot label images with
    different sizes without running into issues.

Setup

The Label Studio SAM backend works best if you have Local
Storage enabled for
your project. It is also possible to set up shared local storage, but is not
recommended. Currently, the backend does not work with cloud storage (S3,
Azure, GCP).

Setting up the Label Studio server

Enabling local storage file serving

You can enable local storage file serving by setting the following variables:

LABEL_STUDIO_LOCAL_FILES_DOCUMENT_ROOT=<path_to_image_data>
LABEL_STUDIO_LOCAL_FILES_SERVING_ENABLED=true

For example, if you’re launching Label Studio with Docker, you can enable these variables with

docker run -it -p 8080:8080 \
               -v $(pwd)/mydata:/label-studio/data \
               --env LABEL_STUDIO_LOCAL_FILES_SERVING_ENABLED=true \
               --env LABEL_STUDIO_LOCAL_FILES_DOCUMENT_ROOT=/label-studio/data/images \
               heartexlabs/label-studio:latest

Note the IP address that you are running your Label Studio instance as the
LABEL_STUDIO_HOST. This will be necessary for setting up the connection to your
SAM model.

Because you are hosting both Label Studio and the ML backend in
Docker containers, the hostname localhost will not resolve to the correct
address. There are a number of ways to determine your host IP address. These
can include calling either ip a or ifconfig from the command line and
inspecting the output, or finding the address that has been assigned to your
computer through the system network configuration settings.

Obtain your API token

Log into the Label Studio interface (in the example above, at
http://<LABEL_STUDIO_HOST>:8080).

Go to the Account & Settings
page, and make a note of the Access Token, which we will use later as
the LABEL_STUDIO_ACCESS_TOKEN.

Setting up the SAM backend

Clone the repository

Make a clone of this repository on your host system and move it into the working
directory.

git clone https://github.com/humansignal/label-studio-ml-backend
cd label-studio-ml-backend/label_studio_ml/examples/segment_anything_model

Using Docker Compose (recommended)

We suggest using Docker Compose to host and
run the backend. For GPU support, please consult the Docker Compose GPU Access
Guide to understand how to pass
through GPU resources to services.

Edit the docker-compose.yml file and fill in the values for the
LABEL_STUDIO_HOST and LABEL_STUDIO_ACCESS_TOKEN variables for your particular
installation. Be sure to append the port that Label Studio is running on in
your LABEL_STUDIO_HOST variable, for example http://192.168.1.36:8080 if
Label Studio is running on port 8080.

Run the command docker compose up --build to build the container and run it
locally.

Setting up the backend manually

Download model weights

This step is only necessary if you are not using the Docker build for this model.

For MobileSAM install the weights using this
link
and place them in a folder (along with the advanced_sam.py and onnx_sam.py files)
For using regular SAM and/or ONNX- Follow SAM installation instructions with
pip. Then, install
the ViT-H SAM model
For the ONNX model install using python onnxconverter.py

You can download all weights and models using the following command:

./download_models.sh

Installation requirements

Change your directory to this folder and then install all of the python requirements.

pip install -r requirements.txt

Adjust variables and `_wsgi.py` depending on your choice of model

You can set the following environment variables to change the behavior of the model.

LABEL_STUDIO_HOST sets the endpoint of the Label Studio host. Must begin with http://
LABEL_STUDIO_ACCESS_TOKEN sets the API access token for the Label Studio host.
SAM_CHOICE selects which model to use.
- SAM_CHOICE=MobileSAM to use MobileSAM (default)
- SAM_CHOICE=SAM to use the original SAM model.
- SAM_CHOICE=ONNX to use the ONNX model.

Start the Backend

You can now manually start the ML backend.

python _wsgi.py

docker-compose up

to start the backend in a Docker container

MOBILESAM_CHECKPOINT=path/to/mobile_sam.pt label-studio-ml start segment_anything_model/

Note: If you see an error on MacOS, try set the environment variable KMP_DUPLICATE_LIB_OK=True

Set up a project in Label Studio for Segment Anything

Log into your Label Studio instance and perform the following steps.

Create a new project.
Under the Labeling Setup step when creating the project, or under
Labeling Interface in the project settings, paste the sample
template into the code dialog. Save the interface.
Go to the Model page in the project settings and click Connect Model.
Enter a title for the model, and the URL for the instance of the model you
just created.

If you’re running Label Studio in Docker or on another host, you
should use the direct IP address of where the model is hosted (localhost
will not work). Be sure to include the port number that the model is hosted on
(the default is 9090). For example, if the model is hosted on 192.168.1.36,
the URL for the model would be http://192.168.1.36:9090
Click Validate and Save.

You can now upload images into your project and begin annotating.

The video also goes over this process, but does part of it while in the newly created project menu.

Creating annotations

See this video tutorial
to get a better understanding of the workflow when annotating with SAM.

Use the Alt hotkey to alter keypoint positive and negative labels.

Notes for AdvancedSAM

Please watch this video first
For the best experience, follow the video tutorial above and uncheck ‘Auto
accept annotation suggestions’ when running predictions.
After generating the prediction from an assortment of inputs, make sure you click the
checkmark that is outside of the image to finalize the region (this should either
be above or below the image. Watch the video for a visual guide).
There may be a checkmark inside the image next to a generated prediction,
but do not use that one.

For some reason, the checkmark that is not on the
image itself cleans the other input prompts used for generating
the region, and only leaves the predicted region after being clicked (this is
the most compatible way to use the backend).
You may run into problems creating instances of the same class if you click
the checkmark on the image and it leaves the labels used to guide the
region.
After labeling your object, select the label in the menu and select the type
of brush label you would like to give it at the top of your label keys below
your image. This allows you to change the class of your prediction. See the
video for a better explanation.
Only the negative keypoints can be used for subtracting from prediction
areas for the model. Positive keypoints and rectangles tell the model
areas of interest to make positive predictions.
Multiple keypoints may be used to provide areas for the model where predictions
should be extended. Only one rectangle label may be used when generating
a prediction as an area where the model prediction should occur/be extended.

If you place multiple rectangle labels, the model will use the newest
rectangle label along with all other keypoints when aiding the model
prediction.

Notes for ONNX

The ONNX model uses the orig_img_size in onnx_converter.py that defines
an image ratio for the ONNX model. Change this to the ratio of the images
that you are labeling before generating the model. If you are labeling images
of different sizes, use Advanced SAM instead, or generate a new ONNX model
for different image groups with different sizes. If you do not adjust
orig_img_size, and your image aspect ratios do not match what is
already defined, then your predictions will be offset from the image. If
using docker compose to launch the model, be sure to rebuild the host
container.
Make sure you adjust orig_img_size BEFORE generating the ONNX model when
using onnx_converter.py
Guide on changing the code - "orig_im_size": torch.tensor([#heightofimages, #widthofimages], dtype=torch.float),

Notes for exporting

COCO and YOLO format is not supported (this project exports using brush
labels, so try NumPy or PNG export instead)

Labeling configs

When using the AdvancedSAM

Give one brush label per class you want to annotate.
Hold the Alt hotkey to create negative keypoints.
Add one rectangle label for each of your classes that you want to annotate.
The video reviews these points as well if you are
confused after reading this.

Base example:

<View>
<Style>
  .main {
    font-family: Arial, sans-serif;
    background-color: #f5f5f5;
    margin: 0;
    padding: 40px 5px 5px 5px;
  }
  .container {
    display: flex;
    justify-content: space-between;
    margin-bottom: 20px;
  }
  .column {
    flex: 1;
    padding: 10px;
    margin: 5px; 
    background-color: #fff;
    border-radius: 5px;
    box-shadow: 0 2px 5px rgba(0, 0, 0, 0.1);
    text-align: center;
  }
  .column .title {
    margin: 0;
    color: #333;
  }
  .column .label {
    margin-top: 10px;
    padding: 10px;
    padding-bottom: 7px; 
    background-color: #f9f9f9;
    border-radius: 3px;
  }
  .lsf-labels {
    margin: 5px 0 0 0; 
  }
  .image-container {
    width: 100%;
    height: 300px;
    background-color: #ddd;
    border-radius: 5px;
  }
</Style>
  
<View className="main">
  <View className="container">
    <View className="column">
      <HyperText value="" name="h1" className="help" inline="true">
        Brush for manual labeling
      </HyperText>
      <View className="label">        
        <BrushLabels name="tag" toName="image">
          <Label value="Foreground" background="#FF0000" />
          <Label value="Background" background="#0d14d3" />
        </BrushLabels>
      </View>
    </View>
    
    <View className="column">
      <HyperText value="" name="h2" className="help" inline="true">
        <span title="1. Click purple auto Keypoints/Rectangle icon on toolbar. 2. Click Foreground/Background label here">
          Keypoints for auto-labeling
        </span>
      </HyperText>
      <View className="label">
        <KeyPointLabels name="tag2" toName="image" smart="true">
          <Label value="Foreground" smart="true" background="#FFaa00" showInline="true" />
          <Label value="Background" smart="true" background="#00aaFF" showInline="true" />
        </KeyPointLabels>
      </View>
    </View>
    
    <View className="column">
      <HyperText value="" name="h3" className="help" inline="true">
        <span title="1. Click purple auto Keypoints/Rectangle icon on toolbar. 2. Click Foreground/Background label here">
          Rectangles for auto-labeling
        </span>
      </HyperText>
      <View className="label">
        <RectangleLabels name="tag3" toName="image" smart="true">
          <Label value="Foreground" background="#FF00FF" showInline="true" />
          <Label value="Background" background="#00FF00" showInline="true" />
        </RectangleLabels>
      </View>
    </View>
    
  </View>
  
  <View className="image-container">
    <Image name="image" value="$image" zoom="true" zoomControl="true" />
  </View>
  
</View>
</View>

When using the ONNX model

Label values for the keypoints, rectangle, and brush labels must correspond.
Other than that, make sure that smart="True" for each keypoint label and
rectangle label.

For the ONNX model:

<View>
  <Image name="image" value="$image" zoom="true"/>
  <BrushLabels name="tag" toName="image">
    <Label value="Banana" background="#FF0000"/>
    <Label value="Orange" background="#0d14d3"/>
  </BrushLabels>
  <KeyPointLabels name="tag2" toName="image" smart="true">
    <Label value="Banana" smart="true" background="#000000" showInline="true"/>
    <Label value="Orange" smart="true" background="#000000" showInline="true"/>
  </KeyPointLabels>
  <RectangleLabels name="tag3" toName="image" smart="true">
    <Label value="Banana" background="#000000" showInline="true"/>
    <Label value="Orange" background="#000000" showInline="true"/>
  </RectangleLabels>
</View>

Credits

Original Segment Anything Model paper-

@article{kirillov2023segany,
  title={Segment Anything},
  author={Kirillov, Alexander and Mintun, Eric and Ravi, Nikhila and Mao, Hanzi and Rolland, Chloe and Gustafson, Laura and Xiao, Tete and Whitehead, Spencer and Berg, Alexander C. and Lo, Wan-Yen and Doll{\'a}r, Piotr and Girshick, Ross},
  journal={arXiv:2304.02643},
  year={2023}
}

MobileSAM paper-

@article{mobile_sam,
  title={Faster Segment Anything: Towards Lightweight SAM for Mobile Applications},
  author={Zhang, Chaoning and Han, Dongshen and Qiao, Yu and Kim, Jung Uk and Bae, Sung-Ho and Lee, Seungkyu and Hong, Choong Seon},
  journal={arXiv preprint arXiv:2306.14289},
  year={2023}
}

Interactive annotation with Segment Anything Model

Interactive annotation in Label Studio with Segment Anything Model

Quickstart

Using Docker Compose (recommended)

GPU support

About the models

Advanced Segment Anything Model

ONNX Segment Anything Model

Model configuration options

Setup

Setting up the Label Studio server

Enabling local storage file serving

Obtain your API token

Setting up the SAM backend

Clone the repository

Using Docker Compose (recommended)

Setting up the backend manually

Download model weights

Installation requirements

Adjust variables and _wsgi.py depending on your choice of model

Start the Backend

Set up a project in Label Studio for Segment Anything

Creating annotations

Notes for AdvancedSAM

Notes for ONNX

Notes for exporting

Labeling configs

When using the AdvancedSAM

When using the ONNX model

Credits

In this article

Adjust variables and `_wsgi.py` depending on your choice of model