clarena.stl_datasets.fgvc_aircraft

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