Source code for utils.geometry

from typing import Tuple, List
import numpy as np
import numpy.typing as npt
from matplotlib.path import Path


def rotate_and_move(xy: npt.NDArray, rotation: float = 0.,
                    offset_before_rotation: npt.NDArray = np.array([0, 0]),
                    offset_before_rotation_lidar: npt.NDArray = np.array([0, 0]),
                    is_lidar_index: List = None,
                    offset_after_rotation: npt.NDArray = np.array([0, 0])) -> npt.NDArray:
    """
    Function to move and rotate given coordinates.

    First moves the coordinates, then rotates them around the origin and then move them again.
    Using a numpy built rotation matrix.

    Parameters
    ----------
    xy : npt.NDArray
        Coordinates to move, rotate and move again.
    rotation : float, optional
        Rotation angle to rotate the coordinates around the origin, by default 0.
    offset_before_rotation : npt.NDArray, optional
        Offset (Camera) to move the coordinates before rotating them, by default np.array([0, 0]).
    offset_before_rotation_lidar : npt.NDArray, optional
        Offset (LiDAR) to move the coordinates before rotating them, by default np.array([0, 0]).
    is_lidar_index : List, optional
        List of booleans to indicate if the point is from the lidar or camera, by default None.
    offset_after_rotation : npt.NDArray, optional
        Offset to move the coordinates after rotating them, by default np.array([0, 0]).

    Returns
    -------
    npt.NDArray
        Moved, rotated and moved coordinates.
    """
    # Get correct offset for points with the information if they are from the lidar or not
    if is_lidar_index is not None:
        offset_before_rotation_apply = [offset_before_rotation_lidar if x == True else offset_before_rotation for x in is_lidar_index]
        offset_before_rotation_apply = np.array(offset_before_rotation_apply)
    else:
        offset_before_rotation_apply = offset_before_rotation

    xy += offset_before_rotation_apply
    c, s = np.cos(-rotation), np.sin(-rotation)
    j = np.matrix([[c, s], [-s, c]])
    m = np.dot(j, [xy[:, 0], xy[:, 1]])
    points = np.array(m).T
    points += offset_after_rotation
    return points


def tranform_cartesian_to_polar_cone_coords(cone_coordinates: npt.NDArray) -> npt.NDArray:
    """
    Transforms 2d cartesian coords to polar coords.

    Parameters
    ----------
    cone_coordinates : npt.NDArray
        2d position of cones in cartesian coordinates.

    Returns
    -------
    npt.NDArray
        2d position of cones in polar coordinates.
    """
    return np.array([np.arctan2(cone_coordinates[:, 1], cone_coordinates[:, 0]),
                     np.hypot(cone_coordinates[:, 1], cone_coordinates[:, 0])]).T


def generate_fov_path(fov_min_distance: float = 0.,
                      fov_max_distance: float = 12.,
                      fov_angle: float = np.radians(100),
                      fov_scaling: float = 1.0,
                      offset_after_rotation: npt.NDArray[np.floating] = np.array([0, 0]),
                      offset_before_rotation: npt.NDArray[np.floating] = np.array([0, 0]),
                      rotation: float = 0.,
                      point_n: int = 11) -> Path:
    """
    Generate matplotlib path to represent the FoV of a camera sensor.

    Parameters
    ----------
    fov_min_distance : float, optional
        Minimal distance of the FOV of the camera, by default 0m.
    fov_max_distance : float, optional
        Maximal distance of the field of view of the camera, by default 12,.
    fov_angle : float, optional
        Angle of the field of view of the camera, by default np.radians(100).
    fov_scaling : float, optional
        Scaling factor for the field of view of the camera, by default 1.0.
    offset_after_rotation : npt.NDArray[np.floating], optional
        Offset to move the field of view after rotating it, by default np.array([0, 0]).
    offset_before_rotation : npt.NDArray[np.floating], optional
        Offset to move the field of view before rotating it, by default np.array([0, 0]).
    rotation : float, optional
        Rotation angle to rotate the field of view, by default 0.
    point_n : int, optional
        Count of points to use for constructing the path of the field of view, by default 11.

    Returns
    -------
    Path
        Matplotlib path to describe the field of view of the camera.
    """
    path_points = []

    fov_min_distance /= fov_scaling
    fov_max_distance *= fov_scaling
    fov_angle *= fov_scaling

    for angle in np.linspace(- fov_angle / 2, fov_angle / 2, point_n):
        y = fov_min_distance * np.sin(angle)
        x = fov_min_distance * (np.cos(angle))
        path_points.append((x, y))

    for angle in np.linspace(fov_angle / 2, - fov_angle / 2, point_n):
        y = fov_max_distance * np.sin(angle)
        x = fov_max_distance * (np.cos(angle))
        path_points.append((x, y))

    path_points = rotate_and_move(np.array(path_points),
                                  rotation=rotation,
                                  offset_after_rotation=offset_after_rotation,
                                  offset_before_rotation=offset_before_rotation)
    return Path(path_points, closed=True)


def limit_coordinates_inside_path(coordinates: npt.NDArray, path: Path) -> Tuple[npt.NDArray, npt.NDArray]:
    """
    Limit the given set of coordinates to those that lie within the specified matplotlib path.

    Parameters
    ----------
    coordinates : npt.NDArray
        Specified coordinates to be limited.
    path : Path
        Specified path to limit the coordinates.

    Returns
    -------
    Tuple[npt.NDArray, npt.NDArray]
        Coordinates lieing inside the matplotlib Path,
        Indices of returned coordinates in respect to the input data.
    """
    indices = np.argwhere(path.contains_points(coordinates))
    indices = indices.reshape(indices.shape[0])
    return coordinates[indices], indices