PhenoMENON2025/ScaleNet-AIUpscaler

1from torch.utils.data import Dataset2from torchvision import transforms3from PIL import Image4import os5 6class SRDataset(Dataset):7    def __init__(self, lr_dirs, hr_dir, resize_to=(1920, 1080), file_types=('.png', '.jpg', '.jpeg')):8        if isinstance(lr_dirs, str):9            lr_dirs = [lr_dirs]10 11        self.lr_paths = []12        self.hr_dir = hr_dir13        self.resize_to = resize_to14        self.to_tensor = transforms.ToTensor()15 16        for lr_dir in lr_dirs:17            for fname in os.listdir(lr_dir):18                if fname.lower().endswith(file_types):19                    self.lr_paths.append((os.path.join(lr_dir, fname), fname))20 21    def __len__(self):22        return len(self.lr_paths)23 24    def __getitem__(self, idx):25        lr_path, fname = self.lr_paths[idx]26        hr_path = os.path.join(self.hr_dir, fname)27 28        lr = Image.open(lr_path).convert("RGB")29        hr = Image.open(hr_path).convert("RGB")30 31        if self.resize_to:32            hr = hr.resize(self.resize_to, Image.BICUBIC)33            lr = lr.resize((self.resize_to[0] // 4, self.resize_to[1] // 4), Image.BICUBIC)34 35        return self.to_tensor(lr), self.to_tensor(hr)36