Pytorch 创建带有增强图像的数据集

Question

有没有办法通过创建 X 个增强图像来生成 DataLoader？我目前的代码只创建了一个增强图像

class ImageDataset(data.Dataset):
    
    def __init__(self, root_dir, transform=None):
        
        self.root_dir = root_dir
        self.img_names = os.listdir()
        self.transform = transform
        
    def __getitem__(self, index):
        
        img = Image.open(os.path.join(self.root, self.img_names[index])).convert('RGB')
        if self.transform is not None:
            img = self.transform(img)
            
        return img
            
    def __len__(self):
        
        return len(self.img_names)

另外，我想添加标签，其中来自同一图像的增强图像具有相同的标签

Answer 1

这里有几个等效的方法可以做到这一点。

选项 1

我们可以更改数据集类本身以多次提供相同的数据。这可以通过报告更长的长度并使用索引（mod 长度）选择图像名称来实现。

class ImageDataset(data.Dataset):
    def __init__(self, root_dir, repetitions=1, transform=None):
        self.root_dir = root_dir
        self.img_names = os.listdir()
        self.transform = transform
        self.repetitions = repetitions
        
    def __getitem__(self, index):
        img = Image.open(os.path.join(self.root,
            self.img_names[index % len(self.img_names)])).convert('RGB')
        if self.transform is not None:
            img = self.transform(img)
        return img
            
    def __len__(self):
        return len(self.img_names) * self.repetitions

选项 2

或者，我们可以使用 torch.utils.data.Subset 并多次指定相同的索引。

# using your original implementation for ImageDataset
dataset = ImageDataset(root, transforms)
dataset = torch.utils.data.Subset(dataset, list(range(len(dataset))) * repetitions)

---