论文地址：1707.06543.pdf (arxiv.org)
项目地址：github.com/kivenyangmi…
关键词：除雾网络 、 CVPR 、 ORT推理 、 AODNET 、 自适应去模糊、 深度卷积网络、 多尺度
前言
由于AOD-Net网络是2017年的顶刊论文，在网上也有相当多的讲解，因此本文中将不会过多赘述论文中的内容，仅摘取精华进行讲解。
导读
AOD-Net 是一种用于图像去雾的深度神经网络模型，旨在通过对图像进行学习和分析，从中提取出场景中的深度信息，从而去除图像中的雾霾。AOD-Net 的全称为 “All-in-One Dehazing Network”，是由学者们在研究图像去雾的过程中提出的一种全新的深度神经网络模型。该模型采用了基于 CNN 的网络结构，在学习和分析图像的同时，还能够提取出场景中的深度信息，从而能够更加准确地去除图像中的雾霾。
相较于传统的图像去雾算法，AOD-Net 具有更高的去雾效果和更快的速度，并且能够适应各种不同的场景和光照条件。因此，AOD-Net 在图像去雾领域具有广泛的应用前景，并且已经被广泛地应用于实际的图像处理任务中。
原理
AOD-Net其原理基于自适应的最优估计理论和深度学习技术。AOD-Net的设计思想是通过将自适应最优估计理论引入深度学习框架中，实现图像去噪的自适应性和最优性。
结构组成
具体来说，AOD-Net由两个子网络组成，即噪声估计网络和噪声去除网络。噪声估计网络用于估计图像中的噪声水平，而噪声去除网络则根据噪声估计网络输出的噪声水平信息，以及图像本身的特征，对图像进行去噪处理。
AOD-Net的噪声估计网络使用一个类似于U-Net的架构，从原始图像中提取特征并预测噪声的分布。在噪声去除网络中，使用残差网络的结构进行去噪处理，同时使用一个可学习的标准化层来适应预测的噪声水平。
AOD-Net的优点在于它能够自适应地估计噪声水平，从而在不同的噪声水平下进行图像去噪，同时具有较高的去噪效果和计算效率。
代码
下面是使用torch框架实现的AOD-Net网络结构
import torch
import torch.nn as nn
class AODNet(nn.Module):
def __init__(self):
super(AODNet, self).__init__()
self.relu = nn.ReLU(inplace=True)
self.conv1 = nn.Conv2d(3, 3, 1, 1, 0, bias=True)
self.conv2 = nn.Conv2d(3, 3, 3, 1, 1, bias=True)
self.conv3 = nn.Conv2d(6, 3, 5, 1, 2, bias=True)
self.conv4 = nn.Conv2d(6, 3, 7, 1, 3, bias=True)
self.conv5 = nn.Conv2d(12, 3, 3, 1, 1, bias=True)
def forward(self, x):
x1 = self.relu(self.conv1(x))
x2 = self.relu(self.conv2(x1))
concat1 = torch.cat((x1, x2), 1)
x3 = self.relu(self.conv3(concat1))
concat2 = torch.cat((x2, x3), 1)
x4 = self.relu(self.conv4(concat2))
concat3 = torch.cat((x1, x2, x3, x4), 1)
x5 = self.relu(self.conv5(concat3))
clean_image = self.relu((x5 * x) - x5 + 1)
return clean_image
结构示例图
精髓部分
AOD-Net主要精髓是AOD-Net通过自适应去模糊和多尺度处理的结合，能够在不同场景下有效地提高图像去模糊的效果，具有较高的实用性和普适性。
具体来说，AOD-Net采用了一个深度卷积神经网络，通过输入模糊图像和一个预先设定的点扩散函数（PSF），输出一张更加清晰的图像。网络架构包含两个主要组成部分：自适应去模糊模块和多尺度处理模块。
AOD-Net自适应去模糊模块是AOD-Net的核心，它可以根据输入的模糊图像和PSF自适应地学习一个去模糊滤波器，用于恢复更清晰的图像。这个模块中采用了自适应残差学习机制，能够使网络在学习过程中自动选择最优的滤波器。此外，为了保证网络的泛化能力和稳定性，自适应去模糊模块还采用了一些特殊的技巧，例如局部标准化和正则化。
AOD-Net多尺度处理模块则可以在不同的尺度上处理输入图像，从而有效地提高去模糊效果。具体来说，AOD-Net使用了一个金字塔结构，将输入图像按照不同的尺度进行下采样，并分别输入自适应去模糊模块进行处理。之后，网络将不同尺度的输出图像进行上采样和融合，得到最终的去模糊结果。
实操
在实操中我将分模块的讲解如何运用网络训练一个自己的除雾网络以及进行ORT推理，在GITHUB上也有相当多的关于AOD-NET的代码，我这里也有一份代码放置在github上，欢迎大家自取。这里主要讲解数据集的准备以及使用ort推理。
数据集的准备
由于我们需要使用该网络，因此数据集是必不可少的，所以我们可以将本地的一些图像进行生成雾图加雾处理，然后将原图和含有雾的图像分成两个文件夹放置起来即可。
import math
import os
import random
import cv2
import numpy as np
from numba import jit
@jit(nopython=True)
def fog(j, center0, l, center1, size, beta):
d = -0.04 * math.sqrt((j - center0) ** 2 + (l - center1) ** 2) + size
td = math.exp(-beta * d)
return d, td
def put_fog(raw_img_path, fog_img_path):
img = cv2.imread(raw_img_path)
img_f = img / 255.0
(row, col, chs) = img.shape
A = random.uniform(0.45, 0.65) # 亮度
beta = random.uniform(0.03, 0.075) # 亮度
max_rc = max(row, col)
size = math.sqrt(max_rc) # 雾化尺寸
center = (row // 2, col // 2) # 雾化中心
for j in range(row):
for l in range(col):
center0 = center[0]
center1 = center[1]
d, td = fog(j, center0, l, center1, size, beta)
img_f[j][l][:] = img_f[j][l][:] * td + A * (1 - td)
out_img = (img_f * 255).astype(np.uint8)
cv2.imwrite(fog_img_path, out_img)
print("---------------------------------")
def ChangeImgName(ImgName, NewName):
old_img = cv2.imread(ImgName)
cv2.imwrite(NewName, old_img)
if __name__ == '__main__':
# demo()
path = "C:/Users/kiven/Desktop/AOD//data/hazy/"
fog_path = "C:/Users/kiven/Desktop/AOD//data/1/"
name_list = os.listdir(path)
we = "_1_1_2"
for name in name_list:
imgname = name[0:-4] + we + ".jpg"
ChangeImgName(path + name, fog_path+ imgname)
# put_fog(path+name, fog_path+name)
pth2onnx
我们新建一个.py文件对训练得到的pth权重进行转化，为了后期能够摆脱torch的安装以及可在边缘设备上运行，这里我们需要明确的是网络的输入大小以及输入输出的名字。
import torch
import torch.onnx
def pth_to_onnx(input, checkpoint, onnx_path, input_names,output_names):
if not onnx_path.endswith('.onnx'):
print('Warning!')
return 0
model = torch.load('saved_models/dehaze_net_epoch_17.pth', map_location=torch.device('cpu'))
# #指定模型的输入，以及onnx的输出路径
torch.onnx.export(model, input, onnx_path, verbose=True, input_names=input_names,
output_names=output_names) # 指定模型的输入，以及onnx的输出路径
if __name__ == '__main__':
checkpoint = './saved_models/dehaze_net_epoch_17.pth'
onnx_path = './dehaze_net.onnx'
input = torch.randn(1, 3, 450, 600)
input_names = ['input']
output_names = ['output']
pth_to_onnx(input, checkpoint, onnx_path, input_names, output_names)
ORT推理
在ORT中我们主要工作是将torch中的一些操作替换为numpy的操作：
- 采用pillow读取图像；
- 转换为numpy然后归一化0~1之间；
- 数据转换为float32格式；
- 在位置0维度上添加一个维度；
- 使用ort读取网络并推理（推理结果是一个list）
- 获取推理结果后取出结果numpy(shape= 1,3,450,600)
- 去除shape[0]的地方
- 显示结果图
import numpy as np
from PIL import Image
import onnxruntime as ort
import matplotlib.pyplot as plt
def dehaze_image(image_name):
data_hazy = Image.open(image_name)
data_hazy = np.array(data_hazy) / 255.0
original_img = data_hazy.copy()
data_hazy = np.array(data_hazy, dtype=np.float32)
data_hazy = np.transpose(data_hazy, (2, 0, 1))
data_hazy = np.expand_dims(data_hazy, 0)
bolb = data_hazy
input_name = 'input'
output_name = 'output'
dehaze_net = ort.InferenceSession("./dehaze_net.onnx", providers=ort.get_available_providers())
netOutputImg = dehaze_net.run([output_name], {input_name: bolb})
pdata = netOutputImg[0]
clean_image = pdata.squeeze()
clean_image = np.swapaxes(clean_image, 0, 1)
clean_image = np.swapaxes(clean_image, 1, 2)
plt.subplot(1, 2, 1)
plt.imshow(original_img)
plt.axis('off')
plt.title('Original Image')
plt.subplot(1, 2, 2)
plt.imshow(clean_image)
plt.axis('off')
plt.title('Dehaze Image')
plt.show()
if __name__ == '__main__':
img_name = './test_images/test0.png'
dehaze_image(img_name)
结语
虽然AOD-Net已经取得了很好的去雾效果，但是还有很多挑战需要克服，例如在复杂场景下的去雾效果提升，以及对运动模糊等噪声的抵抗能力。未来的研究将不断推动图像去雾技术的发展，希望这篇文章能为读者们对图像去雾技术的理解和应用提供一些帮助。
本文正在参加「金石计划」
