nrtk-albumentations 2.1.0__py3-none-any.whl

albumentations/augmentations/other/lambda_transform.py

@@ -0,0 +1,180 @@
+"""Lambda transform module for creating custom user-defined transformations.
+
+This module provides a flexible transform class that allows users to define their own
+custom transformation functions for different targets (image, mask, keypoints, bboxes).
+It's particularly useful for implementing custom logic that isn't available in the
+standard transforms.
+
+The Lambda transform accepts different callable functions for each target type and
+applies them when the transform is executed. This allows for maximum flexibility
+while maintaining compatibility with the Albumentations pipeline structure.
+
+Key features:
+- Apply different custom functions to different target types
+- Compatible with all Albumentations pipeline features
+- Support for all image types and formats
+- Ability to handle any number of channels
+- Warning system for lambda expressions and multiprocessing compatibility
+
+Note that using actual lambda expressions (rather than named functions) can cause
+issues with multiprocessing, as lambdas cannot be properly pickled.
+"""
+
+from __future__ import annotations
+
+import warnings
+from types import LambdaType
+from typing import Any, Callable
+
+import numpy as np
+
+from albumentations.augmentations.pixel import functional as fpixel
+from albumentations.core.transforms_interface import NoOp
+from albumentations.core.utils import format_args
+
+__all__ = ["Lambda"]
+
+
+class Lambda(NoOp):
+    """A flexible transformation class for using user-defined transformation functions per targets.
+    Function signature must include **kwargs to accept optional arguments like interpolation method, image size, etc:
+
+    Args:
+        image (Callable[..., Any] | None): Image transformation function.
+        mask (Callable[..., Any] | None): Mask transformation function.
+        keypoints (Callable[..., Any] | None): Keypoints transformation function.
+        bboxes (Callable[..., Any] | None): BBoxes transformation function.
+        p (float): probability of applying the transform. Default: 1.0.
+
+    Targets:
+        image, mask, bboxes, keypoints, volume, mask3d
+
+    Image types:
+        uint8, float32
+
+    Number of channels:
+        Any
+
+    """
+
+    def __init__(
+        self,
+        image: Callable[..., Any] | None = None,
+        mask: Callable[..., Any] | None = None,
+        keypoints: Callable[..., Any] | None = None,
+        bboxes: Callable[..., Any] | None = None,
+        name: str | None = None,
+        p: float = 1.0,
+    ):
+        super().__init__(p=p)
+
+        self.name = name
+        self.custom_apply_fns = dict.fromkeys(("image", "mask", "keypoints", "bboxes"), fpixel.noop)
+        for target_name, custom_apply_fn in {
+            "image": image,
+            "mask": mask,
+            "keypoints": keypoints,
+            "bboxes": bboxes,
+        }.items():
+            if custom_apply_fn is not None:
+                if isinstance(custom_apply_fn, LambdaType) and custom_apply_fn.__name__ == "<lambda>":
+                    warnings.warn(
+                        "Using lambda is incompatible with multiprocessing. "
+                        "Consider using regular functions or partial().",
+                        stacklevel=2,
+                    )
+
+                self.custom_apply_fns[target_name] = custom_apply_fn
+
+    def apply(self, img: np.ndarray, **params: Any) -> np.ndarray:
+        """Apply the Lambda transform to the input image.
+
+        Args:
+            img (np.ndarray): The input image to apply the Lambda transform to.
+            **params (Any): Additional parameters (not used in this transform).
+
+        Returns:
+            np.ndarray: The image with the applied Lambda transform.
+
+        """
+        fn = self.custom_apply_fns["image"]
+        return fn(img, **params)
+
+    def apply_to_mask(self, mask: np.ndarray, **params: Any) -> np.ndarray:
+        """Apply the Lambda transform to the input mask.
+
+        Args:
+            mask (np.ndarray): The input mask to apply the Lambda transform to.
+            **params (Any): Additional parameters (not used in this transform).
+
+        Returns:
+            np.ndarray: The mask with the applied Lambda transform.
+
+        """
+        fn = self.custom_apply_fns["mask"]
+        return fn(mask, **params)
+
+    def apply_to_bboxes(self, bboxes: np.ndarray, **params: Any) -> np.ndarray:
+        """Apply the Lambda transform to the input bounding boxes.
+
+        Args:
+            bboxes (np.ndarray): The input bounding boxes to apply the Lambda transform to.
+            **params (Any): Additional parameters (not used in this transform).
+
+        Returns:
+            np.ndarray: The bounding boxes with the applied Lambda transform.
+
+        """
+        fn = self.custom_apply_fns["bboxes"]
+        return fn(bboxes, **params)
+
+    def apply_to_keypoints(self, keypoints: np.ndarray, **params: Any) -> np.ndarray:
+        """Apply the Lambda transform to the input keypoints.
+
+        Args:
+            keypoints (np.ndarray): The input keypoints to apply the Lambda transform to.
+            **params (Any): Additional parameters (not used in this transform).
+
+        Returns:
+            np.ndarray: The keypoints with the applied Lambda transform.
+
+        """
+        fn = self.custom_apply_fns["keypoints"]
+        return fn(keypoints, **params)
+
+    @classmethod
+    def is_serializable(cls) -> bool:
+        """Check if the Lambda transform is serializable.
+
+        Returns:
+            bool: True if the transform is serializable, False otherwise.
+
+        """
+        return False
+
+    def to_dict_private(self) -> dict[str, Any]:
+        """Convert the Lambda transform to a dictionary.
+
+        Returns:
+            dict[str, Any]: The dictionary representation of the transform.
+
+        """
+        if self.name is None:
+            msg = (
+                "To make a Lambda transform serializable you should provide the `name` argument, "
+                "e.g. `Lambda(name='my_transform', image=<some func>, ...)`."
+            )
+            raise ValueError(msg)
+        return {"__class_fullname__": self.get_class_fullname(), "__name__": self.name}
+
+    def __repr__(self) -> str:
+        """Return the string representation of the Lambda transform.
+
+        Returns:
+            str: The string representation of the Lambda transform.
+
+        """
+        state = {"name": self.name}
+        state.update(self.custom_apply_fns.items())  # type: ignore[arg-type]
+        state.update(self.get_base_init_args())
+        return f"{self.__class__.__name__}({format_args(state)})"

albumentations/augmentations/other/type_transform.py

@@ -0,0 +1,261 @@
+"""Transforms for type conversion between float and other data types.
+
+This module provides transform classes for converting image data types, primarily
+for converting between floating point and integer representations. These transforms
+are useful for preprocessing before neural network input (ToFloat) and for converting
+network outputs back to standard image formats (FromFloat).
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import numpy as np
+from albucore import (
+    from_float,
+    get_max_value,
+    to_float,
+)
+from pydantic import (
+    model_validator,
+)
+from typing_extensions import Literal, Self
+
+from albumentations.core.transforms_interface import (
+    BaseTransformInitSchema,
+    ImageOnlyTransform,
+)
+
+__all__ = [
+    "FromFloat",
+    "ToFloat",
+]
+
+
+class ToFloat(ImageOnlyTransform):
+    """Convert the input image to a floating-point representation.
+
+    This transform divides pixel values by `max_value` to get a float32 output array
+    where all values lie in the range [0, 1.0]. It's useful for normalizing image data
+    before feeding it into neural networks or other algorithms that expect float input.
+
+    Args:
+        max_value (float | None): The maximum possible input value. If None, the transform
+            will try to infer the maximum value by inspecting the data type of the input image:
+            - uint8: 255
+            - uint16: 65535
+            - uint32: 4294967295
+            - float32: 1.0
+            Default: None.
+        p (float): Probability of applying the transform. Default: 1.0.
+
+    Targets:
+        image, volume
+
+    Image types:
+        uint8, uint16, uint32, float32
+
+    Returns:
+        np.ndarray: Image in floating point representation, with values in range [0, 1.0].
+
+    Note:
+        - If the input image is already float32 with values in [0, 1], it will be returned unchanged.
+        - For integer types (uint8, uint16, uint32), the function will scale the values to [0, 1] range.
+        - The output will always be float32, regardless of the input type.
+        - This transform is often used as a preprocessing step before applying other transformations
+          or feeding the image into a neural network.
+
+    Raises:
+        TypeError: If the input image data type is not supported.
+
+    Examples:
+        >>> import numpy as np
+        >>> import albumentations as A
+        >>>
+        # Convert uint8 image to float
+        >>> image = np.random.randint(0, 256, (100, 100, 3), dtype=np.uint8)
+        >>> transform = A.ToFloat(max_value=None)
+        >>> float_image = transform(image=image)['image']
+        >>> assert float_image.dtype == np.float32
+        >>> assert 0 <= float_image.min() <= float_image.max() <= 1.0
+        >>>
+        # Convert uint16 image to float with custom max_value
+        >>> image = np.random.randint(0, 4096, (100, 100, 3), dtype=np.uint16)
+        >>> transform = A.ToFloat(max_value=4095)
+        >>> float_image = transform(image=image)['image']
+        >>> assert float_image.dtype == np.float32
+        >>> assert 0 <= float_image.min() <= float_image.max() <= 1.0
+
+    See Also:
+        FromFloat: The inverse operation, converting from float back to the original data type.
+
+    """
+
+    class InitSchema(BaseTransformInitSchema):
+        max_value: float | None
+
+    def __init__(
+        self,
+        max_value: float | None = None,
+        p: float = 1.0,
+    ):
+        super().__init__(p=p)
+        self.max_value = max_value
+
+    def apply(self, img: np.ndarray, **params: Any) -> np.ndarray:
+        """Apply the ToFloat transform to the input image.
+
+        Args:
+            img (np.ndarray): The input image to apply the ToFloat transform to.
+            **params (Any): Additional parameters (not used in this transform).
+
+        Returns:
+            np.ndarray: The image with the applied ToFloat transform.
+
+        """
+        return to_float(img, self.max_value)
+
+    def apply_to_images(self, images: np.ndarray, **params: Any) -> np.ndarray:
+        """Apply the ToFloat transform to the input batch of images.
+
+        Args:
+            images (np.ndarray): The input batch of images to apply the ToFloat transform to.
+            **params (Any): Additional parameters (not used in this transform).
+
+        Returns:
+            np.ndarray: The batch of images with the applied ToFloat transform.
+
+        """
+        return to_float(images, self.max_value)
+
+    def apply_to_volume(self, volume: np.ndarray, **params: Any) -> np.ndarray:
+        """Apply the ToFloat transform to the input volume.
+
+        Args:
+            volume (np.ndarray): The input volume to apply the ToFloat transform to.
+            **params (Any): Additional parameters (not used in this transform).
+
+        Returns:
+            np.ndarray: The volume with the applied ToFloat transform.
+
+        """
+        return self.apply_to_images(volume, **params)
+
+    def apply_to_volumes(self, volumes: np.ndarray, **params: Any) -> np.ndarray:
+        """Apply the ToFloat transform to the input volumes.
+
+        Args:
+            volumes (np.ndarray): The input volumes to apply the ToFloat transform to.
+            **params (Any): Additional parameters (not used in this transform).
+
+        Returns:
+            np.ndarray: The volumes with the applied ToFloat transform.
+
+        """
+        return self.apply_to_images(volumes, **params)
+
+
+class FromFloat(ImageOnlyTransform):
+    """Convert an image from floating point representation to the specified data type.
+
+    This transform is designed to convert images from a normalized floating-point representation
+    (typically with values in the range [0, 1]) to other data types, scaling the values appropriately.
+
+    Args:
+        dtype (str): The desired output data type. Supported types include 'uint8', 'uint16',
+                     'uint32'. Default: 'uint8'.
+        max_value (float | None): The maximum value for the output dtype. If None, the transform
+                                  will attempt to infer the maximum value based on the dtype.
+                                  Default: None.
+        p (float): Probability of applying the transform. Default: 1.0.
+
+    Targets:
+        image, volume
+
+    Image types:
+        float32, float64
+
+    Note:
+        - This is the inverse transform for ToFloat.
+        - Input images are expected to be in floating point format with values in the range [0, 1].
+        - For integer output types (uint8, uint16, uint32), the function will scale the values
+          to the appropriate range (e.g., 0-255 for uint8).
+        - For float output types (float32, float64), the values will remain in the [0, 1] range.
+        - The transform uses the `from_float` function internally, which ensures output values
+          are within the valid range for the specified dtype.
+
+    Examples:
+        >>> import numpy as np
+        >>> import albumentations as A
+        >>> transform = A.FromFloat(dtype='uint8', max_value=None, p=1.0)
+        >>> image = np.random.rand(100, 100, 3).astype(np.float32)  # Float image in [0, 1] range
+        >>> result = transform(image=image)
+        >>> uint8_image = result['image']
+        >>> assert uint8_image.dtype == np.uint8
+        >>> assert uint8_image.min() >= 0 and uint8_image.max() <= 255
+
+    """
+
+    class InitSchema(BaseTransformInitSchema):
+        dtype: Literal["uint8", "uint16", "uint32"]
+        max_value: float | None
+
+        @model_validator(mode="after")
+        def _update_max_value(self) -> Self:
+            if self.max_value is None:
+                self.max_value = get_max_value(np.dtype(self.dtype))
+
+            return self
+
+    def __init__(
+        self,
+        dtype: Literal["uint8", "uint16", "uint32"] = "uint8",
+        max_value: float | None = None,
+        p: float = 1.0,
+    ):
+        super().__init__(p=p)
+        self.dtype = dtype
+        self.max_value = max_value
+
+    def apply(self, img: np.ndarray, **params: Any) -> np.ndarray:
+        """Apply the FromFloat transform to the input image.
+
+        Args:
+            img (np.ndarray): The input image to apply the FromFloat transform to.
+            **params (Any): Additional parameters (not used in this transform).
+
+        Returns:
+            np.ndarray: The image with the applied FromFloat transform.
+
+        """
+        return from_float(img, np.dtype(self.dtype), self.max_value)
+
+    def apply_to_images(self, images: np.ndarray, **params: Any) -> np.ndarray:
+        """Apply the FromFloat transform to the input images.
+
+        Args:
+            images (np.ndarray): The input images to apply the FromFloat transform to.
+            **params (Any): Additional parameters (not used in this transform).
+
+        """
+        return from_float(images, np.dtype(self.dtype), self.max_value)
+
+    def apply_to_volume(self, volume: np.ndarray, **params: Any) -> np.ndarray:
+        """Apply the FromFloat transform to the input volume.
+
+        Args:
+            volume (np.ndarray): The input volume to apply the FromFloat transform to.
+            **params (Any): Additional parameters (not used in this transform).
+
+        """
+        return self.apply_to_images(volume, **params)
+
+    def apply_to_volumes(self, volumes: np.ndarray, **params: Any) -> np.ndarray:
+        """Apply the FromFloat transform to the input volumes.
+
+        Args:
+            volumes (np.ndarray): The input volumes to apply the FromFloat transform to.
+            **params (Any): Additional parameters (not used in this transform).
+
+        """
+        return self.apply_to_images(volumes, **params)