dataeval 0.83.0__py3-none-any.whl → 0.84.0__py3-none-any.whl

dataeval/utils/data/_embeddings.py

@@ -3,14 +3,14 @@ from __future__ import annotations
 __all__ = []
 
 import math
-from typing import Any, Iterator, Sequence
+from typing import Any, Iterator, Sequence, cast
 
 import torch
 from torch.utils.data import DataLoader, Subset
 from tqdm import tqdm
 
 from dataeval.config import DeviceLike, get_device
-from dataeval.typing import Array, Dataset
+from dataeval.typing import Array, Dataset, Transform
 from dataeval.utils.torch.models import SupportsEncode
 
 
@@ -26,11 +26,15 @@ class Embeddings:
         Dataset to access original images from.
     batch_size : int
         Batch size to use when encoding images.
+    transforms : Transform or Sequence[Transform] or None, default None
+        Transforms to apply to images before encoding.
     model : torch.nn.Module or None, default None
         Model to use for encoding images.
     device : DeviceLike or None, default None
         The hardware device to use if specified, otherwise uses the DataEval
         default or torch default.
+    cache : bool, default False
+        Whether to cache the embeddings in memory.
     verbose : bool, default False
         Whether to print progress bar when encoding images.
     """
@@ -43,19 +47,27 @@ class Embeddings:
         self,
         dataset: Dataset[tuple[Array, Any, Any]],
         batch_size: int,
+        transforms: Transform[torch.Tensor] | Sequence[Transform[torch.Tensor]] | None = None,
         model: torch.nn.Module | None = None,
         device: DeviceLike | None = None,
+        cache: bool = False,
         verbose: bool = False,
     ) -> None:
         self.device = get_device(device)
-        self.batch_size = batch_size
+        self.cache = cache
+        self.batch_size = batch_size if batch_size > 0 else 1
         self.verbose = verbose
 
         self._dataset = dataset
+        self._length = len(dataset)
         model = torch.nn.Flatten() if model is None else model
+        self._transforms = [transforms] if isinstance(transforms, Transform) else transforms
         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]
+        self._cached_idx = set()
+        self._embeddings: torch.Tensor = torch.empty(())
+        self._shallow: bool = False
 
     def to_tensor(self, indices: Sequence[int] | None = None) -> torch.Tensor:
         """
@@ -79,30 +91,96 @@ class Embeddings:
         else:
             return self[:]
 
-    # Reduce overhead cost by not tracking tensor gradients
-    @torch.no_grad
+    @classmethod
+    def from_array(cls, array: Array, device: DeviceLike | None = None) -> Embeddings:
+        """
+        Instantiates a shallow Embeddings object using an array.
+
+        Parameters
+        ----------
+        array : Array
+            The array to convert to embeddings.
+        device : DeviceLike or None, default None
+            The hardware device to use if specified, otherwise uses the DataEval
+            default or torch default.
+
+        Returns
+        -------
+        Embeddings
+
+        Example
+        -------
+        >>> import numpy as np
+        >>> from dataeval.utils.data._embeddings import Embeddings
+        >>> array = np.random.randn(100, 3, 224, 224)
+        >>> embeddings = Embeddings.from_array(array)
+        >>> print(embeddings.to_tensor().shape)
+        torch.Size([100, 3, 224, 224])
+        """
+        embeddings = Embeddings([], 0, None, None, device, True, False)
+        embeddings._length = len(array)
+        embeddings._cached_idx = set(range(len(array)))
+        embeddings._embeddings = torch.as_tensor(array).to(get_device(device))
+        embeddings._shallow = True
+        return embeddings
+
+    def _encode(self, images: list[torch.Tensor]) -> torch.Tensor:
+        if self._transforms:
+            images = [transform(image) for transform in self._transforms for image in images]
+        return self._encoder(torch.stack(images).to(self.device))
+
+    @torch.no_grad()  # Reduce overhead cost by not tracking tensor gradients
     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)  # type: ignore
-        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
+        dataset = cast(torch.utils.data.Dataset[tuple[Array, Any, Any]], self._dataset)
+        total_batches = math.ceil(len(indices) / self.batch_size)
+
+        # If not caching, process all indices normally
+        if not self.cache:
+            for images in tqdm(
+                DataLoader(Subset(dataset, indices), self.batch_size, collate_fn=self._collate_fn),
+                total=total_batches,
+                desc="Batch embedding",
+                disable=not self.verbose,
+            ):
+                yield self._encode(images)
+            return
+
+        # If caching, process each batch of indices at a time, preserving original order
+        for i in tqdm(range(0, len(indices), self.batch_size), desc="Batch embedding", disable=not self.verbose):
+            batch = indices[i : i + self.batch_size]
+            uncached = [idx for idx in batch if idx not in self._cached_idx]
+
+            if uncached:
+                # Process uncached indices as as single batch
+                for images in DataLoader(Subset(dataset, uncached), len(uncached), collate_fn=self._collate_fn):
+                    embeddings = self._encode(images)
+
+                    if not self._embeddings.shape:
+                        full_shape = (len(self._dataset), *embeddings.shape[1:])
+                        self._embeddings = torch.empty(full_shape, dtype=embeddings.dtype, device=self.device)
+
+                    self._embeddings[uncached] = embeddings
+                    self._cached_idx.update(uncached)
+
+            yield self._embeddings[batch]
 
     def __getitem__(self, key: int | slice, /) -> torch.Tensor:
-        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.")
+        if not isinstance(key, slice) and not hasattr(key, "__int__"):
+            raise TypeError("Invalid argument type.")
+
+        if self._shallow:
+            if not self._embeddings.shape:
+                raise ValueError("Embeddings not initialized.")
+            return self._embeddings[key]
+
+        indices = list(range(len(self._dataset))[key]) if isinstance(key, slice) else [int(key)]
+        result = torch.vstack(list(self._batch(indices))).to(self.device)
+        return result.squeeze(0) if len(indices) == 1 else result
 
     def __iter__(self) -> Iterator[torch.Tensor]:
         # process in batches while yielding individual embeddings
-        for batch in self._batch(range(len(self._dataset))):
+        for batch in self._batch(range(self._length)):
             yield from batch
 
     def __len__(self) -> int:
-        return len(self._dataset)
+        return self._length

dataeval/utils/data/_images.py

@@ -2,11 +2,15 @@ from __future__ import annotations
 
 __all__ = []
 
-from typing import Any, Generic, Iterator, Sequence, TypeVar, cast, overload
+from typing import TYPE_CHECKING, Any, Generic, Iterator, Sequence, TypeVar, cast, overload
 
-from dataeval.typing import Dataset
+from dataeval.typing import Array, Dataset
+from dataeval.utils._array import as_numpy, channels_first_to_last
 
-T = TypeVar("T")
+if TYPE_CHECKING:
+    from matplotlib.figure import Figure
+
+T = TypeVar("T", bound=Array)
 
 
 class Images(Generic[T]):
@@ -21,7 +25,10 @@ class Images(Generic[T]):
         Dataset to access images from.
     """
 
-    def __init__(self, dataset: Dataset[tuple[T, Any, Any] | T]) -> None:
+    def __init__(
+        self,
+        dataset: Dataset[tuple[T, Any, Any] | T],
+    ) -> None:
         self._is_tuple_datum = isinstance(dataset[0], tuple)
         self._dataset = dataset
 
@@ -40,25 +47,41 @@ class Images(Generic[T]):
         """
         return self[:]
 
+    def plot(
+        self,
+        indices: Sequence[int],
+        images_per_row: int = 3,
+        figsize: tuple[int, int] = (10, 10),
+    ) -> Figure:
+        import matplotlib.pyplot as plt
+
+        num_images = len(indices)
+        num_rows = (num_images + images_per_row - 1) // images_per_row
+        fig, axes = plt.subplots(num_rows, images_per_row, figsize=figsize)
+        for i, ax in enumerate(axes.flatten()):
+            image = channels_first_to_last(as_numpy(self[i]))
+            ax.imshow(image)
+            ax.axis("off")
+        plt.tight_layout()
+        return fig
+
     @overload
     def __getitem__(self, key: int, /) -> T: ...
     @overload
     def __getitem__(self, key: slice, /) -> Sequence[T]: ...
 
     def __getitem__(self, key: int | slice, /) -> Sequence[T] | T:
+        if isinstance(key, slice):
+            return [self._get_image(k) for k in range(len(self._dataset))[key]]
+        elif hasattr(key, "__int__"):
+            return self._get_image(int(key))
+        raise TypeError(f"Key must be integers or slices, not {type(key)}")
+
+    def _get_image(self, index: int) -> T:
         if self._is_tuple_datum:
-            dataset = cast(Dataset[tuple[T, Any, Any]], self._dataset)
-            if isinstance(key, slice):
-                return [dataset[k][0] for k in range(len(self._dataset))[key]]
-            elif isinstance(key, int):
-                return dataset[key][0]
+            return cast(Dataset[tuple[T, Any, Any]], self._dataset)[index][0]
         else:
-            dataset = cast(Dataset[T], self._dataset)
-            if isinstance(key, slice):
-                return [dataset[k] for k in range(len(self._dataset))[key]]
-            elif isinstance(key, int):
-                return dataset[key]
-        raise TypeError(f"Key must be integers or slices, not {type(key)}")
+            return cast(Dataset[T], self._dataset)[index]
 
     def __iter__(self) -> Iterator[T]:
         for i in range(len(self._dataset)):

dataeval/utils/data/_selection.py

@@ -5,7 +5,7 @@ __all__ = []
 from enum import IntEnum
 from typing import Generic, Iterator, Sequence, TypeVar
 
-from dataeval.typing import AnnotatedDataset, DatasetMetadata, Transform
+from dataeval.typing import AnnotatedDataset, DatasetMetadata
 
 _TDatum = TypeVar("_TDatum")
 
@@ -35,8 +35,6 @@ class Select(AnnotatedDataset[_TDatum]):
         The dataset to wrap.
     selections : Selection or list[Selection], optional
         The selection criteria to apply to the dataset.
-    transforms : Transform or list[Transform], optional
-        The transforms to apply to the dataset.
 
     Examples
     --------
@@ -70,16 +68,12 @@ class Select(AnnotatedDataset[_TDatum]):
         self,
         dataset: AnnotatedDataset[_TDatum],
         selections: Selection[_TDatum] | Sequence[Selection[_TDatum]] | None = None,
-        transforms: Transform[_TDatum] | Sequence[Transform[_TDatum]] | None = None,
     ) -> None:
         self.__dict__.update(dataset.__dict__)
         self._dataset = dataset
         self._size_limit = len(dataset)
         self._selection = list(range(self._size_limit))
         self._selections = self._sort(selections)
-        self._transforms = (
-            [] if transforms is None else [transforms] if isinstance(transforms, Transform) else transforms
-        )
 
         # Ensure metadata is populated correctly as DatasetMetadata TypedDict
         _metadata = getattr(dataset, "metadata", {})
@@ -98,8 +92,7 @@ class Select(AnnotatedDataset[_TDatum]):
         title = f"{self.__class__.__name__} Dataset"
         sep = "-" * len(title)
         selections = f"Selections: [{', '.join([str(s) for s in self._selections])}]"
-        transforms = f"Transforms: [{', '.join([str(t) for t in self._transforms])}]"
-        return f"{title}\n{sep}{nt}{selections}{nt}{transforms}{nt}Selected Size: {len(self)}\n\n{self._dataset}"
+        return f"{title}\n{sep}{nt}{selections}{nt}Selected Size: {len(self)}\n\n{self._dataset}"
 
     def _sort(self, selections: Selection[_TDatum] | Sequence[Selection[_TDatum]] | None) -> list[Selection]:
         if not selections:
@@ -117,13 +110,8 @@ class Select(AnnotatedDataset[_TDatum]):
             selection(self)
         self._selection = self._selection[: self._size_limit]
 
-    def _transform(self, datum: _TDatum) -> _TDatum:
-        for t in self._transforms:
-            datum = t(datum)
-        return datum
-
     def __getitem__(self, index: int) -> _TDatum:
-        return self._transform(self._dataset[self._selection[index]])
+        return self._dataset[self._selection[index]]
 
     def __iter__(self) -> Iterator[_TDatum]:
         for i in range(len(self)):

dataeval/utils/data/_split.py

@@ -2,19 +2,22 @@ from __future__ import annotations
 
 __all__ = []
 
+import logging
 import warnings
-from typing import Any, Iterator, Protocol
+from typing import Any, Iterator, Protocol, Sequence
 
 import numpy as np
 from numpy.typing import NDArray
-from sklearn.cluster import KMeans
-from sklearn.metrics import silhouette_score
 from sklearn.model_selection import GroupKFold, KFold, StratifiedGroupKFold, StratifiedKFold
 from sklearn.utils.multiclass import type_of_target
 
-from dataeval.config import get_seed
+from dataeval.config import EPSILON
 from dataeval.outputs._base import set_metadata
 from dataeval.outputs._utils import SplitDatasetOutput, TrainValSplit
+from dataeval.typing import AnnotatedDataset
+from dataeval.utils.data._metadata import Metadata
+
+_logger = logging.getLogger(__name__)
 
 
 class KFoldSplitter(Protocol):
@@ -85,7 +88,7 @@ def calculate_validation_fraction(num_folds: int, test_frac: float, val_frac: fl
     return val_base * (1.0 / num_folds) * (1.0 - test_frac)
 
 
-def _validate_labels(labels: NDArray[np.intp], total_partitions: int) -> None:
+def validate_labels(labels: NDArray[np.intp], total_partitions: int) -> None:
     """
     Check to make sure there is more input data than the total number of partitions requested
 
@@ -116,7 +119,7 @@ def _validate_labels(labels: NDArray[np.intp], total_partitions: int) -> None:
         raise ValueError("Detected continuous labels. Labels must be discrete for proper stratification")
 
 
-def is_stratifiable(labels: NDArray[np.intp], num_partitions: int) -> bool:
+def validate_stratifiable(labels: NDArray[np.intp], num_partitions: int) -> None:
     """
     Check if the dataset can be stratified by class label over the given number of partitions
 
@@ -132,26 +135,23 @@ def is_stratifiable(labels: NDArray[np.intp], num_partitions: int) -> bool:
     bool
         True if dataset can be stratified else False
 
-    Warns
-    -----
-    UserWarning
-        Warns user if the dataset cannot be stratified due to the total number of [train, val, test]
+    Raises
+    ------
+    ValueError
+        If the dataset cannot be stratified due to the total number of [train, val, test]
         partitions exceeding the number of instances of the rarest class label.
     """
 
     # Get the minimum count of all labels
     lowest_label_count = np.unique(labels, return_counts=True)[1].min()
     if lowest_label_count < num_partitions:
-        warnings.warn(
+        raise ValueError(
             f"Unable to stratify due to label frequency. The lowest label count ({lowest_label_count}) is fewer "
-            f"than the total number of partitions ({num_partitions}) requested.",
-            UserWarning,
+            f"than the total number of partitions ({num_partitions}) requested."
         )
-        return False
-    return True
 
 
-def is_groupable(group_ids: NDArray[np.intp], num_partitions: int) -> bool:
+def validate_groupable(groups: NDArray[np.intp], num_partitions: int) -> None:
     """
     Warns user if the number of unique group_ids is incompatible with a grouped partition containing
     num_folds folds. If this is the case, returns groups=None, which tells the partitioner not to
@@ -159,7 +159,7 @@ def is_groupable(group_ids: NDArray[np.intp], num_partitions: int) -> bool:
 
     Parameters
     ----------
-    group_ids : NDArray of ints
+    groups : NDArray of ints
         The id of the group each sample at the corresponding index belongs to
     num_partitions : int
         Total number of train, val, and test splits requested
@@ -169,60 +169,24 @@ def is_groupable(group_ids: NDArray[np.intp], num_partitions: int) -> bool:
     bool
         True if the dataset can be grouped by the given group ids else False
 
-    Warns
-    -----
-    UserWarning
-        Warns if there are fewer groups than the requested number of partitions plus one
+    Raises
+    ------
+    ValueError
+        If there are is only one unique group.
+    ValueError
+        If there are fewer groups than the requested number of partitions plus one
     """
 
-    num_unique_groups = len(np.unique(group_ids))
+    num_unique_groups = len(np.unique(groups))
     # Cannot separate if only one group exists
     if num_unique_groups == 1:
-        return False
+        raise ValueError(f"Unique groups ({num_unique_groups}) must be greater than 1.")
 
     if num_unique_groups < num_partitions:
-        warnings.warn(
-            f"Groups must be greater than num partitions. Got {num_unique_groups} and {num_partitions}. "
-            "Reverting to ungrouped partitioning",
-            UserWarning,
-        )
-        return False
-    return True
-
-
-def bin_kmeans(array: NDArray[Any]) -> NDArray[np.intp]:
-    """
-    Find bins of continuous data by iteratively applying k-means clustering, and keeping the
-    clustering with the highest silhouette score.
-
-    Parameters
-    ----------
-    array : NDArray
-        continuous data to bin
+        raise ValueError(f"Unique groups ({num_unique_groups}) must be greater than num partitions ({num_partitions}).")
 
-    Returns
-    -------
-    NDArray[int]:
-        bin numbers assigned by the kmeans best clusterer.
-    """
 
-    if array.ndim == 1:
-        array = array.reshape([-1, 1])
-        best_score = 0.60
-    else:
-        best_score = 0.50
-    bin_index = np.zeros(len(array), dtype=np.intp)
-    for k in range(2, 20):
-        clusterer = KMeans(n_clusters=k, random_state=get_seed())
-        cluster_labels = clusterer.fit_predict(array)
-        score = silhouette_score(array, cluster_labels, sample_size=25_000, random_state=get_seed())
-        if score > best_score:
-            best_score = score
-            bin_index = cluster_labels.astype(np.intp)
-    return bin_index
-
-
-def get_group_ids(metadata: dict[str, Any], group_names: list[str], num_samples: int) -> NDArray[np.intp]:
+def get_groups(metadata: Metadata, split_on: Sequence[str] | None) -> NDArray[np.intp] | None:
     """
     Returns individual group numbers based on a subset of metadata defined by groupnames
 
@@ -232,32 +196,20 @@ def get_group_ids(metadata: dict[str, Any], group_names: list[str], num_samples:
         dictionary containing all metadata
     groupnames : list
         which groups from the metadata dictionary to consider for dataset grouping
-    num_samples : int
-        number of labels. Used to ensure agreement between input data/labels and metadata entries.
-
-    Raises
-    ------
-    IndexError
-        raised if an entry in the metadata dictionary doesn't have the same length as num_samples
 
     Returns
     -------
     np.ndarray
         group identifiers from metadata
     """
-    features2group = {k: np.array(v) for k, v in metadata.items() if k in group_names}
-    if not features2group:
-        return np.zeros(num_samples, dtype=np.intp)
-    for name, feature in features2group.items():
-        if len(feature) != num_samples:
-            raise ValueError(
-                f"Feature length does not match number of labels. Got {len(feature)} features and {num_samples} samples"
-            )
-
-        if type_of_target(feature) == "continuous":
-            features2group[name] = bin_kmeans(feature)
-    binned_features = np.stack(list(features2group.values()), axis=1)
-    _, group_ids = np.unique(binned_features, axis=0, return_inverse=True)
+    # get only the factors that are present in the metadata
+    if split_on is None:
+        return None
+
+    split_set = set(split_on)
+    indices = [i for i, name in enumerate(metadata.discrete_factor_names) if name in split_set]
+    binned_features = metadata.discrete_data[:, indices]
+    group_ids = np.unique(binned_features, axis=0, return_inverse=True)[1]
     return group_ids
 
 
@@ -294,10 +246,18 @@ def make_splits(
     split_defs: list[TrainValSplit] = []
     n_labels = len(np.unique(labels))
     splitter = KFOLD_GROUP_STRATIFIED_MAP[(groups is not None, stratified)](n_folds)
+    _logger.log(logging.DEBUG, f"splitter={splitter.__class__.__name__}(n_splits={n_folds})")
     good = False
     attempts = 0
     while not good and attempts < 3:
         attempts += 1
+        _logger.log(
+            logging.DEBUG,
+            f"attempt={attempts}: splitter.split("
+            + f"index=arr(len={len(index)}, unique={np.unique(index)}), "
+            + f"labels=arr(len={len(index)}, unique={np.unique(index)}), "
+            + ("groups=None" if groups is None else f"groups=arr(len={len(groups)}, unique={np.unique(groups)}))"),
+        )
         splits = splitter.split(index, labels, groups)
         split_defs.clear()
         for train_idx, eval_idx in splits:
@@ -341,20 +301,20 @@ def find_best_split(
         counts = np.bincount(arr, minlength=minlength)
         return counts / np.sum(counts)
 
-    def weight(arr: NDArray, class_freq: NDArray) -> np.float64:
-        return np.sum(np.abs(freq(arr, len(class_freq)) - class_freq))
+    def weight(arr: NDArray, class_freq: NDArray) -> float:
+        return float(np.sum(np.abs(freq(arr, len(class_freq)) - class_freq)))
 
-    def class_freq_diff(split: TrainValSplit) -> np.float64:
+    def class_freq_diff(split: TrainValSplit) -> float:
         class_freq = freq(labels)
         return weight(labels[split.train], class_freq) + weight(labels[split.val], class_freq)
 
-    def split_ratio(split: TrainValSplit) -> np.float64:
-        return np.float64(len(split.val) / (len(split.val) + len(split.train)))
+    def split_ratio(split: TrainValSplit) -> float:
+        return len(split.val) / (len(split.val) + len(split.train))
 
-    def split_diff(split: TrainValSplit) -> np.float64:
+    def split_diff(split: TrainValSplit) -> float:
         return abs(split_frac - split_ratio(split))
 
-    def split_inv_diff(split: TrainValSplit) -> np.float64:
+    def split_inv_diff(split: TrainValSplit) -> float:
         return abs(1 - split_frac - split_ratio(split))
 
     # Selects minimization function based on inputs
@@ -399,11 +359,12 @@ def single_split(
         Indices of data partitioned for training and evaluation
     """
 
-    _, label_counts = np.unique(labels, return_counts=True)
-    max_folds = label_counts.min()
-    min_folds = np.unique(groups).shape[0] if groups is not None else 2
-    divisor = split_frac + 1e-06 if split_frac <= 2 / 3 else 1 - split_frac - 1e-06
-    n_folds = round(min(max(1 / divisor, min_folds), max_folds))  # Clips value between min_folds and max_folds
+    unique_groups = 2 if groups is None else len(np.unique(groups))
+    max_folds = min(min(np.unique(labels, return_counts=True)[1]), unique_groups) if stratified else unique_groups
+
+    divisor = split_frac if split_frac <= 2 / 3 else 1 - split_frac
+    n_folds = min(max(round(1 / (divisor + EPSILON)), 2), max_folds)  # Clips value between 2 and max_folds
+    _logger.log(logging.DEBUG, f"n_folds={n_folds} clipped between[2, {max_folds}]")
 
     split_candidates = make_splits(index, labels, n_folds, groups, stratified)
     return find_best_split(labels, split_candidates, stratified, split_frac)
@@ -411,22 +372,20 @@ def single_split(
 
 @set_metadata
 def split_dataset(
-    labels: list[int] | NDArray[np.intp],
+    dataset: AnnotatedDataset[Any] | Metadata,
     num_folds: int = 1,
     stratify: bool = False,
-    split_on: list[str] | None = None,
-    metadata: dict[str, Any] | None = None,
+    split_on: Sequence[str] | None = None,
     test_frac: float = 0.0,
     val_frac: float = 0.0,
 ) -> SplitDatasetOutput:
     """
-    Top level splitting function. Returns a dataclass containing a list of train and validation indices.
-    Indices for a test holdout may also be optionally included
+    Dataset splitting function. Returns a dataclass containing a list of train and validation indices.
 
     Parameters
     ----------
-    labels : list or NDArray of ints
-        Classification Labels used to generate splits. Determines the size of the dataset
+    dataset : AnnotatedDataset or Metadata
+        Dataset to split.
     num_folds : int, default 1
         Number of [train, val] folds. If equal to 1, val_frac must be greater than 0.0
     stratify : bool, default False
@@ -436,8 +395,6 @@ def split_dataset(
         Keys of the metadata dictionary upon which to group the dataset.
         A grouped partition is divided such that no group is present within both the training and
         validation set. Split_on groups should be selected to mitigate validation bias
-    metadata : dict or None, default None
-        Dict containing data for potential dataset grouping. See split_on above
     test_frac : float, default 0.0
         Fraction of data to be optionally held out for test set
     val_frac : float, default 0.0
@@ -450,13 +407,8 @@ def split_dataset(
         Output class containing a list of indices of training
         and validation data for each fold and optional test indices
 
-    Raises
-    ------
-    TypeError
-        Raised if split_on is passed, but metadata is None or empty
-
-    Note
-    ----
+    Notes
+    -----
     When specifying groups and/or stratification, ratios for test and validation splits can vary
     as the stratification and grouping take higher priority than the percentages
     """
@@ -464,30 +416,25 @@ def split_dataset(
     val_frac = calculate_validation_fraction(num_folds, test_frac, val_frac)
     total_partitions = num_folds + 1 if test_frac else num_folds
 
-    if isinstance(labels, list):
-        labels = np.array(labels, dtype=np.intp)
+    metadata = dataset if isinstance(dataset, Metadata) else Metadata(dataset)
+    labels = metadata.class_labels
 
-    label_length: int = len(labels)
+    validate_labels(labels, total_partitions)
+    if stratify:
+        validate_stratifiable(labels, total_partitions)
 
-    _validate_labels(labels, total_partitions)
-    stratify &= is_stratifiable(labels, total_partitions)
-    groups = None
-    if split_on:
-        if metadata is None or metadata == {}:
-            raise TypeError("If split_on is specified, metadata must also be provided, got None")
-        possible_groups = get_group_ids(metadata, split_on, label_length)
+    groups = get_groups(metadata, split_on)
+    if groups is not None:
         # Accounts for a test set that is 100 % of the data
         group_partitions = total_partitions + 1 if val_frac else total_partitions
-        if is_groupable(possible_groups, group_partitions):
-            groups = possible_groups
+        validate_groupable(groups, group_partitions)
 
-    index = np.arange(label_length)
+    index = np.arange(len(labels))
 
-    tvs = (
-        single_split(index=index, labels=labels, split_frac=test_frac, groups=groups, stratified=stratify)
-        if test_frac
-        else TrainValSplit(index, np.array([], dtype=np.intp))
-    )
+    if test_frac:
+        tvs = single_split(index, labels, test_frac, groups, stratify)
+    else:
+        tvs = TrainValSplit(index, np.array([], dtype=np.intp))
 
     tv_labels = labels[tvs.train]
     tv_groups = groups[tvs.train] if groups is not None else None