PyTorch框架——基于深度学习EfficientDeRain神经网络AI去雨滴图像增强系统

第一步：EfficientDeRain介绍

        EfficientDeRain 是一个针对单张图像去雨的开源项目，该项目由清华大学的研究团队提出，主要用于处理图像中的雨水干扰，恢复图像的真实场景

核心功能

        图像去雨：EfficientDeRain 通过学习像素级的膨胀滤波，有效去除图像中的雨水干扰，恢复清晰图像。
        高效率：项目设计考虑到了效率，能够在较短的时间内处理大量图像，适用于需要快速处理的应用场景。
        可扩展性：项目提供了多种数据集的预训练模型，支持自定义数据集的训练，方便用户根据具体需求进行优化。

第二步：LYT-Net网络结构

        该算法的原理非常简单，最重要的思想是把去雨看为图像的逐像素滤波问题。而滤波操作是高度优化的操作，在GPU上的实现必定非常快。

看懂下面这张图，即可完全理解作者的算法思想：

        图像经深度卷积网络学习逐像素的卷积核参数，然后与原图做卷积即得最终的去雨后图像，训练的时候需要（有雨、无雨）的图像对。

        作者指出，尽管上述思想没有问题，但因为逐像素卷积核大小的问题，如果只学习普通卷积核（即每个像素预测三个通道的3x(3x3)个参数）如上图中的（a）部分，对于雨条较大的图像很难取得满意的效果，因为毕竟卷积的过程是寻找周围非雨条像素赋以高权重的加权，卷积核如果没有覆盖到非雨条像素，肯定效果不好。

        为在尺度上应对大雨条，所以作者做了改进，让神经网络预测多尺度的空洞卷积核，如（b）子图中是预测4个尺度的空洞卷积核，空洞卷积后再把结果加权，获得最终的去雨图像。

        所以算法的核心思路可总结为：学习多尺度空洞卷积+图像加权融合

第三步：模型代码展示

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np# ----------------------------------------
#         Initialize the networks
# ----------------------------------------
def weights_init(net, init_type = 'normal', init_gain = 0.02):"""Initialize network weights.Parameters:net (network)   -- network to be initializedinit_type (str) -- the name of an initialization method: normal | xavier | kaiming | orthogonalinit_gain (float)    -- scaling factor for normal, xavier and orthogonalIn our paper, we choose the default setting: zero mean Gaussian distribution with a standard deviation of 0.02"""def init_func(m):classname = m.__class__.__name__if hasattr(m, 'weight') and classname.find('Conv') != -1:if init_type == 'normal':torch.nn.init.normal_(m.weight.data, 0.0, init_gain)elif init_type == 'xavier':torch.nn.init.xavier_normal_(m.weight.data, gain = init_gain)elif init_type == 'kaiming':torch.nn.init.kaiming_normal_(m.weight.data, a = 0, mode = 'fan_in')elif init_type == 'orthogonal':torch.nn.init.orthogonal_(m.weight.data, gain = init_gain)else:raise NotImplementedError('initialization method [%s] is not implemented' % init_type)elif classname.find('BatchNorm2d') != -1:torch.nn.init.normal_(m.weight.data, 1.0, 0.02)torch.nn.init.constant_(m.bias.data, 0.0)# apply the initialization function <init_func>print('initialize network with %s type' % init_type)net.apply(init_func)# ----------------------------------------
#      Kernel Prediction Network (KPN)
# ----------------------------------------
class Basic(nn.Module):def __init__(self, in_ch, out_ch, g=16, channel_att=False, spatial_att=False):super(Basic, self).__init__()self.channel_att = channel_attself.spatial_att = spatial_attself.conv1 = nn.Sequential(nn.Conv2d(in_channels=in_ch, out_channels=out_ch, kernel_size=3, stride=1, padding=1),# nn.BatchNorm2d(out_ch),nn.ReLU(),nn.Conv2d(in_channels=out_ch, out_channels=out_ch, kernel_size=3, stride=1, padding=1),# nn.BatchNorm2d(out_ch),nn.ReLU(),nn.Conv2d(in_channels=out_ch, out_channels=out_ch, kernel_size=3, stride=1, padding=1),# nn.BatchNorm2d(out_ch),nn.ReLU())if channel_att:self.att_c = nn.Sequential(nn.Conv2d(2*out_ch, out_ch//g, 1, 1, 0),nn.ReLU(),nn.Conv2d(out_ch//g, out_ch, 1, 1, 0),nn.Sigmoid())if spatial_att:self.att_s = nn.Sequential(nn.Conv2d(in_channels=2, out_channels=1, kernel_size=7, stride=1, padding=3),nn.Sigmoid())def forward(self, data):"""Forward function.:param data::return: tensor"""fm = self.conv1(data)if self.channel_att:# fm_pool = F.adaptive_avg_pool2d(fm, (1, 1)) + F.adaptive_max_pool2d(fm, (1, 1))fm_pool = torch.cat([F.adaptive_avg_pool2d(fm, (1, 1)), F.adaptive_max_pool2d(fm, (1, 1))], dim=1)att = self.att_c(fm_pool)fm = fm * attif self.spatial_att:fm_pool = torch.cat([torch.mean(fm, dim=1, keepdim=True), torch.max(fm, dim=1, keepdim=True)[0]], dim=1)att = self.att_s(fm_pool)fm = fm * attreturn fmclass KPN(nn.Module):def __init__(self, color=True, burst_length=1, blind_est=True, kernel_size=[5], sep_conv=False,channel_att=False, spatial_att=False, upMode='bilinear', core_bias=False):super(KPN, self).__init__()self.upMode = upModeself.burst_length = burst_lengthself.core_bias = core_biasself.color_channel = 3 if color else 1in_channel = (3 if color else 1) * (burst_length if blind_est else burst_length+1)out_channel = (3 if color else 1) * (2 * sum(kernel_size) if sep_conv else np.sum(np.array(kernel_size) ** 2)) * burst_lengthif core_bias:out_channel += (3 if color else 1) * burst_length# 各个卷积层定义# 2~5层都是均值池化+3层卷积self.conv1 = Basic(in_channel, 64, channel_att=False, spatial_att=False)self.conv2 = Basic(64, 128, channel_att=False, spatial_att=False)self.conv3 = Basic(128, 256, channel_att=False, spatial_att=False)self.conv4 = Basic(256, 512, channel_att=False, spatial_att=False)self.conv5 = Basic(512, 512, channel_att=False, spatial_att=False)# 6~8层要先上采样再卷积self.conv6 = Basic(512+512, 512, channel_att=channel_att, spatial_att=spatial_att)self.conv7 = Basic(256+512, 256, channel_att=channel_att, spatial_att=spatial_att)self.conv8 = Basic(256+128, out_channel, channel_att=channel_att, spatial_att=spatial_att)self.outc = nn.Conv2d(out_channel, out_channel, 1, 1, 0)self.kernel_pred = KernelConv(kernel_size, sep_conv, self.core_bias)self.conv_final = nn.Conv2d(in_channels=12, out_channels=3, kernel_size=3, stride=1, padding=1)# 前向传播函数def forward(self, data_with_est, data, white_level=1.0):"""forward and obtain pred image directly:param data_with_est: if not blind estimation, it is same as data:param data::return: pred_img_i and img_pred"""conv1 = self.conv1(data_with_est)conv2 = self.conv2(F.avg_pool2d(conv1, kernel_size=2, stride=2))conv3 = self.conv3(F.avg_pool2d(conv2, kernel_size=2, stride=2))conv4 = self.conv4(F.avg_pool2d(conv3, kernel_size=2, stride=2))conv5 = self.conv5(F.avg_pool2d(conv4, kernel_size=2, stride=2))# 开始上采样  同时要进行skip connectionconv6 = self.conv6(torch.cat([conv4, F.interpolate(conv5, scale_factor=2, mode=self.upMode)], dim=1))conv7 = self.conv7(torch.cat([conv3, F.interpolate(conv6, scale_factor=2, mode=self.upMode)], dim=1))#print(conv7.size())conv8 = self.conv8(torch.cat([conv2, F.interpolate(conv7, scale_factor=2, mode=self.upMode)], dim=1))# return channel K*K*Ncore = self.outc(F.interpolate(conv8, scale_factor=2, mode=self.upMode))pred1 = self.kernel_pred(data, core, white_level, rate=1)pred2 = self.kernel_pred(data, core, white_level, rate=2)pred3 = self.kernel_pred(data, core, white_level, rate=3)pred4 = self.kernel_pred(data, core, white_level, rate=4)pred_cat = torch.cat([torch.cat([torch.cat([pred1, pred2], dim=1), pred3], dim=1), pred4], dim=1)pred = self.conv_final(pred_cat)#pred = self.kernel_pred(data, core, white_level, rate=1)return predclass KernelConv(nn.Module):"""the class of computing prediction"""def __init__(self, kernel_size=[5], sep_conv=False, core_bias=False):super(KernelConv, self).__init__()self.kernel_size = sorted(kernel_size)self.sep_conv = sep_convself.core_bias = core_biasdef _sep_conv_core(self, core, batch_size, N, color, height, width):"""convert the sep_conv core to conv2d core2p --> p^2:param core: shape: batch*(N*2*K)*height*width:return:"""kernel_total = sum(self.kernel_size)core = core.view(batch_size, N, -1, color, height, width)if not self.core_bias:core_1, core_2 = torch.split(core, kernel_total, dim=2)else:core_1, core_2, core_3 = torch.split(core, kernel_total, dim=2)# output corecore_out = {}cur = 0for K in self.kernel_size:t1 = core_1[:, :, cur:cur + K, ...].view(batch_size, N, K, 1, 3, height, width)t2 = core_2[:, :, cur:cur + K, ...].view(batch_size, N, 1, K, 3, height, width)core_out[K] = torch.einsum('ijklno,ijlmno->ijkmno', [t1, t2]).view(batch_size, N, K * K, color, height, width)cur += K# it is a dictreturn core_out, None if not self.core_bias else core_3.squeeze()def _convert_dict(self, core, batch_size, N, color, height, width):"""make sure the core to be a dict, generally, only one kind of kernel size is suitable for the func.:param core: shape: batch_size*(N*K*K)*height*width:return: core_out, a dict"""core_out = {}core = core.view(batch_size, N, -1, color, height, width)core_out[self.kernel_size[0]] = core[:, :, 0:self.kernel_size[0]**2, ...]bias = None if not self.core_bias else core[:, :, -1, ...]return core_out, biasdef forward(self, frames, core, white_level=1.0, rate=1):"""compute the pred image according to core and frames:param frames: [batch_size, N, 3, height, width]:param core: [batch_size, N, dict(kernel), 3, height, width]:return:"""if len(frames.size()) == 5:batch_size, N, color, height, width = frames.size()else:batch_size, N, height, width = frames.size()color = 1frames = frames.view(batch_size, N, color, height, width)if self.sep_conv:core, bias = self._sep_conv_core(core, batch_size, N, color, height, width)else:core, bias = self._convert_dict(core, batch_size, N, color, height, width)img_stack = []pred_img = []kernel = self.kernel_size[::-1]for index, K in enumerate(kernel):if not img_stack:padding_num = (K//2) * rateframe_pad = F.pad(frames, [padding_num, padding_num, padding_num, padding_num])for i in range(0, K):for j in range(0, K):img_stack.append(frame_pad[..., i*rate:i*rate + height, j*rate:j*rate + width])img_stack = torch.stack(img_stack, dim=2)else:k_diff = (kernel[index - 1] - kernel[index]) // 2img_stack = img_stack[:, :, k_diff:-k_diff, ...]# print('img_stack:', img_stack.size())pred_img.append(torch.sum(core[K].mul(img_stack), dim=2, keepdim=False))pred_img = torch.stack(pred_img, dim=0)# print('pred_stack:', pred_img.size())pred_img_i = torch.mean(pred_img, dim=0, keepdim=False)#print("pred_img_i", pred_img_i.size())# N = 1pred_img_i = pred_img_i.squeeze(2)#print("pred_img_i", pred_img_i.size())# if bias is permittedif self.core_bias:if bias is None:raise ValueError('The bias should not be None.')pred_img_i += bias# print('white_level', white_level.size())pred_img_i = pred_img_i / white_level#pred_img = torch.mean(pred_img_i, dim=1, keepdim=True)# print('pred_img:', pred_img.size())# print('pred_img_i:', pred_img_i.size())return pred_img_iclass LossFunc(nn.Module):"""loss function of KPN"""def __init__(self, coeff_basic=1.0, coeff_anneal=1.0, gradient_L1=True, alpha=0.9998, beta=100):super(LossFunc, self).__init__()self.coeff_basic = coeff_basicself.coeff_anneal = coeff_annealself.loss_basic = LossBasic(gradient_L1)self.loss_anneal = LossAnneal(alpha, beta)def forward(self, pred_img_i, pred_img, ground_truth, global_step):"""forward function of loss_func:param frames: frame_1 ~ frame_N, shape: [batch, N, 3, height, width]:param core: a dict coverted by ......:param ground_truth: shape [batch, 3, height, width]:param global_step: int:return: loss"""return self.coeff_basic * self.loss_basic(pred_img, ground_truth), self.coeff_anneal * self.loss_anneal(global_step, pred_img_i, ground_truth)class LossBasic(nn.Module):"""Basic loss function."""def __init__(self, gradient_L1=True):super(LossBasic, self).__init__()self.l1_loss = nn.L1Loss()self.l2_loss = nn.MSELoss()self.gradient = TensorGradient(gradient_L1)def forward(self, pred, ground_truth):return self.l2_loss(pred, ground_truth) + \self.l1_loss(self.gradient(pred), self.gradient(ground_truth))class LossAnneal(nn.Module):"""anneal loss function"""def __init__(self, alpha=0.9998, beta=100):super(LossAnneal, self).__init__()self.global_step = 0self.loss_func = LossBasic(gradient_L1=True)self.alpha = alphaself.beta = betadef forward(self, global_step, pred_i, ground_truth):""":param global_step: int:param pred_i: [batch_size, N, 3, height, width]:param ground_truth: [batch_size, 3, height, width]:return:"""loss = 0for i in range(pred_i.size(1)):loss += self.loss_func(pred_i[:, i, ...], ground_truth)loss /= pred_i.size(1)return self.beta * self.alpha ** global_step * lossclass TensorGradient(nn.Module):"""the gradient of tensor"""def __init__(self, L1=True):super(TensorGradient, self).__init__()self.L1 = L1def forward(self, img):w, h = img.size(-2), img.size(-1)l = F.pad(img, [1, 0, 0, 0])r = F.pad(img, [0, 1, 0, 0])u = F.pad(img, [0, 0, 1, 0])d = F.pad(img, [0, 0, 0, 1])if self.L1:return torch.abs((l - r)[..., 0:w, 0:h]) + torch.abs((u - d)[..., 0:w, 0:h])else:return torch.sqrt(torch.pow((l - r)[..., 0:w, 0:h], 2) + torch.pow((u - d)[..., 0:w, 0:h], 2))if __name__ == '__main__':kpn = KPN().cuda()a = torch.randn(4, 3, 224, 224).cuda()b = kpn(a, a)print(b.shape)

第四步：运行

第五步：整个工程的内容

 项目完整文件下载请见演示与介绍视频的简介处给出：➷➷➷

PyTorch框架——基于深度学习EfficientDeRain神经网络AI去雨滴图像增强系统_哔哩哔哩_bilibili