Here you find the code for our paper "Diffusion Models for Earth Observation Use-cases: from cloud removal to urban change detection", accepted as oral at Big Data From Space 2023 (BIDS 2023)

download the following checkpoint in "results" folder and rename it "clouds_best.pt".
You can find a demo at the following notebook EO_Diffusion.ipynb

Conda environment:

• Conda 23.1.0
• CUDA toolkit: 11.7.1
• Pytorch: 11.3.0
• Torchvision + torchaudio: 0.14.0 + 0.13.0
• Tested on an NVIDIA RTX 4000 (49 GB)

GPU utilities installation: I don't recomment using the exported eo_diffusion.yml file. It's better to install it directly from pytorch website with the required versions, as shown in the command below:

conda install pytorch torchvision torchaudio pytorch-cuda=11.7 -c pytorch -c nvidia
conda create -n env_name
conda activate env_name
conda install pip
pip install -r requirements.txt

python train.py

Specify batch size, diffusion steps training hyperparameters from command line (you have default values otherwise)

python train.py --batch_size 4 --timesteps 1000 --lr 1e-05 --epochs 200

Based on the same concept expressed in RePaint https://arxiv.org/abs/2201.09865

python train.py --cond_type "sum"

Save your results

python train.py --dir path/to/your/dir --ckpt ckpt_name

--save option if you want to store images

python inference.py --ckpt path/to/your/ckpt --outdir path/to/your/folder_samples --save

Below you find the two relevant lines to modify concerning U-Net architecture and data loaders, in train.py and inference.py

base_dim, dim_mults, attention_resolutions,num_res_blocks, num_heads=128,[1,2,3,4],[4,8],2,8
train_dataloader,test_dataloader=create_cloud_dataloaders(batch_size=args.batch_size, num_workers=4, size=image_size,
                ratio=0.5, length=-1, num_patches=2000, percents=[99,0,70])

In data.py you find all the available dataloaders with the title create_{dataset_name}_dataloader.
In data_load.py you find all the Dataset classes for the available datasets.

Concerning U-Net and Diffusion, you can modify the parameters at this two lines:

unet = UNetModel(image_size, in_channels=in_channels+cond_channels, model_channels=base_dim, out_channels=out_channels, channel_mult=dim_mults, 
                attention_resolutions=attention_resolutions,num_res_blocks=num_res_blocks, num_heads=num_heads, num_classes=num_classes)
model=EODiffusion(unet,
            timesteps=args.timesteps,
            image_size=image_size,
            in_channels=in_channels
            ).to(device)

Lilian Weng blog on Diffusion Models: https://lilianweng.github.io/posts/2021-07-11-diffusion-models/

The Denoising Diffusion Probabilistic Models paper: https://arxiv.org/pdf/2006.11239.pdf

RePaint paper: https://arxiv.org/abs/2201.09865

This work is the result of a collaboration between ESA, Φ-Lab and Sapienza University of Rome, Alcor Lab for my master thesis (manuscript pdf).