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!

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. Place the dataset as following:

/your/dataset/directory/
      └── coco/
            ├── images/
            │     ├── train2017/
            │     │       ├── xxxxxxxxx.jpg
            │     │       └── ...
            │     └── val2017/
            │             ├── xxxxxxxxx.jpg
            │             └── ...
            └── annotations/
                  ├── 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 to datasets/coco/:

mv curated datasets/coco/

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

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 (To be updated soon)

We have pretrained PiCIE weights.

Visualization (To be updated soon)

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.