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

dataeval/utils/data/_dataset.py

@@ -0,0 +1,217 @@
+from __future__ import annotations
+
+__all__ = []
+
+from typing import Any, Generic, Iterable, Literal, Sequence, TypeVar
+
+from dataeval.typing import (
+    Array,
+    ArrayLike,
+    DatasetMetadata,
+    ImageClassificationDataset,
+    ObjectDetectionDataset,
+)
+from dataeval.utils._array import as_numpy
+
+
+def _validate_data(
+    datum_type: Literal["ic", "od"],
+    images: Array | Sequence[Array],
+    labels: Sequence[int] | Sequence[Sequence[int]],
+    bboxes: Sequence[Sequence[Sequence[float]]] | None,
+    metadata: Sequence[dict[str, Any]] | None,
+) -> None:
+    # Validate inputs
+    dataset_len = len(images)
+
+    if not isinstance(images, (Sequence, Array)) or len(images[0].shape) != 3:
+        raise ValueError("Images must be a sequence or array of 3 dimensional arrays (H, W, C).")
+    if len(labels) != dataset_len:
+        raise ValueError(f"Number of labels ({len(labels)}) does not match number of images ({dataset_len}).")
+    if bboxes is not None and len(bboxes) != dataset_len:
+        raise ValueError(f"Number of bboxes ({len(bboxes)}) does not match number of images ({dataset_len}).")
+    if metadata is not None and len(metadata) != dataset_len:
+        raise ValueError(f"Number of metadata ({len(metadata)}) does not match number of images ({dataset_len}).")
+
+    if datum_type == "ic":
+        if not isinstance(labels, Sequence) or not isinstance(labels[0], int):
+            raise TypeError("Labels must be a sequence of integers for image classification.")
+    elif datum_type == "od":
+        if not isinstance(labels, Sequence) or not isinstance(labels[0], Sequence) or not isinstance(labels[0][0], int):
+            raise TypeError("Labels must be a sequence of sequences of integers for object detection.")
+        if (
+            bboxes is None
+            or not isinstance(bboxes, (Sequence, Array))
+            or not isinstance(bboxes[0], (Sequence, Array))
+            or not isinstance(bboxes[0][0], (Sequence, Array))
+            or not len(bboxes[0][0]) == 4
+        ):
+            raise TypeError("Boxes must be a sequence of sequences of (x0, y0, x1, y1) for object detection.")
+
+
+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)
+
+
+_TLabels = TypeVar("_TLabels", Sequence[int], Sequence[Sequence[int]])
+
+
+class BaseAnnotatedDataset(Generic[_TLabels]):
+    def __init__(
+        self,
+        datum_type: Literal["ic", "od"],
+        images: Array | Sequence[Array],
+        labels: _TLabels,
+        metadata: Sequence[dict[str, Any]] | None,
+        classes: Sequence[str] | None,
+        name: str | None = None,
+    ) -> None:
+        self._classes = classes if classes is not None else [str(i) for i in range(_find_max(labels) + 1)]
+        self._index2label = dict(enumerate(self._classes))
+        self._images = images
+        self._labels = labels
+        self._metadata = metadata
+        self._id = name or f"{len(self._images)}_image_{len(self._index2label)}_class_{datum_type}_dataset"
+
+    @property
+    def metadata(self) -> DatasetMetadata:
+        return DatasetMetadata(id=self._id, index2label=self._index2label)
+
+    def __len__(self) -> int:
+        return len(self._images)
+
+
+class CustomImageClassificationDataset(BaseAnnotatedDataset[Sequence[int]], ImageClassificationDataset):
+    def __init__(
+        self,
+        images: Array | Sequence[Array],
+        labels: Sequence[int],
+        metadata: Sequence[dict[str, Any]] | None,
+        classes: Sequence[str] | None,
+        name: str | None = None,
+    ) -> None:
+        super().__init__("ic", images, labels, metadata, classes)
+        if name is not None:
+            self.__name__ = name
+            self.__class__.__name__ = name
+            self.__class__.__qualname__ = name
+
+    def __getitem__(self, idx: int, /) -> tuple[Array, Array, dict[str, Any]]:
+        one_hot = [0.0] * len(self._index2label)
+        one_hot[self._labels[idx]] = 1.0
+        return (
+            self._images[idx],
+            as_numpy(one_hot),
+            self._metadata[idx] if self._metadata is not None else {},
+        )
+
+
+class CustomObjectDetectionDataset(BaseAnnotatedDataset[Sequence[Sequence[int]]], ObjectDetectionDataset):
+    class ObjectDetectionTarget:
+        def __init__(self, labels: Sequence[int], bboxes: Sequence[Sequence[float]]) -> None:
+            self._labels = labels
+            self._bboxes = bboxes
+            self._scores = [1.0] * len(labels)
+
+        @property
+        def labels(self) -> Sequence[int]:
+            return self._labels
+
+        @property
+        def boxes(self) -> Sequence[Sequence[float]]:
+            return self._bboxes
+
+        @property
+        def scores(self) -> Sequence[float]:
+            return self._scores
+
+    def __init__(
+        self,
+        images: Array | Sequence[Array],
+        labels: Sequence[Sequence[int]],
+        bboxes: Sequence[Sequence[Sequence[float]]],
+        metadata: Sequence[dict[str, Any]] | None,
+        classes: Sequence[str] | None,
+        name: str | None = None,
+    ) -> None:
+        super().__init__("od", images, labels, metadata, classes)
+        if name is not None:
+            self.__name__ = name
+            self.__class__.__name__ = name
+            self.__class__.__qualname__ = name
+        self._bboxes = bboxes
+
+    @property
+    def metadata(self) -> DatasetMetadata:
+        return DatasetMetadata(id=self._id, index2label=self._index2label)
+
+    def __getitem__(self, idx: int, /) -> tuple[Array, ObjectDetectionTarget, dict[str, Any]]:
+        return (
+            self._images[idx],
+            self.ObjectDetectionTarget(self._labels[idx], self._bboxes[idx]),
+            self._metadata[idx] if self._metadata is not None else {},
+        )
+
+
+def to_image_classification_dataset(
+    images: Array | Sequence[Array],
+    labels: Sequence[int],
+    metadata: Sequence[dict[str, Any]] | None,
+    classes: Sequence[str] | None,
+    name: str | None = None,
+) -> ImageClassificationDataset:
+    """
+    Helper function to create custom ImageClassificationDataset classes.
+
+    Parameters
+    ----------
+    images : Array | Sequence[Array]
+        The images to use in the dataset.
+    labels : Sequence[int]
+        The labels to use in the dataset.
+    metadata : Sequence[dict[str, Any]] | None
+        The metadata to use in the dataset.
+    classes : Sequence[str] | None
+        The classes to use in the dataset.
+
+    Returns
+    -------
+    ImageClassificationDataset
+    """
+    _validate_data("ic", images, labels, None, metadata)
+    return CustomImageClassificationDataset(images, labels, metadata, classes, name)
+
+
+def to_object_detection_dataset(
+    images: Array | Sequence[Array],
+    labels: Sequence[Sequence[int]],
+    bboxes: Sequence[Sequence[Sequence[float]]],
+    metadata: Sequence[dict[str, Any]] | None,
+    classes: Sequence[str] | None,
+    name: str | None = None,
+) -> ObjectDetectionDataset:
+    """
+    Helper function to create custom ObjectDetectionDataset classes.
+
+    Parameters
+    ----------
+    images : Array | Sequence[Array]
+        The images to use in the dataset.
+    labels : Sequence[Sequence[int]]
+        The labels to use in the dataset.
+    bboxes : Sequence[Sequence[Sequence[float]]]
+        The bounding boxes (x0,y0,x1,y0) to use in the dataset.
+    metadata : Sequence[dict[str, Any]] | None
+        The metadata to use in the dataset.
+    classes : Sequence[str] | None
+        The classes to use in the dataset.
+
+    Returns
+    -------
+    ObjectDetectionDataset
+    """
+    _validate_data("od", images, labels, bboxes, metadata)
+    return CustomObjectDetectionDataset(images, labels, bboxes, metadata, classes, name)

dataeval/utils/data/_embeddings.py

@@ -9,9 +9,8 @@ import torch
 from torch.utils.data import DataLoader, Subset
 from tqdm import tqdm
 
-from dataeval.config import get_device
-from dataeval.typing import TArray
-from dataeval.utils.data._types import Dataset
+from dataeval.config import DeviceLike, get_device
+from dataeval.typing import Array, Dataset
 from dataeval.utils.torch.models import SupportsEncode
 
 
@@ -25,13 +24,14 @@ class Embeddings:
     ----------
     dataset : ImageClassificationDataset or ObjectDetectionDataset
         Dataset to access original images from.
-    batch_size : int, optional
+    batch_size : int
         Batch size to use when encoding images.
-    model : torch.nn.Module, optional
+    model : torch.nn.Module or None, default None
         Model to use for encoding images.
-    device : torch.device, optional
-        Device to use for encoding images.
-    verbose : bool, optional
+    device : DeviceLike or None, default None
+        The hardware device to use if specified, otherwise uses the DataEval
+        default or torch default.
+    verbose : bool, default False
         Whether to print progress bar when encoding images.
     """
 
@@ -41,11 +41,10 @@ class Embeddings:
 
     def __init__(
         self,
-        dataset: Dataset[TArray, Any],
+        dataset: Dataset[tuple[Array, Any, Any]],
         batch_size: int,
-        indices: Sequence[int] | None = None,
         model: torch.nn.Module | None = None,
-        device: torch.device | str | None = None,
+        device: DeviceLike | None = None,
         verbose: bool = False,
     ) -> None:
         self.device = get_device(device)
@@ -53,7 +52,6 @@ class Embeddings:
         self.verbose = verbose
 
         self._dataset = dataset
-        self._indices = indices if indices is not None else range(len(dataset))
         model = torch.nn.Flatten() if model is None else model
         self._model = model.to(self.device).eval()
         self._encoder = model.encode if isinstance(model, SupportsEncode) else model
@@ -78,7 +76,7 @@ class Embeddings:
     @torch.no_grad
     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)
+        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
@@ -87,7 +85,7 @@ class Embeddings:
             embeddings = self._encoder(torch.stack(images).to(self.device))
             yield embeddings
 
-    def __getitem__(self, key: int | slice | list[int]) -> torch.Tensor:
+    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):

dataeval/utils/data/_images.py

@@ -2,13 +2,14 @@ from __future__ import annotations
 
 __all__ = []
 
-from typing import Any, Generic, Iterator, Sequence, overload
+from typing import Any, Generic, Iterator, Sequence, TypeVar, cast, overload
 
-from dataeval.typing import TArray
-from dataeval.utils.data._types import Dataset
+from dataeval.typing import Dataset
 
+T = TypeVar("T")
 
-class Images(Generic[TArray]):
+
+class Images(Generic[T]):
     """
     Collection of image data from a dataset.
 
@@ -16,17 +17,15 @@ class Images(Generic[TArray]):
 
     Parameters
     ----------
-    dataset : ImageClassificationDataset or ObjectDetectionDataset
+    dataset : Dataset[tuple[T, ...]] or Dataset[T]
         Dataset to access images from.
     """
 
-    def __init__(
-        self,
-        dataset: Dataset[TArray, Any],
-    ) -> None:
+    def __init__(self, dataset: Dataset[tuple[T, Any, Any] | T]) -> None:
+        self._is_tuple_datum = isinstance(dataset[0], tuple)
         self._dataset = dataset
 
-    def to_list(self) -> Sequence[TArray]:
+    def to_list(self) -> Sequence[T]:
         """
         Converts entire dataset to a sequence of images.
 
@@ -37,29 +36,33 @@ class Images(Generic[TArray]):
 
         Returns
         -------
-        list[TArray]
+        list[T]
         """
         return self[:]
 
     @overload
-    def __getitem__(self, key: slice | list[int]) -> Sequence[TArray]: ...
-
+    def __getitem__(self, key: int, /) -> T: ...
     @overload
-    def __getitem__(self, key: int) -> TArray: ...
-
-    def __getitem__(self, key: int | slice | list[int]) -> Sequence[TArray] | TArray:
-        if isinstance(key, list):
-            return [self._dataset[i][0] for i in key]
-        if isinstance(key, slice):
-            indices = list(range(len(self._dataset))[key])
-            return [self._dataset[i][0] for i in indices]
-        elif isinstance(key, int):
-            return self._dataset[key][0]
-        raise TypeError("Invalid argument type.")
-
-    def __iter__(self) -> Iterator[TArray]:
+    def __getitem__(self, key: slice, /) -> Sequence[T]: ...
+
+    def __getitem__(self, key: int | slice, /) -> Sequence[T] | 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]
+        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)}")
+
+    def __iter__(self) -> Iterator[T]:
         for i in range(len(self._dataset)):
-            yield self._dataset[i][0]
+            yield self[i]
 
     def __len__(self) -> int:
         return len(self._dataset)

dataeval/utils/data/_metadata.py

@@ -3,18 +3,19 @@ from __future__ import annotations
 __all__ = []
 
 import warnings
-from typing import TYPE_CHECKING, Any, Literal, Mapping, Sequence
+from typing import TYPE_CHECKING, Any, Literal, Mapping, Sequence, cast
 
 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,
+from dataeval.typing import (
+    AnnotatedDataset,
+    Array,
+    ArrayLike,
     ObjectDetectionTarget,
 )
+from dataeval.utils._array import as_numpy, to_numpy
+from dataeval.utils._bin import bin_data, digitize_data, is_continuous
 from dataeval.utils.metadata import merge
 
 if TYPE_CHECKING:
@@ -65,7 +66,7 @@ class Metadata:
 
     def __init__(
         self,
-        dataset: Dataset[Any, Any],
+        dataset: AnnotatedDataset[tuple[Any, Any, dict[str, Any]]],
         *,
         continuous_factor_bins: Mapping[str, int | Sequence[float]] | None = None,
         auto_bin_method: Literal["uniform_width", "uniform_count", "clusters"] = "uniform_width",
@@ -276,12 +277,12 @@ class Metadata:
         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
 
+        # Validate the metadata dimensions
+        self._validate()
+
         # Include specified metadata keys
         if self.include:
             metadata = {i: self.merged[i] for i in self.include if i in self.merged}
@@ -341,7 +342,11 @@ class Metadata:
 
         # 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._discrete_data = (
+            np.stack(list(discrete_metadata.values()), axis=-1, dtype=np.int64)
+            if discrete_metadata
+            else np.array([], 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)
@@ -350,3 +355,15 @@ class Metadata:
         )
         self._total_num_factors = len(self._discrete_factor_names + self._continuous_factor_names) + 1
         self._processed = True
+
+    def add_factors(self, factors: Mapping[str, ArrayLike]) -> None:
+        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()):
+            raise ValueError(
+                "The lists/arrays in the provided factors have a different length than the current metadata factors."
+            )
+        merged = cast(tuple[dict[str, ArrayLike], dict[str, list[str]]], self._merged)[0]
+        for k, v in factors.items():
+            merged[k] = v

dataeval/utils/data/_selection.py

@@ -3,12 +3,11 @@ from __future__ import annotations
 __all__ = []
 
 from enum import IntEnum
-from typing import Any, Generic, Iterator, Sequence, TypeVar
+from typing import Generic, Iterator, Sequence, TypeVar
 
-from dataeval.utils.data._types import Dataset
+from dataeval.typing import AnnotatedDataset, DatasetMetadata
 
-_TData = TypeVar("_TData")
-_TTarget = TypeVar("_TTarget")
+_TDatum = TypeVar("_TDatum", covariant=True)
 
 
 class SelectionStage(IntEnum):
@@ -17,16 +16,16 @@ class SelectionStage(IntEnum):
     ORDER = 2
 
 
-class Selection(Generic[_TData, _TTarget]):
+class Selection(Generic[_TDatum]):
     stage: SelectionStage
 
-    def __call__(self, dataset: Select[_TData, _TTarget]) -> None: ...
+    def __call__(self, dataset: Select[_TDatum]) -> 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]):
+class Select(AnnotatedDataset[_TDatum]):
     """
     Wraps a dataset and applies selection criteria to it.
 
@@ -60,35 +59,43 @@ class Select(Generic[_TData, _TTarget], Dataset[_TData, _TTarget]):
     (data_20, 0, {'id': 20})
     """
 
-    _dataset: Dataset[_TData, _TTarget]
+    _dataset: AnnotatedDataset[_TDatum]
     _selection: list[int]
-    _selections: Sequence[Selection[_TData, _TTarget]]
+    _selections: Sequence[Selection[_TDatum]]
     _size_limit: int
 
     def __init__(
         self,
-        dataset: Dataset[_TData, _TTarget],
-        selections: Selection[_TData, _TTarget] | list[Selection[_TData, _TTarget]] | None = None,
+        dataset: AnnotatedDataset[_TDatum],
+        selections: Selection[_TDatum] | list[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.__dict__.update(dataset.__dict__)
+
+        # Ensure metadata is populated correctly as DatasetMetadata TypedDict
+        _metadata = getattr(dataset, "metadata", {})
+        if "id" not in _metadata:
+            _metadata["id"] = dataset.__class__.__name__
+        self._metadata = DatasetMetadata(**_metadata)
 
         if self._selections:
             self._apply_selections()
 
+    @property
+    def metadata(self) -> DatasetMetadata:
+        return self._metadata
+
     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}"
+        return f"{title}\n{sep}{nt}{selections}{nt}Selected Size: {len(self)}\n\n{self._dataset}"
 
-    def _sort_selections(
-        self, selections: Selection[_TData, _TTarget] | Sequence[Selection[_TData, _TTarget]] | None
-    ) -> list[Selection]:
+    def _sort_selections(self, selections: Selection[_TDatum] | Sequence[Selection[_TDatum]] | None) -> list[Selection]:
         if not selections:
             return []
 
@@ -104,14 +111,10 @@ class Select(Generic[_TData, _TTarget], Dataset[_TData, _TTarget]):
             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]]:
+    def __getitem__(self, index: int) -> _TDatum:
         return self._dataset[self._selection[index]]
 
-    def __iter__(self) -> Iterator[tuple[_TData, _TTarget, dict[str, Any]]]:
+    def __iter__(self) -> Iterator[_TDatum]:
         for i in range(len(self)):
             yield self[i]
 

dataeval/utils/data/_split.py

@@ -3,7 +3,6 @@ from __future__ import annotations
 __all__ = []
 
 import warnings
-from dataclasses import dataclass
 from typing import Any, Iterator, Protocol
 
 import numpy as np
@@ -13,32 +12,9 @@ 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
-
-
-@dataclass
-class TrainValSplit:
-    """Tuple containing train and validation indices"""
-
-    train: NDArray[np.intp]
-    val: NDArray[np.intp]
-
-
-@dataclass(frozen=True)
-class SplitDatasetOutput(Output):
-    """
-    Output class containing test indices and a list of TrainValSplits.
-
-    Attributes
-    ----------
-    test: NDArray[np.intp]
-        Indices for the test set
-    folds: list[TrainValSplit]
-        List where each index contains the indices for the train and validation splits
-    """
-
-    test: NDArray[np.intp]
-    folds: list[TrainValSplit]
+from dataeval.config import get_seed
+from dataeval.outputs._base import set_metadata
+from dataeval.outputs._utils import SplitDatasetOutput, TrainValSplit
 
 
 class KFoldSplitter(Protocol):
@@ -237,9 +213,9 @@ def bin_kmeans(array: NDArray[Any]) -> NDArray[np.intp]:
         best_score = 0.50
     bin_index = np.zeros(len(array), dtype=np.intp)
     for k in range(2, 20):
-        clusterer = KMeans(n_clusters=k)
+        clusterer = KMeans(n_clusters=k, random_state=get_seed())
         cluster_labels = clusterer.fit_predict(array)
-        score = silhouette_score(array, cluster_labels, sample_size=25_000)
+        score = silhouette_score(array, cluster_labels, sample_size=25_000, random_state=get_seed())
         if score > best_score:
             best_score = score
             bin_index = cluster_labels.astype(np.intp)

dataeval/utils/data/_targets.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+from typing import Iterator
+
 __all__ = []
 
 from dataclasses import dataclass
@@ -52,10 +54,16 @@ class Targets:
                 + f"    source: {None if self.source is None else self.source.shape}\n"
             )
 
+        if self.bboxes is not None and len(self.bboxes) > 0 and self.bboxes.shape[-1] != 4:
+            raise ValueError("Bounding boxes must be in (x0,y0,x1,y1) format.")
+
     def __len__(self) -> int:
-        return len(self.labels)
+        if self.source is None:
+            return len(self.labels)
+        else:
+            return len(np.unique(self.source))
 
-    def at(self, idx: int) -> Targets:
+    def __getitem__(self, idx: int, /) -> Targets:
         if self.source is None or self.bboxes is None:
             return Targets(
                 np.atleast_1d(self.labels[idx]),
@@ -71,3 +79,7 @@ class Targets:
                 np.atleast_2d(self.bboxes[mask]),
                 np.atleast_1d(self.source[mask]),
             )
+
+    def __iter__(self) -> Iterator[Targets]:
+        for i in range(len(self.labels)) if self.source is None else np.unique(self.source):
+            yield self[i]

dataeval/utils/data/datasets/_base.py

@@ -6,8 +6,10 @@ 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,
+from dataeval.utils.data.datasets._fileio import _ensure_exists
+from dataeval.utils.data.datasets._mixin import BaseDatasetMixin
+from dataeval.utils.data.datasets._types import (
+    AnnotatedDataset,
     DatasetMetadata,
     ImageClassificationDataset,
     ObjectDetectionDataset,
@@ -16,8 +18,6 @@ from dataeval.utils.data._types import (
     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")
@@ -31,7 +31,7 @@ class DataLocation(NamedTuple):
     checksum: str
 
 
-class BaseDataset(Dataset[_TArray, _TTarget], Generic[_TArray, _TTarget, _TRawTarget]):
+class BaseDataset(AnnotatedDataset[tuple[_TArray, _TTarget, dict[str, Any]]], Generic[_TArray, _TTarget, _TRawTarget]):
     """
     Base class for internet downloaded datasets.
     """