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.
- 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
-
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
└── ...
To train a model use train.py
, for example:
python train.py --dataset vaihingen --scale_procedure loop --n_scales 3 --n_timesteps 25
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
@InProceedings{Kolbeinsson_2024_WACV,
author = {Kolbeinsson, Benedikt and Mikolajczyk, Krystian},
title = {Multi-Class Segmentation From Aerial Views Using Recursive Noise Diffusion},
booktitle = {Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision (WACV)},
month = {January},
year = {2024},
pages = {8439-8449}
}