stouputils 1.12.1__py3-none-any.whl

stouputils/data_science/dataset/dataset.py

@@ -0,0 +1,292 @@
+"""
+This module contains the Dataset class, which provides an easy way to handle ML datasets.
+
+The Dataset class has the following attributes:
+
+- training_data        (XyTuple):              Training data containing features, labels and file paths
+- test_data            (XyTuple):              Test data containing features, labels and file paths
+- num_classes          (int):                  Number of classes in the dataset
+- name                 (str):                  Name of the dataset
+- grouping_strategy    (GroupingStrategy):     Strategy for grouping images when loading
+- labels               (list[str]):            List of class labels (strings)
+- loading_type         (Literal["image"]):     Type of the dataset (currently only "image" is supported)
+- original_dataset     (Dataset | None):       Original dataset used for data augmentation
+- class_distribution   (dict[str, dict]):      Class distribution counts for train/test sets
+
+It provides methods for:
+
+- Loading image datasets from directories using different grouping strategies
+- Splitting data into train/test sets with stratification (and care for data augmentation)
+- Managing class distributions and dataset metadata
+"""
+# pyright: reportUnknownMemberType=false
+
+# Imports
+from __future__ import annotations
+
+import os
+from collections.abc import Generator, Iterable
+from typing import Any, Literal
+
+import numpy as np
+from numpy.typing import NDArray
+
+from ...decorators import handle_error, LogLevels
+from ...print import warning, progress
+from ...collections import unique_list
+from ..utils import Utils
+from .grouping_strategy import GroupingStrategy
+from .xy_tuple import XyTuple
+
+# Constants
+DEFAULT_IMAGE_KWARGS: dict[str, Any] = {
+	"image_size": (224, 224),
+	"label_mode": "categorical",
+	"color_mode": "rgb",
+	"batch_size": 1
+}
+""" Default image kwargs sent to keras.image_dataset_from_directory """
+
+# Dataset class
+class Dataset:
+	""" Dataset class used for easy data handling. """
+
+	# Class constructors
+	def __init__(
+		self,
+		training_data: XyTuple | list[Any],
+		val_data: XyTuple | list[Any] | None = None,
+		test_data: XyTuple | list[Any] | None = None,
+		name: str = "",
+		grouping_strategy: GroupingStrategy = GroupingStrategy.NONE,
+		labels: tuple[str, ...] = (),
+		loading_type: Literal["image"] = "image"
+	) -> None:
+		""" Initialize the dataset class
+
+		>>> Dataset(training_data=tuple(), test_data=tuple(), name="doctest")
+		Traceback (most recent call last):
+			...
+		AssertionError: data must be a tuple with X and y as iterables
+		"""
+		if val_data is None:
+			val_data = XyTuple.empty()
+		if test_data is None:
+			test_data = XyTuple.empty()
+
+		# Assertions
+		all_data: tuple[Any, ...] = (training_data, val_data, test_data)
+		for data in all_data:
+			if not isinstance(data, XyTuple):
+				assert isinstance(data, Iterable) \
+					and 2 <= len(data) <= 3 \
+					and isinstance(data[0], Iterable) \
+					and isinstance(data[1], Iterable), "data must be a tuple with X and y as iterables"
+
+		# Get training, validation and test data
+		xy_tuples: list[XyTuple] = [XyTuple(*data) if not isinstance(data, XyTuple) else data for data in all_data]
+
+		# Pre-process for attributes initialization
+		num_classes: int = self._get_num_classes(xy_tuples[0].y, xy_tuples[1].y, xy_tuples[2].y)
+		labels = tuple(str(x).replace("_", " ").title() for x in (labels if labels else range(num_classes)))
+
+		# Initialize attributes
+		self._training_data: XyTuple = xy_tuples[0]
+		""" Training data as XyTuple containing X and y as numpy arrays.
+		This is a protected attribute accessed via the public property self.training_data. """
+		self._val_data: XyTuple = xy_tuples[1]
+		""" Validation data as XyTuple containing X and y as numpy arrays.
+		This is a protected attribute accessed via the public property self.val_data. """
+		self._test_data: XyTuple = xy_tuples[2]
+		""" Test data as XyTuple containing X and y as numpy arrays.
+		This is a protected attribute accessed via the public property self.test_data. """
+		self.num_classes: int = num_classes
+		""" Number of classes in the dataset (y) """
+		self.name: str = os.path.basename(name)
+		""" Name of the dataset (path given in the constructor are converted,
+		ex: ".../data/pizza_not_pizza" becomes "pizza_not_pizza") """
+		self.loading_type: Literal["image"] = loading_type
+		""" Type of the dataset """
+		self.grouping_strategy: GroupingStrategy = grouping_strategy
+		""" Grouping strategy for the dataset """
+		self.labels: tuple[str, ...] = labels
+		""" List of class labels (strings) """
+		self.class_distribution: dict[str, dict[int, int]] = {"train": {}, "val": {}, "test": {}}
+		""" Class distribution in the dataset for both training and test sets """
+		self.original_dataset: Dataset | None = None
+		""" Original dataset used for data augmentation (can be None) """
+
+		# Update class distribution
+		self._update_class_distribution()
+
+	def _get_num_classes(self, *values: Any) -> int:
+		""" Get the number of classes in the dataset.
+
+		Args:
+			values (NDArray[Any]): Arrays containing class labels
+		Returns:
+			int: Number of unique classes
+		"""
+		# Handle case where arrays have different dimensions (1D vs 2D)
+		processed_values: list[NDArray[Any]] = []
+		for value in values:
+			value: NDArray[Any] = np.array(value)
+			if len(value.shape) == 2:  # One-hot encoded
+				processed_values.append(Utils.convert_to_class_indices(value))
+			else:
+				processed_values.append(value)
+
+		return len(np.unique(np.concatenate(processed_values)))
+
+	def _update_class_distribution(self, update_num_classes: bool = False) -> None:
+		""" Update the class distribution dictionary for both training and test data. """
+		# For each data type,
+		for data_type, data in (("train", self._training_data), ("val", self._val_data), ("test", self._test_data)):
+			y_data: NDArray[Any] = np.array(data.y)
+			if len(y_data.shape) == 2:  # One-hot encoded
+				y_data = Utils.convert_to_class_indices(y_data)
+
+			# Update the class distribution
+			self.class_distribution[data_type] = {}
+			for class_id in range(self.num_classes):
+				self.class_distribution[data_type][class_id] = np.sum(y_data == class_id)
+
+		# Update the number of classes if needed
+		if update_num_classes:
+			self.num_classes = self._get_num_classes(self._training_data.y, self._val_data.y, self._test_data.y)
+
+	def exclude_augmented_images_from_val_test(self, original_dataset: Dataset) -> None:
+		""" Exclude augmented versions of validation and test images from the training set.
+
+		This ensures that augmented versions of images in the validation and test sets are not present in the training set,
+		which would cause data leakage.
+
+		Args:
+			original_dataset (Dataset): The original dataset containing the test images to exclude
+		"""
+		# Get base filenames from original test set
+		progress("Excluding augmented versions of validation and test images from training set...")
+		val_test_base_names: list[list[str]] = [
+			[os.path.splitext(os.path.basename(f))[0] for f in filepaths]
+			for filepaths in (*original_dataset.val_data.filepaths, *original_dataset.test_data.filepaths)
+		]
+		val_test_base_names = unique_list(val_test_base_names, method="str")
+
+		# Get base filenames from training set
+		train_base_names: list[list[str]] = [
+			[os.path.splitext(os.path.basename(f))[0] for f in filepaths]
+			for filepaths in self.training_data.filepaths
+		]
+
+		# Remove augmented versions of test images from training set
+		# Get indices of training samples that are not augmented versions of test samples
+		# For each training sample, check if any of its filenames start with any test filename
+		train_indices: list[int] = [
+			i for i, train_names in enumerate(train_base_names)
+			if not any(
+				any(train_name.startswith(name) for train_name in train_names)
+				for names in val_test_base_names
+				for name in names
+			)
+		]
+
+		# Update training data to exclude augmented versions
+		self._training_data = XyTuple(
+			[self.training_data.X[i] for i in train_indices],
+			[self.training_data.y[i] for i in train_indices],
+			tuple(self.training_data.filepaths[i] for i in train_indices)
+		)
+
+		# Use original test data
+		self._test_data = original_dataset.test_data     # Impossible to have augmented test_data here
+		self._val_data = original_dataset.val_data       # Impossible to have augmented val_data here
+		self._update_class_distribution(update_num_classes=False)
+
+	@property
+	def training_data(self) -> XyTuple:
+		return self._training_data
+
+	@training_data.setter
+	def training_data(self, value: XyTuple | Any) -> None:
+		warning("Setting training data...", value)
+		self._training_data = XyTuple(*value) if not isinstance(value, XyTuple) else value
+		self._update_class_distribution(update_num_classes=True)
+
+	@property
+	def val_data(self) -> XyTuple:
+		return self._val_data
+
+	@val_data.setter
+	def val_data(self, value: XyTuple | Any) -> None:
+		self._val_data = XyTuple(*value) if not isinstance(value, XyTuple) else value
+		self._update_class_distribution(update_num_classes=True)
+
+	@property
+	def test_data(self) -> XyTuple:
+		return self._test_data
+
+	@test_data.setter
+	def test_data(self, value: XyTuple | Any) -> None:
+		self._test_data = XyTuple(*value) if not isinstance(value, XyTuple) else value
+		self._update_class_distribution(update_num_classes=True)
+
+	# Class methods
+	def __str__(self) -> str:
+		train_dist: dict[int, int] = self.class_distribution["train"]
+		val_dist: dict[int, int] = self.class_distribution["val"]
+		test_dist: dict[int, int] = self.class_distribution["test"]
+		return (
+			f"Dataset {self.name}: "
+			f"{len(self.training_data.X):,} training samples, "
+			f"{len(self.val_data.X):,} validation samples, "
+			f"{len(self.test_data.X):,} test samples, "
+			f"{self.num_classes:,} classes "
+			f"(train: {train_dist}, val: {val_dist}, test: {test_dist})"
+		)
+
+	def __repr__(self) -> str:
+		return (
+			f"Dataset(training_data={self.training_data!r}, "
+			f"val_data={self.val_data!r}, "
+			f"test_data={self.test_data!r}, "
+			f"num_classes={self.num_classes}, "
+			f"name={self.name!r}, "
+			f"grouping_strategy={self.grouping_strategy.name})"
+		)
+
+	def __iter__(self) -> Generator[XyTuple, Any, Any]:
+		""" Allow unpacking of the dataset into train and test sets.
+
+		Returns:
+			Generator[XyTuple], Any, Any]: Generator over the dataset splits
+
+		>>> X, y = [[1]], [2]
+		>>> dataset = Dataset(training_data=(X, y), test_data=(X, y), name="doctest")
+		>>> train, val, test = dataset
+		>>> train == (X, y) and val == () and test == (X, y)
+		True
+		>>> train == XyTuple(X, y) and val == XyTuple.empty() and test == XyTuple(X, y)
+		True
+		"""
+		yield from (self.training_data, self.val_data, self.test_data)
+
+	def get_experiment_name(self, override_name: str = "") -> str:
+		""" Get the experiment name for mlflow, e.g. "DatasetName_GroupingStrategyName"
+
+		Args:
+			override_name (str): Override the Dataset name
+		Returns:
+			str: Experiment name
+		"""
+		if override_name:
+			return f"{override_name}_{self.grouping_strategy.name.title()}"
+		else:
+			return f"{self.name}_{self.grouping_strategy.name.title()}"
+
+
+	# Static methods
+	@staticmethod
+	@handle_error(error_log=LogLevels.ERROR_TRACEBACK)
+	def empty() -> Dataset:
+		return Dataset(XyTuple.empty(), name="empty", grouping_strategy=GroupingStrategy.NONE)
+

stouputils/data_science/dataset/dataset_loader.py

@@ -0,0 +1,135 @@
+"""
+This module contains the DatasetLoader class which handles dataset loading operations.
+
+The DatasetLoader class provides the following key features:
+
+- Loading image datasets from directories using keras.image_dataset_from_directory
+- Handling different grouping strategies (when having multiple images per subject)
+- Preventing data leakage between train/test sets when using data augmentation
+- Ensuring test data consistency when loading an augmented dataset
+"""
+
+# Imports
+from typing import Any, Literal
+
+import numpy as np
+
+from ...decorators import handle_error, LogLevels
+from ..config.get import DataScienceConfig
+from .dataset import Dataset
+from .grouping_strategy import GroupingStrategy
+from .xy_tuple import XyTuple
+
+# Constants
+DEFAULT_IMAGE_KWARGS: dict[str, Any] = {
+    "image_size": (224, 224),
+    "color_mode": "RGB",
+}
+
+class DatasetLoader:
+	""" Handles dataset loading operations """
+
+	@staticmethod
+	@handle_error(error_log=LogLevels.ERROR_TRACEBACK)
+	def from_path(
+		path: str,
+		loading_type: Literal["image"] = "image",
+		seed: int = DataScienceConfig.SEED,
+		test_size: float = 0.2,
+		val_size: float = 0.2,
+		grouping_strategy: GroupingStrategy = GroupingStrategy.NONE,
+		based_of: str = "",
+		**kwargs: Any
+	) -> Dataset:
+		""" Create a balanced dataset from a path.
+
+		Args:
+			path              (str):               Path to the dataset
+			loading_type      (Literal["image"]):  Type of the dataset
+			seed              (int):               Seed for the random generator
+			test_size         (float):             Size of the test dataset (0 means no test set)
+			val_size          (float):             Size of the validation dataset (0 means no validation set)
+			grouping_strategy (GroupingStrategy):  Grouping strategy for the dataset (ex: GroupingStrategy.CONCATENATE)
+			based_of          (str):               Assuming `path` is an augmentation of `based_of`,
+				this parameter is used to load the original dataset and
+				prevent having test_data that have augmented images in the training set
+			**kwargs          (Any):               Keyword arguments for the loading function
+				(ex for image: `keras.src.utils.image_dataset_from_directory(..., **kwargs)`)
+		Returns:
+			Dataset: Dataset object
+
+		Examples:
+			.. code-block:: python
+
+				> dataset = DatasetLoader.from_path(
+					path="data/pizza_augmented",
+					loading_type="image",
+					seed=42,
+					test_size=0.2,
+					val_size=0.2,
+					grouping_strategy=GroupingStrategy.NONE,
+					based_of="data/pizza",
+
+					# Image loading kwargs
+					color_mode="grayscale",
+					image_size=(224, 224),
+				)
+		"""
+		# Assertions
+		assert grouping_strategy in GroupingStrategy, f"Invalid grouping strategy: '{grouping_strategy.name}'"
+		assert loading_type in ("image",), f"Invalid loading type: '{loading_type}'"
+
+		# Set seed
+		np.random.seed(seed)
+
+		# Load the base dataset
+		original_dataset: Dataset = Dataset.empty()
+		if based_of:
+			original_dataset = DatasetLoader.from_path(
+				path=based_of,
+				loading_type=loading_type,
+				seed=seed,
+				test_size=test_size,
+				val_size=val_size,
+				grouping_strategy=grouping_strategy,
+				**kwargs
+			)
+
+		# Load the data
+		all_data: XyTuple = XyTuple.empty()
+		if loading_type == "image":
+			for key in DEFAULT_IMAGE_KWARGS.keys():
+				if not kwargs.get(key):
+					kwargs[key] = DEFAULT_IMAGE_KWARGS[key]
+
+			# Load the data using image_dataset_from_directory
+			# Grouping strategy can be changed by image_dataset_from_directory so we need to save it
+			all_data, all_labels, grouping_strategy = GroupingStrategy.image_dataset_from_directory(
+				grouping_strategy, path, seed, **kwargs
+			)
+
+			# Split the data using stratification
+			real_test_size: float = test_size if not based_of else 0
+			training_data, test_data = all_data.split(real_test_size, seed=DataScienceConfig.SEED)
+			training_data, val_data = training_data.split(val_size, seed=DataScienceConfig.SEED)
+
+		# Create and return the dataset
+		dataset = Dataset(
+			training_data=training_data,
+			val_data=val_data,
+			test_data=test_data,
+			name=path,
+			grouping_strategy=grouping_strategy,
+			labels=all_labels,
+			loading_type=loading_type
+		)
+
+		# If this dataset is based on another dataset, ensure test data consistency
+		if based_of:
+			dataset.exclude_augmented_images_from_val_test(original_dataset)
+
+		# Remember the original dataset
+		dataset.original_dataset = original_dataset
+
+		return dataset
+