使用 pytorch ImageFolder 过滤类/子文件夹

Question

这是我的文件夹结构

\n

image-folders/\n   \xe2\x94\x9c\xe2\x94\x80\xe2\x94\x80 class_0/\n   |   \xe2\x94\x9c\xe2\x94\x80\xe2\x94\x80 001.jpg\n   |   \xe2\x94\x9c\xe2\x94\x80\xe2\x94\x80 002.jpg\n   \xe2\x94\x94\xe2\x94\x80\xe2\x94\x80 class_1/\n   |   \xe2\x94\x9c\xe2\x94\x80\xe2\x94\x80 001.jpg\n   |   \xe2\x94\x94\xe2\x94\x80\xe2\x94\x80 002.jpg\n   \xe2\x94\x94\xe2\x94\x80\xe2\x94\x80 class_2/\n       \xe2\x94\x9c\xe2\x94\x80\xe2\x94\x80 001.jpg\n       \xe2\x94\x94\xe2\x94\x80\xe2\x94\x80 002.jpg\n

\n

通过使用ImageFoldertorchvision，我可以使用以下语法创建数据集：
\ndataset = ImageFolder("image-folders",...)

\n

但这将读取整个子文件夹并创建 3 个目标类。我不想包含 class_2 文件夹，我希望我的数据集仅包含 class_0 和 class_1，除了删除/移动 class_2 文件夹之外，还有什么方法可以实现此目的？

\n

Answer 1

您可以使用torch.utils.data.Subset原始完整ImageFolder数据集来完成此操作：

from torchvision.datasets import ImageFolder
from torch.utils.data import Subset

# construct the full dataset
dataset = ImageFolder("image-folders",...)
# select the indices of all other folders
idx = [i for i in range(len(dataset)) if dataset.imgs[i][1] != dataset.class_to_idx['class_s']]
# build the appropriate subset
subset = Subset(dataset, idx)