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#[cfg(feature = "streaming")]
11use crate::streaming::records_io;
12#[cfg(feature = "streaming")]
13use spatialrust_records::{
14    BoundedSpatialRecordSink, RecordsError, RecordsResult, SchemaDescriptor, SpatialRecordChunk,
15};
16
17/// Output encoding for PLY writers.
18#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
19pub enum PlyWriteFormat {
20    /// ASCII PLY.
21    #[default]
22    Ascii,
23    /// Binary little-endian PLY.
24    BinaryLittleEndian,
25}
26
27/// Writes point clouds to PLY files or streams.
28pub struct PlyWriter<W: Write> {
29    writer: W,
30    format: PlyWriteFormat,
31}
32
33impl<W: Write> PlyWriter<W> {
34    /// Creates a new PLY writer.
35    #[must_use]
36    pub const fn new(writer: W, format: PlyWriteFormat) -> Self {
37        Self { writer, format }
38    }
39}
40
41impl<W: Write> PointWriter for PlyWriter<W> {
42    fn write(
43        &mut self,
44        cloud: &PointCloud,
45        _options: &WriteOptions,
46    ) -> spatialrust_core::SpatialResult<()> {
47        write_ply(&mut self.writer, cloud, self.format)
48            .map_err(|error| spatialrust_core::SpatialError::Io(error.to_string()))
49    }
50}
51
52/// Writes a point cloud to a PLY stream.
53pub fn write_ply<W: Write>(
54    writer: &mut W,
55    cloud: &PointCloud,
56    format: PlyWriteFormat,
57) -> Result<(), IoError> {
58    cloud.validate()?;
59    let properties = ply_properties_from_schema(cloud.schema())?;
60    let header = PlyHeader {
61        format: match format {
62            PlyWriteFormat::Ascii => PlyFormat::Ascii,
63            PlyWriteFormat::BinaryLittleEndian => PlyFormat::BinaryLittleEndian,
64        },
65        vertex_count: cloud.len(),
66        properties,
67    };
68    header.write_header(writer)?;
69
70    match format {
71        PlyWriteFormat::Ascii => write_ascii_vertices(writer, cloud, &header.properties)?,
72        PlyWriteFormat::BinaryLittleEndian => {
73            write_binary_vertices(writer, cloud, &header.properties)?
74        }
75    }
76    Ok(())
77}
78
79/// Writes a point cloud to a PLY file on disk.
80pub fn write_ply_file(
81    path: impl AsRef<std::path::Path>,
82    cloud: &PointCloud,
83    format: PlyWriteFormat,
84) -> Result<(), IoError> {
85    let file = std::fs::File::create(path.as_ref())?;
86    let mut writer = std::io::BufWriter::new(file);
87    write_ply(&mut writer, cloud, format)
88}
89
90fn ply_properties_from_schema(schema: &PointSchema) -> Result<Vec<PlyProperty>, IoError> {
91    schema.fields().iter().map(ply_property_from_field).collect()
92}
93
94/// Sequential PLY sink with an exact upfront vertex-count contract.
95#[cfg(feature = "streaming")]
96pub struct PlyChunkSink<W: Write> {
97    writer: W,
98    schema: SchemaDescriptor,
99    properties: Vec<PlyProperty>,
100    format: PlyWriteFormat,
101    expected_points: u64,
102    written_points: u64,
103    finished: bool,
104}
105
106#[cfg(feature = "streaming")]
107impl<W: Write> PlyChunkSink<W> {
108    /// Writes the PLY header and creates a chunk sink.
109    pub fn new(
110        mut writer: W,
111        schema: SchemaDescriptor,
112        expected_points: u64,
113        format: PlyWriteFormat,
114    ) -> Result<Self, IoError> {
115        let properties = ply_properties_from_schema(schema.point_schema())?;
116        let vertex_count = usize::try_from(expected_points)
117            .map_err(|_| ply_format("PLY vertex count does not fit usize"))?;
118        PlyHeader {
119            format: match format {
120                PlyWriteFormat::Ascii => PlyFormat::Ascii,
121                PlyWriteFormat::BinaryLittleEndian => PlyFormat::BinaryLittleEndian,
122            },
123            vertex_count,
124            properties: properties.clone(),
125        }
126        .write_header(&mut writer)?;
127        Ok(Self {
128            writer,
129            schema,
130            properties,
131            format,
132            expected_points,
133            written_points: 0,
134            finished: false,
135        })
136    }
137}
138
139#[cfg(feature = "streaming")]
140impl PlyChunkSink<std::io::BufWriter<std::fs::File>> {
141    /// Creates a local PLY file sink.
142    pub fn create(
143        path: impl AsRef<std::path::Path>,
144        schema: SchemaDescriptor,
145        expected_points: u64,
146        format: PlyWriteFormat,
147    ) -> Result<Self, IoError> {
148        Self::new(
149            std::io::BufWriter::new(std::fs::File::create(path)?),
150            schema,
151            expected_points,
152            format,
153        )
154    }
155}
156
157#[cfg(feature = "streaming")]
158impl<W: Write> BoundedSpatialRecordSink for PlyChunkSink<W> {
159    fn write_chunk(&mut self, chunk: &SpatialRecordChunk) -> RecordsResult<()> {
160        if self.finished {
161            return Err(RecordsError::InvalidConfiguration("PLY sink is already finished".into()));
162        }
163        let cloud = chunk.record().cloud();
164        if cloud.schema() != self.schema.point_schema() {
165            return Err(RecordsError::SchemaMismatch(
166                "PLY chunk schema differs from sink schema".into(),
167            ));
168        }
169        let next = self
170            .written_points
171            .checked_add(
172                u64::try_from(cloud.len())
173                    .map_err(|_| RecordsError::ReceiptOverflow("PLY chunk point count".into()))?,
174            )
175            .ok_or_else(|| RecordsError::ReceiptOverflow("PLY point count".into()))?;
176        if next > self.expected_points {
177            return Err(RecordsError::InvalidChunk(format!(
178                "PLY sink expected {} points but received at least {next}",
179                self.expected_points
180            )));
181        }
182        match self.format {
183            PlyWriteFormat::Ascii => {
184                write_ascii_vertices(&mut self.writer, cloud, &self.properties)
185            }
186            PlyWriteFormat::BinaryLittleEndian => {
187                write_binary_vertices(&mut self.writer, cloud, &self.properties)
188            }
189        }
190        .map_err(records_io)?;
191        self.written_points = next;
192        Ok(())
193    }
194
195    fn finish(&mut self) -> RecordsResult<()> {
196        if self.written_points != self.expected_points {
197            return Err(RecordsError::InvalidChunk(format!(
198                "PLY sink expected {} points but received {}",
199                self.expected_points, self.written_points
200            )));
201        }
202        self.writer.flush().map_err(IoError::from).map_err(records_io)?;
203        self.finished = true;
204        Ok(())
205    }
206}
207
208fn write_ascii_vertices<W: Write>(
209    writer: &mut W,
210    cloud: &PointCloud,
211    properties: &[PlyProperty],
212) -> Result<(), IoError> {
213    for point_index in 0..cloud.len() {
214        let mut first = true;
215        for property in properties {
216            if !first {
217                write!(writer, " ")?;
218            }
219            first = false;
220            write!(writer, "{}", read_scalar(cloud, property, point_index)?)?;
221        }
222        writeln!(writer)?;
223    }
224    Ok(())
225}
226
227fn write_binary_vertices<W: Write>(
228    writer: &mut W,
229    cloud: &PointCloud,
230    properties: &[PlyProperty],
231) -> Result<(), IoError> {
232    for point_index in 0..cloud.len() {
233        for property in properties {
234            write_binary_scalar(writer, cloud, property, point_index)?;
235        }
236    }
237    Ok(())
238}
239
240fn write_binary_scalar<W: Write>(
241    writer: &mut W,
242    cloud: &PointCloud,
243    property: &PlyProperty,
244    point_index: usize,
245) -> Result<(), IoError> {
246    let field = find_field_for_property(cloud.schema(), property)?;
247    let buffer = cloud.field(&field.name).map_err(IoError::from)?;
248
249    match (field.dtype, property.kind) {
250        (DType::F32 | DType::F16, PlyPropertyKind::Float) => {
251            let value = buffer.as_f32().map_err(IoError::from)?[point_index];
252            writer.write_all(&value.to_le_bytes())?;
253        }
254        (DType::F64, PlyPropertyKind::Double) => {
255            let PointBuffer::F64(values) = buffer else {
256                return Err(IoError::Core(spatialrust_core::SpatialError::UnsupportedDType(
257                    field.dtype,
258                )));
259            };
260            writer.write_all(&values[point_index].to_le_bytes())?;
261        }
262        (DType::I32, PlyPropertyKind::Int) => {
263            let PointBuffer::I32(values) = buffer else {
264                return Err(IoError::Core(spatialrust_core::SpatialError::UnsupportedDType(
265                    field.dtype,
266                )));
267            };
268            writer.write_all(&values[point_index].to_le_bytes())?;
269        }
270        (DType::U32, PlyPropertyKind::UInt) => {
271            let PointBuffer::U32(values) = buffer else {
272                return Err(IoError::Core(spatialrust_core::SpatialError::UnsupportedDType(
273                    field.dtype,
274                )));
275            };
276            writer.write_all(&values[point_index].to_le_bytes())?;
277        }
278        (DType::U8, PlyPropertyKind::UChar) => {
279            let PointBuffer::U8(values) = buffer else {
280                return Err(IoError::Core(spatialrust_core::SpatialError::UnsupportedDType(
281                    field.dtype,
282                )));
283            };
284            writer.write_all(&[values[point_index]])?;
285        }
286        (DType::U16, PlyPropertyKind::UShort) => {
287            let PointBuffer::U16(values) = buffer else {
288                return Err(IoError::Core(spatialrust_core::SpatialError::UnsupportedDType(
289                    field.dtype,
290                )));
291            };
292            writer.write_all(&values[point_index].to_le_bytes())?;
293        }
294        _ => return Err(ply_format(format!("cannot encode field `{}` to PLY", field.name))),
295    }
296    Ok(())
297}
298
299fn read_scalar(
300    cloud: &PointCloud,
301    property: &PlyProperty,
302    point_index: usize,
303) -> Result<f32, IoError> {
304    let field = find_field_for_property(cloud.schema(), property)?;
305    let buffer = cloud.field(&field.name).map_err(IoError::from)?;
306    let value = match field.dtype {
307        DType::F32 | DType::F16 => buffer.as_f32().map_err(IoError::from)?[point_index],
308        DType::F64 => {
309            let PointBuffer::F64(values) = buffer else {
310                return Err(IoError::Core(spatialrust_core::SpatialError::UnsupportedDType(
311                    field.dtype,
312                )));
313            };
314            values[point_index] as f32
315        }
316        DType::U8 => {
317            let PointBuffer::U8(values) = buffer else {
318                return Err(IoError::Core(spatialrust_core::SpatialError::UnsupportedDType(
319                    field.dtype,
320                )));
321            };
322            f32::from(values[point_index])
323        }
324        DType::U16 => {
325            let PointBuffer::U16(values) = buffer else {
326                return Err(IoError::Core(spatialrust_core::SpatialError::UnsupportedDType(
327                    field.dtype,
328                )));
329            };
330            f32::from(values[point_index])
331        }
332        DType::I32 => {
333            let PointBuffer::I32(values) = buffer else {
334                return Err(IoError::Core(spatialrust_core::SpatialError::UnsupportedDType(
335                    field.dtype,
336                )));
337            };
338            values[point_index] as f32
339        }
340        DType::U32 => {
341            let PointBuffer::U32(values) = buffer else {
342                return Err(IoError::Core(spatialrust_core::SpatialError::UnsupportedDType(
343                    field.dtype,
344                )));
345            };
346            values[point_index] as f32
347        }
348    };
349    Ok(value)
350}
351
352fn find_field_for_property<'a>(
353    schema: &'a PointSchema,
354    property: &PlyProperty,
355) -> Result<&'a spatialrust_core::PointField, IoError> {
356    let semantic = infer_property_semantic(&property.name);
357    schema
358        .fields()
359        .iter()
360        .find(|field| {
361            field.name == property.name
362                || (semantic != FieldSemantic::Unknown && field.semantic == semantic)
363        })
364        .ok_or_else(|| ply_format(format!("missing field for PLY property `{}`", property.name)))
365}
spatialrust_io/ply/writer.rs

spatialrust_io/ply/
writer.rs
