通过编码模糊核空间探索图像去模糊

关于项目

我们引入了一种方法,将任意数据集中的锐化-模糊图像对的模糊算子编码到模糊核空间中。假设编码的核空间足够接近实际场景中的模糊算子,我们提出了一种用于盲图像去模糊的交替优化算法。它通过编码空间中的核来近似未知的模糊算子,并搜索相应的清晰图像。由于该方法的设计,编码的核空间是完全可微分的,因此可以很容易地应用于深度神经网络模型中。

方法详情和实验结果可以在我们的论文中找到:

@inproceedings{m_Tran-etal-CVPR21, 
  author = {Phong Tran and Anh Tran and Quynh Phung and Minh Hoai}, 
  title = {Explore Image Deblurring via Encoded Blur Kernel Space}, 
  year = {2021}, 
  booktitle = {Proceedings of the {IEEE} Conference on Computer Vision and Pattern Recognition (CVPR)} 
}

每当使用此代码库产生发表成果或将其整合到其他软件中时,请引用我们的论文。

目录

• 关于项目
• 入门指南
  • 前提条件
  • 安装
  • 使用预训练模型
• 训练和评估
• 模型库

入门指南

前提条件

• Python >= 3.7
• Pytorch >= 1.4.0
• CUDA >= 10.0

安装

git clone https://github.com/VinAIResearch/blur-kernel-space-exploring.git
cd blur-kernel-space-exploring

conda create -n BlurKernelSpace -y python=3.7
conda activate BlurKernelSpace
conda install --file requirements.txt

训练和评估

准备数据集

你可以在模型库部分下载数据集。

要使用自定义数据集,你的数据集必须按以下方式组织:

root
├── blur_imgs
    ├── 000
    ├──── 00000000.png
    ├──── 00000001.png
    ├──── ...
    ├── 001
    ├──── 00000000.png
    ├──── 00000001.png
    ├──── ...
├── sharp_imgs
    ├── 000
    ├──── 00000000.png
    ├──── 00000001.png
    ├──── ...
    ├── 001
    ├──── 00000000.png
    ├──── 00000001.png
    ├──── ...

其中root、blur_imgs和sharp_imgs文件夹可以有任意名称。例如,让root, blur_imgs, sharp_imgs分别为REDS, train_blur, train_sharp(即,你正在使用REDS训练集),然后使用以下脚本创建lmdb数据集:

python create_lmdb.py --H 720 --W 1280 --C 3 --img_folder REDS/train_sharp --name train_sharp_wval --save_path ../datasets/REDS/train_sharp_wval.lmdb
python create_lmdb.py --H 720 --W 1280 --C 3 --img_folder REDS/train_blur --name train_blur_wval --save_path ../datasets/REDS/train_blur_wval.lmdb

其中(H, C, W)是图像的形状(注意数据集中的所有图像必须具有相同的形状),img_folder是包含图像的文件夹,name是数据集的名称,save_path是保存目标(save_path必须以.lmdb结尾)。

脚本完成后,./REDS中将创建两个文件夹train_sharp_wval.lmdb和train_blur_wval.lmdb。

训练

要进行图像去模糊、数据增强和模糊生成,首先需要训练模糊编码网络(论文中的F函数)。这是你唯一需要训练的网络。创建数据集后,将options/kernel_encoding/REDS/woVAE.yml中的dataroot_HQ和dataroot_LQ的值更改为之前创建的清晰和模糊lmdb数据集的路径,然后使用以下脚本训练模型:

python train.py -opt options/kernel_encoding/REDS/woVAE.yml

其中opt是包含训练配置的yaml文件的路径。你可以在options文件夹中找到一些默认配置。检查点、训练状态和日志将保存在experiments/modelName中。你可以在yaml文件中更改配置(学习率、超参数、网络结构等)。

测试

数据增强

要增强给定的数据集,首先使用scripts/create_lmdb.py创建一个lmdb数据集,如前所述。然后使用以下脚本:

python data_augmentation.py --target_H=720 --target_W=1280 \
			    --source_H=720 --source_W=1280\
			    --augmented_H=256 --augmented_W=256\
                            --source_LQ_root=datasets/REDS/train_blur_wval.lmdb \
                            --source_HQ_root=datasets/REDS/train_sharp_wval.lmdb \
			    --target_HQ_root=datasets/REDS/test_sharp_wval.lmdb \
                            --save_path=results/GOPRO_augmented \
                            --num_images=10 \
                            --yml_path=options/data_augmentation/default.yml

(target_H, target_W)、(source_H, source_W)和(augmented_H, augmented_W)分别是目标图像、源图像和增强图像的所需形状。source_LQ_root、source_HQ_root和target_HQ_root是之前创建的参考模糊-清晰对和输入清晰图像的lmdb数据集的路径。num_images是增强数据集的大小。model_path是训练模型的路径。yml_path是模型配置文件的路径。结果将保存在save_path中。

生成新的模糊核

要根据清晰图像生成模糊图像，使用以下命令：

python generate_blur.py --yml_path=options/generate_blur/default.yml \
		        --image_path=imgs/sharp_imgs/mushishi.png \
			--num_samples=10
			--save_path=./res.png

其中model_path是预训练模型的路径，yml_path是配置文件的路径。image_path是清晰图像的路径。运行脚本后，与清晰图像对应的模糊图像将保存在save_path中。以下是一些预期输出： 注意：这只适用于使用--VAE标志训练的模型。输入图像的大小必须可被128整除。

通用去模糊

要对模糊图像进行去模糊处理，使用以下命令：

python generic_deblur.py --image_path imgs/blur_imgs/blur1.png --yml_path options/generic_deblur/default.yml --save_path ./res.png

其中image_path是模糊图像的路径。yml_path是配置文件的路径。去模糊后的图像将保存到save_path。

使用清晰图像先验进行去模糊

首先，你需要下载预训练的styleGAN或styleGAN2网络。如果你想使用styleGAN，下载映射和合成网络，然后将它们重命名并分别复制到experiments/pretrained/stylegan_mapping.pt和experiments/pretrained/stylegan_synthesis.pt。如果你想使用styleGAN2，下载预训练模型，然后将其重命名并复制到experiments/pretrained/stylegan2.pt。

要使用styleGAN潜在空间作为清晰图像先验来对模糊图像进行去模糊处理，你可以使用以下命令之一：

python domain_specific_deblur.py --input_dir imgs/blur_faces \
		    --output_dir experiments/domain_specific_deblur/results \
		    --yml_path options/domain_specific_deblur/stylegan.yml  # 使用stylegan的潜在空间
python domain_specific_deblur.py --input_dir imgs/blur_faces \
		    --output_dir experiments/domain_specific_deblur/results \
		    --yml_path options/domain_specific_deblur/stylegan2.yml  # 使用stylegan2的潜在空间

结果将保存在experiments/domain_specific_deblur/results中。 注意：通常，代码仍然适用于大小可被128整除的图像。然而，由于我们的模糊核不是均匀的，随着图像大小的增加，核的大小也会增加。

模型库

预训练模型和相应的数据集在下表中提供。下载数据集和模型后，按照测试部分的说明进行数据增强、生成模糊图像或图像去模糊。

模型名称	数据集	状态
REDS woVAE	REDS	:heavy_check_mark:
GOPRO woVAE	GOPRO	:heavy_check_mark:
GOPRO wVAE	GOPRO	:heavy_check_mark:
GOPRO + REDS woVAE	GOPRO, REDS	:heavy_check_mark:

注释和参考

训练代码借鉴自EDVR项目：https://github.com/xinntao/EDVR

骨干网络代码借鉴自DeblurGAN项目：https://github.com/KupynOrest/DeblurGAN

styleGAN代码借鉴自PULSE项目：https://github.com/adamian98/pulse

stylegan2代码借鉴自：https://github.com/rosinality/stylegan2-pytorch