README.md 5.25 KB
Newer Older
mashun1's avatar
diffbir  
mashun1 committed
1
# DiffBIR
0x3f3f3f3fun's avatar
0x3f3f3f3fun committed
2

mashun1's avatar
diffbir  
mashun1 committed
3
## 论文
0x3f3f3f3fun's avatar
0x3f3f3f3fun committed
4

mashun1's avatar
diffbir  
mashun1 committed
5
**DiffBIR: Towards Blind Image Restoration with Generative Diffusion Prior**
0x3f3f3f3fun's avatar
0x3f3f3f3fun committed
6

mashun1's avatar
diffbir  
mashun1 committed
7
* https://arxiv.org/abs/2308.15070
0x3f3f3f3fun's avatar
0x3f3f3f3fun committed
8

mashun1's avatar
diffbir  
mashun1 committed
9
## 模型结构
0x3f3f3f3fun's avatar
0x3f3f3f3fun committed
10

mashun1's avatar
diffbir  
mashun1 committed
11
第一阶段模型使用8个Swin Transformer blocks(RSTB),每个RSTB中包含6个Swin Transformer Layers(STL),其中head数为6,window size为8.
0x3f3f3f3fun's avatar
0x3f3f3f3fun committed
12

mashun1's avatar
diffbir  
mashun1 committed
13
第二阶段模型基于Stable Diffusioin 2.1-base,创建了一个与Unet中encoder block与middle block相同的网络。
0x3f3f3f3fun's avatar
0x3f3f3f3fun committed
14

mashun1's avatar
diffbir  
mashun1 committed
15
![Alt text](images/image.png)
0x3f3f3f3fun's avatar
0x3f3f3f3fun committed
16

mashun1's avatar
diffbir  
mashun1 committed
17
## 算法原理
0x3f3f3f3fun's avatar
0x3f3f3f3fun committed
18

mashun1's avatar
diffbir  
mashun1 committed
19
用途:该算法为两阶段算法,可以提升图像的分辨率。
0x3f3f3f3fun's avatar
0x3f3f3f3fun committed
20

mashun1's avatar
diffbir  
mashun1 committed
21
第一阶段使用复原模块,从具有未知和复杂降质的低质量(LQ)图像中恢复清晰图像;
0x3f3f3f3fun's avatar
0x3f3f3f3fun committed
22

mashun1's avatar
diffbir  
mashun1 committed
23
第二阶段使用生成模块来重新生成丢失的信息。
0x3f3f3f3fun's avatar
0x3f3f3f3fun committed
24

mashun1's avatar
diffbir  
mashun1 committed
25
## 环境配置
0x3f3f3f3fun's avatar
0x3f3f3f3fun committed
26
27


mashun1's avatar
diffbir  
mashun1 committed
28
### Docker(方法一)
0x3f3f3f3fun's avatar
0x3f3f3f3fun committed
29

mashun1's avatar
diffbir  
mashun1 committed
30
31
32
    docker pull image.sourcefind.cn:5000/dcu/admin/base/pytorch:1.13.1-centos7.6-dtk-23.04.1-py39-latest
    docker run --shm-size 10g --network=host --name=diffbir --privileged --device=/dev/kfd --device=/dev/dri --group-add video --cap-add=SYS_PTRACE --security-opt seccomp=unconfined -v 项目地址(绝对路径):/home/ -it <your IMAGE ID> bash
    pip install -r requirements.txt
0x3f3f3f3fun's avatar
0x3f3f3f3fun committed
33
34


mashun1's avatar
diffbir  
mashun1 committed
35
### Docker(方法二)
0x3f3f3f3fun's avatar
0x3f3f3f3fun committed
36

mashun1's avatar
diffbir  
mashun1 committed
37
38
39
40
    # 需要在对应的目录下
    docker build -t <IMAGE_NAME>:<TAG> .
    # <your IMAGE ID>用以上拉取的docker的镜像ID替换
    docker run -it --shm-size 10g --network=host --name=diffbir --privileged --device=/dev/kfd --device=/dev/dri --group-add video --cap-add=SYS_PTRACE --security-opt seccomp=unconfined <your IMAGE ID> bash
41

mashun1's avatar
diffbir  
mashun1 committed
42
43
44
### Anaconda (方法三)
1、关于本项目DCU显卡所需的特殊深度学习库可从光合开发者社区下载安装:
https://developer.hpccube.com/tool/
0x3f3f3f3fun's avatar
0x3f3f3f3fun committed
45

mashun1's avatar
diffbir  
mashun1 committed
46
47
48
49
50
51
52
    DTK驱动:dtk23.04.1
    python:python3.9
    torch:1.13.1
    torchvision:0.14.1
    torchaudio:0.13.1
    deepspeed:0.9.2
    apex:0.1
53

mashun1's avatar
diffbir  
mashun1 committed
54
Tips:以上dtk驱动、python、torch等DCU相关工具版本需要严格一一对应
0x3f3f3f3fun's avatar
0x3f3f3f3fun committed
55

mashun1's avatar
diffbir  
mashun1 committed
56
2、其它非特殊库参照requirements.txt安装
0x3f3f3f3fun's avatar
0x3f3f3f3fun committed
57

mashun1's avatar
diffbir  
mashun1 committed
58
    pip install -r requirements.txt
59

0x3f3f3f3fun's avatar
0x3f3f3f3fun committed
60

mashun1's avatar
diffbir  
mashun1 committed
61
## 数据集
0x3f3f3f3fun's avatar
0x3f3f3f3fun committed
62

mashun1's avatar
diffbir  
mashun1 committed
63
下载地址(训练+测试集):https://www.image-net.org/ (imagenet1k)
0x3f3f3f3fun's avatar
0x3f3f3f3fun committed
64

mashun1's avatar
diffbir  
mashun1 committed
65
66
67
68
    datasets
        |- train
            |- n01440764
                |- xxx.JPEG
0x3f3f3f3fun's avatar
0x3f3f3f3fun committed
69

mashun1's avatar
mashun1 committed
70
71
72
73
74
75
76
下载地址(训练+测试集):https://www.robots.ox.ac.uk/~vgg/data/pets/

    datasets
        |- images
            |- xxx.JPG

注意:以上两个数据集都可以使用,可任选其一。
ziyannchen's avatar
ziyannchen committed
77

mashun1's avatar
diffbir  
mashun1 committed
78
## 训练
79

mashun1's avatar
diffbir  
mashun1 committed
80
### 阶段一
81

mashun1's avatar
diffbir  
mashun1 committed
82
83
84
85
86
87
1、数据准备:该操作用于生成训练以及验证数据路径列表

    python scripts/make_file_list.py \  
    --img_folder [hq_dir_path] \         # 包含图片的文件夹
    --val_size [validation_set_size] \   # 验证集大小
    --save_folder [save_dir_path] \      # 路径列表保存文件夹
ziyannchen's avatar
ziyannchen committed
88
    --follow_links
0x3f3f3f3fun's avatar
0x3f3f3f3fun committed
89

mashun1's avatar
diffbir  
mashun1 committed
90
91
92
2、修改配置文件

    修改 `configs/dataset`中相应的yaml配置文件
0x3f3f3f3fun's avatar
0x3f3f3f3fun committed
93

mashun1's avatar
diffbir  
mashun1 committed
94
    修改 `configs/train_swinir.yaml`配置文件
0x3f3f3f3fun's avatar
0x3f3f3f3fun committed
95

mashun1's avatar
diffbir  
mashun1 committed
96
97
98
3、训练

    python train.py --config [training_config_path]
0x3f3f3f3fun's avatar
0x3f3f3f3fun committed
99

mashun1's avatar
diffbir  
mashun1 committed
100
注意:该阶段训练得到的模型将用于第二阶段的训练。
0x3f3f3f3fun's avatar
0x3f3f3f3fun committed
101

mashun1's avatar
diffbir  
mashun1 committed
102
### 阶段二
0x3f3f3f3fun's avatar
0x3f3f3f3fun committed
103

mashun1's avatar
diffbir  
mashun1 committed
104
1、模型准备(Stable Diffusion v2.1): https://huggingface.co/stabilityai/stable-diffusion-2-1-base/resolve/main/v2-1_512-ema-pruned.ckpt
0x3f3f3f3fun's avatar
0x3f3f3f3fun committed
105
106


mashun1's avatar
diffbir  
mashun1 committed
107
108
109
110
111
112
113
114
115
116
117
118
119
2、初始化模型参数

    python scripts/make_stage2_init_weight.py \
    --cldm_config configs/model/cldm.yaml \
    --sd_weight [sd_v2.1_ckpt_path] \
    --swinir_weight [swinir_ckpt_path] \  # 第一阶段训练得到的模型
    --output [init_weight_output_path]    # 初始化模型保存地址

3、修改配置文件

    修改`configs/train_cldm.yaml`配置文件

4、训练
0x3f3f3f3fun's avatar
0x3f3f3f3fun committed
120
121
122

    python train.py --config [training_config_path]

mashun1's avatar
diffbir  
mashun1 committed
123
## 推理
0x3f3f3f3fun's avatar
0x3f3f3f3fun committed
124

mashun1's avatar
diffbir  
mashun1 committed
125
### general Image
0x3f3f3f3fun's avatar
0x3f3f3f3fun committed
126

mashun1's avatar
diffbir  
mashun1 committed
127
128
模型下载地址:
* https://huggingface.co/lxq007/DiffBIR/resolve/main/general_full_v1.ckpt
0x3f3f3f3fun's avatar
0x3f3f3f3fun committed
129

mashun1's avatar
diffbir  
mashun1 committed
130
* https://huggingface.co/lxq007/DiffBIR/resolve/main/general_swinir_v1.ckpt
0x3f3f3f3fun's avatar
0x3f3f3f3fun committed
131

mashun1's avatar
diffbir  
mashun1 committed
132
133
134
135
136
137
138
139
140
141
        python inference.py \
        --input inputs/demo/general \
        --config configs/model/cldm.yaml \
        --ckpt weights/general_full_v1.ckpt \  
        --reload_swinir --swinir_ckpt weights/general_swinir_v1.ckpt \  
        --steps 50 \
        --sr_scale 4 \
        --color_fix_type wavelet \
        --output results/demo/general \
        --device cuda [--tiled --tile_size 512 --tile_stride 256]
0x3f3f3f3fun's avatar
0x3f3f3f3fun committed
142

mashun1's avatar
diffbir  
mashun1 committed
143
注意:方括号中的参数为可选项,模型也可以替换为在训练阶段得到的
0x3f3f3f3fun's avatar
0x3f3f3f3fun committed
144

mashun1's avatar
diffbir  
mashun1 committed
145
### Face Image
0x3f3f3f3fun's avatar
0x3f3f3f3fun committed
146

mashun1's avatar
diffbir  
mashun1 committed
147
模型下载地址:
0x3f3f3f3fun's avatar
0x3f3f3f3fun committed
148

mashun1's avatar
diffbir  
mashun1 committed
149
* https://huggingface.co/lxq007/DiffBIR/resolve/main/face_full_v1.ckpt
0x3f3f3f3fun's avatar
0x3f3f3f3fun committed
150

mashun1's avatar
diffbir  
mashun1 committed
151
for aligned face inputs
0x3f3f3f3fun's avatar
0x3f3f3f3fun committed
152

mashun1's avatar
diffbir  
mashun1 committed
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
    python inference_face.py \
    --input inputs/demo/face/aligned \
    --sr_scale 1 \
    --output results/demo/face/aligned \
    --has_aligned \
    --device cuda

for unaligned face inputs

    python inference_face.py \
    --input inputs/demo/face/whole_img \
    --sr_scale 2 \
    --output results/demo/face/whole_img \
    --bg_upsampler DiffBIR \
    --device cuda

## result

恢复后的图像

![Alt text](images/samples_step-004900_e-000008_b-001203.png)

低质量图像

![Alt text](images/lq_step-004900_e-000008_b-001203.png)

### 精度



## 应用场景
0x3f3f3f3fun's avatar
0x3f3f3f3fun committed
184

mashun1's avatar
diffbir  
mashun1 committed
185
### 算法类别
0x3f3f3f3fun's avatar
0x3f3f3f3fun committed
186

mashun1's avatar
diffbir  
mashun1 committed
187
`图像超分`
0x3f3f3f3fun's avatar
0x3f3f3f3fun committed
188

mashun1's avatar
diffbir  
mashun1 committed
189
### 热点应用行业
0x3f3f3f3fun's avatar
0x3f3f3f3fun committed
190

mashun1's avatar
diffbir  
mashun1 committed
191
`媒体、科研、教育`
0x3f3f3f3fun's avatar
0x3f3f3f3fun committed
192

mashun1's avatar
diffbir  
mashun1 committed
193
## 源码仓库及问题反馈
0x3f3f3f3fun's avatar
0x3f3f3f3fun committed
194

mashun1's avatar
mashun1 committed
195
https://developer.hpccube.com/codes/modelzoo/diffbir_pytorch
0x3f3f3f3fun's avatar
0x3f3f3f3fun committed
196

mashun1's avatar
diffbir  
mashun1 committed
197
## 参考
0x3f3f3f3fun's avatar
0x3f3f3f3fun committed
198

mashun1's avatar
diffbir  
mashun1 committed
199
* https://github.com/XPixelGroup/DiffBIR
mashun1's avatar
mashun1 committed
200
* https://github.com/XPixelGroup/DiffBIR/issues/55