From ede89b7a64d9244cf70c997702ce6836ffca022f Mon Sep 17 00:00:00 2001 From: Yang Zhou Date: Wed, 6 Nov 2024 10:15:21 -0500 Subject: [PATCH] add use_multi_threading for log_depth_image (#12) * add use_multi_threading for log_depth_image * fix docs --- src/lib.rs | 99 ++++++++++++++++++++++++++++++++++++++--------------- src/main.rs | 2 +- 2 files changed, 73 insertions(+), 28 deletions(-) diff --git a/src/lib.rs b/src/lib.rs index 54e3d981..138f023c 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -425,6 +425,8 @@ impl Imu { Ok(()) } /// Simulates IMU readings with added bias and noise + /// + /// The added bias and noise are based on normal distributions /// # Arguments /// * `true_acceleration` - The true acceleration vector /// * `true_angular_velocity` - The true angular velocity vector @@ -457,6 +459,9 @@ impl Imu { } } /// PID controller for quadrotor position and attitude control +/// +/// The kpid_pos and kpid_att gains are following the format of +/// porportional, derivative and integral gains /// # Example /// ``` /// use nalgebra::Vector3; @@ -2130,7 +2135,7 @@ impl Camera { /// # Example /// ``` /// use peng_quad::Camera; - /// let camera = Camera::new((800, 600), 60.0, 0.1, 100.0); + /// let camera = Camera::new((800, 600), 1.0, 5.0, 120.0); /// ``` pub fn new(resolution: (usize, usize), fov_vertical: f32, near: f32, far: f32) -> Self { let (width, height) = resolution; @@ -2165,11 +2170,14 @@ impl Camera { } /// Renders the depth of the scene from the perspective of the quadrotor + /// + /// When the depth value is out of the near and far clipping planes, it is set to infinity + /// When the resolution is larger than 32x24, multi-threading can accelerate the rendering /// # Arguments /// * `quad_position` - The position of the quadrotor /// * `quad_orientation` - The orientation of the quadrotor /// * `maze` - The maze in the scene - /// * `depth_buffer` - The depth buffer to store the depth values + /// * `use_multi_threading` - Whether to use multi-threading to render the depth /// # Errors /// * If the depth buffer is not large enough to store the depth values /// # Example @@ -2180,8 +2188,10 @@ impl Camera { /// let quad_position = Vector3::new(0.0, 0.0, 0.0); /// let quad_orientation = UnitQuaternion::identity(); /// let mut maze = Maze::new([-1.0, -1.0, -1.0], [1.0, 1.0, 1.0], 5, [0.1, 0.1, 0.1], [0.1, 0.5]); - /// let use_multithreading = true; - /// camera.render_depth(&quad_position, &quad_orientation, &maze, use_multithreading); + /// let use_multi_threading = true; + /// camera.render_depth(&quad_position, &quad_orientation, &maze, use_multi_threading); + /// let use_multi_threading = false; + /// camera.render_depth(&quad_position, &quad_orientation, &maze, use_multi_threading); /// ``` pub fn render_depth( &mut self, @@ -2537,7 +2547,7 @@ pub fn log_trajectory( )?; Ok(()) } -/// log mesh data to the rerun recording stream +/// Log mesh data to the rerun recording stream /// # Arguments /// * `rec` - The rerun::RecordingStream instance /// * `division` - The number of divisions in the mesh @@ -2584,10 +2594,14 @@ pub fn log_mesh( )?; Ok(()) } -/// log depth image data to the rerun recording stream +/// Log depth image data to the rerun recording stream +/// +/// When the depth value is `f32::INFINITY`, the pixel is considered invalid and logged as black +/// When the resolution is larger than 32x24, multi-threading can accelerate the rendering /// # Arguments /// * `rec` - The rerun::RecordingStream instance /// * `cam` - The Camera instance +/// * `use_multi_threading` - Whether to use multithreading to log the depth image /// # Errors /// * If the data cannot be logged to the recording stream /// # Example @@ -2595,32 +2609,62 @@ pub fn log_mesh( /// use peng_quad::{log_depth_image, Camera}; /// let rec = rerun::RecordingStreamBuilder::new("log.rerun").connect().unwrap(); /// let camera = Camera::new((640, 480), 0.1, 100.0, 60.0); -/// log_depth_image(&rec, &camera).unwrap(); +/// let use_multi_threading = false; +/// log_depth_image(&rec, &camera, use_multi_threading).unwrap(); +/// let use_multi_threading = true; +/// log_depth_image(&rec, &camera, use_multi_threading).unwrap(); /// ``` -pub fn log_depth_image(rec: &rerun::RecordingStream, cam: &Camera) -> Result<(), SimulationError> { +pub fn log_depth_image( + rec: &rerun::RecordingStream, + cam: &Camera, + use_multi_threading: bool, +) -> Result<(), SimulationError> { let (width, height) = (cam.resolution.0, cam.resolution.1); let (min_depth, max_depth) = (cam.near, cam.far); let depth_image = &cam.depth_buffer; - let mut image = rerun::external::ndarray::Array::zeros((height, width, 3)); + let mut image: rerun::external::ndarray::Array = + rerun::external::ndarray::Array::zeros((height, width, 3)); let depth_range = max_depth - min_depth; - image - .axis_iter_mut(rerun::external::ndarray::Axis(0)) - .enumerate() - .for_each(|(y, mut row)| { - for (x, mut pixel) in row - .axis_iter_mut(rerun::external::ndarray::Axis(0)) - .enumerate() - { - let depth = depth_image[y * width + x]; - let color = if depth.is_finite() { - let normalized_depth = ((depth - min_depth) / depth_range).clamp(0.0, 1.0); - color_map_fn(normalized_depth * 255.0) - } else { - (0, 0, 0) - }; - (pixel[0], pixel[1], pixel[2]) = color; - } - }); + let scale_factor = 255.0 / depth_range; + if use_multi_threading { + const CHUNK_SIZE: usize = 32; + image + .as_slice_mut() + .expect("Failed to get mutable slice of image") + .par_chunks_exact_mut(CHUNK_SIZE * 3) + .enumerate() + .for_each(|(chunk_index, chunk)| { + let start_index = chunk_index * 32; + for (i, pixel) in chunk.chunks_exact_mut(3).enumerate() { + let idx = start_index + i; + let depth = depth_image[idx]; + let color = if depth.is_finite() { + let normalized_depth = + ((depth - min_depth) * scale_factor).clamp(0.0, 255.0); + color_map_fn(normalized_depth) + } else { + (0, 0, 0) + }; + (pixel[0], pixel[1], pixel[2]) = color; + } + }); + } else { + for (index, pixel) in image + .as_slice_mut() + .expect("Failed to get mutable slice of image") + .chunks_exact_mut(3) + .enumerate() + { + let depth = depth_image[index]; + let color = if depth.is_finite() { + let normalized_depth = ((depth - min_depth) * scale_factor).clamp(0.0, 255.0); + color_map_fn(normalized_depth) + } else { + (0, 0, 0) + }; + (pixel[0], pixel[1], pixel[2]) = color; + } + } let rerun_image = rerun::Image::from_color_model_and_tensor(rerun::ColorModel::RGB, image) .map_err(|e| SimulationError::OtherError(format!("Failed to create rerun image: {}", e)))?; rec.log("world/quad/cam/depth", &rerun_image)?; @@ -2646,6 +2690,7 @@ pub fn log_depth_image(rec: &rerun::RecordingStream, cam: &Camera) -> Result<(), /// let cam_transform = [0.0, 0.0, 1.0, -1.0, 0.0, 0.0, 0.0, -1.0, 0.0]; /// let camera = Camera::new((800, 600), 60.0, 0.1, 100.0); /// log_pinhole_depth(&rec, &camera, cam_position, cam_orientation, cam_transform).unwrap(); +/// ``` pub fn log_pinhole_depth( rec: &rerun::RecordingStream, diff --git a/src/main.rs b/src/main.rs index eb3bd6ff..7f3ea4d8 100644 --- a/src/main.rs +++ b/src/main.rs @@ -167,7 +167,7 @@ fn main() -> Result<(), SimulationError> { &measured_gyro, )?; if config.render_depth { - log_depth_image(rec, &camera)?; + log_depth_image(rec, &camera, config.use_multithreading_depth_rendering)?; log_pinhole_depth( rec, &camera,