基于深度学习神经网络的AI图像PSD去雾系统源码

AI洲抿嘴的薯片

已于 2024-05-12 20:11:12 修改

阅读量1.2k

点赞数 10

分类专栏：图像画质项目文章标签：深度学习人工智能神经网络

于 2024-05-12 20:10:49 首次发布

本文链接：https://blog.csdn.net/m0_59023219/article/details/138765064

版权

图像画质项目专栏收录该内容

7 篇文章 0 订阅

订阅专栏

第一步：PSD介绍

以往的研究主要集中在具有合成模糊图像的训练模型上，当模型用于真实世界的模糊图像时，会导致性能下降。

为了解决上述问题，提高去雾的泛化性能，作者提出了一种Principled Synthetic-to-real Dehazing (PSD)框架。

本文提出的PSD适用于将现有的去雾模型推广到实际领域，包括两个阶段：有监督的预训练和无监督的微调。

预训练阶段，作者将选定的去雾模型主干修改为一个基于物理模型的网络，并用合成数据训练该网络。利用设计良好的主干，我们可以得到一个预先训练的模型，在合成域上具有良好的去雾性能。

微调阶段，作者利用真实的模糊图像以无监督的方式训练模型。

本文的贡献：

作者将真实世界的去雾任务重新定义为一个合成到真实的泛化框架：首先一个在合成配对数据上预先训练的去雾模型主干，真实的模糊图像随后将被利用以一种无监督的方式微调模型。PSD易于使用，可以以大多数深度去雾模型为骨干。
由于没有清晰的真实图像作为监督，作者利用几个流行的、有充分根据的物理先验来指导微调。作者将它们合并成一个预先的损失committee，作为具体任务的代理指导，这一部分也是PSD的核心。
性能达到SOTA

第二步：PSD网络结构

首先对两个框架大的方向做一个整体概述：

Pre-training

首先采用目前性能最好的框架之一作为网络的主干

然后我们将主干修改为一个基于物理的网络，根据一个单一的雾图同时生成干净的图像 J，传输图 t 和大气光 A，为了共同优化这三个分量，作者加入了一个重建损失，它引导网络输出服从物理散射模型。

在这个阶段，只使用标记的合成数据进行训练，最终得到一个在合成域上预训练的模型。

Fine-tuning

作者利用未标记的真实数据将预训练模型从合成域推广到真实域。受去雾强物理背景的启发，作者认为一个高质量的无雾图像应该遵循一些特定的统计规则，这些规则可以从图像先验中推导出来。此外，单一先验提供的物理知识并不总是可靠的，所以作者的目标是找到多个先验的组合，希望它们能够相互补充。

基于上述，作者设计了一个先验损失committee来作为任务特定的代理指导，用于训练未标记的真实数据。

此外，作者应用了一种learning without forgetting (LwF)的方法，该方法通过将原始任务的训练数据(即合成的模糊图像)通过网络运转到同真实的模糊数据一起，从而强行使得模型记忆合成领域的知识。

第三步：模型代码展示

import torch
import torch.nn as nn
import torch.nn.functional as F

class BlockUNet1(nn.Module):
    def __init__(self, in_channels, out_channels, upsample=False, relu=False, drop=False, bn=True):
        super(BlockUNet1, self).__init__()

        self.conv = nn.Conv2d(in_channels, out_channels, 4, 2, 1, bias=False)
        self.deconv = nn.ConvTranspose2d(in_channels, out_channels, 4, 2, 1, bias=False)

        self.dropout = nn.Dropout2d(0.5)
        self.batch = nn.InstanceNorm2d(out_channels)

        self.upsample = upsample
        self.relu = relu
        self.drop = drop
        self.bn = bn

    def forward(self, x):
        if self.relu == True:
            y = F.relu(x)
        elif self.relu == False:
            y = F.leaky_relu(x, 0.2)
        if self.upsample == True:
            y = self.deconv(y)
            if self.bn == True:
                y = self.batch(y)
            if self.drop == True:
                y = self.dropout(y)

        elif self.upsample == False:
            y = self.conv(y)
            if self.bn == True:
                y = self.batch(y)
            if self.drop == True:
                y = self.dropout(y)

        return y

class G2(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(G2, self).__init__()

        self.conv = nn.Conv2d(in_channels, 8, 4, 2, 1, bias=False)
        self.layer1 = BlockUNet1(8, 16)
        self.layer2 = BlockUNet1(16, 32)
        self.layer3 = BlockUNet1(32, 64)
        self.layer4 = BlockUNet1(64, 64)
        self.layer5 = BlockUNet1(64, 64)
        self.layer6 = BlockUNet1(64, 64)
        self.layer7 = BlockUNet1(64, 64)
        self.dlayer7 = BlockUNet1(64, 64, True, True, True, False)
        self.dlayer6 = BlockUNet1(128, 64, True, True, True)
        self.dlayer5 = BlockUNet1(128, 64, True, True, True)
        self.dlayer4 = BlockUNet1(128, 64, True, True)
        self.dlayer3 = BlockUNet1(128, 32, True, True)
        self.dlayer2 = BlockUNet1(64, 16, True, True)
        self.dlayer1 = BlockUNet1(32, 8, True, True)
        self.relu = nn.ReLU()
        self.dconv = nn.ConvTranspose2d(16, out_channels, 4, 2, 1, bias=False)
        self.lrelu = nn.LeakyReLU(0.2)

    def forward(self, x):
        y1 = self.conv(x)
        y2 = self.layer1(y1)
        y3 = self.layer2(y2)
        y4 = self.layer3(y3)
        y5 = self.layer4(y4)
        y6 = self.layer5(y5)
        y7 = self.layer6(y6)
        y8 = self.layer7(y7)

        dy8 = self.dlayer7(y8)
        concat7 = torch.cat([dy8, y7], 1)
        dy7 = self.dlayer6(concat7)
        concat6 = torch.cat([dy7, y6], 1)
        dy6 = self.dlayer5(concat6)
        concat5 = torch.cat([dy6, y5], 1)
        dy5 = self.dlayer4(concat5)
        concat4 = torch.cat([dy5, y4], 1)
        dy4 = self.dlayer3(concat4)
        concat3 = torch.cat([dy4, y3], 1)
        dy3 = self.dlayer2(concat3)
        concat2 = torch.cat([dy3, y2], 1)
        dy2 = self.dlayer1(concat2)
        concat1 = torch.cat([dy2, y1], 1)
        out = self.relu(concat1)
        out = self.dconv(out)
        out = self.lrelu(out)

        return F.avg_pool2d(out, (out.shape[2], out.shape[3]))
    
def default_conv(in_channels, out_channels, kernel_size, bias=True):
    return nn.Conv2d(in_channels, out_channels, kernel_size,padding=(kernel_size//2), bias=bias)
    
class PALayer(nn.Module):
    def __init__(self, channel):
        super(PALayer, self).__init__()
        self.pa = nn.Sequential(
                nn.Conv2d(channel, channel // 8, 1, padding=0, bias=True),
                nn.ReLU(inplace=True),
                nn.Conv2d(channel // 8, 1, 1, padding=0, bias=True),
                nn.Sigmoid()
        )
    def forward(self, x):
        y = self.pa(x)
        return x * y

class CALayer(nn.Module):
    def __init__(self, channel):
        super(CALayer, self).__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.ca = nn.Sequential(
                nn.Conv2d(channel, channel // 8, 1, padding=0, bias=True),
                nn.ReLU(inplace=True),
                nn.Conv2d(channel // 8, channel, 1, padding=0, bias=True),
                nn.Sigmoid()
        )

    def forward(self, x):
        y = self.avg_pool(x)
        y = self.ca(y)
        return x * y

class Block(nn.Module):
    def __init__(self, conv, dim, kernel_size,):
        super(Block, self).__init__()
        self.conv1=conv(dim, dim, kernel_size, bias=True)
        self.act1=nn.ReLU(inplace=True)
        self.conv2=conv(dim,dim,kernel_size,bias=True)
        self.calayer=CALayer(dim)
        self.palayer=PALayer(dim)
    def forward(self, x):
        res=self.act1(self.conv1(x))
        res=res+x 
        res=self.conv2(res)
        res=self.calayer(res)
        res=self.palayer(res)
        res += x 
        return res
class Group(nn.Module):
    def __init__(self, conv, dim, kernel_size, blocks):
        super(Group, self).__init__()
        modules = [ Block(conv, dim, kernel_size)  for _ in range(blocks)]
        modules.append(conv(dim, dim, kernel_size))
        self.gp = nn.Sequential(*modules)
    def forward(self, x):
        res = self.gp(x)
        res += x
        return res

class FFANet(nn.Module):
    def __init__(self,gps,blocks,conv=default_conv):
        super(FFANet, self).__init__()
        self.gps=gps
        self.dim=64
        kernel_size=3
        pre_process = [conv(3, self.dim, kernel_size)]
        assert self.gps==3
        self.g1= Group(conv, self.dim, kernel_size,blocks=blocks)
        self.g2= Group(conv, self.dim, kernel_size,blocks=blocks)
        self.g3= Group(conv, self.dim, kernel_size,blocks=blocks)
        self.ca=nn.Sequential(*[
            nn.AdaptiveAvgPool2d(1),
            nn.Conv2d(self.dim*self.gps,self.dim//16,1,padding=0),
            nn.ReLU(inplace=True),
            nn.Conv2d(self.dim//16, self.dim*self.gps, 1, padding=0, bias=True),
            nn.Sigmoid()
            ])
        self.palayer=PALayer(self.dim)

        self.conv_J_1 = nn.Conv2d(64, 64, 3, 1, 1, bias=False)
        self.conv_J_2 = nn.Conv2d(64, 3, 3, 1, 1, bias=False)
        self.conv_T_1 = nn.Conv2d(64, 16, 3, 1, 1, bias=False)
        self.conv_T_2 = nn.Conv2d(16, 1, 3, 1, 1, bias=False)
        
        post_precess = [
            conv(self.dim, self.dim, kernel_size),
            conv(self.dim, 3, kernel_size)]

        self.pre = nn.Sequential(*pre_process)
        self.post = nn.Sequential(*post_precess)
        self.ANet = G2(3, 3)

    def forward(self, x1, x2=0, Val=False):
        x = self.pre(x1)
        res1=self.g1(x)
        res2=self.g2(res1)
        res3=self.g3(res2)
        w=self.ca(torch.cat([res1,res2,res3],dim=1))
        w=w.view(-1,self.gps,self.dim)[:,:,:,None,None]
        out=w[:,0,::]*res1+w[:,1,::]*res2+w[:,2,::]*res3
        out=self.palayer(out)
        
        out_J = self.conv_J_1(out)
        out_J = self.conv_J_2(out_J)
        out_J = out_J + x1
        out_T = self.conv_T_1(out)
        out_T = self.conv_T_2(out_T)
        
        if Val == False:
            out_A = self.ANet(x1)
        else:
            out_A = self.ANet(x2)
            
        out_I = out_T * out_J + (1 - out_T) * out_A
        #x=self.post(out)
        return out, out_J, out_T, out_A, out_I
    
if __name__ == "__main__":
    net=FFA(gps=3,blocks=19)
    print(net)

第四步：运行