From a66cb1e4c234c4eae184b717d650f920dd48181e Mon Sep 17 00:00:00 2001 From: Chinmay J Sultanpuri Date: Sat, 5 Sep 2026 00:52:52 +0530 Subject: [PATCH] Add UnlabeledImageDataset --- docs/source/datasets.rst | 1 + test/test_datasets.py | 85 ++++++++++++++++++++++++++++++++ torchvision/datasets/__init__.py | 3 +- torchvision/datasets/folder.py | 73 +++++++++++++++++++++++++++ 4 files changed, 161 insertions(+), 1 deletion(-) diff --git a/docs/source/datasets.rst b/docs/source/datasets.rst index 3caa7434e20..2e6cc520010 100644 --- a/docs/source/datasets.rst +++ b/docs/source/datasets.rst @@ -180,6 +180,7 @@ Base classes for custom datasets DatasetFolder ImageFolder + UnlabeledImageDataset VisionDataset Transforms v2 diff --git a/test/test_datasets.py b/test/test_datasets.py index 22c14cbc08d..2fec82c7442 100644 --- a/test/test_datasets.py +++ b/test/test_datasets.py @@ -1755,6 +1755,91 @@ def test_classes(self, config): assert all([a == b for a, b in zip(dataset.classes, info["classes"])]) +class UnlabeledImageDatasetTestCase(datasets_utils.ImageDatasetTestCase): + DATASET_CLASS = datasets.UnlabeledImageDataset + FEATURE_TYPES = (PIL.Image.Image,) # just the image, no target + + # Fake tree: 3 png (2 top-level, 1 nested), 1 jpg + 1 upper-case .JPG top-level, 1 jpeg + 1 bmp + # nested, and 1 .txt that must be ignored -> 7 images total. + _NUM_IMAGES = 7 + _NUM_PNG_RECURSIVE = 3 + _NUM_PNG_TOP_LEVEL = 2 + _UPPERCASE_EXT_FILE_NAME = "top_3.JPG" + + def inject_fake_data(self, tmpdir, config): + tmpdir = pathlib.Path(tmpdir) + nested = tmpdir / "nested" / "deep" + nested.mkdir(parents=True) + + # distinct stems so "top_3.JPG" doesn't collide with a .jpg on case-insensitive filesystems + for name in ["top_0.png", "top_1.png", "top_2.jpg", self._UPPERCASE_EXT_FILE_NAME]: + datasets_utils.create_image_file(tmpdir, name) + for name in ["nested_0.png", "nested_1.jpeg", "nested_2.bmp"]: + datasets_utils.create_image_file(nested, name) + (tmpdir / "not_an_image.txt").write_text("definitely not an image") + + return dict(num_examples=self._NUM_IMAGES) + + def test_default_patterns_are_recursive(self): + with self.create_dataset() as (dataset, info): + assert len(dataset) == info["num_examples"] == self._NUM_IMAGES + + def test_custom_patterns_restrict_discovery(self): + with self.create_dataset(patterns=["**/*.png"]) as (dataset, _): + assert len(dataset) == self._NUM_PNG_RECURSIVE + assert all(path.endswith(".png") for path in dataset.samples) + + def test_non_recursive_pattern(self): + with self.create_dataset(patterns=["*.png"]) as (dataset, _): + assert len(dataset) == self._NUM_PNG_TOP_LEVEL + + def test_uppercase_extensions_are_discovered(self): + with self.create_dataset() as (dataset, _): + assert any(pathlib.Path(path).name == self._UPPERCASE_EXT_FILE_NAME for path in dataset.samples) + + def test_samples_are_unique_and_sorted(self): + # On case-insensitive filesystems the lower- and upper-case default patterns match the same files. + with self.create_dataset() as (dataset, info): + assert len(dataset.samples) == info["num_examples"] == len(set(dataset.samples)) + assert dataset.samples == sorted(dataset.samples) + + def test_getitem_returns_bare_image(self): + with self.create_dataset() as (dataset, _): + sample = dataset[0] + assert isinstance(sample, PIL.Image.Image) + assert not isinstance(sample, tuple) + + def test_transform_is_applied_to_every_sample(self): + transform = unittest.mock.Mock(wraps=lambda image: image) + with self.create_dataset(transform=transform) as (dataset, info): + samples = list(dataset) + assert len(samples) == info["num_examples"] + assert transform.call_count == info["num_examples"] + + def test_custom_loader_is_respected(self): + loader = unittest.mock.Mock(wraps=datasets.folder.default_loader) + with self.create_dataset(loader=loader) as (dataset, _): + dataset[0] + loader.assert_called_once_with(dataset.samples[0]) + + def test_root_as_pathlib_path(self): + with self.create_dataset() as (dataset, info): + from_pathlib = datasets.UnlabeledImageDataset(pathlib.Path(dataset.root)) + assert len(from_pathlib) == info["num_examples"] + + def test_raises_when_nothing_is_found(self): + with self.create_dataset() as (dataset, _): + empty_dir = pathlib.Path(dataset.root) / "empty" + empty_dir.mkdir() + # missing root, existing but empty root, and a pattern that matches nothing + with pytest.raises(FileNotFoundError): + datasets.UnlabeledImageDataset(os.path.join(dataset.root, "i_do_not_exist")) + with pytest.raises(FileNotFoundError): + datasets.UnlabeledImageDataset(empty_dir) + with pytest.raises(FileNotFoundError): + datasets.UnlabeledImageDataset(dataset.root, patterns=["**/*.no_such_extension"]) + + class KittiTestCase(datasets_utils.ImageDatasetTestCase): DATASET_CLASS = datasets.Kitti FEATURE_TYPES = (PIL.Image.Image, (list, type(None))) # test split returns None as target diff --git a/torchvision/datasets/__init__.py b/torchvision/datasets/__init__.py index 4005b4a9072..eba8993abb6 100644 --- a/torchvision/datasets/__init__.py +++ b/torchvision/datasets/__init__.py @@ -25,7 +25,7 @@ from .fgvc_aircraft import FGVCAircraft from .flickr import Flickr30k, Flickr8k from .flowers102 import Flowers102 -from .folder import DatasetFolder, ImageFolder +from .folder import DatasetFolder, ImageFolder, UnlabeledImageDataset from .food101 import Food101 from .gtsrb import GTSRB from .hmdb51 import HMDB51 @@ -131,6 +131,7 @@ "ETH3DStereo", "wrap_dataset_for_transforms_v2", "Imagenette", + "UnlabeledImageDataset", ) diff --git a/torchvision/datasets/folder.py b/torchvision/datasets/folder.py index 387439c0433..fbc569eadd4 100644 --- a/torchvision/datasets/folder.py +++ b/torchvision/datasets/folder.py @@ -1,5 +1,6 @@ import os import os.path +from collections.abc import Sequence from pathlib import Path from typing import Any, Callable, cast, Optional, Union @@ -335,3 +336,75 @@ def __init__( allow_empty=allow_empty, ) self.imgs = self.samples + + +class UnlabeledImageDataset(VisionDataset): + """A generic data loader for a flat or nested directory of unlabeled images. + + Unlike :class:`ImageFolder`, images do not need to be organized into per-class subfolders, and + ``__getitem__`` returns just the (optionally transformed) image, not a ``(sample, target)`` tuple. + + Args: + root (str or ``pathlib.Path``): Root directory path. + patterns (str or sequence of str, optional): Glob pattern(s), relative to ``root`` and passed + to :meth:`pathlib.Path.glob`, used to discover the images. Defaults to ``None``, which + matches every :const:`IMG_EXTENSIONS` extension recursively, in both lower- and upper-case. + A file matched by several patterns is still listed once. + transform (callable, optional): A function/transform that takes in a PIL image or torch.Tensor, depends + on the given loader, and returns a transformed version. E.g, ``transforms.RandomCrop``. + loader (callable, optional): A function to load an image given its path. Defaults to + :func:`default_loader`. + + Attributes: + samples (list): Sorted list of the discovered image file paths. + """ + + def __init__( + self, + root: Union[str, Path], + patterns: Optional[Union[str, Sequence[str]]] = None, + transform: Optional[Callable] = None, + loader: Callable[[str], Any] = default_loader, + ) -> None: + super().__init__(root, transform=transform) + + self._default_patterns = patterns is None + if patterns is None: + # Path.glob() is case-sensitive on Linux, so also match e.g. "img.JPG"; duplicates are dropped below. + patterns = [f"**/*{ext}" for ext in IMG_EXTENSIONS] + [f"**/*{ext.upper()}" for ext in IMG_EXTENSIONS] + elif isinstance(patterns, str): + patterns = [patterns] + + root_path = Path(self.root) + if not root_path.is_dir(): + raise FileNotFoundError(f"The root directory {root_path} does not exist (or is not a directory).") + + samples = sorted({str(path) for pattern in patterns for path in root_path.glob(pattern) if path.is_file()}) + if not samples: + raise FileNotFoundError( + f"Found no image file in {root_path} matching any of the patterns {list(patterns)}." + ) + + self.loader = loader + self.patterns = list(patterns) + self.samples = samples + + def __getitem__(self, index: int) -> Any: + """ + Args: + index (int): Index + + Returns: + (Any): The image at ``index``, transformed by ``transform`` if one was given. No target is returned. + """ + sample = self.loader(self.samples[index]) + if self.transform is not None: + sample = self.transform(sample) + return sample + + def __len__(self) -> int: + return len(self.samples) + + def extra_repr(self) -> str: + patterns = "default (all IMG_EXTENSIONS, case-insensitive)" if self._default_patterns else self.patterns + return f"patterns={patterns}"