vis4d/data/transforms/flip.py

"""Horizontal flip augmentation."""

import numpy as np
import torch

from vis4d.common.typing import NDArrayF32, NDArrayUI8
from vis4d.data.const import AxisMode
from vis4d.data.const import CommonKeys as K
from vis4d.op.geometry.rotation import (
    euler_angles_to_matrix,
    matrix_to_euler_angles,
    matrix_to_quaternion,
    quaternion_to_matrix,
)

from .base import Transform


@Transform(K.images, K.images)
class FlipImages:
    """Flip a list of numpy image array of shape [N, H, W, C]."""

    def __init__(self, direction: str = "horizontal"):
        """Creates an instance of FlipImage.

        Args:
            direction (str, optional): Either vertical or horizontal.
                Defaults to "horizontal".

        Raises:
            ValueError: If direction is not horizontal or vertical.
        """
        if direction not in ["horizontal", "vertical"]:
            raise ValueError(f"Direction {direction} not known!")
        self.direction = direction

    def __call__(self, images: list[NDArrayF32]) -> list[NDArrayF32]:
        """Execute flipping op.

        Args:
            image (NDArrayF32): [N, H, W, C] array of image.

        Returns:
            list[NDArrayF32]: [N, H, W, C] array of flipped image.
        """
        for i, image in enumerate(images):
            image_ = torch.from_numpy(image)
            if self.direction == "horizontal":
                images[i] = image_.flip(2).numpy()
            if self.direction == "vertical":
                images[i] = image_.flip(1).numpy()
        return images


@Transform(in_keys=(K.boxes2d, K.images), out_keys=(K.boxes2d,))
class FlipBoxes2D:
    """Flip a list of 2D bounding boxes."""

    def __init__(self, direction: str = "horizontal"):
        """Creates an instance of FlipBoxes2D.

        Args:
            direction (str, optional): Either vertical or horizontal.
                Defaults to "horizontal".

        Raises:
            ValueError: If direction is not horizontal or vertical.
        """
        if direction not in ["horizontal", "vertical"]:
            raise ValueError(f"Direction {direction} not known!")
        self.direction = direction

    def __call__(
        self, boxes_list: list[NDArrayF32], images: list[NDArrayF32]
    ) -> list[NDArrayF32]:
        """Execute flipping op.

        Args:
            boxes (list[NDArrayF32]): List of [M, 4] array of boxes.
            image (list[NDArrayF32]): List of [N, H, W, C] array of image.

        Returns:
            list[NDArrayF32]: List of [M, 4] array of flipped boxes.
        """
        for i, (boxes, image) in enumerate(zip(boxes_list, images)):
            if self.direction == "horizontal":
                im_width = image.shape[2]
                tmp = im_width - boxes[..., 2::4]
                boxes[..., 2::4] = im_width - boxes[..., 0::4]
                boxes[..., 0::4] = tmp
            elif self.direction == "vertical":
                im_height = image.shape[1]
                tmp = im_height - boxes[..., 3::4]
                boxes[..., 3::4] = im_height - boxes[..., 1::4]
                boxes[..., 1::4] = tmp
            boxes_list[i] = boxes
        return boxes_list


@Transform(K.seg_masks, K.seg_masks)
class FlipSegMasks:
    """Flip segmentation masks."""

    def __init__(self, direction: str = "horizontal"):
        """Creates an instance of FlipSemanticMasks.

        Args:
            direction (str, optional): Either vertical or horizontal.
                Defaults to "horizontal".

        Raises:
            ValueError: If direction is not horizontal or vertical.
        """
        if direction not in ["horizontal", "vertical"]:
            raise ValueError(f"Direction {direction} not known!")
        self.direction = direction

    def __call__(self, masks: list[NDArrayUI8]) -> list[NDArrayUI8]:
        """Execute flipping op.

        Args:
            masks (NDArrayUI8): [H, W] array of masks.

        Returns:
            list[NDArrayUI8]: [H, W] array of flipped masks.
        """
        for i, mask in enumerate(masks):
            mask_ = torch.from_numpy(mask)
            if self.direction == "horizontal":
                mask = mask_.flip(1).numpy()
            if self.direction == "vertical":
                mask = mask_.flip(0).numpy()
            masks[i] = mask
        return masks


@Transform(K.depth_maps, K.depth_maps)
class FlipDepthMaps:
    """Flip depth map."""

    def __init__(self, direction: str = "horizontal"):
        """Creates an instance of FlipDepth.

        Args:
            direction (str, optional): Either vertical or horizontal.
                Defaults to "horizontal".
        """
        self.direction = direction
        if direction not in ["horizontal", "vertical"]:
            raise ValueError(f"Direction {self.direction} not known!")

    def __call__(self, depths: list[NDArrayF32]) -> list[NDArrayF32]:
        """Execute flipping op.

        Args:
            depths (list[NDArrayF32]): Each is a [H, W] array of depth.

        Returns:
            list[NDArrayF32]: Each is a [H, W] array of flipped depth.
        """
        for i, depth in enumerate(depths):
            depth_ = torch.from_numpy(depth)
            if self.direction == "horizontal":
                depths[i] = depth_.flip(1).numpy()
            if self.direction == "vertical":
                depths[i] = depth_.flip(0).numpy()

        return depths


@Transform(K.optical_flows, K.optical_flows)
class FlipOpticalFlows:
    """Flip optical flow map."""

    def __init__(self, direction: str = "horizontal"):
        """Creates an instance of FlipOpticalFlow.

        Args:
            direction (str, optional): Either vertical or horizontal.
                Defaults to "horizontal".
        """
        self.direction = direction
        if direction not in ["horizontal", "vertical"]:
            raise ValueError(f"Direction {self.direction} not known!")

    def __call__(self, flows: list[NDArrayF32]) -> list[NDArrayF32]:
        """Execute flipping op.

        Args:
            flows (NDArrayF32): Each is a [H, W, 2] array of optical flow.

        Returns:
            list[NDArrayF32]: Each is a [H, W, 2] array of flipped optical
                flow.
        """
        for i, flow in enumerate(flows):
            flow_ = torch.from_numpy(flow)
            if self.direction == "horizontal":
                image_flipped = flow_.flip(1).numpy()
                image_flipped[..., 0] *= -1
                flows[i] = image_flipped
            if self.direction == "vertical":
                image_flipped = flow_.flip(0).numpy()
                image_flipped[..., 1] *= -1
                flows[i] = image_flipped
        return flows


@Transform(K.instance_masks, K.instance_masks)
class FlipInstanceMasks:
    """Flip instance masks."""

    def __init__(self, direction: str = "horizontal"):
        """Creates an instance of FlipInstanceMasks.

        Args:
            direction (str, optional): Either vertical or horizontal.
                Defaults to "horizontal".

        Raises:
            ValueError: If direction is not horizontal or vertical.
        """
        if direction not in ["horizontal", "vertical"]:
            raise ValueError(f"Direction {direction} not known!")
        self.direction = direction

    def __call__(self, masks: list[NDArrayUI8]) -> list[NDArrayUI8]:
        """Execute flipping op.

        Args:
            masks (list[NDArrayUI8]): List of [N, H, W] array of masks.

        Returns:
            list[NDArrayUI8]: List of [N, H, W] array of flipped masks.
        """
        for i, mask in enumerate(masks):
            mask_ = torch.from_numpy(mask)
            if self.direction == "horizontal":
                mask = mask_.flip(2).numpy()
            if self.direction == "vertical":
                mask = mask_.flip(1).numpy()
            masks[i] = mask
        return masks


def get_axis(direction: str, axis_mode: AxisMode) -> int:
    """Get axis number of certain direction given axis mode.

    Args:
        direction (str): One of horizontal, vertical and lateral.
        axis_mode (AxisMode): axis mode.

    Returns:
        int: Number of axis in certain direction.
    """
    if direction not in {"horizontal", "lateral", "vertical"}:
        raise ValueError(f"Direction {direction} not known!")
    coord_mapping = {
        AxisMode.ROS: {"horizontal": 0, "lateral": 1, "vertical": 2},
        AxisMode.OPENCV: {"horizontal": 0, "vertical": 1, "lateral": 2},
    }
    return coord_mapping[axis_mode][direction]


@Transform(in_keys=(K.boxes3d, K.axis_mode), out_keys=(K.boxes3d,))
class FlipBoxes3D:
    """Flip 3D bounding box array."""

    def __init__(self, direction: str = "horizontal"):
        """Creates an instance of FlipBoxes3D.

        Args:
            direction (str, optional): Either vertical or horizontal.
                Defaults to "horizontal".
        """
        self.direction = direction

    def __call__(
        self, boxes_list: list[NDArrayF32], axis_mode_list: list[AxisMode]
    ) -> list[NDArrayF32]:
        """Execute flipping."""
        for i, (boxes, axis_mode) in enumerate(
            zip(boxes_list, axis_mode_list)
        ):
            axis = get_axis(self.direction, axis_mode)
            angle_dir = (
                "vertical" if self.direction == "horizontal" else "lateral"
            )
            angles_axis = get_axis(angle_dir, axis_mode)
            boxes[:, axis] *= -1.0
            angles = matrix_to_euler_angles(
                quaternion_to_matrix(torch.from_numpy(boxes[:, 6:]))
            )
            angles[:, angles_axis] = np.pi - angles[:, angles_axis]
            boxes[:, 6:] = matrix_to_quaternion(
                euler_angles_to_matrix(angles)
            ).numpy()

            boxes_list[i] = boxes

        return boxes_list


@Transform(in_keys=(K.points3d, K.axis_mode), out_keys=(K.points3d,))
class FlipPoints3D:
    """Flip pointcloud array."""

    def __init__(self, direction: str = "horizontal"):
        """Creates an instance of FlipBoxes2D.

        Args:
            direction (str, optional): Either vertical or horizontal.
                Defaults to "horizontal".
        """
        self.direction = direction

    def __call__(
        self, points3d_list: list[NDArrayF32], axis_mode_list: list[AxisMode]
    ) -> list[NDArrayF32]:
        """Execute flipping."""
        for i, (points3d, axis_mode) in enumerate(
            zip(points3d_list, axis_mode_list)
        ):
            points3d[:, get_axis(self.direction, axis_mode)] *= -1.0
            points3d_list[i] = points3d
        return points3d_list


@Transform(in_keys=(K.intrinsics, K.images), out_keys=(K.intrinsics,))
class FlipIntrinsics:
    """Modify intrinsics for image flip."""

    def __init__(self, direction: str = "horizontal"):
        """Creates an instance of FlipIntrinsics.

        Args:
            direction (str, optional): Either vertical or horizontal.
                Defaults to "horizontal".

        Raises:
            ValueError: If direction is not horizontal or vertical.
        """
        if direction not in ["horizontal", "vertical"]:
            raise ValueError(f"Direction {direction} not known!")
        self.direction = direction

    def __call__(
        self, intrinsics_list: list[NDArrayF32], images: list[NDArrayF32]
    ) -> list[NDArrayF32]:
        """Execute flipping."""
        for i, (intrinsics, image) in enumerate(zip(intrinsics_list, images)):
            if self.direction == "horizontal":
                center = image.shape[2] / 2
                intrinsics[0, 2] = center - intrinsics[0, 2] + center
            elif self.direction == "vertical":
                center = image.shape[1] / 2
                intrinsics[1, 2] = center - intrinsics[1, 2] + center
            intrinsics_list[i] = intrinsics
        return intrinsics_list