小编Ami*_*lah的帖子

我使用 PyTorch 数据集类开发了一个自定义数据集。代码是这样的：

class CustomDataset(torch.utils.data.Dataset):

    def __init__(self, root_path, transform=None):
        self.path = root_path
        self.mean = mean
        self.std = std
        self.transform = transform
        self.images = []
        self.masks = []

        for add in os.listdir(self.path):
            # Some script to load file from directory and appending address to relative array
            ...

        self.masks.sort()
        self.images.sort()

    def __len__(self):
        return len(self.images)

    def __getitem__(self, item):
        image_address = self.images[item]
        mask_address = self.masks[item]



        if self.transform is not None:
            augment = self.transform(image=np.asarray(Image.open(image_address, 'r', None)),
                                     mask=np.asarray(Image.open(mask_address, 'r', None)))
            image = Image.fromarray(augment['image'])
            mask = augment['mask'] …