clarena.stl_datasets.oxford_iiit_pet

  1r"""
  2The submodule in `stl_datasets` for Oxford-IIIT Pet dataset.
  3"""
  4
  5__all__ = ["OxfordIIITPet"]
  6
  7import logging
  8from typing import Callable
  9
 10import torch
 11from torch.utils.data import Dataset, random_split
 12from torchvision.datasets import OxfordIIITPet as OxfordIIITPetRaw
 13from torchvision.transforms import transforms
 14
 15from clarena.stl_datasets.base import STLDatasetFromRaw
 16from clarena.stl_datasets.raw import OxfordIIITPet2, OxfordIIITPet37
 17
 18# always get logger for built-in logging in each module
 19pylogger = logging.getLogger(__name__)
 20
 21
 22class OxfordIIITPet(STLDatasetFromRaw):
 23    r"""Permuted Oxford-IIIT Pet dataset. The [Oxford-IIIT Pet dataset](https://www.robots.ox.ac.uk/~vgg/data/pets/) is a collection of cat and dog pictures. It consists of 7,349 images of 37 breeds (classes), each color image. It also provides a binary classification version with 2 classes (cat or dog). We support both versions in Permuted Oxford-IIIT Pet."""
 24
 25    def __init__(
 26        self,
 27        root: str,
 28        target_type: str,
 29        validation_percentage: float,
 30        batch_size: int = 1,
 31        num_workers: int = 0,
 32        custom_transforms: Callable | transforms.Compose | None = None,
 33        repeat_channels: int | None = None,
 34        to_tensor: bool = True,
 35        resize: tuple[int, int] | None = None,
 36    ) -> None:
 37        r"""
 38        **Args:**
 39        - **root** (`str`): the root directory where the original Oxford-IIIT Pet data 'OxfordIIITPet/' live.
 40        - **target_type** (`str`): the target type; one of:
 41            1. 'category': Label for one of the 37 pet categories.
 42            2. 'binary-category': Binary label for cat or dog.
 43        - **validation_percentage** (`float`): the percentage to randomly split some training data into validation data.
 44        - **batch_size** (`int`): The batch size in train, val, test dataloader.
 45        - **num_workers** (`int`): the number of workers for dataloaders.
 46        - **custom_transforms** (`transform` or `transforms.Compose` or `None`): the custom transforms to apply to ONLY TRAIN dataset. Can be a single transform, composed transforms or no transform. `ToTensor()`, normalize and so on are not included.
 47        - **repeat_channels** (`int` | `None`): the number of channels to repeat. Default is None, which means no repeat. If not None, it should be an integer.
 48        - **to_tensor** (`bool`): whether to include `ToTensor()` transform. Default is True.
 49        - **resize** (`tuple[int, int]` | `None` or list of them): the size to resize the images to. Default is None, which means no resize. If not None, it should be a tuple of two integers.
 50        """
 51
 52        if target_type == "category":
 53            self.original_dataset_python_class: type[Dataset] = OxfordIIITPet37
 54        elif target_type == "binary-category":
 55            self.original_dataset_python_class: type[Dataset] = OxfordIIITPet2
 56            r"""The original dataset class."""
 57
 58        super().__init__(
 59            root=root,
 60            batch_size=batch_size,
 61            num_workers=num_workers,
 62            custom_transforms=custom_transforms,
 63            repeat_channels=repeat_channels,
 64            to_tensor=to_tensor,
 65            resize=resize,
 66        )
 67
 68        self.target_type: str = target_type
 69        r"""The target type. """
 70
 71        self.validation_percentage: float = validation_percentage
 72        r"""The percentage to randomly split some training data into validation data."""
 73
 74    def prepare_data(self) -> None:
 75        r"""Download the original Oxford-IIIT Pet dataset if haven't."""
 76
 77        OxfordIIITPetRaw(
 78            root=self.root,
 79            split="trainval",
 80            target_types=self.target_type,
 81            download=True,
 82        )
 83        OxfordIIITPetRaw(
 84            root=self.root,
 85            split="test",
 86            target_types=self.target_type,
 87            download=True,
 88        )
 89
 90        pylogger.debug(
 91            "The original Oxford-IIIT Pet dataset has been downloaded to %s.",
 92            self.root,
 93        )
 94
 95    def train_and_val_dataset(self) -> tuple[Dataset, Dataset]:
 96        """Get the training and validation dataset.
 97
 98        **Returns:**
 99        - **train_and_val_dataset** (`tuple[Dataset, Dataset]`): the train and validation dataset.
100        """
101        dataset_train_and_val = OxfordIIITPetRaw(
102            root=self.root,
103            split="trainval",
104            target_types=self.target_type,
105            transform=self.train_and_val_transforms(),
106            target_transform=self.target_transform(),
107            download=False,
108        )
109
110        return random_split(
111            dataset_train_and_val,
112            lengths=[1 - self.validation_percentage, self.validation_percentage],
113            generator=torch.Generator().manual_seed(
114                42
115            ),  # this must be set fixed to make sure the datasets across experiments are the same. Don't handle it to global seed as it might vary across experiments
116        )
117
118    def test_dataset(self) -> Dataset:
119        r"""Get the test dataset.
120
121        **Returns:**
122        - **test_dataset** (`Dataset`): the test dataset.
123        """
124        dataset_test = OxfordIIITPetRaw(
125            root=self.root,
126            split="test",
127            target_types=self.target_type,
128            transform=self.test_transforms(),
129            target_transform=self.target_transform(),
130            download=False,
131        )
132
133        return dataset_test