基于Pytorch实现对偶生成对抗网络来实现图像去雾python源码+项目说明+代码注释.zip

import itertools import os import random import argparse import cv2 import numpy as np import matplotlib.pyplot as plt import torch from torch import nn, optim from torch.utils.data import Dataset, DataLoader from torchvision import transforms from torchvision.utils import save_image device = torch.device("cuda" if torch.cuda.is_available() else "cpu") if torch.cuda.is_available(): print(" -- 使用GPU进行训练 -- ") def parseArgs(): parser = argparse.ArgumentParser(description='manual to this script') parser.add_argument("--save_path", type=str, default="./result/") parser.add_argument("--data_path", type=str, default="../Data/") parser.add_argument("--origin", type=str, default="GT") parser.add_argument("--hazy", type=str, default="hazy") parser.add_argument("--batch_size", type=int, default=4) args = parser.parse_args() return args ## 生成器 U-Net（输入照片为256*256） ## class Generator(nn.Module): def __init__(self, in_ch, out_ch, ngf=64): """ 定义生成器的网络结构 :param in_ch: 输入数据的通道数 :param out_ch: 输出数据的通道数 :param ngf: 第一层卷积的通道数 number of generator's first conv filters """ super(Generator, self).__init__() # 下面的激活函数都放在下一个模块的第一步 是为了skip-connect方便 # 左半部分 U-Net encoder # 每层输入大小折半，从输入图片大小256开始 # 256 * 256（输入） self.en1 = nn.Sequential( nn.Conv2d(in_ch, ngf, kernel_size=4, stride=2, padding=1), # 输入图片已正则化 不需BatchNorm ) # 128 * 128 self.en2 = nn.Sequential( nn.LeakyReLU(0.2, inplace=True), nn.Conv2d(ngf, ngf * 2, kernel_size=4, stride=2, padding=1), nn.BatchNorm2d(ngf * 2) ) # 64 * 64 self.en3 = nn.Sequential( nn.LeakyReLU(0.2, inplace=True), nn.Conv2d(ngf * 2, ngf * 4, kernel_size=4, stride=2, padding=1), nn.BatchNorm2d(ngf * 4) ) # 32 * 32 self.en4 = nn.Sequential( nn.LeakyReLU(0.2, inplace=True), nn.Conv2d(ngf * 4, ngf * 8, kernel_size=4, stride=2, padding=1), nn.BatchNorm2d(ngf * 8) ) # 16 * 16 self.en5 = nn.Sequential( nn.LeakyReLU(0.2, inplace=True), nn.Conv2d(ngf * 8, ngf * 8, kernel_size=4, stride=2, padding=1), nn.BatchNorm2d(ngf * 8) ) # 8 * 8 self.en6 = nn.Sequential( nn.LeakyReLU(0.2, inplace=True), nn.Conv2d(ngf * 8, ngf * 8, kernel_size=4, stride=2, padding=1), nn.BatchNorm2d(ngf * 8) ) # 4 * 4 self.en7 = nn.Sequential( nn.LeakyReLU(0.2, inplace=True), nn.Conv2d(ngf * 8, ngf * 8, kernel_size=4, stride=2, padding=1), nn.BatchNorm2d(ngf * 8) ) # 2 * 2 self.en8 = nn.Sequential( nn.LeakyReLU(0.2, inplace=True), nn.Conv2d(ngf * 8, ngf * 8, kernel_size=4, stride=2, padding=1) # Encoder输出不用BatchNorm ) # 右半部分 U-Net decoder # skip-connect: 前一层的输出+对称的卷积层 # 1 * 1（输入） self.de1 = nn.Sequential( nn.ReLU(inplace=True), nn.ConvTranspose2d(ngf * 8, ngf * 8, kernel_size=4, stride=2, padding=1), nn.BatchNorm2d(ngf * 8), nn.Dropout(p=0.5) ) # 2 * 2 self.de2 = nn.Sequential( nn.ReLU(inplace=True), # skip-connect 所以输入管道数是之前输出的2倍 nn.ConvTranspose2d(ngf * 8 * 2, ngf * 8, kernel_size=4, stride=2, padding=1), nn.BatchNorm2d(ngf * 8), nn.Dropout(p=0.5) ) # 4 * 4 self.de3 = nn.Sequential( nn.ReLU(inplace=True), nn.ConvTranspose2d(ngf * 8 * 2, ngf * 8, kernel_size=4, stride=2, padding=1), nn.BatchNorm2d(ngf * 8), nn.Dropout(p=0.5) ) # 8 * 8 self.de4 = nn.Sequential( nn.ReLU(inplace=True), nn.ConvTranspose2d(ngf * 8 * 2, ngf * 8, kernel_size=4, stride=2, padding=1), nn.BatchNorm2d(ngf * 8), nn.Dropout(p=0.5) ) # 16 * 16 self.de5 = nn.Sequential( nn.ReLU(inplace=True), nn.ConvTranspose2d(ngf * 8 * 2, ngf * 4, kernel_size=4, stride=2, padding=1), nn.BatchNorm2d(ngf * 4), nn.Dropout(p=0.5) ) # 32 * 32 self.de6 = nn.Sequential( nn.ReLU(inplace=True), nn.ConvTranspose2d(ngf * 4 * 2, ngf * 2, kernel_size=4, stride=2, padding=1), nn.BatchNorm2d(ngf * 2), nn.Dropout(p=0.5) ) # 64 * 64 self.de7 = nn.Sequential( nn.ReLU(inplace=True), nn.ConvTranspose2d(ngf * 2 * 2, ngf, kernel_size=4, stride=2, padding=1), nn.BatchNorm2d(ngf), nn.Dropout(p=0.5) ) # 128 * 128 self.de8 = nn.Sequential( nn.ReLU(inplace=True), nn.ConvTranspose2d(ngf * 2, out_ch, kernel_size=4, stride=2, padding=1), # Encoder输出不用BatchNorm nn.Tanh() ) def forward(self, X): """ 生成器模块前向传播 :param X: 输入生成器的数据 :return: 生成器的输出 """ # Encoder en1_out = self.en1(X) en2_out = self.en2(en1_out) en3_out = self.en3(en2_out) en4_out = self.en4(en3_out) en5_out = self.en5(en4_out) en6_out = self.en6(en5_out) en7_out = self.en7(en6_out) en8_out = self.en8(en7_out) # Decoder de1_out = self.de1(en8_out) de1_cat = torch.cat([de1_out, en7_out], dim=1) # cat by channel de2_out = self.de2(de1_cat) de2_cat = torch.cat([de2_out, en6_out], 1) de3_out = self.de3(de2_cat) de3_cat = torch.cat([de3_out, en5_out], 1) de4_out = self.de4(de3_cat) de4_cat = torch.cat([de4_out, en4_out], 1) de5_out = self.de5(de4_cat) de5_cat = torch.cat([de5_out, en3_out], 1) de6_out = self.de6(de5_cat) de6_cat = torch.cat([de6_out, en2_out], 1) de7_out = self.de7(de6_cat) de7_cat = torch.cat([de7_out, en1_out], 1) de8_out = self.de8(de7_cat) return de8_out ## 辨别器 PatchGAN（其实就是卷积网络而已） ## class Discriminator(nn.Module): def __init__(self, in_ch, ndf=64): """ 定义判别器的网络结构 :param in_ch: 输入数据的通道数 :param ndf: 第一层卷积的通道数 number of discriminator's first conv filters """ super(Discriminator, self).__init__() # 不是输出一个表示真假概率的实数，而是一个N*N的Patch矩阵（此处为30*30），其中每一块对应输入数据的一小块 # in_ch + out_ch 是为将对应真假数据同时输入 # 256 * 256（输入） self.layer1 = nn.Sequential( nn.Conv2d(in_ch, ndf, kernel_size=4, stride=2, padding=1), # 输入图片已正则化 不需BatchNorm nn.LeakyReLU(0.2, inplace=True) ) # 128 * 128 self.layer2 = nn.Sequential( nn.Conv2d(ndf, ndf * 2, kernel_size=4, stride=2, padding=1), nn.BatchNorm2d(ndf * 2), nn.LeakyReLU(0.2, inplace=True) ) # 64 * 64 self.layer3 = nn.Sequential( nn.Conv2d(ndf * 2, ndf * 4, kernel_size=4, stride=2, padding=1), nn.BatchNorm2d(ndf * 4), nn.LeakyReLU(0.2, inplace=True) ) # 32 * 32 self.layer4 = nn.Sequential( nn.Conv2d(ndf * 4, ndf * 8