深度学习（超分辨率）

简单训练了一个模型，可以实现超分辨率效果。模型在这里。

模型用了一些卷积层，最后接一个PixelShuffle算子。

训练数据是原始图像resize后的亮度通道。

标签是原始图像的亮度通道。

损失函数设为MSE。

代码如下：

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from torchvision.transforms import Compose, CenterCrop, ToTensor, Resize
from PIL import Image
from os import listdir
from os.path import join
import numpy as np

crop_size = 256
upscale_factor = 3
crop_size = crop_size - (crop_size % upscale_factor)

input_transformer= Compose([
        CenterCrop(crop_size),
        Resize(crop_size // upscale_factor),
        ToTensor()])

target_transform =Compose([
        CenterCrop(crop_size),
        ToTensor()])

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()

        self.relu = nn.ReLU()
        self.conv1 = nn.Conv2d(1, 64, 5, 1, 2)
        self.conv2 = nn.Conv2d(64, 64, 3, 1, 1)
        self.conv3 = nn.Conv2d(64, 32, 3, 1, 1)
        self.conv4 = nn.Conv2d(32, upscale_factor ** 2, 3, 1, 1)
        self.pixel_shuffle = nn.PixelShuffle(upscale_factor)

    def forward(self, x):
        x = self.relu(self.conv1(x))
        x = self.relu(self.conv2(x))
        x = self.relu(self.conv3(x))
        x = self.pixel_shuffle(self.conv4(x))       
        return x

class SRData(Dataset):
    def __init__(self, image_dir):
        self.image_filenames = [join(image_dir, x) for x in listdir(image_dir)]

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

    def __getitem__(self, index):
        image = Image.open(self.image_filenames[index]).convert('YCbCr')
        y, _, _ = image.split()

        img = input_transformer(y)
        lab = target_transform(y)
        return img, lab

def train():
    num_epochs = 2

    model = Net()
    optimizer = optim.Adam(model.parameters(), lr=0.01)
    criterion = nn.MSELoss()

    train_dataset = SRData('./dataset')
    train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    model.to(device)
    model.train()

    for epoch in range(num_epochs):
        running_loss = 0.0
        for images, labels in train_loader:

            images = images.to(device)
            labels = labels.to(device)

            outputs = model(images)
            loss = criterion(outputs, labels)

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            running_loss += loss.item()

        print(f"Epoch [{epoch+1}/{num_epochs}], Loss: {running_loss/len(train_loader):.4f}")

    torch.save(model, 'super_res.pth')

def test():

    img = Image.open("test.jpg").convert('YCbCr')
    y, cb, cr = img.split()

    model = torch.load("super_res.pth")
    img_to_tensor = ToTensor()
    input = img_to_tensor(y).view(1, 1, y.size[1], y.size[0])

    model = model.cuda()
    input = input.cuda()

    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_cb = cb.resize(out_img_y.size, Image.BICUBIC)
    out_img_cr = cr.resize(out_img_y.size, Image.BICUBIC)
    out_img = Image.merge('YCbCr', [out_img_y, out_img_cb, out_img_cr]).convert('RGB')

    out_img.save("out.jpg")

if __name__ == "__main__":
  #  train()
    test()