使用 pytorch ImageFolder 过滤类/子文件夹 [英] filter class/subfolder with pytorch ImageFolder

问题描述

这是我的文件夹结构

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

通过使用来自 torchvision 的 ImageFolder，我可以使用以下语法创建数据集:
dataset = ImageFolder("image-folders",...)

By using ImageFolder from torchvision, I can create dataset with this syntax :
dataset = ImageFolder("image-folders",...)

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

But this will read the entire subfolder and create 3 target classes. I don't want to include the class_2 folder, I want my dataset to only contains class_0 and class_1 only, is there any way to achieve this besides delete/move the class_2 folder?

推荐答案

您可以使用 torch.utils.data.Subset:

You can do this by using torch.utils.data.Subset of the original full ImageFolder dataset:

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)

这篇关于使用 pytorch ImageFolder 过滤类/子文件夹的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！