Source code for robotblockset.cameras.image_transform

"""Composable transforms for image arrays and pixel coordinates.

The transforms support two-dimensional grayscale images and channel-last color
images. Each transform can be applied both to image data and to corresponding
``(x, y)`` coordinates.
"""

from abc import ABC
from typing import Sequence, Tuple, TypeAlias, Union

import cv2
import numpy as np

from robotblockset.rbs_typing import NumpyFloatImageType, NumpyIntImageType, OpenCVIntImageType

ImageArrayType: TypeAlias = Union[OpenCVIntImageType, NumpyFloatImageType, NumpyIntImageType]
"""Grayscale ``(H, W)`` or channel-last color ``(H, W, C)`` image array."""

# Retained for compatibility with code importing the former public alias.
HWCImageType: TypeAlias = ImageArrayType
"""Compatibility alias for :data:`ImageArrayType`."""

ImageShapeType: TypeAlias = Union[Tuple[int, int], Tuple[int, int, int]]
ImageCoordinateType: TypeAlias = Union[int, float]
ImagePointType: TypeAlias = Tuple[ImageCoordinateType, ImageCoordinateType]




class ImageTransform(ABC):
    """Base interface for related image and point-coordinate transforms."""



    def __init__(self, input_shape: ImageShapeType) -> None:
        """Initialize a transform for images with a known input shape.

        Parameters
        ----------
        input_shape : ImageShapeType
            Grayscale ``(height, width)`` or color
            ``(height, width, channels)`` input shape.
        """
        self._input_shape = input_shape


    @property
    def _input_h(self) -> int:
        """Input image height in pixels."""
        return self._input_shape[0]

    @property
    def _input_w(self) -> int:
        """Input image width in pixels."""
        return self._input_shape[1]

    @property
    def shape(self) -> ImageShapeType:
        """Return the shape of the transformed image.

        Returns
        -------
        ImageShapeType
            Grayscale or channel-last output shape.

        Raises
        ------
        NotImplementedError
            If a subclass does not implement the property.
        """
        raise NotImplementedError



    def transform_image(self, image: ImageArrayType) -> ImageArrayType:
        """Apply the transform to an image array.

        Parameters
        ----------
        image : ImageArrayType
            Source image. Implementations return a transformed image without
            intentionally modifying this array.

        Returns
        -------
        ImageArrayType
            Transformed image array.

        Raises
        ------
        NotImplementedError
            If a subclass does not implement the method.
        """
        raise NotImplementedError




    def transform_point(self, point: ImagePointType) -> ImagePointType:
        """Map a source-image point into transformed-image coordinates.

        Parameters
        ----------
        point : ImagePointType
            Source ``(x, y)`` coordinate.

        Returns
        -------
        ImagePointType
            Coordinate in the transformed image.

        Raises
        ------
        NotImplementedError
            If a subclass does not implement the method.
        """
        raise NotImplementedError




    def reverse_transform_point(self, point: ImagePointType) -> ImagePointType:
        """Map a transformed-image point back to source-image coordinates.

        Parameters
        ----------
        point : ImagePointType
            Transformed ``(x, y)`` coordinate.

        Returns
        -------
        ImagePointType
            Coordinate in the source image.

        Raises
        ------
        NotImplementedError
            If a subclass does not implement the method.
        """
        raise NotImplementedError


    def __call__(self, image: ImageArrayType) -> ImageArrayType:
        """Apply :meth:`transform_image` to ``image``."""
        return self.transform_image(image)





class ComposedTransform(ImageTransform):
    """Apply a non-empty sequence of image transforms in order."""



    def __init__(self, transforms: Sequence[ImageTransform]) -> None:
        """Create a transform composition.

        Parameters
        ----------
        transforms : Sequence[ImageTransform]
            Ordered, non-empty sequence of transforms.

        Raises
        ------
        ValueError
            If ``transforms`` is empty.
        """
        if len(transforms) == 0:
            raise ValueError("transforms must be a non-empty sequence.")

        super().__init__(transforms[0]._input_shape)
        self.transforms = transforms


    @property
    def shape(self) -> ImageShapeType:
        """Return the output shape of the final transform."""
        return self.transforms[-1].shape



    def transform_image(self, image: ImageArrayType) -> ImageArrayType:
        """Apply every transform to an image in sequence."""
        for transform in self.transforms:
            image = transform.transform_image(image)
        return image




    def transform_point(self, point: ImagePointType) -> ImagePointType:
        """Map a point through every transform in sequence."""
        for transform in self.transforms:
            point = transform.transform_point(point)
            print(point)
        return point




    def reverse_transform_point(self, point: ImagePointType) -> ImagePointType:
        """Map a point backward through the transforms in reverse order."""
        for transform in reversed(self.transforms):
            point = transform.reverse_transform_point(point)
        return point






def crop(image: ImageArrayType, x: int, y: int, w: int, h: int) -> ImageArrayType:
    """Return a copy of a rectangular image region.

    Parameters
    ----------
    image : ImageArrayType
        Grayscale or channel-last source image.
    x : int
        X-coordinate of the top-left crop corner.
    y : int
        Y-coordinate of the top-left crop corner.
    w : int
        Crop width in pixels.
    h : int
        Crop height in pixels.

    Returns
    -------
    ImageArrayType
        Copy of the selected region. Coordinates outside the image follow NumPy
        slicing semantics.
    """
    # The first array index is y because it selects rows, while the second is x.
    if len(image.shape) == 2:
        return image[y : y + h, x : x + w].copy()

    return image[y : y + h, x : x + w, :].copy()





class Crop(ImageTransform):
    """Crop an image and translate corresponding point coordinates."""



    def __init__(self, input_shape: ImageShapeType, x: int, y: int, w: int, h: int) -> None:
        """Create a rectangular crop transform.

        Parameters
        ----------
        input_shape : ImageShapeType
            Shape of the source image.
        x : int
            X-coordinate of the crop's top-left corner.
        y : int
            Y-coordinate of the crop's top-left corner.
        w : int
            Crop width in pixels.
        h : int
            Crop height in pixels.
        """
        super().__init__(input_shape)
        self.x = x
        self.y = y
        self.w = w
        self.h = h


    @property
    def shape(self) -> ImageShapeType:
        """Return the declared crop shape."""
        if len(self._input_shape) == 2:
            return self.h, self.w

        c = self._input_shape[2]
        return self.h, self.w, c



    def transform_image(self, image: ImageArrayType) -> ImageArrayType:
        """Return a copy of the cropped image region."""
        return crop(image, self.x, self.y, self.w, self.h)




    def transform_point(self, point: ImagePointType) -> ImagePointType:
        """Translate a source point into crop-local coordinates.

        Raises
        ------
        ValueError
            If the source point lies outside the crop rectangle.
        """
        x, y = point
        if not (x >= self.x and x < self.x + self.w):
            raise ValueError(f"x-coordinate {x} is outside of the crop range [{self.x}, {self.x + self.w})")
        if not (y >= self.y and y < self.y + self.h):
            raise ValueError(f"y-coordinate {y} is outside of the crop range [{self.y}, {self.y + self.h})")
        return x - self.x, y - self.y




    def reverse_transform_point(self, point: ImagePointType) -> ImagePointType:
        """Translate a crop-local point into source-image coordinates.

        Raises
        ------
        ValueError
            If the point lies outside the crop dimensions.
        """
        x, y = point
        if not (x >= 0 and x < self.w):
            raise ValueError(f"x-coordinate {x} is outside of the crop range [0, {self.w})")
        if not (y >= 0 and y < self.h):
            raise ValueError(f"y-coordinate {y} is outside of the crop range [0, {self.h})")
        return x + self.x, y + self.y






class Resize(ImageTransform):
    """Resize an image and scale corresponding point coordinates."""



    def __init__(self, input_shape: ImageShapeType, h: int, w: int, round_transformed_points: bool = True) -> None:
        """Create an image resize transform.

        Transforming points to or from a resized image can produce fractional
        coordinates. By default these coordinates are rounded to the nearest
        integer. Set ``round_transformed_points`` to ``False`` to preserve the
        floating-point result.

        Parameters
        ----------
        input_shape : ImageShapeType
            Shape of the source image.
        h : int
            Output height in pixels.
        w : int
            Output width in pixels.
        round_transformed_points : bool, optional
            Whether point transformations return rounded coordinates. The
            default is ``True``.
        """
        super().__init__(input_shape)
        self.h = h
        self.w = w
        self.round_transformed_points = round_transformed_points


    @property
    def shape(self) -> ImageShapeType:
        """Return the resized image shape."""
        if len(self._input_shape) == 2:
            return self.h, self.w

        c = self._input_shape[2]
        return self.h, self.w, c



    def transform_image(self, image: ImageArrayType) -> ImageArrayType:
        """Resize an image using OpenCV's default interpolation."""
        return cv2.resize(image, (self.w, self.h))




    def transform_point(self, point: ImagePointType) -> ImagePointType:
        """Scale a source point into resized-image coordinates.

        Raises
        ------
        ValueError
            If the point lies outside the source image dimensions.
        """
        x, y = point
        if not (x >= 0 and x < self._input_w):
            raise ValueError(f"x-coordinate {x} is outside of the input image range [0, {self._input_w})")
        if not (y >= 0 and y < self._input_h):
            raise ValueError(f"y-coordinate {y} is outside of the input image range [0, {self._input_h})")

        w_scale = self.w / self._input_w
        h_scale = self.h / self._input_h

        x_float = w_scale * x
        y_float = h_scale * y

        if self.round_transformed_points:
            return round(x_float), round(y_float)

        return x_float, y_float




    def reverse_transform_point(self, point: ImagePointType) -> ImagePointType:
        """Scale a resized-image point back into source coordinates.

        Raises
        ------
        ValueError
            If the point lies outside the resized image dimensions.
        """
        x, y = point
        if not (x >= 0 and x < self.w):
            raise ValueError(f"x-coordinate {x} is outside of the resized image range [0, {self.w})")
        if not (y >= 0 and y < self.h):
            raise ValueError(f"y-coordinate {y} is outside of the resized image range [0, {self.h})")
        w_scale_inverse = self._input_w / self.w
        h_scale_inverse = self._input_h / self.h

        x_float = w_scale_inverse * x
        y_float = h_scale_inverse * y

        if self.round_transformed_points:
            return round(x_float), round(y_float)

        return x_float, y_float






class Rotate90(ImageTransform):
    """Rotate images and corresponding points by multiples of 90 degrees."""



    def __init__(
        self,
        input_shape: ImageShapeType,
        num_rotations: int = 1,
    ) -> None:
        """Create a counter-clockwise right-angle rotation transform.

        Parameters
        ----------
        input_shape : ImageShapeType
            Shape of the source image.
        num_rotations : int, optional
            Number of counter-clockwise 90-degree rotations. Values are reduced
            modulo four. The default is ``1``.

        Raises
        ------
        TypeError
            If ``num_rotations`` is not an integer.
        """
        super().__init__(input_shape)

        if not isinstance(num_rotations, int):
            raise TypeError("num_rotations must be an int")

        self._num_rotations = num_rotations % 4


    @property
    def shape(self) -> ImageShapeType:
        """Return the rotated image shape."""
        if self._num_rotations % 2 == 0:
            h, w = self._input_shape[:2]
        else:
            w, h = self._input_shape[:2]

        if len(self._input_shape) == 2:
            return h, w

        c = self._input_shape[2]
        return h, w, c



    def transform_image(self, image: ImageArrayType) -> ImageArrayType:
        """Return a copied array rotated counter-clockwise by 90-degree steps."""
        # The copy here ensure the result is not a view into the original image.
        return np.rot90(image, self._num_rotations).copy()




    def transform_point(self, point: ImagePointType) -> ImagePointType:
        """Rotate a source point into output-image coordinates.

        Raises
        ------
        ValueError
            If the point lies outside the source image dimensions.
        """
        x, y = point
        if not (x >= 0 and x < self._input_w):
            raise ValueError(f"x-coordinate {x} is outside of the input image range [0, {self._input_w})")
        if not (y >= 0 and y < self._input_h):
            raise ValueError(f"y-coordinate {y} is outside of the input image range [0, {self._input_h})")

        if self._num_rotations == 1:
            return y, self._input_w - x - 1
        elif self._num_rotations == 2:
            return self._input_w - x - 1, self._input_h - y - 1
        elif self._num_rotations == 3:
            return self._input_h - y - 1, x
        return x, y




    def reverse_transform_point(self, point: ImagePointType) -> ImagePointType:
        """Rotate an output-image point back into source coordinates."""
        x, y = point
        if self._num_rotations == 1:
            return self._input_w - y - 1, x
        elif self._num_rotations == 2:
            return self._input_w - x - 1, self._input_h - y - 1
        elif self._num_rotations == 3:
            return y, self._input_h - x - 1
        return x, y