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

2023-12-20 20:15

本文主要是介绍pytorch实现DCP暗通道先验去雾算法及其onnx导出,希望对大家解决编程问题提供一定的参考价值,需要的开发者们随着小编来一起学习吧!

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

  • 简介
  • 实现
  • ONNX导出
    • 导出
    • 测试

简介

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

实现

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

import torchdef 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],因为要匹配三个通道,所以需要expanddark = 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)# 代公式算txtx =  1 - omega * torch.min(img / a.view(-1, 3, 1, 1), dim=1, keepdim=True)[0]# 代公式算jxreturn (img - a) / torch.clamp_min(tx, 0.1) + a

函数有两个参数:

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

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

class DCP(torch.nn.Module):def __init__(self, omega):self._omega = omegadef forward(self, x):return dcp(x, self._omega)

ONNX导出

导出

既然能封装成Module,那么就顺便试了一下导出ONNX。
导出onnx需要安装onnx和onnxsim:

pip install onnx onnxsim

导出代码如下:

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],因为要匹配三个通道,所以需要expanddark = 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)# 代公式算txtx =  1 - omega * torch.min(img / a.view(-1, 3, 1, 1), dim=1, keepdim=True)[0]# 代公式算jxreturn (img - a) / torch.clamp_min(tx, 0.1) + aclass 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
导出后的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都可以:

import numpy as np
import cv2
from PIL import Image
from onnxruntime import InferenceSessionmodel = 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')

效果:

onnx效果

这篇关于pytorch实现DCP暗通道先验去雾算法及其onnx导出的文章就介绍到这儿,希望我们推荐的文章对编程师们有所帮助!



http://www.chinasem.cn/article/517347

相关文章

python设置环境变量路径实现过程

《python设置环境变量路径实现过程》本文介绍设置Python路径的多种方法:临时设置(Windows用`set`,Linux/macOS用`export`)、永久设置(系统属性或shell配置文件... 目录设置python路径的方法临时设置环境变量(适用于当前会话)永久设置环境变量(Windows系统

Python对接支付宝支付之使用AliPay实现的详细操作指南

《Python对接支付宝支付之使用AliPay实现的详细操作指南》支付宝没有提供PythonSDK,但是强大的github就有提供python-alipay-sdk,封装里很多复杂操作,使用这个我们就... 目录一、引言二、准备工作2.1 支付宝开放平台入驻与应用创建2.2 密钥生成与配置2.3 安装ali

Spring Security 单点登录与自动登录机制的实现原理

《SpringSecurity单点登录与自动登录机制的实现原理》本文探讨SpringSecurity实现单点登录(SSO)与自动登录机制,涵盖JWT跨系统认证、RememberMe持久化Token... 目录一、核心概念解析1.1 单点登录(SSO)1.2 自动登录(Remember Me)二、代码分析三、

PyCharm中配置PyQt的实现步骤

《PyCharm中配置PyQt的实现步骤》PyCharm是JetBrains推出的一款强大的PythonIDE,结合PyQt可以进行pythion高效开发桌面GUI应用程序,本文就来介绍一下PyCha... 目录1. 安装China编程PyQt1.PyQt 核心组件2. 基础 PyQt 应用程序结构3. 使用 Q

Python实现批量提取BLF文件时间戳

《Python实现批量提取BLF文件时间戳》BLF(BinaryLoggingFormat)作为Vector公司推出的CAN总线数据记录格式,被广泛用于存储车辆通信数据,本文将使用Python轻松提取... 目录一、为什么需要批量处理 BLF 文件二、核心代码解析:从文件遍历到数据导出1. 环境准备与依赖库

linux下shell脚本启动jar包实现过程

《linux下shell脚本启动jar包实现过程》确保APP_NAME和LOG_FILE位于目录内,首次启动前需手动创建log文件夹,否则报错,此为个人经验,供参考,欢迎支持脚本之家... 目录linux下shell脚本启动jar包样例1样例2总结linux下shell脚本启动jar包样例1#!/bin

go动态限制并发数量的实现示例

《go动态限制并发数量的实现示例》本文主要介绍了Go并发控制方法,通过带缓冲通道和第三方库实现并发数量限制,文中通过示例代码介绍的非常详细,对大家的学习或者工作具有一定的参考学习价值,需要的朋友们下面... 目录带有缓冲大小的通道使用第三方库其他控制并发的方法因为go从语言层面支持并发,所以面试百分百会问到

PyTorch中的词嵌入层(nn.Embedding)详解与实战应用示例

《PyTorch中的词嵌入层(nn.Embedding)详解与实战应用示例》词嵌入解决NLP维度灾难,捕捉语义关系,PyTorch的nn.Embedding模块提供灵活实现,支持参数配置、预训练及变长... 目录一、词嵌入(Word Embedding)简介为什么需要词嵌入?二、PyTorch中的nn.Em

Go语言并发之通知退出机制的实现

《Go语言并发之通知退出机制的实现》本文主要介绍了Go语言并发之通知退出机制的实现,文中通过示例代码介绍的非常详细,对大家的学习或者工作具有一定的参考学习价值,需要的朋友们下面随着小编来一起学习学习吧... 目录1、通知退出机制1.1 进程/main函数退出1.2 通过channel退出1.3 通过cont

Python实现PDF按页分割的技术指南

《Python实现PDF按页分割的技术指南》PDF文件处理是日常工作中的常见需求,特别是当我们需要将大型PDF文档拆分为多个部分时,下面我们就来看看如何使用Python创建一个灵活的PDF分割工具吧... 目录需求分析技术方案工具选择安装依赖完整代码实现使用说明基本用法示例命令输出示例技术亮点实际应用场景扩