如何将标签从DataFrame分配给PyTorch ImageFolder数据集

Question

我正在使用ImageFolder从文件夹中加载图像，我想为DataFrame中的图像分配标签。 DataFrame在文件夹中包含图像名称并标记图像

文件夹中的所有图像均按字母排序。

我的代码：

def load_data(data_path, targets):
    train_data = torchvision.datasets.ImageFolder(
        root=data_path,
        transform=torchvision.transforms.ToTensor()
)
    train_data.classes = ['FAKE', 'REAL']
    train_data.class_to_idx = {k: v for k, v in zip(train_data.classes, [0, 1])}
    train_data.targets = targets
    return train_data

metadata = metadata.sort_values(by='image_name') # metadata it's my dataframe with labels
targets = np.array(metadata.label.astype('category').cat.codes)
# And i'm try to reshape my array with labels
# targets = targets.reshape(targets.shape[0], 1)
targets = torch.from_numpy(targets)

train_dataset = load_data('../input/deepfake-faces/', targets)

例如，在我的数据集中，此代码的预期输出应为“ Label：1”，但我得到0。

img, label = train_dataset[10]
print('Label: ', label)
plt.imshow(img.permute(1, 2, 0))

但是，当我计算带有标签1的图像时，我得到了正确的结果。

np.count_nonzero(train_dataset.targets == 1)