init

app.py

@@ -0,0 +1,40 @@
+'''
+Author: Egrt
+Date: 2022-01-13 13:34:10
+LastEditors: Egrt
+LastEditTime: 2022-01-13 13:48:57
+FilePath: \LicenseGAN\app.py
+'''
+import os
+os.system('pip install pytorch')
+os.system('pip install gradio==2.5.3')
+from PIL import Image
+from esrgan import ESRGAN
+import gradio as gr
+
+esrgan = ESRGAN()
+
+# --------模型推理---------- #
+def inference(img):
+    lr_shape = [12, 24]
+    img = img.resize((lr_shape[1], lr_shape[0]), Image.BICUBIC)
+    r_image = esrgan.generate_1x1_image(img)
+    return r_image
+
+# --------网页信息---------- #  
+title = "车牌超分辨率重建"
+description = "使用生成对抗网络对低分辨率车牌图片进行八倍的超分辨率重建，能够有效的恢复出车牌号。  @西南科技大学智能控制与图像处理研究室"
+article = "<p style='text-align: center'><a href='https://arxiv.org/abs/2108.10257' target='_blank'>LicenseGAN: Image Restoration Using Swin Transformer</a> | <a href='https://github.com/JingyunLiang/SwinIR' target='_blank'>Github Repo</a></p>"
+example_img_dir  = 'img'
+example_img_name = os.listdir(example_img_dir)
+examples=[[os.path.join(example_img_dir, image_path)] for image_path in example_img_name if image_path.endswith('.jpg')]
+gr.Interface(
+    inference, 
+    [gr.inputs.Image(type="pil", label="Input")], 
+    gr.outputs.Image(type="pil", label="Output"),
+    title=title,
+    description=description,
+    article=article,
+    enable_queue=True,
+    examples=examples
+    ).launch(debug=True)

esrgan.py

@@ -0,0 +1,79 @@
+import numpy as np
+import torch
+import torch.backends.cudnn as cudnn
+from PIL import Image
+import cv2
+from nets.esrgan import Generator
+from utils.utils import cvtColor, preprocess_input
+
+
+class ESRGAN(object):
+    #-----------------------------------------#
+    #   注意修改model_path
+    #-----------------------------------------#
+    _defaults = {
+        #-----------------------------------------------#
+        #   model_path指向logs文件夹下的权值文件
+        #-----------------------------------------------#
+        "model_path"        : 'model_data/Generator_ESRGAN.pth',
+        #-----------------------------------------------#
+        #   上采样的倍数，和训练时一样
+        #-----------------------------------------------#
+        "scale_factor"      : 8, 
+        #-------------------------------#
+        #   是否使用Cuda
+        #   没有GPU可以设置成False
+        #-------------------------------#
+        "cuda"              : False,
+    }
+
+    #---------------------------------------------------#
+    #   初始化SRGAN
+    #---------------------------------------------------#
+    def __init__(self, **kwargs):
+        self.__dict__.update(self._defaults)
+        for name, value in kwargs.items():
+            setattr(self, name, value)  
+        self.generate()
+
+    def generate(self):
+        self.net    = Generator(self.scale_factor)
+
+        device      = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        self.net.load_state_dict(torch.load(self.model_path, map_location=device))
+        self.net    = self.net.eval()
+        print('{} model, and classes loaded.'.format(self.model_path))
+
+        if self.cuda:
+            self.net = torch.nn.DataParallel(self.net)
+            cudnn.benchmark = True
+            self.net = self.net.cuda()
+        
+    def generate_1x1_image(self, image):
+        #---------------------------------------------------------#
+        #   在这里将图像转换成RGB图像，防止灰度图在预测时报错。
+        #   代码仅仅支持RGB图像的预测，所有其它类型的图像都会转化成RGB
+        #---------------------------------------------------------#
+        image       = cvtColor(image)
+        #---------------------------------------------------------#
+        #   添加上batch_size维度，并进行归一化
+        #---------------------------------------------------------#
+        image_data  = np.expand_dims(np.transpose(preprocess_input(np.array(image, dtype=np.float32), [0.5,0.5,0.5], [0.5,0.5,0.5]), [2,0,1]), 0)
+        
+        with torch.no_grad():
+            image_data = torch.from_numpy(image_data).type(torch.FloatTensor)
+            if self.cuda:
+                image_data = image_data.cuda()
+
+            #---------------------------------------------------------#
+            #   将图像输入网络当中进行预测！
+            #---------------------------------------------------------#
+            hr_image = self.net(image_data)[0]
+            #---------------------------------------------------------#
+            #   将归一化的结果再转成rgb格式
+            #---------------------------------------------------------#
+            hr_image = (hr_image.cpu().data.numpy().transpose(1, 2, 0) * 0.5 + 0.5)
+            hr_image = (hr_image-np.min(hr_image))/(np.max(hr_image)-np.min(hr_image)) * 255
+
+            hr_image = Image.fromarray(np.uint8(hr_image))
+            return hr_image

model_data/Generator_ESRGAN.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4c137b3da480f7ad251641ace39d73b2adb30ec0c40662cadce2bf0e80b8fca8
+size 28697247

nets/esrgan.py

@@ -0,0 +1,140 @@
+import math
+import torch
+from torch import nn
+
+
+class DenseResidualBlock(nn.Module):
+    """
+    密集连接型残差网络
+    """
+
+    def __init__(self, filters, res_scale=0.2):
+        super(DenseResidualBlock, self).__init__()
+        self.res_scale = res_scale
+
+        def block(in_features, non_linearity=True):
+            layers = [nn.Conv2d(in_features, filters, 3, 1, 1, bias=True)]
+            if non_linearity:
+                layers += [nn.GELU()]
+            return nn.Sequential(*layers)
+
+        self.b1 = block(in_features=1 * filters)
+        self.b2 = block(in_features=2 * filters)
+        self.b3 = block(in_features=3 * filters)
+        self.b4 = block(in_features=4 * filters)
+        self.b5 = block(in_features=5 * filters, non_linearity=False)
+        self.blocks = [self.b1, self.b2, self.b3, self.b4, self.b5]
+
+    def forward(self, x):
+        inputs = x
+        for block in self.blocks:
+            out = block(inputs)
+            inputs = torch.cat([inputs, out], 1)
+        return out.mul(self.res_scale) + x
+
+class ResidualInResidualDenseBlock(nn.Module):
+    def __init__(self, filters, res_scale=0.2):
+        super(ResidualInResidualDenseBlock, self).__init__()
+        self.res_scale = res_scale
+        self.dense_blocks = nn.Sequential(
+            DenseResidualBlock(filters), DenseResidualBlock(filters), DenseResidualBlock(filters)
+        )
+
+    def forward(self, x):
+        return self.dense_blocks(x).mul(self.res_scale) + x
+
+class UpsampleBLock(nn.Module):
+    def __init__(self, in_channels, up_scale):
+        super(UpsampleBLock, self).__init__()
+        self.conv = nn.Conv2d(in_channels, in_channels * up_scale ** 2, kernel_size=3, padding=1)
+        self.pixel_shuffle = nn.PixelShuffle(up_scale)
+        self.gelu = nn.GELU()
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.pixel_shuffle(x)
+        x = self.gelu(x)
+        return x
+
+class Generator(nn.Module):
+    def __init__(self, scale_factor, channels=3, filters=64, num_res_blocks=4):
+        super(Generator, self).__init__()
+        upsample_block_num = int(math.log(scale_factor, 2))
+        # 第一个卷积层
+        self.conv1 = nn.Conv2d(channels, filters, kernel_size=3, stride=1, padding=1)
+        # 密集残差连接块
+        self.res_blocks = nn.Sequential(*[ResidualInResidualDenseBlock(filters) for _ in range(num_res_blocks)])
+        # 第二个卷积层
+        self.conv2 = nn.Conv2d(filters, filters, kernel_size=3, stride=1, padding=1)
+        self.upsample = [UpsampleBLock(filters, 2) for _ in range(upsample_block_num)]
+        self.upsample = nn.Sequential(*self.upsample)
+        # 输出卷积层
+        self.conv3 = nn.Sequential(
+                        nn.Conv2d(filters, filters, kernel_size=3, stride=1, padding=1),
+                        nn.GELU(),
+                        nn.Conv2d(filters, channels, kernel_size=3, stride=1, padding=1)
+                        )
+
+    def forward(self, x):
+        out1 = self.conv1(x)
+        out  = self.res_blocks(out1)
+        out2 = self.conv2(out)
+        out  = torch.add(out1, out2)
+        upsample = self.upsample(out)
+        out  = self.conv3(upsample)
+        return out
+
+
+class Discriminator(nn.Module):
+    def __init__(self):
+        super(Discriminator, self).__init__()
+        self.net = nn.Sequential(
+            nn.Conv2d(3, 64, kernel_size=3, padding=1),
+            nn.GELU(),
+
+            nn.Conv2d(64, 64, kernel_size=3, stride=2, padding=1),
+            nn.BatchNorm2d(64),
+            nn.GELU(),
+
+            nn.Conv2d(64, 128, kernel_size=3, padding=1),
+            nn.BatchNorm2d(128),
+            nn.GELU(),
+
+            nn.Conv2d(128, 128, kernel_size=3, stride=2, padding=1),
+            nn.BatchNorm2d(128),
+            nn.GELU(),
+
+            nn.Conv2d(128, 256, kernel_size=3, padding=1),
+            nn.BatchNorm2d(256),
+            nn.GELU(),
+
+            nn.Conv2d(256, 256, kernel_size=3, stride=2, padding=1),
+            nn.BatchNorm2d(256),
+            nn.GELU(),
+
+            nn.Conv2d(256, 512, kernel_size=3, padding=1),
+            nn.BatchNorm2d(512),
+            nn.GELU(),
+
+            nn.Conv2d(512, 512, kernel_size=3, stride=2, padding=1),
+            nn.BatchNorm2d(512),
+            nn.GELU(),
+
+            nn.AdaptiveAvgPool2d(1),
+            nn.Conv2d(512, 1024, kernel_size=1),
+            nn.GELU(),
+            nn.Conv2d(1024, 1, kernel_size=1)
+        )
+
+    def forward(self, x):
+        batch_size = x.size(0)
+        return torch.sigmoid(self.net(x).view(batch_size))
+
+if __name__ == "__main__":
+    from torchsummary import summary
+
+    # 需要使用device来指定网络在GPU还是CPU运行
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    model = Generator(8).to(device)
+    summary(model, input_size=(3,12,24))
+

utils/__init__.py

@@ -0,0 +1 @@
+from .utils import *

utils/dataloader.py

@@ -0,0 +1,157 @@
+from random import randint
+
+import cv2
+import numpy as np
+from PIL import Image
+from torch.utils.data.dataset import Dataset
+
+from utils import cvtColor, preprocess_input
+from torch.utils.data import DataLoader
+
+def get_new_img_size(width, height, img_min_side=600):
+    if width <= height:
+        f = float(img_min_side) / width
+        resized_height = int(f * height)
+        resized_width = int(img_min_side)
+    else:
+        f = float(img_min_side) / height
+        resized_width = int(f * width)
+        resized_height = int(img_min_side)
+
+    return resized_width, resized_height
+
+class SRGANDataset(Dataset):
+    def __init__(self, train_lines, lr_shape, hr_shape):
+        super(SRGANDataset, self).__init__()
+
+        self.train_lines    = train_lines
+        self.train_batches  = len(train_lines)
+
+        self.lr_shape       = lr_shape
+        self.hr_shape       = hr_shape
+
+    def __len__(self):
+        return self.train_batches
+
+    def __getitem__(self, index):
+        index = index % self.train_batches
+
+        image_origin = Image.open(self.train_lines[index].split()[0])
+        if self.rand()<.5:
+            img_h = self.get_random_data(image_origin, self.hr_shape)
+        else:
+            img_h = self.random_crop(image_origin, self.hr_shape[1], self.hr_shape[0])
+        img_l = img_h.resize((self.lr_shape[1], self.lr_shape[0]), Image.BICUBIC)
+
+        img_h = np.transpose(preprocess_input(np.array(img_h, dtype=np.float32), [0.5,0.5,0.5], [0.5,0.5,0.5]), [2,0,1])
+        img_l = np.transpose(preprocess_input(np.array(img_l, dtype=np.float32), [0.5,0.5,0.5], [0.5,0.5,0.5]), [2,0,1])
+        return np.array(img_l), np.array(img_h)
+
+    def rand(self, a=0, b=1):
+        return np.random.rand()*(b-a) + a
+
+    def get_random_data(self, image, input_shape, jitter=.3, hue=.1, sat=1.5, val=1.5, random=True):
+        #------------------------------#
+        #   读取图像并转换成RGB图像
+        #------------------------------#
+        image   = cvtColor(image)
+        #------------------------------#
+        #   获得图像的高宽与目标高宽
+        #------------------------------#
+        iw, ih  = image.size
+        h, w    = input_shape
+
+        if not random:
+            scale = min(w/iw, h/ih)
+            nw = int(iw*scale)
+            nh = int(ih*scale)
+            dx = (w-nw)//2
+            dy = (h-nh)//2
+
+            #---------------------------------#
+            #   将图像多余的部分加上灰条
+            #---------------------------------#
+            image       = image.resize((nw,nh), Image.BICUBIC)
+            new_image   = Image.new('RGB', (w,h), (128,128,128))
+            new_image.paste(image, (dx, dy))
+            image_data  = np.array(new_image, np.float32)
+
+            return image_data
+
+        #------------------------------------------#
+        #   对图像进行缩放并且进行长和宽的扭曲
+        #------------------------------------------#
+        new_ar = w/h * self.rand(1-jitter,1+jitter)/self.rand(1-jitter,1+jitter)
+        scale = self.rand(1, 1.5)
+        if new_ar < 1:
+            nh = int(scale*h)
+            nw = int(nh*new_ar)
+        else:
+            nw = int(scale*w)
+            nh = int(nw/new_ar)
+        image = image.resize((nw,nh), Image.BICUBIC)
+
+        #------------------------------------------#
+        #   将图像多余的部分加上灰条
+        #------------------------------------------#
+        dx = int(self.rand(0, w-nw))
+        dy = int(self.rand(0, h-nh))
+        new_image = Image.new('RGB', (w,h), (128,128,128))
+        new_image.paste(image, (dx, dy))
+        image = new_image
+
+        #------------------------------------------#
+        #   翻转图像
+        #------------------------------------------#
+        flip = self.rand()<.5
+        if flip: image = image.transpose(Image.FLIP_LEFT_RIGHT)
+        
+        rotate = self.rand()<.5
+        if rotate: 
+            angle = np.random.randint(-15,15)
+            a,b = w/2,h/2
+            M = cv2.getRotationMatrix2D((a,b),angle,1)
+            image = cv2.warpAffine(np.array(image), M, (w,h), borderValue=[128,128,128]) 
+
+        #------------------------------------------#
+        #   色域扭曲
+        #------------------------------------------#
+        hue = self.rand(-hue, hue)
+        sat = self.rand(1, sat) if self.rand()<.5 else 1/self.rand(1, sat)
+        val = self.rand(1, val) if self.rand()<.5 else 1/self.rand(1, val)
+        x = cv2.cvtColor(np.array(image,np.float32)/255, cv2.COLOR_RGB2HSV)
+        x[..., 1] *= sat
+        x[..., 2] *= val
+        x[x[:,:, 0]>360, 0] = 360
+        x[:, :, 1:][x[:, :, 1:]>1] = 1
+        x[x<0] = 0
+        image_data = cv2.cvtColor(x, cv2.COLOR_HSV2RGB)*255
+        return Image.fromarray(np.uint8(image_data))
+
+    def random_crop(self, image, width, height):
+        #--------------------------------------------#
+        #   如果图像过小无法截取，先对图像进行放大
+        #--------------------------------------------#
+        if image.size[0] < self.hr_shape[1] or image.size[1] < self.hr_shape[0]:
+            resized_width, resized_height = get_new_img_size(width, height, img_min_side=np.max(self.hr_shape))
+            image = image.resize((resized_width, resized_height), Image.BICUBIC)
+
+        #--------------------------------------------#
+        #   随机截取一部分
+        #--------------------------------------------#
+        width1  = randint(0, image.size[0] - width)
+        height1 = randint(0, image.size[1] - height)
+
+        width2  = width1 + width
+        height2 = height1 + height
+
+        image   = image.crop((width1, height1, width2, height2))
+        return image
+        
+def SRGAN_dataset_collate(batch):
+    images_l = []
+    images_h = []
+    for img_l, img_h in batch:
+        images_l.append(img_l)
+        images_h.append(img_h)
+    return np.array(images_l), np.array(images_h)

utils/preprocess.py

@@ -0,0 +1,151 @@
+import numpy as np
+import os
+import matplotlib.image as mpimage
+import argparse
+import functools
+from utils import add_arguments, print_arguments
+from dask.distributed import LocalCluster
+from dask import bag as dbag
+from dask.diagnostics import ProgressBar
+from typing import Tuple
+from PIL import Image
+
+
+
+# Dataset statistics that I gathered in development
+#-----------------------------------#
+#   用于过滤感知质量较低的不良图片
+#-----------------------------------#
+IMAGE_MEAN = 0.5
+IMAGE_MEAN_STD = 0.028
+
+IMG_STD = 0.28
+IMG_STD_STD = 0.01
+
+
+def readImage(fileName: str) -> np.ndarray:
+    image = mpimage.imread(fileName)
+    return image
+
+
+#-----------------------------------#
+#   从文件名中提取车牌的坐标
+#-----------------------------------#
+
+
+def parseLabel(label: str) -> Tuple[np.ndarray, np.ndarray]:
+    annotation = label.split('-')[3].split('_')
+    coor1 = [int(i) for i in annotation[0].split('&')]
+    coor2 = [int(i) for i in annotation[1].split('&')]
+    coor3 = [int(i) for i in annotation[2].split('&')]
+    coor4 = [int(i) for i in annotation[3].split('&')]
+    coor = np.array([coor1, coor2, coor3, coor4])
+    center = np.mean(coor, axis=0)
+    return coor, center.astype(int)
+
+
+#-----------------------------------#
+#   根据车牌坐标裁剪出车牌图像
+#-----------------------------------#
+
+
+def cropImage(image: np.ndarray, coor: np.ndarray, center: np.ndarray) -> np.ndarray:
+    maxW = np.max(coor[:, 0] - center[0])  # max plate width
+    maxH = np.max(coor[:, 1] - center[1])  # max plate height
+
+    xWanted = [64, 128, 192, 256]
+    yWanted = [32, 64, 96, 128]
+
+    found = False
+    for w, h in zip(xWanted, yWanted):
+        if maxW < w//2 and maxH < h//2:
+            maxH = h//2
+            maxW = w//2
+            found = True
+            break
+    if not found:  # 车牌太大则丢弃
+        return np.array([])
+    elif center[1]-maxH < 0 or center[1]+maxH >= image.shape[1] or \
+            center[0]-maxW < 0 or center[0] + maxW >= image.shape[0]:
+        return np.array([])
+    else:
+        return image[center[1]-maxH:center[1]+maxH, center[0]-maxW:center[0]+maxW]
+
+#-----------------------------------#
+#           保存车牌图片
+#-----------------------------------#
+
+
+def saveImage(image: np.ndarray, fileName: str, outDir: str) -> int:
+    if image.shape[0] == 0:
+        return 0
+    else:
+        imgShape = image.shape
+        if imgShape[1] == 64:
+            mpimage.imsave(os.path.join(outDir, '64_32', fileName), image)
+        elif imgShape[1] == 128:
+            mpimage.imsave(os.path.join(outDir, '128_64', fileName), image)
+        elif imgShape[1] == 208:
+            mpimage.imsave(os.path.join(outDir, '192_96', fileName), image)
+        else: #resize large images
+            image = Image.fromarray(image).resize((192, 96))
+            image = np.asarray(image) # back to numpy array
+            mpimage.imsave(os.path.join(outDir, '192_96', fileName), image)
+        return 1
+
+
+#-----------------------------------#
+# 包装成一个函数，以便将处理区分到不同目录
+#-----------------------------------#
+
+def processImage(file: str, inputDir: str, outputDir: str, subFolder: str) -> int:
+    result = parseLabel(file)
+    filePath = os.path.join(inputDir,subFolder, file)
+    image = readImage(filePath)
+    plate = cropImage(image, result[0], result[1])
+    if plate.shape[0] == 0:
+        return 0
+    mean = np.mean(plate/255.0)
+    std = np.std(plate/255.0)
+    # 亮度不好的
+    if mean <= IMAGE_MEAN - 10*IMAGE_MEAN_STD or mean >= IMAGE_MEAN + 10*IMAGE_MEAN_STD:
+        return 0
+    # 低对比度的
+    if std <= IMG_STD - 10*IMG_STD_STD:
+        return 0
+    status = saveImage(plate, file, outputDir)
+    return status
+
+
+def main(argv):
+    jobNum = int(argv.jobNum)
+    outputDir = argv.outputDir
+    inputDir = argv.inputDir
+    try:
+        os.mkdir(outputDir)
+        for shape in ['64_32', '128_64', '192_96']:
+            os.mkdir(os.path.join(outputDir, shape))
+    except OSError:
+        pass  # 地址已经存在
+    client = LocalCluster(n_workers=jobNum, threads_per_worker=5)  # 开启多线程
+    for subFolder in ['ccpd_base', 'ccpd_db', 'ccpd_fn', 'ccpd_rotate', 'ccpd_tilt', 'ccpd_weather']:
+        fileList = os.listdir(os.path.join(inputDir, subFolder))
+        print('* {} images found in {}. Start processing ...'.format(len(fileList), subFolder))
+        toDo = dbag.from_sequence(fileList, npartitions=jobNum*30).persist()  # persist the bag in memory
+        toDo = toDo.map(processImage, inputDir, outputDir, subFolder)
+        pbar = ProgressBar(minimum=2.0)
+        pbar.register()  # 登记所有的计算，以便更好地跟踪
+        result = toDo.compute()
+        print('* image cropped: {}. Done ...'.format(sum(result)))
+    client.close()  # 关闭集群
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description=__doc__)
+    add_arg = functools.partial(add_arguments, argparser=parser)
+    add_arg('jobNum',           int,    4,                                         '处理图片的线程数')
+    add_arg('inputDir',         str,    'datasets/CCPD2019',                       '输入图片目录')
+    add_arg('outputDir',        str,    'datasets/CCPD2019_new',                   '保存图片目录')
+    args = parser.parse_args()
+    print_arguments(args)
+    main(args)

utils/utils.py

@@ -0,0 +1,60 @@
+import itertools
+import numpy as np
+import matplotlib.pyplot as plt
+import torch
+import distutils.util
+
+def show_result(num_epoch, G_net, imgs_lr, imgs_hr):
+    with torch.no_grad():
+        test_images = G_net(imgs_lr)
+
+        fig, ax = plt.subplots(1, 2)
+
+        for j in itertools.product(range(2)):
+            ax[j].get_xaxis().set_visible(False)
+            ax[j].get_yaxis().set_visible(False)
+        
+        ax[0].cla()
+        ax[0].imshow(np.transpose(test_images.cpu().numpy()[0] * 0.5 + 0.5, [1,2,0]))
+
+        ax[1].cla()
+        ax[1].imshow(np.transpose(imgs_hr.cpu().numpy()[0] * 0.5 + 0.5, [1,2,0]))
+
+        label = 'Epoch {0}'.format(num_epoch)
+        fig.text(0.5, 0.04, label, ha='center')
+        plt.savefig("results/train_out/epoch_" + str(num_epoch) + "_results.png")
+        plt.close('all')  #避免内存泄漏
+
+#---------------------------------------------------------#
+#   将图像转换成RGB图像，防止灰度图在预测时报错。
+#   代码仅仅支持RGB图像的预测，所有其它类型的图像都会转化成RGB
+#---------------------------------------------------------#
+def cvtColor(image):
+    if len(np.shape(image)) == 3 and np.shape(image)[2] == 3:
+        return image 
+    else:
+        image = image.convert('RGB')
+        return image 
+
+def preprocess_input(image, mean, std):
+    image = (image/255 - mean)/std
+    return image
+
+def get_lr(optimizer):
+    for param_group in optimizer.param_groups:
+        return param_group['lr']
+
+def print_arguments(args):
+    print("-----------  Configuration Arguments -----------")
+    for arg, value in sorted(vars(args).items()):
+        print("%s: %s" % (arg, value))
+    print("------------------------------------------------")
+
+
+def add_arguments(argname, type, default, help, argparser, **kwargs):
+    type = distutils.util.strtobool if type == bool else type
+    argparser.add_argument("--" + argname,
+                           default=default,
+                           type=type,
+                           help=help + ' 默认: %(default)s.',
+                           **kwargs)

utils/utils_fit.py

@@ -0,0 +1,85 @@
+import torch
+from tqdm import tqdm
+
+from .utils import show_result, get_lr
+from .utils_metrics import PSNR, SSIM
+
+
+def fit_one_epoch(G_model_train, D_model_train, G_model, D_model, VGG_feature_model, G_optimizer, D_optimizer, BCE_loss, MSE_loss, epoch, epoch_size, gen, Epoch, cuda, batch_size, save_interval):
+    G_total_loss = 0
+    D_total_loss = 0
+    G_total_PSNR = 0
+    G_total_SSIM = 0
+
+    with tqdm(total=epoch_size,desc=f'Epoch {epoch + 1}/{Epoch}',postfix=dict,mininterval=0.3) as pbar:
+        for iteration, batch in enumerate(gen):
+            if iteration >= epoch_size:
+                break
+            
+            with torch.no_grad():
+                lr_images, hr_images    = batch
+                lr_images, hr_images    = torch.from_numpy(lr_images).type(torch.FloatTensor), torch.from_numpy(hr_images).type(torch.FloatTensor)
+                y_real, y_fake          = torch.ones(batch_size), torch.zeros(batch_size)
+                if cuda:
+                    lr_images, hr_images, y_real, y_fake  = lr_images.cuda(), hr_images.cuda(), y_real.cuda(), y_fake.cuda()
+            
+            #-------------------------------------------------#
+            #   训练判别器
+            #-------------------------------------------------#
+            D_optimizer.zero_grad()
+
+            D_result                = D_model_train(hr_images)
+            D_real_loss             = BCE_loss(D_result, y_real)
+            D_real_loss.backward()
+
+            G_result                = G_model_train(lr_images)
+            D_result                = D_model_train(G_result).squeeze()
+            D_fake_loss             = BCE_loss(D_result, y_fake)
+            D_fake_loss.backward()
+
+            D_optimizer.step()
+
+            D_train_loss            = D_real_loss + D_fake_loss
+
+            #-------------------------------------------------#
+            #   训练生成器
+            #-------------------------------------------------#
+            G_optimizer.zero_grad()
+
+            G_result                = G_model_train(lr_images)
+            image_loss              = MSE_loss(G_result, hr_images)
+
+            D_result                = D_model_train(G_result).squeeze()
+            adversarial_loss        = BCE_loss(D_result, y_real)
+
+            perception_loss         = MSE_loss(VGG_feature_model(G_result), VGG_feature_model(hr_images))
+
+            G_train_loss            = image_loss + 1e-3 * adversarial_loss + 2e-6 * perception_loss 
+
+            G_train_loss.backward()
+            G_optimizer.step()
+            
+            G_total_loss            += G_train_loss.item()
+            D_total_loss            += D_train_loss.item()
+
+            with torch.no_grad():
+                G_total_PSNR        += PSNR(G_result, hr_images).item()
+                G_total_SSIM        += SSIM(G_result, hr_images).item()
+                
+            pbar.set_postfix(**{'G_loss'    : G_total_loss / (iteration + 1), 
+                                'D_loss'    : D_total_loss / (iteration + 1), 
+                                'G_PSNR'    : G_total_PSNR / (iteration + 1), 
+                                'G_SSIM'    : G_total_SSIM / (iteration + 1), 
+                                'lr'        : get_lr(G_optimizer)})
+            pbar.update(1)
+
+            if iteration % save_interval == 0:
+                show_result(epoch + 1, G_model_train, lr_images, hr_images)
+
+    print('Epoch:'+ str(epoch + 1) + '/' + str(Epoch))
+    print('G Loss: %.4f || D Loss: %.4f ' % (G_total_loss / epoch_size, D_total_loss / epoch_size))
+    print('Saving state, iter:', str(epoch+1))
+    
+    if (epoch + 1) % 10==0:
+        torch.save(G_model.state_dict(), 'logs/G_Epoch%d-GLoss%.4f-DLoss%.4f.pth'%((epoch + 1), G_total_loss / epoch_size, D_total_loss / epoch_size))
+        torch.save(D_model.state_dict(), 'logs/D_Epoch%d-GLoss%.4f-DLoss%.4f.pth'%((epoch + 1), G_total_loss / epoch_size, D_total_loss / epoch_size))

utils/utils_metrics.py

@@ -0,0 +1,69 @@
+import torch
+import torch.nn.functional as F
+from math import exp
+import numpy as np
+ 
+def gaussian(window_size, sigma):
+    gauss = torch.Tensor([exp(-(x - window_size//2)**2/float(2*sigma**2)) for x in range(window_size)])
+    return gauss/gauss.sum()
+ 
+def create_window(window_size, channel=1):
+    _1D_window = gaussian(window_size, 1.5).unsqueeze(1)
+    _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0)
+    window = _2D_window.expand(channel, 1, window_size, window_size).contiguous()
+    return window
+
+def SSIM(img1, img2, window_size=11, window=None, size_average=True, full=False):
+    img1 = (img1 * 0.5 + 0.5) * 255
+    img2 = (img2 * 0.5 + 0.5) * 255
+    min_val = 0
+    max_val = 255
+    L = max_val - min_val
+    img2 = torch.clamp(img2, 0.0, 255.0)
+ 
+    padd = 0
+    (_, channel, height, width) = img1.size()
+    if window is None:
+        real_size = min(window_size, height, width)
+        window = create_window(real_size, channel=channel).to(img1.device)
+ 
+    mu1 = F.conv2d(img1, window, padding=padd, groups=channel)
+    mu2 = F.conv2d(img2, window, padding=padd, groups=channel)
+ 
+    mu1_sq = mu1.pow(2)
+    mu2_sq = mu2.pow(2)
+    mu1_mu2 = mu1 * mu2
+ 
+    sigma1_sq = F.conv2d(img1 * img1, window, padding=padd, groups=channel) - mu1_sq
+    sigma2_sq = F.conv2d(img2 * img2, window, padding=padd, groups=channel) - mu2_sq
+    sigma12 = F.conv2d(img1 * img2, window, padding=padd, groups=channel) - mu1_mu2
+ 
+    C1 = (0.01 * L) ** 2
+    C2 = (0.03 * L) ** 2
+ 
+    v1 = 2.0 * sigma12 + C2
+    v2 = sigma1_sq + sigma2_sq + C2
+    cs = torch.mean(v1 / v2)  # contrast sensitivity
+ 
+    ssim_map = ((2 * mu1_mu2 + C1) * v1) / ((mu1_sq + mu2_sq + C1) * v2)
+ 
+    if size_average:
+        ret = ssim_map.mean()
+    else:
+        ret = ssim_map.mean(1).mean(1).mean(1)
+ 
+    if full:
+        return ret, cs
+    return ret
+ 
+def tf_log10(x):
+    numerator = torch.log(x)
+    denominator = torch.log(torch.tensor(10.0))
+    return numerator / denominator
+
+def PSNR(img1, img2):
+    img1 = (img1 * 0.5 + 0.5) * 255
+    img2 = (img2 * 0.5 + 0.5) * 255
+    max_pixel = 255.0
+    img2 = torch.clamp(img2, 0.0, 255.0)
+    return 10.0 * tf_log10((max_pixel ** 2) / (torch.mean(torch.pow(img2 - img1, 2))))