[Refactoring] depth_buffer | log_depth_image | log_pinhole_depth (#11)

* refactor depth_buffer into camera

* rename pinhole_depth to log_pinhole_depth

* parallel computation of rotated_rays in multithreading

* optimize aabb test of ray_cast, change api of log_pinhole_depth

* change api of log_depth_image

* fix docs of log_pinhole_depth

src/lib.rs

@@ -2114,6 +2114,8 @@ pub struct Camera {
     pub aspect_ratio: f32,
     /// The ray directions of each pixel in the camera
     pub ray_directions: Vec<Vector3<f32>>,
+    /// Depth buffer
+    pub depth_buffer: Vec<f32>,
 }
 /// Implementation of the camera
 impl Camera {
@@ -2146,6 +2148,8 @@ impl Camera {
             (width as f32 / height as f32 * (fov_vertical / 2.0).tan()).atan() * 2.0;
         let horizontal_focal_length = (width as f32 / 2.0) / ((fov_horizontal / 2.0).tan());
         let vertical_focal_length = (height as f32 / 2.0) / ((fov_vertical / 2.0).tan());
+        let depth_buffer = vec![0.0; width * height];
+
         Self {
             resolution,
             fov_vertical,
@@ -2156,8 +2160,10 @@ impl Camera {
             far,
             aspect_ratio,
             ray_directions,
+            depth_buffer,
         }
     }
+
     /// Renders the depth of the scene from the perspective of the quadrotor
     /// # Arguments
     /// * `quad_position` - The position of the quadrotor
@@ -2170,20 +2176,18 @@ impl Camera {
     /// ```
     /// use peng_quad::{Camera, Maze};
     /// use nalgebra::{Vector3, UnitQuaternion};
-    /// let camera = Camera::new((800, 600), 60.0, 0.1, 100.0);
+    /// let mut camera = Camera::new((800, 600), 60.0, 0.1, 100.0);
     /// 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 mut depth_buffer = vec![0.0; 800 * 600];
     /// let use_multithreading = true;
-    /// camera.render_depth(&quad_position, &quad_orientation, &maze, &mut depth_buffer, use_multithreading);
+    /// camera.render_depth(&quad_position, &quad_orientation, &maze, use_multithreading);
     /// ```
     pub fn render_depth(
-        &self,
+        &mut self,
         quad_position: &Vector3<f32>,
         quad_orientation: &UnitQuaternion<f32>,
         maze: &Maze,
-        depth_buffer: &mut Vec<f32>,
         use_multi_threading: bool,
     ) -> Result<(), SimulationError> {
         let (width, height) = self.resolution;
@@ -2194,8 +2198,7 @@ impl Camera {
 
         const CHUNK_SIZE: usize = 64;
         if use_multi_threading {
-            depth_buffer.reserve((total_pixels - depth_buffer.capacity()).max(0));
-            depth_buffer
+            self.depth_buffer
                 .par_chunks_mut(CHUNK_SIZE)
                 .enumerate()
                 .try_for_each(|(chunk_idx, chunk)| {
@@ -2205,101 +2208,113 @@ impl Camera {
                         if ray_idx >= total_pixels {
                             break;
                         }
-                        let direction = rotation_camera_to_world * self.ray_directions[ray_idx];
-                        *depth = self.ray_cast(
+                        *depth = ray_cast(
                             quad_position,
                             &rotation_world_to_camera,
-                            &direction,
+                            &(rotation_camera_to_world * self.ray_directions[ray_idx]),
                             maze,
+                            self.near,
+                            self.far,
                         )?;
                     }
                     Ok::<(), SimulationError>(())
                 })?;
         } else {
-            depth_buffer.clear();
-            depth_buffer.reserve((total_pixels - depth_buffer.capacity()).max(0));
             for i in 0..total_pixels {
-                depth_buffer.push(self.ray_cast(
+                self.depth_buffer[i] = ray_cast(
                     quad_position,
                     &rotation_world_to_camera,
                     &(rotation_camera_to_world * self.ray_directions[i]),
                     maze,
-                )?);
+                    self.near,
+                    self.far,
+                )?;
             }
         }
         Ok(())
     }
-    /// Casts a ray from the camera origin in the given direction
-    /// # Arguments
-    /// * `origin` - The origin of the ray
-    /// * `rotation_world_to_camera` - The rotation matrix from world to camera coordinates
-    /// * `direction` - The direction of the ray
-    /// * `maze` - The maze in the scene
-    /// # Returns
-    /// * The distance to the closest obstacle hit by the ray
-    /// # Errors
-    /// * If the ray does not hit any obstacles
-    /// # Example
-    /// ```
-    /// use peng_quad::{Camera, Maze};
-    /// use nalgebra::{Vector3, Matrix3};
-    /// let camera = Camera::new((800, 600), 60.0, 0.1, 100.0);
-    /// let origin = Vector3::new(0.0, 0.0, 0.0);
-    /// let rotation_world_to_camera = Matrix3::identity();
-    /// let direction = Vector3::new(1.0, 0.0, 0.0);
-    /// 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 distance = camera.ray_cast(&origin, &rotation_world_to_camera, &direction, &maze);
-    /// ```
-    pub fn ray_cast(
-        &self,
-        origin: &Vector3<f32>,
-        rotation_world_to_camera: &Matrix3<f32>,
-        direction: &Vector3<f32>,
-        maze: &Maze,
-    ) -> Result<f32, SimulationError> {
-        let mut closest_hit = self.far;
-        // Inline tube intersection
-        for axis in 0..3 {
-            if direction[axis].abs() > f32::EPSILON {
-                for &bound in &[maze.lower_bounds[axis], maze.upper_bounds[axis]] {
-                    let t = (bound - origin[axis]) / direction[axis];
-                    if t > self.near && t < closest_hit {
-                        let intersection_point = origin + direction * t;
-                        if (0..3).all(|i| {
-                            i == axis
-                                || (intersection_point[i] >= maze.lower_bounds[i]
-                                    && intersection_point[i] <= maze.upper_bounds[i])
-                        }) {
-                            closest_hit = t;
-                        }
-                    }
+}
+
+/// Casts a ray from the camera origin in the given direction
+/// # Arguments
+/// * `origin` - The origin of the ray
+/// * `rotation_world_to_camera` - The rotation matrix from world to camera coordinates
+/// * `direction` - The direction of the ray
+/// * `maze` - The maze in the scene
+/// * `near` - The minimum distance to consider
+/// * `far` - The maximum distance to consider
+/// # Returns
+/// * The distance to the closest obstacle hit by the ray
+/// # Errors
+/// * If the ray does not hit any obstacles
+/// # Example
+/// ```
+/// use peng_quad::{ray_cast, Maze};
+/// use nalgebra::{Vector3, Matrix3};
+/// let origin = Vector3::new(0.0, 0.0, 0.0);
+/// let rotation_world_to_camera = Matrix3::identity();
+/// let direction = Vector3::new(0.0, 0.0, 1.0);
+/// let 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 near = 0.1;
+/// let far = 100.0;
+/// let distance = ray_cast(&origin, &rotation_world_to_camera, &direction, &maze, near, far);
+/// ```
+pub fn ray_cast(
+    origin: &Vector3<f32>,
+    rotation_world_to_camera: &Matrix3<f32>,
+    direction: &Vector3<f32>,
+    maze: &Maze,
+    near: f32,
+    far: f32,
+) -> Result<f32, SimulationError> {
+    let mut closest_hit = far;
+    // Inline tube intersection
+    for axis in 0..3 {
+        let t_near = if direction[axis] > 0.0 {
+            (maze.upper_bounds[axis] - origin[axis]) / direction[axis]
+        } else {
+            (maze.lower_bounds[axis] - origin[axis]) / direction[axis]
+        };
+        if t_near > near && t_near < closest_hit {
+            let intersection_point = origin + direction * t_near;
+            let mut valid = true;
+            for i in 0..3 {
+                if i != axis
+                    && (intersection_point[i] < maze.lower_bounds[i]
+                        || intersection_point[i] > maze.upper_bounds[i])
+                {
+                    valid = false;
+                    break;
                 }
             }
-        }
-        // Early exit if we've hit a wall closer than any possible obstacle
-        if closest_hit <= self.near {
-            return Ok(f32::INFINITY);
-        }
-        // Inline sphere intersection
-        for obstacle in &maze.obstacles {
-            let oc = origin - obstacle.position;
-            let b = oc.dot(direction);
-            let c = oc.dot(&oc) - obstacle.radius * obstacle.radius;
-            let discriminant = b * b - c;
-            if discriminant >= 0.0 {
-                // let t = -b - discriminant.sqrt();
-                let t = -b - fast_sqrt(discriminant);
-                if t > self.near && t < closest_hit {
-                    closest_hit = t;
-                }
+            if valid {
+                closest_hit = t_near;
             }
         }
-        if closest_hit < self.far {
-            Ok((rotation_world_to_camera * direction * closest_hit).x)
-        } else {
-            Ok(f32::INFINITY)
+    }
+    // Early exit if we've hit a wall closer than any possible obstacle
+    if closest_hit <= near {
+        return Ok(f32::INFINITY);
+    }
+    // Inline sphere intersection
+    for obstacle in &maze.obstacles {
+        let oc = origin - obstacle.position;
+        let b = oc.dot(direction);
+        let c = oc.dot(&oc) - obstacle.radius * obstacle.radius;
+        let discriminant = b * b - c;
+        if discriminant >= 0.0 {
+            // let t = -b - discriminant.sqrt();
+            let t = -b - fast_sqrt(discriminant);
+            if t > near && t < closest_hit {
+                closest_hit = t;
+            }
         }
     }
+    if closest_hit < far {
+        Ok((rotation_world_to_camera * direction * closest_hit).x)
+    } else {
+        Ok(f32::INFINITY)
+    }
 }
 /// Logs simulation data to the rerun recording stream
 /// # Arguments
@@ -2572,27 +2587,20 @@ pub fn log_mesh(
 /// log depth image data to the rerun recording stream
 /// # Arguments
 /// * `rec` - The rerun::RecordingStream instance
-/// * `depth_image` - The depth image data
-/// * `resolution` - The width and height of the depth image
-/// * `min_depth` - The minimum depth value
-/// * `max_depth` - The maximum depth value
+/// * `cam` - The Camera instance
 /// # Errors
 /// * If the data cannot be logged to the recording stream
 /// # Example
 /// ```no_run
-/// use peng_quad::log_depth_image;
+/// use peng_quad::{log_depth_image, Camera};
 /// let rec = rerun::RecordingStreamBuilder::new("log.rerun").connect().unwrap();
-/// let depth_image = vec![0.0; 640 * 480];
-/// log_depth_image(&rec, &depth_image, (640usize, 480usize), 0.0, 1.0).unwrap();
+/// let camera = Camera::new((640, 480), 0.1, 100.0, 60.0);
+/// log_depth_image(&rec, &camera).unwrap();
 /// ```
-pub fn log_depth_image(
-    rec: &rerun::RecordingStream,
-    depth_image: &[f32],
-    resolution: (usize, usize),
-    min_depth: f32,
-    max_depth: f32,
-) -> Result<(), SimulationError> {
-    let (width, height) = resolution;
+pub fn log_depth_image(rec: &rerun::RecordingStream, cam: &Camera) -> 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 depth_range = max_depth - min_depth;
     image
@@ -2625,29 +2633,28 @@ pub fn log_depth_image(
 /// * `cam_position` - The position vector of the camera (aligns with the quad)
 /// * `cam_orientation` - The orientation quaternion of quad
 /// * `cam_transform` - The transform matrix between quad and camera alignment
-/// * `depth_image` - The depth image data
 /// # Errors
 /// * If the data cannot be logged to the recording stream
 /// # Example
 /// ```no_run
-/// use peng_quad::{pinhole_depth, Camera};
+/// use peng_quad::{log_pinhole_depth, Camera};
 /// use nalgebra::{Vector3, UnitQuaternion};
 /// let rec = rerun::RecordingStreamBuilder::new("log.rerun").connect().unwrap();
 /// let depth_image = vec![ 0.0f32 ; 640 * 480];
 /// let cam_position = Vector3::new(0.0,0.0,0.0);
 /// let cam_orientation = UnitQuaternion::identity();
 /// 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);
-/// pinhole_depth(&rec, &camera, cam_position, cam_orientation, cam_transform, &depth_image).unwrap();
+/// log_pinhole_depth(&rec, &camera, cam_position, cam_orientation, cam_transform).unwrap();
 
-pub fn pinhole_depth(
+pub fn log_pinhole_depth(
     rec: &rerun::RecordingStream,
     cam: &Camera,
     cam_position: Vector3<f32>,
     cam_orientation: UnitQuaternion<f32>,
     cam_transform: [f32; 9],
-    depth_image: &[f32],
 ) -> Result<(), SimulationError> {
+    let depth_image = &cam.depth_buffer;
     let (width, height) = cam.resolution;
     let pinhole_camera = rerun::Pinhole::from_focal_length_and_resolution(
         (cam.horizontal_focal_length, cam.vertical_focal_length),

src/main.rs

@@ -52,15 +52,14 @@ fn main() -> Result<(), SimulationError> {
         config.maze.obstacles_velocity_bounds,
         config.maze.obstacles_radius_bounds,
     );
-    let camera = Camera::new(
+    let mut camera = Camera::new(
         config.camera.resolution,
         config.camera.fov_vertical.to_radians(),
         config.camera.near,
         config.camera.far,
     );
     let mut planner_manager = PlannerManager::new(Vector3::zeros(), 0.0);
     let mut trajectory = Trajectory::new(Vector3::new(0.0, 0.0, 0.0));
-    let mut depth_buffer: Vec<f32> = vec![0.0; camera.resolution.0 * camera.resolution.1];
     let mut previous_thrust = 0.0;
     let mut previous_torque = Vector3::zeros();
     let planner_config: Vec<PlannerStepConfig> = config
@@ -144,14 +143,13 @@ fn main() -> Result<(), SimulationError> {
         }
         imu.update(quad.time_step)?;
         let (true_accel, true_gyro) = quad.read_imu()?;
-        let (_measured_accel, _measured_gyro) = imu.read(true_accel, true_gyro)?;
+        let (measured_accel, measured_gyro) = imu.read(true_accel, true_gyro)?;
         if i % (config.simulation.simulation_frequency / config.simulation.log_frequency) == 0 {
             if config.render_depth {
                 camera.render_depth(
                     &quad.position,
                     &quad.orientation,
                     &maze,
-                    &mut depth_buffer,
                     config.use_multithreading_depth_rendering,
                 )?;
             }
@@ -165,24 +163,17 @@ fn main() -> Result<(), SimulationError> {
                     &quad,
                     &desired_position,
                     &desired_velocity,
-                    &_measured_accel,
-                    &_measured_gyro,
+                    &measured_accel,
+                    &measured_gyro,
                 )?;
                 if config.render_depth {
-                    log_depth_image(
-                        rec,
-                        &depth_buffer,
-                        camera.resolution,
-                        camera.near,
-                        camera.far,
-                    )?;
-                    pinhole_depth(
+                    log_depth_image(rec, &camera)?;
+                    log_pinhole_depth(
                         rec,
                         &camera,
                         quad.position,
                         quad.orientation,
                         config.camera.rotation_transform,
-                        &depth_buffer,
                     )?;
                 }
                 log_maze_obstacles(rec, &maze)?;