有哪些做品牌特卖的网站,企业网站建设公司哪家好,产品销售型企业网站,龙山建设工程有限公司网站pytorch实现DCP暗通道先验去雾算法及其onnx导出 简介实现ONNX导出导出测试 简介
最近在做图像去雾#xff0c;于是在Pytorch上复现了一下dcp算法。暗通道先验去雾算法是大神何恺明2009年发表在CVPR上的一篇论文#xff0c;还获得了当年的CVPR最佳论文。
实现
具体原理就不… pytorch实现DCP暗通道先验去雾算法及其onnx导出 简介实现ONNX导出导出测试 简介
最近在做图像去雾于是在Pytorch上复现了一下dcp算法。暗通道先验去雾算法是大神何恺明2009年发表在CVPR上的一篇论文还获得了当年的CVPR最佳论文。
实现
具体原理就不阐述了网上的解析多的是这里直接把用pytorch复现的代码贴出来
import torchdef dcp(img, omega0.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, dim1, keepdimTrue)[0]indices torch.topk(dark.view(-1, imsz), knumpx, dim1)[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), dim1, keepdimTrue)[0]# 代公式算jxreturn (img - a) / torch.clamp_min(tx, 0.1) a函数有两个参数
img经归一化后的(N,C,H,W)布局的图像omegaDCP算法的一个参数ω数值越大效果越强
如果想在模型训练时引入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, omega0.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, dim1, keepdimTrue)[0]indices torch.topk(dark.view(-1, imsz), knumpx, dim1)[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), dim1, keepdimTrue)[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(outputdcp.onnx):torch.onnx.export(DCPExport(), (torch.randn(1, 3, 255, 255, dtypetorch.float32), torch.tensor(0.75, dtypetorch.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输入输出如下
输入 fog_image[float32]形状为NCHW且归一化的有雾图像其中通道数C必须为3omega[float32]dcp的参数类型为浮点数 输出 clear_image[float32]形状为NCHW且归一化的无雾图像其中通道数C为3
下载链接https://pan.baidu.com/s/1A1jSJQBFCGTeM8vbHOrysQ?pwdtl6p
测试
用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, dtypenp.float32)})[0][0]
res np.transpose(res, (1, 2, 0))
res np.clip(res*2550.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, dtypenp.float32)})[0][0]
res np.transpose(res, (1, 2, 0))
res np.clip(res*2550.5, 0, 255).astype(np.uint8)
Image.fromarray(np.concatenate((image, res), 1)).save(onnx-pil.png)效果
