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#[cfg(feature = "streaming")]
10use crate::streaming::records_io;
11#[cfg(feature = "streaming")]
12use spatialrust_records::{
13 BoundedSpatialRecordSink, RecordsError, RecordsResult, SchemaDescriptor, SpatialRecordChunk,
14};
15
16#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
18pub enum PcdWriteFormat {
19 #[default]
21 Ascii,
22 Binary,
24}
25
26pub struct PcdWriter<W: Write> {
28 writer: W,
29 format: PcdWriteFormat,
30}
31
32impl<W: Write> PcdWriter<W> {
33 #[must_use]
35 pub const fn new(writer: W, format: PcdWriteFormat) -> Self {
36 Self { writer, format }
37 }
38}
39
40impl<W: Write> PointWriter for PcdWriter<W> {
41 fn write(
42 &mut self,
43 cloud: &PointCloud,
44 _options: &WriteOptions,
45 ) -> spatialrust_core::SpatialResult<()> {
46 write_pcd(&mut self.writer, cloud, self.format)
47 .map_err(|error| spatialrust_core::SpatialError::Io(error.to_string()))
48 }
49}
50
51pub fn write_pcd<W: Write>(
53 writer: &mut W,
54 cloud: &PointCloud,
55 format: PcdWriteFormat,
56) -> Result<(), IoError> {
57 cloud.validate()?;
58 let specs = pcd_specs_from_schema(cloud.schema())?;
59 write_header(writer, cloud, &specs, format)?;
60
61 match format {
62 PcdWriteFormat::Ascii => write_ascii_payload(writer, cloud, &specs)?,
63 PcdWriteFormat::Binary => write_binary_payload(writer, cloud, &specs)?,
64 }
65 Ok(())
66}
67
68pub fn write_pcd_file(
70 path: impl AsRef<std::path::Path>,
71 cloud: &PointCloud,
72 format: PcdWriteFormat,
73) -> Result<(), IoError> {
74 let file = std::fs::File::create(path)?;
75 let mut writer = std::io::BufWriter::new(file);
76 write_pcd(&mut writer, cloud, format)
77}
78
79fn write_header<W: Write>(
80 writer: &mut W,
81 cloud: &PointCloud,
82 specs: &[PcdFieldSpec],
83 format: PcdWriteFormat,
84) -> Result<(), IoError> {
85 let data = match format {
86 PcdWriteFormat::Ascii => "ascii",
87 PcdWriteFormat::Binary => "binary",
88 };
89
90 writeln!(writer, "# .PCD v0.7 - Point Cloud Data file format")?;
91 writeln!(writer, "VERSION 0.7")?;
92 write_list_line(writer, "FIELDS", specs.iter().map(|spec| spec.name.as_str()))?;
93 write_list_line(writer, "SIZE", specs.iter().map(|spec| spec.size.to_string()))?;
94 write_list_line(
95 writer,
96 "TYPE",
97 specs.iter().map(|spec| match spec.kind {
98 PcdType::I => "I",
99 PcdType::U => "U",
100 PcdType::F => "F",
101 }),
102 )?;
103 write_list_line(writer, "COUNT", specs.iter().map(|spec| spec.count.to_string()))?;
104 writeln!(writer, "WIDTH {}", cloud.len())?;
105 writeln!(writer, "HEIGHT 1")?;
106 writeln!(writer, "VIEWPOINT 0 0 0 1 0 0 0")?;
107 writeln!(writer, "POINTS {}", cloud.len())?;
108 writeln!(writer, "DATA {data}")?;
109 Ok(())
110}
111
112#[cfg(feature = "streaming")]
113fn write_stream_header<W: Write>(
114 writer: &mut W,
115 schema: &PointSchema,
116 point_count: u64,
117 specs: &[PcdFieldSpec],
118 format: PcdWriteFormat,
119) -> Result<(), IoError> {
120 let count = usize::try_from(point_count)
121 .map_err(|_| pcd_format("PCD point count does not fit usize"))?;
122 let data = match format {
123 PcdWriteFormat::Ascii => "ascii",
124 PcdWriteFormat::Binary => "binary",
125 };
126 schema.validate_positions()?;
127 writeln!(writer, "# .PCD v0.7 - Point Cloud Data file format")?;
128 writeln!(writer, "VERSION 0.7")?;
129 write_list_line(writer, "FIELDS", specs.iter().map(|spec| spec.name.as_str()))?;
130 write_list_line(writer, "SIZE", specs.iter().map(|spec| spec.size.to_string()))?;
131 write_list_line(
132 writer,
133 "TYPE",
134 specs.iter().map(|spec| match spec.kind {
135 PcdType::I => "I",
136 PcdType::U => "U",
137 PcdType::F => "F",
138 }),
139 )?;
140 write_list_line(writer, "COUNT", specs.iter().map(|spec| spec.count.to_string()))?;
141 writeln!(writer, "WIDTH {count}")?;
142 writeln!(writer, "HEIGHT 1")?;
143 writeln!(writer, "VIEWPOINT 0 0 0 1 0 0 0")?;
144 writeln!(writer, "POINTS {count}")?;
145 writeln!(writer, "DATA {data}")?;
146 Ok(())
147}
148
149#[cfg(feature = "streaming")]
151pub struct PcdChunkSink<W: Write> {
152 writer: W,
153 schema: SchemaDescriptor,
154 specs: Vec<PcdFieldSpec>,
155 format: PcdWriteFormat,
156 expected_points: u64,
157 written_points: u64,
158 finished: bool,
159}
160
161#[cfg(feature = "streaming")]
162impl<W: Write> PcdChunkSink<W> {
163 pub fn new(
165 mut writer: W,
166 schema: SchemaDescriptor,
167 expected_points: u64,
168 format: PcdWriteFormat,
169 ) -> Result<Self, IoError> {
170 let specs = pcd_specs_from_schema(schema.point_schema())?;
171 write_stream_header(&mut writer, schema.point_schema(), expected_points, &specs, format)?;
172 Ok(Self {
173 writer,
174 schema,
175 specs,
176 format,
177 expected_points,
178 written_points: 0,
179 finished: false,
180 })
181 }
182}
183
184#[cfg(feature = "streaming")]
185impl PcdChunkSink<std::io::BufWriter<std::fs::File>> {
186 pub fn create(
188 path: impl AsRef<std::path::Path>,
189 schema: SchemaDescriptor,
190 expected_points: u64,
191 format: PcdWriteFormat,
192 ) -> Result<Self, IoError> {
193 Self::new(
194 std::io::BufWriter::new(std::fs::File::create(path)?),
195 schema,
196 expected_points,
197 format,
198 )
199 }
200}
201
202#[cfg(feature = "streaming")]
203impl<W: Write> BoundedSpatialRecordSink for PcdChunkSink<W> {
204 fn write_chunk(&mut self, chunk: &SpatialRecordChunk) -> RecordsResult<()> {
205 if self.finished {
206 return Err(RecordsError::InvalidConfiguration("PCD sink is already finished".into()));
207 }
208 let cloud = chunk.record().cloud();
209 if cloud.schema() != self.schema.point_schema() {
210 return Err(RecordsError::SchemaMismatch(
211 "PCD chunk schema differs from sink schema".into(),
212 ));
213 }
214 let next = self
215 .written_points
216 .checked_add(
217 u64::try_from(cloud.len())
218 .map_err(|_| RecordsError::ReceiptOverflow("PCD chunk point count".into()))?,
219 )
220 .ok_or_else(|| RecordsError::ReceiptOverflow("PCD point count".into()))?;
221 if next > self.expected_points {
222 return Err(RecordsError::InvalidChunk(format!(
223 "PCD sink expected {} points but received at least {next}",
224 self.expected_points
225 )));
226 }
227 match self.format {
228 PcdWriteFormat::Ascii => write_ascii_payload(&mut self.writer, cloud, &self.specs),
229 PcdWriteFormat::Binary => write_binary_payload(&mut self.writer, cloud, &self.specs),
230 }
231 .map_err(records_io)?;
232 self.written_points = next;
233 Ok(())
234 }
235
236 fn finish(&mut self) -> RecordsResult<()> {
237 if self.written_points != self.expected_points {
238 return Err(RecordsError::InvalidChunk(format!(
239 "PCD sink expected {} points but received {}",
240 self.expected_points, self.written_points
241 )));
242 }
243 self.writer.flush().map_err(IoError::from).map_err(records_io)?;
244 self.finished = true;
245 Ok(())
246 }
247}
248
249fn write_list_line<W: Write, I, S>(writer: &mut W, key: &str, values: I) -> Result<(), IoError>
250where
251 I: IntoIterator<Item = S>,
252 S: AsRef<str>,
253{
254 write!(writer, "{key}")?;
255 for value in values {
256 write!(writer, " {}", value.as_ref())?;
257 }
258 writeln!(writer)?;
259 Ok(())
260}
261
262fn write_ascii_payload<W: Write>(
263 writer: &mut W,
264 cloud: &PointCloud,
265 specs: &[PcdFieldSpec],
266) -> Result<(), IoError> {
267 for point_index in 0..cloud.len() {
268 let mut first = true;
269 for spec in specs {
270 if !first {
271 write!(writer, " ")?;
272 }
273 first = false;
274 if spec.name.eq_ignore_ascii_case("rgb") {
275 let r = read_scalar(cloud, "r", point_index)? as u32;
276 let g = read_scalar(cloud, "g", point_index)? as u32;
277 let b = read_scalar(cloud, "b", point_index)? as u32;
278 let packed = (r << 16) | (g << 8) | b;
279 write!(writer, "{packed}")?;
280 continue;
281 }
282 write!(writer, "{}", read_scalar(cloud, &spec.name, point_index)?)?;
283 }
284 writeln!(writer)?;
285 }
286 Ok(())
287}
288
289fn write_binary_payload<W: Write>(
290 writer: &mut W,
291 cloud: &PointCloud,
292 specs: &[PcdFieldSpec],
293) -> Result<(), IoError> {
294 for point_index in 0..cloud.len() {
295 for spec in specs {
296 if spec.name.eq_ignore_ascii_case("rgb") {
297 let r = read_scalar(cloud, "r", point_index)? as u32;
298 let g = read_scalar(cloud, "g", point_index)? as u32;
299 let b = read_scalar(cloud, "b", point_index)? as u32;
300 let packed = (r << 16) | (g << 8) | b;
301 writer.write_all(&packed.to_le_bytes())?;
302 continue;
303 }
304 write_binary_scalar(writer, cloud, spec, point_index)?;
305 }
306 }
307 Ok(())
308}
309
310fn write_binary_scalar<W: Write>(
311 writer: &mut W,
312 cloud: &PointCloud,
313 spec: &PcdFieldSpec,
314 point_index: usize,
315) -> Result<(), IoError> {
316 let field = cloud
317 .schema()
318 .fields()
319 .iter()
320 .find(|field| field.name == spec.name)
321 .ok_or_else(|| pcd_format(format!("missing field `{}`", spec.name)))?;
322 let buffer = cloud.field(&field.name).map_err(IoError::from)?;
323
324 match (field.dtype, spec.kind, spec.size) {
325 (DType::F32 | DType::F16, PcdType::F, 4) => {
326 let value = buffer.as_f32().map_err(IoError::from)?[point_index];
327 writer.write_all(&value.to_le_bytes())?;
328 }
329 (DType::F64, PcdType::F, 8) => {
330 let PointBuffer::F64(values) = buffer else {
331 return Err(IoError::Core(spatialrust_core::SpatialError::UnsupportedDType(
332 field.dtype,
333 )));
334 };
335 writer.write_all(&values[point_index].to_le_bytes())?;
336 }
337 (DType::I32, PcdType::I, 4) => {
338 let PointBuffer::I32(values) = buffer else {
339 return Err(IoError::Core(spatialrust_core::SpatialError::UnsupportedDType(
340 field.dtype,
341 )));
342 };
343 writer.write_all(&values[point_index].to_le_bytes())?;
344 }
345 (DType::U8, PcdType::U, 1) => {
346 let PointBuffer::U8(values) = buffer else {
347 return Err(IoError::Core(spatialrust_core::SpatialError::UnsupportedDType(
348 field.dtype,
349 )));
350 };
351 writer.write_all(&[values[point_index]])?;
352 }
353 (DType::U16, PcdType::U, 2) => {
354 let PointBuffer::U16(values) = buffer else {
355 return Err(IoError::Core(spatialrust_core::SpatialError::UnsupportedDType(
356 field.dtype,
357 )));
358 };
359 writer.write_all(&values[point_index].to_le_bytes())?;
360 }
361 (DType::U32, PcdType::U, 4) => {
362 let PointBuffer::U32(values) = buffer else {
363 return Err(IoError::Core(spatialrust_core::SpatialError::UnsupportedDType(
364 field.dtype,
365 )));
366 };
367 writer.write_all(&values[point_index].to_le_bytes())?;
368 }
369 _ => return Err(pcd_format(format!("cannot encode field `{}` to PCD", field.name))),
370 }
371 Ok(())
372}
373
374use spatialrust_core::PointBuffer;
375
376fn read_scalar(cloud: &PointCloud, name: &str, point_index: usize) -> Result<f32, IoError> {
377 let field = cloud
378 .schema()
379 .fields()
380 .iter()
381 .find(|field| field.name == name)
382 .ok_or_else(|| pcd_format(format!("missing field `{name}`")))?;
383 let buffer = cloud.field(name).map_err(IoError::from)?;
384 let value = match field.dtype {
385 DType::F32 | DType::F16 => buffer.as_f32().map_err(IoError::from)?[point_index],
386 DType::F64 => {
387 let PointBuffer::F64(values) = buffer else {
388 return Err(IoError::Core(spatialrust_core::SpatialError::UnsupportedDType(
389 field.dtype,
390 )));
391 };
392 values[point_index] as f32
393 }
394 DType::U8 => {
395 let PointBuffer::U8(values) = buffer else {
396 return Err(IoError::Core(spatialrust_core::SpatialError::UnsupportedDType(
397 field.dtype,
398 )));
399 };
400 f32::from(values[point_index])
401 }
402 DType::U16 => {
403 let PointBuffer::U16(values) = buffer else {
404 return Err(IoError::Core(spatialrust_core::SpatialError::UnsupportedDType(
405 field.dtype,
406 )));
407 };
408 f32::from(values[point_index])
409 }
410 DType::I32 => {
411 let PointBuffer::I32(values) = buffer else {
412 return Err(IoError::Core(spatialrust_core::SpatialError::UnsupportedDType(
413 field.dtype,
414 )));
415 };
416 values[point_index] as f32
417 }
418 DType::U32 => {
419 let PointBuffer::U32(values) = buffer else {
420 return Err(IoError::Core(spatialrust_core::SpatialError::UnsupportedDType(
421 field.dtype,
422 )));
423 };
424 values[point_index] as f32
425 }
426 };
427 Ok(value)
428}
429
430fn pcd_specs_from_schema(schema: &PointSchema) -> Result<Vec<PcdFieldSpec>, IoError> {
431 let mut specs = Vec::new();
432 let mut index = 0;
433 while index < schema.len() {
434 let field = &schema.fields()[index];
435 if matches!(
436 field.semantic,
437 FieldSemantic::ColorR | FieldSemantic::ColorG | FieldSemantic::ColorB
438 ) && field.semantic == FieldSemantic::ColorR
439 && index + 2 < schema.len()
440 && schema.fields()[index + 1].semantic == FieldSemantic::ColorG
441 && schema.fields()[index + 2].semantic == FieldSemantic::ColorB
442 {
443 specs.push(PcdFieldSpec { name: "rgb".into(), size: 4, kind: PcdType::F, count: 1 });
444 index += 3;
445 continue;
446 }
447
448 specs.push(pcd_spec_from_field(field)?);
449 index += 1;
450 }
451 Ok(specs)
452}
453
454fn pcd_spec_from_field(field: &PointField) -> Result<PcdFieldSpec, IoError> {
455 let name = match field.semantic {
456 FieldSemantic::PositionX => "x",
457 FieldSemantic::PositionY => "y",
458 FieldSemantic::PositionZ => "z",
459 FieldSemantic::NormalX => "normal_x",
460 FieldSemantic::NormalY => "normal_y",
461 FieldSemantic::NormalZ => "normal_z",
462 FieldSemantic::Intensity => "intensity",
463 FieldSemantic::Curvature => "curvature",
464 FieldSemantic::Ring => "ring",
465 FieldSemantic::TimeOffset => "timestamp",
466 FieldSemantic::Label => "label",
467 _ => field.name.as_str(),
468 }
469 .to_owned();
470
471 let (kind, size) = match field.dtype {
472 DType::F32 | DType::F16 => (PcdType::F, 4),
473 DType::F64 => (PcdType::F, 8),
474 DType::I32 => (PcdType::I, 4),
475 DType::U8 => (PcdType::U, 1),
476 DType::U16 => (PcdType::U, 2),
477 DType::U32 => (PcdType::U, 4),
478 };
479
480 let _semantic = infer_field_semantic(&name);
481 Ok(PcdFieldSpec { name, size, kind, count: field.components })
482}