By Tsung-Wei Ke*, Sangwoo Mo* and Stella X. Yu
Official implementation of "Learning Hierarchical Image Segmentation For Recognition and By Recognition". This work is selected as a spotlight paper in ICLR 2024.
Image segmentation and recognition occur simultaneously, with recognition relying on the underlying segmentation to form a continuous visual grouping hierarchy. For example, the same object can be parsed into different part-to-whole structures, resulting in varying recognitions. Despite this, most prior works treated segmentation and recognition as separate tasks. In this paper, we aim to devise a learning framework that involves segmentation in the recognition process, utilizing hierarchical segmentation for recognition, which is learned by recognition. Specifically, we propose CAST, which realizes this concept through designs inspired by vision transformers, enabling concurrent segmentation and recognition with a single model. The core idea of CAST is to employ adaptive segment tokens that group the finest pixels into coarser segments, using the latest embedding to represent the entire image for recognition. Trained solely on image recognition objectives, CAST automatically discovers the hierarchy of segments. Our experiments demonstrate that CAST provides consistent hierarchical segmentation and recognition, which is impossible with state-of-the-art segmentation methods such as SAM. Additionally, CAST offers several advantages over the standard ViT, including improved semantic segmentation, computational efficiency, and object-centric attention.
Create a conda environment with the following command:
# initiate conda env
> conda update conda
> conda env create -f environment.yaml
> conda activate cast
# install dgl (https://www.dgl.ai/pages/start.html)
> pip install dgl==1.1.3+cu116 -f https://data.dgl.ai/wheels/cu116/dgl-1.1.3%2Bcu116-cp38-cp38-manylinux1_x86_64.whl
To facilitate fast development on top of our model, we provide here an overview of our implementation of CAST.
The model can be indenpendently installed and used as a stand-alone package.
> pip install -e .
# import the model
> from cast_models.cast import cast_small
> model = cast_small()
See Preparing datasets.
See Getting started with self-supervised learning of CAST
See Getting started with fully-supervised learning of CAST
See Getting started with fine-tuning of CAST on Pascal Context and ADE20K
See Getting started with unsupervised segementation and fine-tuning of CAST on VOC
See Getting started with part segmentation of CAST on PartImageNet
See our jupyter notebook
We host the model weights on hugging face.
- Self-supervised trained models on ImageNet-1K. We use a smaller batch size (256) and shorter training epochs (100) compared to the original setup of MoCo-v3.
| Acc. | pre-trained weights | |
|---|---|---|
| ViT-S | 67.9 | link |
| CAST-S | 68.1 | link |
| CAST-SD | 69.1 | link |
| ViT-B | 72.5 | link |
| CAST-B | 72.4 | link |
- Self-supervised trained models on COCO. We evaluate sementic segmentation before and after fine-tuning on Pascal VOC.
| CAST-S | mIoU | pre-trained weights |
|---|---|---|
| before fine-tuning | 38.4 | link |
| after fine-tuning | 67.6 | link |
We pre-train CAST for 400 epochs on COCO. Segmentation before fine-tuning on VOC uses the checkpoint at the 149th epoch. Segmentation after fine-tuning on VOC uses the final checkpoint.
- Supervised trained models on ImageNet-1K. We evaluate semantic segmentaiton by fine-tuning on ADE20K and Pascal Context.
| ImageNet-1K (Acc.) | ADE20K (mIoU) | Pascal Context (mIoU) | pre-trained weights | |
|---|---|---|---|---|
| CAST-S | 80.4 | 43.1 | 49.1 | link |
We also provide a compressed zip file, which include most of the trained model weights, except those fine-tuned on ADE20K and Pascal Context.
If you find this code useful for your research, please consider citing our paper "Learning Hierarchical Image Segmentation For Recognition and By Recognition".
@inproceedings{ke2023cast,
title={Learning Hierarchical Image Segmentation For Recognition and By Recognition},
author={Ke, Tsung-Wei and Mo, Sangwoo and Stella, X Yu},
booktitle={The Twelfth International Conference on Learning Representations},
year={2023}
}
CAST is released under the MIT License (refer to the LICENSE file for details).
This release of code is based on MOCO-v3, DeiT, SegFormer, SPML.
