dataeval 0.83.0__py3-none-any.whl → 0.84.1__py3-none-any.whl

dataeval/typing.py

@@ -21,9 +21,10 @@ __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
 
-from typing_extensions import NotRequired, Required
+import numpy.typing
+from typing_extensions import NotRequired, ReadOnly, Required
 
 if sys.version_info >= (3, 10):
     from typing import TypeAlias
@@ -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):
@@ -91,8 +94,8 @@ class DatasetMetadata(TypedDict, total=False):
         A lookup table converting label value to class name
     """
 
-    id: Required[str]
-    index2label: NotRequired[dict[int, str]]
+    id: Required[ReadOnly[str]]
+    index2label: NotRequired[ReadOnly[dict[int, str]]]
 
 
 @runtime_checkable
@@ -140,12 +143,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 +183,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 +224,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.
 """
@@ -237,7 +240,7 @@ Type alias for an :class:`AnnotatedDataset` of :class:`SegmentationDatum` elemen
 
 
 @runtime_checkable
-class Transform(Generic[T], Protocol):
+class Transform(Generic[_T], Protocol):
     """
     Protocol defining a transform function.
 
@@ -262,4 +265,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/_array.py

@@ -13,7 +13,7 @@ import torch
 from numpy.typing import NDArray
 
 from dataeval._log import LogMessage
-from dataeval.typing import ArrayLike
+from dataeval.typing import Array, ArrayLike
 
 _logger = logging.getLogger(__name__)
 
@@ -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,18 +152,54 @@ 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)
+    """
+    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)
+
+
+def channels_first_to_last(array: _TArray) -> _TArray:
     """
-    nparr = as_numpy(array)
-    return nparr.reshape((nparr.shape[0], -1))
+    Converts array from channels first to channels last format
+
+    Parameters
+    ----------
+    array : ArrayLike
+        Input array
+
+    Returns
+    -------
+    ArrayLike
+        Converted array
+    """
+    if isinstance(array, np.ndarray):
+        return np.transpose(array, (1, 2, 0))
+    elif isinstance(array, torch.Tensor):
+        return torch.permute(array, (1, 2, 0))
+    else:
+        raise TypeError(f"Unsupported array type {type(array)}.")

dataeval/utils/data/_dataset.py

@@ -47,13 +47,17 @@ def _validate_data(
             or not len(bboxes[0][0]) == 4
         ):
             raise TypeError("Boxes must be a sequence of sequences of (x0, y0, x1, y1) for object detection.")
+    else:
+        raise ValueError(f"Unknown datum type '{datum_type}'. Must be 'ic' or 'od'.")
 
 
 def _find_max(arr: ArrayLike) -> Any:
-    if isinstance(arr[0], (Iterable, Sequence, Array)):
-        return max([_find_max(x) for x in arr])  # type: ignore
-    else:
-        return max(arr)
+    if isinstance(arr, (Iterable, Sequence, Array)):
+        if isinstance(arr[0], (Iterable, Sequence, Array)):
+            return max([_find_max(x) for x in arr])  # type: ignore
+        else:
+            return max(arr)
+    return arr
 
 
 _TLabels = TypeVar("_TLabels", Sequence[int], Sequence[Sequence[int]])

dataeval/utils/data/_embeddings.py

@@ -3,14 +3,16 @@ from __future__ import annotations
 __all__ = []
 
 import math
-from typing import Any, Iterator, Sequence
+from typing import Any, Iterator, Sequence, cast
 
 import torch
+from numpy.typing import NDArray
 from torch.utils.data import DataLoader, Subset
 from tqdm import tqdm
 
 from dataeval.config import DeviceLike, get_device
-from dataeval.typing import Array, Dataset
+from dataeval.typing import Array, ArrayLike, Dataset, Transform
+from dataeval.utils._array import as_numpy
 from dataeval.utils.torch.models import SupportsEncode
 
 
@@ -26,11 +28,15 @@ class Embeddings:
         Dataset to access original images from.
     batch_size : int
         Batch size to use when encoding images.
+    transforms : Transform or Sequence[Transform] or None, default None
+        Transforms to apply to images before encoding.
     model : torch.nn.Module or None, default None
         Model to use for encoding images.
     device : DeviceLike or None, default None
         The hardware device to use if specified, otherwise uses the DataEval
         default or torch default.
+    cache : bool, default False
+        Whether to cache the embeddings in memory.
     verbose : bool, default False
         Whether to print progress bar when encoding images.
     """
@@ -41,21 +47,29 @@ class Embeddings:
 
     def __init__(
         self,
-        dataset: Dataset[tuple[Array, Any, Any]],
+        dataset: Dataset[tuple[ArrayLike, Any, Any]] | Dataset[ArrayLike],
         batch_size: int,
+        transforms: Transform[torch.Tensor] | Sequence[Transform[torch.Tensor]] | None = None,
         model: torch.nn.Module | None = None,
         device: DeviceLike | None = None,
+        cache: bool = False,
         verbose: bool = False,
     ) -> None:
         self.device = get_device(device)
-        self.batch_size = batch_size
+        self.cache = cache
+        self.batch_size = batch_size if batch_size > 0 else 1
         self.verbose = verbose
 
         self._dataset = dataset
+        self._length = len(dataset)
         model = torch.nn.Flatten() if model is None else model
-        self._model = model.to(self.device).eval()
+        self._transforms = [transforms] if isinstance(transforms, Transform) else transforms
+        self._model = model.to(self.device).eval() if isinstance(model, torch.nn.Module) else model
         self._encoder = model.encode if isinstance(model, SupportsEncode) else model
-        self._collate_fn = lambda datum: [torch.as_tensor(i) for i, _, _ in datum]
+        self._collate_fn = lambda datum: [torch.as_tensor(d[0] if isinstance(d, tuple) else d) for d in datum]
+        self._cached_idx = set()
+        self._embeddings: torch.Tensor = torch.empty(())
+        self._shallow: bool = False
 
     def to_tensor(self, indices: Sequence[int] | None = None) -> torch.Tensor:
         """
@@ -79,30 +93,133 @@ class Embeddings:
         else:
             return self[:]
 
-    # Reduce overhead cost by not tracking tensor gradients
-    @torch.no_grad
+    def to_numpy(self, indices: Sequence[int] | None = None) -> NDArray[Any]:
+        """
+        Converts dataset to embeddings as numpy array.
+
+        Parameters
+        ----------
+        indices : Sequence[int] or None, default None
+            The indices to convert to embeddings
+
+        Returns
+        -------
+        NDArray[Any]
+
+        Warning
+        -------
+        Processing large quantities of data can be resource intensive.
+        """
+        return self.to_tensor(indices).cpu().numpy()
+
+    def new(self, dataset: Dataset[tuple[ArrayLike, Any, Any]] | Dataset[ArrayLike]) -> Embeddings:
+        """
+        Creates a new Embeddings object with the same parameters but a different dataset.
+
+        Parameters
+        ----------
+        dataset : ImageClassificationDataset or ObjectDetectionDataset
+            Dataset to access original images from.
+
+        Returns
+        -------
+        Embeddings
+        """
+        return Embeddings(
+            dataset, self.batch_size, self._transforms, self._model, self.device, self.cache, self.verbose
+        )
+
+    @classmethod
+    def from_array(cls, array: ArrayLike, device: DeviceLike | None = None) -> Embeddings:
+        """
+        Instantiates a shallow Embeddings object using an array.
+
+        Parameters
+        ----------
+        array : ArrayLike
+            The array to convert to embeddings.
+        device : DeviceLike or None, default None
+            The hardware device to use if specified, otherwise uses the DataEval
+            default or torch default.
+
+        Returns
+        -------
+        Embeddings
+
+        Example
+        -------
+        >>> import numpy as np
+        >>> from dataeval.utils.data._embeddings import Embeddings
+        >>> array = np.random.randn(100, 3, 224, 224)
+        >>> embeddings = Embeddings.from_array(array)
+        >>> print(embeddings.to_tensor().shape)
+        torch.Size([100, 3, 224, 224])
+        """
+        embeddings = Embeddings([], 0, None, None, device, True, False)
+        array = array if isinstance(array, Array) else as_numpy(array)
+        embeddings._length = len(array)
+        embeddings._cached_idx = set(range(len(array)))
+        embeddings._embeddings = torch.as_tensor(array).to(get_device(device))
+        embeddings._shallow = True
+        return embeddings
+
+    def _encode(self, images: list[torch.Tensor]) -> torch.Tensor:
+        if self._transforms:
+            images = [transform(image) for transform in self._transforms for image in images]
+        return self._encoder(torch.stack(images).to(self.device))
+
+    @torch.no_grad()  # Reduce overhead cost by not tracking tensor gradients
     def _batch(self, indices: Sequence[int]) -> Iterator[torch.Tensor]:
-        # manual batching
-        dataloader = DataLoader(Subset(self._dataset, indices), batch_size=self.batch_size, collate_fn=self._collate_fn)  # type: ignore
-        for i, images in (
-            tqdm(enumerate(dataloader), total=math.ceil(len(indices) / self.batch_size), desc="Batch processing")
-            if self.verbose
-            else enumerate(dataloader)
-        ):
-            embeddings = self._encoder(torch.stack(images).to(self.device))
-            yield embeddings
+        dataset = cast(torch.utils.data.Dataset, self._dataset)
+        total_batches = math.ceil(len(indices) / self.batch_size)
+
+        # If not caching, process all indices normally
+        if not self.cache:
+            for images in tqdm(
+                DataLoader(Subset(dataset, indices), self.batch_size, collate_fn=self._collate_fn),
+                total=total_batches,
+                desc="Batch embedding",
+                disable=not self.verbose,
+            ):
+                yield self._encode(images)
+            return
+
+        # If caching, process each batch of indices at a time, preserving original order
+        for i in tqdm(range(0, len(indices), self.batch_size), desc="Batch embedding", disable=not self.verbose):
+            batch = indices[i : i + self.batch_size]
+            uncached = [idx for idx in batch if idx not in self._cached_idx]
+
+            if uncached:
+                # Process uncached indices as as single batch
+                for images in DataLoader(Subset(dataset, uncached), len(uncached), collate_fn=self._collate_fn):
+                    embeddings = self._encode(images)
+
+                    if not self._embeddings.shape:
+                        full_shape = (len(self._dataset), *embeddings.shape[1:])
+                        self._embeddings = torch.empty(full_shape, dtype=embeddings.dtype, device=self.device)
+
+                    self._embeddings[uncached] = embeddings
+                    self._cached_idx.update(uncached)
+
+            yield self._embeddings[batch]
 
     def __getitem__(self, key: int | slice, /) -> torch.Tensor:
-        if isinstance(key, slice):
-            return torch.vstack(list(self._batch(range(len(self._dataset))[key]))).to(self.device)
-        elif isinstance(key, int):
-            return self._encoder(torch.as_tensor(self._dataset[key][0]).to(self.device))
-        raise TypeError("Invalid argument type.")
+        if not isinstance(key, slice) and not hasattr(key, "__int__"):
+            raise TypeError("Invalid argument type.")
+
+        if self._shallow:
+            if not self._embeddings.shape:
+                raise ValueError("Embeddings not initialized.")
+            return self._embeddings[key]
+
+        indices = list(range(len(self._dataset))[key]) if isinstance(key, slice) else [int(key)]
+        result = torch.vstack(list(self._batch(indices))).to(self.device)
+        return result.squeeze(0) if len(indices) == 1 else result
 
     def __iter__(self) -> Iterator[torch.Tensor]:
         # process in batches while yielding individual embeddings
-        for batch in self._batch(range(len(self._dataset))):
+        for batch in self._batch(range(self._length)):
             yield from batch
 
     def __len__(self) -> int:
-        return len(self._dataset)
+        return self._length

dataeval/utils/data/_images.py

@@ -2,11 +2,15 @@ from __future__ import annotations
 
 __all__ = []
 
-from typing import Any, Generic, Iterator, Sequence, TypeVar, cast, overload
+from typing import TYPE_CHECKING, Any, Generic, Iterator, Sequence, TypeVar, cast, overload
 
-from dataeval.typing import Dataset
+from dataeval.typing import Array, ArrayLike, Dataset
+from dataeval.utils._array import as_numpy, channels_first_to_last
 
-T = TypeVar("T")
+if TYPE_CHECKING:
+    from matplotlib.figure import Figure
+
+T = TypeVar("T", Array, ArrayLike)
 
 
 class Images(Generic[T]):
@@ -21,7 +25,10 @@ class Images(Generic[T]):
         Dataset to access images from.
     """
 
-    def __init__(self, dataset: Dataset[tuple[T, Any, Any] | T]) -> None:
+    def __init__(
+        self,
+        dataset: Dataset[tuple[T, Any, Any] | T],
+    ) -> None:
         self._is_tuple_datum = isinstance(dataset[0], tuple)
         self._dataset = dataset
 
@@ -40,25 +47,41 @@ class Images(Generic[T]):
         """
         return self[:]
 
+    def plot(
+        self,
+        indices: Sequence[int],
+        images_per_row: int = 3,
+        figsize: tuple[int, int] = (10, 10),
+    ) -> Figure:
+        import matplotlib.pyplot as plt
+
+        num_images = len(indices)
+        num_rows = (num_images + images_per_row - 1) // images_per_row
+        fig, axes = plt.subplots(num_rows, images_per_row, figsize=figsize)
+        for i, ax in enumerate(axes.flatten()):
+            image = channels_first_to_last(as_numpy(self[i]))
+            ax.imshow(image)
+            ax.axis("off")
+        plt.tight_layout()
+        return fig
+
     @overload
     def __getitem__(self, key: int, /) -> T: ...
     @overload
     def __getitem__(self, key: slice, /) -> Sequence[T]: ...
 
     def __getitem__(self, key: int | slice, /) -> Sequence[T] | T:
+        if isinstance(key, slice):
+            return [self._get_image(k) for k in range(len(self._dataset))[key]]
+        elif hasattr(key, "__int__"):
+            return self._get_image(int(key))
+        raise TypeError(f"Key must be integers or slices, not {type(key)}")
+
+    def _get_image(self, index: int) -> T:
         if self._is_tuple_datum:
-            dataset = cast(Dataset[tuple[T, Any, Any]], self._dataset)
-            if isinstance(key, slice):
-                return [dataset[k][0] for k in range(len(self._dataset))[key]]
-            elif isinstance(key, int):
-                return dataset[key][0]
+            return cast(Dataset[tuple[T, Any, Any]], self._dataset)[index][0]
         else:
-            dataset = cast(Dataset[T], self._dataset)
-            if isinstance(key, slice):
-                return [dataset[k] for k in range(len(self._dataset))[key]]
-            elif isinstance(key, int):
-                return dataset[key]
-        raise TypeError(f"Key must be integers or slices, not {type(key)}")
+            return cast(Dataset[T], self._dataset)[index]
 
     def __iter__(self) -> Iterator[T]:
         for i in range(len(self._dataset)):

dataeval/utils/data/_metadata.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 __all__ = []
 
 import warnings
-from typing import TYPE_CHECKING, Any, Literal, Mapping, Sequence, cast
+from typing import TYPE_CHECKING, Any, Literal, Mapping, Sequence, Sized, cast
 
 import numpy as np
 from numpy.typing import NDArray
@@ -208,8 +208,9 @@ class Metadata:
             raw.append(metadata)
 
             if is_od_target := isinstance(target, ObjectDetectionTarget):
-                target_len = len(target.labels)
-                labels.extend(as_numpy(target.labels).tolist())
+                target_labels = as_numpy(target.labels)
+                target_len = len(target_labels)
+                labels.extend(target_labels.tolist())
                 bboxes.extend(as_numpy(target.boxes).tolist())
                 scores.extend(as_numpy(target.scores).tolist())
                 srcidx.extend([i] * target_len)
@@ -360,7 +361,7 @@ class Metadata:
         self._merge()
         self._processed = False
         target_len = len(self.targets.source) if self.targets.source is not None else len(self.targets)
-        if any(len(v) != target_len for v in factors.values()):
+        if any(len(v if isinstance(v, Sized) else as_numpy(v)) != target_len for v in factors.values()):
             raise ValueError(
                 "The lists/arrays in the provided factors have a different length than the current metadata factors."
             )

dataeval/utils/data/_selection.py

@@ -5,7 +5,7 @@ __all__ = []
 from enum import IntEnum
 from typing import Generic, Iterator, Sequence, TypeVar
 
-from dataeval.typing import AnnotatedDataset, DatasetMetadata, Transform
+from dataeval.typing import AnnotatedDataset, DatasetMetadata
 
 _TDatum = TypeVar("_TDatum")
 
@@ -35,8 +35,6 @@ class Select(AnnotatedDataset[_TDatum]):
         The dataset to wrap.
     selections : Selection or list[Selection], optional
         The selection criteria to apply to the dataset.
-    transforms : Transform or list[Transform], optional
-        The transforms to apply to the dataset.
 
     Examples
     --------
@@ -70,16 +68,12 @@ class Select(AnnotatedDataset[_TDatum]):
         self,
         dataset: AnnotatedDataset[_TDatum],
         selections: Selection[_TDatum] | Sequence[Selection[_TDatum]] | None = None,
-        transforms: Transform[_TDatum] | Sequence[Transform[_TDatum]] | None = None,
     ) -> None:
         self.__dict__.update(dataset.__dict__)
         self._dataset = dataset
         self._size_limit = len(dataset)
         self._selection = list(range(self._size_limit))
         self._selections = self._sort(selections)
-        self._transforms = (
-            [] if transforms is None else [transforms] if isinstance(transforms, Transform) else transforms
-        )
 
         # Ensure metadata is populated correctly as DatasetMetadata TypedDict
         _metadata = getattr(dataset, "metadata", {})
@@ -98,8 +92,7 @@ class Select(AnnotatedDataset[_TDatum]):
         title = f"{self.__class__.__name__} Dataset"
         sep = "-" * len(title)
         selections = f"Selections: [{', '.join([str(s) for s in self._selections])}]"
-        transforms = f"Transforms: [{', '.join([str(t) for t in self._transforms])}]"
-        return f"{title}\n{sep}{nt}{selections}{nt}{transforms}{nt}Selected Size: {len(self)}\n\n{self._dataset}"
+        return f"{title}\n{sep}{nt}{selections}{nt}Selected Size: {len(self)}\n\n{self._dataset}"
 
     def _sort(self, selections: Selection[_TDatum] | Sequence[Selection[_TDatum]] | None) -> list[Selection]:
         if not selections:
@@ -117,13 +110,8 @@ class Select(AnnotatedDataset[_TDatum]):
             selection(self)
         self._selection = self._selection[: self._size_limit]
 
-    def _transform(self, datum: _TDatum) -> _TDatum:
-        for t in self._transforms:
-            datum = t(datum)
-        return datum
-
     def __getitem__(self, index: int) -> _TDatum:
-        return self._transform(self._dataset[self._selection[index]])
+        return self._dataset[self._selection[index]]
 
     def __iter__(self) -> Iterator[_TDatum]:
         for i in range(len(self)):