This is the official GitHub project repository of the research paper Cross-stitched Multi-task Unified Dual Resursive Networks for Single Image Deraining and Desnowing authored by Karavarsamis Sotiris, Doumanoglou Alexandros, Konstantoudakis Konstantinos and Zarpalas Dimitrios.
This code and data repository has been made publicly available for researchers and practitioners to evaluate the CMUDRN model, or use it in their own work. This repository also makes available pre-trained network weights for reproducibility or for demo purposes.
Our paper is available at https://arxiv.org/abs/2211.08290.
Please take the steps to understand README.md before you run experiments. Our codebase is a little bit complicated, and it is worth reading the instructions we put forward here.
The codebase presented in this repository is based on the implementation of the Dual Recursive Network (DRN) model (available at https://github.com/csdwren/DRN).
In the CMUDRN model, we extended the equations of the DRN model in order to design an image restoration model that can unify image denoising when the noise source is environmental rain or snow. In unified image restoration, we typically do not know which type of noise the input is contaminated by. A unified image restoration algorithm has to find out the type of noise on its own. Our paper documents these equation derivations.
To download this repository on your system, run the command
$ git clone https://github.com/brigr/CMUDRN
To run CMUDRN (i.e., train a model and then do inference, or to make inferences for inputs images using a pre-trained model), you need to install the following Python libraries on your system (we here ignore library versions):
torch
torchvision
tqdm
visdom
numpy
h5py
opencv-python
scikit-image
Note: Since in 2022 we recombined serveral versions of the abovementioned libraries to develop this repository, we follow this version recommendation policy: we supply a requirements_2022.txt file which defines library versions up to the year 2022. However, at the time we released this repository, we tested the code with the latest libraries released in the year 2026. Therefore, we also supply a requirements_2026.txt file which defines the newest library versions that we used to prepare and test our code.
Python versions: To install library versions from requirements_2022.txt or from requirements_2026.txt, we used Python 3.10. We recommend that you follow this convention to prevent that something breaks in the installations or during run-time.
We assume you first cloned the CMUDRN repisitory as indicated above. Then, to create a Python virtual environment named myenv, change to a directory of your choice (for example, change to /home/myuser) and then run
$ python3.10 -m venv myenv
Then, enter the Python environment using the command
$ source myenv/bin/activate
After running this command line, you should see a command prompt that looks like
(myenv) $
To install the library dependencies within a Python environment, first decide if you want to use the 2022 or the 2026 suggested versions.
If you want to set up those from the year 2022, issue the command
(myenv) $ python3.10 -m pip install -r requirements_2022.txt
If instead you need to use the latest library versions (of the year 2026), issue instead the command
(myenv) $ python3.10 -m pip install -r requirements_2026.txt
This section contains information on getting the augmented CSD dataset, preparing it on your machine before being able to train a model using the dataset.
To train and/or evaluate a CMUDRN model, you need examples of rainy and snowy image pairs. You can create your own dataset if you want, but you can of course use the evaluation dataset that we created and used in our publication.
We call our dataset the augmented CSD dataset, since this is based on the CSD dataset introduced in the ICCV 2021 paper of Chen, et al. titled "ALL Snow Removed: Single Image Desnowing Algorithm Using Hierarchical Dual-Tree Complex Wavelet Representation and Contradict Channel Loss". The CSD dataset contains 8,000 image pairs of synthetically generated snowy images and their associated clean images.
We decided to use the existing snowy images in the CSD dataset and generate rainy images given the 8,000 clean images in the CSD dataset. We utilized the rain simulating algorithm by Garg and Nayar (see our paper for more information).
If you want to access the augmented CSD dataset, please use this link.
As prompted in the previous section, download the augmented CSD dataset using this link. You will download the file augmented_CSD.zip. Move this .zip file into your local clone of this repository. Uncompress the .zip file, creating the local folder augmented_CSD. Rename the augmented_CSD folder to train.
Under the test/ folder that exists in your repository clone, create the folders rain/ and snow/. For example, you can open up a shell at the root folder of the repository and run the following two commands
$ mkdir -p test/rain
$ mkdir -p test/snow
From the augmented_CSD/rain folder, copy the files 5601.tiff up to 8000.tiff and paste them into the folder test/rain. From the augmented_CSD/snow folder, copy the files 5601.tiff up to 8000.tiff and paste them into the folder test/snow.
To evaluate CMUDRN using the 8,000 image triplets available in the augmented CSD dataset, we keep the first 70% of the samples (amounting to 5600 samples) in the dataset for model training, and keep the rest 30% of the total samples (amounting to 2400 samples) for model evaluation. If you need to compare your model with CMUDRN, please use this convention.
Once you get the augmented CSD dataset, it is mandatory that you generate the necessary .h5 files before you proceed to train a new model. Again, please read carefully the following notes so that you have an understanding of the instructions before you attempt to run any command. These files store the augmented CSD dataset in .h5 format. To build these files, change into the folder containing your local clone of this repository, and run the command
(myenv) $ python3.10 train_CSD.py --preprocess
Keep in mind that (myenv) may not appear in your shell if your did not make a Python virtual environment or if you did not activate it.
Note that it may take some time for this command to complete. Moreover, we do not include these .h5 files in this code repository due to them being very large (for your notice, they weigh 88 Gibabytes). You may need at least this free space to generate the .h5 on your own.
To train a CMUDRN model, run the train.py file with the following arguments. train.py takes in general the following arguments
usage: train.py [-h] [--preprocess PREPROCESS] --batchSize BATCHSIZE --epochs EPOCHS [--intra_iter INTRA_ITER] [--inter_iter INTER_ITER] [--a_S A_S] [--enable_ssim_loss] [--enable_fro_loss]
[--enable_headonly_loss] [--use_stitching] [--milestone MILESTONE] --lr LR --save_folder SAVE_FOLDER [--save_freq SAVE_FREQ] --data_path DATA_PATH [--use_visdom]
[--visdom_port VISDOM_PORT] [--use_gpu] [--gpu_id GPU_ID]
CMUDRN training script
options:
-h, --help show this help message and exit
--preprocess PREPROCESS
prepares .h5 dataset files required prior to model training
--batchSize BATCHSIZE
size of the batch size during training
--epochs EPOCHS number of model training epochs
--intra_iter INTRA_ITER
number of intra iterations in model
--inter_iter INTER_ITER
number of inter iteration in model
--a_S A_S coefficient of the cross-stitch units in model
--enable_ssim_loss enable SSIM loss for model training
--enable_fro_loss enable Frobenius loss for model training
--enable_headonly_loss
enforce baseline CMUDRN training with global-only loss
--use_stitching enable stitching (for model ablation)
--milestone MILESTONE
when to decay learning rate; should be less than epochs
--lr LR initial learning rate
--save_folder SAVE_FOLDER
path of log files
--save_freq SAVE_FREQ
model save frequency every epoch count
--data_path DATA_PATH
path to training data
--use_visdom enable Visdom logging
--visdom_port VISDOM_PORT
set Visdom port for expriment monitoring
--use_gpu to use GPU or not
--gpu_id GPU_ID GPU id
To train a CMUDRN model over the augmented CSD data, first decide which type of CMUDRN model you would like to create (see Table 1 in our paper).
Below we give 7 different command lines to run a CMUDRN model of a particular type based on the rows of Table 1 in our paper.
Use of Visdom: Note that in the command lines below, we set --use_visdom and --visdom_port 9999. This enforces the use of visdom for logging experiment metrics, and running a Visdom server on port 9999 (you can change port 9999 to any other porn number).
(myenv) $ python3.10 -m visdom.server -port 9999
(myenv) $ python3.10 train.py --batchSize 28 --epochs 5 --intra_iter 3 --inter_iter 3 --a_S 0.5 --enable_ssim_loss --enable_fro_loss --use_stitching --lr 0.001 --data_path augmented_CSD/rain,augmented_CSD/snow --save_folder logs/mynewmodel_type1 --save_freq 1 --use_visdom --visdom_port 9999 --use_gpu --gpu_id 0
(myenv) $ python3.10 train.py --batchSize 28 --epochs 5 --intra_iter 3 --inter_iter 3 --a_S 0.5 --enable_ssim_loss --enable_fro_loss --lr 0.001 --data_path augmented_CSD/rain,augmented_CSD/snow --save_folder logs/mynewmodel_type2 --save_freq 1 --use_visdom --visdom_port 9999 --use_gpu --gpu_id 0
(myenv) $ python3.10 train.py --batchSize 28 --epochs 5 --intra_iter 3 --inter_iter 3 --a_S 0.5 --enable_fro_loss --use_stitching --lr 0.001 --data_path augmented_CSD/rain,augmented_CSD/snow --save_folder logs/mynewmodel_type3 --save_freq 1 --use_visdom --visdom_port 9999 --use_gpu --gpu_id 0
(myenv) $ python3.10 train.py --batchSize 28 --epochs 5 --intra_iter 3 --inter_iter 3 --a_S 0.5 --enable_fro_loss --lr 0.001 --data_path augmented_CSD/rain,augmented_CSD/snow --save_folder logs/mynewmodel_type4 --save_freq 1 --use_visdom --visdom_port 9999 --use_gpu --gpu_id 0
(myenv) $ python3.10 train.py --batchSize 28 --epochs 5 --intra_iter 3 --inter_iter 3 --a_S 0.5 --enable_ssim_loss --use_stitching --lr 0.001 --data_path augmented_CSD/rain,augmented_CSD/snow --save_folder logs/mynewmodel_type5 --save_freq 1 --use_visdom --visdom_port 9999 --use_gpu --gpu_id 0
(myenv) $ python3.10 train.py --batchSize 28 --epochs 5 --intra_iter 3 --inter_iter 3 --a_S 0.5 --enable_ssim_loss --lr 0.001 --data_path augmented_CSD/rain,augmented_CSD/snow --save_folder logs/mynewmodel_type6 --save_freq 1 --use_visdom --visdom_port 9999 --use_gpu --gpu_id 0
(myenv) $ python3.10 train.py --batchSize 28 --epochs 5 --intra_iter 3 --inter_iter 3 --a_S 0.5 --enable_ssim_loss --enable_fro_loss --enable_headonly_loss --use_stitching --lr -1.001 --data_path augmented_CSD/rain,augmented_CSD/snow --save_folder logs/mynewmodel_type7 --save_freq 1 --use_visdom --visdom_port 9999 --use_gpu --gpu_id 0
The available test.py file is used to make inferences. This script has the following parameters:
usage: test.py [-h] --data_path DATA_PATH --save_path SAVE_PATH [--checkpoint_path CHECKPOINT_PATH] [--use_gpu] [--gpu_id GPU_ID] [--intra_iter INTRA_ITER] [--inter_iter INTER_ITER]
[--image_extension IMAGE_EXTENSION] [--require_chan_reordering]
CMUDRN inference script
options:
-h, --help show this help message and exit
--data_path DATA_PATH
path to test images to do inference on
--save_path SAVE_PATH
path to folder for saving denoised images
--checkpoint_path CHECKPOINT_PATH
--use_gpu use GPU or not
--gpu_id GPU_ID GPU identifier (id) to use for inference
--intra_iter INTRA_ITER
count of intra iterations
--inter_iter INTER_ITER
count of inter iterations
--image_extension IMAGE_EXTENSION
extension of target files (default: .tiff)
--require_chan_reordering
require channel reordering before inference
Discussion: We make some clarifications here. By defining --checkpoint_path and setting a particular path to a .pth file, you also need to set --intra_iter X and --inter_iter Y to the particular values X and Y which were used to generate the checkpoint via train.py. Moreover, you also need to define the extension of the probe files that you need to do inference on via the switch --image_extension .tiff. For instance, if you need to apply the model on PNG files, you need to set the switch --image_extension .png.
Here is an example command line invoking test.py using the pre-trained model under logs/stitching_ssim_fro_inter3_intra3/
(myenv) $ mkdir ./loc
(myenv) $ python3.10 test.py --data_path test/snow --save_path ./loc --checkpoint_path logs/stitching_ssim_fro_inter3_intra3/net_latest.pth --use_gpu --intra_iter 3 --inter_iter 3 --image_extension .tiff --require_chan_reordering
If you have any problem with using this repository, or if would like to report an error in this repository, please contact Alexandros Doumanoglou (aldoum@iti.gr).
If you find our paper and/or code useful in your work, please cite our paper published in WFIOT 2022. You can use the following bibtex entry
@inproceedings{karavarsamis2022cross,
title={Cross-Stitched Multi-task Dual Recursive Networks for Unified Single Image Deraining and Desnowing},
author={Karavarsamis, Sotiris and Doumanoglou, Alexandros and Konstantoudakis, Konstantinos and Zarpalas, Dimitrios},
booktitle={2022 IEEE 8th World Forum on Internet of Things (WF-IoT)},
pages={1--6},
year={2022},
organization={IEEE}
}