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, mut point_cloud: PointCloud) -> Result<(), Error> {
80        if let Some(frame_id) = &self.frame_id {
81            point_cloud.resolve_to_frame(frame_id.clone())?;
82        }
83        /*let mut resulting_point_cloud: PointCloud =
84        self.frame_id
85            .clone()
86            .map_or(point_cloud.to_owned(), |f: FrameId| {
87                point_cloud.resolve_to_frame(f)?;
88                point_cloud
89            });*/
90
91        if point_cloud.contains_colors() {
92            match self.color_depth {
93                ColorDepth::EightBit => {
94                    let converted_colors: Vec<Srgb<u8>> = point_cloud
95                        .point_data
96                        .get_all_colors()?
97                        .into_par_iter()
98                        .map(|x| x.into_format())
99                        .collect();
100
101                    let color_red_series = Series::new(
102                        PointDataColumnType::X.into(),
103                        converted_colors.iter().map(|c| c.red).collect::<Vec<u8>>(),
104                    );
105                    point_cloud
106                        .point_data
107                        .data_frame
108                        .replace(PointDataColumnType::ColorRed.as_str(), color_red_series)?;
109
110                    let color_green_series = Series::new(
111                        PointDataColumnType::Y.into(),
112                        converted_colors
113                            .iter()
114                            .map(|c| c.green)
115                            .collect::<Vec<u8>>(),
116                    );
117                    point_cloud
118                        .point_data
119                        .data_frame
120                        .replace(PointDataColumnType::ColorGreen.as_str(), color_green_series)?;
121
122                    let color_blue_series = Series::new(
123                        PointDataColumnType::Z.into(),
124                        converted_colors.iter().map(|c| c.blue).collect::<Vec<u8>>(),
125                    );
126                    point_cloud
127                        .point_data
128                        .data_frame
129                        .replace(PointDataColumnType::ColorBlue.as_str(), color_blue_series)?;
130                }
131                ColorDepth::SixteenBit => {}
132            }
133        }
134
135        let writer: Box<dyn Write> = if let Some(compression_level) = &self.compression_level {
136            let buf_writer = BufWriter::with_capacity(
137                zstd::stream::Encoder::<Vec<u8>>::recommended_input_size(),
138                zstd::stream::Encoder::new(self.writer, *compression_level)?.auto_finish(),
139            );
140            Box::new(buf_writer)
141        } else {
142            Box::new(self.writer)
143        };
144
145        CsvWriter::new(writer)
146            .with_separator(self.separator)
147            .with_null_value(self.null_value)
148            .finish(&mut point_cloud.point_data.data_frame)?;
149
150        Ok(())
151    }
152}
153
154impl XyzWriter<File> {
155    pub fn from_path(path: impl AsRef<Path>) -> Result<Self, Error> {
156        let file_name_str = path
157            .as_ref()
158            .file_name()
159            .ok_or(NoFileName())?
160            .to_string_lossy()
161            .to_lowercase();
162        if !file_name_str.ends_with(FILE_EXTENSION_XYZ_ZST_FORMAT)
163            && !file_name_str.ends_with(FILE_EXTENSION_XYZ_FORMAT)
164        {
165            return Err(InvalidFileExtension(file_name_str.to_string()));
166        }
167
168        let file = OpenOptions::new()
169            .create(true)
170            .write(true)
171            .truncate(true)
172            .open(path)?;
173        Ok(Self::new(file))
174    }
175}
epoint_io/xyz/write.rs

epoint_io/xyz/
write.rs
