Add UnlabeledImageDataset

docs/source/datasets.rst

@@ -180,6 +180,7 @@ Base classes for custom datasets
 
     DatasetFolder
     ImageFolder
+    UnlabeledImageDataset
     VisionDataset
 
 Transforms v2

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

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",
 )
 
 

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}"