深度学习与图像超分辨率

import torch

import torch.nn as nn

import torch.optim as optim

from torchvision.transforms import ToTensor, Resize

from torch.utils.data import DataLoader, Dataset

from PIL import Image

import os

## 定义SRCNN网络结构

class SRCNN(nn.Module):

def __init__(self):

super(SRCNN, self).__init__()

self.conv1 = nn.Conv2d(1, 64, kernel_size=9, padding=4)

self.conv2 = nn.Conv2d(64, 32, kernel_size=5, padding=2)

self.conv3 = nn.Conv2d(32, 1, kernel_size=5, padding=2)

self.relu = nn.ReLU()

def forward(self, x):

x = self.relu(self.conv1(x))

x = self.relu(self.conv2(x))

x = self.conv3(x)

return x

## 创建自定义数据集类

class ImageDataset(Dataset):

def __init__(self, image_dir, upscale_factor):

self.image_filenames = [os.path.join(image_dir, x) for x in os.listdir(image_dir)]

self.upscale_factor = upscale_factor

self.transform_lr = Resize((int(28 / upscale_factor), int(28 / upscale_factor)), interpolation=Image.BICUBIC)

self.transform_hr = ToTensor()

def __getitem__(self, index):

hr_image = Image.open(self.image_filenames[index])

lr_image = self.transform_lr(hr_image)

hr_tensor = self.transform_hr(hr_image)

lr_tensor = self.transform_hr(lr_image)

return lr_tensor, hr_tensor

def __len__(self):

return len(self.image_filenames)

## 初始化参数

batch_size = 16

epochs = 100

learning_rate = 0.001

upscale_factor = 3

image_dir = 'path_to_images'

## 加载数据集

dataset = ImageDataset(image_dir=image_dir, upscale_factor=upscale_factor)

data_loader = DataLoader(dataset=dataset, batch_size=batch_size, shuffle=True)

## 实例化模型和优化器

model = SRCNN()

criterion = nn.MSELoss()

optimizer = optim.Adam(model.parameters(), lr=learning_rate)

## 训练循环

for epoch in range(epochs):

for data in data_loader:

lr_images, hr_images = data

# 前向传播

sr_images = model(lr_images)

loss = criterion(sr_images, hr_images)

# 反向传播

optimizer.zero_grad()

loss.backward()

optimizer.step()

print(f'Epoch [{epoch+1}/{epochs}], Loss: {loss.item():.4f}')

## 保存训练好的模型

torch.save(model.state_dict(), 'srcnn.pth')