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

dataeval/utils/{plot.py → _plot.py}

@@ -6,9 +6,9 @@ import contextlib
 from typing import Any
 
 import numpy as np
-from numpy.typing import ArrayLike
 
-from dataeval.interop import to_numpy
+from dataeval.typing import ArrayLike
+from dataeval.utils._array import to_numpy
 
 with contextlib.suppress(ImportError):
     from matplotlib.figure import Figure
@@ -171,7 +171,7 @@ def histogram_plot(
         data_dict,
     ):
         # Plot the histogram for the chosen metric
-        ax.hist(data_dict[metric], bins=20, log=log)
+        ax.hist(data_dict[metric].astype(np.float64), bins=20, log=log)
 
         # Add labels to the histogram
         ax.set_title(metric)
@@ -229,7 +229,7 @@ def channel_histogram_plot(
         # Plot the histogram for the chosen metric
         data = data_dict[metric][ch_mask].reshape(-1, max_channels)
         ax.hist(
-            data,
+            data.astype(np.float64),
             bins=20,
             density=True,
             log=log,

dataeval/utils/data/__init__.py

@@ -0,0 +1,22 @@
+"""Provides utility functions for interacting with Computer Vision datasets."""
+
+__all__ = [
+    "collate",
+    "datasets",
+    "Embeddings",
+    "Images",
+    "Metadata",
+    "Select",
+    "SplitDatasetOutput",
+    "Targets",
+    "split_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 SplitDatasetOutput, split_dataset
+from dataeval.utils.data._targets import Targets
+
+from . import collate, datasets

dataeval/utils/data/_embeddings.py

@@ -0,0 +1,105 @@
+from __future__ import annotations
+
+__all__ = []
+
+import math
+from typing import Any, Iterator, Sequence
+
+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.utils.torch.models import SupportsEncode
+
+
+class Embeddings:
+    """
+    Collection of image embeddings from a dataset.
+
+    Embeddings are accessed by index or slice and are only loaded on-demand.
+
+    Parameters
+    ----------
+    dataset : ImageClassificationDataset or ObjectDetectionDataset
+        Dataset to access original images from.
+    batch_size : int, optional
+        Batch size to use when encoding images.
+    model : torch.nn.Module, optional
+        Model to use for encoding images.
+    device : torch.device, optional
+        Device to use for encoding images.
+    verbose : bool, optional
+        Whether to print progress bar when encoding images.
+    """
+
+    device: torch.device
+    batch_size: int
+    verbose: bool
+
+    def __init__(
+        self,
+        dataset: Dataset[TArray, Any],
+        batch_size: int,
+        indices: Sequence[int] | None = None,
+        model: torch.nn.Module | None = None,
+        device: torch.device | str | None = None,
+        verbose: bool = False,
+    ) -> None:
+        self.device = get_device(device)
+        self.batch_size = batch_size
+        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
+        self._collate_fn = lambda datum: [torch.as_tensor(i) for i, _, _ in datum]
+
+    def to_tensor(self) -> torch.Tensor:
+        """
+        Converts entire dataset to embeddings.
+
+        Warning
+        -------
+        Will process the entire dataset in batches and return
+        embeddings as a single Tensor in memory.
+
+        Returns
+        -------
+        torch.Tensor
+        """
+        return self[:]
+
+    # Reduce overhead cost by not tracking tensor gradients
+    @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)
+        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.")
+
+    def __iter__(self) -> Iterator[torch.Tensor]:
+        # process in batches while yielding individual embeddings
+        for batch in self._batch(range(len(self._dataset))):
+            yield from batch
+
+    def __len__(self) -> int:
+        return len(self._dataset)

dataeval/utils/data/_images.py

@@ -0,0 +1,65 @@
+from __future__ import annotations
+
+__all__ = []
+
+from typing import Any, Generic, Iterator, Sequence, overload
+
+from dataeval.typing import TArray
+from dataeval.utils.data._types import Dataset
+
+
+class Images(Generic[TArray]):
+    """
+    Collection of image data from a dataset.
+
+    Images are accessed by index or slice and are only loaded on-demand.
+
+    Parameters
+    ----------
+    dataset : ImageClassificationDataset or ObjectDetectionDataset
+        Dataset to access images from.
+    """
+
+    def __init__(
+        self,
+        dataset: Dataset[TArray, Any],
+    ) -> None:
+        self._dataset = dataset
+
+    def to_list(self) -> Sequence[TArray]:
+        """
+        Converts entire dataset to a sequence of images.
+
+        Warning
+        -------
+        Will load the entire dataset and return the images as a
+        single sequence of images in memory.
+
+        Returns
+        -------
+        list[TArray]
+        """
+        return self[:]
+
+    @overload
+    def __getitem__(self, key: slice | list[int]) -> Sequence[TArray]: ...
+
+    @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]:
+        for i in range(len(self._dataset)):
+            yield self._dataset[i][0]
+
+    def __len__(self) -> int:
+        return len(self._dataset)

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)