I would also consider going one level above the ImageFolder Class which inherits from DatasetFolder.
DatasetFolder uses a method to index the folder subdirectories for each class:
def find_classes(directory: str) -> Tuple[List[str], Dict[str, int]]:
"""Finds the class folders in a dataset.
See :class:`DatasetFolder` for details.
"""
classes = sorted(entry.name for entry in os.scandir(directory) if entry.is_dir())
if not classes:
raise FileNotFoundError(f"Couldn't find any class folder in {directory}.")
class_to_idx = {cls_name: i for i, cls_name in enumerate(classes)}
return classes, class_to_idx
See here: torchvision.datasets.folder — Torchvision 0.10.0 documentation
You may simply create your own DatasetFolder (which inherits from VisionDataset, don’t forget to inherit from that) and then let your own Image Folder class inherit from DatasetFolder (if even needed).
By creating your own DatasetFolder, create a new find_classes method, which only scans for subdirectories in your dir, with your desired class name
def find_classes(directory: str, desired_class_names: List) -> Tuple[List[str], Dict[str, int]]:
"""Finds the class folders in a dataset."""
.......
Hope that makes sense and helps! Also just came across this question and tomorrow I am going to solve it!