图像去模糊代码—NAFNet运行记录
NAFNet论文是ECCV 2022的一篇文章,是计算机低级视觉-图像重建方面的工作,文章做的实验有图像去模糊、图像去噪、图像超分辨率等,实验效果非常棒,代码也写的非常好。
本博文教你如何用该论文的代码做图像去模糊的推理。实际上代码的readme.md已经写的很清楚了,但是有些人github也上不去,也挺闹心。
关于训练,因为前期准备数据集比较麻烦,暂时就不讲训练了。
1. 电脑环境
Centos (服务器操作系统)+miniconda (用于创建虚拟环境)+ NAFNet (代码) + torch (1.13.1+cu116)
2. 了解工作
论文题目:Simple Baselines for Image Restoration
论文地址:https://arxiv.org/pdf/2204.04676v4
论文github地址:https://github.com/megvii-research/NAFNet
强烈建议给个star,毕竟代码写的确实很棒,阅读对自己很有用。
官方给出的各个模型预训练权重结果如下,可以从谷歌云盘或者百度云下载。
Results and Pre-trained Models
name | Dataset | PSNR | SSIM | pretrained models |
---|---|---|---|---|
NAFNet-GoPro-width32 | GoPro | 32.8705 | 0.9606 | gdrive | 百度网盘 |
NAFNet-GoPro-width64 | GoPro | 33.7103 | 0.9668 | gdrive | 百度网盘 |
NAFNet-SIDD-width32 | SIDD | 39.9672 | 0.9599 | gdrive | 百度网盘 |
NAFNet-SIDD-width64 | SIDD | 40.3045 | 0.9614 | gdrive | 百度网盘 |
NAFNet-REDS-width64 | REDS | 29.0903 | 0.8671 | gdrive | 百度网盘 |
NAFSSR-L_4x | Flickr1024 | 24.17 | 0.7589 | gdrive | 百度网盘 |
NAFSSR-L_2x | Flickr1024 | 29.68 | 0.9221 | gdrive | 百度网盘 |
Baseline-GoPro-width32 | GoPro | 32.4799 | 0.9575 | gdrive | 百度网盘 |
Baseline-GoPro-width64 | GoPro | 33.3960 | 0.9649 | gdrive | 百度网盘 |
Baseline-SIDD-width32 | SIDD | 39.8857 | 0.9596 | gdrive | 百度网盘 |
Baseline-SIDD-width64 | SIDD | 40.2970 | 0.9617 | gdrive | 百度网盘 |
如果不能访问github,可以从我百度云下载代码和预训练权重,experiments/pretrained_models路径中有所有的预训练模型权重。
百度云链接:
链接:https://pan.baidu.com/s/1m_i54ioEZsde8euky7qK6Q?pwd=rcu1
提取码:rcu1
备注:本文只进行推理预测,如果想要训练自己的数据集,需要对数据集进行相应的处理。
3. 前提准备工作
已经安装过anaconda或者miniconda,电脑存在conda环境。
4. 配置环境
- 新建环境
conda create -n nafnet python=3.10
- 激活进入环境
conda activate nafnet
- 安装必要的依赖
pip instll -r requirement.txt
如果无法安装tb-nightly是因为清华、bfsu镜像源的问题,临时更改为阿里源:
python -m pip install tb-nightly -i https://mirrors.aliyun.com/pypi/simple
- 安装torch
pip install torch==1.13.1+cu116 torchvision==0.14.1+cu116 torchaudio==0.13.1 --extra-index-url https://download.pytorch.org/whl/cu116
安装torch之后,记得查看torch是否可用cuda,使用如下命令,首先进入python环境,接着导入torch,然后使用如下命令查看cuda是否可用,如果可用,返回True,如果cuda不可用,返回False。
import torch
torch.cuda.is_available()
- 安装basicsr
python setup.py develop --no_cuda_ext
这个安装可能容易出现问题,记得查看安装是否成功。
显示如下表示安装成功:
如果没有Cython,可能会报错如下,
需要安装Cython
pip install Cython
5. 预测
如果是直接预测,不训练,使用如下代码就可以直接对图像去模糊。
python basicsr/demo.py -opt options/test/REDS/NAFNet-width64.yml --input_path ./demo/blurry.jpg --output_path ./demo/deblur_img.png
图像去模糊如下,效果非常好:
6. 训练
准备数据集可能比较麻烦了,在这就不讲了,就说一下训练的代码。如果有人需要,到时候再写一篇训练的博客。
训练时间一般比较长,一块卡,600张图片对,batch设置为4,训练大概也需要1天多了。
如果是单卡,使用如下命令训练,前提是数据集已经准备好,且放入了对应的文件中:
python -m torch.distributed.launch --nproc_per_node=1 --master_port=4321 basicsr/train.py -opt options/train/GoPro/NAFNet-width32.yml --launcher pytorch
训练成功,有类似如下界面