如何将 PyTorch 中的数据集同时合并为 return 图像和 numpy 文件

Question

我正在尝试构建一个可以拍摄图像和姿势的数据加载器。图像以 .jpg 文件的形式保存，姿势以 .npy 文件的形式保存。图像和姿势位于不同的文件夹中，但具有相同的子文件夹结构和名称。子文件夹的形式是classes，即每个class都有一个对应的文件夹。我想应用图像转换，然后 return 图像（我正在使用 torchvision datasets.ImageFolder）。对于姿势，我使用 torchvision datasets.DatasetFolder。如何组合这两个数据集，以便同时获得同名的姿势和图像？

class ReIDFolder_images(datasets.ImageFolder):

    def __init__(self, root, transform):
        super().__init__(root, transform)
        targets = np.asarray([s[1] for s in self.samples])
        self.targets = targets
        self.img_num = len(self.samples)
        print(self.img_num)

    def _get_cam_id(self, path):
        camera_id = []
        filename = os.path.basename(path)
        camera_id = filename.split('c')[1][0]
        return int(camera_id)-1

    def _get_pos_sample(self, target, index, path):
        pos_index = np.argwhere(self.targets == target)
        pos_index = pos_index.flatten()
        pos_index = np.setdiff1d(pos_index, index)
        if len(pos_index)==0:  # in the query set, only one sample
            return path
        else:
            rand = random.randint(0,len(pos_index)-1)
        return self.samples[pos_index[rand]][0]

    def _get_neg_sample(self, target):
        neg_index = np.argwhere(self.targets != target)
        neg_index = neg_index.flatten()
        rand = random.randint(0,len(neg_index)-1)
        return self.samples[neg_index[rand]]

    def __getitem__(self, index):
        path, target = self.samples[index]
        sample = self.loader(path)

        pos_path = self._get_pos_sample(target, index, path)
        pos = self.loader(pos_path)

        if self.transform is not None:
            sample = self.transform(sample)
            pos = self.transform(pos)

        if self.target_transform is not None:
            target = self.target_transform(target)

        return sample, target, pos

class ReIDFolder_poses(datasets.DatasetFolder):

    def __init__(self, root):
        super().__init__(root, loader=self.npy_loader, extensions='.npy')

        targets = np.asarray([s[1] for s in self.samples])
        self.targets = targets  
        self.img_num = len(self.samples)
        print(self.img_num)

    def npy_loader(self, path):
        sample = torch.Tensor(np.load(path))
        return sample

    def _get_cam_id(self, path):
        camera_id = []
        filename = os.path.basename(path)
        camera_id = filename.split('c')[1][0]
        return int(camera_id)-1

    def _get_pos_sample(self, target, index, path):
        pos_index = np.argwhere(self.targets == target)
        pos_index = pos_index.flatten()
        pos_index = np.setdiff1d(pos_index, index)
        if len(pos_index)==0:  # in the query set, only one sample
            return path
        else:
            rand = random.randint(0,len(pos_index)-1)
        return self.samples[pos_index[rand]][0]

    def _get_neg_sample(self, target):
        neg_index = np.argwhere(self.targets != target)
        neg_index = neg_index.flatten()
        rand = random.randint(0,len(neg_index)-1)
        return self.samples[neg_index[rand]]

    def __getitem__(self, index):
        path, target = self.samples[index]
        sample = self.loader(path)

        pos_path = self._get_pos_sample(target, index, path)
        pos = self.loader(pos_path)

        return sample, target, pos

Answer 1

我能够解决这个问题！原来我不必继承datasets.DatasetFolder。由于标签相同，我只是创建了一个继承 datasets.ImageFolder 的 class，并将修改后的路径提供给函数 npy_loader。

Answer 2

如何使用 ConcatDatset，以便您可以使用 dataset.ImageFolder

连接两个数据集

如何将 PyTorch 中的数据集同时合并为 return 图像和 numpy 文件

How to combine datasets in PyTorch to return image and numpy file simultaneously

python

computer-vision

pytorch

torchvision