1use std::io::Write;
2
3use spatialrust_core::{DType, FieldSemantic, PointCloud, PointField, PointSchema};
4
5use crate::error::{pcd_format, IoError};
6use crate::pcd::schema::{infer_field_semantic, PcdFieldSpec, PcdType};
7use crate::{PointWriter, WriteOptions};
8
9#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
11pub enum PcdWriteFormat {
12 #[default]
14 Ascii,
15 Binary,
17}
18
19pub struct PcdWriter<W: Write> {
21 writer: W,
22 format: PcdWriteFormat,
23}
24
25impl<W: Write> PcdWriter<W> {
26 #[must_use]
28 pub const fn new(writer: W, format: PcdWriteFormat) -> Self {
29 Self { writer, format }
30 }
31}
32
33impl<W: Write> PointWriter for PcdWriter<W> {
34 fn write(
35 &mut self,
36 cloud: &PointCloud,
37 _options: &WriteOptions,
38 ) -> spatialrust_core::SpatialResult<()> {
39 write_pcd(&mut self.writer, cloud, self.format)
40 .map_err(|error| spatialrust_core::SpatialError::Io(error.to_string()))
41 }
42}
43
44pub fn write_pcd<W: Write>(
46 writer: &mut W,
47 cloud: &PointCloud,
48 format: PcdWriteFormat,
49) -> Result<(), IoError> {
50 cloud.validate()?;
51 let specs = pcd_specs_from_schema(cloud.schema())?;
52 write_header(writer, cloud, &specs, format)?;
53
54 match format {
55 PcdWriteFormat::Ascii => write_ascii_payload(writer, cloud, &specs)?,
56 PcdWriteFormat::Binary => write_binary_payload(writer, cloud, &specs)?,
57 }
58 Ok(())
59}
60
61pub fn write_pcd_file(
63 path: impl AsRef<std::path::Path>,
64 cloud: &PointCloud,
65 format: PcdWriteFormat,
66) -> Result<(), IoError> {
67 let file = std::fs::File::create(path)?;
68 let mut writer = std::io::BufWriter::new(file);
69 write_pcd(&mut writer, cloud, format)
70}
71
72fn write_header<W: Write>(
73 writer: &mut W,
74 cloud: &PointCloud,
75 specs: &[PcdFieldSpec],
76 format: PcdWriteFormat,
77) -> Result<(), IoError> {
78 let data = match format {
79 PcdWriteFormat::Ascii => "ascii",
80 PcdWriteFormat::Binary => "binary",
81 };
82
83 writeln!(writer, "# .PCD v0.7 - Point Cloud Data file format")?;
84 writeln!(writer, "VERSION 0.7")?;
85 write_list_line(writer, "FIELDS", specs.iter().map(|spec| spec.name.as_str()))?;
86 write_list_line(writer, "SIZE", specs.iter().map(|spec| spec.size.to_string()))?;
87 write_list_line(
88 writer,
89 "TYPE",
90 specs.iter().map(|spec| match spec.kind {
91 PcdType::I => "I",
92 PcdType::U => "U",
93 PcdType::F => "F",
94 }),
95 )?;
96 write_list_line(writer, "COUNT", specs.iter().map(|spec| spec.count.to_string()))?;
97 writeln!(writer, "WIDTH {}", cloud.len())?;
98 writeln!(writer, "HEIGHT 1")?;
99 writeln!(writer, "VIEWPOINT 0 0 0 1 0 0 0")?;
100 writeln!(writer, "POINTS {}", cloud.len())?;
101 writeln!(writer, "DATA {data}")?;
102 Ok(())
103}
104
105fn write_list_line<W: Write, I, S>(writer: &mut W, key: &str, values: I) -> Result<(), IoError>
106where
107 I: IntoIterator<Item = S>,
108 S: AsRef<str>,
109{
110 write!(writer, "{key}")?;
111 for value in values {
112 write!(writer, " {}", value.as_ref())?;
113 }
114 writeln!(writer)?;
115 Ok(())
116}
117
118fn write_ascii_payload<W: Write>(
119 writer: &mut W,
120 cloud: &PointCloud,
121 specs: &[PcdFieldSpec],
122) -> Result<(), IoError> {
123 for point_index in 0..cloud.len() {
124 let mut first = true;
125 for spec in specs {
126 if !first {
127 write!(writer, " ")?;
128 }
129 first = false;
130 if spec.name.eq_ignore_ascii_case("rgb") {
131 let r = read_scalar(cloud, "r", point_index)? as u32;
132 let g = read_scalar(cloud, "g", point_index)? as u32;
133 let b = read_scalar(cloud, "b", point_index)? as u32;
134 let packed = (r << 16) | (g << 8) | b;
135 write!(writer, "{packed}")?;
136 continue;
137 }
138 write!(writer, "{}", read_scalar(cloud, &spec.name, point_index)?)?;
139 }
140 writeln!(writer)?;
141 }
142 Ok(())
143}
144
145fn write_binary_payload<W: Write>(
146 writer: &mut W,
147 cloud: &PointCloud,
148 specs: &[PcdFieldSpec],
149) -> Result<(), IoError> {
150 for point_index in 0..cloud.len() {
151 for spec in specs {
152 if spec.name.eq_ignore_ascii_case("rgb") {
153 let r = read_scalar(cloud, "r", point_index)? as u32;
154 let g = read_scalar(cloud, "g", point_index)? as u32;
155 let b = read_scalar(cloud, "b", point_index)? as u32;
156 let packed = (r << 16) | (g << 8) | b;
157 writer.write_all(&packed.to_le_bytes())?;
158 continue;
159 }
160 write_binary_scalar(writer, cloud, spec, point_index)?;
161 }
162 }
163 Ok(())
164}
165
166fn write_binary_scalar<W: Write>(
167 writer: &mut W,
168 cloud: &PointCloud,
169 spec: &PcdFieldSpec,
170 point_index: usize,
171) -> Result<(), IoError> {
172 let field = cloud
173 .schema()
174 .fields()
175 .iter()
176 .find(|field| field.name == spec.name)
177 .ok_or_else(|| pcd_format(format!("missing field `{}`", spec.name)))?;
178 let buffer = cloud.field(&field.name).map_err(IoError::from)?;
179
180 match (field.dtype, spec.kind, spec.size) {
181 (DType::F32 | DType::F16, PcdType::F, 4) => {
182 let value = buffer.as_f32().map_err(IoError::from)?[point_index];
183 writer.write_all(&value.to_le_bytes())?;
184 }
185 (DType::F64, PcdType::F, 8) => {
186 let PointBuffer::F64(values) = buffer else {
187 return Err(IoError::Core(spatialrust_core::SpatialError::UnsupportedDType(
188 field.dtype,
189 )));
190 };
191 writer.write_all(&values[point_index].to_le_bytes())?;
192 }
193 (DType::I32, PcdType::I, 4) => {
194 let PointBuffer::I32(values) = buffer else {
195 return Err(IoError::Core(spatialrust_core::SpatialError::UnsupportedDType(
196 field.dtype,
197 )));
198 };
199 writer.write_all(&values[point_index].to_le_bytes())?;
200 }
201 (DType::U8, PcdType::U, 1) => {
202 let PointBuffer::U8(values) = buffer else {
203 return Err(IoError::Core(spatialrust_core::SpatialError::UnsupportedDType(
204 field.dtype,
205 )));
206 };
207 writer.write_all(&[values[point_index]])?;
208 }
209 (DType::U16, PcdType::U, 2) => {
210 let PointBuffer::U16(values) = buffer else {
211 return Err(IoError::Core(spatialrust_core::SpatialError::UnsupportedDType(
212 field.dtype,
213 )));
214 };
215 writer.write_all(&values[point_index].to_le_bytes())?;
216 }
217 (DType::U32, PcdType::U, 4) => {
218 let PointBuffer::U32(values) = buffer else {
219 return Err(IoError::Core(spatialrust_core::SpatialError::UnsupportedDType(
220 field.dtype,
221 )));
222 };
223 writer.write_all(&values[point_index].to_le_bytes())?;
224 }
225 _ => return Err(pcd_format(format!("cannot encode field `{}` to PCD", field.name))),
226 }
227 Ok(())
228}
229
230use spatialrust_core::PointBuffer;
231
232fn read_scalar(cloud: &PointCloud, name: &str, point_index: usize) -> Result<f32, IoError> {
233 let field = cloud
234 .schema()
235 .fields()
236 .iter()
237 .find(|field| field.name == name)
238 .ok_or_else(|| pcd_format(format!("missing field `{name}`")))?;
239 let buffer = cloud.field(name).map_err(IoError::from)?;
240 let value = match field.dtype {
241 DType::F32 | DType::F16 => buffer.as_f32().map_err(IoError::from)?[point_index],
242 DType::F64 => {
243 let PointBuffer::F64(values) = buffer else {
244 return Err(IoError::Core(spatialrust_core::SpatialError::UnsupportedDType(
245 field.dtype,
246 )));
247 };
248 values[point_index] as f32
249 }
250 DType::U8 => {
251 let PointBuffer::U8(values) = buffer else {
252 return Err(IoError::Core(spatialrust_core::SpatialError::UnsupportedDType(
253 field.dtype,
254 )));
255 };
256 f32::from(values[point_index])
257 }
258 DType::U16 => {
259 let PointBuffer::U16(values) = buffer else {
260 return Err(IoError::Core(spatialrust_core::SpatialError::UnsupportedDType(
261 field.dtype,
262 )));
263 };
264 f32::from(values[point_index])
265 }
266 DType::I32 => {
267 let PointBuffer::I32(values) = buffer else {
268 return Err(IoError::Core(spatialrust_core::SpatialError::UnsupportedDType(
269 field.dtype,
270 )));
271 };
272 values[point_index] as f32
273 }
274 DType::U32 => {
275 let PointBuffer::U32(values) = buffer else {
276 return Err(IoError::Core(spatialrust_core::SpatialError::UnsupportedDType(
277 field.dtype,
278 )));
279 };
280 values[point_index] as f32
281 }
282 };
283 Ok(value)
284}
285
286fn pcd_specs_from_schema(schema: &PointSchema) -> Result<Vec<PcdFieldSpec>, IoError> {
287 let mut specs = Vec::new();
288 let mut index = 0;
289 while index < schema.len() {
290 let field = &schema.fields()[index];
291 if matches!(
292 field.semantic,
293 FieldSemantic::ColorR | FieldSemantic::ColorG | FieldSemantic::ColorB
294 ) && field.semantic == FieldSemantic::ColorR
295 && index + 2 < schema.len()
296 && schema.fields()[index + 1].semantic == FieldSemantic::ColorG
297 && schema.fields()[index + 2].semantic == FieldSemantic::ColorB
298 {
299 specs.push(PcdFieldSpec { name: "rgb".into(), size: 4, kind: PcdType::F, count: 1 });
300 index += 3;
301 continue;
302 }
303
304 specs.push(pcd_spec_from_field(field)?);
305 index += 1;
306 }
307 Ok(specs)
308}
309
310fn pcd_spec_from_field(field: &PointField) -> Result<PcdFieldSpec, IoError> {
311 let name = match field.semantic {
312 FieldSemantic::PositionX => "x",
313 FieldSemantic::PositionY => "y",
314 FieldSemantic::PositionZ => "z",
315 FieldSemantic::NormalX => "normal_x",
316 FieldSemantic::NormalY => "normal_y",
317 FieldSemantic::NormalZ => "normal_z",
318 FieldSemantic::Intensity => "intensity",
319 FieldSemantic::Curvature => "curvature",
320 FieldSemantic::Ring => "ring",
321 FieldSemantic::TimeOffset => "timestamp",
322 FieldSemantic::Label => "label",
323 _ => field.name.as_str(),
324 }
325 .to_owned();
326
327 let (kind, size) = match field.dtype {
328 DType::F32 | DType::F16 => (PcdType::F, 4),
329 DType::F64 => (PcdType::F, 8),
330 DType::I32 => (PcdType::I, 4),
331 DType::U8 => (PcdType::U, 1),
332 DType::U16 => (PcdType::U, 2),
333 DType::U32 => (PcdType::U, 4),
334 };
335
336 let _semantic = infer_field_semantic(&name);
337 Ok(PcdFieldSpec { name, size, kind, count: field.components })
338}