Skip to main content

spatialrust_io/ply/
writer.rs

1use std::io::Write;
2
3use spatialrust_core::{DType, FieldSemantic, PointBuffer, PointCloud, PointSchema};
4
5use crate::error::{ply_format, IoError};
6use crate::ply::header::{PlyFormat, PlyHeader, PlyProperty, PlyPropertyKind};
7use crate::ply::schema::{infer_property_semantic, ply_property_from_field};
8use crate::{PointWriter, WriteOptions};
9
10/// Output encoding for PLY writers.
11#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
12pub enum PlyWriteFormat {
13    /// ASCII PLY.
14    #[default]
15    Ascii,
16    /// Binary little-endian PLY.
17    BinaryLittleEndian,
18}
19
20/// Writes point clouds to PLY files or streams.
21pub struct PlyWriter<W: Write> {
22    writer: W,
23    format: PlyWriteFormat,
24}
25
26impl<W: Write> PlyWriter<W> {
27    /// Creates a new PLY writer.
28    #[must_use]
29    pub const fn new(writer: W, format: PlyWriteFormat) -> Self {
30        Self { writer, format }
31    }
32}
33
34impl<W: Write> PointWriter for PlyWriter<W> {
35    fn write(
36        &mut self,
37        cloud: &PointCloud,
38        _options: &WriteOptions,
39    ) -> spatialrust_core::SpatialResult<()> {
40        write_ply(&mut self.writer, cloud, self.format)
41            .map_err(|error| spatialrust_core::SpatialError::Io(error.to_string()))
42    }
43}
44
45/// Writes a point cloud to a PLY stream.
46pub fn write_ply<W: Write>(
47    writer: &mut W,
48    cloud: &PointCloud,
49    format: PlyWriteFormat,
50) -> Result<(), IoError> {
51    cloud.validate()?;
52    let properties = ply_properties_from_schema(cloud.schema())?;
53    let header = PlyHeader {
54        format: match format {
55            PlyWriteFormat::Ascii => PlyFormat::Ascii,
56            PlyWriteFormat::BinaryLittleEndian => PlyFormat::BinaryLittleEndian,
57        },
58        vertex_count: cloud.len(),
59        properties,
60    };
61    header.write_header(writer)?;
62
63    match format {
64        PlyWriteFormat::Ascii => write_ascii_vertices(writer, cloud, &header.properties)?,
65        PlyWriteFormat::BinaryLittleEndian => {
66            write_binary_vertices(writer, cloud, &header.properties)?
67        }
68    }
69    Ok(())
70}
71
72/// Writes a point cloud to a PLY file on disk.
73pub fn write_ply_file(
74    path: impl AsRef<std::path::Path>,
75    cloud: &PointCloud,
76    format: PlyWriteFormat,
77) -> Result<(), IoError> {
78    let file = std::fs::File::create(path.as_ref())?;
79    let mut writer = std::io::BufWriter::new(file);
80    write_ply(&mut writer, cloud, format)
81}
82
83fn ply_properties_from_schema(schema: &PointSchema) -> Result<Vec<PlyProperty>, IoError> {
84    schema.fields().iter().map(ply_property_from_field).collect()
85}
86
87fn write_ascii_vertices<W: Write>(
88    writer: &mut W,
89    cloud: &PointCloud,
90    properties: &[PlyProperty],
91) -> Result<(), IoError> {
92    for point_index in 0..cloud.len() {
93        let mut first = true;
94        for property in properties {
95            if !first {
96                write!(writer, " ")?;
97            }
98            first = false;
99            write!(writer, "{}", read_scalar(cloud, property, point_index)?)?;
100        }
101        writeln!(writer)?;
102    }
103    Ok(())
104}
105
106fn write_binary_vertices<W: Write>(
107    writer: &mut W,
108    cloud: &PointCloud,
109    properties: &[PlyProperty],
110) -> Result<(), IoError> {
111    for point_index in 0..cloud.len() {
112        for property in properties {
113            write_binary_scalar(writer, cloud, property, point_index)?;
114        }
115    }
116    Ok(())
117}
118
119fn write_binary_scalar<W: Write>(
120    writer: &mut W,
121    cloud: &PointCloud,
122    property: &PlyProperty,
123    point_index: usize,
124) -> Result<(), IoError> {
125    let field = find_field_for_property(cloud.schema(), property)?;
126    let buffer = cloud.field(&field.name).map_err(IoError::from)?;
127
128    match (field.dtype, property.kind) {
129        (DType::F32 | DType::F16, PlyPropertyKind::Float) => {
130            let value = buffer.as_f32().map_err(IoError::from)?[point_index];
131            writer.write_all(&value.to_le_bytes())?;
132        }
133        (DType::F64, PlyPropertyKind::Double) => {
134            let PointBuffer::F64(values) = buffer else {
135                return Err(IoError::Core(spatialrust_core::SpatialError::UnsupportedDType(
136                    field.dtype,
137                )));
138            };
139            writer.write_all(&values[point_index].to_le_bytes())?;
140        }
141        (DType::I32, PlyPropertyKind::Int) => {
142            let PointBuffer::I32(values) = buffer else {
143                return Err(IoError::Core(spatialrust_core::SpatialError::UnsupportedDType(
144                    field.dtype,
145                )));
146            };
147            writer.write_all(&values[point_index].to_le_bytes())?;
148        }
149        (DType::U32, PlyPropertyKind::UInt) => {
150            let PointBuffer::U32(values) = buffer else {
151                return Err(IoError::Core(spatialrust_core::SpatialError::UnsupportedDType(
152                    field.dtype,
153                )));
154            };
155            writer.write_all(&values[point_index].to_le_bytes())?;
156        }
157        (DType::U8, PlyPropertyKind::UChar) => {
158            let PointBuffer::U8(values) = buffer else {
159                return Err(IoError::Core(spatialrust_core::SpatialError::UnsupportedDType(
160                    field.dtype,
161                )));
162            };
163            writer.write_all(&[values[point_index]])?;
164        }
165        (DType::U16, PlyPropertyKind::UShort) => {
166            let PointBuffer::U16(values) = buffer else {
167                return Err(IoError::Core(spatialrust_core::SpatialError::UnsupportedDType(
168                    field.dtype,
169                )));
170            };
171            writer.write_all(&values[point_index].to_le_bytes())?;
172        }
173        _ => return Err(ply_format(format!("cannot encode field `{}` to PLY", field.name))),
174    }
175    Ok(())
176}
177
178fn read_scalar(
179    cloud: &PointCloud,
180    property: &PlyProperty,
181    point_index: usize,
182) -> Result<f32, IoError> {
183    let field = find_field_for_property(cloud.schema(), property)?;
184    let buffer = cloud.field(&field.name).map_err(IoError::from)?;
185    let value = match field.dtype {
186        DType::F32 | DType::F16 => buffer.as_f32().map_err(IoError::from)?[point_index],
187        DType::F64 => {
188            let PointBuffer::F64(values) = buffer else {
189                return Err(IoError::Core(spatialrust_core::SpatialError::UnsupportedDType(
190                    field.dtype,
191                )));
192            };
193            values[point_index] as f32
194        }
195        DType::U8 => {
196            let PointBuffer::U8(values) = buffer else {
197                return Err(IoError::Core(spatialrust_core::SpatialError::UnsupportedDType(
198                    field.dtype,
199                )));
200            };
201            f32::from(values[point_index])
202        }
203        DType::U16 => {
204            let PointBuffer::U16(values) = buffer else {
205                return Err(IoError::Core(spatialrust_core::SpatialError::UnsupportedDType(
206                    field.dtype,
207                )));
208            };
209            f32::from(values[point_index])
210        }
211        DType::I32 => {
212            let PointBuffer::I32(values) = buffer else {
213                return Err(IoError::Core(spatialrust_core::SpatialError::UnsupportedDType(
214                    field.dtype,
215                )));
216            };
217            values[point_index] as f32
218        }
219        DType::U32 => {
220            let PointBuffer::U32(values) = buffer else {
221                return Err(IoError::Core(spatialrust_core::SpatialError::UnsupportedDType(
222                    field.dtype,
223                )));
224            };
225            values[point_index] as f32
226        }
227    };
228    Ok(value)
229}
230
231fn find_field_for_property<'a>(
232    schema: &'a PointSchema,
233    property: &PlyProperty,
234) -> Result<&'a spatialrust_core::PointField, IoError> {
235    let semantic = infer_property_semantic(&property.name);
236    schema
237        .fields()
238        .iter()
239        .find(|field| {
240            field.name == property.name
241                || (semantic != FieldSemantic::Unknown && field.semantic == semantic)
242        })
243        .ok_or_else(|| ply_format(format!("missing field for PLY property `{}`", property.name)))
244}