Pytorch搭建SRGAN平臺提升圖片超分辨率

更新時間：2022年04月29日 17:37:32 作者：Bubbliiiing

這篇文章主要為大家介紹了Pytorch搭建SRGAN平臺提升圖片超分辨率，有需要的朋友可以借鑒參考下，希望能夠有所幫助，祝大家多多進(jìn)步，早日升職加薪

網(wǎng)絡(luò)構(gòu)建

一、什么是SRGAN

SRGAN出自論文Photo-Realistic Single Image Super-Resolution Using a Generative Adversarial Network。

如果將SRGAN看作一個黑匣子，其主要的功能就是輸入一張低分辨率圖片，生成高分辨率圖片。

該文章提到，普通的超分辨率模型訓(xùn)練網(wǎng)絡(luò)時只用到了均方差作為損失函數(shù)，雖然能夠獲得很高的峰值信噪比，但是恢復(fù)出來的圖像通常會丟失高頻細(xì)節(jié)。

SRGAN利用感知損失(perceptual loss)和對抗損失(adversarial loss)來提升恢復(fù)出的圖片的真實感。

二、生成網(wǎng)絡(luò)的構(gòu)建

生成網(wǎng)絡(luò)的構(gòu)成如上圖所示，生成網(wǎng)絡(luò)的作用是輸入一張低分辨率圖片，生成高分辨率圖片。：

SRGAN的生成網(wǎng)絡(luò)由三個部分組成。

1、低分辨率圖像進(jìn)入后會經(jīng)過一個卷積+RELU函數(shù)。

2、然后經(jīng)過B個殘差網(wǎng)絡(luò)結(jié)構(gòu)，每個殘差結(jié)構(gòu)都包含兩個卷積+標(biāo)準(zhǔn)化+RELU，還有一個殘差邊。

3、然后進(jìn)入上采樣部分，在經(jīng)過兩次上采樣后，原圖的高寬變?yōu)樵瓉淼?倍，實現(xiàn)分辨率的提升。

前兩個部分用于特征提取，第三部分用于提高分辨率。

import math
import torch
from torch import nn
class ResidualBlock(nn.Module):
    def __init__(self, channels):
        super(ResidualBlock, self).__init__()
        self.conv1 = nn.Conv2d(channels, channels, kernel_size=3, padding=1)
        self.bn1 = nn.BatchNorm2d(channels)
        self.prelu = nn.PReLU(channels)
        self.conv2 = nn.Conv2d(channels, channels, kernel_size=3, padding=1)
        self.bn2 = nn.BatchNorm2d(channels)
    def forward(self, x):
        short_cut = x
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.prelu(x)
        x = self.conv2(x)
        x = self.bn2(x)
        return x + short_cut
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.prelu = nn.PReLU(in_channels)
    def forward(self, x):
        x = self.conv(x)
        x = self.pixel_shuffle(x)
        x = self.prelu(x)
        return x
class Generator(nn.Module):
    def __init__(self, scale_factor, num_residual=16):
        upsample_block_num = int(math.log(scale_factor, 2))
        super(Generator, self).__init__()
        self.block_in = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=9, padding=4),
            nn.PReLU(64)
        )
        self.blocks = []
        for _ in range(num_residual):
            self.blocks.append(ResidualBlock(64))
        self.blocks = nn.Sequential(*self.blocks)
        self.block_out = nn.Sequential(
            nn.Conv2d(64, 64, kernel_size=3, padding=1),
            nn.BatchNorm2d(64)
        )
        self.upsample = [UpsampleBLock(64, 2) for _ in range(upsample_block_num)]
        self.upsample.append(nn.Conv2d(64, 3, kernel_size=9, padding=4))
        self.upsample = nn.Sequential(*self.upsample)
    def forward(self, x):
        x = self.block_in(x)
        short_cut = x
        x = self.blocks(x)
        x = self.block_out(x)
        upsample = self.upsample(x + short_cut)
        return torch.tanh(upsample)

三、判別網(wǎng)絡(luò)的構(gòu)建

判別網(wǎng)絡(luò)的構(gòu)成如上圖所示：

SRGAN的判別網(wǎng)絡(luò)由不斷重復(fù)的卷積+LeakyRELU和標(biāo)準(zhǔn)化組成。
對于判斷網(wǎng)絡(luò)來講，它的目的是判斷輸入圖片的真假，它的輸入是圖片，輸出是判斷結(jié)果。

判斷結(jié)果處于0-1之間，利用接近1代表判斷為真圖片，接近0代表判斷為假圖片。

判斷網(wǎng)絡(luò)的構(gòu)建和普通卷積網(wǎng)絡(luò)差距不大，都是不斷的卷積對圖片進(jìn)行下采用，在多次卷積后，最終接一次全連接判斷結(jié)果。

實現(xiàn)代碼如下：

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.LeakyReLU(0.2),
            nn.Conv2d(64, 64, kernel_size=3, stride=2, padding=1),
            nn.BatchNorm2d(64),
            nn.LeakyReLU(0.2),
            nn.Conv2d(64, 128, kernel_size=3, padding=1),
            nn.BatchNorm2d(128),
            nn.LeakyReLU(0.2),
            nn.Conv2d(128, 128, kernel_size=3, stride=2, padding=1),
            nn.BatchNorm2d(128),
            nn.LeakyReLU(0.2),
            nn.Conv2d(128, 256, kernel_size=3, padding=1),
            nn.BatchNorm2d(256),
            nn.LeakyReLU(0.2),
            nn.Conv2d(256, 256, kernel_size=3, stride=2, padding=1),
            nn.BatchNorm2d(256),
            nn.LeakyReLU(0.2),
            nn.Conv2d(256, 512, kernel_size=3, padding=1),
            nn.BatchNorm2d(512),
            nn.LeakyReLU(0.2),
            nn.Conv2d(512, 512, kernel_size=3, stride=2, padding=1),
            nn.BatchNorm2d(512),
            nn.LeakyReLU(0.2),
            nn.AdaptiveAvgPool2d(1),
            nn.Conv2d(512, 1024, kernel_size=1),
            nn.LeakyReLU(0.2),
            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))