# MaskedDenoising
## 论文
[Masked Image Training for Generalizable Deep Image Denoising](https://arxiv.org/abs/2303.13132)

## 模型结构

本文对模型修改较小，主要基于SwinIR模型结构增加了input mask 和 attention masks。

<div align=center>
    <img src="./doc/method.jpg"/>
</div>

## 算法原理

传统的去噪模型是靠识别噪声本身来起去噪作用的,而不是真正理解图像内容模型。本文的模型在特征提取之后，会对输入图像进行随机大比例遮盖（input mask），比如遮盖75%~85%的像素，迫使网络学习重构被遮盖的内容，增强对图像本身分布的建模能力。遮挡训练的方法可以使模型学习理解和重构图像的内容，而不仅仅依赖于噪声特征，从而获得更好的泛化能力。

## 环境配置

### Docker（方法一）

-v 路径、docker_name和imageID根据实际情况修改

```image.sourcefind.cn:5000/dcu/admin/base/pytorch:1.13.1-centos7.6-dtk-23.04.1-py38-latest
docker pull image.sourcefind.cn:5000/dcu/admin/base/pytorch:1.13.1-centos7.6-dtk-23.04.1-py38-latest
docker run -it -v /path/your_code_data/:/path/ your_code_data/ --shm-size=32G --privileged=true --device=/dev/kfd --device=/dev/dri/ --group-add video --name docker_name imageID bash

cd /your_code_path/maskeddenoising_pytorch
pip install -r requirement.txt
```

### Dockerfile（方法二）

-v 路径、docker_name和imageID根据实际情况修改

```
cd ./docker
cp ../requirement.txt requirement.txt
docker build --no-cache -t maskeddenoising:latest .
docker run -it -v /path/your_code_data/:/path/your_code_data/ --shm-size=32G --privileged=true --device=/dev/kfd --device=/dev/dri/ --group-add video --name docker_name imageID bash

cd /your_code_path/maskeddenoising_pytorch
pip install -r requirement.txt
```

### Anaconda（方法三）

1、关于本项目DCU显卡所需的特殊深度学习库可从光合开发者社区下载安装： https://developer.hpccube.com/tool/

```
DTK软件栈：dtk23.04.1
python：python3.8
torch：1.13.1
torchvision：0.14.1
```

Tips：以上dtk软件栈、python、torch等DCU相关工具版本需要严格一一对应

2、其他非特殊库直接按照requirement.txt安装

```
pip install -r requirement.txt
```

## 数据集

将待训练数据Train400/DIV2K/Flickr2K放入 trainset 文件夹中

Train400：https://github.com/cszn/DnCNN/tree/master/TrainingCodes/DnCNN_TrainingCodes_v1.0/data

DIV2K官方地址：https://data.vision.ee.ethz.ch/cvl/DIV2K/

[Train Data (HR images)](http://data.vision.ee.ethz.ch/cvl/DIV2K/DIV2K_train_HR.zip)

如果官方地址无法下载，可以选择AiStudio公开数据集里的DIV2K下载DIV2K_train_HR.zip：https://aistudio.baidu.com/aistudio/datasetdetail/104667

Flickr2K：https://cv.snu.ac.kr/research/EDSR/Flickr2K.tar

数据整理完成后，保存在trainset文件夹下，执行gen_data.py获得预处理图像文件夹 trainsets/trainH。

```
gen_data.py
```

测试数据集(已在项目中预置)：BSD68、CBSD68、Kodak24、McMaster
https://github.com/cszn/FFDNet/tree/master/testsets

数据集的目录结构如下：

```
├── trainset
│   ├── DIV2K
│   ├── Flickr2K
│   └── Train400
├── trainsets
│   └── trainH
├── testsets
│   ├── BSD68
│   ├── CBSD68
│   ├── Kodak24
│   └── McMaster
```

## 训练

修改配置文件 options/masked_denoising/input_mask_80_90.json 中参数为实际训练数据，主要参数如下：

"gpu_ids": [0,1,2,3] 训练的卡号

"dataroot_H": "trainsets/trainH" 数据地址

input mask: 设置 "if_mask"和"mask1", "mask2"(line 32-34), 制作比例将在mask1和mask2之间随机采样。

attention mask: 设置 "use_mask" 和 "mask_ratio1", "mask_ratio2" (line 68-70). attention mask ratio 可以是一个范围或者一个固定值。

### 单机多卡
# 普通训练
```
bash train.sh
```

# 分布式训练
```
bash train_multi.sh
```

## 推理
如需使用自己的模型，请修改：

--model_path 训练模型地址

--opt 训练模型对应的json文件

--name 结果保存路径results/{name}

#### 单卡推理

```
bash test.sh
```

## result

本地测试集测试结果单张展示：

<div align=center>
    <img src="./doc/origin.png"/>
</div>

<div align=center>
    <img src="./doc/results.png"/>
</div>

### 精度

基于项目提供的测试数据，得到单卡测试结果如下：

|       | PSNR | SSIM | LPIPS  |
| :------: | :------: | :------: | :------: |
| ours | 29.04 | 0.7615 | 0.1294  |
| paper | 30.13 | 0/7981 | 0.1031 |


## 应用场景
### 算法类别
图像降噪

### 热点应用行业
交通,公安,制造

## 源码仓库及问题反馈
http://developer.hpccube.com/codes/modelzoo/maskeddenoising_pytorch.git

## 参考资料
https://github.com/haoyuc/MaskedDenoising.git