"""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
[docs]
@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
[docs]
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
[docs]
@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
[docs]
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
[docs]
@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
[docs]
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
[docs]
@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!")
[docs]
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
[docs]
@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!")
[docs]
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
[docs]
@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
[docs]
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
[docs]
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]
[docs]
@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
[docs]
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
[docs]
@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
[docs]
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
[docs]
@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
[docs]
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