We present U2Seg, a unified framework for Unsupervised Universal image Segmentation that consistently outperforms previous state-of-the-art methods designed for individual tasks: CutLER for unsupervised instance segmentation, STEGO for unsupervised semantic segmentation, and the naive combination of CutLER and STEGO for unsupervised panoptic segmentation.
Unsupervised Universal Image Segmentation
Dantong Niu*, Xudong Wang*, Xinyang Han*, Long Lian, Roei Herzig, Trevor Darrell
Berkeley AI Research, UC Berkeley
Preprint
project page | arxiv | bibtex
See installation instructions.
See Preparing Datasets for U2Seg.
U2Seg has 4 stages: 1) generating pseudo instance masks with MaskCut and clustering, 2) generating pseudo semantic masks with STEGO and 3) merging instance masks, semantic masks to get pseudo labels for panoptic segmentation and 4) learning unsupervised universal segmentor from pseudo-masks of unlabeled data.
This part includes MaskCut+Clustering, which we use to generate the pseudo for training of U2Seg, additional information will come later. For implementers who wants to play with out models, we provide well-processed annotations in Data Preparation.
After you prepare the dataset following the above instructions, you should be able to train the U2Seg universal segmentation model by:
python ./tools/train_net.py --config-file ./configs/COCO-PanopticSegmentation/u2seg_R50_800.yaml
Note: you need to download the pre-trained dino backbone and change the path of the corresponding yaml file.
To train U2Seg model with different clustering number (e.g. 300), you can use configs/COCO-PanopticSegmentation/u2seg_R50_300.yaml config file and set the environment variable by export CLUSTER_NUM=300. (This variable would be used in detectron2/data/datasets/builtin.py and detectron2/data/datasets/builtin_meta.py)
We provide models trained with different cluster numbers and training sets. Each cell in the table below contains a link to the corresponding model checkpoint. Place the downloaded ckpts under ckpts folder.
| Cluster Num | ImageNet | COCO | ImageNet + COCO |
|---|---|---|---|
| 800 | CheckPoint | CheckPoint | CheckPoint |
| 300 | CheckPoint | CheckPoint | CheckPoint |
To run inference on images, you should first assign a checkpoint in the u2seg_eval.yaml, then:
python ./demo/u2seg_demo.py --config-file configs/COCO-PanopticSegmentation/u2seg_eval_800.yaml --input demo/images/*jpg --output results/demo_800
To test model trained with different clustering number (e.g. 300), you can use config file like this configs/COCO-PanopticSegmentation/u2seg_R50_300.yaml.
We give a few demo images in demo/images and the corresponding visualizations of the panoptic segmentation with U2Seg:
Coming soon.
Coming soon.
U2Seg is licensed under the Apache, however portions of the project are available under separate license terms: CutLER, Detectron2 and DINO are licensed under the CC-BY-NC license; TokenCut, Bilateral Solver and CRF are licensed under the MIT license; If you later add other third party code, please keep this license info updated, and please let us know if that component is licensed under something other than CC-BY-NC, MIT, or CC0.
U2Seg's wide range of detection capabilities may introduce similar challenges to many other visual recognition methods. As the image can contain arbitrary instances, it may impact the model output.
If you have any general questions, feel free to email us at Dantong Niu, Xudong Wang, Xinyang Han. If you have code or implementation-related questions, please feel free to send emails to us or open an issue in this codebase (We recommend that you open an issue in this codebase, because your questions may help others).
If you find our work inspiring or use our codebase in your research, please consider giving a star ⭐ and a citation.
@misc{niu2023unsupervised,
title={Unsupervised Universal Image Segmentation},
author={Dantong Niu and Xudong Wang and Xinyang Han and Long Lian and Roei Herzig and Trevor Darrell},
year={2023},
eprint={2312.17243},
archivePrefix={arXiv},
primaryClass={cs.CV}
}