Recursive Noise Diffusion

This repo is the official implementation of Multi-Class Segmentation from Aerial Views using Recursive Noise Diffusion

The core idea of Recursive Noise Diffusion is the recursive denoising process, as shown in the figure above. Training with recursive denoising involves progressing through each time step t from T to 1, recursively (as the name suggests), which allows a portion of the predicted error to propagate. This process is initialised with pure noise. The noise function diffuses the previous predicted segmentation, then the model denoises this diffused segmentation given a conditioning RGB image. The denoised predicted segmentation is compared to the ground truth. Notably, the ground truth segmentation is never used as part of the input to the model. This process is agnostic to the choice of noise function, diffusion model and loss.

Getting Started

Setup

Clone this repo:

git clone https://github.com/benediktkol/recursive-noise-diffusion.git
cd recursive-noise-diffusion

Install requirements:

pip install -r requirements.txt

[Optional] Create a conda environment:

conda env create -f rnd.yml
conda activate rnd

Download data
- Vaihingen Buildings
- UAVid
File structure

  data
  ├── UAVid
  │   ├── uavid_test
  │   │   ├── seq21
  │   │   │   └── Images
  │   │   │       ├── 000000.png
  │   │   │       └── ...
  │   │   └── ...
  │   ├── uavid_train
  │   │   ├── seq1
  │   │   │   ├── Images
  │   │   │   │   ├── 000000.png
  │   │   │   │   └── ...
  │   │   │   └── Labels
  │   │   │       ├── 000000.png
  │   │   │       └── ...
  │   │   └── ...
  │   └── uavid_val
  │       ├── seq16
  │       │   ├── Images
  │       │   │   ├── 000000.png
  │       │   │   └── ...
  │       │   └── Labels
  │       │       ├── 000000.png
  │       │       └── ...
  │       └── ...
  └── Vaihingen_buildings
      ├── all_buildings_mask_001.png
      ├── ...
      ├── building_001.png
      ├── ...
      ├── building_gt_001.png
      ├── ...
      ├── building_mask_001.png
      └── ...
 recursive-noise-diffusion (this repo)
 ├── test.py
 ├── train.py
 └── ...

Train

To train a model use train.py, for example:

python train.py --dataset vaihingen --scale_procedure loop --n_scales 3 --n_timesteps 25

Evaluation

To evaluate a model use test.py, for example:

python test.py --load_checkpoint /path/to/checkpoint.pt --dataset vaihingen --scale_procedure loop --n_scales 3 --n_timesteps 25

Cite

@article{kolbeinsson2022multi,
  title={Multi-Class Segmentation from Aerial Views using Recursive Noise Diffusion},
  author={Kolbeinsson, Benedikt and Mikolajczyk, Krystian},
  journal={arXiv preprint arXiv:2212.00787},
  year={2022}
}