SAM Annotation Tool (SAMAT)

SAMAT is a hackable PyQt5 based GUI labelling tool for semantic segmentation task enhanced by an assistance feature, reminiscent of the functionality offered by the Photoshop Magic Wand tool.

Note: Magic Wand is optional tool and uses SAM masks that can be generated by python script in scripts directory of this repo.

Showcase

Animation below shows annotation speed in real-time (SAM mask used).

Workflow

After environment setup:

• Organize your data following this structure.
• By design example_dataset/classes.json is used for the SAM GUI to maintain consistency across usages. So, define the file properly once and forget it :)
• (optional) Generate SAM masks from images via scripts/preprocess_dataset.py. Necessary for SAM assistance option.

python scripts/preprocess_dataset.py

• Specify path to your data in config.toml.
• Run GUI via __main__.py (prerequisites should be satisfied).

python __main__.py

• Annotate using brush (and/or) SAM assistance (label is saved on sample switch).
• See Shortcuts section for brush and panel controls (ex. changing image sample etc.)
• (optional) Generate semanticData dataset from saved labels for training class-aware semantic model finetune-anything via scripts/preprocess_dataset.py.

python scripts/postprocess_dataset.py

Getting started

Prerequisites: Environment setup and SAM Model

Annotation tool itself requires only:

• Python 3.11
• PyQt5
• numpy

(optional) In order to generate SAM masks for Magic Wand, you will need to install:

• PyTorch to run SAM model inference
• segment-anything itself + related libs
• Download ViT-H SAM model from https://github.com/facebookresearch/segment-anything and place the subsequent sam_vit_h_4b8939.pth file in assets folder.

(Hint) Setting up conda environment with requirements.txt should cover the prerequisites. But downloading ViT-H model should still be done manually.

Example setup (assuming Miniconda/Anaconda installed):

git clone https://github.com/ayrus144/sam_annotator.git
cd sam_annotator

conda create -n samat python=3.11
conda activate samat

## Environment setup with requirements.txt
pip install -r requirements.txt
pip install torch==2.7.1 torchvision==0.22.1 torchaudio==2.7.1 --index-url https://download.pytorch.org/whl/cu118 # cuda support
## works for both "cuda" and "cpu" devices
## In future, check if SAM2 works as expected and make changes to requirements file accordingly.

Dataset folder structure

Your input data MUST follow this structure:

── my_dataset
   ├── images (input)
   |   ├── 000001.{jpg/png}
   |   ├── 000002.{jpg/png}
   |   └── ...
   ├── labels (optional) - previously saved annotations
   |   ├── 000001.png
   |   ├── 000002.png
   |   └── ...
   └── sam (optional) - output of preprocess_dataset.py
       ├── 000001.png
       ├── 000002.png
       └── ...

Your output structure may look like this:

── my_dataset
   ├── images (input)
   ├── labels (output) - output of samAnnotator.exe / __main__.py
   |   ├── 000001.png (updated if already present)
   |   ├── 000002.png (updated if already present)
   |   └── ...
   ├── sam (optional)
   └── semanticData (extra) - output of postprocess_dataset.py
       ├── img - images folder split as train and val
       |   ├── train
       |   └── val
       ├── ann - labels formated to (class_id, H, W)
       |   ├── train
       |   └── val
       ├── ann_vis - ann folder visualized as color masks
       |   ├── train
       |   └── val
       └── metainfo.yaml - has list of class_names listed in classes.json
           Example metainfo.yaml
           class_names
           - human
           - car
           - road
           ...

• images contains image files of any format you want to label.
• labels contains colored masks files with labels. Automatically created (if no labels yet) and updated real-time (after sample switch).
• sam contains mask files from SAM annotations (8-bit grayscale). Run scripts/preprocess_dataset.py to generate this folder.
• semanticData contains all the file/folders required for training the class-aware semantic model with finetune-anything. Run scripts/postprocess_dataset.py to generate this folder.

Class Labels:

• To maintain similarity among labels used for annotation, a common classes.json at example_dataset folder is used.
• classes.json contains classes description that will be used for labeling.

Example classes.json:

{
    "classes": [
        { "id": 1, "name": "human", "color": "#FF0000" },
        { "id": 2, "name": "car", "color": "#00FF00" }
    ]
}

where:

• id field must coincide with number keys on keyboard, so start with 1 (not 0). Any number of classes allowed, but only first 9 have their shortcuts.
• name field is arbitrary and used only for display in GUI.
• color field specifies the hexadecimal-color this class would be displayed in GUI.

Note:

• Specify path to your data inside config.toml.
• Image files can have any valid format. But all output files are saved as .png files to avoid compression loss.
• Path to SAM weights is already specified in config.toml.

Future Ideas

These are some future ideas that can be implemented to improve the labels:

• Improve mask label quality with existing masks labels. For each object (human, car, road) in the mask, img pair:
  • Sample k random points from the object's mask and save as np.array in point_coords.
  • Re-run the SAM model on the img with sampled points as point_prompt.
  • Use sam.set_image(img) and sam.predict(point_coords, point_labels=np.ones(k)).
• It is important to annotate masks for all images consistently when using GUI, especially to use the masks for training a model.
  • For example, if image 1 (with human, car and road), image 2 (with human and road) and image 3 (with car only), and you intend to train a custom model only for human and car, then:

ANNOTATION 1 - Incorrect 
── label 1 (corresponds to image 1)
   ├── human mask
   └── road mask
── label 2 (corresponds to image 2)
   └── road mask
── label 3 (corresponds to image 3)
   └── car mask
Incorrect because car mask is missing in label 1 and human mask is missing in label 2. 
If a model was trained over this data, the model will have difficult time because
MODEL is able to see    --->    But is learning
    image 1                         image 1
    - human                         - human
    - car                           - NO car (mask missing)
    image 2                         image 2
    - human                         - NO human (mask missing)
    image 3                         image 3
    - car                           - car
Presence of road mask is not the issue.

ANNOTATION 2 - Correct
── label 1 (corresponds to image 1)
   ├── human mask
   ├── car mask
── label 2 (corresponds to image 2)
   ├── human mask
── label 3 (corresponds to image 3)
   └── car mask
Correct because all labels have masks of object we are interested in. Absence of road masks is not an issue. 

• Given you have multiple images from same location, ideally masks should overlap for same objects, so:
  • For each object in the image, calculate the IoU (Intersection over Union) with a reference.
  • Retain only those masks that have object IoU's greater than a defined threshold.

Shortcuts

Shortcut	Description
Left Mouse Button	Draw with brush + fill region (in SAM mode)
Right Mouse Button	Pan motion on zoomed-in image
Mouse Wheel	Zoom in/out
Ctrl + Mouse Wheel	Change brush size
1-9	Select class (color to draw on label layer)
E	Eraser tool (transparent brush)
Space	Reset zoom
C	Clear label
S	Switch SAM assistance mode on/off
,/.	Previous/Next sample