Add interior point computation to pytorch data loader.

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::<Float32Type>().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()) {

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());
 

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<f32, Ix2>) -> Array<f32, Ix3> {
+/// Convert (N,10) cuboids to vertices (N,8,3).
+pub fn cuboids_to_vertices(cuboids: &ArrayView<f32, Ix2>) -> Array<f32, Ix3> {
     let num_cuboids = cuboids.shape()[0];
     let mut polygons = Array::<f32, Ix3>::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<f32, Ix1>) -> Array<f32, Ix2> {
+/// Convert a single cuboid to vertices.
+fn _cuboid_to_vertices(cuboid: &ArrayView<f32, Ix1>) -> Array<f32, Ix2> {
     let center_xyz = cuboid.slice(s![0..3]);
     let dims_lwh = cuboid.slice(s![3..6]);
     let quat_wxyz = cuboid.slice(s![6..10]);

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<f32, Ix2>,
+) -> &'py PyArray<f32, Ix3> {
+    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<f32, Ix2>,
+    cuboids: PyReadonlyArray<f32, Ix3>,
+) -> &'py PyArray<bool, Ix2> {
+    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::<DataLoader>()?;
     m.add_class::<Sweep>()?;
+    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)?)?;

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: ...

src/av2/torch/structures/cuboids.py

@@ -5,21 +5,24 @@
 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):
     """Cuboid parameterization modes."""
 
     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)

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.
 

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)