dataeval 0.84.0__py3-none-any.whl → 1.0.0__py3-none-any.whl

dataeval/detectors/drift/updates.py

@@ -8,11 +8,12 @@ from __future__ import annotations
 __all__ = ["LastSeenUpdate", "ReservoirSamplingUpdate"]
 
 from abc import ABC, abstractmethod
-from typing import Any
 
 import numpy as np
 from numpy.typing import NDArray
 
+from dataeval.utils._array import flatten
+
 
 class BaseUpdateStrategy(ABC):
     """
@@ -28,8 +29,7 @@ class BaseUpdateStrategy(ABC):
         self.n = n
 
     @abstractmethod
-    def __call__(self, x_ref: NDArray[Any], x: NDArray[Any], count: int) -> NDArray[Any]:
-        """Abstract implementation of update strategy"""
+    def __call__(self, x_ref: NDArray[np.float32], x_new: NDArray[np.float32], count: int) -> NDArray[np.float32]: ...
 
 
 class LastSeenUpdate(BaseUpdateStrategy):
@@ -42,9 +42,8 @@ class LastSeenUpdate(BaseUpdateStrategy):
         Update with last n instances seen by the detector.
     """
 
-    def __call__(self, x_ref: NDArray[Any], x: NDArray[Any], count: int) -> NDArray[Any]:
-        x_updated = np.concatenate([x_ref, x], axis=0)
-        return x_updated[-self.n :]
+    def __call__(self, x_ref: NDArray[np.float32], x_new: NDArray[np.float32], count: int) -> NDArray[np.float32]:
+        return np.concatenate([x_ref, flatten(x_new)], axis=0)[-self.n :]
 
 
 class ReservoirSamplingUpdate(BaseUpdateStrategy):
@@ -57,16 +56,18 @@ class ReservoirSamplingUpdate(BaseUpdateStrategy):
         Update with last n instances seen by the detector.
     """
 
-    def __call__(self, x_ref: NDArray[Any], x: NDArray[Any], count: int) -> NDArray[Any]:
-        if x.shape[0] + count <= self.n:
-            return np.concatenate([x_ref, x], axis=0)
+    def __call__(self, x_ref: NDArray[np.float32], x_new: NDArray[np.float32], count: int) -> NDArray[np.float32]:
+        if x_new.shape[0] + count <= self.n:
+            return np.concatenate([x_ref, flatten(x_new)], axis=0)
 
         n_ref = x_ref.shape[0]
-        output_size = min(self.n, n_ref + x.shape[0])
-        shape = (output_size,) + x.shape[1:]
+        output_size = min(self.n, n_ref + x_new.shape[0])
+        shape = (output_size,) + x_new.shape[1:]
+
         x_reservoir = np.zeros(shape, dtype=x_ref.dtype)
         x_reservoir[:n_ref] = x_ref
-        for item in x:
+
+        for item in x_new:
             count += 1
             if n_ref < self.n:
                 x_reservoir[n_ref, :] = item

dataeval/detectors/linters/duplicates.py

@@ -4,13 +4,13 @@ __all__ = []
 
 from typing import Any, Sequence, overload
 
+from dataeval.data._images import Images
 from dataeval.metrics.stats import hashstats
 from dataeval.metrics.stats._base import combine_stats, get_dataset_step_from_idx
 from dataeval.outputs import DuplicatesOutput, HashStatsOutput
 from dataeval.outputs._base import set_metadata
 from dataeval.outputs._linters import DatasetDuplicateGroupMap, DuplicateGroup
-from dataeval.typing import Array, Dataset
-from dataeval.utils.data._images import Images
+from dataeval.typing import ArrayLike, Dataset
 
 
 class Duplicates:
@@ -110,13 +110,15 @@ class Duplicates:
         return DuplicatesOutput(**duplicates)
 
     @set_metadata(state=["only_exact"])
-    def evaluate(self, data: Dataset[Array] | Dataset[tuple[Array, Any, Any]]) -> DuplicatesOutput[DuplicateGroup]:
+    def evaluate(
+        self, data: Dataset[ArrayLike] | Dataset[tuple[ArrayLike, Any, Any]]
+    ) -> DuplicatesOutput[DuplicateGroup]:
         """
         Returns duplicate image indices for both exact matches and near matches
 
         Parameters
         ----------
-        data : Iterable[Array], shape - (N, C, H, W) | Dataset[tuple[Array, Any, Any]]
+        data : Iterable[ArrayLike], shape - (N, C, H, W) | Dataset[tuple[ArrayLike, Any, Any]]
             A dataset of images in an Array format or the output(s) from a hashstats analysis
 
         Returns

dataeval/detectors/linters/outliers.py

@@ -7,14 +7,14 @@ from typing import Any, Literal, Sequence, overload
 import numpy as np
 from numpy.typing import NDArray
 
+from dataeval.data._images import Images
 from dataeval.metrics.stats._base import combine_stats, get_dataset_step_from_idx
 from dataeval.metrics.stats._imagestats import imagestats
 from dataeval.outputs import DimensionStatsOutput, ImageStatsOutput, OutliersOutput, PixelStatsOutput, VisualStatsOutput
 from dataeval.outputs._base import set_metadata
 from dataeval.outputs._linters import IndexIssueMap, OutlierStatsOutput
 from dataeval.outputs._stats import BOX_COUNT, SOURCE_INDEX
-from dataeval.typing import Array, Dataset
-from dataeval.utils.data._images import Images
+from dataeval.typing import ArrayLike, Dataset
 
 
 def _get_outlier_mask(
@@ -197,7 +197,7 @@ class Outliers:
         return OutliersOutput(output_list)
 
     @set_metadata(state=["use_dimension", "use_pixel", "use_visual", "outlier_method", "outlier_threshold"])
-    def evaluate(self, data: Dataset[Array] | Dataset[tuple[Array, Any, Any]]) -> OutliersOutput[IndexIssueMap]:
+    def evaluate(self, data: Dataset[ArrayLike] | Dataset[tuple[ArrayLike, Any, Any]]) -> OutliersOutput[IndexIssueMap]:
         """
         Returns indices of Outliers with the issues identified for each
 

dataeval/detectors/ood/ae.py

@@ -81,7 +81,7 @@ class OOD_AE(OODBase):
 
     def _score(self, X: NDArray[np.float32], batch_size: int = int(1e10)) -> OODScoreOutput:
         # reconstruct instances
-        X_recon = predict_batch(X, self.model, batch_size=batch_size)
+        X_recon = predict_batch(X, self.model, batch_size=batch_size).detach().cpu().numpy()
 
         # compute feature and instance level scores
         fscore = np.power(X - X_recon, 2)

dataeval/metadata/_distance.py

@@ -9,11 +9,11 @@ import numpy as np
 from scipy.stats import iqr, ks_2samp
 from scipy.stats import wasserstein_distance as emd
 
+from dataeval.data import Metadata
 from dataeval.metadata._utils import _compare_keys, _validate_factors_and_data
 from dataeval.outputs import MetadataDistanceOutput, MetadataDistanceValues
 from dataeval.outputs._base import set_metadata
 from dataeval.typing import ArrayLike
-from dataeval.utils.data import Metadata
 
 
 class KSType(NamedTuple):

dataeval/metadata/_ood.py

@@ -9,10 +9,10 @@ from numpy.typing import NDArray
 from sklearn.feature_selection import mutual_info_classif
 
 from dataeval.config import get_seed
+from dataeval.data import Metadata
 from dataeval.metadata._utils import _compare_keys, _validate_factors_and_data
 from dataeval.outputs import MostDeviatedFactorsOutput, OODOutput, OODPredictorOutput
 from dataeval.outputs._base import set_metadata
-from dataeval.utils.data import Metadata
 
 
 def _combine_discrete_continuous(metadata: Metadata) -> tuple[list[str], NDArray[np.float64]]:
@@ -201,7 +201,7 @@ def find_most_deviated_factors(
     MostDeviatedFactorsOutput([])
     """
 
-    ood_mask: NDArray[np.bool] = ood.is_ood
+    ood_mask: NDArray[np.bool_] = ood.is_ood
 
     # No metadata correlated with out of distribution data
     if not any(ood_mask):
@@ -303,7 +303,7 @@ def find_ood_predictors(
     OODPredictorOutput({})
     """
 
-    ood_mask: NDArray[np.bool] = ood.is_ood
+    ood_mask: NDArray[np.bool_] = ood.is_ood
 
     discrete_features_count = len(metadata.discrete_factor_names)
     factors, data = _combine_discrete_continuous(metadata)  # (F, ), (S, F) => F = Fd + Fc
@@ -320,7 +320,7 @@ def find_ood_predictors(
     # Calculate mean, std of each factor over all samples
     scaled_data = (data - np.mean(data, axis=0)) / np.std(data, axis=0, ddof=1)  # (S, F)
 
-    discrete_features = np.zeros_like(factors, dtype=np.bool)
+    discrete_features = np.zeros_like(factors, dtype=np.bool_)
     discrete_features[:discrete_features_count] = True
 
     mutual_info_values = (

dataeval/metrics/bias/_balance.py

@@ -9,10 +9,10 @@ import scipy as sp
 from sklearn.feature_selection import mutual_info_classif, mutual_info_regression
 
 from dataeval.config import EPSILON, get_seed
+from dataeval.data import Metadata
 from dataeval.outputs import BalanceOutput
 from dataeval.outputs._base import set_metadata
 from dataeval.utils._bin import get_counts
-from dataeval.utils.data import Metadata
 
 
 def _validate_num_neighbors(num_neighbors: int) -> int:

dataeval/metrics/bias/_diversity.py

@@ -8,11 +8,11 @@ import numpy as np
 import scipy as sp
 from numpy.typing import NDArray
 
+from dataeval.data import Metadata
 from dataeval.outputs import DiversityOutput
 from dataeval.outputs._base import set_metadata
 from dataeval.utils._bin import get_counts
 from dataeval.utils._method import get_method
-from dataeval.utils.data import Metadata
 
 
 def diversity_shannon(

dataeval/metrics/bias/_parity.py

@@ -10,11 +10,11 @@ from numpy.typing import NDArray
 from scipy.stats import chisquare
 from scipy.stats.contingency import chi2_contingency, crosstab
 
+from dataeval.data import Metadata
 from dataeval.outputs import LabelParityOutput, ParityOutput
 from dataeval.outputs._base import set_metadata
 from dataeval.typing import ArrayLike
 from dataeval.utils._array import as_numpy
-from dataeval.utils.data import Metadata
 
 
 def normalize_expected_dist(expected_dist: NDArray[Any], observed_dist: NDArray[Any]) -> NDArray[Any]:

dataeval/metrics/stats/_base.py

@@ -10,7 +10,7 @@ from copy import deepcopy
 from dataclasses import dataclass
 from functools import partial
 from multiprocessing import Pool
-from typing import Any, Callable, Generic, Iterable, Sequence, TypeVar, cast
+from typing import Any, Callable, Generic, Iterable, Sequence, TypeVar
 
 import numpy as np
 import tqdm
@@ -19,7 +19,7 @@ from numpy.typing import NDArray
 from dataeval.config import get_max_processes
 from dataeval.outputs._stats import BaseStatsOutput, SourceIndex
 from dataeval.typing import Array, ArrayLike, Dataset, ObjectDetectionTarget
-from dataeval.utils._array import to_numpy
+from dataeval.utils._array import as_numpy, to_numpy
 from dataeval.utils._image import normalize_image_shape, rescale
 
 DTYPE_REGEX = re.compile(r"NDArray\[np\.(.*?)\]")
@@ -138,19 +138,19 @@ def process_stats(
 
 
 def process_stats_unpack(
-    args: tuple[int, Array, list[BoundingBox] | None],
+    args: tuple[int, ArrayLike, list[BoundingBox] | None],
     per_channel: bool,
     stats_processor_cls: Iterable[type[StatsProcessor[TStatsOutput]]],
 ) -> StatsProcessorOutput:
     return process_stats(*args, per_channel=per_channel, stats_processor_cls=stats_processor_cls)
 
 
-def _enumerate(dataset: Dataset[Array] | Dataset[tuple[Array, Any, Any]], per_box: bool):
+def _enumerate(dataset: Dataset[ArrayLike] | Dataset[tuple[ArrayLike, Any, Any]], per_box: bool):
     for i in range(len(dataset)):
         d = dataset[i]
         image = d[0] if isinstance(d, tuple) else d
         if per_box and isinstance(d, tuple) and isinstance(d[1], ObjectDetectionTarget):
-            boxes = cast(Array, d[1].boxes)
+            boxes = d[1].boxes if isinstance(d[1].boxes, Array) else as_numpy(d[1].boxes)
             target = [BoundingBox(float(box[i]) for i in range(4)) for box in boxes]
         else:
             target = None
@@ -159,7 +159,7 @@ def _enumerate(dataset: Dataset[Array] | Dataset[tuple[Array, Any, Any]], per_bo
 
 
 def run_stats(
-    dataset: Dataset[Array] | Dataset[tuple[Array, Any, Any]],
+    dataset: Dataset[ArrayLike] | Dataset[tuple[ArrayLike, Any, Any]],
     per_box: bool,
     per_channel: bool,
     stats_processor_cls: Iterable[type[StatsProcessor[TStatsOutput]]],
@@ -173,7 +173,7 @@ def run_stats(
 
     Parameters
     ----------
-    data : Dataset[Array] | Dataset[tuple[Array, Any, Any]]
+    data : Dataset[ArrayLike] | Dataset[tuple[ArrayLike, Any, Any]]
         A dataset of images and targets to compute statistics on.
     per_box : bool
         A flag which determines if the statistics should be evaluated on a per-box basis or not.

dataeval/metrics/stats/_dimensionstats.py

@@ -9,7 +9,7 @@ import numpy as np
 from dataeval.metrics.stats._base import StatsProcessor, run_stats
 from dataeval.outputs import DimensionStatsOutput
 from dataeval.outputs._base import set_metadata
-from dataeval.typing import Array, Dataset
+from dataeval.typing import ArrayLike, Dataset
 from dataeval.utils._image import get_bitdepth
 
 
@@ -34,7 +34,7 @@ class DimensionStatsProcessor(StatsProcessor[DimensionStatsOutput]):
 
 @set_metadata
 def dimensionstats(
-    dataset: Dataset[Array] | Dataset[tuple[Array, Any, Any]],
+    dataset: Dataset[ArrayLike] | Dataset[tuple[ArrayLike, Any, Any]],
     *,
     per_box: bool = False,
 ) -> DimensionStatsOutput:

dataeval/metrics/stats/_hashstats.py

@@ -14,7 +14,7 @@ from scipy.fftpack import dct
 from dataeval.metrics.stats._base import StatsProcessor, run_stats
 from dataeval.outputs import HashStatsOutput
 from dataeval.outputs._base import set_metadata
-from dataeval.typing import Array, ArrayLike, Dataset
+from dataeval.typing import ArrayLike, Dataset
 from dataeval.utils._array import as_numpy
 from dataeval.utils._image import normalize_image_shape, rescale
 
@@ -105,7 +105,7 @@ class HashStatsProcessor(StatsProcessor[HashStatsOutput]):
 
 @set_metadata
 def hashstats(
-    dataset: Dataset[Array] | Dataset[tuple[Array, Any, Any]],
+    dataset: Dataset[ArrayLike] | Dataset[tuple[ArrayLike, Any, Any]],
     *,
     per_box: bool = False,
 ) -> HashStatsOutput:

dataeval/metrics/stats/_imagestats.py

@@ -10,12 +10,12 @@ from dataeval.metrics.stats._pixelstats import PixelStatsProcessor
 from dataeval.metrics.stats._visualstats import VisualStatsProcessor
 from dataeval.outputs import ChannelStatsOutput, ImageStatsOutput
 from dataeval.outputs._base import set_metadata
-from dataeval.typing import Array, Dataset
+from dataeval.typing import ArrayLike, Dataset
 
 
 @overload
 def imagestats(
-    dataset: Dataset[Array] | Dataset[tuple[Array, Any, Any]],
+    dataset: Dataset[ArrayLike] | Dataset[tuple[ArrayLike, Any, Any]],
     *,
     per_box: bool = False,
     per_channel: Literal[True],
@@ -24,7 +24,7 @@ def imagestats(
 
 @overload
 def imagestats(
-    dataset: Dataset[Array] | Dataset[tuple[Array, Any, Any]],
+    dataset: Dataset[ArrayLike] | Dataset[tuple[ArrayLike, Any, Any]],
     *,
     per_box: bool = False,
     per_channel: Literal[False] = False,
@@ -33,7 +33,7 @@ def imagestats(
 
 @set_metadata
 def imagestats(
-    dataset: Dataset[Array] | Dataset[tuple[Array, Any, Any]],
+    dataset: Dataset[ArrayLike] | Dataset[tuple[ArrayLike, Any, Any]],
     *,
     per_box: bool = False,
     per_channel: bool = False,

dataeval/metrics/stats/_labelstats.py

@@ -5,10 +5,10 @@ __all__ = []
 from collections import Counter, defaultdict
 from typing import Any, Mapping, TypeVar
 
+from dataeval.data._metadata import Metadata
 from dataeval.outputs import LabelStatsOutput
 from dataeval.outputs._base import set_metadata
 from dataeval.typing import AnnotatedDataset
-from dataeval.utils.data._metadata import Metadata
 
 TValue = TypeVar("TValue")
 
@@ -38,7 +38,7 @@ def labelstats(dataset: Metadata | AnnotatedDataset[Any]) -> LabelStatsOutput:
     --------
     Calculate basic :term:`statistics<Statistics>` on labels for a dataset.
 
-    >>> from dataeval.utils.data import Metadata
+    >>> from dataeval.data import Metadata
     >>> stats = labelstats(Metadata(dataset))
     >>> print(stats.to_table())
     Class Count: 5

dataeval/metrics/stats/_pixelstats.py

@@ -10,7 +10,7 @@ from scipy.stats import entropy, kurtosis, skew
 from dataeval.metrics.stats._base import StatsProcessor, run_stats
 from dataeval.outputs import PixelStatsOutput
 from dataeval.outputs._base import set_metadata
-from dataeval.typing import Array, Dataset
+from dataeval.typing import ArrayLike, Dataset
 
 
 class PixelStatsProcessor(StatsProcessor[PixelStatsOutput]):
@@ -37,7 +37,7 @@ class PixelStatsProcessor(StatsProcessor[PixelStatsOutput]):
 
 @set_metadata
 def pixelstats(
-    dataset: Dataset[Array] | Dataset[tuple[Array, Any, Any]],
+    dataset: Dataset[ArrayLike] | Dataset[tuple[ArrayLike, Any, Any]],
     *,
     per_box: bool = False,
     per_channel: bool = False,

dataeval/metrics/stats/_visualstats.py

@@ -9,7 +9,7 @@ import numpy as np
 from dataeval.metrics.stats._base import StatsProcessor, run_stats
 from dataeval.outputs import VisualStatsOutput
 from dataeval.outputs._base import set_metadata
-from dataeval.typing import Array, Dataset
+from dataeval.typing import ArrayLike, Dataset
 from dataeval.utils._image import edge_filter
 
 QUARTILES = (0, 25, 50, 75, 100)
@@ -44,7 +44,7 @@ class VisualStatsProcessor(StatsProcessor[VisualStatsOutput]):
 
 @set_metadata
 def visualstats(
-    dataset: Dataset[Array] | Dataset[tuple[Array, Any, Any]],
+    dataset: Dataset[ArrayLike] | Dataset[tuple[ArrayLike, Any, Any]],
     *,
     per_box: bool = False,
     per_channel: bool = False,

dataeval/outputs/_bias.py

@@ -13,11 +13,11 @@ with contextlib.suppress(ImportError):
     import pandas as pd
     from matplotlib.figure import Figure
 
+from dataeval.data._images import Images
 from dataeval.outputs._base import Output
 from dataeval.typing import ArrayLike, Dataset
 from dataeval.utils._array import as_numpy, channels_first_to_last
 from dataeval.utils._plot import heatmap
-from dataeval.utils.data._images import Images
 
 TData = TypeVar("TData", np.float64, NDArray[np.float64])
 

dataeval/typing.py

@@ -21,8 +21,9 @@ __all__ = [
 
 
 import sys
-from typing import Any, Generic, Iterator, Protocol, Sequence, TypedDict, TypeVar, Union, runtime_checkable
+from typing import Any, Generic, Iterator, Protocol, TypedDict, TypeVar, runtime_checkable
 
+import numpy.typing
 from typing_extensions import NotRequired, ReadOnly, Required
 
 if sys.version_info >= (3, 10):
@@ -31,6 +32,16 @@ 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):
     """
@@ -67,16 +78,8 @@ class Array(Protocol):
     def __len__(self) -> int: ...
 
 
-T = TypeVar("T")
+_T = TypeVar("_T")
 _T_co = TypeVar("_T_co", covariant=True)
-_ScalarType = Union[int, float, bool, str]
-ArrayLike: TypeAlias = Union[Sequence[_ScalarType], Sequence[Sequence[_ScalarType]], Sequence[Array], Array]
-"""
-Type alias for array-like objects used for interoperability with DataEval.
-
-This includes native Python sequences, as well as objects that conform to
-the :class:`Array` protocol.
-"""
 
 
 class DatasetMetadata(TypedDict, total=False):
@@ -95,6 +98,22 @@ class DatasetMetadata(TypedDict, total=False):
     index2label: NotRequired[ReadOnly[dict[int, str]]]
 
 
+class ModelMetadata(TypedDict, total=False):
+    """
+    Model metadata required for all `AnnotatedModel` classes.
+
+    Attributes
+    ----------
+    id : Required[str]
+        A unique identifier for the model
+    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):
     """
@@ -140,12 +159,12 @@ class AnnotatedDataset(Dataset[_T_co], Generic[_T_co], Protocol):
 # ========== IMAGE CLASSIFICATION DATASETS ==========
 
 
-ImageClassificationDatum: TypeAlias = tuple[Array, Array, dict[str, Any]]
+ImageClassificationDatum: TypeAlias = tuple[ArrayLike, ArrayLike, dict[str, Any]]
 """
 Type alias for an image classification datum tuple.
 
-- :class:`Array` of shape (C, H, W) - Image data in channel, height, width format.
-- :class:`Array` of shape (N,) - Class label as one-hot encoded ground-truth or prediction confidences.
+- :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.
 """
 
@@ -180,11 +199,11 @@ class ObjectDetectionTarget(Protocol):
     def scores(self) -> ArrayLike: ...
 
 
-ObjectDetectionDatum: TypeAlias = tuple[Array, ObjectDetectionTarget, dict[str, Any]]
+ObjectDetectionDatum: TypeAlias = tuple[ArrayLike, ObjectDetectionTarget, dict[str, Any]]
 """
 Type alias for an object detection datum tuple.
 
-- :class:`Array` of shape (C, H, W) - Image data in channel, height, width format.
+- :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.
 """
@@ -221,11 +240,11 @@ class SegmentationTarget(Protocol):
     def scores(self) -> ArrayLike: ...
 
 
-SegmentationDatum: TypeAlias = tuple[Array, SegmentationTarget, dict[str, Any]]
+SegmentationDatum: TypeAlias = tuple[ArrayLike, SegmentationTarget, dict[str, Any]]
 """
 Type alias for an image classification datum tuple.
 
-- :class:`Array` of shape (C, H, W) - Image data in channel, height, width format.
+- :class:`ArrayLike` of shape (C, H, W) - Image data in channel, height, width format.
 - :class:`SegmentationTarget` - Segmentation target information for the image.
 - dict[str, Any] - Datum level metadata.
 """
@@ -235,9 +254,24 @@ SegmentationDataset: TypeAlias = AnnotatedDataset[SegmentationDatum]
 Type alias for an :class:`AnnotatedDataset` of :class:`SegmentationDatum` elements.
 """
 
+# ========== MODEL ==========
+
+
+@runtime_checkable
+class AnnotatedModel(Protocol):
+    """
+    Protocol for an annotated model.
+    """
+
+    @property
+    def metadata(self) -> ModelMetadata: ...
+
+
+# ========== TRANSFORM ==========
+
 
 @runtime_checkable
-class Transform(Generic[T], Protocol):
+class Transform(Generic[_T], Protocol):
     """
     Protocol defining a transform function.
 
@@ -262,4 +296,4 @@ class Transform(Generic[T], Protocol):
     array([0.004, 0.008, 0.012])
     """
 
-    def __call__(self, data: T, /) -> T: ...
+    def __call__(self, data: _T, /) -> _T: ...

dataeval/utils/__init__.py

@@ -4,6 +4,6 @@ in setting up data and architectures that are guaranteed to work with applicable
 DataEval metrics.
 """
 
-__all__ = ["data", "metadata", "torch"]
+__all__ = ["data", "datasets", "torch"]
 
-from . import data, metadata, torch
+from . import data, datasets, torch

dataeval/utils/_array.py

@@ -92,7 +92,7 @@ def ensure_embeddings(
 @overload
 def ensure_embeddings(
     embeddings: T,
-    dtype: None,
+    dtype: None = None,
     unit_interval: Literal[True, False, "force"] = False,
 ) -> T: ...
 
@@ -152,21 +152,32 @@ def ensure_embeddings(
         return arr
 
 
-def flatten(array: ArrayLike) -> NDArray[Any]:
+@overload
+def flatten(array: torch.Tensor) -> torch.Tensor: ...
+@overload
+def flatten(array: ArrayLike) -> NDArray[Any]: ...
+
+
+def flatten(array: ArrayLike) -> NDArray[Any] | torch.Tensor:
     """
     Flattens input array from (N, ... ) to (N, -1) where all samples N have all data in their last dimension
 
     Parameters
     ----------
-    X : NDArray, shape - (N, ... )
+    array : ArrayLike
         Input array
 
     Returns
     -------
-    NDArray, shape - (N, -1)
+    np.ndarray or torch.Tensor, shape: (N, -1)
     """
-    nparr = as_numpy(array)
-    return nparr.reshape((nparr.shape[0], -1))
+    if isinstance(array, np.ndarray):
+        nparr = as_numpy(array)
+        return nparr.reshape((nparr.shape[0], -1))
+    elif isinstance(array, torch.Tensor):
+        return torch.flatten(array, start_dim=1)
+    else:
+        raise TypeError(f"Unsupported array type {type(array)}.")
 
 
 _TArray = TypeVar("_TArray", bound=Array)
@@ -191,4 +202,4 @@ def channels_first_to_last(array: _TArray) -> _TArray:
     elif isinstance(array, torch.Tensor):
         return torch.permute(array, (1, 2, 0))
     else:
-        raise TypeError(f"Unsupported array type {type(array)} for conversion.")
+        raise TypeError(f"Unsupported array type {type(array)}.")

dataeval/utils/data/__init__.py

@@ -1,26 +1,11 @@
-"""Provides utility functions for interacting with Computer Vision datasets."""
+"""Provides access to common Computer Vision datasets."""
+
+from dataeval.utils.data import collate, metadata
+from dataeval.utils.data._dataset import to_image_classification_dataset, to_object_detection_dataset
 
 __all__ = [
     "collate",
-    "datasets",
-    "Embeddings",
-    "Images",
-    "Metadata",
-    "Select",
-    "SplitDatasetOutput",
-    "Targets",
-    "split_dataset",
+    "metadata",
     "to_image_classification_dataset",
     "to_object_detection_dataset",
 ]
-
-from dataeval.outputs._utils import SplitDatasetOutput
-from dataeval.utils.data._dataset import to_image_classification_dataset, to_object_detection_dataset
-from dataeval.utils.data._embeddings import Embeddings
-from dataeval.utils.data._images import Images
-from dataeval.utils.data._metadata import Metadata
-from dataeval.utils.data._selection import Select
-from dataeval.utils.data._split import split_dataset
-from dataeval.utils.data._targets import Targets
-
-from . import collate, datasets