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

Question

这是我的文件夹结构

image-folders/
   ├── class_0/
   |   ├── 001.jpg
   |   ├── 002.jpg
   └── class_1/
   |   ├── 001.jpg
   |   └── 002.jpg
   └── class_2/
       ├── 001.jpg
       └── 002.jpg

通过使用 ImageFolder从 torchvision，我可以使用以下语法创建数据集: dataset = ImageFolder("image-folders",...)但这将读取整个子文件夹并创建 3 个目标类。我不想包含 class_2 文件夹，我希望我的数据集只包含 class_0 和 class_1，除了删除/移动 class_2 文件夹之外还有什么方法可以实现吗？

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)