from __future__ import annotations import functools from functools import wraps from typing import TYPE_CHECKING, Any, Callable, TypeVar, cast import cv2 import numpy as np from albucore.utils import ( is_grayscale_image, is_multispectral_image, is_rgb_image, ) from typing_extensions import Concatenate, ParamSpec from albumentations.core.keypoints_utils import angle_to_2pi_range if TYPE_CHECKING: from pathlib import Path __all__ = [ "read_bgr_image", "read_rgb_image", "read_grayscale", "angle_2pi_range", "non_rgb_error", ] P = ParamSpec("P") T = TypeVar("T", bound=np.ndarray) F = TypeVar("F", bound=Callable[..., Any]) def read_bgr_image(path: str | Path) -> np.ndarray: return cv2.imread(str(path), cv2.IMREAD_COLOR) def read_rgb_image(path: str | Path) -> np.ndarray: image = read_bgr_image(path) return cv2.cvtColor(image, cv2.COLOR_BGR2RGB) def read_grayscale(path: str | Path) -> np.ndarray: return cv2.imread(str(path), cv2.IMREAD_GRAYSCALE) def angle_2pi_range( func: Callable[Concatenate[np.ndarray, P], np.ndarray], ) -> Callable[Concatenate[np.ndarray, P], np.ndarray]: @wraps(func) def wrapped_function(keypoints: np.ndarray, *args: P.args, **kwargs: P.kwargs) -> np.ndarray: result = func(keypoints, *args, **kwargs) if len(result) > 0 and result.shape[1] > 2: # noqa: PLR2004 result[:, 2] = angle_to_2pi_range(result[:, 2]) return result return wrapped_function def non_rgb_error(image: np.ndarray) -> None: """Check if the input image is RGB and raise a ValueError if it's not. This function is used to ensure that certain transformations are only applied to RGB images. It provides helpful error messages for grayscale and multi-spectral images. Args: image (np.ndarray): The input image to check. Expected to be a numpy array representing an image. Raises: ValueError: If the input image is not an RGB image (i.e., does not have exactly 3 channels). The error message includes specific instructions for grayscale images and a note about incompatibility with multi-spectral images. Note: - RGB images are expected to have exactly 3 channels. - Grayscale images (1 channel) will trigger an error with conversion instructions. - Multi-spectral images (more than 3 channels) will trigger an error stating incompatibility. Example: >>> import numpy as np >>> rgb_image = np.random.randint(0, 256, (100, 100, 3), dtype=np.uint8) >>> non_rgb_error(rgb_image) # No error raised >>> >>> grayscale_image = np.random.randint(0, 256, (100, 100), dtype=np.uint8) >>> non_rgb_error(grayscale_image) # Raises ValueError with conversion instructions >>> >>> multispectral_image = np.random.randint(0, 256, (100, 100, 5), dtype=np.uint8) >>> non_rgb_error(multispectral_image) # Raises ValueError stating incompatibility """ if not is_rgb_image(image): message = "This transformation expects 3-channel images" if is_grayscale_image(image): message += "\nYou can convert your grayscale image to RGB using cv2.cvtColor(image, cv2.COLOR_GRAY2RGB))" if is_multispectral_image(image): # Any image with a number of channels other than 1 and 3 message += "\nThis transformation cannot be applied to multi-spectral images" raise ValueError(message) def check_range(value: tuple[float, float], lower_bound: float, upper_bound: float, name: str | None) -> None: """Checks if the given value is within the specified bounds Args: value: The value to check and convert. Can be a single float or a tuple of floats. lower_bound: The lower bound for the range check. upper_bound: The upper bound for the range check. name: The name of the parameter being checked. Used for error messages. Raises: ValueError: If the value is outside the bounds or if the tuple values are not ordered correctly. """ if not all(lower_bound <= x <= upper_bound for x in value): raise ValueError(f"All values in {name} must be within [{lower_bound}, {upper_bound}] for tuple inputs.") if not value[0] <= value[1]: raise ValueError(f"{name!s} tuple values must be ordered as (min, max). Got: {value}") class PCA: def __init__(self, n_components: int | None = None) -> None: if n_components is not None and n_components <= 0: raise ValueError("Number of components must be greater than zero.") self.n_components = n_components self.mean: np.ndarray | None = None self.components_: np.ndarray | None = None self.explained_variance_: np.ndarray | None = None def fit(self, x: np.ndarray) -> None: x = x.astype(np.float64) n_samples, n_features = x.shape # Determine the number of components if not set if self.n_components is None: self.n_components = min(n_samples, n_features) self.mean, eigenvectors, eigenvalues = cv2.PCACompute2(x, mean=None, maxComponents=self.n_components) self.components_ = eigenvectors self.explained_variance_ = eigenvalues.flatten() def transform(self, x: np.ndarray) -> np.ndarray: if self.components_ is None: raise ValueError( "This PCA instance is not fitted yet. " "Call 'fit' with appropriate arguments before using this estimator.", ) x = x.astype(np.float64) return cv2.PCAProject(x, self.mean, self.components_) def fit_transform(self, x: np.ndarray) -> np.ndarray: self.fit(x) return self.transform(x) def inverse_transform(self, x: np.ndarray) -> np.ndarray: if self.components_ is None: raise ValueError( "This PCA instance is not fitted yet. " "Call 'fit' with appropriate arguments before using this estimator.", ) return cv2.PCABackProject(x, self.mean, self.components_) def explained_variance_ratio(self) -> np.ndarray: if self.explained_variance_ is None: raise ValueError( "This PCA instance is not fitted yet. " "Call 'fit' with appropriate arguments before using this method.", ) total_variance = np.sum(self.explained_variance_) return self.explained_variance_ / total_variance def cumulative_explained_variance_ratio(self) -> np.ndarray: return np.cumsum(self.explained_variance_ratio()) def handle_empty_array(func: F) -> F: @functools.wraps(func) def wrapper(array: T, *args: Any, **kwargs: Any) -> Any: if len(array) == 0: return array return func(array, *args, **kwargs) return cast(F, wrapper)