Add augmentation database + random object scaling.

rust/src/bin/export_augmentation_database.rs

@@ -0,0 +1,166 @@
+//! # export_augmentation_database
+//!
+//! Exports cropped annotations from the `train` dataset for augmentation during training.
+
+use std::{fs, path::PathBuf};
+
+use av2::{
+    data_loader::DataLoader,
+    geometry::polytope::{compute_interior_points_mask, cuboids_to_polygons},
+    io::write_feather_eager,
+};
+use indicatif::ProgressBar;
+
+#[macro_use]
+extern crate log;
+use itertools::Itertools;
+use ndarray::{s, Array, Axis, Ix2};
+use once_cell::sync::Lazy;
+use polars::{
+    df,
+    prelude::{DataFrame, Float32Type, NamedFrom},
+    series::Series,
+};
+use std::collections::HashMap;
+
+/// Constants can be changed to fit your directory structure.
+/// However, it's recommend to place the datasets in the default folders.
+
+/// Root directory to datasets.
+static ROOT_DIR: Lazy<PathBuf> = Lazy::new(|| dirs::home_dir().unwrap().join("data/datasets/"));
+
+/// Dataset name.
+static DATASET_NAME: &str = "av2";
+
+/// Dataset type. This will either be "lidar" or "sensor".
+static DATASET_TYPE: &str = "sensor";
+
+/// Split names for the dataset.
+static SPLIT_NAMES: Lazy<Vec<&str>> = Lazy::new(|| vec!["train"]);
+
+/// Number of accumulated sweeps.
+const NUM_ACCUMULATED_SWEEPS: usize = 1;
+
+/// Memory maps the sweeps for fast pre-processing. Requires .feather files to be uncompressed.
+const MEMORY_MAPPED: bool = false;
+
+static DST_DATASET_NAME: Lazy<String> =
+    Lazy::new(|| format!("{DATASET_NAME}_{NUM_ACCUMULATED_SWEEPS}_database"));
+static SRC_PREFIX: Lazy<PathBuf> = Lazy::new(|| ROOT_DIR.join(DATASET_NAME).join(DATASET_TYPE));
+static DST_PREFIX: Lazy<PathBuf> =
+    Lazy::new(|| ROOT_DIR.join(DST_DATASET_NAME.clone()).join(DATASET_TYPE));
+
+static EXPORTED_COLUMN_NAMES: Lazy<Vec<&str>> = Lazy::new(|| {
+    vec![
+        "x",
+        "y",
+        "z",
+        "intensity",
+        "laser_number",
+        "offset_ns",
+        "timedelta_ns",
+    ]
+});
+
+/// Script entrypoint.
+pub fn main() {
+    env_logger::init();
+    for split_name in SPLIT_NAMES.clone() {
+        let split_path = SRC_PREFIX.join(split_name);
+        if !split_path.exists() {
+            error!("Cannot find `{split_path:?}`. Skipping ...");
+            continue;
+        }
+        let data_loader = DataLoader::new(
+            ROOT_DIR.clone().to_str().unwrap(),
+            DATASET_NAME,
+            DATASET_TYPE,
+            split_name,
+            NUM_ACCUMULATED_SWEEPS,
+            MEMORY_MAPPED,
+        );
+
+        let mut category_counter: HashMap<String, u64> = HashMap::new();
+        let bar = ProgressBar::new(data_loader.len() as u64);
+        for sweep in data_loader {
+            let lidar = &sweep.lidar.0;
+            let lidar_ndarray = lidar.to_ndarray::<Float32Type>().unwrap();
+
+            let cuboids = sweep.cuboids.unwrap().0;
+            let category = cuboids["category"]
+                .utf8()
+                .unwrap()
+                .into_iter()
+                .map(|x| x.unwrap())
+                .collect_vec()
+                .clone();
+
+            let cuboids = cuboids.clone().to_ndarray::<Float32Type>().unwrap();
+            let cuboid_vertices = cuboids_to_polygons(&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()) {
+                let indices = m
+                    .iter()
+                    .enumerate()
+                    .filter_map(|(i, x)| match *x {
+                        true => Some(i),
+                        _ => None,
+                    })
+                    .collect_vec();
+
+                let points_i = lidar_ndarray.select(Axis(0), &indices);
+                let data_frame_i = _build_data_frame(points_i, EXPORTED_COLUMN_NAMES.clone());
+
+                category_counter
+                    .entry(c.to_string())
+                    .and_modify(|count| *count += 1)
+                    .or_insert(0);
+                let count = category_counter.get(&c.to_string()).unwrap();
+
+                let dst = DST_PREFIX.join(c).join(format!("{count:08}.feather"));
+                fs::create_dir_all(dst.parent().unwrap()).unwrap();
+                write_feather_eager(&dst, data_frame_i);
+            }
+            bar.inc(1);
+        }
+
+        let category = category_counter.keys().cloned().collect_vec();
+        let count = category_counter.values().cloned().collect_vec();
+        let num_padding = category_counter.values().map(|_| 8_u8).collect_vec();
+        let index =
+            df!("category" => category, "count" => count, "num_padding" => num_padding).unwrap();
+
+        let dst = DST_PREFIX.join("_index.feather");
+        write_feather_eager(&dst, index);
+    }
+}
+
+// Helper method to build exported `DataFrame`.
+fn _build_data_frame(arr: Array<f32, Ix2>, column_names: Vec<&str>) -> DataFrame {
+    let series_vec = arr
+        .columns()
+        .into_iter()
+        .zip(column_names.into_iter())
+        .map(|(column, column_name)| match column_name {
+            "x" => Series::new("x", column.to_owned().into_raw_vec()),
+            "y" => Series::new("y", column.to_owned().into_raw_vec()),
+            "z" => Series::new("z", column.to_owned().into_raw_vec()),
+            "intensity" => Series::new(
+                "intensity",
+                column.to_owned().mapv(|x| x as u8).into_raw_vec(),
+            ),
+            "laser_number" => Series::new(
+                "laser_number",
+                column.to_owned().mapv(|x| x as u8).into_raw_vec(),
+            ),
+            "offset_ns" => Series::new(
+                "offset_ns",
+                column.to_owned().mapv(|x| x as u32).into_raw_vec(),
+            ),
+            "timedelta_ns" => Series::new("timedelta_ns", column.to_owned().into_raw_vec()),
+            _ => panic!(),
+        })
+        .collect_vec();
+    DataFrame::from_iter(series_vec)
+}

rust/src/data_loader.rs

@@ -42,7 +42,7 @@ const MAX_CAM_LIDAR_TOL_NS: f32 = 50000000.;
 
 /// Data associated with a single lidar sweep.
 #[pyclass]
-#[derive(Debug)]
+#[derive(Clone, Debug)]
 pub struct Sweep {
     /// Ego-vehicle city pose.
     #[pyo3(get, set)]

rust/src/geometry/augmentations.rs

@@ -3,18 +3,24 @@
 //! Geometric augmentations.
 
 use itertools::Itertools;
-use ndarray::{concatenate, Axis};
+use ndarray::{azip, concatenate, s, Axis};
 use polars::{
-    lazy::dsl::{col, GetOutput},
+    lazy::dsl::{col, cols, GetOutput},
     prelude::{DataFrame, DataType, IntoLazy},
     series::Series,
 };
-use rand_distr::{Bernoulli, Distribution};
+use rand_distr::{Bernoulli, Distribution, Uniform};
 
-use crate::share::{data_frame_to_ndarray_f32, ndarray_to_expr_vec};
+use crate::{
+    io::ndarray_from_frame,
+    share::{data_frame_to_ndarray_f32, ndarray_to_expr_vec},
+};
 
-use super::so3::{
-    reflect_orientation_x, reflect_orientation_y, reflect_translation_x, reflect_translation_y,
+use super::{
+    polytope::{compute_interior_points_mask, cuboids_to_polygons},
+    so3::{
+        reflect_orientation_x, reflect_orientation_y, reflect_translation_x, reflect_translation_y,
+    },
 };
 
 /// Sample a scene reflection.
@@ -116,3 +122,57 @@ pub fn sample_scene_reflection_y(
         (lidar, cuboids)
     }
 }
+
+/// Sample a scene with random object scaling.
+pub fn sample_random_object_scale(
+    lidar: DataFrame,
+    cuboids: DataFrame,
+    low_inclusive: f64,
+    high_inclusive: f64,
+) -> (DataFrame, DataFrame) {
+    let cuboid_column_names = [
+        "tx_m", "ty_m", "tz_m", "length_m", "width_m", "height_m", "qw", "qx", "qy", "qz",
+    ];
+    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 interior_points_mask =
+        compute_interior_points_mask(&lidar_ndarray.view(), &cuboid_vertices.view());
+
+    let distribution = Uniform::new_inclusive(low_inclusive, high_inclusive);
+
+    azip!((mut c in cuboids_ndarray.outer_iter_mut(), m in interior_points_mask.outer_iter()) {
+        let scale_factor = distribution.sample(&mut rand::thread_rng()) as f32;
+        let indices = m
+            .iter()
+            .enumerate()
+            .filter_map(|(i, x)| match *x {
+                true => Some(i),
+                _ => None,
+            })
+            .collect_vec();
+        let interior_points = lidar_ndarray.select(Axis(0), &indices);
+
+        // TODO: Handle with iterators.
+        for (i, j) in indices.into_iter().enumerate() {
+            let scaled_lidar_ndarray_object = (&interior_points.slice(s![i, ..])
+                - &c.slice(s![..3]))
+                * scale_factor;
+            let scaled_lidar_ndarray_egovehicle =
+                &scaled_lidar_ndarray_object + &c.slice(s![..3]);
+
+            lidar_ndarray
+                .slice_mut(s![j, ..])
+                .assign(&scaled_lidar_ndarray_egovehicle);
+        }
+        c.slice_mut(s![3..6]).mapv_inplace(|x| x * scale_factor);
+    });
+
+    let lidar_column_names = vec!["x", "y", "z"];
+    let series_vec = ndarray_to_expr_vec(lidar_ndarray, lidar_column_names);
+    let augmented_lidar = lidar.lazy().with_columns(series_vec).collect().unwrap();
+
+    let series_vec = ndarray_to_expr_vec(cuboids_ndarray, cuboid_column_names.to_vec());
+    let augmented_cuboids = cuboids.lazy().with_columns(series_vec).collect().unwrap();
+    (augmented_lidar, augmented_cuboids)
+}

rust/src/geometry/mod.rs

@@ -6,6 +6,8 @@
 pub mod augmentations;
 /// Camera models.
 pub mod camera;
+/// Geometric algorithms for polytopes.
+pub mod polytope;
 /// Special Euclidean Group 3.
 pub mod se3;
 /// Special Orthogonal Group 3.

rust/src/geometry/polytope.rs

@@ -0,0 +1,90 @@
+//! # polygons
+//!
+//! Geometric algorithms for polygon geometries.
+
+use ndarray::{concatenate, par_azip, s, Array, ArrayView, Axis, Ix1, Ix2, Ix3, Slice};
+use once_cell::sync::Lazy;
+
+use super::so3::quat_to_mat3;
+
+// Safety: 24 elements (8 * 3 = 24) are defined.
+static VERTS: Lazy<Array<f32, Ix2>> = Lazy::new(|| unsafe {
+    Array::<f32, Ix2>::from_shape_vec_unchecked(
+        (8, 3),
+        vec![
+            1., 1., 1., 1., -1., 1., 1., -1., -1., 1., 1., -1., -1., 1., 1., -1., -1., 1., -1.,
+            -1., -1., -1., 1., -1.,
+        ],
+    )
+});
+
+/// Compute a boolean mask indicating which points are interior to the cuboid geometry.
+pub fn compute_interior_points_mask(
+    points: &ArrayView<f32, Ix2>,
+    cuboid_vertices: &ArrayView<f32, Ix3>,
+) -> Array<bool, Ix2> {
+    let num_points = points.shape()[0];
+    let num_cuboids = cuboid_vertices.shape()[0];
+
+    let a = cuboid_vertices.slice_axis(Axis(1), Slice::from(6..7));
+    let b = cuboid_vertices.slice_axis(Axis(1), Slice::from(3..4));
+    let c = cuboid_vertices.slice_axis(Axis(1), Slice::from(1..2));
+    let vertices = concatenate![Axis(1), a, b, c];
+
+    let reference_index = cuboid_vertices
+        .slice_axis(Axis(1), Slice::from(2..3))
+        .to_owned();
+
+    let uvw = reference_index.clone() - vertices.clone();
+    let reference_index = reference_index.into_shape((num_cuboids, 3)).unwrap();
+
+    let mut dot_uvw_reference = Array::<f32, Ix2>::zeros((num_cuboids, 3));
+    par_azip!((mut a in dot_uvw_reference.outer_iter_mut(), b in uvw.outer_iter(), c in reference_index.outer_iter()) a.assign(&b.dot(&c.t())) );
+
+    let mut dot_uvw_vertices = Array::<f32, Ix2>::zeros((num_cuboids, 3));
+    par_azip!((mut a in dot_uvw_vertices.outer_iter_mut(), b in uvw.outer_iter(), c in vertices.outer_iter()) a.assign(&b.dot(&c.t()).diag()) );
+
+    let dot_uvw_points = uvw
+        .into_shape((num_cuboids * 3, 3))
+        .unwrap()
+        .as_standard_layout()
+        .dot(&points.t().as_standard_layout())
+        .into_shape((num_cuboids, 3, num_points))
+        .unwrap();
+
+    let shape = (num_cuboids, num_points);
+    let mut is_interior =
+        Array::<_, Ix2>::from_shape_vec(shape, vec![false; num_cuboids * num_points]).unwrap();
+    par_azip!((mut a in is_interior.outer_iter_mut(), b in dot_uvw_reference.outer_iter(), c in dot_uvw_points.outer_iter(), d in dot_uvw_vertices.outer_iter()) {
+
+        let c0 = c.slice(s![0, ..]).mapv(|x| ((b[0] <= x) & (x <= d[0])) | ((b[0] >= x) & (x >= d[0])));
+        let c1 = c.slice(s![1, ..]).mapv(|x| ((b[1] <= x) & (x <= d[1])) | ((b[1] >= x) & (x >= d[1])));
+        let c2 = c.slice(s![2, ..]).mapv(|x| ((b[2] <= x) & (x <= d[2])) | ((b[2] >= x) & (x >= d[2])));
+
+        let is_interior_i = &c0 & &c1 & &c2;
+        a.assign(&is_interior_i);
+    });
+
+    is_interior
+}
+
+/// Convert (N,10) cuboids to polygons.
+pub fn cuboids_to_polygons(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))
+    });
+    polygons
+}
+
+/// Convert a single cuboid to a polygon.
+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 verts = &VERTS.clone() * &dims_lwh / 2.;
+    let verts = verts.dot(&mat.t()) + center_xyz;
+    verts.as_standard_layout().to_owned()
+}