Skip to main content

spatialrust_io/pcd/
writer.rs

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/// Output encoding for PCD writers.
17#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
18pub enum PcdWriteFormat {
19    /// ASCII PCD.
20    #[default]
21    Ascii,
22    /// Binary little-endian PCD.
23    Binary,
24}
25
26/// Writes point clouds to PCD files or streams.
27pub struct PcdWriter<W: Write> {
28    writer: W,
29    format: PcdWriteFormat,
30}
31
32impl<W: Write> PcdWriter<W> {
33    /// Creates a new PCD writer.
34    #[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
51/// Writes a point cloud to a PCD stream.
52pub 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
68/// Writes a point cloud to a PCD file on disk.
69pub 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/// Sequential PCD sink with an exact upfront point-count contract.
150#[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    /// Writes the PCD header and creates a chunk sink.
164    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    /// Creates a local PCD file sink.
187    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}