README.md 4.56 KB
Newer Older
chenych's avatar
chenych committed
1
2
# MaskedDenoising
## 论文
Rayyyyy's avatar
Rayyyyy committed
3
4
`Masked Image Training for Generalizable Deep Image Denoising`
- https://arxiv.org/abs/2303.13132
chenych's avatar
chenych committed
5

chenych's avatar
chenych committed
6
## 模型结构
Rayyyyy's avatar
Rayyyyy committed
7
本文对模型修改较小,主要基于SwinIR模型结构增加了`input mask``attention masks`
chenych's avatar
chenych committed
8
9
10
11
12
13
14
15

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

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

16
17
18
19
<div align=center>
    <img src="./doc/progress.png"/>
</div>

chenych's avatar
chenych committed
20
21
22
23
24
## 环境配置

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

dcuai's avatar
dcuai committed
25
26
27
```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
chenych's avatar
chenych committed
28
29
30
31
32
33
34
35
36
37
38
39

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 .
dcuai's avatar
dcuai committed
40
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
chenych's avatar
chenych committed
41
42
43
44
45
46

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

### Anaconda(方法三)
chenzk's avatar
chenzk committed
47
1、关于本项目DCU显卡所需的特殊深度学习库可从光合开发者社区下载安装: https://developer.sourcefind.cn/tool/
chenych's avatar
chenych committed
48
49

```
dcuai's avatar
dcuai committed
50
51
52
53
DTK软件栈:dtk24.04.1
python:python3.10
torch:2.1.0
torchvision:0.16.0
chenych's avatar
chenych committed
54
55
56
57
58
59
60
61
62
63
64
```

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

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

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

## 数据集
chenzk's avatar
chenzk committed
65
根据源github文档下载数据,将待训练数据`Train400/DIV2K/Flickr2K`放入`trainset`文件夹中
chenych's avatar
chenych committed
66

chenzk's avatar
chenzk committed
67
[Train400]
chenych's avatar
chenych committed
68

chenzk's avatar
chenzk committed
69
[DIV2K(HR images)](https://data.vision.ee.ethz.ch/cvl/DIV2K/)
chenych's avatar
chenych committed
70

chenzk's avatar
chenzk committed
71
[Flickr2K](https://cv.snu.ac.kr/research/EDSR/Flickr2K.tar)
chenych's avatar
chenych committed
72

Rayyyyy's avatar
Rayyyyy committed
73
数据整理完成后,保存在`trainset`文件夹下,执行`gen_data.py`获得预处理图像文件夹`trainsets/trainH`
chenych's avatar
chenych committed
74
75

```
Rayyyyy's avatar
Rayyyyy committed
76
python gen_data.py
chenych's avatar
chenych committed
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
```

测试数据集(已在项目中预置):BSD68、CBSD68、Kodak24、McMaster

数据集的目录结构如下:

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

## 训练
Rayyyyy's avatar
Rayyyyy committed
98
修改配置文件`options/masked_denoising/input_mask_80_90.json`中参数为实际训练数据,主要参数如下:
chenych's avatar
chenych committed
99
100
101
102
103

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

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

Rayyyyy's avatar
Rayyyyy committed
104
input mask: 设置 "if_mask"和"mask1","mask2"(line 32-34), 制作比例将在mask1和mask2之间随机采样。
chenych's avatar
chenych committed
105

Rayyyyy's avatar
Rayyyyy committed
106
attention mask: 设置 "use_mask" 和 "mask_ratio1","mask_ratio2" (line 68-70),`attention mask ratio`可以是一个范围或者一个固定值。
chenych's avatar
chenych committed
107
108

### 单机多卡
109
#### 普通训练
Rayyyyy's avatar
Rayyyyy committed
110
```bash
chenych's avatar
chenych committed
111
112
113
bash train.sh
```

114
#### 分布式训练
Rayyyyy's avatar
Rayyyyy committed
115
```bash
chenych's avatar
chenych committed
116
117
118
119
120
121
bash train_multi.sh
```

## 推理
如需使用自己的模型,请修改:

Rayyyyy's avatar
Rayyyyy committed
122
`--model_path` 训练模型地址
chenych's avatar
chenych committed
123

Rayyyyy's avatar
Rayyyyy committed
124
`--opt` 训练模型对应的json文件
chenych's avatar
chenych committed
125

Rayyyyy's avatar
Rayyyyy committed
126
`--name` 结果保存路径`results/{name}`
chenych's avatar
chenych committed
127
128

#### 单卡推理
Rayyyyy's avatar
Rayyyyy committed
129
```bash
chenych's avatar
chenych committed
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
bash test.sh
```

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

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

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

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

Rayyyyy's avatar
Rayyyyy committed
147
| DEVICE | PSNR | SSIM | LPIPS |
chenych's avatar
chenych committed
148
| :------: | :------: | :------: | :------: |
Rayyyyy's avatar
Rayyyyy committed
149
150
| Z100L | 29.04 | 0.7615 | 0.1294 |
| V100S | 30.13 | 0.7981 | 0.1031 |
chenych's avatar
chenych committed
151
152
153
154
155
156

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

### 热点应用行业
chenych's avatar
chenych committed
157
交通,公安,制造
chenych's avatar
chenych committed
158
159

## 源码仓库及问题反馈
chenzk's avatar
chenzk committed
160
- http://developer.sourcefind.cn/codes/modelzoo/maskeddenoising_pytorch.git
chenych's avatar
chenych committed
161
162

## 参考资料
Rayyyyy's avatar
Rayyyyy committed
163
- https://github.com/haoyuc/MaskedDenoising.git