PiCIE: Unsupervised Semantic Segmentation using Invariance and Equivariance in Clustering
Jang Hyun Cho1, Utkarsh Mall2, Kavita Bala2, Bharath Hariharan2
1 The University of Texas at Austin, 2 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 janghyuncho7@utexas.edu.
Please feel free to reach out for any questions or discussions!
News
[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.
Method | Dataset | Pre-trained weight | Train log |
---|---|---|---|
PiCIE | COCO | weight | log |
MDC | COCO | weight | log |
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.