testx8.py

from torchvision.transforms import Compose, ToTensor, Resize
from torchvision.transforms.functional import to_pil_image
from PIL import Image
import torch
import matplotlib.pyplot as plt
import matplotlib
matplotlib.use('TkAgg')
from EDSR_model import EDSR

# 自动选择推理设备
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 加载训练好的模型参数
model = EDSR(
    num_channels=1, # 输入通道为单通道
    base_channel=32, # 残差模块通道数
    upscale_factor=8, # 超分放大倍数
    num_residuals=16 # 残差模块数量
).to(device)
model.load_state_dict(torch.load('weights/EDSRx8.pth', map_location=device))

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

transform = Compose([
            # Resize((32, 32)),
            ToTensor(),
        ])

@torch.no_grad()
def infer(model, img_path):
    y, cb, cr = load_img(img_path)
    y = transform(y)
    y = y.unsqueeze(0)
    y = y.to(device)
    y = model(y)# 亮度通道Y进行超分辨率重建
    y = y.detach().cpu().squeeze(0)
    y = y.clamp(0, 1)
    y = to_pil_image(y).convert('L')

    cb = cb.resize(y.size, Image.BICUBIC)# 调整色度通道的大小，与超分后的亮度通道尺寸匹配
    cr = cr.resize(y.size, Image.BICUBIC)# 合并所有通道并转换为RGB格式
    output = Image.merge('YCbCr', [y, cb, cr]).convert('RGB')
    return output



input_file = 'data/test/LR/0845x2.png'
target_file = 'data/train/HR/0845.png'
predict = infer(model, input_file) # 将目标图片输入模型进行超分辨率预测
print('超分完成')
plt.imshow(predict)
predict.save('result/SR1.png')


if __name__ == '__main__':
    import sys
    import os

    if len(sys.argv) != 3:
        print("Usage: python testx8.py <input_path> <output_path>")
        sys.exit(1)

    input_file = sys.argv[1]
    output_path = sys.argv[2]

    # Create output directory if it doesn't exist
    os.makedirs(os.path.dirname(output_path), exist_ok=True)

    # Process the image
    predict = infer(model, input_file)
    predict.save(output_path)
    print(f"Super-resolution completed. Image saved to {output_path}")