epoint_core/
octree.rs

1use crate::{Error, PointCloud, PointData};
2use ecoord::HasAabb;
3use ecoord::octree::{OctantIndex, Octree};
4use itertools::Itertools;
5use nalgebra::Point3;
6use polars::prelude::NewChunkedArray;
7use std::collections::HashSet;
8
9#[derive(Debug, Clone, Copy, PartialEq)]
10pub struct PointWithIndex {
11    index: usize,
12    point: Point3<f64>,
13}
14
15impl HasAabb for PointWithIndex {
16    fn center(&self) -> Point3<f64> {
17        self.point
18    }
19
20    fn min(&self) -> Point3<f64> {
21        self.point
22    }
23
24    fn max(&self) -> Point3<f64> {
25        self.point
26    }
27}
28
29impl PointData {
30    pub fn compute_octree(
31        &mut self,
32        max_items_per_octant: usize,
33        shuffle_seed_number: Option<u64>,
34    ) -> Result<(), Error> {
35        let all_points: Vec<PointWithIndex> = self
36            .get_all_points()
37            .into_iter()
38            .enumerate()
39            .map(|(index, point)| PointWithIndex { index, point })
40            .collect();
41
42        let octree = Octree::new(all_points, max_items_per_octant, shuffle_seed_number)?;
43
44        let octant_indices: Vec<(usize, OctantIndex)> = octree
45            .cells()
46            .iter()
47            .flat_map(|(octant_index, points)| points.iter().map(move |p| (p.index, *octant_index)))
48            .sorted_by_key(|(i, _)| *i)
49            .collect();
50        self.add_octant_indices(octant_indices.into_iter().map(|p| p.1).collect())?;
51
52        Ok(())
53    }
54}
55
56#[derive(Debug, Clone)]
57pub struct PointCloudOctree {
58    pub point_cloud: PointCloud,
59    pub octree: Octree<PointWithIndex>,
60}
61
62impl PointCloudOctree {
63    pub fn new(
64        point_cloud: PointCloud,
65        max_points_per_octant: usize,
66        shuffle_seed_number: Option<u64>,
67    ) -> Result<Self, Error> {
68        let all_points: Vec<PointWithIndex> = point_cloud
69            .point_data
70            .get_all_points()
71            .into_iter()
72            .enumerate()
73            .map(|(index, point)| PointWithIndex { index, point })
74            .collect();
75        let octree = Octree::new(all_points, max_points_per_octant, shuffle_seed_number)?;
76
77        let point_cloud_octree = PointCloudOctree {
78            point_cloud,
79            octree,
80        };
81        Ok(point_cloud_octree)
82    }
83
84    /// Returns the set of octant indices that contain data.
85    pub fn cell_indices(&self) -> HashSet<OctantIndex> {
86        self.octree.cell_indices()
87    }
88
89    pub fn extract_octant(&self, index: OctantIndex) -> Result<PointCloud, Error> {
90        let cell_content = self
91            .octree
92            .cell(index)
93            .ok_or(Error::NoData("selected octant"))?;
94        let indices = cell_content.iter().map(|p| p.index as u32);
95        let idx_series = polars::prelude::UInt32Chunked::from_iter_values("idx".into(), indices);
96        let filtered_data_frame = self.point_cloud.point_data.data_frame.take(&idx_series)?;
97
98        PointCloud::from_data_frame(
99            filtered_data_frame,
100            self.point_cloud.info.clone(),
101            self.point_cloud.reference_frames.clone(),
102        )
103    }
104}
epoint_core/octree.rs

epoint_core/
octree.rs
