PhenoMENON2025/ScaleNet-AIUpscaler

1import torch2from torch.utils.data import DataLoader, Dataset3from torchvision.utils import save_image4from torchvision import transforms5from fullsrgan import load_generator6from PIL import Image7import os8 9 10class SRDataset(Dataset):11    def __init__(self, lr_dir, resize_to=(1280, 720)):12        self.lr_dir = lr_dir13        self.filenames = [14            f for f in os.listdir(lr_dir)15            if f.lower().endswith(('.png', '.jpg', '.jpeg'))16        ]17        self.to_tensor = transforms.ToTensor()18        self.resize_to = resize_to19 20    def __len__(self):21        return len(self.filenames)22 23    def __getitem__(self, idx):24        filename = self.filenames[idx]25        lr = Image.open(os.path.join(self.lr_dir, filename)).convert('RGB')26        lr = lr.resize(self.resize_to, Image.BICUBIC)27        return self.to_tensor(lr), filename28 29 30 31lr_dir = 'data/lr720-1080ver'32output_dir = 'outputs_test'33os.makedirs(output_dir, exist_ok=True)34 35device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')36 37 38model = load_generator(39    weights_path='srgan_generator4.pth',40    upscale=4,41    device=device,42    refinement=True43)44 45 46dataset = SRDataset(lr_dir, resize_to=(1280, 720))47loader = DataLoader(dataset, batch_size=1, shuffle=False)48 49 50with torch.no_grad():51    for i, (lr_img, filename) in enumerate(loader):52        lr_img = lr_img.to(device)53        sr_img = model(lr_img)54        save_path = os.path.join(output_dir, f"sr_{filename[0]}")55        save_image(sr_img, save_path)56 57print(f" Testing complete. Saved results to: {output_dir}")58