nuscenes_data_pcd/
lib.rs

1use anyhow::{ensure, Result};
2use nuscenes_data::{dataset::SampleDataRef, serializable::FileFormat};
3use pcd_rs::{PcdDeserialize, PcdSerialize};
4use raw_parts::RawParts;
5use std::{
6    fs::File,
7    io::{prelude::*, BufReader},
8    mem,
9};
10
11pub mod prelude {
12    pub use super::SampleDataRefPcdExt;
13}
14
15#[derive(Debug, Clone, PartialEq)]
16pub enum PointCloud {
17    Pcd(Vec<PcdPoint>),
18    Bin(Vec<BinPoint>),
19    NotSupported,
20}
21
22#[derive(Debug, Clone, PartialEq, PcdSerialize, PcdDeserialize)]
23pub struct PcdPoint {
24    pub x: f32,
25    pub y: f32,
26    pub z: f32,
27    pub dyn_prop: i8,
28    pub id: i16,
29    pub rcs: f32,
30    pub vx: f32,
31    pub vy: f32,
32    pub vx_comp: f32,
33    pub vy_comp: f32,
34    pub is_quality_valid: i8,
35    pub ambig_state: i8,
36    pub x_rms: i8,
37    pub y_rms: i8,
38    pub invalid_state: i8,
39    pub pdh0: i8,
40    pub vx_rms: i8,
41    pub vy_rms: i8,
42}
43
44#[derive(Debug, Clone, PartialEq)]
45#[repr(packed)]
46pub struct BinPoint {
47    pub x: f32,
48    pub y: f32,
49    pub z: f32,
50    pub intensity: f32,
51    pub ring_index: i32,
52}
53
54pub trait SampleDataRefPcdExt {
55    fn load_pcd(&self) -> Result<PointCloud>;
56}
57
58impl SampleDataRefPcdExt for SampleDataRef {
59    fn load_pcd(&self) -> Result<PointCloud> {
60        if self.fileformat != FileFormat::Pcd {
61            return Ok(PointCloud::NotSupported);
62        }
63
64        let Some(ext) = self.filename.extension() else {
65            return Ok(PointCloud::NotSupported)
66        };
67        let path = self.path();
68
69        let pcd = if ext == "pcd" {
70            let reader = pcd_rs::Reader::open(path)?;
71            let points: Result<Vec<_>> = reader.collect();
72            PointCloud::Pcd(points?)
73        } else if ext == "bin" {
74            let point_len = mem::size_of::<BinPoint>();
75
76            let buf = {
77                let mut reader = BufReader::new(File::open(&path)?);
78                let mut buf = vec![];
79                let buf_len = reader.read_to_end(&mut buf)?;
80                ensure!(buf_len % point_len == 0, "Unable to load this file {}. The file size is {buf_len}, which is not multiple of {point_len}", path.display());
81                buf
82            };
83
84            // Transmute the byte vec to vec of points
85            let points: Vec<BinPoint> = unsafe {
86                // make sure the capacity is equal to the length of the buffer.
87                let buf = buf.into_boxed_slice().into_vec();
88
89                // transmute the vec
90                let RawParts {
91                    ptr,
92                    length,
93                    capacity,
94                } = RawParts::from_vec(buf);
95                debug_assert_eq!(length, capacity);
96
97                RawParts {
98                    ptr: ptr as *mut BinPoint,
99                    length: length / point_len,
100                    capacity: capacity / point_len,
101                }
102                .into_vec()
103            };
104
105            PointCloud::Bin(points)
106        } else {
107            PointCloud::NotSupported
108        };
109
110        Ok(pcd)
111    }
112}