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 cv2

BATCH_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_str


class DehazeNet(nn.Module):
	def __init__(self, input=16, groups=4):
		super(DehazeNet, self).__init__()
		self.input = input
		self.groups = groups
		self.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就返回true
			if isinstance(m, nn.Conv2d):
				# 初始化 weight 和 bias
				nn.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 x
	
	def 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 y


loader = 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 = root
		self.labels = torch.cuda.FloatTensor(labels)
		self.augment = augment   # 是否需要图像增强
		# self.transform = transform

	def __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.00001
	optimizer = torch.optim.Adam(net.parameters(), lr=0.0000005)
	loss_func = nn.MSELoss().cuda()
	for epoch in range(EPOCH):
		total_loss = 0
		for 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.shape
	patch_size = 16
	num_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 = 0
	for 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')