dataeval 0.69.4__py3-none-any.whl → 0.70.1__py3-none-any.whl

dataeval/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "0.69.4"
+__version__ = "0.70.1"
 
 from importlib.util import find_spec
 
@@ -7,16 +7,16 @@ _IS_TENSORFLOW_AVAILABLE = find_spec("tensorflow") is not None and find_spec("te
 
 del find_spec
 
-from . import detectors, flags, metrics  # noqa: E402
+from . import detectors, metrics  # noqa: E402
 
-__all__ = ["detectors", "flags", "metrics"]
+__all__ = ["detectors", "metrics"]
 
 if _IS_TORCH_AVAILABLE:  # pragma: no cover
-    from . import torch, utils, workflows
+    from . import workflows
 
-    __all__ += ["torch", "utils", "workflows"]
+    __all__ += ["workflows"]
 
-if _IS_TENSORFLOW_AVAILABLE:  # pragma: no cover
-    from . import tensorflow
+if _IS_TENSORFLOW_AVAILABLE or _IS_TORCH_AVAILABLE:  # pragma: no cover
+    from . import utils
 
-    __all__ += ["tensorflow"]
+    __all__ += ["utils"]

dataeval/_internal/datasets.py

@@ -4,18 +4,39 @@ import hashlib
 import os
 import zipfile
 from pathlib import Path
-from typing import Literal
-from urllib.error import HTTPError, URLError
-from urllib.request import urlretrieve
+from typing import Literal, TypeVar
+from warnings import warn
 
 import numpy as np
+import requests
 from numpy.typing import NDArray
 from torch.utils.data import Dataset
 from torchvision.datasets import CIFAR10, VOCDetection  # noqa: F401
 
-
-def _validate_file(fpath, file_md5, chunk_size=65535):
-    hasher = hashlib.md5()
+ClassStringMap = Literal["zero", "one", "two", "three", "four", "five", "six", "seven", "eight", "nine"]
+TClassMap = TypeVar("TClassMap", ClassStringMap, int, list[ClassStringMap], list[int])
+CorruptionStringMap = Literal[
+    "identity",
+    "shot_noise",
+    "impulse_noise",
+    "glass_blur",
+    "motion_blur",
+    "shear",
+    "scale",
+    "rotate",
+    "brightness",
+    "translate",
+    "stripe",
+    "fog",
+    "spatter",
+    "dotted_line",
+    "zigzag",
+    "canny_edges",
+]
+
+
+def _validate_file(fpath, file_md5, md5=False, chunk_size=65535):
+    hasher = hashlib.md5() if md5 else hashlib.sha256()
     with open(fpath, "rb") as fpath_file:
         while chunk := fpath_file.read(chunk_size):
             hasher.update(chunk)
@@ -26,44 +47,74 @@ def _get_file(
     root: str | Path,
     fname: str,
     origin: str,
-    file_md5: str | None = None,
+    file_hash: str | None = None,
+    verbose: bool = True,
+    md5: bool = False,
 ):
-    fname = os.fspath(fname) if isinstance(fname, os.PathLike) else fname
     fpath = os.path.join(root, fname)
-
-    download = False
-    if os.path.exists(fpath):
-        if file_md5 is not None and not _validate_file(fpath, file_md5):
-            download = True
-        else:
-            print("Files already downloaded and verified")
-    else:
-        download = True
+    download = True
+    if os.path.exists(fpath) and file_hash is not None and _validate_file(fpath, file_hash, md5):
+        download = False
+        if verbose:
+            print("File already downloaded and verified.")
+            if md5:
+                print("Extracting zip file...")
 
     if download:
         try:
             error_msg = "URL fetch failure on {}: {} -- {}"
             try:
-                urlretrieve(origin, fpath)
-            except HTTPError as e:
-                raise Exception(error_msg.format(origin, e.code, e.msg)) from e
-            except URLError as e:
-                raise Exception(error_msg.format(origin, e.errno, e.reason)) from e
+                with requests.get(origin, stream=True, timeout=60) as r:
+                    r.raise_for_status()
+                    with open(fpath, "wb") as f:
+                        for chunk in r.iter_content(chunk_size=8192):
+                            if chunk:
+                                f.write(chunk)
+            except requests.exceptions.HTTPError as e:
+                raise Exception(f"{error_msg.format(origin, e.response.status_code)} -- {e.response.reason}") from e
+            except requests.exceptions.RequestException as e:
+                raise Exception(f"{error_msg.format(origin, 'Unknown error')} -- {str(e)}") from e
         except (Exception, KeyboardInterrupt):
             if os.path.exists(fpath):
                 os.remove(fpath)
             raise
 
-        if os.path.exists(fpath) and file_md5 is not None and not _validate_file(fpath, file_md5):
+        if os.path.exists(fpath) and file_hash is not None and not _validate_file(fpath, file_hash, md5):
             raise ValueError(
                 "Incomplete or corrupted file detected. "
-                f"The md5 file hash does not match the provided value "
-                f"of {file_md5}.",
+                f"The file hash does not match the provided value "
+                f"of {file_hash}.",
             )
+
     return fpath
 
 
-def download_dataset(url: str, root: str | Path, fname: str, md5: str) -> str:
+def check_exists(
+    folder: str | Path,
+    url: str,
+    root: str | Path,
+    fname: str,
+    file_hash: str,
+    download: bool = True,
+    verbose: bool = True,
+    md5: bool = False,
+):
+    """Determine if the dataset has already been downloaded."""
+    location = str(folder)
+    if not os.path.exists(folder):
+        if download:
+            location = download_dataset(url, root, fname, file_hash, verbose, md5)
+        else:
+            raise RuntimeError("Dataset not found. You can use download=True to download it")
+    else:
+        if verbose:
+            print("Files already downloaded and verified")
+    return location
+
+
+def download_dataset(
+    url: str, root: str | Path, fname: str, file_hash: str, verbose: bool = True, md5: bool = False
+) -> str:
     """Code to download mnist and corruptions, originates from tensorflow_datasets (tfds):
     https://github.com/tensorflow/datasets/blob/master/tensorflow_datasets/image_classification/mnist_corrupted.py
     """
@@ -71,21 +122,24 @@ def download_dataset(url: str, root: str | Path, fname: str, md5: str) -> str:
     folder = os.path.join(root, name)
     os.makedirs(folder, exist_ok=True)
 
-    path = _get_file(
-        root,
+    fpath = _get_file(
+        folder,
         fname,
         origin=url + fname,
-        file_md5=md5,
+        file_hash=file_hash,
+        verbose=verbose,
+        md5=md5,
     )
-    extract_archive(path, remove_finished=True)
-    return path
+    if md5:
+        folder = extract_archive(fpath, root, remove_finished=True)
+    return folder
 
 
 def extract_archive(
     from_path: str | Path,
     to_path: str | Path | None = None,
     remove_finished: bool = False,
-):
+) -> str:
     """Extract an archive.
 
     The archive type and a possible compression is automatically detected from the file name.
@@ -94,8 +148,11 @@ def extract_archive(
     if not from_path.is_absolute():
         from_path = from_path.resolve()
 
-    if to_path is None:
+    if to_path is None or not os.path.exists(to_path):
         to_path = os.path.dirname(from_path)
+    to_path = Path(to_path)
+    if not to_path.is_absolute():
+        to_path = to_path.resolve()
 
     # Extracting zip
     with zipfile.ZipFile(from_path, "r", compression=zipfile.ZIP_STORED) as zzip:
@@ -103,6 +160,13 @@ def extract_archive(
 
     if remove_finished:
         os.remove(from_path)
+    return str(to_path)
+
+
+def subselect(arr: NDArray, count: int, from_back: bool = False):
+    if from_back:
+        return arr[-count:]
+    return arr[:count]
 
 
 class MNIST(Dataset):
@@ -133,40 +197,42 @@ class MNIST(Dataset):
             'motion_blur' | 'shear' | 'scale' | 'rotate' | 'brightness' | 'translate' | 'stripe' |
             'fog' | 'spatter' | 'dotted_line' | 'zigzag' | 'canny_edges'] | None, default None
             The desired corruption style or None.
+        classes : Literal["zero", "one", "two", "three", "four", "five", "six", "seven", "eight", "nine"]
+            | int | list[int] | list[Literal["zero", "one", "two", "three", "four", "five", "six", "seven",
+            "eight", "nine"]] | None, default None
+            Option to select specific classes from dataset.
+        balance : bool, default True
+            If True, returns equal number of samples for each class.
+        randomize : bool, default False
+            If True, shuffles the data prior to selection - uses a set seed for reproducibility.
+        slice_back : bool, default False
+            If True and size has a value greater than 0, then grabs selection starting at the last image.
+        verbose : bool, default True
+            If True, outputs print statements.
     """
 
-    mirror = "https://zenodo.org/record/3239543/files/"
-
-    resources = ("mnist_c.zip", "4b34b33045869ee6d424616cd3a65da3")
-
-    classes = [
-        "0 - zero",
-        "1 - one",
-        "2 - two",
-        "3 - three",
-        "4 - four",
-        "5 - five",
-        "6 - six",
-        "7 - seven",
-        "8 - eight",
-        "9 - nine",
+    mirror = [
+        "https://storage.googleapis.com/tensorflow/tf-keras-datasets/",
+        "https://zenodo.org/record/3239543/files/",
     ]
 
-    @property
-    def train_labels(self):
-        return self.targets
-
-    @property
-    def test_labels(self):
-        return self.targets
-
-    @property
-    def train_data(self):
-        return self.data
+    resources = [
+        ("mnist.npz", "731c5ac602752760c8e48fbffcf8c3b850d9dc2a2aedcf2cc48468fc17b673d1"),
+        ("mnist_c.zip", "4b34b33045869ee6d424616cd3a65da3"),
+    ]
 
-    @property
-    def test_data(self):
-        return self.data
+    class_dict = {
+        "zero": 0,
+        "one": 1,
+        "two": 2,
+        "three": 3,
+        "four": 4,
+        "five": 5,
+        "six": 6,
+        "seven": 7,
+        "eight": 8,
+        "nine": 9,
+    }
 
     def __init__(
         self,
@@ -179,25 +245,12 @@ class MNIST(Dataset):
         channels: Literal["channels_first", "channels_last"] | None = None,
         flatten: bool = False,
         normalize: tuple[float, float] | None = None,
-        corruption: Literal[
-            "identity",
-            "shot_noise",
-            "impulse_noise",
-            "glass_blur",
-            "motion_blur",
-            "shear",
-            "scale",
-            "rotate",
-            "brightness",
-            "translate",
-            "stripe",
-            "fog",
-            "spatter",
-            "dotted_line",
-            "zigzag",
-            "canny_edges",
-        ]
-        | None = None,
+        corruption: CorruptionStringMap | None = None,
+        classes: TClassMap | None = None,
+        balance: bool = True,
+        randomize: bool = False,
+        slice_back: bool = False,
+        verbose: bool = True,
     ) -> None:
         if isinstance(root, str):
             root = os.path.expanduser(root)
@@ -209,64 +262,113 @@ class MNIST(Dataset):
         self.channels = channels
         self.flatten = flatten
         self.normalize = normalize
-
-        if corruption is None:
-            corruption = "identity"
-        elif corruption == "identity":
-            print("Identity is not a corrupted dataset but the original MNIST dataset")
         self.corruption = corruption
-
-        if os.path.exists(self.mnist_folder):
-            print("Files already downloaded and verified")
-        elif download:
-            download_dataset(self.mirror, self.root, self.resources[0], self.resources[1])
+        self.balance = balance
+        self.randomize = randomize
+        self.from_back = slice_back
+        self.verbose = verbose
+
+        self.class_set = []
+        if classes is not None:
+            if not isinstance(classes, list):
+                classes = [classes]  # type: ignore
+
+            for val in classes:  # type: ignore
+                if isinstance(val, int) and 0 <= val < 10:
+                    self.class_set.append(val)
+                elif isinstance(val, str):
+                    self.class_set.append(self.class_dict[val])
+            self.class_set = set(self.class_set)
+
+        if not self.class_set:
+            self.class_set = set(self.class_dict.values())
+
+        self.num_classes = len(self.class_set)
+
+        if self.corruption is None:
+            file_resource = self.resources[0]
+            mirror = self.mirror[0]
+            md5 = False
         else:
-            raise RuntimeError("Dataset not found. You can use download=True to download it")
+            if self.corruption == "identity" and verbose:
+                print("Identity is not a corrupted dataset but the original MNIST dataset.")
+            file_resource = self.resources[1]
+            mirror = self.mirror[1]
+            md5 = True
+        check_exists(self.mnist_folder, mirror, self.root, file_resource[0], file_resource[1], download, verbose, md5)
 
         self.data, self.targets = self._load_data()
 
+        self._augmentations()
+
     def _load_data(self):
-        image_file = f"{'train' if self.train else 'test'}_images.npy"
-        data = self._read_image_file(os.path.join(self.mnist_folder, image_file))
-
-        label_file = f"{'train' if self.train else 'test'}_labels.npy"
-        targets = self._read_label_file(os.path.join(self.mnist_folder, label_file))
-
-        if self.size >= 1 and self.size >= len(self.classes):
-            final_data = []
-            final_targets = []
-            for label in range(len(self.classes)):
-                indices = np.where(targets == label)[0]
-                selected_indices = indices[: int(self.size / len(self.classes))]
-                final_data.append(data[selected_indices])
-                final_targets.append(targets[selected_indices])
-            data = np.concatenate(final_data)
-            targets = np.concatenate(final_targets)
-            shuffled_indices = np.random.permutation(data.shape[0])
-            data = data[shuffled_indices]
-            targets = targets[shuffled_indices]
-        elif self.size >= 1:
-            data = data[: self.size]
-            targets = targets[: self.size]
+        if self.corruption is None:
+            image_file = self.resources[0][0]
+            data, targets = self._read_normal_file(os.path.join(self.mnist_folder, image_file))
+        else:
+            image_file = f"{'train' if self.train else 'test'}_images.npy"
+            data = self._read_corrupt_file(os.path.join(self.mnist_folder, image_file))
+            data = data.squeeze()
+
+            label_file = f"{'train' if self.train else 'test'}_labels.npy"
+            targets = self._read_corrupt_file(os.path.join(self.mnist_folder, label_file))
+
+        return data, targets
+
+    def _augmentations(self):
+        if self.size > self.targets.shape[0] and self.verbose:
+            warn(
+                f"Asked for more samples, {self.size}, than the raw dataset contains, {self.targets.shape[0]}. "
+                "Adjusting down to raw dataset size."
+            )
+            self.size = -1
+
+        if self.randomize:
+            rdm_seed = np.random.default_rng(2023)
+            shuffled_indices = rdm_seed.permutation(self.data.shape[0])
+            self.data = self.data[shuffled_indices]
+            self.targets = self.targets[shuffled_indices]
+
+        if not self.balance and self.num_classes > self.size:
+            if self.size > 0:
+                self.data = subselect(self.data, self.size, self.from_back)
+                self.targets = subselect(self.targets, self.size, self.from_back)
+        else:
+            label_dict = {label: np.where(self.targets == label)[0] for label in self.class_set}
+            min_label_count = min(len(indices) for indices in label_dict.values())
+
+            self.per_class_count = int(np.ceil(self.size / self.num_classes)) if self.size > 0 else min_label_count
+
+            if self.per_class_count > min_label_count:
+                self.per_class_count = min_label_count
+                if not self.balance and self.verbose:
+                    warn(
+                        f"Because of dataset limitations, only {min_label_count*self.num_classes} samples "
+                        f"will be returned, instead of the desired {self.size}."
+                    )
+
+            all_indices = np.empty(shape=(self.num_classes, self.per_class_count), dtype=int)
+            for i, label in enumerate(self.class_set):
+                all_indices[i] = subselect(label_dict[label], self.per_class_count, self.from_back)
+            self.data = np.vstack(self.data[all_indices.T])  # type: ignore
+            self.targets = np.hstack(self.targets[all_indices.T])  # type: ignore
 
         if self.unit_interval:
-            data = data / 255
+            self.data = self.data / 255
 
         if self.normalize:
-            data = (data - self.normalize[0]) / self.normalize[1]
+            self.data = (self.data - self.normalize[0]) / self.normalize[1]
 
         if self.dtype:
-            data = data.astype(self.dtype)
+            self.data = self.data.astype(self.dtype)
 
         if self.channels == "channels_first":
-            data = np.moveaxis(data, -1, 1)
-        elif self.channels is None:
-            data = data[:, :, :, 0]
+            self.data = self.data[:, np.newaxis, :, :]
+        elif self.channels == "channels_last":
+            self.data = self.data[:, :, :, np.newaxis]
 
         if self.flatten and self.channels is None:
-            data = data.reshape(data.shape[0], -1)
-
-        return data, targets
+            self.data = self.data.reshape(self.data.shape[0], -1)
 
     def __getitem__(self, index: int) -> tuple[NDArray, int]:
         """
@@ -285,16 +387,18 @@ class MNIST(Dataset):
 
     @property
     def mnist_folder(self) -> str:
+        if self.corruption is None:
+            return os.path.join(self.root, "mnist")
         return os.path.join(self.root, "mnist_c", self.corruption)
 
-    @property
-    def class_to_idx(self) -> dict[str, int]:
-        return {_class: i for i, _class in enumerate(self.classes)}
-
-    def _read_label_file(self, path: str) -> NDArray:
-        x = np.load(path, allow_pickle=False)
-        return x
+    def _read_normal_file(self, path: str) -> tuple[NDArray, NDArray]:
+        with np.load(path, allow_pickle=True) as f:
+            if self.train:
+                x, y = f["x_train"], f["y_train"]
+            else:
+                x, y = f["x_test"], f["y_test"]
+        return x, y
 
-    def _read_image_file(self, path: str) -> NDArray:
+    def _read_corrupt_file(self, path: str) -> NDArray:
         x = np.load(path, allow_pickle=False)
         return x

dataeval/_internal/detectors/clusterer.py

@@ -16,6 +16,8 @@ from dataeval._internal.output import OutputMetadata, set_metadata
 @dataclass(frozen=True)
 class ClustererOutput(OutputMetadata):
     """
+    Output class for :class:`Clusterer` lint detector
+
     Attributes
     ----------
     outliers : List[int]

dataeval/_internal/detectors/drift/base.py

@@ -16,14 +16,14 @@ from typing import Callable, Literal
 import numpy as np
 from numpy.typing import ArrayLike, NDArray
 
-from dataeval._internal.interop import to_numpy
+from dataeval._internal.interop import as_numpy, to_numpy
 from dataeval._internal.output import OutputMetadata, set_metadata
 
 
 @dataclass(frozen=True)
 class DriftBaseOutput(OutputMetadata):
     """
-    Output class for Drift
+    Base output class for Drift detector classes
 
     Attributes
     ----------
@@ -42,7 +42,7 @@ class DriftBaseOutput(OutputMetadata):
 @dataclass(frozen=True)
 class DriftOutput(DriftBaseOutput):
     """
-    Output class for DriftCVM and DriftKS
+    Output class for :class:`DriftCVM`, :class:`DriftKS`, and :class:`DriftUncertainty` drift detectors
 
     Attributes
     ----------
@@ -234,7 +234,7 @@ class BaseDrift:
         if correction not in ["bonferroni", "fdr"]:
             raise ValueError("`correction` must be `bonferroni` or `fdr`.")
 
-        self._x_ref = x_ref
+        self._x_ref = to_numpy(x_ref)
         self.x_ref_preprocessed = x_ref_preprocessed
 
         # Other attributes
@@ -242,7 +242,7 @@ class BaseDrift:
         self.update_x_ref = update_x_ref
         self.preprocess_fn = preprocess_fn
         self.correction = correction
-        self.n = len(self._x_ref)  # type: ignore
+        self.n = len(self._x_ref)
 
         # Ref counter for preprocessed x
         self._x_refcount = 0
@@ -260,9 +260,8 @@ class BaseDrift:
         if not self.x_ref_preprocessed:
             self.x_ref_preprocessed = True
             if self.preprocess_fn is not None:
-                self._x_ref = self.preprocess_fn(self._x_ref)
+                self._x_ref = as_numpy(self.preprocess_fn(self._x_ref))
 
-        self._x_ref = to_numpy(self._x_ref)
         return self._x_ref
 
     def _preprocess(self, x: ArrayLike) -> ArrayLike:
@@ -380,7 +379,7 @@ class BaseDriftUnivariate(BaseDrift):
                 self._n_features = self.x_ref.reshape(self.x_ref.shape[0], -1).shape[-1]
             else:
                 # infer number of features after applying preprocessing step
-                x = to_numpy(self.preprocess_fn(self._x_ref[0:1]))  # type: ignore
+                x = as_numpy(self.preprocess_fn(self._x_ref[0:1]))  # type: ignore
                 self._n_features = x.reshape(x.shape[0], -1).shape[-1]
 
         return self._n_features

dataeval/_internal/detectors/drift/mmd.py

@@ -14,7 +14,7 @@ from typing import Callable
 import torch
 from numpy.typing import ArrayLike
 
-from dataeval._internal.interop import to_numpy
+from dataeval._internal.interop import as_numpy
 from dataeval._internal.output import set_metadata
 
 from .base import BaseDrift, DriftBaseOutput, UpdateStrategy, preprocess_x, update_x_ref
@@ -24,7 +24,7 @@ from .torch import GaussianRBF, get_device, mmd2_from_kernel_matrix
 @dataclass(frozen=True)
 class DriftMMDOutput(DriftBaseOutput):
     """
-    Output class for DriftMMD
+    Output class for :class:`DriftMMD` drift detector
 
     Attributes
     ----------
@@ -110,7 +110,7 @@ class DriftMMD(BaseDrift):
         self.device = get_device(device)
 
         # initialize kernel
-        sigma_tensor = torch.from_numpy(to_numpy(sigma)).to(self.device) if sigma is not None else None
+        sigma_tensor = torch.from_numpy(as_numpy(sigma)).to(self.device) if sigma is not None else None
         self.kernel = kernel(sigma_tensor).to(self.device) if kernel == GaussianRBF else kernel
 
         # compute kernel matrix for the reference data
@@ -147,7 +147,7 @@ class DriftMMD(BaseDrift):
             p-value obtained from the permutation test, MMD^2 between the reference and test set,
             and MMD^2 threshold above which drift is flagged
         """
-        x = to_numpy(x)
+        x = as_numpy(x)
         x_ref = torch.from_numpy(self.x_ref).to(self.device)
         n = x.shape[0]
         kernel_mat = self._kernel_matrix(x_ref, torch.from_numpy(x).to(self.device))