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/// Output encoding for PCD writers.
10#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
11pub enum PcdWriteFormat {
12    /// ASCII PCD.
13    #[default]
14    Ascii,
15    /// Binary little-endian PCD.
16    Binary,
17}
18
19/// Writes point clouds to PCD files or streams.
20pub struct PcdWriter<W: Write> {
21    writer: W,
22    format: PcdWriteFormat,
23}
24
25impl<W: Write> PcdWriter<W> {
26    /// Creates a new PCD writer.
27    #[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
44/// Writes a point cloud to a PCD stream.
45pub 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
61/// Writes a point cloud to a PCD file on disk.
62pub 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}