Semantic-SAM

Semantic-SAM: Segment and Recognize Anything at Any Granularity

In this work, we introduce Semantic-SAM, a universal image segmentation model to enable segment and recognize anything at any desired granularity. We have trained on the whole SA-1B dataset and our model can reproduce SAM and beyond it.

:grapes: [Read our arXiv Paper]  

:apple: [Try Auto Generation with Controllable Granularity Demo]   :apple: [Try Interactive Multi-Granularity Demo]  

:rocket: Features

:fire: Reproduce SAM. SAM training is a sub-task of ours. We have released the training code to reproduce SAM training.

:fire: Beyond SAM. Our newly proposed model offers the following attributes from instance to part level:

• Granularity Abundance. Our model can produce all possible segmentation granularities for a user click with high quality, which enables more controllable and user-friendly interactive segmentation.
• Semantic Awareness. We jointly train SA-1B with semantically labeled datasets to learn the semantics at both object-level and part-level.
• High Quality. We base on the DETR-based model to implement both generic and interactive segmentation, and validate that SA-1B helps generic and part segmentation. The mask quality of multi-granularity is high.

:rocket: News

:fire: We release the training and inference code and demo link of DINOv, which can handle in-context visual prompts for open-set and referring detection & segmentation. Check it out!

:fire: We release the demo code for controllable mask auto-generation with different granularity prompts!

Segment everything for one image. We output controllable granularity masks from semantic, instance to part level when using different granularity prompts.

:fire: We release the demo code for mask auto-generation!

Segment everything for one image. We output more masks with more granularity.

:fire: We release the demo code for interactive segmentation!

One click to output up to 6 granularity masks. Try it in our demo!

:fire: We release the training and inference code and checkpoints (SwinT, SwinL) trained on SA-1B!

:fire: We release the training code to reproduce SAM!

Our model supports a wide range of segmentation tasks and their related applications, including:

• Generic Segmentation
• Part Segmentation
• Interactive Multi-Granularity Segmentation with Semantics
• Multi-Granularity Image Editing

: Related projects:

• Mask DINO: We build upon Mask DINO which is a unified detection and segmentation model to implement our model.
• OpenSeeD: Strong open-set segmentation methods based on Mask DINO. We base on it to implement our open-vocabulary segmentation.
• SEEM: Segment using a wide range of user prompts.
• VLPart: Going denser with open-vocabulary part segmentation.

:unicorn: Getting Started

:hammer_and_wrench: Installation

pip3 install torch==1.13.1 torchvision==0.14.1 --extra-index-url https://download.pytorch.org/whl/cu113
python -m pip install 'git+https://github.com/MaureenZOU/detectron2-xyz.git'
pip install git+https://github.com/cocodataset/panopticapi.git
git clone https://github.com/UX-Decoder/Semantic-SAM
cd Semantic-SAM
python -m pip install -r requirements.txt

export DATASET=/pth/to/dataset  # path to your coco data

:star: A few lines to get generated results

First download a checkpoint from model zoo.

• For interactive multi-granularity segmentation

…

• For mask auto generation

from semantic_sam import prepare_image, plot_results, build_semantic_sam, SemanticSamAutomaticMaskGenerator
original_image, input_image = prepare_image(image_pth='examples/dog.jpg')  # change the image path to your image
mask_generator = SemanticSamAutomaticMaskGenerator(build_semantic_sam(model_type='<model_type>', ckpt='</your/ckpt/path>')) # model_type: 'L' / 'T', depends on your checkpint
masks = mask_generator.generate(input_image)
plot_results(masks, original_image, save_path='../vis/')  # results and original images will be saved at save_path

Advanced usage:

• Level is set to [1,2,3,4,5,6] to use all six prompts by default
• You can change the input prompt for controllable mask auto-generation to get the granularity results you want. An example is shown in here
• Here are some examples of mask_generator for generating different granularity results

mask_generator = SemanticSamAutomaticMaskGenerator(semantic_sam, level=[1]) # [1] and [2] for semantic level.
mask_generator = SemanticSamAutomaticMaskGenerator(semantic_sam, level=[3]) # [3] for instance level.
mask_generator = SemanticSamAutomaticMaskGenerator(semantic_sam, level=[6]) # [4], [5], [6] for different part level.

:mosque: Data preparation

Please refer to prepare SA-1B data. Let us know if you need more instructions about it.

:volcano: Model Zoo

The currently released checkpoints are only trained with SA-1B data.

Name Training Dataset Backbone 1-IoU@Multi-Granularity 1-IoU@COCO(Max|Oracle) download Semantic-SAM | config SA-1B SwinT 88.1 54.5|73.8 model Semantic-SAM | config SA-1B SwinL 89.0 55.1|74.1 model

:arrow_forward: Demo

For interactive segmentation.

python demo.py --ckpt /your/ckpt/path

For mask auto-generation.

python demo_auto_generation.py --ckpt /your/ckpt/path

:sunflower: Evaluation

We do zero-shot evaluation on COCO val2017. $n is the number of gpus you use

For SwinL backbone

python train_net.py --eval_only --resume --num-gpus $n --config-file configs/semantic_sam_only_sa-1b_swinL.yaml COCO.TEST.BATCH_SIZE_TOTAL=$n  MODEL.WEIGHTS=/path/to/weights

For SwinT backbone

python train_net.py --eval_only --resume --num-gpus $n --config-file configs/semantic_sam_only_sa-1b_swinT.yaml COCO.TEST.BATCH_SIZE_TOTAL=$n  MODEL.WEIGHTS=/path/to/weights

:star: Training

We currently release the code of training on SA-1B only. Complete training with semantics will be released later. $n is the number of gpus you use before running the training code, you need to specify your training data of SA-1B.

export SAM_DATASETS=/pth/to/dataset
export SAM_SUBSET_START=$start
export SAM_SUBSET_END=$end

We convert SA-1B data into 100 tsv files. start(int, 0-99) is the start of your SA-1B data index and end(int, 0-99) is the end of your data index. If you are not using the tsv data formats, you can refer to this json registration for SAM for a reference.

For SwinL backbone

python train_net.py --resume --num-gpus $n  --config-file configs/semantic_sam_only_sa-1b_swinL.yaml COCO.TEST.BATCH_SIZE_TOTAL=$n  SAM.TEST.BATCH_SIZE_TOTAL=$n  SAM.TRAIN.BATCH_SIZE_TOTAL=$n

For SwinT backbone

python train_net.py --resume --num-gpus $n  --config-file configs/semantic_sam_only_sa-1b_swinT.yaml COCO.TEST.BATCH_SIZE_TOTAL=$n  SAM.TEST.BATCH_SIZE_TOTAL=$n  SAM.TRAIN.BATCH_SIZE_TOTAL=$n
**We also support training to reproduce SAM**
```shell
python train_net.py --resume --num-gpus $n  --config-file configs/semantic_sam_reproduce_sam_swinL.yaml COCO.TEST.BATCH_SIZE_TOTAL=$n  SAM.TEST.BATCH_SIZE_TOTAL=$n  SAM.TRAIN.BATCH_SIZE_TOTAL=$n

…

bibtex
@article{li2023semantic,
  title={Semantic-SAM: Segment and Recognize Anything at Any Granularity},
  author={Li, Feng and Zhang, Hao and Sun, Peize and Zou, Xueyan and Liu, Shilong and Yang, Jianwei and Li, Chunyuan and Zhang, Lei and Gao, Jianfeng},
  journal={arXiv preprint arXiv:2307.04767},
  year={2023}
}
}