diff --git a/rust/src/bin/export_augmentation_database.rs b/rust/src/bin/export_augmentation_database.rs index 932dfead..0cb8fba0 100644 --- a/rust/src/bin/export_augmentation_database.rs +++ b/rust/src/bin/export_augmentation_database.rs @@ -6,7 +6,7 @@ use std::{fs, path::PathBuf}; use av2::{ data_loader::DataLoader, - geometry::polytope::{compute_interior_points_mask, cuboids_to_polygons}, + geometry::polytope::{compute_interior_points_mask, cuboids_to_vertices}, io::write_feather_eager, }; use indicatif::ProgressBar; @@ -96,7 +96,7 @@ pub fn main() { .clone(); let cuboids = cuboids.clone().to_ndarray::().unwrap(); - let cuboid_vertices = cuboids_to_polygons(&cuboids.view()); + let cuboid_vertices = cuboids_to_vertices(&cuboids.view()); let points = lidar_ndarray.slice(s![.., ..3]); let mask = compute_interior_points_mask(&points.view(), &cuboid_vertices.view()); for (c, m) in category.into_iter().zip(mask.outer_iter()) { diff --git a/rust/src/geometry/augmentations.rs b/rust/src/geometry/augmentations.rs index 193b7437..b0079a25 100644 --- a/rust/src/geometry/augmentations.rs +++ b/rust/src/geometry/augmentations.rs @@ -17,7 +17,7 @@ use crate::{ }; use super::{ - polytope::{compute_interior_points_mask, cuboids_to_polygons}, + polytope::{compute_interior_points_mask, cuboids_to_vertices}, so3::{ reflect_orientation_x, reflect_orientation_y, reflect_translation_x, reflect_translation_y, }, @@ -135,7 +135,7 @@ pub fn sample_random_object_scale( ]; let mut lidar_ndarray = ndarray_from_frame(&lidar, cols(["x", "y", "z"])); let mut cuboids_ndarray = ndarray_from_frame(&cuboids, cols(cuboid_column_names)); - let cuboid_vertices = cuboids_to_polygons(&cuboids_ndarray.view()); + let cuboid_vertices = cuboids_to_vertices(&cuboids_ndarray.view()); let interior_points_mask = compute_interior_points_mask(&lidar_ndarray.view(), &cuboid_vertices.view()); diff --git a/rust/src/geometry/polytope.rs b/rust/src/geometry/polytope.rs index 85ab98fd..fe778f1e 100644 --- a/rust/src/geometry/polytope.rs +++ b/rust/src/geometry/polytope.rs @@ -68,18 +68,18 @@ pub fn compute_interior_points_mask( is_interior } -/// Convert (N,10) cuboids to polygons. -pub fn cuboids_to_polygons(cuboids: &ArrayView) -> Array { +/// Convert (N,10) cuboids to vertices (N,8,3). +pub fn cuboids_to_vertices(cuboids: &ArrayView) -> Array { let num_cuboids = cuboids.shape()[0]; let mut polygons = Array::::zeros([num_cuboids, 8, 3]); par_azip!((mut p in polygons.outer_iter_mut(), c in cuboids.outer_iter()) { - p.assign(&_cuboid_to_polygon(&c)) + p.assign(&_cuboid_to_vertices(&c)) }); polygons } -/// Convert a single cuboid to a polygon. -fn _cuboid_to_polygon(cuboid: &ArrayView) -> Array { +/// Convert a single cuboid to vertices. +fn _cuboid_to_vertices(cuboid: &ArrayView) -> Array { let center_xyz = cuboid.slice(s![0..3]); let dims_lwh = cuboid.slice(s![3..6]); let quat_wxyz = cuboid.slice(s![6..10]); diff --git a/rust/src/lib.rs b/rust/src/lib.rs index f8540610..5ca3c625 100644 --- a/rust/src/lib.rs +++ b/rust/src/lib.rs @@ -15,7 +15,8 @@ pub mod share; pub mod structures; use data_loader::{DataLoader, Sweep}; -use ndarray::{Dim, Ix1, Ix2}; +use geometry::polytope::{compute_interior_points_mask, cuboids_to_vertices}; +use ndarray::{Dim, Ix1, Ix2, Ix3}; use numpy::PyReadonlyArray; use numpy::{IntoPyArray, PyArray}; use pyo3::prelude::*; @@ -84,11 +85,35 @@ fn py_yaw_to_quat<'py>( yaw_to_quat(&quat_wxyz.as_array().view()).into_pyarray(py) } +#[pyfunction] +#[pyo3(name = "cuboids_to_vertices")] +#[allow(clippy::type_complexity)] +fn py_cuboids_to_vertices<'py>( + py: Python<'py>, + cuboids: PyReadonlyArray, +) -> &'py PyArray { + cuboids_to_vertices(&cuboids.as_array().view()).into_pyarray(py) +} + +#[pyfunction] +#[pyo3(name = "compute_interior_points_mask")] +#[allow(clippy::type_complexity)] +fn py_compute_interior_points_mask<'py>( + py: Python<'py>, + points_m: PyReadonlyArray, + cuboids: PyReadonlyArray, +) -> &'py PyArray { + compute_interior_points_mask(&points_m.as_array().view(), &cuboids.as_array().view()) + .into_pyarray(py) +} + /// A Python module implemented in Rust. #[pymodule] fn _r(_py: Python, m: &PyModule) -> PyResult<()> { m.add_class::()?; m.add_class::()?; + m.add_function(wrap_pyfunction!(py_compute_interior_points_mask, m)?)?; + m.add_function(wrap_pyfunction!(py_cuboids_to_vertices, m)?)?; m.add_function(wrap_pyfunction!(py_quat_to_mat3, m)?)?; m.add_function(wrap_pyfunction!(py_quat_to_yaw, m)?)?; m.add_function(wrap_pyfunction!(py_voxelize, m)?)?; diff --git a/src/av2/_r.pyi b/src/av2/_r.pyi index 2000aba6..50dfb473 100644 --- a/src/av2/_r.pyi +++ b/src/av2/_r.pyi @@ -6,6 +6,8 @@ from typing import List, Optional, Tuple import polars as pl import torch +from av2.utils.typing import NDArrayFloat32 + @dataclass class DataLoader: root_dir: str @@ -27,3 +29,8 @@ class Sweep: lidar: pl.DataFrame sweep_uuid: Tuple[str, int] cuboids: Optional[pl.DataFrame] + +def compute_interior_points_mask( + points_xyz_m: NDArrayFloat32, cuboid_vertices: NDArrayFloat32 +) -> NDArrayFloat32: ... +def cuboids_to_vertices(cuboids: NDArrayFloat32) -> NDArrayFloat32: ... diff --git a/src/av2/torch/structures/cuboids.py b/src/av2/torch/structures/cuboids.py index d38118e8..8de016ff 100644 --- a/src/av2/torch/structures/cuboids.py +++ b/src/av2/torch/structures/cuboids.py @@ -5,14 +5,16 @@ from dataclasses import dataclass from enum import Enum from functools import cached_property -from typing import List +from typing import List, Tuple import pandas as pd import torch from kornia.geometry.conversions import euler_from_quaternion +import av2._r as rust + from .. import XYZLWH_QWXYZ_COLUMNS -from .utils import tensor_from_frame +from .utils import ndarray_from_frame, tensor_from_frame class CuboidMode(str, Enum): @@ -20,6 +22,7 @@ class CuboidMode(str, Enum): XYZLWH_T = "XYZLWH_T" # 1-DOF orientation. XYZLWH_QWXYZ = "XYZLWH_QWXYZ" # 3-DOF orientation. + XYZ_VERTICES = "XYZ_VERTICES" # Vertex representation (eight vertices). @dataclass(frozen=True) @@ -60,8 +63,8 @@ def as_tensor(self, cuboid_mode: CuboidMode = CuboidMode.XYZLWH_T) -> torch.Tens Raises: NotImplementedError: Raised if the cuboid mode is not supported. """ - xyzlwh_qwxyz = tensor_from_frame(self._frame, list(XYZLWH_QWXYZ_COLUMNS)) if cuboid_mode == CuboidMode.XYZLWH_T: + xyzlwh_qwxyz = tensor_from_frame(self._frame, list(XYZLWH_QWXYZ_COLUMNS)) quat_wxyz = xyzlwh_qwxyz[:, 6:10] w, x, y, z = ( quat_wxyz[:, 0], @@ -72,8 +75,32 @@ def as_tensor(self, cuboid_mode: CuboidMode = CuboidMode.XYZLWH_T) -> torch.Tens _, _, yaw = euler_from_quaternion(w, x, y, z) return torch.concat([xyzlwh_qwxyz[:, :6], yaw[:, None]], dim=-1) elif cuboid_mode == CuboidMode.XYZLWH_QWXYZ: + xyzlwh_qwxyz = tensor_from_frame(self._frame, list(XYZLWH_QWXYZ_COLUMNS)) return xyzlwh_qwxyz + elif cuboid_mode == CuboidMode.XYZ_VERTICES: + xyzlwh_qwxyz = ndarray_from_frame(self._frame, list(XYZLWH_QWXYZ_COLUMNS)) + xyz_vertices = rust.cuboids_to_vertices(xyzlwh_qwxyz) + return torch.as_tensor(xyz_vertices) else: raise NotImplementedError( f"{cuboid_mode} orientation mode is not implemented." ) + + def compute_interior_points_mask(self, points_xyz_m: torch.Tensor) -> torch.Tensor: + """Compute a boolean mask for each cuboid indicating whether each point is interior to the geometry. + + Notation: + N: Number of points. + M: Number of cuboids. + + Args: + points_xyz_m: (N,3) Points to filter. + + Returns: + (N,M) Boolean array indicating which points are interior. + """ + vertices_dst_xyz = self.as_tensor(CuboidMode.XYZ_VERTICES) + is_interior = rust.compute_interior_points_mask( + points_xyz_m[:, :3].numpy(), vertices_dst_xyz.numpy() + ) + return torch.as_tensor(is_interior) diff --git a/src/av2/torch/structures/utils.py b/src/av2/torch/structures/utils.py index 4b73eb47..7e2e7dc5 100644 --- a/src/av2/torch/structures/utils.py +++ b/src/av2/torch/structures/utils.py @@ -14,7 +14,7 @@ def tensor_from_frame(frame: pd.DataFrame, columns: List[str]) -> torch.Tensor: - """Build lidar `torch` tensor from `pandas` dataframe. + """Build lidar `torch` tensor from a `pandas` dataframe. Notation: N: Number of rows. @@ -27,10 +27,28 @@ def tensor_from_frame(frame: pd.DataFrame, columns: List[str]) -> torch.Tensor: Returns: (N,K) tensor containing the frame data. """ - frame_npy = frame.loc[:, columns].to_numpy().astype(np.float32) + frame_npy = ndarray_from_frame(frame, columns) return torch.as_tensor(frame_npy) +def ndarray_from_frame(frame: pd.DataFrame, columns: List[str]) -> torch.Tensor: + """Build lidar `numpy` ndarray from a `pandas` dataframe. + + Notation: + N: Number of rows. + K: Number of columns. + + Args: + frame: (N,K) Pandas DataFrame containing N rows with K columns. + columns: List of DataFrame columns. + + Returns: + (N,K) ndarray containing the frame data. + """ + frame_npy = frame.loc[:, columns].to_numpy().astype(np.float32) + return frame_npy + + def SE3_from_frame(frame: pd.DataFrame) -> Se3: """Build SE(3) object from `pandas` DataFrame. diff --git a/tutorials/3d_object_detection.py b/tutorials/3d_object_detection.py index ff6d9cc5..e5048a16 100644 --- a/tutorials/3d_object_detection.py +++ b/tutorials/3d_object_detection.py @@ -7,7 +7,9 @@ from kornia.geometry.linalg import transform_points from tqdm import tqdm +from av2.geometry.geometry import compute_interior_points_mask from av2.torch.data_loaders.detection import DetectionDataLoader +from av2.torch.structures.cuboids import CuboidMode logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) @@ -47,6 +49,11 @@ def main( # Lidar (x,y,z) in meters and intensity (i). lidar_tensor = sweep.lidar.as_tensor() + # (N,M) Mask indicating which points (rows) are interior to each cuboid (columns). + interior_points_mask = sweep.cuboids.compute_interior_points_mask( + lidar_tensor[:, :3] + ) + # Synchronized ring cameras. synchronized_images = data_loader.get_synchronized_images(i)