maite-datasets 0.0.1__py3-none-any.whl

maite_datasets/_protocols.py

@@ -0,0 +1,224 @@
+"""
+Common type protocols used for interoperability with MAITE.
+"""
+
+import sys
+from typing import (
+    Any,
+    Generic,
+    Iterator,
+    Mapping,
+    Protocol,
+    TypedDict,
+    TypeVar,
+    runtime_checkable,
+)
+
+import numpy.typing
+from typing_extensions import NotRequired, ReadOnly, Required
+
+if sys.version_info >= (3, 10):
+    from typing import TypeAlias
+else:
+    from typing_extensions import TypeAlias
+
+
+ArrayLike: TypeAlias = numpy.typing.ArrayLike
+"""
+Type alias for a `Union` representing objects that can be coerced into an array.
+
+See Also
+--------
+`NumPy ArrayLike <https://numpy.org/doc/stable/reference/typing.html#numpy.typing.ArrayLike>`_
+"""
+
+
+@runtime_checkable
+class Array(Protocol):
+    """
+    Protocol for array objects providing interoperability with DataEval.
+
+    Supports common array representations with popular libraries like
+    PyTorch, Tensorflow and JAX, as well as NumPy arrays.
+
+    Example
+    -------
+    >>> import numpy as np
+    >>> import torch
+    >>> from maite_datasets._typing import Array
+
+    Create array objects
+
+    >>> ndarray = np.random.random((10, 10))
+    >>> tensor = torch.tensor([1, 2, 3])
+
+    Check type at runtime
+
+    >>> isinstance(ndarray, Array)
+    True
+
+    >>> isinstance(tensor, Array)
+    True
+    """
+
+    @property
+    def shape(self) -> tuple[int, ...]: ...
+    def __array__(self) -> Any: ...
+    def __getitem__(self, key: Any, /) -> Any: ...
+    def __iter__(self) -> Iterator[Any]: ...
+    def __len__(self) -> int: ...
+
+
+_T = TypeVar("_T")
+_T_co = TypeVar("_T_co", covariant=True)
+
+
+class DatasetMetadata(TypedDict, total=False):
+    """
+    Dataset level metadata required for all `AnnotatedDataset` classes.
+
+    Attributes
+    ----------
+    id : Required[str]
+        A unique identifier for the dataset
+    index2label : NotRequired[dict[int, str]]
+        A lookup table converting label value to class name
+    """
+
+    id: Required[ReadOnly[str]]
+    index2label: NotRequired[ReadOnly[dict[int, str]]]
+
+
+@runtime_checkable
+class Dataset(Generic[_T_co], Protocol):
+    """
+    Protocol for a generic `Dataset`.
+
+    Methods
+    -------
+    __getitem__(index: int)
+        Returns datum at specified index.
+    __len__()
+        Returns dataset length.
+    """
+
+    def __getitem__(self, index: int, /) -> _T_co: ...
+    def __len__(self) -> int: ...
+
+
+@runtime_checkable
+class AnnotatedDataset(Dataset[_T_co], Generic[_T_co], Protocol):
+    """
+    Protocol for a generic `AnnotatedDataset`.
+
+    Attributes
+    ----------
+    metadata : :class:`.DatasetMetadata` or derivatives.
+
+    Methods
+    -------
+    __getitem__(index: int)
+        Returns datum at specified index.
+    __len__()
+        Returns dataset length.
+
+    Notes
+    -----
+    Inherits from :class:`.Dataset`.
+    """
+
+    @property
+    def metadata(self) -> DatasetMetadata: ...
+
+
+# ========== IMAGE CLASSIFICATION DATASETS ==========
+
+
+ImageClassificationDatum: TypeAlias = tuple[ArrayLike, ArrayLike, Mapping[str, Any]]
+"""
+Type alias for an image classification datum tuple.
+
+- :class:`ArrayLike` of shape (C, H, W) - Image data in channel, height, width format.
+- :class:`ArrayLike` of shape (N,) - Class label as one-hot encoded ground-truth or prediction confidences.
+- dict[str, Any] - Datum level metadata.
+"""
+
+
+ImageClassificationDataset: TypeAlias = AnnotatedDataset[ImageClassificationDatum]
+"""
+Type alias for an :class:`AnnotatedDataset` of :class:`ImageClassificationDatum` elements.
+"""
+
+# ========== OBJECT DETECTION DATASETS ==========
+
+
+@runtime_checkable
+class ObjectDetectionTarget(Protocol):
+    """
+    Protocol for targets in an Object Detection dataset.
+
+    Attributes
+    ----------
+    boxes : :class:`ArrayLike` of shape (N, 4)
+    labels : :class:`ArrayLike` of shape (N,)
+    scores : :class:`ArrayLike` of shape (N, M)
+    """
+
+    @property
+    def boxes(self) -> ArrayLike: ...
+
+    @property
+    def labels(self) -> ArrayLike: ...
+
+    @property
+    def scores(self) -> ArrayLike: ...
+
+
+ObjectDetectionDatum: TypeAlias = tuple[
+    ArrayLike, ObjectDetectionTarget, Mapping[str, Any]
+]
+"""
+Type alias for an object detection datum tuple.
+
+- :class:`ArrayLike` of shape (C, H, W) - Image data in channel, height, width format.
+- :class:`ObjectDetectionTarget` - Object detection target information for the image.
+- dict[str, Any] - Datum level metadata.
+"""
+
+
+ObjectDetectionDataset: TypeAlias = AnnotatedDataset[ObjectDetectionDatum]
+"""
+Type alias for an :class:`AnnotatedDataset` of :class:`ObjectDetectionDatum` elements.
+"""
+
+
+# ========== TRANSFORM ==========
+
+
+@runtime_checkable
+class Transform(Generic[_T], Protocol):
+    """
+    Protocol defining a transform function.
+
+    Requires a `__call__` method that returns transformed data.
+
+    Example
+    -------
+    >>> from typing import Any
+    >>> from numpy.typing import NDArray
+
+    >>> class MyTransform:
+    ...     def __init__(self, divisor: float) -> None:
+    ...         self.divisor = divisor
+    ...
+    ...     def __call__(self, data: NDArray[Any], /) -> NDArray[Any]:
+    ...         return data / self.divisor
+
+    >>> my_transform = MyTransform(divisor=255.0)
+    >>> isinstance(my_transform, Transform)
+    True
+    >>> my_transform(np.array([1, 2, 3]))
+    array([0.004, 0.008, 0.012])
+    """
+
+    def __call__(self, data: _T, /) -> _T: ...

maite_datasets/_types.py

@@ -0,0 +1,54 @@
+from __future__ import annotations
+
+__all__ = []
+
+from dataclasses import dataclass
+from typing import Generic, TypedDict, TypeVar
+
+from typing_extensions import NotRequired, Required
+
+_T_co = TypeVar("_T_co", covariant=True)
+
+
+class Dataset(Generic[_T_co]):
+    """Abstract generic base class for PyTorch style Dataset"""
+
+    def __getitem__(self, index: int) -> _T_co: ...
+    def __add__(self, other: Dataset[_T_co]) -> Dataset[_T_co]: ...
+
+
+class DatasetMetadata(TypedDict):
+    id: Required[str]
+    index2label: NotRequired[dict[int, str]]
+    split: NotRequired[str]
+
+
+class DatumMetadata(TypedDict, total=False):
+    id: Required[str]
+
+
+_TDatum = TypeVar("_TDatum")
+_TArray = TypeVar("_TArray")
+
+
+class AnnotatedDataset(Dataset[_TDatum]):
+    metadata: DatasetMetadata
+
+    def __len__(self) -> int: ...
+
+
+class ImageClassificationDataset(
+    AnnotatedDataset[tuple[_TArray, _TArray, DatumMetadata]]
+): ...
+
+
+@dataclass
+class ObjectDetectionTarget(Generic[_TArray]):
+    boxes: _TArray
+    labels: _TArray
+    scores: _TArray
+
+
+class ObjectDetectionDataset(
+    AnnotatedDataset[tuple[_TArray, ObjectDetectionTarget[_TArray], DatumMetadata]]
+): ...

maite_datasets/image_classification/__init__.py

@@ -0,0 +1,11 @@
+"""Module for MAITE compliant Image Classification datasets."""
+
+from maite_datasets.image_classification._cifar10 import CIFAR10
+from maite_datasets.image_classification._mnist import MNIST
+from maite_datasets.image_classification._ships import Ships
+
+__all__ = [
+    "CIFAR10",
+    "MNIST",
+    "Ships",
+]

maite_datasets/image_classification/_cifar10.py

@@ -0,0 +1,233 @@
+from __future__ import annotations
+
+__all__ = []
+
+from pathlib import Path
+from typing import Any, Literal, Sequence, TypeVar
+
+import numpy as np
+from numpy.typing import NDArray
+
+from maite_datasets._base import BaseICDataset, DataLocation
+from maite_datasets._mixin._numpy import BaseDatasetNumpyMixin
+from maite_datasets._protocols import Transform
+
+CIFARClassStringMap = Literal[
+    "airplane",
+    "automobile",
+    "bird",
+    "cat",
+    "deer",
+    "dog",
+    "frog",
+    "horse",
+    "ship",
+    "truck",
+]
+TCIFARClassMap = TypeVar(
+    "TCIFARClassMap", CIFARClassStringMap, int, list[CIFARClassStringMap], list[int]
+)
+
+
+class CIFAR10(BaseICDataset[NDArray[np.number[Any]]], BaseDatasetNumpyMixin):
+    """
+    `CIFAR10 <https://www.cs.toronto.edu/~kriz/cifar.html>`_ Dataset as NumPy arrays.
+
+    Parameters
+    ----------
+    root : str or pathlib.Path
+        Root directory where the data should be downloaded to or the ``cifar10`` folder of the already downloaded data.
+    image_set : "train", "test" or "base", default "train"
+        If "base", returns all of the data to allow the user to create their own splits.
+    transforms : Transform, Sequence[Transform] or None, default None
+        Transform(s) to apply to the data.
+    download : bool, default False
+        If True, downloads the dataset from the internet and puts it in root directory.
+        Class checks to see if data is already downloaded to ensure it does not create a duplicate download.
+    verbose : bool, default False
+        If True, outputs print statements.
+
+    Attributes
+    ----------
+    path : pathlib.Path
+        Location of the folder containing the data.
+    image_set : "train", "test" or "base"
+        The selected image set from the dataset.
+    transforms : Sequence[Transform]
+        The transforms to be applied to the data.
+    size : int
+        The size of the dataset.
+    index2label : dict[int, str]
+        Dictionary which translates from class integers to the associated class strings.
+    label2index : dict[str, int]
+        Dictionary which translates from class strings to the associated class integers.
+    metadata : DatasetMetadata
+        Typed dictionary containing dataset metadata, such as `id` which returns the dataset class name.
+    """
+
+    _resources = [
+        DataLocation(
+            url="https://www.cs.toronto.edu/~kriz/cifar-10-binary.tar.gz",
+            filename="cifar-10-binary.tar.gz",
+            md5=True,
+            checksum="c32a1d4ab5d03f1284b67883e8d87530",
+        ),
+    ]
+
+    index2label: dict[int, str] = {
+        0: "airplane",
+        1: "automobile",
+        2: "bird",
+        3: "cat",
+        4: "deer",
+        5: "dog",
+        6: "frog",
+        7: "horse",
+        8: "ship",
+        9: "truck",
+    }
+
+    def __init__(
+        self,
+        root: str | Path,
+        image_set: Literal["train", "test", "base"] = "train",
+        transforms: Transform[NDArray[np.number[Any]]]
+        | Sequence[Transform[NDArray[np.number[Any]]]]
+        | None = None,
+        download: bool = False,
+        verbose: bool = False,
+    ) -> None:
+        super().__init__(
+            root,
+            image_set,
+            transforms,
+            download,
+            verbose,
+        )
+
+    def _load_bin_data(
+        self, data_folder: list[Path]
+    ) -> tuple[list[str], list[int], dict[str, Any]]:
+        batch_nums = np.zeros(60000, dtype=np.uint8)
+        all_labels = np.zeros(60000, dtype=np.uint8)
+        all_images = np.zeros((60000, 3, 32, 32), dtype=np.uint8)
+        # Process each batch file, skipping .meta and .html files
+        for batch_file in data_folder:
+            # Get batch parameters
+            batch_type = "test" if "test" in batch_file.stem else "train"
+            batch_num = (
+                5 if batch_type == "test" else int(batch_file.stem.split("_")[-1]) - 1
+            )
+
+            # Load data
+            batch_images, batch_labels = self._unpack_batch_files(batch_file)
+
+            # Stack data
+            num_images = batch_images.shape[0]
+            batch_start = batch_num * num_images
+            all_images[batch_start : batch_start + num_images] = batch_images
+            all_labels[batch_start : batch_start + num_images] = batch_labels
+            batch_nums[batch_start : batch_start + num_images] = batch_num
+
+        # Save data
+        self._loaded_data = all_images
+        np.savez(
+            self.path / "cifar10",
+            images=self._loaded_data,
+            labels=all_labels,
+            batches=batch_nums,
+        )
+
+        # Select data
+        image_list = np.arange(all_labels.shape[0]).astype(str)
+        if self.image_set == "train":
+            return (
+                image_list[np.nonzero(batch_nums != 5)[0]].tolist(),
+                all_labels[batch_nums != 5].tolist(),
+                {"batch_num": batch_nums[batch_nums != 5].tolist()},
+            )
+        if self.image_set == "test":
+            return (
+                image_list[np.nonzero(batch_nums == 5)[0]].tolist(),
+                all_labels[batch_nums == 5].tolist(),
+                {"batch_num": batch_nums[batch_nums == 5].tolist()},
+            )
+        return (
+            image_list.tolist(),
+            all_labels.tolist(),
+            {"batch_num": batch_nums.tolist()},
+        )
+
+    def _load_data_inner(self) -> tuple[list[str], list[int], dict[str, Any]]:
+        """Function to load in the file paths for the data and labels and retrieve metadata"""
+        data_file = self.path / "cifar10.npz"
+        if not data_file.exists():
+            data_folder = sorted((self.path / "cifar-10-batches-bin").glob("*.bin"))
+            if not data_folder:
+                raise FileNotFoundError
+            return self._load_bin_data(data_folder)
+
+        # Load data
+        data = np.load(data_file)
+        self._loaded_data = data["images"]
+        all_labels = data["labels"]
+        batch_nums = data["batches"]
+
+        # Select data
+        image_list = np.arange(all_labels.shape[0]).astype(str)
+        if self.image_set == "train":
+            return (
+                image_list[np.nonzero(batch_nums != 5)[0]].tolist(),
+                all_labels[batch_nums != 5].tolist(),
+                {"batch_num": batch_nums[batch_nums != 5].tolist()},
+            )
+        if self.image_set == "test":
+            return (
+                image_list[np.nonzero(batch_nums == 5)[0]].tolist(),
+                all_labels[batch_nums == 5].tolist(),
+                {"batch_num": batch_nums[batch_nums == 5].tolist()},
+            )
+        return (
+            image_list.tolist(),
+            all_labels.tolist(),
+            {"batch_num": batch_nums.tolist()},
+        )
+
+    def _unpack_batch_files(
+        self, file_path: Path
+    ) -> tuple[NDArray[np.uint8], NDArray[np.uint8]]:
+        # Load pickle data with latin1 encoding
+        with file_path.open("rb") as f:
+            buffer = np.frombuffer(f.read(), dtype=np.uint8)
+            # Each entry is 1 byte for label + 3072 bytes for image (3*32*32)
+            entry_size = 1 + 3072
+            num_entries = buffer.size // entry_size
+            # Extract labels (first byte of each entry)
+            labels = buffer[::entry_size]
+
+            # Extract image data and reshape to (N, 3, 32, 32)
+            images = np.zeros((num_entries, 3, 32, 32), dtype=np.uint8)
+            for i in range(num_entries):
+                # Skip the label byte and get image data for this entry
+                start_idx = i * entry_size + 1  # +1 to skip label
+                img_flat = buffer[start_idx : start_idx + 3072]
+
+                # The CIFAR format stores channels in blocks (all R, then all G, then all B)
+                # Each channel block is 1024 bytes (32x32)
+                red_channel = img_flat[0:1024].reshape(32, 32)
+                green_channel = img_flat[1024:2048].reshape(32, 32)
+                blue_channel = img_flat[2048:3072].reshape(32, 32)
+
+                # Stack the channels in the proper C×H×W format
+                images[i, 0] = red_channel  # Red channel
+                images[i, 1] = green_channel  # Green channel
+                images[i, 2] = blue_channel  # Blue channel
+        return images, labels
+
+    def _read_file(self, path: str) -> NDArray[np.number[Any]]:
+        """
+        Function to grab the correct image from the loaded data.
+        Overwrite of the base `_read_file` because data is an all or nothing load.
+        """
+        index = int(path)
+        return self._loaded_data[index]