clarena.stl_datasets.cub2002011

  1r"""
  2The submodule in `stl_datasets` for CUB-200-2011 dataset.
  3"""
  4
  5__all__ = ["CUB2002011"]
  6
  7import logging
  8from typing import Callable
  9
 10import torch
 11from torch.utils.data import Dataset, random_split
 12from torchvision.transforms import transforms
 13
 14from clarena.stl_datasets.base import STLDatasetFromRaw
 15from clarena.stl_datasets.raw import CUB2002011 as CUB2002011Raw
 16
 17# always get logger for built-in logging in each module
 18pylogger = logging.getLogger(__name__)
 19
 20
 21class CUB2002011(STLDatasetFromRaw):
 22    r"""CUB-200-2011 dataset. The [CUB (Caltech-UCSD Birds)-200-2011)](https://www.vision.caltech.edu/datasets/cub_200_2011/) is a bird image dataset. It consists of 100,000 training, 10,000 validation, 10,000 test images of 200 bird species (classes), each 64x64 color image."""
 23
 24    original_dataset_python_class: type[Dataset] = CUB2002011Raw
 25    r"""The original dataset class."""
 26
 27    def __init__(
 28        self,
 29        root: str,
 30        validation_percentage: float,
 31        batch_size: int = 1,
 32        num_workers: int = 0,
 33        custom_transforms: Callable | transforms.Compose | None = None,
 34        repeat_channels: int | None = None,
 35        to_tensor: bool = True,
 36        resize: tuple[int, int] | None = None,
 37    ) -> None:
 38        r"""
 39        **Args:**
 40        - **root** (`str`): the root directory where the original MNIST data 'MNIST/' live.
 41        - **validation_percentage** (`float`): the percentage to randomly split some training data into validation data.
 42        - **batch_size** (`int`): The batch size in train, val, test dataloader.
 43        - **num_workers** (`int`): the number of workers for dataloaders.
 44        - **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.
 45        - **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.
 46        - **to_tensor** (`bool`): whether to include `ToTensor()` transform. Default is True.
 47        - **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.
 48        """
 49        super().__init__(
 50            root=root,
 51            batch_size=batch_size,
 52            num_workers=num_workers,
 53            custom_transforms=custom_transforms,
 54            repeat_channels=repeat_channels,
 55            to_tensor=to_tensor,
 56            resize=resize,
 57        )
 58
 59        self.validation_percentage: float = validation_percentage
 60        r"""The percentage to randomly split some training data into validation data."""
 61
 62    def prepare_data(self) -> None:
 63        r"""Download the original CUB-200-2011 dataset if haven't."""
 64
 65        CUB2002011Raw(root=self.root, train=True, download=True)
 66        CUB2002011Raw(root=self.root, train=False, download=True)
 67
 68        pylogger.debug(
 69            "The original CUB-200-2011 dataset has been downloaded to %s.",
 70            self.root,
 71        )
 72
 73    def train_and_val_dataset(self) -> tuple[Dataset, Dataset]:
 74        """Get the training and validation dataset.
 75
 76        **Returns:**
 77        - **train_and_val_dataset** (`tuple[Dataset, Dataset]`): the train and validation dataset.
 78        """
 79        dataset_train_and_val = CUB2002011Raw(
 80            root=self.root,
 81            train=True,
 82            transform=self.train_and_val_transforms(),
 83            target_transform=self.target_transform(),
 84            download=False,
 85        )
 86
 87        return random_split(
 88            dataset_train_and_val,
 89            lengths=[1 - self.validation_percentage, self.validation_percentage],
 90            generator=torch.Generator().manual_seed(
 91                42
 92            ),  # 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
 93        )
 94
 95    def test_dataset(self) -> Dataset:
 96        r"""Get the test dataset.
 97
 98        **Returns:**
 99        - **test_dataset** (`Dataset`): the test dataset.
100        """
101        dataset_test = CUB2002011Raw(
102            root=self.root,
103            train=False,
104            transform=self.test_transforms(),
105            target_transform=self.target_transform(),
106            download=False,
107        )
108
109        return dataset_test