Add scaling operation. (#187)

* Add scaling operation.

* Add random global rotation.

* Add unit tests + quaternion canonicalization.

rust/Cargo.toml

@@ -21,6 +21,7 @@ image = { git = "https://github.com/Shnatsel/image", branch = "zune-jpeg-sneaky"
 indicatif = "0.17.3"
 itertools = "0.10.5"
 ndarray = { version = "0.15.6", features = [
+    "approx",
     "matrixmultiply-threading",
     "rayon",
 ] }

rust/src/geometry/augmentations.rs

@@ -2,24 +2,27 @@
 //!
 //! Geometric augmentations.
 
+use std::f32::consts::PI;
+
+use crate::{
+    geometry::so3::_mat3_to_quat,
+    io::ndarray_from_frame,
+    share::{data_frame_to_ndarray_f32, ndarray_to_expr_vec},
+};
 use itertools::Itertools;
-use ndarray::{azip, concatenate, s, Axis};
+use ndarray::{azip, concatenate, par_azip, s, Array, Axis, Ix1, Ix2};
 use polars::{
     lazy::dsl::{col, cols, GetOutput},
     prelude::{DataFrame, DataType, IntoLazy},
     series::Series,
 };
 use rand_distr::{Bernoulli, Distribution, Uniform};
 
-use crate::{
-    io::ndarray_from_frame,
-    share::{data_frame_to_ndarray_f32, ndarray_to_expr_vec},
-};
-
 use super::{
     polytope::{compute_interior_points_mask, cuboids_to_polygons},
     so3::{
-        reflect_orientation_x, reflect_orientation_y, reflect_translation_x, reflect_translation_y,
+        _quat_to_mat3, reflect_orientation_x, reflect_orientation_y, reflect_translation_x,
+        reflect_translation_y,
     },
 };
 
@@ -123,6 +126,149 @@ pub fn sample_scene_reflection_y(
     }
 }
 
+/// Sample a scene global scale.
+/// This scales the lidar coordinates (x,y,z) and the cuboid centers (tx_m,ty_m,tz_m).
+pub fn sample_scene_global_scale(
+    lidar: DataFrame,
+    cuboids: DataFrame,
+    low_inclusive: f64,
+    upper_inclusive: f64,
+) -> (DataFrame, DataFrame) {
+    let distribution = Uniform::new_inclusive(low_inclusive, upper_inclusive);
+    let scale_factor = distribution.sample(&mut rand::thread_rng()) as f32;
+    let augmented_lidar = lidar
+        .lazy()
+        .with_column(col("x").map(
+            move |x| {
+                Ok(Some(
+                    x.f32()
+                        .unwrap()
+                        .into_no_null_iter()
+                        .map(|x| scale_factor * x)
+                        .collect::<Series>(),
+                ))
+            },
+            GetOutput::from_type(DataType::Float32),
+        ))
+        .with_column(col("y").map(
+            move |y| {
+                Ok(Some(
+                    y.f32()
+                        .unwrap()
+                        .into_no_null_iter()
+                        .map(|y| scale_factor * y)
+                        .collect::<Series>(),
+                ))
+            },
+            GetOutput::from_type(DataType::Float32),
+        ))
+        .with_column(col("z").map(
+            move |z| {
+                Ok(Some(
+                    z.f32()
+                        .unwrap()
+                        .into_no_null_iter()
+                        .map(|z| scale_factor * z)
+                        .collect::<Series>(),
+                ))
+            },
+            GetOutput::from_type(DataType::Float32),
+        ))
+        .collect()
+        .unwrap();
+
+    let augmented_cuboids = cuboids
+        .lazy()
+        .with_column(col("tx_m").map(
+            move |x| {
+                Ok(Some(
+                    x.f32()
+                        .unwrap()
+                        .into_no_null_iter()
+                        .map(|x| scale_factor * x)
+                        .collect::<Series>(),
+                ))
+            },
+            GetOutput::from_type(DataType::Float32),
+        ))
+        .with_column(col("ty_m").map(
+            move |y| {
+                Ok(Some(
+                    y.f32()
+                        .unwrap()
+                        .into_no_null_iter()
+                        .map(|y| scale_factor * y)
+                        .collect::<Series>(),
+                ))
+            },
+            GetOutput::from_type(DataType::Float32),
+        ))
+        .with_column(col("tz_m").map(
+            move |z| {
+                Ok(Some(
+                    z.f32()
+                        .unwrap()
+                        .into_no_null_iter()
+                        .map(|z| scale_factor * z)
+                        .collect::<Series>(),
+                ))
+            },
+            GetOutput::from_type(DataType::Float32),
+        ))
+        .collect()
+        .unwrap();
+    (augmented_lidar, augmented_cuboids)
+}
+
+/// Sample a scene global rotation.
+/// This rotates the lidar coordinates (x,y,z) and the cuboid centers (tx_m,ty_m,tz_m).
+pub fn sample_scene_global_rotation(
+    lidar: DataFrame,
+    cuboids: DataFrame,
+    low_inclusive: f64,
+    upper_inclusive: f64,
+) -> (DataFrame, DataFrame) {
+    let distribution = Uniform::new_inclusive(low_inclusive, upper_inclusive);
+    let theta = distribution.sample(&mut rand::thread_rng()) as f32;
+    let rotation = Array::<f32, Ix1>::from_vec(vec![
+        f32::cos(2.0 * PI * theta),
+        f32::sin(2.0 * PI * theta),
+        0.0,
+        -f32::sin(2.0 * PI * theta),
+        f32::cos(2.0 * PI * theta),
+        0.0,
+        0.0,
+        0.0,
+        1.0,
+    ])
+    .into_shape((3, 3))
+    .unwrap();
+
+    let column_names = ["x", "y", "z"];
+    let lidar_ndarray = ndarray_from_frame(&lidar, cols(column_names));
+    let augmented_lidar_ndarray = lidar_ndarray.dot(&rotation);
+    let series_vec = ndarray_to_expr_vec(augmented_lidar_ndarray, column_names.to_vec());
+    let augmented_lidar = lidar.lazy().with_columns(series_vec).collect().unwrap();
+
+    let cuboid_column_names = ["tx_m", "ty_m", "tz_m", "qw", "qx", "qy", "qz"];
+    let cuboids_ndarray = ndarray_from_frame(&cuboids, cols(cuboid_column_names));
+
+    let num_cuboids = cuboids_ndarray.shape()[0];
+    let mut augmented_cuboids = Array::<f32, Ix2>::zeros((num_cuboids, 7));
+    par_azip!((mut ac in augmented_cuboids.outer_iter_mut(), c in cuboids_ndarray.outer_iter()) {
+        let augmented_translation = c.slice(s![..3]).dot(&rotation);
+        let augmented_mat3 = _quat_to_mat3(&c.slice(s![3..7])).dot(&rotation.t());
+        let augmented_rotation = _mat3_to_quat(&augmented_mat3.view());
+
+        ac.slice_mut(s![..3]).assign(&augmented_translation);
+        ac.slice_mut(s![3..7]).assign(&augmented_rotation);
+    });
+
+    let series_vec = ndarray_to_expr_vec(augmented_cuboids, cuboid_column_names.to_vec());
+    let augmented_cuboids = cuboids.lazy().with_columns(series_vec).collect().unwrap();
+    (augmented_lidar, augmented_cuboids)
+}
+
 /// Sample a scene with random object scaling.
 pub fn sample_random_object_scale(
     lidar: DataFrame,

rust/src/geometry/camera/pinhole_camera.rs

@@ -7,7 +7,7 @@ use polars::{
 };
 
 use crate::{
-    geometry::se3::SE3, geometry::so3::quat_to_mat3, geometry::utils::cart_to_hom,
+    geometry::se3::SE3, geometry::so3::_quat_to_mat3, geometry::utils::cart_to_hom,
     io::read_feather_eager,
 };
 
@@ -128,7 +128,7 @@ impl PinholeCamera {
         let tz_m = extract_f32_from_frame(&extrinsics, "tz_m");
 
         let quat_wxyz = Array::<f32, Ix1>::from_vec(vec![qw, qx, qy, qz]);
-        let rotation = quat_to_mat3(&quat_wxyz.view());
+        let rotation = _quat_to_mat3(&quat_wxyz.view());
         let translation = Array::<f32, Ix1>::from_vec(vec![tx_m, ty_m, tz_m]);
         let ego_se3_cam = SE3 {
             rotation,

rust/src/geometry/polytope.rs

@@ -5,7 +5,7 @@
 use ndarray::{concatenate, par_azip, s, Array, ArrayView, Axis, Ix1, Ix2, Ix3, Slice};
 use once_cell::sync::Lazy;
 
-use super::so3::quat_to_mat3;
+use super::so3::_quat_to_mat3;
 
 // Safety: 24 elements (8 * 3 = 24) are defined.
 static VERTS: Lazy<Array<f32, Ix2>> = Lazy::new(|| unsafe {
@@ -83,7 +83,7 @@ fn _cuboid_to_polygon(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]);
-    let mat = quat_to_mat3(&quat_wxyz);
+    let mat = _quat_to_mat3(&quat_wxyz);
     let verts = &VERTS.clone() * &dims_lwh / 2.;
     let verts = verts.dot(&mat.t()) + center_xyz;
     verts.as_standard_layout().to_owned()

rust/src/geometry/so3.rs

@@ -4,10 +4,22 @@
 
 use std::f32::consts::PI;
 
-use ndarray::{par_azip, s, Array, Array2, ArrayView, Ix1, Ix2};
+use ndarray::{par_azip, s, Array, Array2, ArrayView, Ix1, Ix2, Ix3};
+use rand_distr::{Distribution, StandardNormal};
 
 /// Convert a quaternion in scalar-first format to a 3x3 rotation matrix.
-pub fn quat_to_mat3(quat_wxyz: &ArrayView<f32, Ix1>) -> Array<f32, Ix2> {
+/// Parallelized for batch processing.
+pub fn quat_to_mat3(quat_wxyz: &ArrayView<f32, Ix2>) -> Array<f32, Ix3> {
+    let num_quats = quat_wxyz.shape()[0];
+    let mut mat3 = Array::<f32, Ix3>::zeros((num_quats, 3, 3));
+    par_azip!((mut m in mat3.outer_iter_mut(), q in quat_wxyz.outer_iter()) {
+        m.assign(&_quat_to_mat3(&q));
+    });
+    mat3
+}
+
+/// Convert a quaternion in scalar-first format to a 3x3 rotation matrix.
+pub fn _quat_to_mat3(quat_wxyz: &ArrayView<f32, Ix1>) -> Array<f32, Ix2> {
     let w = quat_wxyz[0];
     let x = quat_wxyz[1];
     let y = quat_wxyz[2];
@@ -34,6 +46,57 @@ pub fn quat_to_mat3(quat_wxyz: &ArrayView<f32, Ix1>) -> Array<f32, Ix2> {
     }
 }
 
+/// Convert a 3x3 rotation matrix to a scalar-first quaternion.
+/// Parallelized for batch processing.
+pub fn mat3_to_quat(mat3: &ArrayView<f32, Ix3>) -> Array<f32, Ix2> {
+    let num_transformations = mat3.shape()[0];
+    let mut quat_wxyz = Array::<f32, Ix2>::zeros((num_transformations, 4));
+    par_azip!((mut q in quat_wxyz.outer_iter_mut(), m in mat3.outer_iter()) {
+        q.assign(&_mat3_to_quat(&m))
+    });
+    quat_wxyz
+}
+
+/// Convert a 3x3 rotation matrix to a scalar-first quaternion.
+pub fn _mat3_to_quat(mat3: &ArrayView<f32, Ix2>) -> Array<f32, Ix1> {
+    let trace = mat3[[0, 0]] + mat3[[1, 1]] + mat3[[2, 2]];
+    let mut quat_wxyz = if trace > 0.0 {
+        let s = 0.5 / f32::sqrt(trace + 1.0);
+        let qw = 0.25 / s;
+        let qx = (mat3[[2, 1]] - mat3[[1, 2]]) * s;
+        let qy = (mat3[[0, 2]] - mat3[[2, 0]]) * s;
+        let qz = (mat3[[1, 0]] - mat3[[0, 1]]) * s;
+        Array::<f32, Ix1>::from_vec(vec![qw, qx, qy, qz])
+    } else if mat3[[0, 0]] > mat3[[1, 1]] && mat3[[0, 0]] > mat3[[2, 2]] {
+        let s = 2.0 * f32::sqrt(1.0 + mat3[[0, 0]] - mat3[[1, 1]] - mat3[[2, 2]]);
+        let qw = (mat3[[2, 1]] - mat3[[1, 2]]) / s;
+        let qx = 0.25 * s;
+        let qy = (mat3[[0, 1]] + mat3[[1, 0]]) / s;
+        let qz = (mat3[[0, 2]] + mat3[[2, 0]]) / s;
+        Array::<f32, Ix1>::from_vec(vec![qw, qx, qy, qz])
+    } else if mat3[[1, 1]] > mat3[[2, 2]] {
+        let s = 2.0 * f32::sqrt(1.0 + mat3[[1, 1]] - mat3[[0, 0]] - mat3[[2, 2]]);
+        let qw = (mat3[[0, 2]] - mat3[[2, 0]]) / s;
+        let qx = (mat3[[0, 1]] + mat3[[1, 0]]) / s;
+        let qy = 0.25 * s;
+        let qz = (mat3[[1, 2]] + mat3[[2, 1]]) / s;
+        Array::<f32, Ix1>::from_vec(vec![qw, qx, qy, qz])
+    } else {
+        let s = 2.0 * f32::sqrt(1.0 + mat3[[2, 2]] - mat3[[0, 0]] - mat3[[1, 1]]);
+        let qw = (mat3[[1, 0]] - mat3[[0, 1]]) / s;
+        let qx = (mat3[[0, 2]] + mat3[[2, 0]]) / s;
+        let qy = (mat3[[1, 2]] + mat3[[2, 1]]) / s;
+        let qz = 0.25 * s;
+        Array::<f32, Ix1>::from_vec(vec![qw, qx, qy, qz])
+    };
+
+    // Canonicalize the quaternion.
+    if quat_wxyz[0] < 0.0 {
+        quat_wxyz *= -1.0;
+    }
+    quat_wxyz
+}
+
 /// Convert a scalar-first quaternion to yaw.
 /// In the Argoverse 2 coordinate system, this is counter-clockwise rotation about the +z axis.
 /// Parallelized for batch processing.
@@ -107,3 +170,38 @@ pub fn reflect_translation_y(xyz_m: &ArrayView<f32, Ix2>) -> Array<f32, Ix2> {
         .par_mapv_inplace(|x| -x);
     augmented_xyz_m
 }
+
+/// Sample a random quaternion.
+pub fn sample_random_quat_wxyz() -> Array<f32, Ix1> {
+    let distribution = StandardNormal;
+    let qw: f32 = distribution.sample(&mut rand::thread_rng());
+    let qx: f32 = distribution.sample(&mut rand::thread_rng());
+    let qy: f32 = distribution.sample(&mut rand::thread_rng());
+    let qz: f32 = distribution.sample(&mut rand::thread_rng());
+    let quat_wxyz = Array::<f32, Ix1>::from_vec(vec![qw, qx, qy, qz]);
+    let norm = quat_wxyz.dot(&quat_wxyz).sqrt();
+    let mut versor_wxyz = quat_wxyz / norm;
+
+    // Canonicalize the quaternion.
+    if versor_wxyz[0] < 0.0 {
+        versor_wxyz *= -1.0;
+    }
+    versor_wxyz
+}
+
+#[cfg(test)]
+mod tests {
+    use super::{_mat3_to_quat, _quat_to_mat3, sample_random_quat_wxyz};
+
+    #[test]
+    fn test_quat_to_mat3_round_trip() {
+        let num_trials = 100000;
+        let epsilon = 1e-6;
+        for _ in 0..num_trials {
+            let quat_wxyz = sample_random_quat_wxyz();
+            let mat3 = _quat_to_mat3(&quat_wxyz.view());
+            let _quat_wxyz = _mat3_to_quat(&mat3.view());
+            assert!(quat_wxyz.abs_diff_eq(&_quat_wxyz, epsilon));
+        }
+    }
+}