Effective Cloud Removal for Remote Sensing Images by an Improved Mean-Reverting Denoising Model with Elucidated Design Space
This is the official repository for EMRDM, which is a novel cloud removal methods improved from mean-reverting diffusion models, with a modular framework, a well-elucidated design space, and enhanced modules, demonstrating superior capabilities in both mono-temporal and multi-temporal cloud removal tasks.
- 🎯We have released the implementation of EMRDM.
- 🎯We have made the weights of EMRDM trained on CUHK-CR1, CUHK-CR2, SEN12MS-CR, and Sen2_MTC_New datasets public.
- 🎯We have provided a readme file in Chinese.
- 🎯We have added a bibtex example for citing our work.
We provide the requirements.txt in the root path. All the packages in the virtual environment we used are detailed in this file. But it is not recommended to directly use pip install -r requirements.txt as there are complex dependencies among these packages. We suggest to download some large and pivotal packages, like torch, flash_attn, natten, pytorch-lightning first.
Specifically, it is highly suggested to first create a virtual enviroment using conda, as follows:
conda create --name emrdm python=3.10
conda activate emrdmThen, you can download numpy and torch as follows:
pip install torch==2.2.1 torchaudio==2.2.1 torchvision==0.17.1 numpy==1.26.4The version of CUDA we used is CUDA 12.1.
Next, you can download flash_attn as follows:
MAX_JOBS=4 pip install flash_attn==2.5.9.post1 --no-build-isolationIt may take a lot of time to build and compile the flash_attn package. For accelerating, you can also download the wheel from the official repository of flash_attn and use the following instruction to install from the wheel:
pip install [name_of_flash_attn_wheel].whlAfter that, you can download the neighborhood attention package, i.e., natten, as follows:
pip install natten==0.17.1This may also take a lot of time for compilation. Thus, you can also use the wheel from the offficial repository of natten for fast downloading. Or you can directly use the following instrcution to download natten from its wheel:
pip install natten==0.17.1+torch220cu121 -f https://shi-labs.com/natten/wheelsFor downloading pytorch-lightning, you can use the following instruction:
pip install pytorch-lightning==2.3.0After downloading all the aforementioned pivotal packages, you can use the following instructions to download other packages:
pip install wandb==0.17.8 matplotlib==3.9.2 natsort==8.4.0 \
omegaconf==2.3.0 scipy==1.14.0 dctorch==0.1.2 rasterio==1.3.11
pip install pandas==2.2.3 opencv-python==4.10.0.84 lpips==0.1.4
pip install tifffile==2024.7.24 s2cloudless==1.7.2 \
albumentations==1.4.10 albucore==0.0.12(Optional) If you still encounter the package missing problem, you can refer to the requirements.txt file to download the packages you need.
If you meet other enviroment setting problems we have not mentioned in this readme file, please contack us to report you problems or throw a issue.
We use four datasets: CUHK-CR1, CUHK-CR2, SEN12MS-CR and Sen2_MTC_New. You need to download these datasets first.
We provide the downloading URLs of these datasets as follows:
| Dataset | Type | URL |
|---|---|---|
| CUHK-CR1 | Mono-Temporal | https://github.com/littlebeen/DDPM-Enhancement-for-Cloud-Removal |
| CUHK-CR2 | Mono-Temporal | https://github.com/littlebeen/DDPM-Enhancement-for-Cloud-Removal |
| SEN12MS-CR | Mono-Temporal | https://patricktum.github.io/cloud_removal/sen12mscr/ |
| Sen2_MTC_New | Multi-Temporal | https://github.com/come880412/CTGAN |
For fast starting, you can only download the testing dataset and run the testing instructions given below.
We provide our configuration files, i.e., *.yaml, in the ./configs/example_training/ folder. The code automatically reads the yaml file and sets the configuration. You can change the settings, such as the data path, batch size, number of workers, in the data part of each yaml file to adapt to your expectations. Read the yaml file in ./configs/example_training/ for more details.
We have also included the yaml files for our ablation experiments in the ./configs/example_training/ablation/ directory.
You can use the following instruction in the root path of this repository to run the training process:
python main.py --base configs/example_training/[yaml_file_name].yaml --enable_tf32Here, cuhk.yaml is for training on the CUHK-CR1 dataset, cuhkv2.yaml is for training on the CUHK-CR2 dataset, sen2_mtc_new.yaml is for training on the Sen2_MTC_New dataset, and sentinel.yaml is for training on the SEN12MS-CR dataset. Note that you should modify the data.params.train part in the yaml file according to your dataset path.
You can also use the -l parameter to change the save path of logs, with ./logs as the default path:
python main.py --base configs/example_training/[yaml_file_name].yaml --enable_tf32 -l [path_to_your_logs]If you want to resume from a previous training checkpoint, you can use the follow instruction:
python main.py --base configs/example_training/[yaml_file_name].yaml --enable_tf32 -r [path_to_your_ckpt]If you want to initiate the model from an existing checkpoint and restart the training process, you should modify the value of model.ckpt_path in your yaml file to the path of your checkpoint.
Run the following instruction for testing:
python main.py --base configs/example_training/[yaml_file_name].yaml --enable_tf32 -t falseThe [yaml_file_name].yaml files are the same as those in training process. Note that
- You should set the
data.params.testpart, otherwise the test dataloader will not be implemented. - You should modify he value of
model.ckpt_pathin youryamlfile to the path of your checkpoint. We have provided checkpoints of EMRDM trained on the four datasets used by us in the Sec. Models.
The predicting process will output all cloud removed images. This process only support using one GPU (by setting lightning.trainer.devices to only one device). You can run predicting process using:
python main.py --base configs/example_training/[yaml_file_name].yaml --enable_tf32 -t false --no-test true --predict trueThe [yaml_file_name].yaml files are the same as those in the testing process. Note that you should set the data.params.predict part and the model.ckpt_path part (the same way as testing), otherwise you will not obtain the correct results.
We have released all the training results, including the training configuration files, the weights of trained models, and the log files of the training process.
We have also released all the testing results, including the testing configuration files and the log files of the testing process.
All the files can be downloaded via the following URLs:
In the shared files, the train folder contains the training results, while the test folder contains the testing results. The sharing code for Baidu Drive is 6161.
The code is based on the official implementations of the generative-models, k-diffusion, utae-paps and other repositories, as follows:
| Repository | URL | Relationship to this Repository |
|---|---|---|
generative-models |
https://github.com/Stability-AI/generative-models | The main body of our code is based on this repository. |
k-diffusion |
https://github.com/crowsonkb/k-diffusion | The implementation of our denoising network is based on this repository. |
utae-paps |
https://github.com/VSainteuf/utae-paps | Our temporal fusion attention mechanism is based on this repository. |
UnCRtainTS |
https://github.com/PatrickTUM/UnCRtainTS | Our dataloader of the SEN12MS-CR dataset is based on this repository. |
SEN12MS-CR-TS |
https://github.com/PatrickTUM/SEN12MS-CR-TS | Our dataloader of the SEN12MS-CR dataset is based on this repository. |
CTGAN |
https://github.com/come880412/CTGAN | Our dataloader of the Sen2_MTC_New dataset is based on this repository. |
DDPM-Enhancement-for-Cloud-Removal |
https://github.com/littlebeen/DDPM-Enhancement-for-Cloud-Removal | Our dataloader of the CUHK-CR1 and CUHK-CR2 datasets are based on this repository. |
We express our sincere gratitude for their significant contributions to the open-source community.
If you have encountered any problems, feel free to contact me via my email liuyi2052697@foxmail.com (preferred) or liuyi61@tongji.edu.cn.
@inproceedings{liu2025effective,
title={Effective cloud removal for remote sensing images by an improved mean-reverting denoising model with elucidated design space},
author={Liu, Yi and Li, Wengen and Guan, Jihong and Zhou, Shuigeng and Zhang, Yichao},
booktitle={Proceedings of the Computer Vision and Pattern Recognition Conference},
pages={17851--17861},
year={2025}
}