zoom-learn-zoom
CVPR 2019论文《Zoom to Learn, Learn to Zoom》的代码
本文表明,在将机器学习应用于摄影数字变焦时,使用真实的RAW传感器数据进行训练是有益的。此代码基于tensorflow(在V 1.13.1上测试)。已在Ubuntu 16.04 LTS上测试通过。
SR-RAW数据集
使用SR-RAW
SR-RAW训练和测试数据现在可以在这里下载。
(如果你想在不下载完整训练/测试数据集的情况下尝试,请参阅快速推理部分)
要下载测试数据集(7 GB),运行:
bash ./scripts/download.sh 19zlN1fqRRm7E_6i5J3B1OskJocVeuvzG test.zip
unzip test.zip
rm test.zip
我们使用35mm图像(主要在序列中命名为'00006')进行测试。
要下载训练数据集(58 GB),运行:
bash ./scripts/download.sh 1qp6z3F4Ru9srwq1lNZr3pQ4kcVN-AOlM train.zip
unzip train.zip
rm train.zip
百度网盘上的训练数据集: @llp1996 已经很好心地将数据集上传到了百度网盘。提取码是:wi02。原始问题在这里提出。
尝试使用自己的数据
我们的模型是在索尼数码相机RAW格式的原始数据上训练的。如果你使用其他类型的原始数据格式,比如iPhone使用的DNG(你可以使用应用Halide从iPhone存储原始数据),有必要用该格式的原始数据对模型进行微调。
快速推理
我们将下载预训练模型和示例原始数据。
git clone https://github.com/ceciliavision/zoom-learn-zoom.git
cd zoom-learn-zoom
bash ./scripts/download.sh 1iForbFhhWqrq22FA1xIusfUpdi8td4Kq model.zip
unzip model.zip
bash ./scripts/download.sh 1WVSGaKIJVHwphTKhcr9ajolEnBh3aUkR quick_inference.zip
unzip quick_inference.zip
rm *.zip
python3 inference.py
关于config/inference.yaml
的注意事项
- 要对文件夹进行推理,将
mode
设置为inference
,并设置inference_root
(例如./quick_inference/
) - 要对单个图像进行推理,将
mode
设置为inference_single
,并设置inference_path
(例如./quick_inference/00134.ARW
) - 设置
task_folder
(例如./restore_4x
) - 结果保存在
./[task_folder]/[mode]
中
训练
CoBi损失
论文中提出的CoBi损失的实现可以在./CX
目录中找到。它是基于原始上下文损失实现修改的。参考./loss.py
查看它的使用方法。完整的训练流程正在准备中,将在10月左右发布。
数据预处理
我们提供了对齐函数和脚本,以解决在捕获数据集时的手部运动问题。这是一个可选步骤,因为CoBi损失不需要像素级对齐的数据对进行训练。然而,我们注意到进行初步(不精确)对齐步骤可以加快收敛速度。总之,我们提供:
./scripts/run_align.sh
是调用./main_crop.py
和./main_align_camera.py
的脚本,它首先对视场进行对齐,然后解决图像之间的手部运动错位问题./scripts/run_wb.sh
是调用./main_wb.py
的脚本,用于计算应用于相机ISP中处理过的图像的白平衡
要运行这些脚本,请在[TRAIN_PATH]中填入你的本地训练数据路径,在[TEST_PATH]中填入你的本地测试数据路径。如果你使用自己收集的数据进行训练,你可以遵循我们的数据目录结构或修改这些脚本。
bash ./scripts/run_align.sh [TRAIN_PATH]
bash ./scripts/run_wb.sh [TRAIN_PATH]
bash ./scripts/run_align.sh [TEST_PATH]
bash ./scripts/run_wb.sh [TEST_PATH]
运行这些脚本后,你可以在训练期间使用每个序列内的tform.txt
和wb.txt
。名为./cropped
、./compare
和./aligned
的文件夹仅用于可视化保存。
[2020年5月15日更新] 准备RAW-RGB对演示
我们添加了一个有文档说明的Jupyter笔记本demo_train_patch_pair.ipynb
,带你完成准备RAW-RGB对的过程。在utils.py
中添加了一些函数,请查看最新的提交。
引用
如果你发现这项工作对你的研究有用,请引用:
@inproceedings{zhang2019zoom
title={Zoom to Learn, Learn to Zoom},
author={Zhang, Xuaner and Chen, Qifeng and Ng, Ren and Koltun, Vladlen},
booktitle={Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition},
year={2019}
}
联系
如有任何问题,请联系我(Cecilia Zhang cecilia77@berkeley.edu)。