PiCIE: Unsupervised Semantic Segmentation using Invariance and Equivariance in Clustering

Jang Hyun Cho¹, Utkarsh Mall², Kavita Bala², Bharath Hariharan²

¹ The University of Texas at Austin, ² Cornell University

[paper] [supp] [project page]

This repository is the official implementation of PiCIE: Unsupervised Semantic Segmentation using Invariance and Equivariance in Clustering, CVPR 2021.

Contact: Jang Hyun Cho [email protected].

Please feel free to reach out for any questions or discussions!

News

[12/10/2021] Fixed visualization notebook.

[12/10/2021] Training and evaluation code for cityscapes dataset for PiCIE and MDC with trained weights and logs.

[07/31/2021] Minor update on the training dataset module (please see this issue).

[07/14/2021] Added "fine_to_coarse_dict.pickle" file.

[07/06/2021] Unstable training bug fixed + weights and logs uploaded.

Setup

Setting up for this project involves installing dependencies and preparing the datasets.

Installing dependencies

To install all the dependencies, please run the following:

conda env create -f env.yml

Preparing Dataset

Please download the trainset and the validset of COCO dataset as well as the annotations and the json files. Place the dataset as following:

/your/dataset/directory/
      └── coco/
            ├── images/
            │     ├── train2017/
            │     │       ├── xxxxxxxxx.jpg
            │     │       └── ...
            │     └── val2017/
            │             ├── xxxxxxxxx.jpg
            │             └── ...
            └── annotations/
                  ├── train2017/
                  │       ├── xxxxxxxxx.png
                  │       └── ...
                  ├── val2017/
                  │       ├── xxxxxxxxx.png
                  │       └── ...
                  ├── COCO_2017_train.json
                  └── COCO_2017_val.json

Then, create a symbolic link as following:

cd PiCIE
ln -s /your/dataset/directory/ datasets 

Similarly, setup a symbolic link for the save directory as following:

ln -s /your/save/directory/ results

Finally, move curated folder and fine_to_coarse_dict.pickle file to datasets/coco/:

mv curated datasets/coco/
mv fine_to_coarse_dict.pickle datasets/coco/

This will setup the dataset that contains the same set of images with IIC (config).

Running PiCIE

Below are training and testing commands to train PiCIE.

Training

Below line will run the training code with default setting in the background.

nohup ./sh_files/train_picie.sh > logs/picie_train.out & 

Below line will run the testing code with default setting in the background.

Testing

nohup ./sh_files/test_picie.sh > logs/picie_test.out &

Pretrained Models

We have pretrained PiCIE weights.

Visualization

We prepared a jupyter notebook for visualization.

Citation

If you find PiCIE useful in your research, please consider citing:

@inproceedings{Cho2021PiCIE,
  title = {PiCIE: Unsupervised Semantic Segmentation using Invariance and Equivariance in Clustering},
  author = {Jang Hyun Cho and  Utkarsh Mall and  Kavita Bala and  Bharath Hariharan},
  year = {2021},
  booktitle = {CVPR}
}

Acknowledgements

We thank Facebook AI Research for the open-soource library Faiss. Also, our implementation largely borrows from DeepCluster and DeeperCluster for clustering with Faiss.

TODO's

• Dependency & dataset setup.
• Clear up and add complete train & test codes.
• Baseline MDC code.
• Weights and logs.
• Make visualization notebook easier to use + better colors.