Python，类数据集，如何在pytorch中将图像与其各自的标签连接起来

Posted 2023-03-29

【中文标题】Python，类数据集，如何在pytorch中将图像与其各自的标签连接起来

【英文标题】：Python, class dataset, how to concatenate images with their respective labels in pytorch

【发布时间】：2020-10-07 13:39:57

【问题描述】：

我是 PyTorch 的新手，在过去的几天里，我一直在努力使用可让您构建自定义数据集的类 Dataset。

我正在使用这个数据集 (https://www.kaggle.com/ianmoone0617/flower-goggle-tpu-classification/kernels)，问题是图像及其标签位于不同的文件夹中，我不知道如何连接它们。

这是我正在使用的代码：

class MyDataset(Dataset):

    def __init__(self, csv_file, root_dir, transform=None):
        self.labels = pd.read_csv(csv_file)
        self.root_dir = root_dir
        self.transform = transform

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

    def __getitem__(self, index):
        if torch.is_tensor(index):
            index = index.tolist()

        image_name = os.path.join(self.root_dir, self.labels.iloc[index, 0])
        image = io.imread(image_name)

        if self.transform:
            image = self.transform(image)

        return (image, labels)

虽然文件夹的结构如下：

我真的很想了解这一点，所以提前谢谢你们！

【参考方案1】：

看来你快到了。有很多方法可以解决这个问题。例如，您可以在初始化期间读取两个 csv 文件以构建一个字典，将flowers_idx.csv 中的标签字符串映射到flowers_label.csv 中指定的标签索引。

import os
import pandas as pd
import torch
from torchvision.datasets.folder import default_loader
from torch.utils.data import Dataset

class MyDataset(Dataset):
    def __init__(self, data_csv, label_csv, root_dir, transform=None):
        self.data_entries = pd.read_csv(data_csv)
        self.root_dir = root_dir
        self.transform = transform

        label_map = pd.read_csv(label_csv)
        self.label_str_to_idx = label_str: label_idx for label_idx, label_str in label_map.iloc

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

    def __getitem__(self, index):
        if torch.is_tensor(index):
            index = index.item()

        label = self.label_str_to_idx[self.data_entries.iloc[index, 1]] 
        image_path = os.path.join(self.root_dir, f'self.data_entries.iloc[index, 0].jpeg')

        # torchvision datasets generally return PIL image rather than numpy ndarray
        image = default_loader(image_path)

        # alternative to load ndarray using skimage.io
        # image = io.imread(image_path)

        if self.transform:
            image = self.transform(image)

        return (image, label)

请注意，这将返回 PIL 图像而不是 ndarrays，因为这通常是 torchvision 数据集返回的内容。这也很好，因为许多 torchvision 变换只能应用于 PIL 图像。

现在一个简单的用例可能是：

import torchvision.transforms as tt

dataset_dir = '/home/jodag/datasets/527293_966816_bundle_archive'
# TODO add more transforms/data-augmentation etc...
transform = tt.Compose((
    tt.ToTensor(),
))

dataset = MyDataset(
    os.path.join(dataset_dir, 'flowers_idx.csv'),
    os.path.join(dataset_dir, 'flowers_label.csv'),
    os.path.join(dataset_dir, 'flower_tpu/flower_tpu/flowers_google/flowers_google'),
    transform)

image, label = dataset[0]

在训练或验证期间，您可能会使用 DataLoader 对数据集进行采样。

【讨论】：

嘿，非常感谢你，我非常感谢你！感谢您的建议，我了解了如何构建课程！最后一件事我尝试在不使用您的代码的情况下实现该类并且它有效，但最后当我尝试评估模型时，它给了我一个错误，说“元组对象没有属性大小”。我认为问题在于，在类数据集中，作为标签，我没有用字典映射它们，我只是使用了 flowers_idx.csv 中的标签，它们是字符串。有没有办法解决这个问题，以防我没有另一个文件，flowers_label.csv？