admin管理员组

文章数量:1530808

本文介绍一个去雾算法ChaIR的使用方法,可以完成图像去雾,也可以用于图像去雨、去噪音等任务。本文不涉及论文原理,只包含源代码的跑通和使用。

先展示一下效果:

原图去雾

论文:Exploring the potential of channel interactions for image restoration

代码地址:https://github/c-yn/ChaIR/tree/main/Dehazing

本文的代码及数据集、训练好的权重:图像去雾代码-SOTS划分好的8:2数据集-训练好的去雾权重-包含推理代码

一、准备数据集

作者在github中给出了去雾数据集reside-indoor,reside-outdoor ,SOTS的地址,因为reside-indoor/outdoor太大了,本文介绍SOTS数据集的使用方法。

数据集地址:

数据集google drive百度云
reside-indoorhttps://drive.google/drive/folders/1pbtfTp29j7Ip-mRzDpMpyopCfXd-ZJhC?usp=sharing百度网盘 请输入提取码
reside-outdoorhttps://drive.google/drive/folders/1eL4Qs-WNj7PzsKwDRsgUEzmysdjkRs22?usp=sharing-
SOTShttps://drive.google/file/d/16j2dwVIa9q_0RtpIXMzhu-7Q6dwz_D1N/view?usp=sharing百度网盘 请输入提取码

这里请注意,使用SOTS需要将其转换为如下格式:

SOTS数据集也有indoor和outdoor,本文只使用outdoor,本文将SOTS outdoor数据集按照8:2划分训练集和测试集,并提供转换好的数据连接:SOTS数据集8:2划分训练和验证集,可用于训练去雾模型

准备好了数据集之后,按照如下目录结构放置即可:

至此数据集准备完成。

二、安装环境

接下来安装conda环境,首先下载代码:

git clone https://github/c-yn/ChaIR.git
cd ChaIR

创建虚拟环境:

conda create -n chair python=3.10
conda activate chair 

按照官网教程安装pytorch,我安装的是torch 2.3.1 cuda118(可以跳过):

conda install pytorch torchvision torchaudio pytorch-cuda=11.8 -c pytorch -c nvidia

安装pytorch-gradual-warmup-lr:

cd pytorch-gradual-warmup-lr
python setup.py install

后续使用过程中会提示缺少相关库,因为源代码未提供requirements.txt库,缺少的库需要自行安装,本文不做赘述。

三、训练验证

环境安装完成,因为本文只介绍去雾模型,所以进入ChaIR/Dehazing/OTS目录:

cd ChaIR/Dehazing/OTS

开始训练,如果显存不够,可以降低batch size,我用的rtx3060,我把batch size降低为2:

python main.py --mode train --data_dir  SOTS/outdoor

训练完成后,results/ChaIR目录的ots里面有权重,如下图:

使用best权重进行验证:

python main.py --data_dir  SOTS/outdoor --test_model  results/ChaIR/ots/Best.pkl

可以看到我训练的精度:

四、推理

因为ChaIR没有推理代码,所以自己写了一个推理代码进行推理,效果如下图:

原图标签ChaIR推理结果(自己的权重)

本文提供训练好的模型和推理代码以及数据集,地址:图像去雾代码-SOTS划分好的8:2数据集-训练好的去雾权重-包含推理代码

本文标签: 深度算法图像ChaIR