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#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
12pub enum PlyWriteFormat {
13 #[default]
15 Ascii,
16 BinaryLittleEndian,
18}
19
20pub struct PlyWriter<W: Write> {
22 writer: W,
23 format: PlyWriteFormat,
24}
25
26impl<W: Write> PlyWriter<W> {
27 #[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
45pub 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
72pub 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}