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

dataeval/utils/data/_dataset.py

@@ -47,6 +47,8 @@ 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:

dataeval/utils/data/_embeddings.py

@@ -3,14 +3,14 @@ from __future__ import annotations
 __all__ = []
 
 import math
-from typing import Any, Iterator, Sequence
+from typing import Any, Iterator, Sequence, cast
 
 import torch
 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, Dataset, Transform
 from dataeval.utils.torch.models import SupportsEncode
 
 
@@ -26,11 +26,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.
     """
@@ -43,61 +47,140 @@ class Embeddings:
         self,
         dataset: Dataset[tuple[Array, Any, Any]],
         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._transforms = [transforms] if isinstance(transforms, Transform) else transforms
         self._model = model.to(self.device).eval()
         self._encoder = model.encode if isinstance(model, SupportsEncode) else model
         self._collate_fn = lambda datum: [torch.as_tensor(i) for i, _, _ in datum]
+        self._cached_idx = set()
+        self._embeddings: torch.Tensor = torch.empty(())
+        self._shallow: bool = False
 
-    def to_tensor(self) -> torch.Tensor:
+    def to_tensor(self, indices: Sequence[int] | None = None) -> torch.Tensor:
         """
-        Converts entire dataset to embeddings.
+        Converts dataset to embeddings.
 
-        Warning
-        -------
-        Will process the entire dataset in batches and return
-        embeddings as a single Tensor in memory.
+        Parameters
+        ----------
+        indices : Sequence[int] or None, default None
+            The indices to convert to embeddings
 
         Returns
         -------
         torch.Tensor
+
+        Warning
+        -------
+        Processing large quantities of data can be resource intensive.
+        """
+        if indices is not None:
+            return torch.vstack(list(self._batch(indices))).to(self.device)
+        else:
+            return self[:]
+
+    @classmethod
+    def from_array(cls, array: Array, device: DeviceLike | None = None) -> Embeddings:
         """
-        return self[:]
+        Instantiates a shallow Embeddings object using an array.
 
-    # Reduce overhead cost by not tracking tensor gradients
-    @torch.no_grad
+        Parameters
+        ----------
+        array : Array
+            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)
+        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
-
-    def __getitem__(self, key: int | slice | list[int], /) -> torch.Tensor:
-        if isinstance(key, list):
-            return torch.vstack(list(self._batch(key))).to(self.device)
-        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.")
+        dataset = cast(torch.utils.data.Dataset[tuple[Array, Any, Any]], 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 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, 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", bound=Array)
 
 
 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/_selection.py

@@ -7,7 +7,7 @@ from typing import Generic, Iterator, Sequence, TypeVar
 
 from dataeval.typing import AnnotatedDataset, DatasetMetadata
 
-_TDatum = TypeVar("_TDatum", covariant=True)
+_TDatum = TypeVar("_TDatum")
 
 
 class SelectionStage(IntEnum):
@@ -67,13 +67,13 @@ class Select(AnnotatedDataset[_TDatum]):
     def __init__(
         self,
         dataset: AnnotatedDataset[_TDatum],
-        selections: Selection[_TDatum] | list[Selection[_TDatum]] | None = None,
+        selections: Selection[_TDatum] | Sequence[Selection[_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(selections)
+        self._selections = self._sort(selections)
 
         # Ensure metadata is populated correctly as DatasetMetadata TypedDict
         _metadata = getattr(dataset, "metadata", {})
@@ -81,8 +81,7 @@ class Select(AnnotatedDataset[_TDatum]):
             _metadata["id"] = dataset.__class__.__name__
         self._metadata = DatasetMetadata(**_metadata)
 
-        if self._selections:
-            self._apply_selections()
+        self._select()
 
     @property
     def metadata(self) -> DatasetMetadata:
@@ -92,10 +91,10 @@ class Select(AnnotatedDataset[_TDatum]):
         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)])}]"
+        selections = f"Selections: [{', '.join([str(s) for s in self._selections])}]"
         return f"{title}\n{sep}{nt}{selections}{nt}Selected Size: {len(self)}\n\n{self._dataset}"
 
-    def _sort_selections(self, selections: Selection[_TDatum] | Sequence[Selection[_TDatum]] | None) -> list[Selection]:
+    def _sort(self, selections: Selection[_TDatum] | Sequence[Selection[_TDatum]] | None) -> list[Selection]:
         if not selections:
             return []
 
@@ -106,7 +105,7 @@ class Select(AnnotatedDataset[_TDatum]):
         selection_list = [selection for category in sorted(grouped) for selection in grouped[category]]
         return selection_list
 
-    def _apply_selections(self) -> None:
+    def _select(self) -> None:
         for selection in self._selections:
             selection(self)
         self._selection = self._selection[: self._size_limit]