diff --git a/rust/spark-lib/src/ply.rs b/rust/spark-lib/src/ply.rs index 57681c85..a957cfa8 100644 --- a/rust/spark-lib/src/ply.rs +++ b/rust/spark-lib/src/ply.rs @@ -10,6 +10,8 @@ pub const PLY_MAGIC: u32 = 0x00796c70; // "ply" const MAX_SPLAT_CHUNK: usize = 65536; const SH_C0: f32 = 0.28209479177387814; const SUPER_CHUNK_SIZE: usize = 256; +const POINT_CLOUD_PROPERTIES: [&str; 6] = ["x", "y", "z", "red", "green", "blue"]; +const DEFAULT_POINT_SCALE: f32 = 0.001; pub struct PlyDecoder { splats: T, @@ -69,6 +71,14 @@ impl PlyDecoder { lod_tree: false, })?; PlyState::SuperSplat(state) + } else if parsed.is_pointcloud { + let state = PointCloudDecoderState::new(parsed.num_splats, parsed.vertex.record_size, parsed.vertex.properties.clone())?; + self.splats.init_splats(&SplatInit { + num_splats: parsed.num_splats, + max_sh_degree: 0, + lod_tree: false, + })?; + PlyState::PointCloud(state) } else { let state = PlyDecoderState::new(parsed.num_splats, parsed.vertex.record_size, parsed.vertex.properties.clone())?; self.splats.init_splats(&SplatInit { @@ -86,12 +96,65 @@ impl PlyDecoder { fn poll_data(&mut self) -> anyhow::Result<()> { match self.state { + Some(PlyState::PointCloud(_)) => self.poll_data_pointcloud(), Some(PlyState::Standard(_)) => self.poll_data_standard(), Some(PlyState::SuperSplat(_)) => self.poll_data_supersplat(), None => unreachable!(), } } + fn poll_data_pointcloud(&mut self) -> anyhow::Result<()> { + let Some(PlyState::PointCloud(state)) = self.state.as_mut() else { unreachable!() }; + let mut offset = 0; + loop { + let available = (self.buffer.len() - offset) / state.record_size; + let remaining = state.num_splats.saturating_sub(state.next_splat); + let count = remaining.min(available).min(MAX_SPLAT_CHUNK); + if count == 0 { + break; + } + + state.ensure_out(count); + + for i in 0..count { + let [i3, i4] = [i * 3, i * 4]; + let base = offset + i * state.record_size; + + for d in 0..3 { + state.out_center[i3 + d] = state.xyz[d].get_f32(&self.buffer, base); + } + state.out_opacity[i] = match state.alpha { + Some(alpha) => alpha.get_f32(&self.buffer, base), + None => 1.0 + }; + for d in 0..3 { + state.out_rgb[i3 + d] = state.rgb[d].get_f32(&self.buffer, base); + } + state.out_scale.splice(i3..i3+3, [DEFAULT_POINT_SCALE, DEFAULT_POINT_SCALE, DEFAULT_POINT_SCALE]); + state.out_quat.splice(i4..i4+4, [0.0, 0.0, 0.0, 1.0]); + } + + self.splats.set_batch(state.next_splat, count, &SplatProps { + center: &state.out_center[..count * 3], + opacity: &state.out_opacity[..count], + rgb: &state.out_rgb[..count * 3], + scale: &state.out_scale[..count * 3], + quat: &state.out_quat[..count * 4], + sh1: &Vec::new(), + sh2: &Vec::new(), + sh3: &Vec::new(), + ..Default::default() + }); + + state.next_splat += count; + offset += count * state.record_size; + } + + self.buffer.drain(..offset); + Ok(()) + } + + fn poll_data_standard(&mut self) -> anyhow::Result<()> { let Some(PlyState::Standard(state)) = self.state.as_mut() else { unreachable!() }; let mut offset = 0; @@ -258,6 +321,11 @@ impl ChunkReceiver for PlyDecoder { // As long as we've read the correct number of splats, we're good. match state { + PlyState::PointCloud(state) => { + if state.next_splat != state.num_splats { + return Err(anyhow!("Expected {} splats, got {}", state.num_splats, state.next_splat)); + } + }, PlyState::Standard(state) => { if state.next_splat != state.num_splats { return Err(anyhow!("Expected {} splats, got {}", state.num_splats, state.next_splat)); @@ -284,6 +352,7 @@ impl ChunkReceiver for PlyDecoder { #[derive(Debug)] enum PlyState { + PointCloud(PointCloudDecoderState), Standard(PlyDecoderState), SuperSplat(SuperSplatState), } @@ -335,6 +404,7 @@ struct ParsedHeader { chunk: Option, sh: Option, num_splats: usize, + is_pointcloud: bool, is_supersplat: bool, } @@ -421,12 +491,14 @@ fn parse_header(header: &str) -> anyhow::Result { let vertex = elements.iter().find(|e| e.name == "vertex").cloned().ok_or(anyhow!("Missing vertex element"))?; let chunk = elements.iter().find(|e| e.name == "chunk").cloned(); let sh = elements.iter().find(|e| e.name == "sh").cloned(); + let is_pointcloud = POINT_CLOUD_PROPERTIES.iter().all(|&p| vertex.properties.contains_key(p)); Ok(ParsedHeader { num_splats: vertex.count, vertex, chunk, sh, + is_pointcloud, is_supersplat: elements.iter().any(|e| e.name == "chunk"), elements, }) @@ -955,6 +1027,74 @@ impl PlyDecoderState { } } +#[derive(Debug)] +struct PointCloudDecoderState { + num_splats: usize, + record_size: usize, + next_splat: usize, + + #[allow(unused)] + properties: HashMap, + xyz: [PlyProperty; 3], + rgb: [PlyProperty; 3], + alpha: Option, + + out_center: Vec, + out_opacity: Vec, + out_rgb: Vec, + out_scale: Vec, + out_quat: Vec, +} + +impl PointCloudDecoderState { + fn new(num_splats: usize, record_size: usize, properties: HashMap) -> anyhow::Result { + let xyz = [ + *properties.get("x").ok_or(anyhow!("Missing x property"))?, + *properties.get("y").ok_or(anyhow!("Missing y property"))?, + *properties.get("z").ok_or(anyhow!("Missing z property"))?, + ]; + let rgb = [ + *properties.get("red").ok_or(anyhow!("Missing red property"))?, + *properties.get("green").ok_or(anyhow!("Missing green property"))?, + *properties.get("blue").ok_or(anyhow!("Missing blue property"))?, + ]; + let alpha = properties.get("alpha").map(|p| *p); + + Ok(Self { + num_splats, + record_size, + next_splat: 0, + properties, + xyz, + rgb, + alpha, + out_center: Vec::new(), + out_opacity: Vec::new(), + out_rgb: Vec::new(), + out_scale: Vec::new(), + out_quat: Vec::new(), + }) + } + + fn ensure_out(&mut self, count: usize) { + if self.out_center.len() < (count * 3) { + self.out_center.resize(count * 3, 0.0); + } + if self.out_opacity.len() < count { + self.out_opacity.resize(count, 0.0); + } + if self.out_rgb.len() < (count * 3) { + self.out_rgb.resize(count * 3, 0.0); + } + if self.out_scale.len() < (count * 3) { + self.out_scale.resize(count * 3, 0.0); + } + if self.out_quat.len() < (count * 4) { + self.out_quat.resize(count * 4, 0.0); + } + } +} + #[derive(Debug, Clone, Copy)] pub enum PlyPropertyType { Char,