class TestDataset(Dataset): def __init__(self, folder_path, hr, lr): super(TestDataset, self).__init__() self.folder_path = folder_path self.hr = hr self.lr = lr self.hr_path = os.path.join(self.folder_path, self.hr) self.lr_path = os.path.join(self.folder_path, self.lr) self.hr_list = [entry.name for entry in os.scandir(self.hr_path) if entry.is_file()] self.lr_list = [entry.name for entry in os.scandir(self.lr_path) if entry.is_file()] self.filenames = self.hr_list self.transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) def __getitem__(self, index): self.hr_name = self.hr_list[index] self.lr_name = self.lr_list[index] self.hr_item_path = os.path.join(self.hr_path, self.hr_name) self.lr_item_path = os.path.join(self.lr_path, self.lr_name) hr_img = Image.open(self.hr_item_path).convert("RGB") lr_img = Image.open(self.lr_item_path).convert("RGB") hr_img = self.transform(hr_img) lr_img = self.transform(lr_img) return lr_img, hr_img def __len__(self): return len(self.filenames)
时间: 2023-03-29 14:01:32 浏览: 211
这是一个 Python 代码段,用于定义一个名为 TestDataset 的数据集类。该类接受三个参数:文件夹路径、高分辨率图像文件夹名称和低分辨率图像文件夹名称。它通过扫描指定文件夹中的文件来获取高分辨率和低分辨率图像的文件名列表,并将其存储在 hr_list 和 lr_list 中。在 __getitem__ 方法中,它会打开指定文件夹中的图像文件,并将其转换为张量,并返回低分辨率和高分辨率图像的张量。在 __len__ 方法中,它返回文件名列表的长度。
阅读全文