简单训练了一个模型，可以实现超分辨率效果。模型在这里。
模型用了一些卷积层，最后接一个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 npcrop_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 xclass 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, labdef train():num_epochs = 2model = 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.0for 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.0out_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()
效果如下：
原图：
结果：
