def __init__(self, root_dir, transform=None): self.root_dir = root_dir self.transform = transform self.vi_files = sorted(os.listdir(os.path.join(root_dir, 'vi'))) def __len__(self): return len(self.vi_files) def __getitem__(self, idx): vi_filename = self.vi_files[idx] vi_path = os.path.join(self.root_dir, 'vi', vi_filename) with open(vi_path, 'rb') as f: vi_image = Image.open(f).convert('RGB') if self.transform: vi_image = self.transform(vi_image) return vi_image

这是一个 PyTorch 中 Dataset 类的实现，用于读取一个包含图像数据的文件夹。其中，root_dir 是包含数据集的根目录，transform 是可选的数据增强操作，vi_files 是一个包含文件名的列表，通过 sorted 函数对其进行排序，便于按顺序加载数据。__len__() 方法返回数据集的大小，__getitem__() 方法实现了数据的加载和预处理操作。在这个实现中，每个图像被解码为 RGB 格式，并且根据 transform 执行一些数据增强操作（如果有的话），最后返回处理后的图像数据。