DehazeNet单图像去雾的端到端系统（学习笔记）

1、图像去雾算法分类：图像去雾分为图像恢复和图像增强，在图像恢复中又分为单图像去雾和多图像去雾，图像去雾中使用的物理模型是大气散射模型，可参考暗通道先验中对大气散射模型的解释。

2、与雾相关的特征

（1）暗通道先验：清晰图像块的RGB颜色空间中有一个通道很暗（数值很低甚至接近于零）。暗通道为局部区域中所有像素颜色的最小值:：

暗通道先验可直接用于估计透射率t(x)，t(x)正比于1-D(x)。

（2）最大化对比度：根据大气散射模型，图像的对比度因雾度而降低为，基于该观察，局部对比度是s×s局部pitch中像素强度相对于中心像素的方差，局部最大值为r×r区域r中的局部对比度值定义为C(x)：

通过最大化局部对比度来增强图像的可见性。

（3）颜色衰减：雾会导致图像饱和度的降低和亮度的增加，整体上表现为颜色的衰减。根据颜色衰减先验，亮度和饱和度的差值记为A(x)，景深和A(x)成正比，所以A(x)可以用于估计透射率 ：

 （4）色调差异：原始图像I(x)与其半逆图像(x)之间的色调差异已用于检测雾度，

对于无雾图像，其半逆图像的三个通道中的像素值不会全部翻转，从而导致(x)和I(x)之间的色调变化较大。色调差异特征被定义为：

上标“h”表示HSV颜色空间中图像的色调通道，透射率t(x)向H(x)反向传播。

3、DeHazeNet主要贡献

（1）F1层：Maxout激活函数学习与雾相关的特征；

（2）F4层：用BReLu替代ReLu和Sigmoid函数，Sigmoid函数会出现梯度消失问题，导致收敛速度慢，ReLu函数更好的适用于分类问题而非回归问题，BReLu可以保持双边约束和局部线性来进行图像恢复，提高收敛性。

（3）端到端系统：输入是有雾图像，输出是其对应的透射率。

4、网络结构

（1）F1特征提取层，即提取有雾图像特征。根据不同的假设与先验设计不同的滤波器。举的例子中有16个滤波器。其中每四个是上述一种先验特征滤波器。通过maxout unit的激活函数，每四个输出一张图。这里不padding，输入是3*16*16三通道的块。输出的是四个16*12*12，每一个代表一种特征。

（2）F2使用多尺度的平行卷积操作。由于多尺度特征被证明有利于去雾并且在inception的模型中也用到了平行卷积，即同一张图用不同尺度的卷积核进行卷积。分别用16个3*3、16个5*5和16个7*7的卷积核进行卷积，每一种尺度产生16个，并且通过padding每张图大小应该是一致的。总共获得48个48*10*10。

（3）F3Maxpooling对局部数据敏感，另外根据假设透射率有局部不变性，所以用一个7*7局部最大值滤波替代maxpooling。输出是48个48*6*6。

（4）通过1个4*4的卷积核，产生1*1的标量，并且使用的激活函数为BReLU。因为ReLU抑制了小于0的数，只适用于图像分类等方面，并不适合图像复原。因为最后的透射率图允许高于1或者低于0。所以提出了BReLU，既保持了局部线性，又保持了双边的限制。输出的是一个标量，即输入块中心点的透射率值。

5、与传统去雾方法的联系：DehazeNet的第一层特征F1设计用于有雾图像的特征提取。以暗通道先验为例，如果权重W1是一个相反的滤波器（在一个通道的中心有值为−1的稀疏矩阵），而B1是一个单位偏差，那么特征图的最大输出相当于颜色通道的最小输出，这类似于暗通道。同样，当权重为圆形滤波器时，F1与最大对比度相似；当W1包含全通滤波器和相反滤波器时，F1与最大和最小特征图相似，这是颜色空间从RGB到HSV转换的运算，然后提取颜色衰减和色调色差特征。

综上所述，在如上图所示的滤波器学习成功后，第2节中提到的与雾相关的特征可以从DehazeNet的第一层中提取出来。另一方面，Maxout激活函数可以看作是对任意凸函数的分段线性逼近。在本文中，在四个特征映射(k=4)中选择最大值来近似一个任意的凸函数。

6、训练细节：

（1）使用深度学习的方法去雾需要有ground truth，由于自然场景中有雾图像和无雾图像无法同时存在，所以使用合成数据集，故此算法对自然场景中的有雾图像效果不太好。

（2）合成数据过程：给定一个无雾图像J(x)、大气光α和一个随机透射率t∈（0,1），合成一个模糊图像为I(x)=J(x)t+α(1−t)。为了减少变量学习中的不确定性，将大气光α设置为1。

（3）损失函数：MSE损失函数，随机梯度下降

（4）图像去雾：网络训练完成之后，得到初始透射率，再通过引导滤波进行细化，然后根据大气散射模型复原图像。

7、待改进的地方

（1）把大气光α当成了全局常量，所以这个算法在雾度均匀的情况下比较好，在不均匀雾度下效果不太好；

（2）有雾图像和无雾图像之间可以直接进行端到端映射，而不用估计透射率。

8、python代码实现

import torch
import torch.nn as nn
from torch.utils.data.dataset import Dataset
from PIL import Image
import torchvision
from torchvision import transforms
import torch.utils.data as data
#import torchsnooper
import cv2BATCH_SIZE = 128
EPOCH = 10# BRelu used for GPU. Need to add that reference in pytorch source file.
class BRelu(nn.Hardtanh):def __init__(self, inplace=False):super(BRelu, self).__init__(0., 1., inplace)def extra_repr(self):inplace_str = 'inplace=True' if self.inplace else ''return inplace_strclass DehazeNet(nn.Module):def __init__(self, input=16, groups=4):super(DehazeNet, self).__init__()self.input = inputself.groups = groupsself.conv1 = nn.Conv2d(in_channels=3, out_channels=self.input, kernel_size=5)self.conv2 = nn.Conv2d(in_channels=4, out_channels=16, kernel_size=3, padding=1)self.conv3 = nn.Conv2d(in_channels=4, out_channels=16, kernel_size=5, padding=2)self.conv4 = nn.Conv2d(in_channels=4, out_channels=16, kernel_size=7, padding=3)self.maxpool = nn.MaxPool2d(kernel_size=7, stride=1)self.conv5 = nn.Conv2d(in_channels=48, out_channels=1, kernel_size=6)#self.brelu = nn.BReLU()for name, m in self.named_modules():# lambda : 定义简单的函数    lambda x: 表达式# map(func, iter)  iter 依次调用 func# any : 有一个是true就返回trueif isinstance(m, nn.Conv2d):# 初始化 weight 和 biasnn.init.normal(m.weight, mean=0,std=0.001)if m.bias is not None:nn.init.constant_(m.bias, 0)def Maxout(self, x, groups):x = x.reshape(x.shape[0], groups, x.shape[1]//groups, x.shape[2], x.shape[3])x, y = torch.max(x, dim=2, keepdim=True)out = x.reshape(x.shape[0],-1, x.shape[3], x.shape[4])return out#BRelu used to CPU. It can't work on GPU.def BRelu(self, x):x = torch.max(x, torch.zeros(x.shape[0],x.shape[1],x.shape[2],x.shape[3]))x = torch.min(x, torch.ones(x.shape[0],x.shape[1],x.shape[2],x.shape[3]))return xdef forward(self, x):out = self.conv1(x)out = self.Maxout(out, self.groups)out1 = self.conv2(out)out2 = self.conv3(out)out3 = self.conv4(out)y = torch.cat((out1,out2,out3), dim=1)#print(y.shape[0],y.shape[1],y.shape[2],y.shape[3],)y = self.maxpool(y)#print(y.shape[0],y.shape[1],y.shape[2],y.shape[3],)y = self.conv5(y)# y = self.relu(y)# y = self.BRelu(y)#y = torch.min(y, torch.ones(y.shape[0],y.shape[1],y.shape[2],y.shape[3]))y = self.BRelu(y)y = y.reshape(y.shape[0],-1)return yloader = torchvision.transforms.Compose([transforms.ToTensor(),transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
augmentation = torchvision.transforms.Compose([transforms.RandomHorizontalFlip(0.5),transforms.RandomVerticalFlip(0.5),transforms.RandomRotation(30),transforms.ToTensor(),transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])class FogData(Dataset):# root：图像存放地址根路径# augment：是否需要图像增强def __init__(self, root, labels, augment=True):# 初始化 可以定义图片地址 标签 是否变换 变换函数self.image_files = rootself.labels = torch.cuda.FloatTensor(labels)self.augment = augment   # 是否需要图像增强# self.transform = transformdef __getitem__(self, index):# 读取图像数据并返回if self.augment:img = Image.open(self.image_files[index])img = augmentation(img)img = img.cuda()return img, self.labels[index]else:img = Image.open(self.image_files[index])img = loader(img)img = img.cuda()return img, self.labels[index]def __len__(self):# 返回图像的数量return len(self.image_files)path_train = []
file = open('path_train.txt', mode='r')
content = file.readlines()
for i in range(len(content)):path_train.append(content[i][:-1])label_train = []
file = open('label_train.txt', mode='r')
content = file.readlines()
for i in range(len(content)):label_train.append(float(content[i][:-1]))#print(float(content[i][:-1]))train_data = FogData(path_train, label_train, False)
train_loader = data.DataLoader(dataset=train_data, batch_size=BATCH_SIZE, shuffle=True, )net = DehazeNet()
net.load_state_dict(torch.load(r'defog4_noaug.pth', map_location='cpu'))#@torchsnooper.snoop()
def train():lr = 0.00001optimizer = torch.optim.Adam(net.parameters(), lr=0.0000005)loss_func = nn.MSELoss().cuda()for epoch in range(EPOCH):total_loss = 0for i, (x, y) in enumerate(train_loader):# 输入训练数据# 清空上一次梯度optimizer.zero_grad()output = net(x)# 计算误差loss = loss_func(output, y)total_loss = total_loss+loss# 误差反向传递loss.backward()# 优化器参数更新optimizer.step()if i % 10 == 5:print('Epoch', epoch, '|step ', i, 'loss: %.4f' % loss.item(), )print('Epoch', epoch, 'total_loss', total_loss.item())torch.save(net.state_dict(), r'defog4_noaug.pth')#train()def defog(pic_dir):img = Image.open(pic_dir)img1 = loader(img)img2 = transforms.ToTensor()(img)c, h, w = img1.shapepatch_size = 16num_w = int(w / patch_size)num_h = int(h / patch_size)t_list = []for i in range(0, num_w):for j in range(0, num_h):patch = img1[:, 0 + j * patch_size:patch_size + j * patch_size,0 + i * patch_size:patch_size + i * patch_size]patch = torch.unsqueeze(patch, dim=0)t = net(patch)t_list.append([i,j,t])t_list = sorted(t_list, key=lambda t_list:t_list[2])a_list = t_list[:len(t_list)//100]a0 = 0for k in range(0,len(a_list)):patch = img2[:, 0 + a_list[k][1] * patch_size:patch_size + a_list[k][1] * patch_size,0 + a_list[k][0] * patch_size:patch_size + a_list[k][0] * patch_size]a = torch.max(patch)if a0 < a.item():a0 = a.item()for k in range(0,len(t_list)):img2[:, 0 + t_list[k][1] * patch_size:patch_size + t_list[k][1] * patch_size,0 + t_list[k][0] * patch_size:patch_size + t_list[k][0] * patch_size] = (img2[:,0 + t_list[k][1] * patch_size:patch_size + t_list[k][1] * patch_size,0 + t_list[k][0] * patch_size:patch_size + t_list[k][0] * patch_size] - a0*(1-t_list[k][2]))/t_list[k][2]defog_img = transforms.ToPILImage()(img2)defog_img.save('./test21-1.jpg')defog('./21-1.jpg')