STEGO: Unsupervised Semantic Segmentation by Distilling Feature Correspondences

Pytorch implementation of the STEGO unsupervised semantic segmentation system. [Paper] [Video]

Install

Clone this repository:

git clone https://github.com/mhamilton723/STEGO.git
cd STEGO

Install Conda Environment

Please visit the Anaconda install page if you do not already have conda installed

conda env create -f environment.yml
conda activate stego

Download Pre-Trained Models

cd src
python download_models.py

Download Datasets

First, change the pytorch_data_dir variable to your systems pytorch data directory where datasets are stored.

python download_data.py
cd /YOUR/PYTORCH/DATA/DIR
unzip cocostuff.zip
unzip cityscapes.zip
unzip potsdam.zip
unzip potsdamraw.zip

Evaluation

From inside STEGO/src please run the following:

python eval_segmentation.py

Training

From inside STEGO/src please run the following:

python train_segmentation.py

Citation

@article{hamilton2022unsupervised,
  title={Unsupervised Semantic Segmentation by Distilling Feature Correspondences},
  author={Hamilton, Mark and Zhang, Zhoutong and Hariharan, Bharath and Snavely, Noah and Freeman, William T},
  journal={arXiv preprint arXiv:2203.08414},
  year={2022}
}