Learning a Deep Convolutional Network for Image Super-Resolution論文分析與pytorch代碼

論文地址

簡介

超分辨率技術（Super-Resolution）是指從觀測到的低分辨率圖像重建出相應的高分辨率圖像。
SR可分爲兩類:從多張低分辨率圖像重建出高分辨率圖像和從單張低分辨率圖像重建出高分辨率圖像。基於深度學習的SR，主要是基於單張低分辨率的重建方法，即Single Image Super-Resolution (SISR)。該論文創建了一種深度學習的方法，來實現單張低分辨率圖像的重建。
SISR是一個逆問題，對於一個低分辨率圖像，可能存在許多不同的高分辨率圖像與之對應，因此通常在求解高分辨率圖像時會加一個先驗信息進行規範化約束。在傳統的方法中，這個先驗信息可以通過若干成對出現的低-高分辨率圖像的實例中學到。而基於深度學習的SR通過神經網絡直接學習分辨率圖像到高分辨率圖像的端到端的映射函數。

模型圖

模型框架

SRCNN算法的框架，SRCNN將深度學習與傳統稀疏編碼之間的關係作爲依據，將3層網絡劃分爲圖像塊提取和表徵(Patch extraction and representation)、非線性映射(Non-linear mapping)以及最終的重建(Reconstruction)。

SRCNN的流程爲：

（1）先將低分辨率圖像使用雙立次差值放大至目標尺寸（如放大至2倍、3倍、4倍），此時仍然稱放大至目標尺寸後的圖像爲低分辨率圖像(Low-resolution image)，即圖中的輸入(input)；

（2）將低分辨率圖像輸入三層卷積神經網絡，（舉例：在論文中的其中一實驗相關設置，對YCrCb顏色空間中的Y通道進行重建，網絡形式爲(conv1+relu1)—(conv2+relu2)—(conv3)）
第一層卷積：卷積核尺寸9×9 (f1×f1)，卷積核數目64 (n1)，輸出64張特徵圖；
第二層卷積：卷積核尺寸1×1(f2×f2)，卷積核數目32(n2)，輸出32張特徵圖；
第三層卷積：卷積核尺寸5×5(f3×f3)，卷積核數目1(n3)，輸出1張特徵圖即爲最終重建高分辨率圖像。

我們首先使用雙立方插值將其放大到所需的大小，這是我們執行的唯一預處理。讓我們將插值圖像表示爲Y.我們的目標是從Y中恢復與地面實況高分辨率圖像X儘可能相似的圖像F（Y）。爲了便於呈現，我們仍稱Y爲“低分辨率的“圖像，雖然它與X具有相同的大小。我們希望學習映射F，它在概念上由三個操作組成：

1）補丁提取和表示：該操作從低分辨率圖像Y中提取（重疊）補丁，並將每個補丁表示爲高維向量。這些向量包括一組特徵圖，其數量等於向量的維數。
2）非線性映射：該操作是每個高維向量到另一個高維向量的非線性映射。每個映射的矢量在概念上是高分辨率補丁的表示。這些向量包括另一組特徵映射。
3）重建：該操作聚合高分辨率補丁表示以生成最終的高分辨率圖像。該圖像應該與真實標記X相似。

算法流程

Patch extraction and representation

圖像恢復中的流行策略是密集地提取補丁，然後通過一組預先訓練的基礎（例如PCA，DCT，Haar等）來表示它們。這相當於用一組濾波器對圖像進行卷積，每個濾波器都是一組基。將這些基礎的優化包括在網絡的優化中。第一層表示爲操作F1 ： F1​(Y)=max( 0, W1​∗Y+B1​ )
W1：表示濾波器，大小爲c × f1​ × f1 ​× n1
B1：表示偏置， 是n1維向量
c：輸入圖像的通道數
f1：濾波器的空間大小
n1：濾波器的個數。
很明顯，W1在圖像上應用n1​個卷積，每個卷積的內核大小爲c×f1×f1
輸出由n1​個特徵圖組成
B1是n1​維向量，其每個元素與濾波器相關聯。
我們在濾波器響應上應用整流線性單元（ReLU，max（0，x））

non-linear mapping 非線性映射

第一層爲每個補丁提取n1​維特徵。在第二操作中，我們將這些n1​維向量中的每一個映射爲n2​維向量。這相當於應用具有平凡空間支持1 x 1的n2​濾波器。
此解釋僅適用於1 x 1過濾器。但是很容易推廣到像3 x 3或5 x 5這樣的大型濾波器 。在那種情況下，非線性映射不是輸入圖像的補丁; 相反，它是在3 x 3或5 x 5特徵圖的“補丁”。第二層的操作是：F2(Y)=max(0,W2∗F1(Y)+B2)
W2：表示濾波器，大小爲​n1×1×1×n2
B2​：表示偏置，大小是n2​維。
每個輸出n2​維向量在概念上是將用於重建的高分辨率補丁的表示。

Reconstruction

在傳統方法中，經常對預測的重疊高分辨率補丁進行平均以產生最終的完整圖像。平均可以被認爲是一組特徵圖上的預定義濾波器（其中每個位置是高分辨率補片的“扁平”矢量形式）。由此推動，我們定義卷積層以產生最終的高分辨率圖像：F(Y)=W3∗F2(Y)+B3
W3：表示濾波器​，大小是n2×f3×f3×c
B3：​是c維矢量。

訓練

（1）訓練數據集：論文中某一實驗採用91張自然圖像作爲訓練數據集，對訓練集中的圖像先使用雙三次差值縮小到低分辨率尺寸，再將其放大到目標放大尺寸，最後切割成諸多33×33圖像塊作爲訓練數據，作爲標籤數據的則爲圖像中心的21×21圖像塊（與卷積層細節設置相關）；
（2）損失函數：採用MSE函數作爲卷積神經網絡損失函數；

（3）卷積層細節設置：第一層卷積核9×9，得到特徵圖尺寸爲(33-9)/1+1=25，第二層卷積核1×1，得到特徵圖尺寸不變，第三層卷積核5×5，得到特徵圖尺寸爲(25-5)/1+1=21。訓練時得到的尺寸爲21×21，因此圖像中心的21×21圖像塊作爲標籤數據。（卷積訓練時不進行padding）

測試

（1）全卷積網絡：所用網絡爲全卷積網絡，因此作爲實際測試時，直接輸入完整圖像即可；
（2）Padding：訓練時得到的實際上是除去四周(33-21)/2=6像素外的圖像，若直接採用訓練時的設置（無padding），得到的圖像最後會減少四周各6像素（如插值放大後輸入512×512，輸出500×500）。因此在測試時每一層卷積都進行了padding（卷積核尺寸爲1×1的不需要進行padding）。這樣保證插值放大後輸入與輸出尺寸的一致性。

實驗結果

Pytorch代碼實現

由於沒有GPU，在Github上找了相關較爲簡便的代碼加以修改，並不是完全對論文的復現，而是通過其思路實現低分辨率圖像輸入，得到對應的高分辨率圖像輸出。

使用說明

To train the model with a zoom factor of 2, for 200 epochs and on GPU:

python main.py --zoom_factor 2 --nb_epoch 200 --cuda

At each epoch, a .pth model file will be saved.

To use the model on an image: (the zoom factor must be the same the one used to train the model)

python run.py --zoom_factor 2 --model model_199.pth --image example.jpg --cuda

文件存放

運行代碼

model.py

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


class SRCNN(nn.Module):
    def __init__(self):
        super(SRCNN, self).__init__()
        self.conv1 = nn.Conv2d(1, 64, kernel_size=9, padding=4)  # color_channel=1, n1=64,f1=9,
        self.conv2 = nn.Conv2d(64, 32, kernel_size=1, padding=0)  # channel=64,n2=64,f2=9,
        self.conv3 = nn.Conv2d(32, 1, kernel_size=5, padding=2)  # channel=32,n3=1,f3=9,
        
    def forward(self, x):
        out = F.relu(self.conv1(x))
        out = F.relu(self.conv2(out))
        out = self.conv3(out)

        return out

data.py

from os import listdir
from os.path import join

from torch.utils.data import Dataset
import torchvision.transforms as transforms
from PIL import Image, ImageFilter

def is_image_file(filename):
    return any(filename.endswith(extension) for extension in [".png", ".jpg", ".jpeg"])


def load_img(filepath):
    img = Image.open(filepath).convert('YCbCr')
    y, _, _ = img.split()
    return y


CROP_SIZE = 32


class DatasetFromFolder(Dataset):
    def __init__(self, image_dir, zoom_factor):
        super(DatasetFromFolder, self).__init__()
        self.image_filenames = [join(image_dir, x) for x in listdir(image_dir) if is_image_file(x)]

        crop_size = CROP_SIZE - (CROP_SIZE % zoom_factor) # Valid crop size
        self.input_transform = transforms.Compose([transforms.CenterCrop(crop_size),  # cropping the image
                                      transforms.Resize(crop_size//zoom_factor),  # subsampling the image (half size)
                                      transforms.Resize(crop_size, interpolation=Image.BICUBIC),  # bicubic upsampling to get back the original size
                                      transforms.ToTensor()])
        self.target_transform = transforms.Compose([transforms.CenterCrop(crop_size),  # since it's the target, we keep its original quality
                                       transforms.ToTensor()])

    def __getitem__(self, index):
        input = load_img(self.image_filenames[index])
        target = input.copy()
        
        # input = input.filter(ImageFilter.GaussianBlur(1)) 
        input = self.input_transform(input)
        target = self.target_transform(target)

        return input, target

    def __len__(self):
        return len(self.image_filenames)

main.py

import argparse
from math import log10

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader

from data import DatasetFromFolder
from model import SRCNN

parser = argparse.ArgumentParser(description='SRCNN training parameters')
parser.add_argument('--zoom_factor', type=int, required=True)
parser.add_argument('--nb_epochs', type=int, default=200)
parser.add_argument('--cuda', action='store_true')
args = parser.parse_args()

device = torch.device("cuda:0" if (torch.cuda.is_available() and args.cuda) else "cpu")
torch.manual_seed(0)
torch.cuda.manual_seed(0)

# Parameters
BATCH_SIZE = 4
NUM_WORKERS = 0  # on Windows, set this variable to 0

trainset = DatasetFromFolder("F:\\Py_WorkSpace\\papers_code\\SRCNN-master\\data\\train", zoom_factor=args.zoom_factor)
testset = DatasetFromFolder("F:\\Py_WorkSpace\\papers_code\\SRCNN-master\\data\\test", zoom_factor=args.zoom_factor)

trainloader = DataLoader(dataset=trainset, batch_size=BATCH_SIZE, shuffle=True, num_workers=NUM_WORKERS)
testloader = DataLoader(dataset=testset, batch_size=BATCH_SIZE, shuffle=False, num_workers=NUM_WORKERS)

model = SRCNN().to(device)
criterion = nn.MSELoss()
optimizer = optim.Adam(  # we use Adam instead of SGD like in the paper, because it's faster
    [
        {"params": model.conv1.parameters(), "lr": 0.0001},  
        {"params": model.conv2.parameters(), "lr": 0.0001},
        {"params": model.conv3.parameters(), "lr": 0.00001},
    ], lr=0.00001,
)

for epoch in range(args.nb_epochs):

    # Train
    epoch_loss = 0
    for iteration, batch in enumerate(trainloader):
        input, target = batch[0].to(device), batch[1].to(device)
        optimizer.zero_grad()

        out = model(input)
        loss = criterion(out, target)
        loss.backward()
        optimizer.step()
        epoch_loss += loss.item()

    print(f"Epoch {epoch}. Training loss: {epoch_loss / len(trainloader)}")

    # Test
    avg_psnr = 0
    with torch.no_grad():
        for batch in testloader:
            input, target = batch[0].to(device), batch[1].to(device)

            out = model(input)
            loss = criterion(out, target)
            psnr = 10 * log10(1 / loss.item())
            avg_psnr += psnr
    print(f"Average PSNR: {avg_psnr / len(testloader)} dB.")

    # Save model
    torch.save(model, f"model_{epoch}.pth")

run.py

import argparse

import numpy as np
import torch
import torchvision.transforms as transforms
from PIL import Image

parser = argparse.ArgumentParser(description='SRCNN run parameters')
parser.add_argument('--model', type=str, required=True)
parser.add_argument('--image', type=str, required=True)
parser.add_argument('--zoom_factor', type=int, required=True)
parser.add_argument('--cuda', action='store_true')
args = parser.parse_args()

img = Image.open(args.image).convert('YCbCr')
img = img.resize((int(img.size[0]*args.zoom_factor), int(img.size[1]*args.zoom_factor)), Image.BICUBIC)   # first, we upscale the image via bicubic interpolation
img.save("F:\\Py_WorkSpace\\papers_code\\SRCNN-master\\bicubic_img.jpg")
y, cb, cr = img.split()

img_to_tensor = transforms.ToTensor()
input = img_to_tensor(y).view(1, -1, y.size[1], y.size[0])  # we only work with the "Y" channel

device = torch.device("cuda:0" if (torch.cuda.is_available() and args.cuda) else "cpu")
print(device)
model = torch.load(args.model).to(device)
input = input.to(device)

out = model(input)
out = out.cpu()
out_img_y = out[0].detach().numpy()
out_img_y *= 255.0
out_img_y = out_img_y.clip(0, 255)
out_img_y = Image.fromarray(np.uint8(out_img_y[0]), mode='L')

out_img = Image.merge('YCbCr', [out_img_y, cb, cr]).convert('RGB')  # we merge the output of our network with the upscaled Cb and Cr from before
                                                                    # before converting the result in RGB
out_img.save("F:\\Py_WorkSpace\\papers_code\\SRCNN-master\\zoomed_img.jpg")

運行操作

打開Anaconda Prompt,進入虛擬環境VE,執行python文件

圖片對比

Original image

Bicubic image

SRCNN image

後續工作

1.代碼中 torch.save(model, f"model_{epoch}.pth")不知道.pth文件存在哪了,希望有大神可以告知
2.進一步理解原理,並總結
3.對代碼模塊進行理解,熟悉pytorch模型,argparse模塊等,並總結

參考文章

(SRCNN)及pytorch實現_Learning a Deep Convolutional Network for Image Super-Resolution
SRCNN

。