CVPR21 《A Two-stage Deep Network for High Dynamic Range Image Reconstruction》HDR高动态范围图像重建论文复现

CVPR21 《A Two-stage Deep Network for High Dynamic Range Image Reconstruction》HDR高动态范围图像重建论文复现

Linux环境配置

docker
Pytorch 1.7.1 - cuda 11.0 - cudnn8-devel
torchvision
tensorboard
opencv-python
etaprogress
torchsummary
ptflops
matplotlib
scikit-learn
scikit-image

代码结构

获取训练集

图像训练集来源于HdM HDR dataset，该数据集一共有1289组场景。每一组场景所包含的图像有：长曝光8bit LDR影像，中等曝光8bit LDR影像，短曝光8bit LDR影像，16bit HDR 影像(ground-truth).
我们使用其中的1000组场景进行训练，余下场景用于测试验证。
从网站下载训练集 https://competitions.codalab.org/competitions/28161#participate
下载的7个图像压缩包统一解压到train文件夹

训练集下载链接

训练集图像预处理

该训练集图像的尺寸为19001060，用于训练的卷积层设定尺寸为255255，因此需要对训练集进行预处理切割。以255为边长对图像进行随机切割，最终得到7551 图像patches。
对于该训练模型，规定输入的一组对照集为： 一张8bit LDR 图像， 一张8 bit HDR 图像， 一张16bit HDR 图像，（由于原始训练集只有16bit HDR图像，因此8bit HDR 图像是由16 bit HDR图像切割压缩归一化得到的）

指令

// An highlighted block
cd extras //进入extras文件夹
python processHDMDHR.py -r ../train -t ../patch/ -p 256   //运行训练集处理代码

将存在train文件夹里的原始训练集图片（19001060）分割成256256大小的patches

训练卷积，得到训练模型

指令

cd .. //回到main.py 所在路径
python main.py -ts -e 25 -b 8

参数

-ts 开始训练
-tr 加载预训练模型，继续训练
-e 设置epoch值 默认25
-b 设置batch值 默认8

使用GPU训练一次耗时约16小时
得到/weights/Two-stageHDR_checkpoint.pth权重模型

训练卷积，获得训练模型

得到/weights/Two-stageHDR_checkpoint.pth权重模型

测试训练模型

下载测试集 测试集下载链接

解压到inputimage文件夹
代码训练对象为中度曝光图像，因此输入测试图像%04d_medium.png

指令

python main.py -i -s ./inputimage -d ./output

输出结果保存在/outputTwo-stageHDR/inputimage路径下
因为普通显示屏不支持16bit HDR图像，所以图像亮度很低，需要支持16bit显示屏可以观看高清图像

16bit HDR图像可视化

为了验证我们的模型对于LDR 转换HDR的效果，需要对得到16bit HDR图像进行可视化，从而在普通显示屏上与8bit LDR进行对比观察验证。因此对16bit HDR 图像进行压缩剪裁为8bit HDR，能够与8bit LDR图像在普通显示屏上直接观察对比性能

指令

cd visual
python visual16to8.py -r ../outputTwo-stageHDR/inputimage -t ./vout/ -p 1060

得到由16bit HDR压缩而来的8bit HDR图像，色彩亮度正常

输入8 bit LDR 与输出8 bit HDR 图像的直观细节对比

（左LDR 右HDR )

图像转视频处理

指令

cd visual
python getvideo.py

将vout文件夹里的图片转化为视频输出到visual文件夹中
生成HDR8output.mp4视频

总思路

8bit LDR视频 → 8 bit LDR 图像 →16 bit HDR 图像 →8 bit HDR 图像 → 8 bit HDR 视频

最终实现8 bit LDR 到16bit HDR（8 bit HDR 可视化）视频转换

优缺点分析

优点：

• 能够在未获得照片拍摄设备信息的前提下，将8bit LDR图像输入后获得16bit HDR图像
• 比特扩展图像重建性能较高，16bit高动态范围图像
• 不需要多余的图像预处理步骤，集成度较高

缺点：

• 对算力要求较高，需要较高性能GPU支持
• 运算速度有待提高，目前无法达到实时视频转化，暂考虑本地播放器异步移植
• 目前大多屏幕不支持16bitHDR，需要考虑色调映射或硬件更新

补充

该论文目的是研究8bit LDR -16bit HDR 图像重建，关于视频生成以及16bit HDR 可视化为8 bit HDR图像的功能是我补充的代码完成的。整个项目完整跑通的代码放在我的github 上，方便大家下载尝试~
https://github.com/xypu98/two-stage-HDR-video/tree/master
github复现代码