epoint_core/
octree.rs

1use crate::{Error, PointCloud, PointData};
2use ecoord::HasAabb;
3use ecoord::octree::{OctantIndex, Octree, StorageMode};
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        storage_mode: StorageMode,
34        shuffle_seed_number: Option<u64>,
35    ) -> Result<(), Error> {
36        let all_points: Vec<PointWithIndex> = self
37            .get_all_points()
38            .into_iter()
39            .enumerate()
40            .map(|(index, point)| PointWithIndex { index, point })
41            .collect();
42
43        let octree = Octree::new(
44            all_points,
45            max_items_per_octant,
46            storage_mode,
47            shuffle_seed_number,
48        )?;
49
50        let octant_indices: Vec<(usize, OctantIndex)> = octree
51            .cells()
52            .iter()
53            .flat_map(|(octant_index, points)| points.iter().map(move |p| (p.index, *octant_index)))
54            .sorted_by_key(|(i, _)| *i)
55            .collect();
56        self.add_octant_indices(octant_indices.into_iter().map(|p| p.1).collect())?;
57
58        Ok(())
59    }
60}
61
62#[derive(Debug, Clone)]
63pub struct PointCloudOctree {
64    pub point_cloud: PointCloud,
65    pub octree: Octree<PointWithIndex>,
66}
67
68impl PointCloudOctree {
69    pub fn new(
70        point_cloud: PointCloud,
71        max_points_per_octant: usize,
72        storage_mode: StorageMode,
73        shuffle_seed_number: Option<u64>,
74    ) -> Result<Self, Error> {
75        let all_points: Vec<PointWithIndex> = point_cloud
76            .point_data
77            .get_all_points()
78            .into_iter()
79            .enumerate()
80            .map(|(index, point)| PointWithIndex { index, point })
81            .collect();
82        let octree = Octree::new(
83            all_points,
84            max_points_per_octant,
85            storage_mode,
86            shuffle_seed_number,
87        )?;
88
89        let point_cloud_octree = PointCloudOctree {
90            point_cloud,
91            octree,
92        };
93        Ok(point_cloud_octree)
94    }
95
96    /// Returns the set of octant indices that contain data.
97    pub fn cell_indices(&self) -> HashSet<OctantIndex> {
98        self.octree.cell_indices()
99    }
100
101    pub fn extract_octant(&self, index: OctantIndex) -> Result<PointCloud, Error> {
102        let cell_content = self
103            .octree
104            .cell(index)
105            .ok_or(Error::NoData("selected octant"))?;
106        let indices = cell_content.iter().map(|p| p.index as u32);
107        let idx_series = polars::prelude::UInt32Chunked::from_iter_values("idx".into(), indices);
108        let filtered_data_frame = self.point_cloud.point_data.data_frame.take(&idx_series)?;
109
110        PointCloud::from_data_frame(
111            filtered_data_frame,
112            self.point_cloud.info.clone(),
113            self.point_cloud.transform_tree.clone(),
114        )
115    }
116}
epoint_core/octree.rs

epoint_core/
octree.rs
