Masked Image Training for Generalizable Deep Image Denoising

Haoyu Chen, Jinjin Gu, Yihao Liu, Salma Abdel Magid, Chao Dong, Qiong Wang, Hanspeter Pfister, Lei Zhu

Abstract: When capturing and storing images, devices inevitably introduce noise. Reducing this noise is a critical task called image denoising. Deep learning has become the de facto method for image denoising, especially with the emergence of Transformer-based models that have achieved notable state-of-the-art results on various image tasks. However, deep learning-based methods often suffer from a lack of generalization ability. For example, deep models trained on Gaussian noise may perform poorly when tested on other noise distributions. To address this issue, We present a novel approach to enhance the generalization performance of denoising networks, known as masked training. Our method involves masking random pixels of the input image and reconstructing the missing information during training. We also mask out the features in the self-attention layers to avoid the impact of training-testing inconsistency. Our approach exhibits better generalization ability than other deep learning models and is directly applicable to real- world scenarios. Additionally, our interpretability analysis demonstrates the superiority of our method.

Masked Image Training

The transformer architecture of our proposed masked image training. We only make the minimal change to the original SwinIR architecture – the input mask operation and the attention masks.

Clone repo

git clone https://github.com/haoyuc/MaskedDenoising.git

pip install -r requirement.txt

This code is built on KAIR. We thank the authors for sharing their codes. For more detailed code information, please refer to KAIR.

---

The core code we used:

Input mask:

MaskedDenoising/data/dataset_masked_denoising.py

Line 60 in b4cb8ed

img_L, img_H = utils_mask.input_mask_with_noise(img_H,

Attention mask:

MaskedDenoising/models/network_swinir.py

Line 273 in b4cb8ed

def mask_image(self, image, x_size=None):

Training

❗️ This project is more oriented towards exploring how to make the model have better generalization capabilities when the degraded distribution of the training data is limited.
We train the model only on Gaussian noise (sigma=15), and test the model on other noise types and noise levels. So this project may not be suitable for certain practical applications.

You should modify the json file from options first, for example,

• setting "gpu_ids": [0,1,2,3] if 4 GPUs are used,
• setting "dataroot_H": "trainsets/trainH" if path of the high quality dataset is trainsets/trainH, more images are better.
• input mask: setting "if_mask" and "mask1", "mask2"(line 32-34), the making ratio will randomly sample between mask1 and mask2.
• attention mask: setting "use_mask" and "mask_ratio1", "mask_ratio2" (line 68-70). The attention mask ratio can be a range or a fixed value.

---

• Training with DataParallel - PSNR

python main_train_psnr.py --opt options/masked_denoising/input_mask_80_90.json

• Training with DistributedDataParallel - PSNR - 4 GPUs

python -m torch.distributed.launch --nproc_per_node=4 --master_port=1234 main_train_psnr.py --opt options/masked_denoising/input_mask_80_90.json  --dist True

• Kill distributed training processes of main_train_psnr.py

kill $(ps aux | grep main_train_psnr.py | grep -v grep | awk '{print $2}')

Testing

The pre-trained models have already been included in this repository (model_zoo/input_mask_80_90.pth).
❗️ Please note that this model was trained solely on Gaussian noise (sigma=15).

python main_test_swinir.py \
        --model_path model_zoo/input_mask_80_90.pth  \
        --name input_mask_80_90/McM_poisson_20  \
        --opt model_zoo/input_mask_80_90.json \
        --folder_gt testset/McM/HR  \
        --folder_lq testset/McM/McM_poisson_20

trainsets

You can use more images for better porformance.

• DIV2K
• Flickr2K
• BSD500
• WED

testsets

• cbsd68
• kodak24
• urban100
• McMaster

References

@InProceedings{Chen_2023_CVPR,
    author    = {Chen, Haoyu and Gu, Jinjin and Liu, Yihao and Magid, Salma Abdel and Dong, Chao and Wang, Qiong and Pfister, Hanspeter and Zhu, Lei},
    title     = {Masked Image Training for Generalizable Deep Image Denoising},
    booktitle = {Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
    month     = {June},
    year      = {2023},
    pages     = {1692-1703}
}