| import os | |
| import pathlib | |
| from typing import Any, Callable, Optional, Tuple, Union | |
| from .folder import default_loader | |
| from .utils import download_and_extract_archive, verify_str_arg | |
| from .vision import VisionDataset | |
| class DTD(VisionDataset): | |
| """`Describable Textures Dataset (DTD) <https://www.robots.ox.ac.uk/~vgg/data/dtd/>`_. | |
| Args: | |
| root (str or ``pathlib.Path``): Root directory of the dataset. | |
| split (string, optional): The dataset split, supports ``"train"`` (default), ``"val"``, or ``"test"``. | |
| partition (int, optional): The dataset partition. Should be ``1 <= partition <= 10``. Defaults to ``1``. | |
| .. note:: | |
| The partition only changes which split each image belongs to. Thus, regardless of the selected | |
| partition, combining all splits will result in all images. | |
| 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`` | |
| target_transform (callable, optional): A function/transform that takes in the target and transforms it. | |
| download (bool, optional): If True, downloads the dataset from the internet and | |
| puts it in root directory. If dataset is already downloaded, it is not | |
| downloaded again. Default is False. | |
| loader (callable, optional): A function to load an image given its path. | |
| By default, it uses PIL as its image loader, but users could also pass in | |
| ``torchvision.io.decode_image`` for decoding image data into tensors directly. | |
| """ | |
| _URL = "https://www.robots.ox.ac.uk/~vgg/data/dtd/download/dtd-r1.0.1.tar.gz" | |
| _MD5 = "fff73e5086ae6bdbea199a49dfb8a4c1" | |
| def __init__( | |
| self, | |
| root: Union[str, pathlib.Path], | |
| split: str = "train", | |
| partition: int = 1, | |
| transform: Optional[Callable] = None, | |
| target_transform: Optional[Callable] = None, | |
| download: bool = False, | |
| loader: Callable[[Union[str, pathlib.Path]], Any] = default_loader, | |
| ) -> None: | |
| self._split = verify_str_arg(split, "split", ("train", "val", "test")) | |
| if not isinstance(partition, int) and not (1 <= partition <= 10): | |
| raise ValueError( | |
| f"Parameter 'partition' should be an integer with `1 <= partition <= 10`, " | |
| f"but got {partition} instead" | |
| ) | |
| self._partition = partition | |
| super().__init__(root, transform=transform, target_transform=target_transform) | |
| self._base_folder = pathlib.Path(self.root) / type(self).__name__.lower() | |
| self._data_folder = self._base_folder / "dtd" | |
| self._meta_folder = self._data_folder / "labels" | |
| self._images_folder = self._data_folder / "images" | |
| if download: | |
| self._download() | |
| if not self._check_exists(): | |
| raise RuntimeError("Dataset not found. You can use download=True to download it") | |
| self._image_files = [] | |
| classes = [] | |
| with open(self._meta_folder / f"{self._split}{self._partition}.txt") as file: | |
| for line in file: | |
| cls, name = line.strip().split("/") | |
| self._image_files.append(self._images_folder.joinpath(cls, name)) | |
| classes.append(cls) | |
| self.classes = sorted(set(classes)) | |
| self.class_to_idx = dict(zip(self.classes, range(len(self.classes)))) | |
| self._labels = [self.class_to_idx[cls] for cls in classes] | |
| self.loader = loader | |
| def __len__(self) -> int: | |
| return len(self._image_files) | |
| def __getitem__(self, idx: int) -> Tuple[Any, Any]: | |
| image_file, label = self._image_files[idx], self._labels[idx] | |
| image = self.loader(image_file) | |
| if self.transform: | |
| image = self.transform(image) | |
| if self.target_transform: | |
| label = self.target_transform(label) | |
| return image, label | |
| def extra_repr(self) -> str: | |
| return f"split={self._split}, partition={self._partition}" | |
| def _check_exists(self) -> bool: | |
| return os.path.exists(self._data_folder) and os.path.isdir(self._data_folder) | |
| def _download(self) -> None: | |
| if self._check_exists(): | |
| return | |
| download_and_extract_archive(self._URL, download_root=str(self._base_folder), md5=self._MD5) | |