dataeval 0.76.0__py3-none-any.whl → 0.81.0__py3-none-any.whl

dataeval/utils/{dataset/split.py → data/_split.py}

@@ -4,7 +4,7 @@ __all__ = []
 
 import warnings
 from dataclasses import dataclass
-from typing import Any, Iterator, NamedTuple, Protocol
+from typing import Any, Iterator, Protocol
 
 import numpy as np
 from numpy.typing import NDArray
@@ -13,10 +13,11 @@ from sklearn.metrics import silhouette_score
 from sklearn.model_selection import GroupKFold, KFold, StratifiedGroupKFold, StratifiedKFold
 from sklearn.utils.multiclass import type_of_target
 
-from dataeval.output import Output, set_metadata
+from dataeval._output import Output, set_metadata
 
 
-class TrainValSplit(NamedTuple):
+@dataclass
+class TrainValSplit:
     """Tuple containing train and validation indices"""
 
     train: NDArray[np.intp]
@@ -274,8 +275,7 @@ def get_group_ids(metadata: dict[str, Any], group_names: list[str], num_samples:
     for name, feature in features2group.items():
         if len(feature) != num_samples:
             raise ValueError(
-                f"Feature length does not match number of labels. "
-                f"Got {len(feature)} features and {num_samples} samples"
+                f"Feature length does not match number of labels. Got {len(feature)} features and {num_samples} samples"
             )
 
         if type_of_target(feature) == "continuous":
@@ -505,23 +505,22 @@ def split_dataset(
         if is_groupable(possible_groups, group_partitions):
             groups = possible_groups
 
-    test_indices: NDArray[np.intp]
     index = np.arange(label_length)
 
-    tv_indices, test_indices = (
+    tvs = (
         single_split(index=index, labels=labels, split_frac=test_frac, groups=groups, stratified=stratify)
         if test_frac
-        else (index, np.array([], dtype=np.intp))
+        else TrainValSplit(index, np.array([], dtype=np.intp))
     )
 
-    tv_labels = labels[tv_indices]
-    tv_groups = groups[tv_indices] if groups is not None else None
+    tv_labels = labels[tvs.train]
+    tv_groups = groups[tvs.train] if groups is not None else None
 
     if num_folds == 1:
-        tv_splits = [single_split(tv_indices, tv_labels, val_frac, tv_groups, stratify)]
+        tv_splits = [single_split(tvs.train, tv_labels, val_frac, tv_groups, stratify)]
     else:
-        tv_splits = make_splits(tv_indices, tv_labels, num_folds, tv_groups, stratify)
+        tv_splits = make_splits(tvs.train, tv_labels, num_folds, tv_groups, stratify)
 
-    folds: list[TrainValSplit] = [TrainValSplit(tv_indices[split.train], tv_indices[split.val]) for split in tv_splits]
+    folds: list[TrainValSplit] = [TrainValSplit(tvs.train[split.train], tvs.train[split.val]) for split in tv_splits]
 
-    return SplitDatasetOutput(test_indices, folds)
+    return SplitDatasetOutput(tvs.val, folds)

dataeval/utils/data/_targets.py

@@ -0,0 +1,73 @@
+from __future__ import annotations
+
+__all__ = []
+
+from dataclasses import dataclass
+
+import numpy as np
+from numpy.typing import NDArray
+
+
+def _len(arr: NDArray, dim: int) -> int:
+    return 0 if len(arr) == 0 else len(np.atleast_1d(arr) if dim == 1 else np.atleast_2d(arr))
+
+
+@dataclass(frozen=True)
+class Targets:
+    """
+    Dataclass defining targets for image classification or object detection.
+
+    Attributes
+    ----------
+    labels : NDArray[np.intp]
+        Labels (N,) for N images or objects
+    scores : NDArray[np.float32]
+        Probability scores (N,M) for N images of M classes or confidence score (N,) of objects
+    bboxes : NDArray[np.float32] | None
+        Bounding boxes (N,4) for N objects in (x0,y0,x1,y1) format
+    source : NDArray[np.intp] | None
+        Source image index (N,) for N objects
+    """
+
+    labels: NDArray[np.intp]
+    scores: NDArray[np.float32]
+    bboxes: NDArray[np.float32] | None
+    source: NDArray[np.intp] | None
+
+    def __post_init__(self) -> None:
+        if (self.bboxes is None) != (self.source is None):
+            raise ValueError("Either both bboxes and source must be provided or neither.")
+
+        labels = _len(self.labels, 1)
+        scores = _len(self.scores, 2) if self.bboxes is None else _len(self.scores, 1)
+        bboxes = labels if self.bboxes is None else _len(self.bboxes, 2)
+        source = labels if self.source is None else _len(self.source, 1)
+
+        if labels != scores or labels != bboxes or labels != source:
+            raise ValueError(
+                "Labels, scores, bboxes and source must be the same length (if provided).\n"
+                + f"    labels: {self.labels.shape}\n"
+                + f"    scores: {self.scores.shape}\n"
+                + f"    bboxes: {None if self.bboxes is None else self.bboxes.shape}\n"
+                + f"    source: {None if self.source is None else self.source.shape}\n"
+            )
+
+    def __len__(self) -> int:
+        return len(self.labels)
+
+    def at(self, idx: int) -> Targets:
+        if self.source is None or self.bboxes is None:
+            return Targets(
+                np.atleast_1d(self.labels[idx]),
+                np.atleast_2d(self.scores[idx]),
+                None,
+                None,
+            )
+        else:
+            mask = np.where(self.source == idx, True, False)
+            return Targets(
+                np.atleast_1d(self.labels[mask]),
+                np.atleast_1d(self.scores[mask]),
+                np.atleast_2d(self.bboxes[mask]),
+                np.atleast_1d(self.source[mask]),
+            )

dataeval/utils/data/_types.py

@@ -0,0 +1,58 @@
+from __future__ import annotations
+
+__all__ = []
+
+import sys
+from dataclasses import dataclass
+from typing import Any, Generic, Protocol, TypedDict, TypeVar
+
+if sys.version_info >= (3, 11):
+    from typing import NotRequired, Required
+else:
+    from typing_extensions import NotRequired, Required
+
+from torch.utils.data import Dataset as _Dataset
+
+_TArray = TypeVar("_TArray")
+_TData = TypeVar("_TData", covariant=True)
+_TTarget = TypeVar("_TTarget", covariant=True)
+
+
+class DatasetMetadata(TypedDict):
+    id: Required[str]
+    index2label: NotRequired[dict[int, str]]
+    split: NotRequired[str]
+
+
+class Dataset(_Dataset[tuple[_TData, _TTarget, dict[str, Any]]]):
+    metadata: DatasetMetadata
+
+    def __getitem__(self, index: Any) -> tuple[_TData, _TTarget, dict[str, Any]]: ...
+    def __len__(self) -> int: ...
+
+
+class ImageClassificationDataset(Dataset[_TArray, _TArray]): ...
+
+
+@dataclass
+class ObjectDetectionTarget(Generic[_TArray]):
+    boxes: _TArray
+    labels: _TArray
+    scores: _TArray
+
+
+class ObjectDetectionDataset(Dataset[_TArray, ObjectDetectionTarget[_TArray]]): ...
+
+
+@dataclass
+class SegmentationTarget(Generic[_TArray]):
+    mask: _TArray
+    labels: _TArray
+    scores: _TArray
+
+
+class SegmentationDataset(Dataset[_TArray, SegmentationTarget[_TArray]]): ...
+
+
+class Transform(Generic[_TArray], Protocol):
+    def __call__(self, data: _TArray, /) -> _TArray: ...

dataeval/utils/data/collate.py

@@ -0,0 +1,103 @@
+"""
+Collate functions used with a PyTorch DataLoader to load data from MAITE compliant datasets.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Iterable, Sequence, TypeVar
+
+import numpy as np
+import torch
+from numpy.typing import NDArray
+
+from dataeval.typing import ArrayLike
+from dataeval.utils._array import as_numpy
+
+T_in = TypeVar("T_in")
+T_tgt = TypeVar("T_tgt")
+T_md = TypeVar("T_md")
+
+
+def list_collate_fn(
+    batch_data_as_singles: Iterable[tuple[T_in, T_tgt, T_md]],
+) -> tuple[Sequence[T_in], Sequence[T_tgt], Sequence[T_md]]:
+    """
+    A collate function that takes a batch of individual data points in the format
+    (input, target, metadata) and returns three lists: the input batch, the target batch,
+    and the metadata batch. This is useful for loading data with torch.utils.data.DataLoader
+    when the target and metadata are not tensors.
+
+    Parameters
+    ----------
+    batch_data_as_singles : An iterable of (input, target, metadata) tuples.
+
+    Returns
+    -------
+    tuple[Sequence[T_in], Sequence[T_tgt], Sequence[T_md]]
+        A tuple of three lists: the input batch, the target batch, and the metadata batch.
+    """
+    input_batch: list[T_in] = []
+    target_batch: list[T_tgt] = []
+    metadata_batch: list[T_md] = []
+    for input_datum, target_datum, metadata_datum in batch_data_as_singles:
+        input_batch.append(input_datum)
+        target_batch.append(target_datum)
+        metadata_batch.append(metadata_datum)
+
+    return input_batch, target_batch, metadata_batch
+
+
+def numpy_collate_fn(
+    batch_data_as_singles: Iterable[tuple[ArrayLike, T_tgt, T_md]],
+) -> tuple[NDArray[Any], Sequence[T_tgt], Sequence[T_md]]:
+    """
+    A collate function that takes a batch of individual data points in the format
+    (input, target, metadata) and returns the batched input as a single NumPy array with two
+    lists: the target batch, and the metadata batch. The inputs must be homogeneous arrays.
+
+    Parameters
+    ----------
+    batch_data_as_singles : An iterable of (ArrayLike, target, metadata) tuples.
+
+    Returns
+    -------
+    tuple[NDArray[Any], Sequence[T_tgt], Sequence[T_md]]
+        A tuple of a NumPy array and two lists: the input batch, the target batch, and the metadata batch.
+    """
+    input_batch: list[NDArray[Any]] = []
+    target_batch: list[T_tgt] = []
+    metadata_batch: list[T_md] = []
+    for input_datum, target_datum, metadata_datum in batch_data_as_singles:
+        input_batch.append(as_numpy(input_datum))
+        target_batch.append(target_datum)
+        metadata_batch.append(metadata_datum)
+
+    return np.stack(input_batch) if input_batch else np.array([]), target_batch, metadata_batch
+
+
+def torch_collate_fn(
+    batch_data_as_singles: Iterable[tuple[ArrayLike, T_tgt, T_md]],
+) -> tuple[torch.Tensor, Sequence[T_tgt], Sequence[T_md]]:
+    """
+    A collate function that takes a batch of individual data points in the format
+    (input, target, metadata) and returns the batched input as a single torch Tensor with two
+    lists: the target batch, and the metadata batch. The inputs must be homogeneous arrays.
+
+    Parameters
+    ----------
+    batch_data_as_singles : An iterable of (ArrayLike, target, metadata) tuples.
+
+    Returns
+    -------
+    tuple[torch.Tensor, Sequence[T_tgt], Sequence[T_md]]
+        A tuple of a torch Tensor and two lists: the input batch, the target batch, and the metadata batch.
+    """
+    input_batch: list[torch.Tensor] = []
+    target_batch: list[T_tgt] = []
+    metadata_batch: list[T_md] = []
+    for input_datum, target_datum, metadata_datum in batch_data_as_singles:
+        input_batch.append(torch.as_tensor(input_datum))
+        target_batch.append(target_datum)
+        metadata_batch.append(metadata_datum)
+
+    return torch.stack(input_batch) if input_batch else torch.tensor([]), target_batch, metadata_batch

dataeval/utils/data/datasets/__init__.py

@@ -0,0 +1,17 @@
+"""Provides access to common Computer Vision datasets."""
+
+from dataeval.utils.data.datasets._cifar10 import CIFAR10
+from dataeval.utils.data.datasets._milco import MILCO
+from dataeval.utils.data.datasets._mnist import MNIST
+from dataeval.utils.data.datasets._ships import Ships
+from dataeval.utils.data.datasets._voc import VOCDetection, VOCDetectionTorch, VOCSegmentation
+
+__all__ = [
+    "MNIST",
+    "Ships",
+    "CIFAR10",
+    "MILCO",
+    "VOCDetection",
+    "VOCDetectionTorch",
+    "VOCSegmentation",
+]

dataeval/utils/data/datasets/_base.py

@@ -0,0 +1,254 @@
+from __future__ import annotations
+
+__all__ = []
+
+from abc import abstractmethod
+from pathlib import Path
+from typing import Any, Generic, Iterator, Literal, NamedTuple, Sequence, TypeVar
+
+from dataeval.utils.data._types import (
+    Dataset,
+    DatasetMetadata,
+    ImageClassificationDataset,
+    ObjectDetectionDataset,
+    ObjectDetectionTarget,
+    SegmentationDataset,
+    SegmentationTarget,
+    Transform,
+)
+from dataeval.utils.data.datasets._fileio import _ensure_exists
+from dataeval.utils.data.datasets._mixin import BaseDatasetMixin
+
+_TArray = TypeVar("_TArray")
+_TTarget = TypeVar("_TTarget")
+_TRawTarget = TypeVar("_TRawTarget", list[int], list[str])
+
+
+class DataLocation(NamedTuple):
+    url: str
+    filename: str
+    md5: bool
+    checksum: str
+
+
+class BaseDataset(Dataset[_TArray, _TTarget], Generic[_TArray, _TTarget, _TRawTarget]):
+    """
+    Base class for internet downloaded datasets.
+    """
+
+    # Each subclass should override the attributes below.
+    # Each resource tuple must contain:
+    #    'url': str, the URL to download from
+    #    'filename': str, the name of the file once downloaded
+    #    'md5': boolean, True if it's the checksum value is md5
+    #    'checksum': str, the associated checksum for the downloaded file
+    _resources: list[DataLocation]
+    _resource_index: int = 0
+    index2label: dict[int, str]
+
+    def __init__(
+        self,
+        root: str | Path,
+        download: bool = False,
+        image_set: Literal["train", "val", "test", "base"] = "train",
+        transforms: Transform[_TArray] | Sequence[Transform[_TArray]] | None = None,
+        verbose: bool = False,
+    ) -> None:
+        self._root: Path = root.absolute() if isinstance(root, Path) else Path(root).absolute()
+        transforms = transforms if transforms is not None else []
+        self.transforms: Sequence[Transform[_TArray]] = transforms if isinstance(transforms, Sequence) else [transforms]
+        self.image_set = image_set
+        self._verbose = verbose
+
+        # Internal Attributes
+        self._download = download
+        self._filepaths: list[str]
+        self._targets: _TRawTarget
+        self._datum_metadata: dict[str, list[Any]]
+        self._resource: DataLocation = self._resources[self._resource_index]
+        self._label2index = {v: k for k, v in self.index2label.items()}
+
+        self.metadata: DatasetMetadata = DatasetMetadata(
+            id=self._unique_id(),
+            index2label=self.index2label,
+            split=self.image_set,
+        )
+
+        # Load the data
+        self.path: Path = self._get_dataset_dir()
+        self._filepaths, self._targets, self._datum_metadata = self._load_data()
+        self.size: int = len(self._filepaths)
+
+    def __str__(self) -> str:
+        nt = "\n    "
+        title = f"{self.__class__.__name__} Dataset"
+        sep = "-" * len(title)
+        attrs = [f"{k.capitalize()}: {v}" for k, v in self.__dict__.items() if not k.startswith("_")]
+        return f"{title}\n{sep}{nt}{nt.join(attrs)}"
+
+    @property
+    def label2index(self) -> dict[str, int]:
+        return self._label2index
+
+    def __iter__(self) -> Iterator[tuple[_TArray, _TTarget, dict[str, Any]]]:
+        for i in range(len(self)):
+            yield self[i]
+
+    def _get_dataset_dir(self) -> Path:
+        # Create a designated folder for this dataset (named after the class)
+        if self._root.stem in [
+            self.__class__.__name__.lower(),
+            self.__class__.__name__.upper(),
+            self.__class__.__name__,
+        ]:
+            dataset_dir: Path = self._root
+        else:
+            dataset_dir: Path = self._root / self.__class__.__name__.lower()
+        if not dataset_dir.exists():
+            dataset_dir.mkdir(parents=True, exist_ok=True)
+        return dataset_dir
+
+    def _unique_id(self) -> str:
+        unique_id = f"{self.__class__.__name__}_{self.image_set}"
+        return unique_id
+
+    def _load_data(self) -> tuple[list[str], _TRawTarget, dict[str, Any]]:
+        """
+        Function to determine if data can be accessed or if it needs to be downloaded and/or extracted.
+        """
+        if self._verbose:
+            print(f"Determining if {self._resource.filename} needs to be downloaded.")
+
+        try:
+            result = self._load_data_inner()
+            if self._verbose:
+                print("No download needed, loaded data successfully.")
+        except FileNotFoundError:
+            _ensure_exists(*self._resource, self.path, self._root, self._download, self._verbose)
+            result = self._load_data_inner()
+        return result
+
+    @abstractmethod
+    def _load_data_inner(self) -> tuple[list[str], _TRawTarget, dict[str, Any]]: ...
+
+    def _transform(self, image: _TArray) -> _TArray:
+        """Function to transform the image prior to returning based on parameters passed in."""
+        for transform in self.transforms:
+            image = transform(image)
+        return image
+
+    def __len__(self) -> int:
+        return self.size
+
+
+class BaseICDataset(
+    BaseDataset[_TArray, _TArray, list[int]],
+    BaseDatasetMixin[_TArray],
+    ImageClassificationDataset[_TArray],
+):
+    """
+    Base class for image classification datasets.
+    """
+
+    def __getitem__(self, index: int) -> tuple[_TArray, _TArray, dict[str, Any]]:
+        """
+        Args
+        ----
+        index : int
+            Value of the desired data point
+
+        Returns
+        -------
+        tuple[TArray, TArray, dict[str, Any]]
+            Image, target, datum_metadata - where target is one-hot encoding of class.
+        """
+        # Get the associated label and score
+        label = self._targets[index]
+        score = self._one_hot_encode(label)
+        # Get the image
+        img = self._read_file(self._filepaths[index])
+        img = self._transform(img)
+
+        img_metadata = {key: val[index] for key, val in self._datum_metadata.items()}
+
+        return img, score, img_metadata
+
+
+class BaseODDataset(
+    BaseDataset[_TArray, ObjectDetectionTarget[_TArray], list[str]],
+    BaseDatasetMixin[_TArray],
+    ObjectDetectionDataset[_TArray],
+):
+    """
+    Base class for object detection datasets.
+    """
+
+    def __getitem__(self, index: int) -> tuple[_TArray, ObjectDetectionTarget[_TArray], dict[str, Any]]:
+        """
+        Args
+        ----
+        index : int
+            Value of the desired data point
+
+        Returns
+        -------
+        tuple[TArray, ObjectDetectionTarget[TArray], dict[str, Any]]
+            Image, target, datum_metadata - target.boxes returns boxes in x0, y0, x1, y1 format
+        """
+        # Grab the bounding boxes and labels from the annotations
+        boxes, labels, additional_metadata = self._read_annotations(self._targets[index])
+        # Get the image
+        img = self._read_file(self._filepaths[index])
+        img = self._transform(img)
+
+        target = ObjectDetectionTarget(self._as_array(boxes), self._as_array(labels), self._one_hot_encode(labels))
+
+        img_metadata = {key: val[index] for key, val in self._datum_metadata.items()}
+        img_metadata = img_metadata | additional_metadata
+
+        return img, target, img_metadata
+
+    @abstractmethod
+    def _read_annotations(self, annotation: str) -> tuple[list[list[float]], list[int], dict[str, Any]]: ...
+
+
+class BaseSegDataset(
+    BaseDataset[_TArray, SegmentationTarget[_TArray], list[str]],
+    BaseDatasetMixin[_TArray],
+    SegmentationDataset[_TArray],
+):
+    """
+    Base class for segmentation datasets.
+    """
+
+    _masks: Sequence[str]
+
+    def __getitem__(self, index: int) -> tuple[_TArray, SegmentationTarget[_TArray], dict[str, Any]]:
+        """
+        Args
+        ----
+        index : int
+            Value of the desired data point
+
+        Returns
+        -------
+        tuple[TArray, SegmentationTarget[TArray], dict[str, Any]]
+            Image, target, datum_metadata - target.mask returns the ground truth mask
+        """
+        # Grab the labels from the annotations
+        _, labels, additional_metadata = self._read_annotations(self._targets[index])
+        # Grab the ground truth masks
+        mask = self._read_file(self._masks[index])
+        # Get the image
+        img = self._read_file(self._filepaths[index])
+        img = self._transform(img)
+
+        target = SegmentationTarget(mask, self._as_array(labels), self._one_hot_encode(labels))
+
+        img_metadata = {key: val[index] for key, val in self._datum_metadata.items()}
+        img_metadata = img_metadata | additional_metadata
+
+        return img, target, img_metadata
+
+    @abstractmethod
+    def _read_annotations(self, annotation: str) -> tuple[list[list[float]], list[int], dict[str, Any]]: ...