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复现代码