image - 在 pytorch 中加载 csv 和图像数据集

Question

我正在使用 PyTorch 进行图像分类。我有一个单独的 Images 文件夹，并训练和测试带有图像 ids 和 labels 的 csv 文件。我不知道如何将这些图像和 ID 组合起来并转换为张量。

1. train.csv :包含图像的所有 ID，如 4325.jpg、2345.jpg 等，并包含标签，如猫、狗。
2. Image_data : 包含 ID 名称的所有图像。

Answer 1

您可以通过继承 pytorch 的 torch.utils.data.Dataset 创建自定义数据集类.

以下自定义数据集类的假设是

• csv文件格式为
  
  

<表类="s-表"><头>文件名标签<正文>4325.jpg猫2345.jpg狗

• 所有图片都在里面images folder .

class CustomDataset(torch.utils.data.Dataset):
    def __init__(self, csv_path, images_folder, transform = None):
        self.df = pd.read_csv(csv_path)
        self.images_folder = images_folder
        self.transform = transform
        self.class2index = {"cat":0, "dog":1}

    def __len__(self):
        return len(self.df)
    def __getitem__(self, index):
        filename = self.df[index, "FILENAME"]
        label = self.class2index[self.df[index, "LABEL"]]
        image = PIL.Image.open(os.path.join(self.images_folder, filename))
        if self.transform is not None:
            image = self.transform(image)
        return image, label
        

现在您可以使用此类使用 csv 文件和图像文件夹加载训练和测试数据集。

train_dataset = CustomDataset("path - to - train.csv", "path - to - images - folder"  )
test_dataset = CustomDataset("path - to - test.csv", "path - to - images - folder"  )


image, label = train_dataset[0]