README.md

# 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%的像素，迫使网络学习重构被遮盖的内容，增强对图像本身分布的建模能力。遮挡训练的方法可以使模型学习理解和重构图像的内容，而不仅仅依赖于噪声特征，从而获得更好的泛化能力。

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

## 环境配置

### Docker（方法一）
-v 路径、docker_name和imageID根据实际情况修改

```image.sourcefind.cn:5000/dcu/admin/base/pytorch:2.1.0-ubuntu20.04-dtk24.04.1-py3.10
docker pull image.sourcefind.cn:5000/dcu/admin/base/pytorch:2.1.0-ubuntu20.04-dtk24.04.1-py3.10
docker run -it -v /path/your_code_data/:/path/ your_code_data/ -v /opt/hyhal/:/opt/hyhal/:ro  --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/  -v /opt/hyhal/:/opt/hyhal/:ro --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软件栈：dtk24.04.1
python：python3.10
torch：2.1.0
torchvision：0.16.0
```

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

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

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

## 数据集
数据集下载地址：[SCNet AIDatasets](http://113.200.138.88:18080/aidatasets)，将待训练数据`Train400/DIV2K/Flickr2K`放入`trainset`文件夹中

[Train400](http://113.200.138.88:18080/aidatasets/project-dependency/dncnn/-/tree/master/TrainingCodes/DnCNN_TrainingCodes_v1.0/data/Train400?ref_type=heads)

[Train Data (HR images)](http://113.200.138.88:18080/aidatasets/project-dependency/div2k/-/blob/master/DIV2K_train_HR.zip)

[Flickr2K](http://113.200.138.88:18080/aidatasets/project-dependency/flickr2k/-/blob/master/Flickr2K.tar)

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

```
python gen_data.py
```

测试数据集(已在项目中预置)：BSD68、CBSD68、Kodak24、McMaster
- (http://113.200.138.88:18080/aidatasets/project-dependency/ffdnet/-/tree/master/testsets?ref_type=heads)

数据集的目录结构如下：

```
├── 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
bash train.sh
```

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

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

`--model_path` 训练模型地址

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

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

#### 单卡推理
```bash
bash test.sh
```

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

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

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

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

| DEVICE | PSNR | SSIM | LPIPS |
| :------: | :------: | :------: | :------: |
| Z100L | 29.04 | 0.7615 | 0.1294 |
| V100S | 30.13 | 0.7981 | 0.1031 |

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

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

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

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