diff --git a/Cargo.lock b/Cargo.lock index 4c380fab..696e73f9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -213,6 +213,8 @@ dependencies = [ "polars", "pyo3", "pyo3-polars", + "rand", + "rand_distr", "rayon", "serde", "strum", diff --git a/rust/Cargo.toml b/rust/Cargo.toml index ac74c1fc..80976673 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -44,6 +44,8 @@ pyo3 = { version = "0.18.3", features = ["extension-module"] } pyo3-polars = { git = "https://github.com/benjaminrwilson/pyo3-polars", rev = "993d22a6bf54ceb8f93c5ed9082621330b186a52", features = [ "serde", ] } +rand = "0.8.5" +rand_distr = "0.4.3" rayon = "1.7.0" serde = "1.0.160" strum = "0.24.1" diff --git a/rust/src/geometry/augmentations.rs b/rust/src/geometry/augmentations.rs index 30297422..f4b23b18 100644 --- a/rust/src/geometry/augmentations.rs +++ b/rust/src/geometry/augmentations.rs @@ -2,22 +2,117 @@ //! //! Geometric augmentations. -use std::f32::consts::PI; +use itertools::Itertools; +use ndarray::{concatenate, Axis}; +use polars::{ + lazy::dsl::{col, GetOutput}, + prelude::{DataFrame, DataType, IntoLazy}, + series::Series, +}; +use rand_distr::{Bernoulli, Distribution}; -use ndarray::{Array, ArrayView, Ix2}; +use crate::share::{data_frame_to_ndarray_f32, ndarray_to_expr_vec}; -use crate::geometry::so3::{quat_to_yaw, yaw_to_quat}; +use super::so3::{ + reflect_orientation_x, reflect_orientation_y, reflect_translation_x, reflect_translation_y, +}; -/// Reflect pose across the x-axis. -pub fn reflect_pose_x(quat_wxyz: &ArrayView) -> Array { - let yaw_rad = quat_to_yaw(quat_wxyz); - let reflected_yaw_rad = -yaw_rad; - yaw_to_quat(&reflected_yaw_rad.view()) +/// Sample a scene reflection. +/// This reflects both a point cloud and cuboids across the x-axis. +pub fn sample_scene_reflection_x( + lidar: DataFrame, + cuboids: DataFrame, + p: f64, +) -> (DataFrame, DataFrame) { + let distribution = Bernoulli::new(p).unwrap(); + let is_augmented = distribution.sample(&mut rand::thread_rng()); + if is_augmented { + let augmented_lidar = lidar + .lazy() + .with_column(col("y").map( + move |x| { + Ok(Some( + x.f32() + .unwrap() + .into_no_null_iter() + .map(|y| -y) + .collect::(), + )) + }, + GetOutput::from_type(DataType::Float32), + )) + .collect() + .unwrap(); + + let translation_column_names = vec!["tx_m", "ty_m", "tz_m"]; + let txyz_m = data_frame_to_ndarray_f32(cuboids.clone(), translation_column_names.clone()); + let augmentation_translation = reflect_translation_x(&txyz_m.view()); + + let orientation_column_names = vec!["qw", "qx", "qy", "qz"]; + let quat_wxyz = + data_frame_to_ndarray_f32(cuboids.clone(), orientation_column_names.clone()); + let augmented_orientation = reflect_orientation_x(&quat_wxyz.view()); + let augmented_poses = + concatenate![Axis(1), augmentation_translation, augmented_orientation]; + + let column_names = translation_column_names + .into_iter() + .chain(orientation_column_names) + .collect_vec(); + let series_vec = ndarray_to_expr_vec(augmented_poses, column_names); + let augmented_cuboids = cuboids.lazy().with_columns(series_vec).collect().unwrap(); + (augmented_lidar, augmented_cuboids) + } else { + (lidar, cuboids) + } } -/// Reflect pose across the y-axis. -pub fn reflect_pose_y(quat_wxyz: &ArrayView) -> Array { - let yaw_rad = quat_to_yaw(quat_wxyz); - let reflected_yaw_rad = PI - yaw_rad; - yaw_to_quat(&reflected_yaw_rad.view()) +/// Sample a scene reflection. +/// This reflects both a point cloud and cuboids across the y-axis. +pub fn sample_scene_reflection_y( + lidar: DataFrame, + cuboids: DataFrame, + p: f64, +) -> (DataFrame, DataFrame) { + let distribution: Bernoulli = Bernoulli::new(p).unwrap(); + let is_augmented = distribution.sample(&mut rand::thread_rng()); + if is_augmented { + let augmented_lidar = lidar + .lazy() + .with_column(col("x").map( + move |x| { + Ok(Some( + x.f32() + .unwrap() + .into_no_null_iter() + .map(|x| -x) + .collect::(), + )) + }, + GetOutput::from_type(DataType::Float32), + )) + .collect() + .unwrap(); + + let translation_column_names = vec!["tx_m", "ty_m", "tz_m"]; + let txyz_m = data_frame_to_ndarray_f32(cuboids.clone(), translation_column_names.clone()); + let augmentation_translation = reflect_translation_y(&txyz_m.view()); + + let orientation_column_names = vec!["qw", "qx", "qy", "qz"]; + let quat_wxyz = + data_frame_to_ndarray_f32(cuboids.clone(), orientation_column_names.clone()); + let augmented_orientation = reflect_orientation_y(&quat_wxyz.view()); + let augmented_poses = + concatenate![Axis(1), augmentation_translation, augmented_orientation]; + + let column_names = translation_column_names + .into_iter() + .chain(orientation_column_names) + .collect_vec(); + let series_vec = ndarray_to_expr_vec(augmented_poses, column_names); + let augmented_cuboids = cuboids.lazy().with_columns(series_vec).collect().unwrap(); + (augmented_lidar, augmented_cuboids) + } else { + (lidar, cuboids) + } } diff --git a/rust/src/geometry/se3.rs b/rust/src/geometry/se3.rs index bc3a0207..9fc95d54 100644 --- a/rust/src/geometry/se3.rs +++ b/rust/src/geometry/se3.rs @@ -2,7 +2,8 @@ //! //! Special Euclidean Group 3. -use ndarray::{s, Array1, Array2, ArrayView2}; +use ndarray::ArrayView2; +use ndarray::{s, Array1, Array2}; /// Special Euclidean Group 3 (SE(3)). /// Rigid transformation parameterized by a rotation and translation in $R^3$. diff --git a/rust/src/geometry/so3.rs b/rust/src/geometry/so3.rs index c4f52b96..b3cacc10 100644 --- a/rust/src/geometry/so3.rs +++ b/rust/src/geometry/so3.rs @@ -2,7 +2,9 @@ //! //! Special Orthogonal Group 3 (SO(3)). -use ndarray::{par_azip, Array, Array2, ArrayView, Ix1, Ix2}; +use std::f32::consts::PI; + +use ndarray::{par_azip, s, Array, Array2, ArrayView, Ix1, Ix2}; /// Convert a quaternion in scalar-first format to a 3x3 rotation matrix. pub fn quat_to_mat3(quat_wxyz: &ArrayView) -> Array { @@ -71,3 +73,37 @@ pub fn _yaw_to_quat(yaw_rad: f32) -> Array { let qz = f32::sin(0.5 * yaw_rad); Array::::from_vec(vec![qw, 0.0, 0.0, qz]) } + +/// Reflect orientation across the x-axis. +/// (N,4) `quat_wxyz` orientation of `N` rigid objects. +pub fn reflect_orientation_x(quat_wxyz: &ArrayView) -> Array { + let yaw_rad = quat_to_yaw(quat_wxyz); + let reflected_yaw_rad = -yaw_rad; + yaw_to_quat(&reflected_yaw_rad.view()) +} + +/// Reflect orientation across the y-axis. +/// (N,4) `quat_wxyz` orientation of `N` rigid objects. +pub fn reflect_orientation_y(quat_wxyz: &ArrayView) -> Array { + let yaw_rad = quat_to_yaw(quat_wxyz); + let reflected_yaw_rad = PI - yaw_rad; + yaw_to_quat(&reflected_yaw_rad.view()) +} + +/// Reflect translation across the x-axis. +pub fn reflect_translation_x(xyz_m: &ArrayView) -> Array { + let mut augmented_xyz_m = xyz_m.to_owned(); + augmented_xyz_m + .slice_mut(s![.., 1]) + .par_mapv_inplace(|y| -y); + augmented_xyz_m +} + +/// Reflect translation across the y-axis. +pub fn reflect_translation_y(xyz_m: &ArrayView) -> Array { + let mut augmented_xyz_m = xyz_m.to_owned(); + augmented_xyz_m + .slice_mut(s![.., 0]) + .par_mapv_inplace(|x| -x); + augmented_xyz_m +} diff --git a/rust/src/lib.rs b/rust/src/lib.rs index 0e6b0f5c..f8540610 100644 --- a/rust/src/lib.rs +++ b/rust/src/lib.rs @@ -15,7 +15,6 @@ pub mod share; pub mod structures; use data_loader::{DataLoader, Sweep}; -use geometry::augmentations::{reflect_pose_x, reflect_pose_y}; use ndarray::{Dim, Ix1, Ix2}; use numpy::PyReadonlyArray; use numpy::{IntoPyArray, PyArray}; @@ -85,26 +84,6 @@ fn py_yaw_to_quat<'py>( yaw_to_quat(&quat_wxyz.as_array().view()).into_pyarray(py) } -#[pyfunction] -#[pyo3(name = "reflect_pose_x")] -#[allow(clippy::type_complexity)] -fn py_reflect_pose_x<'py>( - py: Python<'py>, - quat_wxyz: PyReadonlyArray, -) -> &'py PyArray { - reflect_pose_x(&quat_wxyz.as_array().view()).into_pyarray(py) -} - -#[pyfunction] -#[pyo3(name = "reflect_pose_y")] -#[allow(clippy::type_complexity)] -fn py_reflect_pose_y<'py>( - py: Python<'py>, - quat_wxyz: PyReadonlyArray, -) -> &'py PyArray { - reflect_pose_y(&quat_wxyz.as_array().view()).into_pyarray(py) -} - /// A Python module implemented in Rust. #[pymodule] fn _r(_py: Python, m: &PyModule) -> PyResult<()> { @@ -112,8 +91,6 @@ fn _r(_py: Python, m: &PyModule) -> PyResult<()> { m.add_class::()?; 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_reflect_pose_x, m)?)?; - m.add_function(wrap_pyfunction!(py_reflect_pose_y, m)?)?; m.add_function(wrap_pyfunction!(py_voxelize, m)?)?; m.add_function(wrap_pyfunction!(py_yaw_to_quat, m)?)?; Ok(()) diff --git a/rust/src/share.rs b/rust/src/share.rs index db5cf66a..03c58854 100644 --- a/rust/src/share.rs +++ b/rust/src/share.rs @@ -3,10 +3,14 @@ //! Conversion methods between different libraries. use ndarray::{Array, Ix2}; -use polars::{prelude::NamedFrom, series::Series}; +use polars::{ + lazy::dsl::{cols, lit, Expr}, + prelude::{DataFrame, Float32Type, IntoLazy, NamedFrom}, + series::Series, +}; -/// Convert the columns of an `ndarray::Array` into a vector of `polars::series::Series`. -pub fn ndarray_to_series_vec(arr: Array, column_names: Vec<&str>) -> Vec { +/// Convert the columns of an `ndarray::Array` into a vector of `polars` expressions. +pub fn ndarray_to_expr_vec(arr: Array, column_names: Vec<&str>) -> Vec { let num_dims = arr.shape()[1]; if num_dims != column_names.len() { panic!("Number of array columns and column names must match."); @@ -18,7 +22,21 @@ pub fn ndarray_to_series_vec(arr: Array, column_names: Vec<&str>) -> V column_name, column.as_standard_layout().to_owned().into_raw_vec(), ); - series_vec.push(series); + series_vec.push(lit(series)); } series_vec } + +/// Convert a data frame to an `ndarray::Array::`. +pub fn data_frame_to_ndarray_f32( + data_frame: DataFrame, + column_names: Vec<&str>, +) -> Array { + data_frame + .lazy() + .select(&[cols(column_names)]) + .collect() + .unwrap() + .to_ndarray::() + .unwrap() +}