ridcp_pytorch基于codebook先验知识的图像去雾模型

RIDCP

论文

RIDCP: Revitalizing Real Image Dehazing via High-Quality Codebook Priors

• https://arxiv.org/pdf/2304.03994.pdf

模型结构

ridcp先使用Encoder对输入进行特征提取得到Z，特征Z流向Gvq以及G，流向Gvq的Z被替换(CHM方法)为码本中的特征Zvq，然后Zvq通过Gvq获取各个层的特征Fvqi，通过G获取到的特征Fi通过与Fvqi的特征进行融合(NFA方法)得到新的特征Fi，直至G的最后一层，得到目标输出。

算法原理

用途：该算法通过有效利用先验知识（码本）来完成图像去雾功能。

原理：

1. 融合特征(NFA)，消除仅使用码本特征导致的图像失真；
2. 使用CHM（距离计算方式）匹配更好的码本特征。

NFA

CHM

这里 z^代表通过Encoder获取到的特征, zk代表码本中的第k个特征，α 有最佳值21.25，fhk与fck为统计值。

环境配置

Docker (方法一)

docker pull image.sourcefind.cn:5000/dcu/admin/base/pytorch:2.1.0-ubuntu20.04-dtk24.04.1-py3.10
docker run --shm-size 10g --network=host --name=ridcp -v /opt/hyhal:/opt/hyhal:ro --privileged --device=/dev/kfd --device=/dev/dri --group-add video --cap-add=SYS_PTRACE --security-opt seccomp=unconfined -v 项目地址(绝对路径):/home/ -it <your IMAGE ID> bash   
pip install -r requirements.txt
BASICSR_EXT=True python setup.py develop

Dockerfile (方法二)

# 需要在对应的目录下
docker build -t <IMAGE_NAME>:<TAG> .
# <your IMAGE ID>用以上拉取的docker的镜像ID替换
docker run -it --shm-size 10g --network=host --name=ridcp --privileged --device=/dev/kfd --device=/dev/dri --group-add video --cap-add=SYS_PTRACE --security-opt seccomp=unconfined 
<your IMAGE ID> bash
BASICSR_EXT=True python setup.py develop

Anaconda (方法三)

1、关于本项目DCU显卡所需的特殊深度学习库可从光合开发者社区下载安装： https://developer.hpccube.com/tool/

DTK驱动：dtk24.04.1
python：python3.10
torch:2.1.0
torchvision:0.16.0
torchaudio:2.1.2
deepspeed:0.12.3
apex:1.1.0

Tips：以上dtk驱动、python、torch等DCU相关工具版本需要严格一一对应

2、其它非特殊库参照requirements.txt安装

pip install -r requirements.txt

3、其他

BASICSR_EXT=True python setup.py develop

Tips: 使用阿里的pip源

数据集

下载地址（训练集）：百度网盘 请输入提取码 （提取码：qqqo）

datasets
    |- clear_images_no_haze_no_dark_500
        |- xxx.jpg
        |- ...
    |- depth_500
        |- xxx.npy
        |- ...

下载地址（测试集）：https://sites.google.com/view/reside-dehaze-datasets/reside-β

RTTS
    |- Annotations
        |- xxx.xml
        |- ...
    |- ImageSets
        |- Main
            |- xxx.txt
            |- ...
    |- JPEGImages
        |- xxx.png
        |- ...

训练

模型下载：百度网盘 请输入提取码 （HQPs，CHM，提取码：huea）

# 这里的X表示设备编号，如0，1，2，3
CUDA_VISIBLE_DEVICES=X,X,X,X python basicsr/train.py -opt options/RIDCP.yml

推理

模型下载：百度网盘 请输入提取码 （RIDCP，提取码：huea）

python inference_ridcp.py -i examples -w pretrained_models/pretrained_RIDCP.pth -o results --use_weight --alpha -21.25

参数说明：

1. -i <路径>: 输入的图片或者文件夹
2. -w <路径>: 模型参数
3. -o <路径>: 处理完后的结果路径

result

上图是输入图像，下图是输出图像

精度

无

应用场景

算法类别

图像去雾

热点应用行业

交通,医疗,环保,气象

源码仓库及问题反馈

ModelZoo / RIDCP_pytorch · GitLab

参考资料

• https://github.com/RQ-Wu/RIDCP_dehazing