pytorch实现DCP暗通道先验去雾算法及其onnx导出

pytorch实现DCP暗通道先验去雾算法及其onnx导出

简介

最近在做图像去雾,于是在Pytorch上复现了一下dcp算法。暗通道先验去雾算法是大神何恺明2009年发表在CVPR上的一篇论文,还获得了当年的CVPR最佳论文。

实现

具体原理就不阐述了,网上的解析多的是,这里直接把用pytorch复现的代码贴出来:

python 复制代码
import torch


def dcp(img, omega=0.75):
    h, w = img.shape[2:]
    imsz = h * w
    # 要查找的是暗通道中前0.1%的值
    numpx = torch.clamp_min(imsz // 1000, 1)
    # 找到暗通道的索引,弄成[batch, 3, numpx],因为要匹配三个通道,所以需要expand
    dark = torch.min(img, dim=1, keepdim=True)[0]
    indices = torch.topk(dark.view(-1, imsz), k=numpx, dim=1)[1].view(-1, 1, numpx).expand(-1, 3, -1)
    # 用上述索引匹配原图中的3个通道,并求其平均值
    a = (torch.gather(img.view(-1, 3, imsz), 2, indices).sum(2) / numpx).view(-1, 3, 1, 1)
    
    # 代公式算tx
    tx =  1 - omega * torch.min(img / a.view(-1, 3, 1, 1), dim=1, keepdim=True)[0]
    # 代公式算jx
    return (img - a) / torch.clamp_min(tx, 0.1) + a

函数有两个参数:

  1. img:经归一化后的(N,C,H,W)布局的图像
  2. omega:DCP算法的一个参数ω,数值越大效果越强

如果想在模型训练时引入dcp算法,可以用nn.Module封装一下:

python 复制代码
class DCP(torch.nn.Module):
    def __init__(self, omega):
        self._omega = omega

    def forward(self, x):
        return dcp(x, self._omega)

ONNX导出

导出

既然能封装成Module,那么就顺便试了一下导出ONNX。

导出onnx需要安装onnx和onnxsim:

bash 复制代码
pip install onnx onnxsim

导出代码如下:

python 复制代码
import torch
import onnx
from onnxsim import simplify 


def dcp(img, omega=0.75):
    h, w = img.shape[2:]
    imsz = h * w
    # 要查找的是暗通道中前0.1%的值
    numpx = torch.clamp_min(imsz // 1000, 1)
    # 找到暗通道的索引,弄成[batch, 3, numpx],因为要匹配三个通道,所以需要expand
    dark = torch.min(img, dim=1, keepdim=True)[0]
    indices = torch.topk(dark.view(-1, imsz), k=numpx, dim=1)[1].view(-1, 1, numpx).expand(-1, 3, -1)
    # 用上述索引匹配原图中的3个通道,并求其平均值
    a = (torch.gather(img.view(-1, 3, imsz), 2, indices).sum(2) / numpx).view(-1, 3, 1, 1)
    
    # 代公式算tx
    tx =  1 - omega * torch.min(img / a.view(-1, 3, 1, 1), dim=1, keepdim=True)[0]
    # 代公式算jx
    return (img - a) / torch.clamp_min(tx, 0.1) + a

class DCPExport(torch.nn.Module):
    def forward(self, x, omega):
        return dcp(x, omega)

def export(output='dcp.onnx'):
    torch.onnx.export(
        DCPExport(), 
        (torch.randn(1, 3, 255, 255, dtype=torch.float32), torch.tensor(0.75, dtype=torch.float32)), 
        'dcp.onnx', 
        input_names=['fog_image', 'omega'], 
        output_names=['clear_image'], 
        dynamic_axes={
            'fog_image': {0: 'batch', 2: 'height', 3: 'width'},
            'clear_image': {0: 'batch', 2: 'height', 3: 'width'},
        }
    )
    onnx_model = onnx.load(output) 
    model_simp, check = simplify(onnx_model) 
    assert check, "简化模型失败" 
    onnx.save(model_simp, output) 

if __name__ == '__main__':
	export()

导出结果如下:

导出后的onnx输入输出如下:

  • 输入:
    1. fog_image[float32]:形状为NCHW,且归一化的有雾图像,其中通道数C必须为3
    2. omega[float32]:dcp的参数,类型为浮点数
  • 输出:
    1. clear_image[float32]:形状为NCHW,且归一化的无雾图像,其中通道数C为3

下载链接:https://pan.baidu.com/s/1A1jSJQBFCGTeM8vbHOrysQ?pwd=tl6p

测试

用cv2和pil都可以:

python 复制代码
import numpy as np
import cv2
from PIL import Image
from onnxruntime import InferenceSession


model = InferenceSession('dcp.onnx')

# CV2读图
image = cv2.imread('dehaze/dehaze/input/images/indoor1.jpg')
# 这里说明一下,因为dcp对所有通道进行同等变换,所以不用bgr和rgb互转了,出来的结果都是一样的
# x = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
x = np.transpose(image, (2, 0, 1))[None].astype(np.float32) / 255.
res = model.run(['clear_image'], {'fog_image': x, 'omega': np.array(0.75, dtype=np.float32)})[0][0]
res = np.transpose(res, (1, 2, 0))
res = np.clip(res*255+0.5, 0, 255).astype(np.uint8)
# res = cv2.cvtColor(res, cv2.COLOR_RGB2BGR)
cv2.imwrite('onnx-cv.png', np.concatenate((image, res), 1))

# PIL读图
image = Image.open('dehaze/dehaze/input/images/indoor1.jpg')
x = np.transpose(image, (2, 0, 1))[None].astype(np.float32) / 255.
res = model.run(None, {'fog_image': x, 'omega': np.array(0.75, dtype=np.float32)})[0][0]
res = np.transpose(res, (1, 2, 0))
res = np.clip(res*255+0.5, 0, 255).astype(np.uint8)
Image.fromarray(np.concatenate((image, res), 1)).save('onnx-pil.png')

效果:

相关推荐
Qiuner5 分钟前
【源力觉醒 创作者计划】开源、易用、强中文:文心一言4.5或是 普通人/非AI程序员 的第一款中文AI?
人工智能·百度·开源·文心一言·gitcode
未来之窗软件服务17 分钟前
chrome webdrive异常处理-session not created falled opening key——仙盟创梦IDE
前端·人工智能·chrome·仙盟创梦ide·东方仙盟·数据调式
赶紧去巡山25 分钟前
pyhton基础【23】面向对象进阶四
python
AI街潜水的八角34 分钟前
深度学习图像分类数据集—蘑菇识别分类
人工智能·深度学习·分类
旷世奇才李先生1 小时前
PyCharm 安装使用教程
ide·python·pycharm
飞睿科技1 小时前
乐鑫代理商飞睿科技,2025年AI智能语音助手市场发展趋势与乐鑫芯片解决方案分析
人工智能
许泽宇的技术分享1 小时前
从新闻到知识图谱:用大模型和知识工程“八步成诗”打造科技并购大脑
人工智能·科技·知识图谱
这里有鱼汤1 小时前
“对象”?对象你个头!——Python世界观彻底崩塌的一天
后端·python
坤坤爱学习2.01 小时前
求医十年,病因不明,ChatGPT:你看起来有基因突变
人工智能·ai·chatgpt·程序员·大模型·ai编程·大模型学
尘浮7281 小时前
60天python训练计划----day59
开发语言·python