dataeval 0.70.0__py3-none-any.whl → 0.71.0__py3-none-any.whl

dataeval/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "0.70.0"
+__version__ = "0.71.0"
 
 from importlib.util import find_spec
 
@@ -12,11 +12,11 @@ from . import detectors, metrics  # noqa: E402
 __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

@@ -23,7 +23,7 @@ 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
     ----------

dataeval/_internal/detectors/drift/mmd.py

@@ -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
     ----------

dataeval/_internal/detectors/duplicates.py

@@ -17,6 +17,8 @@ TIndexCollection = TypeVar("TIndexCollection", DuplicateGroup, DatasetDuplicateG
 @dataclass(frozen=True)
 class DuplicatesOutput(Generic[TIndexCollection], OutputMetadata):
     """
+    Output class for :class:`Duplicates` lint detector
+
     Attributes
     ----------
     exact : list[list[int] | dict[int, list[int]]]

dataeval/_internal/detectors/ood/ae.py

@@ -15,10 +15,11 @@ import numpy as np
 import tensorflow as tf
 from numpy.typing import ArrayLike
 
-from dataeval._internal.detectors.ood.base import OODBase, OODScore
+from dataeval._internal.detectors.ood.base import OODBase, OODScoreOutput
 from dataeval._internal.interop import as_numpy
 from dataeval._internal.models.tensorflow.autoencoder import AE
 from dataeval._internal.models.tensorflow.utils import predict_batch
+from dataeval._internal.output import set_metadata
 
 
 class OOD_AE(OODBase):
@@ -48,7 +49,8 @@ class OOD_AE(OODBase):
             loss_fn = keras.losses.MeanSquaredError()
         super().fit(as_numpy(x_ref), threshold_perc, loss_fn, optimizer, epochs, batch_size, verbose)
 
-    def score(self, X: ArrayLike, batch_size: int = int(1e10)) -> OODScore:
+    @set_metadata("dataeval.detectors")
+    def score(self, X: ArrayLike, batch_size: int = int(1e10)) -> OODScoreOutput:
         self._validate(X := as_numpy(X))
 
         # reconstruct instances
@@ -62,4 +64,4 @@ class OOD_AE(OODBase):
         sorted_fscore_perc = sorted_fscore[:, -n_score_features:]
         iscore = np.mean(sorted_fscore_perc, axis=1)
 
-        return OODScore(iscore, fscore)
+        return OODScoreOutput(iscore, fscore)

dataeval/_internal/detectors/ood/aegmm.py

@@ -14,12 +14,13 @@ import keras
 import tensorflow as tf
 from numpy.typing import ArrayLike
 
-from dataeval._internal.detectors.ood.base import OODGMMBase, OODScore
+from dataeval._internal.detectors.ood.base import OODGMMBase, OODScoreOutput
 from dataeval._internal.interop import to_numpy
 from dataeval._internal.models.tensorflow.autoencoder import AEGMM
 from dataeval._internal.models.tensorflow.gmm import gmm_energy
 from dataeval._internal.models.tensorflow.losses import LossGMM
 from dataeval._internal.models.tensorflow.utils import predict_batch
+from dataeval._internal.output import set_metadata
 
 
 class OOD_AEGMM(OODGMMBase):
@@ -49,7 +50,8 @@ class OOD_AEGMM(OODGMMBase):
             loss_fn = LossGMM()
         super().fit(x_ref, threshold_perc, loss_fn, optimizer, epochs, batch_size, verbose)
 
-    def score(self, X: ArrayLike, batch_size: int = int(1e10)) -> OODScore:
+    @set_metadata("dataeval.detectors")
+    def score(self, X: ArrayLike, batch_size: int = int(1e10)) -> OODScoreOutput:
         """
         Compute the out-of-distribution (OOD) score for a given dataset.
 
@@ -63,7 +65,7 @@ class OOD_AEGMM(OODGMMBase):
 
         Returns
         -------
-        OODScore
+        OODScoreOutput
             An object containing the instance-level OOD score.
 
         Note
@@ -73,4 +75,4 @@ class OOD_AEGMM(OODGMMBase):
         self._validate(X := to_numpy(X))
         _, z, _ = predict_batch(X, self.model, batch_size=batch_size)
         energy, _ = gmm_energy(z, self.gmm_params, return_mean=False)
-        return OODScore(energy.numpy())  # type: ignore
+        return OODScoreOutput(energy.numpy())  # type: ignore