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

dataeval/utils/metadata.py

@@ -1,12 +1,26 @@
+"""
+Metadata related utility functions that help organize raw metadata into \
+:class:`Metadata` objects for use within `DataEval`.
+"""
+
 from __future__ import annotations
 
-__all__ = ["merge_metadata"]
+__all__ = ["Metadata", "preprocess", "merge", "flatten"]
 
 import warnings
-from typing import Any, Iterable, Mapping, TypeVar, overload
+from dataclasses import dataclass
+from typing import Any, Iterable, Literal, Mapping, TypeVar, overload
 
 import numpy as np
-from numpy.typing import NDArray
+from numpy.typing import ArrayLike, NDArray
+from scipy.stats import wasserstein_distance as wd
+
+from dataeval.interop import as_numpy, to_numpy
+from dataeval.output import Output, set_metadata
+
+DISCRETE_MIN_WD = 0.054
+CONTINUOUS_MIN_SAMPLE_SIZE = 20
+
 
 T = TypeVar("T")
 
@@ -131,9 +145,7 @@ def _flatten_dict_inner(
     return items, size
 
 
-def _flatten_dict(
-    d: Mapping[str, Any], sep: str, ignore_lists: bool, fully_qualified: bool
-) -> tuple[dict[str, Any], int]:
+def flatten(d: Mapping[str, Any], sep: str, ignore_lists: bool, fully_qualified: bool) -> tuple[dict[str, Any], int]:
     """
     Flattens a dictionary and converts values to numeric values when possible.
 
@@ -146,12 +158,12 @@ def _flatten_dict(
     ignore_lists : bool
         Option to skip expanding lists within metadata
     fully_qualified : bool
-        Option to return dictionary keys full qualified instead of minimized
+        Option to return dictionary keys full qualified instead of reduced
 
     Returns
     -------
-    dict[str, Any]
-        A flattened dictionary
+    tuple[dict[str, Any], int]
+        A tuple of the flattened dictionary and the length of detected lists in metadata
     """
     expanded, size = _flatten_dict_inner(d, parent_keys=(), nested=ignore_lists)
 
@@ -185,7 +197,7 @@ def _is_metadata_dict_of_dicts(metadata: Mapping) -> bool:
     return set(metadata[keys[0]]) == set(metadata[keys[1]])
 
 
-def merge_metadata(
+def merge(
     metadata: Iterable[Mapping[str, Any]],
     ignore_lists: bool = False,
     fully_qualified: bool = False,
@@ -222,7 +234,7 @@ def merge_metadata(
     Example
     -------
     >>> list_metadata = [{"common": 1, "target": [{"a": 1, "b": 3, "c": 5}, {"a": 2, "b": 4}], "source": "example"}]
-    >>> reorganized_metadata, image_indicies = merge_metadata(list_metadata)
+    >>> reorganized_metadata, image_indicies = merge(list_metadata)
     >>> reorganized_metadata
     {'common': [1, 1], 'a': [1, 2], 'b': [3, 4], 'source': ['example', 'example']}
     >>> image_indicies
@@ -245,9 +257,7 @@ def merge_metadata(
 
     image_repeats = np.zeros(len(dicts))
     for i, d in enumerate(dicts):
-        flattened, image_repeats[i] = _flatten_dict(
-            d, sep="_", ignore_lists=ignore_lists, fully_qualified=fully_qualified
-        )
+        flattened, image_repeats[i] = flatten(d, sep="_", ignore_lists=ignore_lists, fully_qualified=fully_qualified)
         isect = isect.intersection(flattened.keys()) if isect else set(flattened.keys())
         union = union.union(flattened.keys())
         for k, v in flattened.items():
@@ -276,3 +286,296 @@ def merge_metadata(
         output[k] = np.array(cv) if as_numpy else cv
 
     return output, image_indicies
+
+
+@dataclass(frozen=True)
+class Metadata(Output):
+    """
+    Dataclass containing binned metadata from the :func:`preprocess` function.
+
+    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.int]
+        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.int or np.double] | 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 : NDArray[Any]
+        Array of unique class names (for use with plotting)
+    total_num_factors : int
+        Sum of discrete_factor_names and continuous_factor_names plus 1 for class
+    """
+
+    discrete_factor_names: list[str]
+    discrete_data: NDArray[np.int_]
+    continuous_factor_names: list[str]
+    continuous_data: NDArray[np.int_ | np.double] | None
+    class_labels: NDArray[np.int_]
+    class_names: NDArray[Any]
+    total_num_factors: int
+
+
+@set_metadata
+def preprocess(
+    raw_metadata: Iterable[Mapping[str, Any]],
+    class_labels: ArrayLike | str,
+    continuous_factor_bins: Mapping[str, int | Iterable[float]] | None = None,
+    auto_bin_method: Literal["uniform_width", "uniform_count", "clusters"] = "uniform_width",
+    exclude: Iterable[str] | None = None,
+) -> Metadata:
+    """
+    Restructures the metadata to be in the correct format for the bias functions.
+
+    This identifies whether the incoming metadata is discrete or continuous,
+    and whether the data is already binned or still needs binning.
+    It accepts a list of dictionaries containing the per image metadata and
+    automatically adjusts for multiple targets in an image.
+
+    Parameters
+    ----------
+    raw_metadata : Iterable[Mapping[str, Any]]
+        Iterable collection of metadata dictionaries to flatten and merge.
+    class_labels : ArrayLike or string
+        If arraylike, expects the labels for each image (image classification) or each object (object detection).
+        If the labels are included in the metadata dictionary, pass in the key value.
+    continuous_factor_bins : Mapping[str, int or Iterable[float]] or None, default None
+        User provided dictionary specifying how to bin the continuous metadata factors where the value is either
+        an int to represent the number of bins, or a list of floats representing the edges for each bin.
+    auto_bin_method : "uniform_width" or "uniform_count" or "clusters", default "uniform_width"
+        Method by which the function will automatically bin continuous metadata factors. It is recommended
+        that the user provide the bins through the `continuous_factor_bins`.
+    exclude : Iterable[str] or None, default None
+        User provided collection of metadata keys to exclude when processing metadata.
+
+    Returns
+    -------
+    Metadata
+        Output class containing the binned metadata
+    """
+    # Transform metadata into single, flattened dictionary
+    metadata, image_repeats = merge(raw_metadata)
+
+    continuous_factor_bins = dict(continuous_factor_bins) if continuous_factor_bins else None
+
+    # Drop any excluded metadata keys
+    for k in exclude or ():
+        metadata.pop(k, None)
+        if continuous_factor_bins:
+            continuous_factor_bins.pop(k, None)
+
+    # Get the class label array in numeric form
+    class_array = as_numpy(metadata.pop(class_labels)) if isinstance(class_labels, str) else as_numpy(class_labels)
+    if class_array.ndim > 1:
+        raise ValueError(
+            f"Got class labels with {class_array.ndim}-dimensional "
+            f"shape {class_array.shape}, but expected a 1-dimensional array."
+        )
+    if not np.issubdtype(class_array.dtype, np.int_):
+        unique_classes, numerical_labels = np.unique(class_array, return_inverse=True)
+    else:
+        numerical_labels = class_array
+        unique_classes = np.unique(class_array)
+
+    # Bin according to user supplied bins
+    continuous_metadata = {}
+    discrete_metadata = {}
+    if continuous_factor_bins is not None and 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, image_repeats)
+            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 {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, auto_bin_method)
+        else:
+            _, discrete_metadata[key] = np.unique(data, return_inverse=True)
+
+    # splitting out the dictionaries into the keys and values
+    discrete_factor_names = list(discrete_metadata.keys())
+    discrete_data = np.stack(list(discrete_metadata.values()), axis=-1)
+    continuous_factor_names = list(continuous_metadata.keys())
+    continuous_data = np.stack(list(continuous_metadata.values()), axis=-1) if continuous_metadata else None
+    total_num_factors = len(discrete_factor_names + continuous_factor_names) + 1
+
+    return Metadata(
+        discrete_factor_names,
+        discrete_data,
+        continuous_factor_names,
+        continuous_data,
+        numerical_labels,
+        unique_classes,
+        total_num_factors,
+    )
+
+
+def _digitize_data(data: list[Any] | NDArray[Any], bins: int | Iterable[float]) -> NDArray[np.intp]:
+    """
+    Digitizes a list of values into a given number of bins.
+
+    Parameters
+    ----------
+    data : list | NDArray
+        The values to be digitized.
+    bins : int | Iterable[float]
+        The number of bins or list of bin edges for the discrete values that data will be digitized into.
+
+    Returns
+    -------
+    NDArray[np.intp]
+        The digitized values
+    """
+
+    if not np.all([np.issubdtype(type(n), np.number) for n in data]):
+        raise TypeError(
+            "Encountered a data value with non-numeric type when digitizing a factor. "
+            "Ensure all occurrences of continuous factors are numeric types."
+        )
+    if isinstance(bins, int):
+        _, bin_edges = np.histogram(data, bins=bins)
+        bin_edges[-1] = np.inf
+        bin_edges[0] = -np.inf
+    else:
+        bin_edges = list(bins)
+    return np.digitize(data, bin_edges)
+
+
+def _bin_data(data: NDArray[Any], bin_method: str) -> NDArray[np.int_]:
+    """
+    Bins continuous data through either equal width bins, equal amounts in each bin, or by clusters.
+    """
+    if bin_method == "clusters":
+        # bin_edges = _binning_by_clusters(data)
+        warnings.warn(
+            "Binning by clusters is currently unavailable until changes to the clustering function go through.",
+            UserWarning,
+        )
+        bin_method = "uniform_width"
+
+    # if bin_method != "clusters":  # restore this when clusters bin_method is available
+    counts, bin_edges = np.histogram(data, bins="auto")
+    n_bins = counts.size
+    if counts[counts > 0].min() < 10:
+        counter = 20
+        while counts[counts > 0].min() < 10 and n_bins >= 2 and counter > 0:
+            counter -= 1
+            n_bins -= 1
+            counts, bin_edges = np.histogram(data, bins=n_bins)
+
+    if bin_method == "uniform_count":
+        quantiles = np.linspace(0, 100, n_bins + 1)
+        bin_edges = np.asarray(np.percentile(data, quantiles))
+
+    bin_edges[0] = -np.inf  # type: ignore # until the clusters speed up is merged
+    bin_edges[-1] = np.inf  # type: ignore # and the _binning_by_clusters can be uncommented
+    return np.digitize(data, bin_edges)  # type: ignore
+
+
+def _is_continuous(data: NDArray[np.number], image_indicies: NDArray[np.number]) -> bool:
+    """
+    Determines whether the data is continuous or discrete using the Wasserstein distance.
+
+    Given a 1D sample, we consider the intervals between adjacent points. For a continuous distribution,
+    a point is equally likely to lie anywhere in the interval bounded by its two neighbors. Furthermore,
+    we can put all "between neighbor" locations on the same scale of 0 to 1 by subtracting the smaller
+    neighbor and dividing out the length of the interval. (Duplicates are either assigned to zero or
+    ignored, depending on context). These normalized locations will be much more uniformly distributed
+    for continuous data than for discrete, and this gives us a way to distinguish them. Call this the
+    Normalized Near Neighbor distribution (NNN), defined on the interval [0,1].
+
+    The Wasserstein distance is available in scipy.stats.wasserstein_distance. We can use it to measure
+    how close the NNN is to a uniform distribution over [0,1]. We found that as long as a sample has at
+    least 20 points, and furthermore at least half as many points as there are discrete values, we can
+    reliably distinguish discrete from continuous samples by testing that the Wasserstein distance
+    measured from a uniform distribution is greater or less than 0.054, respectively.
+    """
+    # Check if the metadata is image specific
+    _, data_indicies_unsorted = np.unique(data, return_index=True)
+    if data_indicies_unsorted.size == image_indicies.size:
+        data_indicies = np.sort(data_indicies_unsorted)
+        if (data_indicies == image_indicies).all():
+            data = data[data_indicies]
+
+    # OLD METHOD
+    # uvals = np.unique(data)
+    # pct_unique = uvals.size / data.size
+    # return pct_unique < threshold
+
+    n_examples = len(data)
+
+    if n_examples < CONTINUOUS_MIN_SAMPLE_SIZE:
+        warnings.warn(
+            f"All samples look discrete with so few data points (< {CONTINUOUS_MIN_SAMPLE_SIZE})", UserWarning
+        )
+        return False
+
+    # Require at least 3 unique values before bothering with NNN
+    xu = np.unique(data, axis=None)
+    if xu.size < 3:
+        return False
+
+    Xs = np.sort(data)
+
+    X0, X1 = Xs[0:-2], Xs[2:]  # left and right neighbors
+
+    dx = np.zeros(n_examples - 2)  # no dx at end points
+    gtz = (X1 - X0) > 0  # check for dups; dx will be zero for them
+    dx[np.logical_not(gtz)] = 0.0
+
+    dx[gtz] = (Xs[1:-1] - X0)[gtz] / (X1 - X0)[gtz]  # the core idea: dx is NNN samples.
+
+    shift = wd(dx, np.linspace(0, 1, dx.size))  # how far is dx from uniform, for this feature?
+
+    return shift < DISCRETE_MIN_WD  # if NNN is close enough to uniform, consider the sample continuous.
+
+
+def get_counts(data: NDArray[np.int_], min_num_bins: int | None = None) -> NDArray[np.int_]:
+    """
+    Returns columnwise unique counts for discrete data.
+
+    Parameters
+    ----------
+    data : NDArray
+        Array containing integer values for metadata factors
+    min_num_bins : int | None, default None
+        Minimum number of bins for bincount, helps force consistency across runs
+
+    Returns
+    -------
+    NDArray[np.int_]
+        Bin counts per column of data.
+    """
+    max_value = data.max() + 1 if min_num_bins is None else min_num_bins
+    cnt_array = np.zeros((max_value, data.shape[1]), dtype=np.int_)
+    for idx in range(data.shape[1]):
+        cnt_array[:, idx] = np.bincount(data[:, idx], minlength=max_value)
+
+    return cnt_array

dataeval/{metrics/bias/metadata_utils.py → utils/plot.py}

@@ -6,7 +6,7 @@ import contextlib
 from typing import Any
 
 import numpy as np
-from numpy.typing import ArrayLike, NDArray
+from numpy.typing import ArrayLike
 
 from dataeval.interop import to_numpy
 
@@ -14,30 +14,6 @@ with contextlib.suppress(ImportError):
     from matplotlib.figure import Figure
 
 
-def get_counts(data: NDArray[np.int_], min_num_bins: int | None = None) -> NDArray[np.int_]:
-    """
-    Returns columnwise unique counts for discrete data.
-
-    Parameters
-    ----------
-    data : NDArray
-        Array containing integer values for metadata factors
-    min_num_bins : int | None, default None
-        Minimum number of bins for bincount, helps force consistency across runs
-
-    Returns
-    -------
-    NDArray[np.int_]
-        Bin counts per column of data.
-    """
-    max_value = data.max() + 1 if min_num_bins is None else min_num_bins
-    cnt_array = np.zeros((max_value, data.shape[1]), dtype=np.int_)
-    for idx in range(data.shape[1]):
-        cnt_array[:, idx] = np.bincount(data[:, idx], minlength=max_value)
-
-    return cnt_array
-
-
 def heatmap(
     data: ArrayLike,
     row_labels: list[str] | ArrayLike,
@@ -95,12 +71,17 @@ def heatmap(
     # Rotate the tick labels and set their alignment.
     plt.setp(ax.get_xticklabels(), rotation=45, ha="right", rotation_mode="anchor")
 
-    # Turn spines off and create white grid.
-    ax.spines[:].set_visible(False)
-
-    ax.set_xticks(np.arange(np_data.shape[1] + 1) - 0.5, minor=True)
-    ax.set_yticks(np.arange(np_data.shape[0] + 1) - 0.5, minor=True)
-    ax.grid(which="minor", color="w", linestyle="-", linewidth=3)
+    light_gray = "0.9"
+    # Turn spines on and create light gray easily visible grid.
+    for spine in ax.spines.values():
+        spine.set_visible(True)
+        spine.set_color(light_gray)
+
+    xticks = np.arange(np_data.shape[1] + 1) - 0.5
+    yticks = np.arange(np_data.shape[0] + 1) - 0.5
+    ax.set_xticks(xticks, minor=True)
+    ax.set_yticks(yticks, minor=True)
+    ax.grid(which="minor", color=light_gray, linestyle="-", linewidth=3)
     ax.tick_params(which="minor", bottom=False, left=False)
 
     if xlabel:
@@ -151,79 +132,118 @@ def format_text(*args: str) -> str:
     return f"{x:.2f}".replace("0.00", "0").replace("0.", ".").replace("nan", "")
 
 
-def diversity_bar_plot(labels: NDArray[Any], bar_heights: NDArray[Any]) -> Figure:
+def histogram_plot(
+    data_dict: dict[str, Any],
+    log: bool = True,
+    xlabel: str = "values",
+    ylabel: str = "counts",
+) -> Figure:
     """
-    Plots a formatted bar plot
+    Plots a formatted histogram
 
     Parameters
     ----------
-    labels : NDArray
-        Array containing the labels for each bar
-    bar_heights : NDArray
-        Array containing the values for each bar
+    data_dict : dict
+        Dictionary containing the metrics and their value arrays
+    log : bool, default True
+        If True, plots the histogram on a semi-log scale (y axis)
+    xlabel : str, default "values"
+        X-axis label
+    ylabel : str, default "counts"
+        Y-axis label
 
     Returns
     -------
     matplotlib.figure.Figure
-        Bar plot figure
+        Formatted plot of histograms
     """
     import matplotlib.pyplot as plt
 
-    fig, ax = plt.subplots(figsize=(10, 10))
+    num_metrics = len(data_dict)
+    if num_metrics > 2:
+        rows = int(len(data_dict) / 3)
+        fig, axs = plt.subplots(rows, 3, figsize=(10, rows * 2.5))
+    else:
+        fig, axs = plt.subplots(1, num_metrics, figsize=(4 * num_metrics, 4))
 
-    ax.bar(labels, bar_heights)
-    ax.set_xlabel("Factors")
+    for ax, metric in zip(
+        axs.flat,
+        data_dict,
+    ):
+        # Plot the histogram for the chosen metric
+        ax.hist(data_dict[metric], bins=20, log=log)
 
-    plt.setp(ax.get_xticklabels(), rotation=45, ha="right", rotation_mode="anchor")
+        # Add labels to the histogram
+        ax.set_title(metric)
+        ax.set_ylabel(ylabel)
+        ax.set_xlabel(xlabel)
 
     fig.tight_layout()
     return fig
 
 
-def coverage_plot(images: NDArray[Any], num_images: int) -> Figure:
+def channel_histogram_plot(
+    data_dict: dict[str, Any],
+    log: bool = True,
+    max_channels: int = 3,
+    ch_mask: list[bool] | None = None,
+    xlabel: str = "values",
+    ylabel: str = "counts",
+) -> Figure:
     """
-    Creates a single plot of all of the provided images
+    Plots a formatted heatmap
 
     Parameters
     ----------
-    images : NDArray
-        Array containing only the desired images to plot
+    data_dict : dict
+        Dictionary containing the metrics and their value arrays
+    log : bool, default True
+        If True, plots the histogram on a semi-log scale (y axis)
+    xlabel : str, default "values"
+        X-axis label
+    ylabel : str, default "counts"
+        Y-axis label
 
     Returns
     -------
     matplotlib.figure.Figure
-        Plot of all provided images
+        Formatted plot of histograms
     """
     import matplotlib.pyplot as plt
 
-    num_images = min(num_images, len(images))
+    channelwise_metrics = ["mean", "std", "var", "skew", "zeros", "brightness", "contrast", "darkness", "entropy"]
+    data_keys = [key for key in data_dict if key in channelwise_metrics]
+    label_kwargs = {"label": [f"Channel {i}" for i in range(max_channels)]}
 
-    if images.ndim == 4:
-        images = np.moveaxis(images, 1, -1)
-    elif images.ndim == 3:
-        images = np.repeat(images[:, :, :, np.newaxis], 3, axis=-1)
+    num_metrics = len(data_keys)
+    if num_metrics > 2:
+        rows = int(len(data_keys) / 3)
+        fig, axs = plt.subplots(rows, 3, figsize=(10, rows * 2.5))
     else:
-        raise ValueError(
-            f"Expected a (N,C,H,W) or a (N, H, W) set of images, but got a {images.ndim}-dimensional set of images."
+        fig, axs = plt.subplots(1, num_metrics, figsize=(4 * num_metrics, 4))
+
+    for ax, metric in zip(
+        axs.flat,
+        data_keys,
+    ):
+        # Plot the histogram for the chosen metric
+        data = data_dict[metric][ch_mask].reshape(-1, max_channels)
+        ax.hist(
+            data,
+            bins=20,
+            density=True,
+            log=log,
+            **label_kwargs,
         )
+        # Only plot the labels once for channels
+        if label_kwargs:
+            ax.legend()
+            label_kwargs = {}
 
-    rows = int(np.ceil(num_images / 3))
-    fig, axs = plt.subplots(rows, 3, figsize=(9, 3 * rows))
-
-    if rows == 1:
-        for j in range(3):
-            if j >= len(images):
-                continue
-            axs[j].imshow(images[j])
-            axs[j].axis("off")
-    else:
-        for i in range(rows):
-            for j in range(3):
-                i_j = i * 3 + j
-                if i_j >= len(images):
-                    continue
-                axs[i, j].imshow(images[i_j])
-                axs[i, j].axis("off")
+        # Add labels to the histogram
+        ax.set_title(metric)
+        ax.set_ylabel(ylabel)
+        ax.set_xlabel(xlabel)
 
     fig.tight_layout()
     return fig

dataeval/utils/torch/__init__.py

@@ -5,21 +5,6 @@ While these metrics can take in custom models, DataEval provides utility classes
 to create a seamless integration between custom models and DataEval's metrics.
 """
 
-from dataeval import _IS_TORCH_AVAILABLE, _IS_TORCHVISION_AVAILABLE
+__all__ = ["models", "trainer"]
 
-__all__ = []
-
-if _IS_TORCH_AVAILABLE:
-    from dataeval.utils.torch import models, trainer
-    from dataeval.utils.torch.utils import read_dataset
-
-    __all__ += ["read_dataset", "models", "trainer"]
-
-if _IS_TORCHVISION_AVAILABLE:
-    from dataeval.utils.torch import datasets
-
-    __all__ += ["datasets"]
-
-
-del _IS_TORCH_AVAILABLE
-del _IS_TORCHVISION_AVAILABLE
+from dataeval.utils.torch import models, trainer

dataeval/utils/torch/gmm.py

@@ -1,7 +1,5 @@
 """
-Adapted for Pytorch from:
-
-Source code derived from Alibi-Detect 0.11.4
+Adapted for Pytorch from Alibi-Detect 0.11.4
 https://github.com/SeldonIO/alibi-detect/tree/v0.11.4
 
 Original code Copyright (c) 2023 Seldon Technologies Ltd
@@ -10,13 +8,38 @@ Licensed under Apache Software License (Apache 2.0)
 
 from __future__ import annotations
 
+__all__ = []
+
+from typing import NamedTuple, TypeVar
+
 import numpy as np
 import torch
 
-from dataeval.utils.gmm import GaussianMixtureModelParams
+TGMMData = TypeVar("TGMMData")
+
+
+class GaussianMixtureModelParams(NamedTuple):
+    """
+    phi : torch.Tensor
+        Mixture component distribution weights.
+    mu : torch.Tensor
+        Mixture means.
+    cov : torch.Tensor
+        Mixture covariance.
+    L : torch.Tensor
+        Cholesky decomposition of `cov`.
+    log_det_cov : torch.Tensor
+        Log of the determinant of `cov`.
+    """
+
+    phi: torch.Tensor
+    mu: torch.Tensor
+    cov: torch.Tensor
+    L: torch.Tensor
+    log_det_cov: torch.Tensor
 
 
-def gmm_params(z: torch.Tensor, gamma: torch.Tensor) -> GaussianMixtureModelParams[torch.Tensor]:
+def gmm_params(z: torch.Tensor, gamma: torch.Tensor) -> GaussianMixtureModelParams:
     """
     Compute parameters of Gaussian Mixture Model.
 
@@ -58,7 +81,7 @@ def gmm_params(z: torch.Tensor, gamma: torch.Tensor) -> GaussianMixtureModelPara
 
 def gmm_energy(
     z: torch.Tensor,
-    params: GaussianMixtureModelParams[torch.Tensor],
+    params: GaussianMixtureModelParams,
     return_mean: bool = True,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """