1use crate::compressor::CopcCompressor;
4use crate::copc::{CopcInfo, Entry, HierarchyPage, OctreeNode, VoxelKey};
5
6use las::{Builder, Header};
7
8use std::collections::HashMap;
9use std::fs::File;
10use std::io::{BufWriter, Cursor, Seek, SeekFrom, Write};
11use std::path::Path;
12
13enum UpgradePdrf {
15 From1to6, From3to7, NoUpgrade, }
19
20impl UpgradePdrf {
21 fn log_string(&self) -> &str {
22 match self {
23 UpgradePdrf::From1to6 => "Upgrading LAS PDRF from 1 to 6",
24 UpgradePdrf::From3to7 => "Upgrading LAS PDRF from 3 to 7",
25 UpgradePdrf::NoUpgrade => "COPC supports the given PDRF",
26 }
27 }
28}
29
30pub struct CopcWriter<'a, W: 'a + Write + Seek> {
32 is_closed: bool,
33 start: u64,
34 compressor: CopcCompressor<'a, W>,
36 header: Header,
37 hierarchy: HierarchyPage,
39 min_node_size: i32,
40 max_node_size: i32,
41 copc_info: CopcInfo,
42 root_node: OctreeNode,
44 open_chunks: HashMap<VoxelKey, Cursor<Vec<u8>>>,
46}
47
48impl CopcWriter<'_, BufWriter<File>> {
49 pub fn from_path<P: AsRef<Path>>(
57 path: P,
58 header: Header,
59 min_size: i32,
60 max_size: i32,
61 ) -> crate::Result<Self> {
62 let copc_ext = Path::new(match path.as_ref().file_stem() {
63 Some(copc) => copc,
64 None => return Err(crate::Error::WrongCopcExtension),
65 })
66 .extension();
67
68 match (copc_ext, path.as_ref().extension()) {
69 (Some(copc), Some(laz)) => match (&copc.to_str(), &laz.to_str()) {
70 (Some(copc_str), Some(laz_str)) => {
71 if &copc_str.to_lowercase() != "copc" || &laz_str.to_lowercase() != "laz" {
72 return Err(crate::Error::WrongCopcExtension);
73 }
74 }
75 _ => return Err(crate::Error::WrongCopcExtension),
76 },
77 _ => return Err(crate::Error::WrongCopcExtension),
78 }
79
80 File::create(path)
81 .map_err(crate::Error::from)
82 .and_then(|file| CopcWriter::new(BufWriter::new(file), header, min_size, max_size))
83 }
84}
85
86impl<W: Write + Seek> CopcWriter<'_, W> {
88 pub fn new(mut write: W, header: Header, min_size: i32, max_size: i32) -> crate::Result<Self> {
117 let start = write.stream_position()?;
118
119 let min_node_size = if min_size < 1 {
120 crate::MIN_NODE_SIZE_DEFAULT
121 } else {
122 min_size
123 };
124
125 let max_node_size = if max_size < 1 {
126 crate::MAX_NODE_SIZE_DEFAULT
127 } else {
128 max_size
129 };
130
131 if min_node_size >= max_node_size {
132 return Err(crate::Error::InvalidNodeSize);
133 }
134
135 if header.version() != las::Version::new(1, 4) {
136 log::log!(log::Level::Info, "Old Las version. Upgrading");
137 }
138
139 let mut has_wkt_vlr = false;
140
141 let mut forward_vlrs = Vec::with_capacity(header.vlrs().len());
143 for vlr in header.vlrs() {
144 match (vlr.user_id.to_lowercase().as_str(), vlr.record_id) {
145 ("lasf_projection", 2112) => {
146 has_wkt_vlr = true;
147 forward_vlrs.push(vlr.clone());
148 }
149 ("lasf_projection", 34735..=34737) => (), ("copc", 1 | 1000) => (),
152 ("laszip encoded", 22204) => (),
153 ("lasf_spec", 100..355 | 65535) => (), _ => forward_vlrs.push(vlr.clone()),
156 }
157 }
158
159 let mut forward_evlrs = Vec::with_capacity(header.evlrs().len());
161 for evlr in header.evlrs() {
162 match (evlr.user_id.to_lowercase().as_str(), evlr.record_id) {
163 ("lasf_projection", 2112) => {
164 has_wkt_vlr = true;
165 forward_evlrs.push(evlr.clone());
166 }
167 ("lasf_projection", 34735..=34737) => (), ("copc", 1 | 1000) => (), ("laszip encoded", 22204) => (), ("lasf_spec", 100..355 | 65535) => (), _ => forward_evlrs.push(evlr.clone()),
174 }
175 }
176
177 if !has_wkt_vlr {
187 let epsg = las_crs::parse_las_crs(&header)?;
188 let wkt_data = match crs_definitions::from_code(epsg.horizontal) {
189 Some(wkt) => wkt,
190 None => return Err(crate::Error::InvalidEPSGCode(epsg.horizontal)),
191 }
192 .wkt
193 .as_bytes()
194 .to_owned();
195
196 let mut user_id = [0; 16];
197 for (i, c) in "LASF_Projection".as_bytes().iter().enumerate() {
198 user_id[i] = *c;
199 }
200
201 let crs_vlr = las::raw::Vlr {
202 reserved: 0,
203 user_id,
204 record_id: 2112,
205 record_length_after_header: las::raw::vlr::RecordLength::Vlr(wkt_data.len() as u16),
206 description: [0; 32],
207 data: wkt_data,
208 };
209
210 forward_vlrs.push(las::Vlr::new(crs_vlr));
211 }
212
213 let bounds = header.bounds();
215 if !(bounds.max.x - bounds.min.x).is_normal()
216 || !(bounds.max.y - bounds.min.y).is_normal()
217 || !(bounds.max.z - bounds.min.z).is_normal()
218 {
219 return Err(crate::Error::InvalidBounds(bounds));
220 }
221
222 let mut raw_head = header.into_raw()?;
223
224 let pdrf = raw_head.point_data_record_format & 0b00111111;
226 let upgrade_pdrf = match pdrf {
227 1 => {
228 let upgrade = UpgradePdrf::From1to6;
229
230 log::log!(log::Level::Info, "{}", upgrade.log_string());
231 upgrade
232 }
233 3 => {
234 let upgrade = UpgradePdrf::From3to7;
235
236 log::log!(log::Level::Info, "{}", upgrade.log_string());
237 upgrade
238 }
239 0 | 2 => {
240 return Err(las::Error::InvalidPointFormat(las::point::Format::new(
241 raw_head.point_data_record_format,
242 )?))?;
243 }
244 4..=5 | 9.. => {
245 return Err(las::Error::InvalidPointFormat(las::point::Format::new(
246 raw_head.point_data_record_format,
247 )?))?;
248 }
249 6..=8 => UpgradePdrf::NoUpgrade,
250 };
251
252 raw_head.version = las::Version::new(1, 4);
254 raw_head.point_data_record_format += match upgrade_pdrf {
255 UpgradePdrf::NoUpgrade => 0,
256 UpgradePdrf::From1to6 => 5,
257 UpgradePdrf::From3to7 => 4,
258 };
259 raw_head.point_data_record_format |= 0b11000000; raw_head.point_data_record_length += match upgrade_pdrf {
261 UpgradePdrf::NoUpgrade => 0,
262 _ => 2,
263 };
264 raw_head.global_encoding |= 0b10000; raw_head.number_of_point_records = 0;
266 raw_head.number_of_points_by_return = [0; 5];
267 raw_head.large_file = None;
268 raw_head.evlr = None;
269 raw_head.padding = vec![];
270
271 let mut software_buffer = [0_u8; 32];
272 for (i, byte) in format!("COPC-rs v{}", crate::VERSION).bytes().enumerate() {
273 software_buffer[i] = byte;
274 }
275 raw_head.generating_software = software_buffer;
276
277 let mut builder = Builder::new(raw_head)?;
279 builder.vlrs.push(CopcInfo::default().into_vlr()?);
281
282 let point_format = builder.point_format;
284 let mut laz_items = laz::laszip::LazItemRecordBuilder::new();
285 laz_items.add_item(laz::LazItemType::Point14);
286 if point_format.has_color {
287 if point_format.has_nir {
288 laz_items.add_item(laz::LazItemType::RGBNIR14);
289 } else {
290 laz_items.add_item(laz::LazItemType::RGB14);
291 }
292 }
293 if point_format.extra_bytes > 0 {
294 laz_items.add_item(laz::LazItemType::Byte14(point_format.extra_bytes));
295 }
296
297 let laz_vlr = laz::LazVlrBuilder::new(laz_items.build())
298 .with_variable_chunk_size()
299 .build();
300 let mut cursor = Cursor::new(Vec::<u8>::new());
301 laz_vlr.write_to(&mut cursor)?;
302 let laz_vlr = las::Vlr {
303 user_id: laz::LazVlr::USER_ID.to_owned(),
304 record_id: laz::LazVlr::RECORD_ID,
305 description: laz::LazVlr::DESCRIPTION.to_owned(),
306 data: cursor.into_inner(),
307 };
308 builder.vlrs.push(laz_vlr);
309
310 builder.vlrs.extend(forward_vlrs);
312 builder.evlrs.extend(forward_evlrs);
313 let header = builder.into_header()?;
316
317 header.write_to(&mut write)?;
320
321 let center_point = las::Vector {
322 x: (bounds.min.x + bounds.max.x) / 2.,
323 y: (bounds.min.y + bounds.max.y) / 2.,
324 z: (bounds.min.z + bounds.max.z) / 2.,
325 };
326 let halfsize = (center_point.x - bounds.min.x)
327 .max((center_point.y - bounds.min.y).max(center_point.z - bounds.min.z));
328
329 let mut root_node = OctreeNode::new();
330
331 root_node.bounds = las::Bounds {
332 min: las::Vector {
333 x: center_point.x - halfsize,
334 y: center_point.y - halfsize,
335 z: center_point.z - halfsize,
336 },
337 max: las::Vector {
338 x: center_point.x + halfsize,
339 y: center_point.y + halfsize,
340 z: center_point.z + halfsize,
341 },
342 };
343 root_node.entry.key.level = 0;
344 root_node.entry.offset = write.stream_position()?;
345
346 let copc_info = CopcInfo {
347 center: center_point,
348 halfsize,
349 spacing: 0.,
350 root_hier_offset: 0,
351 root_hier_size: 0,
352 gpstime_minimum: f64::MAX,
353 gpstime_maximum: f64::MIN,
354 };
355
356 Ok(CopcWriter {
357 is_closed: false,
358 start,
359 compressor: CopcCompressor::new(write, header.laz_vlr()?)?,
360 header,
361 hierarchy: HierarchyPage { entries: vec![] },
362 min_node_size,
363 max_node_size,
364 copc_info,
365 root_node,
366 open_chunks: HashMap::default(),
367 })
368 }
369
370 pub fn write<D: IntoIterator<Item = las::Point>>(
400 &mut self,
401 data: D,
402 num_points: i32,
403 ) -> crate::Result<()> {
404 if self.is_closed {
405 return Err(crate::Error::ClosedWriter);
406 }
407
408 let result = if num_points < self.max_node_size + self.min_node_size {
409 self.write_greedy(data)
411 } else {
412 self.write_stochastic(data, num_points as usize)
414 };
415
416 self.close()?;
417 result
418 }
419
420 pub fn is_closed(&self) -> bool {
422 self.is_closed
423 }
424
425 pub fn max_node_size(&self) -> i32 {
427 self.max_node_size
428 }
429
430 pub fn min_node_size(&self) -> i32 {
432 self.min_node_size
433 }
434
435 pub fn header(&self) -> &Header {
437 &self.header
438 }
439
440 pub fn hierarchy_entries(&self) -> &HierarchyPage {
442 &self.hierarchy
443 }
444
445 pub fn copc_info(&self) -> &CopcInfo {
447 &self.copc_info
448 }
449}
450
451impl<W: Write + Seek> CopcWriter<'_, W> {
453 fn write_greedy<D: IntoIterator<Item = las::Point>>(&mut self, data: D) -> crate::Result<()> {
455 let mut invalid_points = Ok(());
456
457 for p in data.into_iter() {
458 if !p.matches(self.header.point_format()) {
459 invalid_points = Err(crate::Error::InvalidPoint(
460 crate::PointAddError::PointAttributesDoNotMatch(*self.header.point_format()),
461 ));
462 continue;
463 }
464 if !bounds_contains_point(&self.root_node.bounds, &p) {
465 if invalid_points.is_ok() {
466 invalid_points = Err(crate::Error::InvalidPoint(
467 crate::PointAddError::PointNotInBounds,
468 ));
469 }
470 continue;
471 }
472
473 self.add_point_greedy(p)?;
474 }
475 invalid_points
476 }
477
478 fn write_stochastic<D: IntoIterator<Item = las::Point>>(
480 &mut self,
481 data: D,
482 num_points: usize,
483 ) -> crate::Result<()> {
484 let mut invalid_points = Ok(());
485
486 let expected_levels =
500 ((((3 * num_points) as f64 / self.max_node_size as f64 + 1.).log2() - 2.) / 2.).ceil()
501 as usize;
502
503 for (i, p) in data.into_iter().enumerate() {
504 if !p.matches(self.header.point_format()) {
505 invalid_points = Err(crate::Error::InvalidPoint(
506 crate::PointAddError::PointAttributesDoNotMatch(*self.header.point_format()),
507 ));
508 continue;
509 }
510 if !bounds_contains_point(&self.root_node.bounds, &p) {
511 if invalid_points.is_ok() {
512 invalid_points = Err(crate::Error::InvalidPoint(
513 crate::PointAddError::PointNotInBounds,
514 ));
515 }
516 continue;
517 }
518
519 if num_points <= i {
522 self.add_point_greedy(p)?;
523 } else {
524 self.add_point_stochastic(p, expected_levels)?;
525 }
526 }
527 invalid_points
528 }
529
530 fn close(&mut self) -> crate::Result<()> {
532 if self.is_closed {
533 return Err(crate::Error::ClosedWriter);
534 }
535 if self.header.number_of_points() < 1 {
536 return Err(crate::Error::EmptyCopcFile);
537 }
538
539 for (key, chunk) in self.open_chunks.drain() {
541 let inner = chunk.into_inner();
542 if inner.is_empty() {
543 continue;
544 }
545 let (chunk_table_entry, chunk_offset) = self.compressor.compress_chunk(inner)?;
546 self.hierarchy.entries.push(Entry {
547 key,
548 offset: chunk_offset,
549 byte_size: chunk_table_entry.byte_count as i32,
550 point_count: chunk_table_entry.point_count as i32,
551 })
552 }
553
554 self.compressor.done()?;
555
556 let start_of_first_evlr = self.compressor.get_mut().stream_position()?;
557
558 let raw_evlrs: Vec<las::Result<las::raw::Vlr>> = self
559 .header
560 .evlrs()
561 .iter()
562 .map(|evlr| evlr.clone().into_raw(true))
563 .collect();
564
565 self.hierarchy
567 .clone()
568 .into_evlr()?
569 .into_raw(true)?
570 .write_to(self.compressor.get_mut())?;
571 for raw_evlr in raw_evlrs {
573 raw_evlr?.write_to(self.compressor.get_mut())?;
574 }
575
576 self.compressor
577 .get_mut()
578 .seek(SeekFrom::Start(self.start))?;
579 self.header.clone().into_raw().and_then(|mut raw_header| {
580 if let Some(mut e) = raw_header.evlr {
581 e.start_of_first_evlr = start_of_first_evlr;
582 e.number_of_evlrs += 1;
583 } else {
584 raw_header.evlr = Some(las::raw::header::Evlr {
585 start_of_first_evlr,
586 number_of_evlrs: 1,
587 });
588 }
589 raw_header.write_to(self.compressor.get_mut())
590 })?;
591
592 self.copc_info.spacing =
594 2. * self.copc_info.halfsize / (self.root_node.entry.point_count as f64);
595 self.copc_info.root_hier_offset = start_of_first_evlr + 60; self.copc_info.root_hier_size = self.hierarchy.byte_size();
597
598 self.copc_info
599 .clone()
600 .into_vlr()?
601 .into_raw(false)?
602 .write_to(self.compressor.get_mut())?;
603
604 self.compressor
605 .get_mut()
606 .seek(SeekFrom::Start(self.start))?;
607
608 self.is_closed = true;
609 Ok(())
610 }
611
612 fn add_point_greedy(&mut self, point: las::Point) -> crate::Result<()> {
616 self.header.add_point(&point);
617
618 if point.gps_time.unwrap() < self.copc_info.gpstime_minimum {
619 self.copc_info.gpstime_minimum = point.gps_time.unwrap();
620 } else if point.gps_time.unwrap() > self.copc_info.gpstime_maximum {
621 self.copc_info.gpstime_maximum = point.gps_time.unwrap();
622 }
623
624 let mut node_key = None;
625 let mut write_chunk = false;
626
627 let root_bounds = self.root_node.bounds;
628
629 let mut nodes_to_check = vec![&mut self.root_node];
632 while let Some(node) = nodes_to_check.pop() {
633 if !bounds_contains_point(&node.bounds, &point) {
634 continue;
636 }
637 if node.is_full(self.max_node_size) {
638 if node.children.is_empty() {
641 let child_keys = node.entry.key.children();
644 for key in child_keys {
645 let child_bounds = key.bounds(&root_bounds);
646 node.children.push(OctreeNode {
647 entry: Entry {
648 key,
649 offset: 0,
650 byte_size: 0,
651 point_count: 0,
652 },
653 bounds: child_bounds,
654 children: Vec::with_capacity(8),
655 })
656 }
657 }
658 for child in node.children.iter_mut() {
660 nodes_to_check.push(child);
661 }
662 } else {
663 node_key = Some(node.entry.key.clone());
665 node.entry.point_count += 1;
666
667 write_chunk = node.is_full(self.max_node_size);
669 break;
670 }
671 }
672 let Some(node_key) = node_key else {
673 return Err(crate::Error::PointNotAddedToAnyNode);
674 };
675
676 let raw_point = point.into_raw(self.header.transforms())?;
677
678 if !self.open_chunks.contains_key(&node_key) {
679 let mut val = Cursor::new(vec![]);
680 raw_point.write_to(&mut val, self.header.point_format())?;
681
682 self.open_chunks.insert(node_key.clone(), val);
683 } else {
684 let buffer = self.open_chunks.get_mut(&node_key).unwrap();
685 raw_point.write_to(buffer, self.header.point_format())?;
686 }
687
688 if write_chunk {
689 let chunk = self.open_chunks.remove(&node_key).unwrap();
690 let (chunk_table_entry, chunk_offset) =
691 self.compressor.compress_chunk(chunk.into_inner())?;
692 self.hierarchy.entries.push(Entry {
693 key: node_key,
694 offset: chunk_offset,
695 byte_size: chunk_table_entry.byte_count as i32,
696 point_count: chunk_table_entry.point_count as i32,
697 });
698 }
699 Ok(())
700 }
701
702 fn add_point_stochastic(
703 &mut self,
704 point: las::Point,
705 expected_levels: usize,
706 ) -> crate::Result<()> {
707 let root_bounds = self.root_node.bounds;
713
714 let mut node_candidates = vec![];
715
716 let mut nodes_to_check = vec![&mut self.root_node];
718 while let Some(node) = nodes_to_check.pop() {
719 if !bounds_contains_point(&node.bounds, &point) {
720 continue;
722 }
723
724 if node.children.is_empty() && node.entry.key.level < expected_levels as i32 {
725 let child_keys = node.entry.key.children();
726 for key in child_keys {
727 let child_bounds = key.bounds(&root_bounds);
728 node.children.push(OctreeNode {
729 entry: Entry {
730 key,
731 offset: 0,
732 byte_size: 0,
733 point_count: 0,
734 },
735 bounds: child_bounds,
736 children: Vec::with_capacity(8),
737 })
738 }
739 }
740 if !node.is_full(self.max_node_size) {
741 node_candidates.push(&mut node.entry);
742 }
743 for child in node.children.iter_mut() {
745 nodes_to_check.push(child);
746 }
747 }
748
749 if node_candidates.is_empty() {
750 return self.add_point_greedy(point);
752 }
753
754 let chosen_index = get_random_weighted_index(&node_candidates);
756
757 let chosen_entry = &mut node_candidates[chosen_index];
758
759 chosen_entry.point_count += 1;
760
761 let write_chunk = chosen_entry.point_count > self.max_node_size;
762
763 let node_key = chosen_entry.key.clone();
764
765 self.header.add_point(&point);
766
767 if point.gps_time.unwrap() < self.copc_info.gpstime_minimum {
768 self.copc_info.gpstime_minimum = point.gps_time.unwrap();
769 } else if point.gps_time.unwrap() > self.copc_info.gpstime_maximum {
770 self.copc_info.gpstime_maximum = point.gps_time.unwrap();
771 }
772
773 let raw_point = point.into_raw(self.header.transforms())?;
774
775 if !self.open_chunks.contains_key(&node_key) {
776 let mut val = Cursor::new(vec![]);
777 raw_point.write_to(&mut val, self.header.point_format())?;
778
779 self.open_chunks.insert(node_key.clone(), val);
780 } else {
781 let buffer = self.open_chunks.get_mut(&node_key).unwrap();
782 raw_point.write_to(buffer, self.header.point_format())?;
783 }
784
785 if write_chunk {
786 let chunk = self.open_chunks.remove(&node_key).unwrap();
787 let (chunk_table_entry, chunk_offset) =
788 self.compressor.compress_chunk(chunk.into_inner())?;
789 self.hierarchy.entries.push(Entry {
790 key: node_key,
791 offset: chunk_offset,
792 byte_size: chunk_table_entry.byte_count as i32,
793 point_count: chunk_table_entry.point_count as i32,
794 });
795 }
796 Ok(())
797 }
798}
799
800fn get_random_weighted_index(entries: &Vec<&mut Entry>) -> usize {
801 let levels: Vec<i32> = entries.iter().map(|e| e.key.level).collect();
803 let zero_level = levels[0];
804
805 let areas: Vec<f64> = levels
807 .iter()
808 .map(|l| (0.25_f64).powi(l - zero_level))
809 .collect();
810 let inv_sum = areas.iter().fold(0., |acc, a| acc + 1. / a);
812
813 let weights: Vec<f64> = areas.iter().map(|a| (1. / a) / inv_sum).collect();
814
815 let random = fastrand::f64();
817 let mut chosen_index = weights.len() - 1;
818
819 for i in 0..weights.len() - 1 {
820 if (weights[i]..=weights[i + 1]).contains(&random) {
821 chosen_index = i;
822 break;
823 }
824 }
825 chosen_index
826}
827
828impl<W: Write + Seek> Drop for CopcWriter<'_, W> {
829 fn drop(&mut self) {
830 if !self.is_closed {
831 self.close()
834 .expect("Error when dropping the writer. No points written.");
835 }
836 }
837}
838
839#[inline]
840fn bounds_contains_point(b: &las::Bounds, p: &las::Point) -> bool {
841 !(b.max.x < p.x
842 || b.max.y < p.y
843 || b.max.z < p.z
844 || b.min.x > p.x
845 || b.min.y > p.y
846 || b.min.z > p.z)
847}