epoint_io/xyz/
write.rs

1use crate::Error::{InvalidFileExtension, NoFileName};
2use crate::FILE_EXTENSION_XYZ_FORMAT;
3use crate::error::Error;
4use crate::xyz::{DEFAULT_XYZ_SEPARATOR, FILE_EXTENSION_XYZ_ZST_FORMAT};
5use ecoord::FrameId;
6use epoint_core::PointDataColumnType;
7use epoint_core::point_cloud::PointCloud;
8use palette::Srgb;
9use polars::prelude::{CsvWriter, NamedFrom, SerWriter, Series};
10use rayon::iter::ParallelIterator;
11use rayon::prelude::IntoParallelIterator;
12use std::fs::{File, OpenOptions};
13use std::io::{BufWriter, Write};
14use std::path::Path;
15
16pub const DEFAULT_COMPRESSION_LEVEL: i32 = 10;
17pub const DEFAULT_NULL_VALUE: &str = "NaN";
18
19/// `XyzWriter` exports a point cloud to a non-native representation.
20///
21#[derive(Debug, Clone)]
22pub struct XyzWriter<W: Write> {
23    writer: W,
24    compression_level: Option<i32>,
25    frame_id: Option<FrameId>,
26    separator: u8,
27    null_value: String,
28    color_depth: ColorDepth,
29}
30
31#[derive(Debug, Clone, Copy, Hash, PartialEq, Eq, Default)]
32pub enum ColorDepth {
33    #[default]
34    EightBit,
35    SixteenBit,
36}
37
38impl<W: Write> XyzWriter<W> {
39    pub fn new(writer: W) -> Self {
40        Self {
41            writer,
42            compression_level: Some(crate::epoint::write::DEFAULT_COMPRESSION_LEVEL),
43            frame_id: None,
44            separator: DEFAULT_XYZ_SEPARATOR,
45            null_value: DEFAULT_NULL_VALUE.to_string(),
46            color_depth: ColorDepth::default(),
47        }
48    }
49
50    pub fn with_compressed(mut self, compressed: bool) -> Self {
51        if compressed {
52            self.compression_level = Some(DEFAULT_COMPRESSION_LEVEL);
53        } else {
54            self.compression_level = None;
55        }
56        self
57    }
58
59    pub fn with_frame_id(mut self, frame_id: FrameId) -> Self {
60        self.frame_id = Some(frame_id);
61        self
62    }
63
64    pub fn with_separator(mut self, separator: u8) -> Self {
65        self.separator = separator;
66        self
67    }
68
69    pub fn with_null_value(mut self, null_value: String) -> Self {
70        self.null_value = null_value;
71        self
72    }
73
74    pub fn with_color_depth(mut self, color_depth: ColorDepth) -> Self {
75        self.color_depth = color_depth;
76        self
77    }
78
79    pub fn finish(self, point_cloud: &PointCloud) -> Result<(), Error> {
80        let mut exported_point_cloud = point_cloud.clone();
81        if let Some(frame_id) = &self.frame_id {
82            exported_point_cloud.resolve_to_frame(frame_id.clone())?;
83        }
84        /*let mut resulting_point_cloud: PointCloud =
85        self.frame_id
86            .clone()
87            .map_or(point_cloud.to_owned(), |f: FrameId| {
88                exported_point_cloud
89                    .resolve_to_frame(f)?;
90                exported_point_cloud
91            });*/
92
93        if exported_point_cloud.contains_colors() {
94            match self.color_depth {
95                ColorDepth::EightBit => {
96                    let converted_colors: Vec<Srgb<u8>> = exported_point_cloud
97                        .point_data
98                        .get_all_colors()?
99                        .into_par_iter()
100                        .map(|x| x.into_format())
101                        .collect();
102
103                    let color_red_series = Series::new(
104                        PointDataColumnType::X.into(),
105                        converted_colors.iter().map(|c| c.red).collect::<Vec<u8>>(),
106                    );
107                    exported_point_cloud
108                        .point_data
109                        .data_frame
110                        .replace(PointDataColumnType::ColorRed.as_str(), color_red_series)?;
111
112                    let color_green_series = Series::new(
113                        PointDataColumnType::Y.into(),
114                        converted_colors
115                            .iter()
116                            .map(|c| c.green)
117                            .collect::<Vec<u8>>(),
118                    );
119                    exported_point_cloud
120                        .point_data
121                        .data_frame
122                        .replace(PointDataColumnType::ColorGreen.as_str(), color_green_series)?;
123
124                    let color_blue_series = Series::new(
125                        PointDataColumnType::Z.into(),
126                        converted_colors.iter().map(|c| c.blue).collect::<Vec<u8>>(),
127                    );
128                    exported_point_cloud
129                        .point_data
130                        .data_frame
131                        .replace(PointDataColumnType::ColorBlue.as_str(), color_blue_series)?;
132                }
133                ColorDepth::SixteenBit => {}
134            }
135        }
136
137        let writer: Box<dyn Write> = if let Some(compression_level) = &self.compression_level {
138            let buf_writer = BufWriter::with_capacity(
139                zstd::stream::Encoder::<Vec<u8>>::recommended_input_size(),
140                zstd::stream::Encoder::new(self.writer, *compression_level)?.auto_finish(),
141            );
142            Box::new(buf_writer)
143        } else {
144            Box::new(self.writer)
145        };
146
147        CsvWriter::new(writer)
148            .with_separator(self.separator)
149            .with_null_value(self.null_value)
150            .finish(&mut exported_point_cloud.point_data.data_frame)?;
151
152        Ok(())
153    }
154}
155
156impl XyzWriter<File> {
157    pub fn from_path(path: impl AsRef<Path>) -> Result<Self, Error> {
158        let file_name_str = path
159            .as_ref()
160            .file_name()
161            .ok_or(NoFileName())?
162            .to_string_lossy()
163            .to_lowercase();
164        if !file_name_str.ends_with(FILE_EXTENSION_XYZ_ZST_FORMAT)
165            && !file_name_str.ends_with(FILE_EXTENSION_XYZ_FORMAT)
166        {
167            return Err(InvalidFileExtension(file_name_str.to_string()));
168        }
169
170        let file = OpenOptions::new()
171            .create(true)
172            .write(true)
173            .truncate(true)
174            .open(path)?;
175        Ok(Self::new(file))
176    }
177}
epoint_io/xyz/write.rs

epoint_io/xyz/
write.rs
