clarena.stl_datasets.flowers102

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