基于深度学习的图像去雨去雾

基于深度学习的图像去雨去雾

文末附有源码下载地址
b站视频地址： https://www.bilibili.com/video/BV1Jr421p7cT/

基于深度学习的图像去雨去雾，使用的网络为unet，
网络代码：

import torch
import torch.nn as nn
from torchsummary import summary
from torchvision import models
from torchvision.models.feature_extraction import create_feature_extractor
import torch.nn.functional as F
from torchstat import statclass Resnet18(nn.Module):def __init__(self):super(Resnet18, self).__init__()self.resnet = models.resnet18(pretrained=False)# self.resnet = create_feature_extractor(self.resnet, {'relu': 'feat320', 'layer1': 'feat160', 'layer2': 'feat80',#                                                'layer3': 'feat40'})def forward(self,x):for name,m in self.resnet._modules.items():x=m(x)if name=='relu':x1=xelif name=='layer1':x2=xelif name=='layer2':x3=xelif name=='layer3':x4=xbreak# x=self.resnet(x)return x1,x2,x3,x4
class Linears(nn.Module):def __init__(self,a,b):super(Linears, self).__init__()self.linear1=nn.Linear(a,b)self.relu1=nn.LeakyReLU()self.linear2 = nn.Linear(b, a)self.sigmoid=nn.Sigmoid()def forward(self,x):x=self.linear1(x)x=self.relu1(x)x=self.linear2(x)x=self.sigmoid(x)return x
class DenseNetBlock(nn.Module):def __init__(self,inplanes=1,planes=1,stride=1):super(DenseNetBlock,self).__init__()self.conv1=nn.Conv2d(inplanes,planes,3,stride,1)self.bn1 = nn.BatchNorm2d(planes)self.relu1=nn.LeakyReLU()self.conv2 = nn.Conv2d(inplanes, planes, 3,stride,1)self.bn2 = nn.BatchNorm2d(planes)self.relu2 = nn.LeakyReLU()self.conv3 = nn.Conv2d(inplanes, planes, 3,stride,1)self.bn3 = nn.BatchNorm2d(planes)self.relu3 = nn.LeakyReLU()def forward(self,x):ins=xx=self.conv1(x)x=self.bn1(x)x=self.relu1(x)x = self.conv2(x)x = self.bn2(x)x = self.relu2(x)x=x+insx2=self.conv3(x)x2 = self.bn3(x2)x2=self.relu3(x2)out=ins+x+x2return out
class SEnet(nn.Module):def __init__(self,chs,reduction=4):super(SEnet,self).__init__()self.average_pooling = nn.AdaptiveAvgPool2d(output_size=(1, 1))self.fc = nn.Sequential(# First reduce dimension, then raise dimension.# Add nonlinear processing to fit the correlation between channelsnn.Linear(chs, chs // reduction),nn.LeakyReLU(inplace=True),nn.Linear(chs // reduction, chs))self.activation = nn.Sigmoid()def forward(self,x):ins=xbatch_size, chs, h, w = x.shapex=self.average_pooling(x)x = x.view(batch_size, chs)x=self.fc(x)x = x.view(batch_size,chs,1,1)return x*ins
class UAFM(nn.Module):def __init__(self):super(UAFM, self).__init__()# self.meanPool_C=torch.max()self.attention=nn.Sequential(nn.Conv2d(4, 8, 3, 1,1),nn.LeakyReLU(),nn.Conv2d(8, 1, 1, 1),nn.Sigmoid())def forward(self,x1,x2):x1_mean_pool=torch.mean(x1,dim=1)x1_max_pool,_=torch.max(x1,dim=1)x2_mean_pool = torch.mean(x2, dim=1)x2_max_pool,_ = torch.max(x2, dim=1)x1_mean_pool=torch.unsqueeze(x1_mean_pool,dim=1)x1_max_pool=torch.unsqueeze(x1_max_pool,dim=1)x2_mean_pool=torch.unsqueeze(x2_mean_pool,dim=1)x2_max_pool=torch.unsqueeze(x2_max_pool,dim=1)cat=torch.cat((x1_mean_pool,x1_max_pool,x2_mean_pool,x2_max_pool),dim=1)a=self.attention(cat)out=x1*a+x2*(1-a)return outclass Net(nn.Module):def __init__(self):super(Net, self).__init__()self.resnet18=Resnet18()self.SENet=SEnet(chs=256)self.UAFM=UAFM()self.DenseNet1=DenseNetBlock(inplanes=256,planes=256)self.transConv1=nn.ConvTranspose2d(256,128,3,2,1,output_padding=1)self.DenseNet2 = DenseNetBlock(inplanes=128, planes=128)self.transConv2 = nn.ConvTranspose2d(128, 64, 3, 2, 1, output_padding=1)self.DenseNet3 = DenseNetBlock(inplanes=64, planes=64)self.transConv3 = nn.ConvTranspose2d(64, 64, 3, 2, 1, output_padding=1)self.transConv4 = nn.ConvTranspose2d(64, 32, 3, 2, 1, output_padding=1)self.DenseNet4=DenseNetBlock(inplanes=32,planes=32)self.out=nn.Sequential(nn.Conv2d(32,3,1,1),nn.Sigmoid())def forward(self,x):"""下采样部分"""x1,x2,x3,x4=self.resnet18(x)# feat320=features['feat320']# feat160=features['feat160']# feat80=features['feat80']# feat40=features['feat40']feat320=x1feat160=x2feat80=x3feat40=x4"""上采样部分"""x=self.SENet(feat40)x=self.DenseNet1(x)x=self.transConv1(x)x=self.UAFM(x,feat80)x=self.DenseNet2(x)x=self.transConv2(x)x=self.UAFM(x,feat160)x = self.DenseNet3(x)x = self.transConv3(x)x = self.UAFM(x, feat320)x=self.transConv4(x)x=self.DenseNet4(x)out=self.out(x)# out=torch.concat((out,out,out),dim=1)*255.return outdef freeze_backbone(self):for param in self.resnet18.parameters():param.requires_grad = Falsedef unfreeze_backbone(self):for param in self.resnet18.parameters():param.requires_grad = Trueif __name__ == '__main__':net=Net()print(net)# stat(net,(3,640,640))summary(net,input_size=(3,512,512),device='cpu')aa=torch.ones((6,3,512,512))out=net(aa)print(out.shape)# ii=torch.zeros((1,3,640,640))# outs=net(ii)# print(outs.shape)

主题界面显示及代码：

from PyQt5.QtGui import *
from PyQt5.QtWidgets import *
from untitled import Ui_Form
import sys
import cv2 as cv
from PyQt5.QtCore import QCoreApplication
import numpy as np
from PyQt5 import QtCore,QtGui
from PIL import Image
from predict import *class My(QMainWindow,Ui_Form):def __init__(self):super(My,self).__init__()self.setupUi(self)self.setWindowTitle('图像去雨去雾')self.setIcon()self.pushButton.clicked.connect(self.pic)self.pushButton_2.clicked.connect(self.pre)self.pushButton_3.clicked.connect(self.pre2)def setIcon(self):palette1 = QPalette()# palette1.setColor(self.backgroundRole(), QColor(192,253,123))   # 设置背景颜色palette1.setBrush(self.backgroundRole(), QBrush(QPixmap('back.png')))  # 设置背景图片self.setPalette(palette1)def pre(self):out=pre(self.img,0)out=self.cv_qt(out)self.label_2.setPixmap(QPixmap.fromImage(out).scaled(self.label.width(),self.label.height(),QtCore.Qt.KeepAspectRatio))def pre2(self):out=pre(self.img,1)out=self.cv_qt(out)self.label_2.setPixmap(QPixmap.fromImage(out).scaled(self.label.width(),self.label.height(),QtCore.Qt.KeepAspectRatio))def pic(self):imgName, imgType = QFileDialog.getOpenFileName(self,"打开图片",""," *.png;;*.jpg;;*.jpeg;;*.bmp;;All Files (*)")#KeepAspectRatiopng = QtGui.QPixmap(imgName).scaled(self.label.width(),self.label.height(),QtCore.Qt.KeepAspectRatio)  # 适应设计label时的大小self.label.setPixmap(png)self.img=Image.open(imgName)self.img=np.array(self.img)def cv_qt(self, src):#src必须为bgr格式图像#src必须为bgr格式图像#src必须为bgr格式图像if len(src.shape)==2:src=np.expand_dims(src,axis=-1)src=np.tile(src,(1,1,3))h, w, d = src.shapeelse:h, w, d = src.shapebytesperline = d * w# self.src=cv.cvtColor(self.src,cv.COLOR_BGR2RGB)qt_image = QImage(src.data, w, h, bytesperline, QImage.Format_RGB888).rgbSwapped()return qt_imageif __name__ == '__main__':QCoreApplication.setAttribute(QtCore.Qt.AA_EnableHighDpiScaling)app=QApplication(sys.argv)my=My()my.show()sys.exit(app.exec_())

项目结构：

直接运行main.py即可弹出交互界面。
项目下载地址：下载地址-列表第19