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 let points: Vec<BinPoint> = unsafe {
86 let buf = buf.into_boxed_slice().into_vec();
88
89 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}