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

dataeval/utils/data/_metadata.py

@@ -0,0 +1,352 @@
+from __future__ import annotations
+
+__all__ = []
+
+import warnings
+from typing import TYPE_CHECKING, Any, Literal, Mapping, Sequence
+
+import numpy as np
+from numpy.typing import NDArray
+
+from dataeval.typing import Array
+from dataeval.utils._array import as_numpy, to_numpy
+from dataeval.utils._bin import bin_data, digitize_data, is_continuous
+from dataeval.utils.data._types import (
+    Dataset,
+    ObjectDetectionTarget,
+)
+from dataeval.utils.metadata import merge
+
+if TYPE_CHECKING:
+    from dataeval.utils.data import Targets
+else:
+    from dataeval.utils.data._targets import Targets
+
+
+class Metadata:
+    """
+    Class containing binned metadata.
+
+    Attributes
+    ----------
+    discrete_factor_names : list[str]
+        List containing factor names for the original data that was discrete and
+        the binned continuous data
+    discrete_data : NDArray[np.int64]
+        Array containing values for the original data that was discrete and the
+        binned continuous data
+    continuous_factor_names : list[str]
+        List containing factor names for the original continuous data
+    continuous_data : NDArray[np.float64] | None
+        Array containing values for the original continuous data or None if there
+        was no continuous data
+    class_labels : NDArray[np.int]
+        Numerical class labels for the images/objects
+    class_names : list[str]
+        List of unique class names
+    total_num_factors : int
+        Sum of discrete_factor_names and continuous_factor_names plus 1 for class
+    image_indices : NDArray[np.intp]
+        Array of the image index that is mapped by the index of the factor
+
+    Parameters
+    ----------
+    dataset : ImageClassificationDataset or ObjectDetectionDataset
+        Dataset to access original targets and metadata from.
+    continuous_factor_bins : Mapping[str, int | Sequence[float]] | None, default None
+        Mapping from continuous factor name to the number of bins or bin edges
+    auto_bin_method : Literal["uniform_width", "uniform_count", "clusters"], default "uniform_width"
+        Method for automatically determining the number of bins for continuous factors
+    exclude : Sequence[str] | None, default None
+        Filter metadata factors to exclude the specified factors, cannot be set with `include`
+    include : Sequence[str] | None, default None
+        Filter metadata factors to include the specified factors, cannot be set with `exclude`
+    """
+
+    def __init__(
+        self,
+        dataset: Dataset[Any, Any],
+        *,
+        continuous_factor_bins: Mapping[str, int | Sequence[float]] | None = None,
+        auto_bin_method: Literal["uniform_width", "uniform_count", "clusters"] = "uniform_width",
+        exclude: Sequence[str] | None = None,
+        include: Sequence[str] | None = None,
+    ) -> None:
+        self._collated = False
+        self._merged = None
+        self._processed = False
+
+        self._dataset = dataset
+        self._continuous_factor_bins = dict(continuous_factor_bins) if continuous_factor_bins else {}
+        self._auto_bin_method = auto_bin_method
+
+        if exclude is not None and include is not None:
+            raise ValueError("Filters for `exclude` and `include` are mutually exclusive.")
+
+        self._exclude = set(exclude or ())
+        self._include = set(include or ())
+
+    @property
+    def targets(self) -> Targets:
+        self._collate()
+        return self._targets
+
+    @property
+    def raw(self) -> list[dict[str, Any]]:
+        self._collate()
+        return self._raw
+
+    @property
+    def exclude(self) -> set[str]:
+        return self._exclude
+
+    @exclude.setter
+    def exclude(self, value: Sequence[str]) -> None:
+        exclude = set(value)
+        if self._exclude != exclude:
+            self._exclude = exclude
+            self._include = set()
+            self._processed = False
+
+    @property
+    def include(self) -> set[str]:
+        return self._include
+
+    @include.setter
+    def include(self, value: Sequence[str]) -> None:
+        include = set(value)
+        if self._include != include:
+            self._include = include
+            self._exclude = set()
+            self._processed = False
+
+    @property
+    def continuous_factor_bins(self) -> Mapping[str, int | Sequence[float]]:
+        return self._continuous_factor_bins
+
+    @continuous_factor_bins.setter
+    def continuous_factor_bins(self, bins: Mapping[str, int | Sequence[float]]) -> None:
+        if self._continuous_factor_bins != bins:
+            self._continuous_factor_bins = dict(bins)
+            self._processed = False
+
+    @property
+    def auto_bin_method(self) -> str:
+        return self._auto_bin_method
+
+    @auto_bin_method.setter
+    def auto_bin_method(self, method: Literal["uniform_width", "uniform_count", "clusters"]) -> None:
+        if self._auto_bin_method != method:
+            self._auto_bin_method = method
+            self._processed = False
+
+    @property
+    def merged(self) -> dict[str, Any]:
+        self._merge()
+        return {} if self._merged is None else self._merged[0]
+
+    @property
+    def dropped_factors(self) -> dict[str, list[str]]:
+        self._merge()
+        return {} if self._merged is None else self._merged[1]
+
+    @property
+    def discrete_factor_names(self) -> list[str]:
+        self._process()
+        return self._discrete_factor_names
+
+    @property
+    def discrete_data(self) -> NDArray[np.int64]:
+        self._process()
+        return self._discrete_data
+
+    @property
+    def continuous_factor_names(self) -> list[str]:
+        self._process()
+        return self._continuous_factor_names
+
+    @property
+    def continuous_data(self) -> NDArray[np.float64]:
+        self._process()
+        return self._continuous_data
+
+    @property
+    def class_labels(self) -> NDArray[np.intp]:
+        self._collate()
+        return self._class_labels
+
+    @property
+    def class_names(self) -> list[str]:
+        self._collate()
+        return self._class_names
+
+    @property
+    def total_num_factors(self) -> int:
+        self._process()
+        return self._total_num_factors
+
+    @property
+    def image_indices(self) -> NDArray[np.intp]:
+        self._process()
+        return self._image_indices
+
+    def _collate(self, force: bool = False):
+        if self._collated and not force:
+            return
+
+        raw: list[dict[str, Any]] = []
+
+        labels = []
+        bboxes = []
+        scores = []
+        srcidx = []
+        is_od = None
+        for i in range(len(self._dataset)):
+            _, target, metadata = self._dataset[i]
+
+            raw.append(metadata)
+
+            if is_od_target := isinstance(target, ObjectDetectionTarget):
+                target_len = len(target.labels)
+                labels.extend(as_numpy(target.labels).tolist())
+                bboxes.extend(as_numpy(target.boxes).tolist())
+                scores.extend(as_numpy(target.scores).tolist())
+                srcidx.extend([i] * target_len)
+            elif isinstance(target, Array):
+                target_len = 1
+                labels.append(int(np.argmax(as_numpy(target))))
+                scores.append(target)
+            else:
+                raise TypeError("Encountered unsupported target type in dataset")
+
+            is_od = is_od_target if is_od is None else is_od
+            if is_od != is_od_target:
+                raise ValueError("Encountered unexpected target type in dataset")
+
+        labels = as_numpy(labels).astype(np.intp)
+        scores = as_numpy(scores).astype(np.float32)
+        bboxes = as_numpy(bboxes).astype(np.float32) if is_od else None
+        srcidx = as_numpy(srcidx).astype(np.intp) if is_od else None
+
+        self._targets = Targets(labels, scores, bboxes, srcidx)
+        self._raw = raw
+
+        index2label = self._dataset.metadata.get("index2label", {})
+        self._class_labels = self._targets.labels
+        self._class_names = [index2label.get(i, str(i)) for i in np.unique(self._class_labels)]
+        self._collated = True
+
+    def _merge(self, force: bool = False):
+        if self._merged is not None and not force:
+            return
+
+        targets_per_image = (
+            None if self.targets.source is None else np.unique(self.targets.source, return_counts=True)[1].tolist()
+        )
+        self._merged = merge(self.raw, return_dropped=True, ignore_lists=False, targets_per_image=targets_per_image)
+
+    def _validate(self) -> None:
+        # Check that metadata is a single, flattened dictionary with uniform array lengths
+        check_length = None
+        if self._targets.labels.ndim > 1:
+            raise ValueError(
+                f"Got class labels with {self._targets.labels.ndim}-dimensional "
+                f"shape {self._targets.labels.shape}, but expected a 1-dimensional array."
+            )
+        for v in self.merged.values():
+            if not isinstance(v, (list, tuple, np.ndarray)):
+                raise TypeError(
+                    "Metadata dictionary needs to be a single dictionary whose values "
+                    "are arraylike containing the metadata on a per image or per object basis."
+                )
+            else:
+                check_length = len(v) if check_length is None else check_length
+                if check_length != len(v):
+                    raise ValueError(
+                        "The lists/arrays in the metadata dict have varying lengths. "
+                        "Metadata requires them to be uniform in length."
+                    )
+        if len(self._class_labels) != check_length:
+            raise ValueError(
+                f"The length of the label array {len(self._class_labels)} is not the same as "
+                f"the length of the metadata arrays {check_length}."
+            )
+
+    def _process(self, force: bool = False) -> None:
+        if self._processed and not force:
+            return
+
+        # Validate the metadata dimensions
+        self._validate()
+
+        # Create image indices from targets
+        self._image_indices = np.arange(len(self.raw)) if self.targets.source is None else self.targets.source
+
+        # Include specified metadata keys
+        if self.include:
+            metadata = {i: self.merged[i] for i in self.include if i in self.merged}
+            continuous_factor_bins = (
+                {i: self.continuous_factor_bins[i] for i in self.include if i in self.continuous_factor_bins}
+                if self.continuous_factor_bins
+                else {}
+            )
+        else:
+            metadata = self.merged
+            continuous_factor_bins = dict(self.continuous_factor_bins) if self.continuous_factor_bins else {}
+            for k in self.exclude:
+                metadata.pop(k, None)
+                continuous_factor_bins.pop(k, None)
+
+        # Remove generated "_image_index" if present
+        if "_image_index" in metadata:
+            metadata.pop("_image_index", None)
+
+        # Bin according to user supplied bins
+        continuous_metadata = {}
+        discrete_metadata = {}
+        if continuous_factor_bins:
+            invalid_keys = set(continuous_factor_bins.keys()) - set(metadata.keys())
+            if invalid_keys:
+                raise KeyError(
+                    f"The keys - {invalid_keys} - are present in the `continuous_factor_bins` dictionary "
+                    "but are not keys in the `metadata` dictionary. Delete these keys from `continuous_factor_bins` "
+                    "or add corresponding entries to the `metadata` dictionary."
+                )
+            for factor, bins in continuous_factor_bins.items():
+                discrete_metadata[factor] = digitize_data(metadata[factor], bins)
+                continuous_metadata[factor] = metadata[factor]
+
+        # Determine category of the rest of the keys
+        remaining_keys = set(metadata.keys()) - set(continuous_metadata.keys())
+        for key in remaining_keys:
+            data = to_numpy(metadata[key])
+            if np.issubdtype(data.dtype, np.number):
+                result = is_continuous(data, self._image_indices)
+                if result:
+                    continuous_metadata[key] = data
+                unique_samples, ordinal_data = np.unique(data, return_inverse=True)
+                if unique_samples.size <= np.max([20, data.size * 0.01]):
+                    discrete_metadata[key] = ordinal_data
+                else:
+                    warnings.warn(
+                        f"A user defined binning was not provided for {key}. "
+                        f"Using the {self.auto_bin_method} method to discretize the data. "
+                        "It is recommended that the user rerun and supply the desired "
+                        "bins using the continuous_factor_bins parameter.",
+                        UserWarning,
+                    )
+                    discrete_metadata[key] = bin_data(data, self.auto_bin_method)
+            else:
+                _, discrete_metadata[key] = np.unique(data, return_inverse=True)
+
+        # Split out the dictionaries into the keys and values
+        self._discrete_factor_names = list(discrete_metadata.keys())
+        self._discrete_data = np.stack(list(discrete_metadata.values()), axis=-1, dtype=np.int64)
+        self._continuous_factor_names = list(continuous_metadata.keys())
+        self._continuous_data = (
+            np.stack(list(continuous_metadata.values()), axis=-1, dtype=np.float64)
+            if continuous_metadata
+            else np.array([], dtype=np.float64)
+        )
+        self._total_num_factors = len(self._discrete_factor_names + self._continuous_factor_names) + 1
+        self._processed = True

dataeval/utils/data/_selection.py

@@ -0,0 +1,119 @@
+from __future__ import annotations
+
+__all__ = []
+
+from enum import IntEnum
+from typing import Any, Generic, Iterator, Sequence, TypeVar
+
+from dataeval.utils.data._types import Dataset
+
+_TData = TypeVar("_TData")
+_TTarget = TypeVar("_TTarget")
+
+
+class SelectionStage(IntEnum):
+    STATE = 0
+    FILTER = 1
+    ORDER = 2
+
+
+class Selection(Generic[_TData, _TTarget]):
+    stage: SelectionStage
+
+    def __call__(self, dataset: Select[_TData, _TTarget]) -> None: ...
+
+    def __str__(self) -> str:
+        return f"{self.__class__.__name__}({', '.join([f'{k}={v}' for k, v in self.__dict__.items()])})"
+
+
+class Select(Generic[_TData, _TTarget], Dataset[_TData, _TTarget]):
+    """
+    Wraps a dataset and applies selection criteria to it.
+
+    Parameters
+    ----------
+    dataset : Dataset
+        The dataset to wrap.
+    selections : Selection or list[Selection], optional
+        The selection criteria to apply to the dataset.
+
+    Examples
+    --------
+    >>> from dataeval.utils.data.selections import ClassFilter, Limit
+
+    >>> # Construct a sample dataset with size of 100 and class count of 10
+    >>> # Elements at index `idx` are returned as tuples:
+    >>> # - f"data_{idx}", one_hot_encoded(idx % class_count), {"id": idx}
+    >>> dataset = SampleDataset(size=100, class_count=10)
+
+    >>> # Apply a selection criteria to the dataset
+    >>> selections = [Limit(size=5), ClassFilter(classes=[0, 2])]
+    >>> selected_dataset = Select(dataset, selections=selections)
+
+    >>> # Iterate over the selected dataset
+    >>> for data, target, meta in selected_dataset:
+    ...     print(f"({data}, {np.argmax(target)}, {meta})")
+    (data_0, 0, {'id': 0})
+    (data_2, 2, {'id': 2})
+    (data_10, 0, {'id': 10})
+    (data_12, 2, {'id': 12})
+    (data_20, 0, {'id': 20})
+    """
+
+    _dataset: Dataset[_TData, _TTarget]
+    _selection: list[int]
+    _selections: Sequence[Selection[_TData, _TTarget]]
+    _size_limit: int
+
+    def __init__(
+        self,
+        dataset: Dataset[_TData, _TTarget],
+        selections: Selection[_TData, _TTarget] | list[Selection[_TData, _TTarget]] | None = None,
+    ) -> None:
+        self._dataset = dataset
+        self._size_limit = len(dataset)
+        self._selection = list(range(self._size_limit))
+        self._selections = self._sort_selections(selections)
+        self.__dict__.update(dataset.__dict__)
+
+        if self._selections:
+            self._apply_selections()
+
+    def __str__(self) -> str:
+        nt = "\n    "
+        title = f"{self.__class__.__name__} Dataset"
+        sep = "-" * len(title)
+        selections = f"Selections: [{', '.join([str(s) for s in self._sort_selections(self._selections)])}]"
+        return f"{title}\n{sep}{nt}{selections}\n\n{self._dataset}"
+
+    def _sort_selections(
+        self, selections: Selection[_TData, _TTarget] | Sequence[Selection[_TData, _TTarget]] | None
+    ) -> list[Selection]:
+        if not selections:
+            return []
+
+        selections = [selections] if isinstance(selections, Selection) else selections
+        grouped: dict[int, list[Selection]] = {}
+        for selection in selections:
+            grouped.setdefault(selection.stage, []).append(selection)
+        selection_list = [selection for category in sorted(grouped) for selection in grouped[category]]
+        return selection_list
+
+    def _apply_selections(self) -> None:
+        for selection in self._selections:
+            selection(self)
+        self._selection = self._selection[: self._size_limit]
+
+    def __getattr__(self, name: str, /) -> Any:
+        selfattr = getattr(self._dataset, name, None)
+        return selfattr if selfattr is not None else getattr(self._dataset, name)
+
+    def __getitem__(self, index: int) -> tuple[_TData, _TTarget, dict[str, Any]]:
+        return self._dataset[self._selection[index]]
+
+    def __iter__(self) -> Iterator[tuple[_TData, _TTarget, dict[str, Any]]]:
+        for i in range(len(self)):
+            yield self[i]
+
+    def __len__(self) -> int:
+        return len(self._selection)

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: ...