From 35f0e67f26af8ee7d6c2b28d930d124d5f4333f9 Mon Sep 17 00:00:00 2001 From: n3ur0m4n Date: Fri, 23 Oct 2026 16:45:06 +0200 Subject: [PATCH] fix(sogs): decode SOG in chunks to avoid wasm32 OOM on large files The decoder kept the zip buffer, the compressed file cache, all decoded RGBA images and full-size f32 arrays alive at the same time, which exceeds the 4 GiB wasm32 limit on large files (~40M splats). - finish() moves the zip buffer into decode_sogs, dropped after preload - cache entries are removed on their last use - compressed bytes are dropped right after each image decode - splats are emitted in MAX_SPLAT_CHUNK blocks via set_batch, with reused buffers; shN centroids/labels are loaded once and decoded per block - tests: 70k-splat v1/v2 SOG zips with 3 SH bands, compared against a per-splat reference Co-Authored-By: Claude Opus 5.5 (1M context) --- rust/spark-lib/src/sogs.rs | 831 ++++++++++++++++++++++--------------- 1 file changed, 486 insertions(+), 345 deletions(-) diff --git a/rust/spark-lib/src/sogs.rs b/rust/spark-lib/src/sogs.rs index 5cae3c3c..79ca08ee 100644 --- a/rust/spark-lib/src/sogs.rs +++ b/rust/spark-lib/src/sogs.rs @@ -130,48 +130,68 @@ impl ChunkReceiver for SogsDecoder { } fn finish(&mut self) -> anyhow::Result<()> { - if self.buffer.len() < 4 { + // Move the zip buffer into decode_sogs, which frees it after preloading + let buf = std::mem::take(&mut self.buffer); + if buf.len() < 4 { return Err(anyhow!("SOGS file too small")); } - let magic = u32::from_le_bytes([self.buffer[0], self.buffer[1], self.buffer[2], self.buffer[3]]); + let magic = u32::from_le_bytes([buf[0], buf[1], buf[2], buf[3]]); if magic != PK_MAGIC { return Err(anyhow!("Not a ZIP/SOGS file")); } - decode_sogs(&self.buffer, &mut self.splats, None)?; + decode_sogs(buf, &mut self.splats, None)?; Ok(()) } } -fn decode_sogs(bytes: &[u8], splats: &mut T, _pathname: Option<&str>) -> anyhow::Result<()> { - let cursor = Cursor::new(bytes); - let mut zip = ZipArchive::new(cursor)?; +fn decode_sogs(bytes: Vec, splats: &mut T, _pathname: Option<&str>) -> anyhow::Result<()> { + let (meta, mut file_cache) = { + let mut zip = ZipArchive::new(Cursor::new(bytes.as_slice()))?; - // Find meta.json entry and its prefix - let mut meta_name: Option = None; - for i in 0..zip.len() { - let name = zip.by_index(i)?.name().to_string(); - if name.ends_with("meta.json") { - meta_name = Some(name); - break; + // Find meta.json entry and its prefix + let mut meta_name: Option = None; + for i in 0..zip.len() { + let name = zip.by_index(i)?.name().to_string(); + if name.ends_with("meta.json") { + meta_name = Some(name); + break; + } } - } - let meta_name = meta_name.ok_or_else(|| anyhow!("meta.json not found in SOGS zip"))?; - let prefix = meta_name.rsplit_once('/').map(|(p, _)| format!("{}/", p)).unwrap_or_default(); + let meta_name = meta_name.ok_or_else(|| anyhow!("meta.json not found in SOGS zip"))?; + let prefix = meta_name.rsplit_once('/').map(|(p, _)| format!("{}/", p)).unwrap_or_default(); + + let meta_bytes = { + let mut meta_file = zip.by_name(&meta_name)?; + let mut meta_bytes = Vec::new(); + std::io::copy(&mut meta_file, &mut meta_bytes)?; + meta_bytes + }; + let meta: PcSogsRoot = serde_json::from_slice(&meta_bytes) + .context("Failed to parse meta.json for SOGS")?; - let meta_bytes = { - let mut meta_file = zip.by_name(&meta_name)?; - let mut meta_bytes = Vec::new(); - std::io::copy(&mut meta_file, &mut meta_bytes)?; - meta_bytes + let mut file_cache: HashMap> = HashMap::new(); + for f in meta_files(&meta) { + preload_file(&mut zip, &prefix, &mut file_cache, f)?; + } + (meta, file_cache) }; - let meta: PcSogsRoot = serde_json::from_slice(&meta_bytes) - .context("Failed to parse meta.json for SOGS")?; - - let mut file_cache: HashMap> = HashMap::new(); - preload_all(&meta, &prefix, &mut zip, &mut file_cache)?; + // Zip no longer needed: all compressed files are in the cache + drop(bytes); + // Each file leaves the cache on its last request (a name can repeat in meta) + let mut uses: HashMap = HashMap::new(); + for f in meta_files(&meta) { + *uses.entry(f.clone()).or_default() += 1; + } let mut get_file = |name: &str| -> anyhow::Result> { - file_cache.get(name).cloned().ok_or_else(|| anyhow!("Missing file {name} in cache")) + let left = uses.get_mut(name).ok_or_else(|| anyhow!("Missing file {name} in cache"))?; + *left -= 1; + if *left == 0 { + file_cache.remove(name) + } else { + file_cache.get(name).cloned() + } + .ok_or_else(|| anyhow!("Missing file {name} in cache")) }; match meta { @@ -180,70 +200,129 @@ fn decode_sogs(bytes: &[u8], splats: &mut T, _pathname: Option } } +fn meta_files(meta: &PcSogsRoot) -> Vec<&String> { + let mut out = Vec::new(); + match meta { + PcSogsRoot::V2(v2) => { + out.extend(v2.means.files.iter()); + out.extend(v2.scales.files.iter()); + out.extend(v2.quats.files.iter()); + out.extend(v2.sh0.files.iter()); + if let Some(shn) = &v2.shn { out.extend(shn.files.iter()); } + } + PcSogsRoot::V1(v1) => { + out.extend(v1.means.files.iter()); + out.extend(v1.scales.files.iter()); + out.extend(v1.quats.files.iter()); + out.extend(v1.sh0.files.iter()); + if let Some(shn) = &v1.shn { out.extend(shn.files.iter()); } + } + } + out +} + +/// Decode an image; the compressed bytes are dropped right after +fn load_image( + get_file: &mut dyn FnMut(&str) -> anyhow::Result>, + name: &str, + what: &str, +) -> anyhow::Result { + let bytes = get_file(name)?; + decode_image(&bytes).with_context(|| format!("decode {what}")) +} + +/// Buffers for one chunk of splats, reused for the whole decode +struct Chunk { + center: Vec, + scale: Vec, + quat: Vec, + rgb: Vec, + opacity: Vec, + sh1: Vec, + sh2: Vec, + sh3: Vec, +} + +/// Call `fill(base, cnt, chunk)` for chunks of MAX_SPLAT_CHUNK splats and pass them to the receiver +fn emit_chunks( + splats: &mut T, + num_splats: usize, + max_sh_degree: usize, + mut fill: impl FnMut(usize, usize, &mut Chunk) -> anyhow::Result<()>, +) -> anyhow::Result<()> { + let n = num_splats.min(MAX_SPLAT_CHUNK); + let sh = |deg: usize, k: usize| if max_sh_degree >= deg { vec![0.0f32; n * k] } else { Vec::new() }; + let mut c = Chunk { + center: vec![0.0; n * 3], + scale: vec![0.0; n * 3], + quat: vec![0.0; n * 4], + rgb: vec![0.0; n * 3], + opacity: vec![0.0; n], + sh1: sh(1, 9), + sh2: sh(2, 15), + sh3: sh(3, 21), + }; + let mut base = 0usize; + while base < num_splats { + let cnt = (num_splats - base).min(MAX_SPLAT_CHUNK); + fill(base, cnt, &mut c)?; + splats.set_batch( + base, + cnt, + &SplatProps { + center: &c.center[..cnt * 3], + opacity: &c.opacity[..cnt], + rgb: &c.rgb[..cnt * 3], + scale: &c.scale[..cnt * 3], + quat: &c.quat[..cnt * 4], + sh1: if max_sh_degree >= 1 { &c.sh1[..cnt * 9] } else { &[] }, + sh2: if max_sh_degree >= 2 { &c.sh2[..cnt * 15] } else { &[] }, + sh3: if max_sh_degree >= 3 { &c.sh3[..cnt * 21] } else { &[] }, + ..Default::default() + }, + ); + base += cnt; + } + splats.finish() +} + fn decode_v2( meta: PcSogsV2, splats: &mut T, get_file: &mut dyn FnMut(&str) -> anyhow::Result>, ) -> anyhow::Result<()> { let _ = meta.version; + let _ = meta.means.shape; let num_splats = meta.count; let max_sh_degree = meta.shn.as_ref().map(|shn| { if shn.bands >= 3 { 3 } else if shn.bands >= 2 { 2 } else if shn.bands >= 1 { 1 } else { 0 } }).unwrap_or(0); splats.init_splats(&SplatInit { num_splats, max_sh_degree, lod_tree: false })?; - let means0 = decode_rgba(&get_file(&meta.means.files[0])?) - .context("decode means[0]")?; - let means1 = decode_rgba(&get_file(&meta.means.files[1])?) - .context("decode means[1]")?; - let scales_img = decode_rgba(&get_file(&meta.scales.files[0])?) - .context("decode scales")?; - let quats_img = decode_rgba(&get_file(&meta.quats.files[0])?) - .context("decode quats")?; - let sh0_img = decode_rgba(&get_file(&meta.sh0.files[0])?) - .context("decode sh0")?; - - let mut center = vec![0.0f32; num_splats * 3]; - let mut scale = vec![0.0f32; num_splats * 3]; - let mut quat = vec![0.0f32; num_splats * 4]; - let mut rgb = vec![0.0f32; num_splats * 3]; - let mut opacity = vec![0.0f32; num_splats]; - let mut sh1: Vec = if max_sh_degree >= 1 { vec![0.0; num_splats * 9] } else { Vec::new() }; - let mut sh2: Vec = if max_sh_degree >= 2 { vec![0.0; num_splats * 15] } else { Vec::new() }; - let mut sh3: Vec = if max_sh_degree >= 3 { vec![0.0; num_splats * 21] } else { Vec::new() }; - - if let Some(shape) = meta.means.shape { - let _ = shape; - } - decode_means(&meta.means.mins, &meta.means.maxs, &means0, &means1, &mut center)?; - decode_scales_v2(&meta.scales.codebook, &scales_img, &mut scale)?; - decode_quats(&quats_img, &mut quat)?; - decode_sh0_v2(&meta.sh0.codebook, &sh0_img, &mut rgb, &mut opacity)?; - - if let Some(shn) = meta.shn { - decode_shn_v2( - shn, - get_file, - num_splats, - &mut sh1, - &mut sh2, - &mut sh3, - )?; - } - - emit_to_receiver( - splats, - num_splats, - max_sh_degree, - ¢er, - &scale, - &quat, - &rgb, - &opacity, - &sh1, - &sh2, - &sh3, - ) + let means0 = load_image(get_file, &meta.means.files[0], "means[0]")?; + let means1 = load_image(get_file, &meta.means.files[1], "means[1]")?; + let scales_img = load_image(get_file, &meta.scales.files[0], "scales")?; + let quats_img = load_image(get_file, &meta.quats.files[0], "quats")?; + let sh0_img = load_image(get_file, &meta.sh0.files[0], "sh0")?; + let shn = match meta.shn { + Some(shn) => Some(ShN::load(get_file, &shn.files, shn.codebook, max_sh_degree)?), + None => None, + }; + + let scale_lookup: Vec = meta.scales.codebook.iter().map(|v| v.exp()).collect(); + let rgb_lookup: Vec = meta.sh0.codebook.iter().map(|v| SH_C0 * v + 0.5).collect(); + let a_lookup: Vec = (0..256).map(|i| i as f32 / 255.0).collect(); + + emit_chunks(splats, num_splats, max_sh_degree, |base, cnt, c| { + decode_means(&meta.means.mins, &meta.means.maxs, &means0, &means1, base, &mut c.center[..cnt * 3]); + decode_scales(&scale_lookup, &scale_lookup, &scale_lookup, &scales_img, base, &mut c.scale[..cnt * 3]); + decode_quats(&quats_img, base, &mut c.quat[..cnt * 4]); + decode_sh0(&rgb_lookup, &rgb_lookup, &rgb_lookup, &a_lookup, &sh0_img, base, &mut c.rgb[..cnt * 3], &mut c.opacity[..cnt]); + if let Some(shn) = &shn { + shn.decode(base, cnt, c); + } + Ok(()) + }) } fn decode_v1( @@ -271,68 +350,61 @@ fn decode_v1( splats.init_splats(&SplatInit { num_splats, max_sh_degree, lod_tree: false })?; - let means0 = decode_rgba(&get_file(&meta.means.files[0])?) - .context("decode means[0]")?; - let means1 = decode_rgba(&get_file(&meta.means.files[1])?) - .context("decode means[1]")?; - let scales_img = decode_rgba(&get_file(&meta.scales.files[0])?) - .context("decode scales")?; - let quats_img = decode_rgba(&get_file(&meta.quats.files[0])?) - .context("decode quats")?; - let sh0_img = decode_rgba(&get_file(&meta.sh0.files[0])?) - .context("decode sh0")?; - - let mut center = vec![0.0f32; num_splats * 3]; - let mut scale = vec![0.0f32; num_splats * 3]; - let mut quat = vec![0.0f32; num_splats * 4]; - let mut rgb = vec![0.0f32; num_splats * 3]; - let mut opacity = vec![0.0f32; num_splats]; - let mut sh1: Vec = if max_sh_degree >= 1 { vec![0.0; num_splats * 9] } else { Vec::new() }; - let mut sh2: Vec = if max_sh_degree >= 2 { vec![0.0; num_splats * 15] } else { Vec::new() }; - let mut sh3: Vec = if max_sh_degree >= 3 { vec![0.0; num_splats * 21] } else { Vec::new() }; - - decode_means(&meta.means.mins, &meta.means.maxs, &means0, &means1, &mut center)?; - decode_scales_v1(&meta.scales.mins, &meta.scales.maxs, &scales_img, &mut scale)?; - decode_quats(&quats_img, &mut quat)?; - decode_sh0_v1(&meta.sh0.mins, &meta.sh0.maxs, &sh0_img, &mut rgb, &mut opacity)?; - - if let Some(shn) = meta.shn { - decode_shn_v1( - shn, - get_file, - num_splats, - max_sh_degree, - &mut sh1, - &mut sh2, - &mut sh3, - )?; - } - - emit_to_receiver( - splats, - num_splats, - max_sh_degree, - ¢er, - &scale, - &quat, - &rgb, - &opacity, - &sh1, - &sh2, - &sh3, - ) + let means0 = load_image(get_file, &meta.means.files[0], "means[0]")?; + let means1 = load_image(get_file, &meta.means.files[1], "means[1]")?; + let scales_img = load_image(get_file, &meta.scales.files[0], "scales")?; + let quats_img = load_image(get_file, &meta.quats.files[0], "quats")?; + let sh0_img = load_image(get_file, &meta.sh0.files[0], "sh0")?; + let shn = match meta.shn { + Some(shn) => { + let lookup: Vec = (0..256) + .map(|i| shn.mins + (shn.maxs - shn.mins) * (i as f32 / 255.0)) + .collect(); + Some(ShN::load(get_file, &shn.files, lookup, max_sh_degree)?) + } + None => None, + }; + + let build_lookup = |min: f32, max: f32, post: fn(f32) -> f32| -> Vec { + (0..256) + .map(|i| min + (max - min) * (i as f32 / 255.0)) + .map(post) + .collect() + }; + let (smin, smax) = (meta.scales.mins, meta.scales.maxs); + let lx = build_lookup(smin[0], smax[0], f32::exp); + let ly = build_lookup(smin[1], smax[1], f32::exp); + let lz = build_lookup(smin[2], smax[2], f32::exp); + let (cmin, cmax) = (meta.sh0.mins, meta.sh0.maxs); + let lr = build_lookup(cmin[0], cmax[0], |v| SH_C0 * v + 0.5); + let lg = build_lookup(cmin[1], cmax[1], |v| SH_C0 * v + 0.5); + let lb = build_lookup(cmin[2], cmax[2], |v| SH_C0 * v + 0.5); + let la = build_lookup(cmin[3], cmax[3], |v| 1.0 / (1.0 + (-v).exp())); + + emit_chunks(splats, num_splats, max_sh_degree, |base, cnt, c| { + decode_means(&meta.means.mins, &meta.means.maxs, &means0, &means1, base, &mut c.center[..cnt * 3]); + decode_scales(&lx, &ly, &lz, &scales_img, base, &mut c.scale[..cnt * 3]); + decode_quats(&quats_img, base, &mut c.quat[..cnt * 4]); + decode_sh0(&lr, &lg, &lb, &la, &sh0_img, base, &mut c.rgb[..cnt * 3], &mut c.opacity[..cnt]); + if let Some(shn) = &shn { + shn.decode(base, cnt, c); + } + Ok(()) + }) } +// The decode_* helpers read the image from `base` and write `out` from 0 (len = cnt * components) + fn decode_means( mins: &[f32; 3], maxs: &[f32; 3], img0: &ImageData, img1: &ImageData, + base: usize, out_center: &mut [f32], -) -> anyhow::Result<()> { - let num_splats = out_center.len() / 3; - for i in 0..num_splats { - let i4 = i * 4; +) { + for i in 0..out_center.len() / 3 { + let i4 = (base + i) * 4; let fx = (img0.rgba[i4] as u32 + ((img1.rgba[i4] as u32) << 8)) as f32 / 65535.0; let fy = (img0.rgba[i4 + 1] as u32 + ((img1.rgba[i4 + 1] as u32) << 8)) as f32 / 65535.0; let fz = (img0.rgba[i4 + 2] as u32 + ((img1.rgba[i4 + 2] as u32) << 8)) as f32 / 65535.0; @@ -347,49 +419,24 @@ fn decode_means( out_center[i3 + 1] = y; out_center[i3 + 2] = z; } - Ok(()) -} - -fn decode_scales_v2(codebook: &[f32], img: &ImageData, out_scale: &mut [f32]) -> anyhow::Result<()> { - let num_splats = out_scale.len() / 3; - let lookup: Vec = codebook.iter().map(|v| v.exp()).collect(); - for i in 0..num_splats { - let i4 = i * 4; - let i3 = i * 3; - out_scale[i3] = lookup[img.rgba[i4] as usize]; - out_scale[i3 + 1] = lookup[img.rgba[i4 + 1] as usize]; - out_scale[i3 + 2] = lookup[img.rgba[i4 + 2] as usize]; - } - Ok(()) } -fn decode_scales_v1(mins: &[f32; 3], maxs: &[f32; 3], img: &ImageData, out_scale: &mut [f32]) -> anyhow::Result<()> { - let num_splats = out_scale.len() / 3; - let build_lookup = |min: f32, max: f32| -> Vec { - (0..256) - .map(|i| min + (max - min) * (i as f32 / 255.0)) - .map(|ln| ln.exp()) - .collect() - }; - let lx = build_lookup(mins[0], maxs[0]); - let ly = build_lookup(mins[1], maxs[1]); - let lz = build_lookup(mins[2], maxs[2]); - for i in 0..num_splats { - let i4 = i * 4; +/// v2: same codebook on all 3 axes; v1: per-axis min/max lookup +fn decode_scales(lx: &[f32], ly: &[f32], lz: &[f32], img: &ImageData, base: usize, out_scale: &mut [f32]) { + for i in 0..out_scale.len() / 3 { + let i4 = (base + i) * 4; let i3 = i * 3; out_scale[i3] = lx[img.rgba[i4] as usize]; out_scale[i3 + 1] = ly[img.rgba[i4 + 1] as usize]; out_scale[i3 + 2] = lz[img.rgba[i4 + 2] as usize]; } - Ok(()) } -fn decode_quats(img: &ImageData, out_quat: &mut [f32]) -> anyhow::Result<()> { - let num_splats = out_quat.len() / 4; +fn decode_quats(img: &ImageData, base: usize, out_quat: &mut [f32]) { const SQRT2: f32 = std::f32::consts::SQRT_2; - let lookup: Vec = (0..256).map(|i| (i as f32 / 255.0 - 0.5) * SQRT2).collect(); - for i in 0..num_splats { - let i4 = i * 4; + let lookup: [f32; 256] = std::array::from_fn(|i| (i as f32 / 255.0 - 0.5) * SQRT2); + for i in 0..out_quat.len() / 4 { + let i4 = (base + i) * 4; let r0 = lookup[img.rgba[i4] as usize]; let r1 = lookup[img.rgba[i4 + 1] as usize]; let r2 = lookup[img.rgba[i4 + 2] as usize]; @@ -405,198 +452,76 @@ fn decode_quats(img: &ImageData, out_quat: &mut [f32]) -> anyhow::Result<()> { out_quat[o + 2] = quat_z; out_quat[o + 3] = quat_w; } - Ok(()) } -fn decode_sh0_v2(codebook: &[f32], img: &ImageData, out_rgb: &mut [f32], out_opacity: &mut [f32]) -> anyhow::Result<()> { - let num_splats = out_opacity.len(); - let lookup_rgb: Vec = codebook.iter().map(|v| SH_C0 * v + 0.5).collect(); - let lookup_a: Vec = (0..256).map(|i| i as f32 / 255.0).collect(); - for i in 0..num_splats { - let i4 = i * 4; - let i3 = i * 3; - out_rgb[i3] = lookup_rgb[img.rgba[i4] as usize]; - out_rgb[i3 + 1] = lookup_rgb[img.rgba[i4 + 1] as usize]; - out_rgb[i3 + 2] = lookup_rgb[img.rgba[i4 + 2] as usize]; - out_opacity[i] = lookup_a[img.rgba[i4 + 3] as usize]; - } - Ok(()) -} - -fn decode_sh0_v1(mins: &[f32; 4], maxs: &[f32; 4], img: &ImageData, out_rgb: &mut [f32], out_opacity: &mut [f32]) -> anyhow::Result<()> { - let num_splats = out_opacity.len(); - let build_lookup = |min: f32, max: f32, post: fn(f32) -> f32| -> Vec { - (0..256) - .map(|i| min + (max - min) * (i as f32 / 255.0)) - .map(post) - .collect() - }; - let lr = build_lookup(mins[0], maxs[0], |v| SH_C0 * v + 0.5); - let lg = build_lookup(mins[1], maxs[1], |v| SH_C0 * v + 0.5); - let lb = build_lookup(mins[2], maxs[2], |v| SH_C0 * v + 0.5); - let la = build_lookup(mins[3], maxs[3], |v| 1.0 / (1.0 + (-v).exp())); - for i in 0..num_splats { - let i4 = i * 4; +fn decode_sh0( + lr: &[f32], + lg: &[f32], + lb: &[f32], + la: &[f32], + img: &ImageData, + base: usize, + out_rgb: &mut [f32], + out_opacity: &mut [f32], +) { + for i in 0..out_opacity.len() { + let i4 = (base + i) * 4; let i3 = i * 3; out_rgb[i3] = lr[img.rgba[i4] as usize]; out_rgb[i3 + 1] = lg[img.rgba[i4 + 1] as usize]; out_rgb[i3 + 2] = lb[img.rgba[i4 + 2] as usize]; out_opacity[i] = la[img.rgba[i4 + 3] as usize]; } - Ok(()) } -fn decode_shn_v2( - shn: ShNV2, - get_file: &mut dyn FnMut(&str) -> anyhow::Result>, - num_splats: usize, - sh1: &mut [f32], - sh2: &mut [f32], - sh3: &mut [f32], -) -> anyhow::Result<()> { - let centroids = decode_image(&get_file(&shn.files[0])?)?; - let labels = decode_image(&get_file(&shn.files[1])?)?; - let lookup = shn.codebook; - let use_sh1 = shn.bands >= 1; - let use_sh2 = shn.bands >= 2; - let use_sh3 = shn.bands >= 3; - - for i in 0..num_splats { - let i4 = i * 4; - let label = labels.rgba[i4] as u16 | ((labels.rgba[i4 + 1] as u16) << 8); - let stride = if use_sh3 { 15 } else if use_sh2 { 8 } else { 3 }; - let col = (label & 63) as usize * stride; - let row = (label >> 6) as usize; - let offset = row * centroids.width + col; - for d in 0..3 { - if use_sh1 { - for k in 0..3 { - sh1[i * 9 + k * 3 + d] = lookup[centroids.rgba[(offset + k) * 4 + d] as usize]; - } - } - if use_sh2 { - for k in 0..5 { - sh2[i * 15 + k * 3 + d] = lookup[centroids.rgba[(offset + 3 + k) * 4 + d] as usize]; - } - } - if use_sh3 { - for k in 0..7 { - sh3[i * 21 + k * 3 + d] = lookup[centroids.rgba[(offset + 8 + k) * 4 + d] as usize]; - } - } - } +/// shN images (centroids + labels) loaded once, decoded per chunk +struct ShN { + centroids: ImageData, + labels: ImageData, + lookup: Vec, + degree: usize, +} + +impl ShN { + fn load( + get_file: &mut dyn FnMut(&str) -> anyhow::Result>, + files: &[String; 2], + lookup: Vec, + degree: usize, + ) -> anyhow::Result { + let centroids = load_image(get_file, &files[0], "shN centroids")?; + let labels = load_image(get_file, &files[1], "shN labels")?; + Ok(Self { centroids, labels, lookup, degree }) } - Ok(()) -} -fn decode_shn_v1( - shn: ShNV1, - get_file: &mut dyn FnMut(&str) -> anyhow::Result>, - num_splats: usize, - max_sh_degree: usize, - sh1: &mut [f32], - sh2: &mut [f32], - sh3: &mut [f32], -) -> anyhow::Result<()> { - let centroids = decode_image(&get_file(&shn.files[0])?)?; - let labels = decode_image(&get_file(&shn.files[1])?)?; - let lookup: Vec = (0..256) - .map(|i| shn.mins + (shn.maxs - shn.mins) * (i as f32 / 255.0)) - .collect(); - - for i in 0..num_splats { - let i4 = i * 4; - let label = labels.rgba[i4] as u16 | ((labels.rgba[i4 + 1] as u16) << 8); - let stride = if max_sh_degree >= 3 { 15 } else if max_sh_degree >= 2 { 8 } else { 3 }; - let col = (label & 63) as usize * stride; - let row = (label >> 6) as usize; - let offset = row * centroids.width + col; - for d in 0..3 { - if max_sh_degree >= 1 { - for k in 0..3 { - sh1[i * 9 + k * 3 + d] = lookup[centroids.rgba[(offset + k) * 4 + d] as usize]; + fn decode(&self, base: usize, cnt: usize, c: &mut Chunk) { + let (centroids, lookup) = (&self.centroids, &self.lookup); + let stride = if self.degree >= 3 { 15 } else if self.degree >= 2 { 8 } else { 3 }; + for i in 0..cnt { + let i4 = (base + i) * 4; + let label = self.labels.rgba[i4] as u16 | ((self.labels.rgba[i4 + 1] as u16) << 8); + let col = (label & 63) as usize * stride; + let row = (label >> 6) as usize; + let offset = row * centroids.width + col; + for d in 0..3 { + if self.degree >= 1 { + for k in 0..3 { + c.sh1[i * 9 + k * 3 + d] = lookup[centroids.rgba[(offset + k) * 4 + d] as usize]; + } } - } - if max_sh_degree >= 2 { - for k in 0..5 { - sh2[i * 15 + k * 3 + d] = lookup[centroids.rgba[(offset + 3 + k) * 4 + d] as usize]; + if self.degree >= 2 { + for k in 0..5 { + c.sh2[i * 15 + k * 3 + d] = lookup[centroids.rgba[(offset + 3 + k) * 4 + d] as usize]; + } } - } - if max_sh_degree >= 3 { - for k in 0..7 { - sh3[i * 21 + k * 3 + d] = lookup[centroids.rgba[(offset + 8 + k) * 4 + d] as usize]; + if self.degree >= 3 { + for k in 0..7 { + c.sh3[i * 21 + k * 3 + d] = lookup[centroids.rgba[(offset + 8 + k) * 4 + d] as usize]; + } } } } } - Ok(()) -} - -fn emit_to_receiver( - splats: &mut T, - num_splats: usize, - max_sh_degree: usize, - center: &[f32], - scale: &[f32], - quat: &[f32], - rgb: &[f32], - opacity: &[f32], - sh1: &[f32], - sh2: &[f32], - sh3: &[f32], -) -> anyhow::Result<()> { - let mut base = 0usize; - while base < num_splats { - let count = (num_splats - base).min(MAX_SPLAT_CHUNK); - let i3 = base * 3; - let i4 = base * 4; - splats.set_batch( - base, - count, - &SplatProps { - center: ¢er[i3..i3 + count * 3], - opacity: &opacity[base..base + count], - rgb: &rgb[i3..i3 + count * 3], - scale: &scale[i3..i3 + count * 3], - quat: &quat[i4..i4 + count * 4], - sh1: if max_sh_degree >= 1 { &sh1[base * 9..base * 9 + count * 9] } else { &[] }, - sh2: if max_sh_degree >= 2 { &sh2[base * 15..base * 15 + count * 15] } else { &[] }, - sh3: if max_sh_degree >= 3 { &sh3[base * 21..base * 21 + count * 21] } else { &[] }, - ..Default::default() - }, - ); - base += count; - } - splats.finish() -} - -fn preload_all( - meta: &PcSogsRoot, - prefix: &str, - zip: &mut ZipArchive>, - cache: &mut HashMap>, -) -> anyhow::Result<()> { - match meta { - PcSogsRoot::V2(v2) => { - for f in v2.means.files.iter() { preload_file(zip, prefix, cache, f)?; } - for f in v2.scales.files.iter() { preload_file(zip, prefix, cache, f)?; } - for f in v2.quats.files.iter() { preload_file(zip, prefix, cache, f)?; } - for f in v2.sh0.files.iter() { preload_file(zip, prefix, cache, f)?; } - if let Some(shn) = &v2.shn { - for f in shn.files.iter() { preload_file(zip, prefix, cache, f)?; } - } - } - PcSogsRoot::V1(v1) => { - for f in v1.means.files.iter() { preload_file(zip, prefix, cache, f)?; } - for f in v1.scales.files.iter() { preload_file(zip, prefix, cache, f)?; } - for f in v1.quats.files.iter() { preload_file(zip, prefix, cache, f)?; } - for f in v1.sh0.files.iter() { preload_file(zip, prefix, cache, f)?; } - if let Some(shn) = &v1.shn { - for f in shn.files.iter() { preload_file(zip, prefix, cache, f)?; } - } - } - } - Ok(()) } fn preload_file( @@ -630,11 +555,6 @@ struct ImageData { height: usize, } -fn decode_rgba(bytes: &[u8]) -> anyhow::Result { - let img = decode_image(bytes)?; - Ok(img) -} - fn decode_image(bytes: &[u8]) -> anyhow::Result { let img = ImageReader::new(Cursor::new(bytes)) .with_guessed_format()? @@ -648,3 +568,224 @@ fn decode_image(bytes: &[u8]) -> anyhow::Result { } + +#[cfg(test)] +mod tests { + use super::*; + + fn png(rgba: Vec, w: u32, h: u32) -> Vec { + let img = image::RgbaImage::from_raw(w, h, rgba).unwrap(); + let mut out = Vec::new(); + DynamicImage::ImageRgba8(img) + .write_to(&mut Cursor::new(&mut out), image::ImageFormat::Png) + .unwrap(); + out + } + + /// Receiver that records every f32 value (GsplatArray stores f16) + #[derive(Default)] + struct Rec { + center: Vec, + scale: Vec, + quat: Vec, + rgb: Vec, + opacity: Vec, + sh: [Vec; 3], + finished: bool, + } + + impl SplatReceiver for Rec { + fn init_splats(&mut self, init: &SplatInit) -> anyhow::Result<()> { + let n = init.num_splats; + (self.center, self.scale, self.quat, self.rgb, self.opacity) = + (vec![0.0; n * 3], vec![0.0; n * 3], vec![0.0; n * 4], vec![0.0; n * 3], vec![0.0; n]); + for (d, k) in [9, 15, 21].into_iter().enumerate() { + self.sh[d] = if init.max_sh_degree > d { vec![0.0; n * k] } else { Vec::new() }; + } + Ok(()) + } + fn finish(&mut self) -> anyhow::Result<()> { self.finished = true; Ok(()) } + fn set_batch(&mut self, base: usize, count: usize, b: &SplatProps) { + let put = |dst: &mut Vec, src: &[f32], k: usize| { + assert_eq!(src.len(), count * k); + dst[base * k..(base + count) * k].copy_from_slice(src); + }; + put(&mut self.center, b.center, 3); + put(&mut self.scale, b.scale, 3); + put(&mut self.quat, b.quat, 4); + put(&mut self.rgb, b.rgb, 3); + put(&mut self.opacity, b.opacity, 1); + for (d, (src, k)) in [(b.sh1, 9), (b.sh2, 15), (b.sh3, 21)].into_iter().enumerate() { + if !self.sh[d].is_empty() { put(&mut self.sh[d], src, k); } + } + } + fn set_center(&mut self, _: usize, _: usize, _: &[f32]) { unimplemented!() } + fn set_opacity(&mut self, _: usize, _: usize, _: &[f32]) { unimplemented!() } + fn set_rgb(&mut self, _: usize, _: usize, _: &[f32]) { unimplemented!() } + fn set_rgba(&mut self, _: usize, _: usize, _: &[f32]) { unimplemented!() } + fn set_scale(&mut self, _: usize, _: usize, _: &[f32]) { unimplemented!() } + fn set_quat(&mut self, _: usize, _: usize, _: &[f32]) { unimplemented!() } + } + + const N: usize = 70_000; // > MAX_SPLAT_CHUNK: at least 2 chunks, the last one partial + const W: usize = 512; + const H: usize = N.div_ceil(W); + const CW: usize = 64 * 15; + const CH: usize = 16; + const MINS: [f32; 3] = [-2.0, -1.0, -3.0]; + const MAXS: [f32; 3] = [2.0, 3.0, 1.5]; + + struct Synth { + ml: Vec, mu: Vec, s: Vec, q: Vec, c: Vec, cent: Vec, lab: Vec, + cb: [Vec; 3], // codebook scales, sh0, shN (v2) + } + + fn xorshift(seed: u32) -> impl FnMut() -> u32 { + let mut x = seed; + move || { x ^= x << 13; x ^= x >> 17; x ^= x << 5; x } + } + + fn synth() -> Synth { + let mut rnd = xorshift(0x1234_5678); + let mut img = |n: usize, f: fn(usize, u32) -> u8| -> Vec { + (0..n * 4).map(|i| f(i % 4, rnd())).collect() + }; + let ml = img(W * H, |_, r| r as u8); + let mu = img(W * H, |_, r| r as u8); + let s = img(W * H, |_, r| r as u8); + let q = img(W * H, |ch, r| if ch == 3 { 252 + (r % 4) as u8 } else { r as u8 }); + let c = img(W * H, |_, r| r as u8); + let cent = img(CW * CH, |_, r| r as u8); + let lab = img(W * H, |ch, r| if ch == 1 { (r % 4) as u8 } else { r as u8 }); + let mut rf = xorshift(0x9e37_79b9); + let mut cb = || -> Vec { (0..256).map(|_| (rf() % 2000) as f32 / 1000.0 - 1.0).collect() }; + let cb = [cb(), cb(), cb()]; + Synth { ml, mu, s, q, c, cent, lab, cb } + } + + /// Reference: formulas evaluated per splat, without chunking + fn reference(sy: &Synth, v1: bool) -> Rec { + let mut r = Rec::default(); + r.init_splats(&SplatInit { num_splats: N, max_sh_degree: 3, lod_tree: false }).unwrap(); + let lin = |min: f32, max: f32, b: u8| min + (max - min) * (b as f32 / 255.0); + for j in 0..N { + let p = j * 4; + for d in 0..3 { + let f = (sy.ml[p + d] as u32 + ((sy.mu[p + d] as u32) << 8)) as f32 / 65535.0; + let v = MINS[d] + (MAXS[d] - MINS[d]) * f; + r.center[j * 3 + d] = v.signum() * (v.abs().exp() - 1.0); + r.scale[j * 3 + d] = if v1 { lin(-3.0, 1.0, sy.s[p + d]).exp() } else { sy.cb[0][sy.s[p + d] as usize].exp() }; + let col = if v1 { lin(-2.0, 2.0, sy.c[p + d]) } else { sy.cb[1][sy.c[p + d] as usize] }; + r.rgb[j * 3 + d] = SH_C0 * col + 0.5; + } + let a = sy.c[p + 3]; + r.opacity[j] = if v1 { 1.0 / (1.0 + (-lin(-4.0, 4.0, a)).exp()) } else { a as f32 / 255.0 }; + let qv = |b: u8| (b as f32 / 255.0 - 0.5) * std::f32::consts::SQRT_2; + let (r0, r1, r2) = (qv(sy.q[p]), qv(sy.q[p + 1]), qv(sy.q[p + 2])); + let rr = (1.0 - (r0 * r0 + r1 * r1 + r2 * r2)).max(0.0).sqrt(); + let quat = match sy.q[p + 3] - 252 { + 0 => [r0, r1, r2, rr], + 1 => [rr, r1, r2, r0], + 2 => [r1, rr, r2, r0], + _ => [r1, r2, rr, r0], + }; + r.quat[j * 4..j * 4 + 4].copy_from_slice(&quat); + let label = sy.lab[p] as usize | ((sy.lab[p + 1] as usize) << 8); + let off = (label >> 6) * CW + (label & 63) * 15; + for (deg, (first, k)) in [(0, 3), (3, 5), (8, 7)].into_iter().enumerate() { + for kk in 0..k { + for d in 0..3 { + let b = sy.cent[(off + first + kk) * 4 + d]; + r.sh[deg][j * k * 3 + kk * 3 + d] = if v1 { lin(-0.5, 0.5, b) } else { sy.cb[2][b as usize] }; + } + } + } + } + r.finished = true; + r + } + + fn assert_same(got: &Rec, exp: &Rec) { + assert!(got.finished); + let pairs = [ + ("center", &got.center, &exp.center), ("scale", &got.scale, &exp.scale), + ("quat", &got.quat, &exp.quat), ("rgb", &got.rgb, &exp.rgb), ("opacity", &got.opacity, &exp.opacity), + ("sh1", &got.sh[0], &exp.sh[0]), ("sh2", &got.sh[1], &exp.sh[1]), ("sh3", &got.sh[2], &exp.sh[2]), + ]; + for (name, g, e) in pairs { + assert_eq!(g.len(), e.len(), "{name} len"); + for (i, (a, b)) in g.iter().zip(e.iter()).enumerate() { + assert!((a - b).abs() <= 1e-6 * (1.0 + b.abs()), "{name}[{i}]: {a} vs {b}"); + } + } + } + + fn synth_files(sy: &Synth) -> HashMap> { + let (w, h) = (W as u32, H as u32); + [ + ("ml.png", png(sy.ml.clone(), w, h)), ("mu.png", png(sy.mu.clone(), w, h)), + ("s.png", png(sy.s.clone(), w, h)), ("q.png", png(sy.q.clone(), w, h)), + ("c.png", png(sy.c.clone(), w, h)), ("lab.png", png(sy.lab.clone(), w, h)), + ("cent.png", png(sy.cent.clone(), CW as u32, CH as u32)), + ].into_iter().map(|(k, v)| (k.to_string(), v)).collect() + } + + fn meta_v2(sy: &Synth) -> Vec { + serde_json::json!({ + "version": 2, "count": N, + "means": {"files": ["ml.png", "mu.png"], "mins": MINS, "maxs": MAXS}, + "scales": {"files": ["s.png"], "codebook": sy.cb[0]}, + "quats": {"files": ["q.png"]}, + "sh0": {"files": ["c.png"], "codebook": sy.cb[1]}, + "shN": {"files": ["cent.png", "lab.png"], "codebook": sy.cb[2], "bands": 3}, + }).to_string().into_bytes() + } + + /// Build an in-memory SOG zip (stored entries, nested under a folder prefix) + fn sog_zip(meta: &[u8], files: HashMap>) -> Vec { + use std::io::Write; + let mut zbuf = Vec::new(); + { + let mut zw = zip::ZipWriter::new(Cursor::new(&mut zbuf)); + let opts = zip::write::SimpleFileOptions::default().compression_method(zip::CompressionMethod::Stored); + zw.start_file("scene/meta.json", opts).unwrap(); + zw.write_all(meta).unwrap(); + for (name, bytes) in files { + zw.start_file(format!("scene/{name}"), opts).unwrap(); + zw.write_all(&bytes).unwrap(); + } + zw.finish().unwrap(); + } + zbuf + } + + /// Feed the zip to SogsDecoder in several pushes, as the streaming loader does + fn decode_zip(zbuf: &[u8]) -> Rec { + let mut dec = SogsDecoder::new(Rec::default(), None); + for part in zbuf.chunks(100_000) { dec.push(part).unwrap(); } + dec.finish().unwrap(); + assert!(dec.buffer.is_empty()); + dec.into_splats() + } + + #[test] + fn decode_v2_chunked_matches_reference() { + let sy = synth(); + let got = decode_zip(&sog_zip(&meta_v2(&sy), synth_files(&sy))); + assert_same(&got, &reference(&sy, false)); + } + + #[test] + fn decode_v1_chunked_matches_reference() { + let sy = synth(); + let meta = serde_json::json!({ + "means": {"files": ["ml.png", "mu.png"], "mins": MINS, "maxs": MAXS, "shape": [N, 3]}, + "scales": {"files": ["s.png"], "mins": [-3.0, -3.0, -3.0], "maxs": [1.0, 1.0, 1.0]}, + "quats": {"files": ["q.png"], "encoding": "quaternion_packed"}, + "sh0": {"files": ["c.png"], "mins": [-2.0, -2.0, -2.0, -4.0], "maxs": [2.0, 2.0, 2.0, 4.0]}, + "shN": {"files": ["cent.png", "lab.png"], "mins": -0.5, "maxs": 0.5, "shape": [N, 45]}, + }).to_string().into_bytes(); + let got = decode_zip(&sog_zip(&meta, synth_files(&sy))); + assert_same(&got, &reference(&sy, true)); + } +}