storytold/spark
2
0
Fork 0
mirror of https://github.com/storytold/spark.git synced 2026-10-09 00:09:53 +00:00
Code Issues Packages Projects Releases Wiki Activity

Remove obsolete code

Browse Source
This commit is contained in:
Andreas Sundquist
2025-05-30 10:49:49 -07:00
parent 7defe60edb
commit 8b5f854848
25 changed files with 5267 additions and 7213 deletions
Download Patch File Download Diff File Expand all files Collapse all files
rust/Cargo.lock
Generated
-56
View File
@@ -8,15 +8,6 @@ version = "1.0.98"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e16d2d3311acee920a9eb8d33b8cbc1787ce4a264e85f964c2404b969bdcd487"
[[package]]
name = "bincode"
version = "1.3.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b1f45e9417d87227c7a56d22e471c6206462cba514c7590c09aff4cf6d1ddcad"
dependencies = [
"serde",
]
[[package]]
name = "bumpalo"
version = "3.17.0"
@@ -43,7 +34,6 @@ dependencies = [
"half",
"js-sys",
"wasm-bindgen",
"wlg",
]
[[package]]
@@ -102,35 +92,6 @@ version = "1.0.20"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "eded382c5f5f786b989652c49544c4877d9f015cc22e145a5ea8ea66c2921cd2"
[[package]]
name = "ruzstd"
version = "0.8.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c581601827da5c717bfae77d7b187e54293d23d8fb6b700b4b5e9b5828a13cc3"
dependencies = [
"twox-hash",
]
[[package]]
name = "serde"
version = "1.0.219"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5f0e2c6ed6606019b4e29e69dbaba95b11854410e5347d525002456dbbb786b6"
dependencies = [
"serde_derive",
]
[[package]]
name = "serde_derive"
version = "1.0.219"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5b0276cf7f2c73365f7157c8123c21cd9a50fbbd844757af28ca1f5925fc2a00"
dependencies = [
"proc-macro2",
"quote",
"syn",
]
[[package]]
name = "syn"
version = "2.0.100"
@@ -142,12 +103,6 @@ dependencies = [
"unicode-ident",
]
[[package]]
name = "twox-hash"
version = "2.1.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e7b17f197b3050ba473acf9181f7b1d3b66d1cf7356c6cc57886662276e65908"
[[package]]
name = "unicode-ident"
version = "1.0.18"
@@ -211,14 +166,3 @@ checksum = "1a05d73b933a847d6cccdda8f838a22ff101ad9bf93e33684f39c1f5f0eece3d"
dependencies = [
"unicode-ident",
]
[[package]]
name = "wlg"
version = "0.1.0"
dependencies = [
"anyhow",
"bincode",
"half",
"ruzstd",
"serde",
]
rust/Cargo.toml
-4
View File
@@ -1,7 +1,6 @@
[workspace]
members = [
"forge-internal-rs",
"wlg",
]
resolver = "2"
@@ -14,9 +13,6 @@ repository = "https://github.com/forge-gfx/forge"
[workspace.dependencies]
anyhow = "1.0.98"
bincode = "1.3.3"
half = "2.6.0"
js-sys = "0.3.77"
ruzstd = "0.8.0"
serde = { version = "1.0.219", features = ["derive"] }
wasm-bindgen = "0.2.100"
rust/forge-internal-rs/Cargo.toml
-1
View File
@@ -15,4 +15,3 @@ anyhow.workspace = true
half.workspace = true
js-sys.workspace = true
wasm-bindgen.workspace = true
wlg = { path = "../wlg" }
rust/forge-internal-rs/src/lib.rs
+2 -77
View File
@@ -1,12 +1,10 @@
use std::cell::RefCell;
use js_sys::{Array, ArrayBuffer, Float32Array, Object, Reflect, Uint16Array, Uint32Array, Uint8Array};
use js_sys::{Float32Array, Uint16Array, Uint32Array};
use wasm_bindgen::prelude::*;
use wlg::decode12_packed;
mod sort;
use sort::{old_sort_internal, OldSortBuffers, sort_internal, SortBuffers};
use sort::{sort_internal, SortBuffers};
mod raycast;
use raycast::{raycast_ellipsoids, raycast_spheres};
@@ -15,82 +13,9 @@ const RAYCAST_BUFFER_COUNT: u32 = 65536;
thread_local! {
static SORT_BUFFERS: RefCell<SortBuffers> = RefCell::new(SortBuffers::default());
static OLD_SORT_BUFFERS: RefCell<OldSortBuffers> = RefCell::new(OldSortBuffers::default());
static RAYCAST_BUFFER: RefCell<Vec<u32>> = RefCell::new(vec![0; RAYCAST_BUFFER_COUNT as usize * 4]);
}
#[wasm_bindgen]
pub fn decode_wlg(bytes: ArrayBuffer, tex_width: u32, tex_height: u32) -> Object {
let mut bytes = Uint8Array::new(&bytes).to_vec();
let (_settings, packed_splats) = match decode12_packed(&mut bytes) {
Ok(result) => result,
Err(err) => {
wasm_bindgen::throw_str(&format!("{}", err));
}
};
let num_splats = packed_splats.0.len() / 4;
let max_splats = if tex_width != 0 && tex_height != 0 {
let width = tex_width as usize;
let height = num_splats.div_ceil(width).min(tex_height as usize);
let depth = num_splats.div_ceil(width * height);
width * height * depth
} else {
num_splats
};
let packed = Uint32Array::new_with_length(max_splats as u32 * 4);
let packed_slice = packed.subarray(0, packed_splats.0.len() as u32);
packed_slice.copy_from(&packed_splats.0);
let result = Object::new();
Reflect::set(&result, &JsValue::from_str("numSplats"), &JsValue::from_f64(num_splats as f64)).unwrap();
Reflect::set(&result, &JsValue::from_str("packedSplats"), &JsValue::from(packed)).unwrap();
result
}
#[wasm_bindgen]
pub fn old_sort_splats(
max_splats: u32, total_splats: u32, readback: Array, ordering: Uint32Array,
) -> u32 {
let max_splats = max_splats as usize;
let total_splats = total_splats as usize;
let num_layers = readback.length() as usize;
let layer_size = readback.get(0).dyn_into::<Uint8Array>().unwrap().length() as usize / 4;
let active_splats = OLD_SORT_BUFFERS.with_borrow_mut(|buffers| {
buffers.ensure_size(max_splats);
// Copy the readback data layers into a contiguous buffer
let mut layer_base = 0;
for layer in 0..num_layers {
let layer_count = layer_size.min(total_splats - layer_base);
if layer_count > 0 {
let layer_buffer = readback.get(layer as u32).dyn_into::<Uint8Array>().unwrap().buffer();
let layer_uint32 = Uint32Array::new_with_byte_offset_and_length(&layer_buffer, 0, layer_count as u32);
layer_uint32.copy_to(&mut buffers.readback[layer_base..layer_base + layer_count]);
}
layer_base += layer_count;
}
let active_splats = match old_sort_internal(buffers, total_splats) {
Ok(active_splats) => active_splats,
Err(err) => {
wasm_bindgen::throw_str(&format!("{}", err));
}
};
if active_splats > 0 {
// Copy out ordering result
let subarray = &buffers.ordering[..active_splats as usize];
ordering.subarray(0, active_splats).copy_from(&subarray);
}
active_splats
});
active_splats
}
#[wasm_bindgen]
pub fn sort_splats(
num_splats: u32, readback: Uint16Array, ordering: Uint32Array,
rust/forge-internal-rs/src/raycast.rs
+21 -1
View File
@@ -1,8 +1,28 @@
use half::f16;
use wlg::decode_scale;
const MIN_OPACITY: f32 = 0.1;
pub const LN_SCALE_MIN: f32 = -9.0;
pub const LN_SCALE_MAX: f32 = 9.0;
pub const LN_RESCALE: f32 = (LN_SCALE_MAX - LN_SCALE_MIN) / 254.0; // 1..=255
// pub fn encode_scale(scale: f32) -> u8 {
// if scale == 0.0 {
// 0
// } else {
// // Allow scales below LN_SCALE_MIN to be encoded as 0, which signifies a 2DGS
// ((scale.ln() - LN_SCALE_MIN) / LN_RESCALE + 1.0).clamp(0.0, 255.0).round() as u8
// }
// }
pub fn decode_scale(scale: u8) -> f32 {
if scale == 0 {
0.0
} else {
(LN_SCALE_MIN + (scale - 1) as f32 * LN_RESCALE).exp()
}
}
pub fn raycast_spheres(
buffer: &[u32], distances: &mut Vec<f32>,
origin: [f32; 3], dir: [f32; 3], near: f32, far: f32,
rust/forge-internal-rs/src/sort.rs
-66
View File
@@ -3,72 +3,6 @@ use anyhow::anyhow;
const DEPTH_INFINITY: u32 = 0x7c00;
const DEPTH_SIZE: usize = DEPTH_INFINITY as usize + 1;
#[derive(Default)]
pub struct OldSortBuffers {
pub readback: Vec<u32>,
pub ordering: Vec<u32>,
pub buckets: Vec<u32>,
}
impl OldSortBuffers {
pub fn ensure_size(&mut self, max_splats: usize) {
if self.readback.len() < max_splats {
self.readback.resize(max_splats, 0);
}
if self.ordering.len() < max_splats {
self.ordering.resize(max_splats, 0);
}
if self.buckets.len() < DEPTH_SIZE {
self.buckets.resize(DEPTH_SIZE, 0);
}
}
}
pub fn old_sort_internal(buffers: &mut OldSortBuffers, total_splats: usize) -> anyhow::Result<u32> {
let OldSortBuffers { readback, ordering, buckets } = buffers;
let readback = &readback[..total_splats];
// Set the bucket counts to zero
buckets.clear();
buckets.resize(DEPTH_SIZE, 0);
// Count the number of splats in each bucket
for &readback_u32 in readback.iter() {
let priority = readback_u32 & 0x7FFF;
if priority < DEPTH_INFINITY {
buckets[priority as usize] += 1;
}
}
// Compute bucket starting offset
let mut active_splats = 0;
for count in buckets.iter_mut() {
let new_total = active_splats + *count;
*count = active_splats;
active_splats = new_total;
}
// Write out splat indices at the right location using bucket offsets
for (index, &readback_u32) in readback.iter().enumerate() {
let priority = readback_u32 & 0x7FFF;
if priority < DEPTH_INFINITY {
ordering[buckets[priority as usize] as usize] = index as u32;
buckets[priority as usize] += 1;
}
}
// Sanity check
if buckets[DEPTH_SIZE - 1] != active_splats {
return Err(anyhow!(
"Expected {} active splats but got {}",
active_splats,
buckets[DEPTH_SIZE - 1]
));
}
Ok(active_splats)
}
#[derive(Default)]
pub struct SortBuffers {
pub readback: Vec<u16>,
rust/wlg/Cargo.toml
-15
View File
@@ -1,15 +0,0 @@
[package]
name = "wlg"
version = "0.1.0"
rust-version.workspace = true
edition.workspace = true
license.workspace = true
authors.workspace = true
repository.workspace = true
[dependencies]
anyhow.workspace = true
bincode.workspace = true
half.workspace = true
serde.workspace = true
ruzstd.workspace = true
rust/wlg/src/lib.rs
-15
View File
@@ -1,15 +0,0 @@
mod ordering;
// WLG0 includes WLGv1 and WLGv2
mod wlg0;
pub use wlg0::{
Wlg0Gaussian, PackedSplats,
Wlg0Settings, Wlg0EncodeSettings,
decode12, decode12_packed,
encode12, encode_scale, decode_scale,
};
// WLG3 will be used for future WLG versions
// mod wlg3;
rust/wlg/src/ordering.rs
-201
View File
@@ -1,201 +0,0 @@
// Compute Morton index for 3D coordinates.
pub fn morton_coord_to_index([x, y, z]: [u16; 3]) -> u64 {
fn expand3(x: u16) -> u64 {
let mut x = x as u64;
x = (x | x << 32) & 0x1f00000000ffff;
x = (x | x << 16) & 0x1f0000ff0000ff;
x = (x | x << 8) & 0x100f00f00f00f00f;
x = (x | x << 4) & 0x10c30c30c30c30c3;
x = (x | x << 2) & 0x1249249249249249;
x
}
(expand3(x) << 0) | (expand3(y) << 1) | (expand3(z) << 2)
}
pub fn morton_coord_to_index_24([x, y, z]: [u32; 3]) -> u128 {
fn expand3_24(x: u32) -> u128 {
let mut x = x as u128;
x = (x | x << 64) & 0x3ff0000000000000000ffffffffu128;
x = (x | x << 32) & 0x3ff00000000ffff00000000ffffu128;
x = (x | x << 16) & 0x30000ff0000ff0000ff0000ff0000ffu128;
x = (x | x << 8) & 0x300f00f00f00f00f00f00f00f00f00fu128;
x = (x | x << 4) & 0x30c30c30c30c30c30c30c30c30c30c3u128;
x = (x | x << 2) & 0x9249249249249249249249249249249u128;
x
}
(expand3_24(x) << 0) | (expand3_24(y) << 1) | (expand3_24(z) << 2)
}
// Compute Hilbert index for 3D coordinates.
// Converts a 48-bit Hilbert index to 3D coordinates (x, y, z).
pub fn _hilbert_index_to_coord(mut index: u64) -> (u16, u16, u16) {
const BITS: u32 = 16;
// Helper function to decode a Hilbert quad to coordinate bits.
fn hilbert_decode_step(quad: u16, s: &mut [u16; 3]) {
let (mut x, mut y, mut z) = (0u16, 0u16, 0u16);
match quad {
0 => { x = 0; y = 0; z = 0; },
1 => { x = 0; y = 0; z = 1; },
2 => { x = 0; y = 1; z = 1; },
3 => { x = 0; y = 1; z = 0; },
4 => { x = 1; y = 1; z = 0; },
5 => { x = 1; y = 1; z = 1; },
6 => { x = 1; y = 0; z = 1; },
7 => { x = 1; y = 0; z = 0; },
_ => {},
}
s[0] = x;
s[1] = y;
s[2] = z;
}
let mut mask = 1u16; // Start with the least significant bit
let mut h = [0u16; 3];
let mut s = [0u16; 3];
for _ in 0..BITS {
let quad = (index & 7) as u16; // Extract the last 3 bits
index >>= 3;
hilbert_decode_step(quad, &mut s);
h[0] |= s[0] * mask;
h[1] |= s[1] * mask;
h[2] |= s[2] * mask;
mask <<= 1; // Move to the next bit
}
(h[0], h[1], h[2])
}
// Converts 3D coordinates (x, y, z) to a 48-bit Hilbert index.
pub fn hilbert_coord_to_index([x, y, z]: [u16; 3]) -> u64 {
const BITS: u32 = 16;
// Helper function to encode coordinate bits to a Hilbert quad.
fn hilbert_encode_step(s: &mut [u16; 3]) -> u16 {
let x = s[0];
let y = s[1];
let z = s[2];
let mut quad = 0u16;
if x == 0 && y == 0 && z == 0 { quad = 0; }
else if x == 0 && y == 0 && z == 1 { quad = 1; }
else if x == 0 && y == 1 && z == 1 { quad = 2; }
else if x == 0 && y == 1 && z == 0 { quad = 3; }
else if x == 1 && y == 1 && z == 0 { quad = 4; }
else if x == 1 && y == 1 && z == 1 { quad = 5; }
else if x == 1 && y == 0 && z == 1 { quad = 6; }
else if x == 1 && y == 0 && z == 0 { quad = 7; }
quad
}
let mut index = 0u64;
let mut h = [x, y, z];
let mut s = [0u16; 3];
for _ in 0..BITS {
s[0] = h[0] & 1;
s[1] = h[1] & 1;
s[2] = h[2] & 1;
let quad = hilbert_encode_step(&mut s);
index <<= 3;
index |= quad as u64;
h[0] >>= 1;
h[1] >>= 1;
h[2] >>= 1;
}
index
}
// Converts a 72-bit Hilbert index to 3D coordinates (x, y, z).
pub fn _hilbert_index_coord_24(index: u128) -> (u32, u32, u32) {
const BITS: u32 = 24;
// Helper function to decode a Hilbert quad to coordinates.
fn hilbert_decode_step(quad: u32, s: &mut [u32; 3]) {
match quad {
0 => { s.swap(0, 1); },
1 => {},
2 => {},
3 => { s[2] ^= 1; },
4 => { s[0] ^= 1; s[1] ^= 1; s[2] ^= 1; },
5 => { s[0] ^= 1; s[1] ^= 1; s[2] ^= 1; s.swap(0, 1); },
6 => { s[0] ^= 1; s[1] ^= 1; s[2] ^= 1; s.swap(1, 2); },
7 => { s.swap(0, 2); },
_ => {},
}
}
let mut idx = index;
let mut mask = 1u32 << (BITS - 1);
let mut h = [0u32; 3];
let mut s = [0u32; 3];
for _ in 0..BITS {
let quad = (idx & 7) as u32;
idx >>= 3;
hilbert_decode_step(quad, &mut s);
h[0] |= s[0] & mask;
h[1] |= s[1] & mask;
h[2] |= s[2] & mask;
mask >>= 1;
}
(h[0], h[1], h[2])
}
// Converts 3D coordinates (x, y, z) to a 72-bit Hilbert index.
pub fn hilbert_coord_to_index_24([x, y, z]: [u32; 3]) -> u128 {
const BITS: u32 = 24;
// Helper function to encode coordinates to a Hilbert quad.
fn hilbert_encode_step(s: &mut [u32; 3]) -> u32 {
let mut quad = 0u32;
let mut bits = [s[0], s[1], s[2]];
if bits[0] > bits[1] { bits.swap(0, 1); quad ^= 1; }
if bits[1] > bits[2] { bits.swap(1, 2); quad ^= 2; }
if bits[0] > bits[1] { bits.swap(0, 1); quad ^= 1; }
s[0] = bits[0];
s[1] = bits[1];
s[2] = bits[2];
quad
}
let mut index = 0u128;
let h = [x, y, z];
let mut s = [0u32; 3];
let mut mask = 1u32 << (BITS - 1);
for _ in 0..BITS {
s[0] = (h[0] & mask) >> (BITS - 1);
s[1] = (h[1] & mask) >> (BITS - 1);
s[2] = (h[2] & mask) >> (BITS - 1);
let quad = hilbert_encode_step(&mut s);
index = (index << 3) | quad as u128;
mask >>= 1;
}
index
}
rust/wlg/src/wlg0/decode.rs
-268
View File
@@ -1,268 +0,0 @@
use anyhow::anyhow;
use std::io::Read;
use half::f16;
use super::{
decompress_data, deobfuscate, encode_scale, first_cumulative,
PackedSplats, Wlg0Gaussian, Wlg0Header, Wlg0Settings,
};
pub(super) trait Wlg0Decoder {
fn init_num_splats(&mut self, num_splats: usize);
fn write_center(&mut self, index: usize, dim: usize, value: f32);
fn write_scale(&mut self, index: usize, dim: usize, value: f32);
fn write_quaternion_f32(&mut self, index: usize, dim: usize, value: f32);
fn write_quaternion_i8(&mut self, index: usize, dim: usize, value: i8);
fn write_quaternion_full_f32(&mut self, index: usize, value: [f32; 4]);
fn write_rgba(&mut self, index: usize, dim: usize, value: u8);
}
impl Wlg0Decoder for Vec<Wlg0Gaussian> {
fn init_num_splats(&mut self, num_splats: usize) {
self.resize(num_splats, Wlg0Gaussian::default());
}
fn write_center(&mut self, index: usize, dim: usize, value: f32) {
self[index].center[dim] = value;
}
fn write_scale(&mut self, index: usize, dim: usize, value: f32) {
self[index].ln_scale[dim] = value.ln();
}
fn write_quaternion_f32(&mut self, index: usize, dim: usize, value: f32) {
self[index].quaternion[dim] = value;
}
fn write_quaternion_i8(&mut self, _index: usize, _dim: usize, _value: i8) {
// No-op sing we use the f32 variant
}
fn write_quaternion_full_f32(&mut self, _index: usize, _value: [f32; 4]) {
// No-op
}
fn write_rgba(&mut self, index: usize, dim: usize, value: u8) {
self[index].color[dim] = value as f32 / 255.0;
}
}
impl Wlg0Decoder for PackedSplats {
fn init_num_splats(&mut self, num_splats: usize) {
self.0.resize(num_splats * 4, 0);
}
fn write_center(&mut self, index: usize, dim: usize, value: f32) {
match dim {
0 => { self.0[index * 4 + 1] |= f16::from_f32(value).to_bits() as u32; },
1 => { self.0[index * 4 + 1] |= (f16::from_f32(value).to_bits() as u32) << 16; },
2 => { self.0[index * 4 + 2] |= f16::from_f32(value).to_bits() as u32; },
_ => unreachable!(),
}
}
fn write_scale(&mut self, index: usize, dim: usize, value: f32) {
let scale8 = encode_scale(value);
match dim {
0 => { self.0[index * 4 + 3] |= scale8 as u32; },
1 => { self.0[index * 4 + 3] |= (scale8 as u32) << 8; },
2 => { self.0[index * 4 + 3] |= (scale8 as u32) << 16; },
_ => unreachable!(),
}
}
fn write_quaternion_f32(&mut self, _index: usize, _dim: usize, _value: f32) {
// No-op since we use the i8 variant
}
fn write_quaternion_i8(&mut self, _index: usize, _dim: usize, _value: i8) {
// Old Qxyz encoding:
// match dim {
// 0 => { self.0[index * 4 + 2] |= ((value as u32) & 0xff) << 16; },
// 1 => { self.0[index * 4 + 2] |= ((value as u32) & 0xff) << 24; },
// 2 => { self.0[index * 4 + 3] |= ((value as u32) & 0xff) << 24; },
// 3 => {
// // Quaternion w is inferred from xyz in PackedSplats
// },
// _ => unreachable!(),
// }
}
fn write_quaternion_full_f32(&mut self, index: usize, value: [f32; 4]) {
// Encode as OctXy88R8
let q = if value[3] < 0.0 { value.map(|v| -v) } else { value };
let theta = 2.0 * q[3].acos();
let xyz_norm = (q[0] * q[0] + q[1] * q[1] + q[2] * q[2]).sqrt();
let axis = if xyz_norm < 1e-6 {
[1.0, 0.0, 0.0]
} else {
[q[0] / xyz_norm, q[1] / xyz_norm, q[2] / xyz_norm]
};
let sum = axis[0].abs() + axis[1].abs() + axis[2].abs();
let p = [axis[0] / sum, axis[1] / sum];
let p = if axis[2] >= 0.0 { p } else {
[
(1.0 - p[1].abs()) * p[0].signum(),
(1.0 - p[0].abs()) * p[1].signum(),
]
};
let uv = [
((p[0] * 0.5 + 0.5) * 255.0).round() as u8,
((p[1] * 0.5 + 0.5) * 255.0).round() as u8,
];
let angle = (theta * (255.0 / 3.14159265359)).round().clamp(0.0, 255.0) as u8;
self.0[index * 4 + 2] |= (uv[0] as u32) << 16;
self.0[index * 4 + 2] |= (uv[1] as u32) << 24;
self.0[index * 4 + 3] |= (angle as u32) << 24;
}
fn write_rgba(&mut self, index: usize, dim: usize, value: u8) {
match dim {
0 => { self.0[index * 4] |= value as u32; },
1 => { self.0[index * 4] |= (value as u32) << 8; },
2 => { self.0[index * 4] |= (value as u32) << 16; },
3 => { self.0[index * 4] |= (value as u32) << 24; },
_ => unreachable!(),
}
}
}
pub(super) fn decode_internal<D: Wlg0Decoder>(
data: &mut [u8],
settings: &Wlg0Settings,
decoder: &mut D,
) -> anyhow::Result<()> {
if settings.enable_obfuscation {
deobfuscate(data);
}
let data = if settings.enable_compression {
decompress_data(data)?
} else {
data.to_vec()
};
let mut reader = std::io::Cursor::new(&data);
let header = Wlg0Header::read(&mut reader)?;
let num_splats = header.num_splats as usize;
decoder.init_num_splats(num_splats);
if !settings.enable_split_dims {
return Err(anyhow!("Unsupported WLG !settings.enable_split_dims"));
}
let mut data_u8: Vec<u8> = vec![0; num_splats];
if settings.enable_split_center_bytes {
let mut data_u32: Vec<u32> = vec![0; num_splats];
for d in 0..3 {
data_u32.iter_mut().for_each(|v| *v = 0);
if settings.enable_24bit_center {
reader.read_exact(&mut data_u8)?;
if settings.enable_first_differences {
first_cumulative(&mut data_u8);
}
for (i, &byte) in data_u8.iter().enumerate() {
data_u32[i] = (byte as u32) << 16;
}
}
reader.read_exact(&mut data_u8)?;
if settings.enable_first_differences {
first_cumulative(&mut data_u8);
}
for (i, &byte) in data_u8.iter().enumerate() {
data_u32[i] |= (byte as u32) << 8;
}
reader.read_exact(&mut data_u8)?;
if settings.enable_first_differences {
first_cumulative(&mut data_u8);
}
for (i, &byte) in data_u8.iter().enumerate() {
data_u32[i] |= byte as u32;
}
let resolution = if settings.enable_24bit_center {
16777215.0
} else {
65535.0
};
let scale = header.center_scale / resolution;
let offset = header.center_offset[d];
for (i, &value) in data_u32.iter().enumerate() {
decoder.write_center(i, d, value as f32 * scale + offset);
}
}
} else {
return Err(anyhow!(
"Unsupported WLG !settings.enable_split_center_bytes"
));
}
for d in 0..3 {
reader.read_exact(&mut data_u8)?;
let scale = header.ln_scale_max - header.ln_scale_min;
let offset = header.ln_scale_min;
for (i, &byte) in data_u8.iter().enumerate() {
let float = ((byte as f32 / 255.0) * scale + offset).exp();
decoder.write_scale(i, d, float);
}
}
{
let mut quaternions: Vec<i8> = vec![0; num_splats * 3];
for d in 0..3 {
reader.read_exact(&mut data_u8)?;
for (i, &byte) in data_u8.iter().enumerate() {
quaternions[i * 3 + d] = byte as i8;
}
}
for i in 0..num_splats {
let quat: [f32; 3] = [
quaternions[i * 3 + 0] as f32 / 127.0,
quaternions[i * 3 + 1] as f32 / 127.0,
quaternions[i * 3 + 2] as f32 / 127.0,
];
let w = (1.0 - quat.iter().map(|&v| v.powi(2)).sum::<f32>())
.max(0.0)
.sqrt();
let w8 = (w * 127.0).round() as i8;
for d in 0..3 {
decoder.write_quaternion_f32(i, d, quat[d]);
decoder.write_quaternion_i8(i, d, quaternions[i * 3 + d]);
}
decoder.write_quaternion_f32(i, 3, w);
decoder.write_quaternion_i8(i, 3, w8);
decoder.write_quaternion_full_f32(i, [quat[0], quat[1], quat[2], w]);
}
}
reader.read_exact(&mut data_u8)?;
if settings.enable_first_differences {
first_cumulative(&mut data_u8);
}
for (i, &byte) in data_u8.iter().enumerate() {
decoder.write_rgba(i, 3, byte);
}
for d in 0..3 {
reader.read_exact(&mut data_u8)?;
if settings.enable_first_differences {
first_cumulative(&mut data_u8);
}
for (i, &byte) in data_u8.iter().enumerate() {
decoder.write_rgba(i, d, byte);
}
}
let position = reader.position() as usize;
if position != data.len() {
return Err(anyhow!("Invalid WLG data size"));
}
Ok(())
}
rust/wlg/src/wlg0/encode.rs
-283
View File
@@ -1,283 +0,0 @@
use std::io::Write;
use super::{compress_data, obfuscate, Wlg0Gaussian, Wlg0Settings, Wlg0Header};
use crate::ordering::{
hilbert_coord_to_index, hilbert_coord_to_index_24, morton_coord_to_index,
morton_coord_to_index_24,
};
// WLG splat representation with fixed-point encoding.
#[derive(Debug)]
struct WlgSplat {
order: u128,
center: [u32; 3],
scale: [u8; 3],
quaternion: [i8; 3],
opacity: u8,
color: [u8; 3],
}
impl WlgSplat {
fn new(
settings: &Wlg0Settings,
index: usize,
center: [f32; 3],
scale: [f32; 3],
mut quaternion: [f32; 4],
opacity: f32,
color: [f32; 3],
) -> Self {
use std::array::from_fn;
let center = if settings.enable_24bit_center {
from_fn(|d| (center[d] * 16777215.0).clamp(0.0, 16777215.0).round() as u32)
} else {
from_fn(|d| (center[d] * 65535.0).clamp(0.0, 65535.0).round() as u32)
};
if quaternion[3] < 0.0 {
quaternion = from_fn(|d| -quaternion[d]);
}
Self {
order: if settings.enable_hilbert_reordering {
// Convert X/Y/Z into a single Hilbert curve index
if settings.enable_24bit_center {
hilbert_coord_to_index_24(center) as u128
} else {
hilbert_coord_to_index(from_fn(|d| center[d] as u16)) as u128
}
} else if settings.enable_morton_reordering {
// Interleave X/Y/Z bits into a single Morton index
if settings.enable_24bit_center {
morton_coord_to_index_24(center) as u128
} else {
morton_coord_to_index(from_fn(|d| center[d] as u16)) as u128
}
} else {
index as u128
},
center,
scale: from_fn(|d| (scale[d] * 255.0).clamp(0.0, 255.0).round() as u8),
quaternion: from_fn(|d| (quaternion[d] * 127.0).clamp(-127.0, 127.0).round() as i8),
opacity: (opacity * 255.0).clamp(0.0, 255.0).round() as u8,
color: from_fn(|d| (color[d] * 255.0).clamp(0.0, 255.0).round() as u8),
}
}
}
// Encoding functions
fn first_differences(settings: &Wlg0Settings, mut data: Vec<u8>) -> Vec<u8> {
if settings.enable_first_differences {
super::first_differences(&mut data);
}
data
}
pub(super) fn encode_internal(
settings: &Wlg0Settings,
gaussians: &[Wlg0Gaussian],
) -> anyhow::Result<Vec<u8>> {
use std::array::from_fn;
let num_splats = gaussians.len();
// Compute min/max bounds for Gsplat centers
let center_min_max = ([f32::INFINITY; 3], [-f32::INFINITY; 3]);
let (center_min, center_max) = gaussians.iter().fold(center_min_max, |(min, max), g| {
let min = from_fn(|i| min[i].min(g.center[i]));
let max = from_fn(|i| max[i].max(g.center[i]));
(min, max)
});
// println!("Center min: {:?}, max: {:?}", center_min, center_max);
// Compute scale factor and offset for center values
let center_ranges: [f32; 3] = from_fn(|i| center_max[i] - center_min[i]);
let center_scale = center_ranges.into_iter().reduce(|a, b| a.max(b)).unwrap();
let center_offset = center_min;
// Compute min/max bounds for Gsplat ln_scales
let ln_scale_min_max = (f32::INFINITY, -f32::INFINITY);
let (ln_scale_min, ln_scale_max) =
gaussians
.iter()
.fold(ln_scale_min_max, |(mut min, mut max), g| {
g.ln_scale.iter().for_each(|&ln_scale| {
let ln_scale = ln_scale.clamp(-10.0, 10.0);
min = min.min(ln_scale);
max = max.max(ln_scale);
});
(min, max)
});
// println!("ln_scale min: {}, max: {}", ln_scale_min, ln_scale_max);
let mut splats: Vec<_> = gaussians
.iter()
.enumerate()
.map(|(index, g)| {
WlgSplat::new(
settings,
index,
from_fn(|d| (g.center[d] - center_offset[d]) / center_scale),
from_fn(|d| (g.ln_scale[d] - ln_scale_min) / (ln_scale_max - ln_scale_min)),
g.quaternion,
g.opacity,
from_fn(|d| g.color[d]),
)
})
.collect();
// Reorder splats using the provided ordering key, either Morton or original index.
splats.sort_by_key(|splat| splat.order);
// Create buffer to write header and splats to
let mut buffer: Vec<u8> = Vec::new();
let header = Wlg0Header {
center_scale,
center_offset,
ln_scale_min,
ln_scale_max,
num_splats: num_splats as u32,
max_sh_order: 0,
num_sh_splats: [num_splats as u32, 0, 0, 0],
};
header.write(&mut buffer)?;
if settings.enable_split_dims {
// Write dimensions separately, i.e. a column-oriented format.
for d in 0..3 {
if settings.enable_split_center_bytes {
// Split center values into separate bytes (MSB order)
if settings.enable_24bit_center {
let data: Vec<u8> = splats
.iter()
.map(|splat| (splat.center[d] >> 16) as u8)
.collect();
buffer.write_all(&first_differences(settings, data))?;
}
let data: Vec<u8> = splats
.iter()
.map(|splat| (splat.center[d] >> 8) as u8)
.collect();
buffer.write_all(&first_differences(settings, data))?;
let data: Vec<u8> = splats.iter().map(|splat| splat.center[d] as u8).collect();
buffer.write_all(&first_differences(settings, data))?;
} else {
if settings.enable_24bit_center {
// Write 24-bit center values into u8 array in LSB order
let data: Vec<u8> = splats
.iter()
.flat_map(|splat| {
let center = splat.center[d];
[center as u8, (center >> 8) as u8, (center >> 16) as u8]
})
.collect();
buffer.write_all(&data)?;
} else {
let data_u16: Vec<u16> =
splats.iter().map(|splat| splat.center[d] as u16).collect();
let data: Vec<u8> = data_u16
.iter()
.flat_map(|&v| v.to_le_bytes().to_vec())
.collect();
buffer.write_all(&data)?;
}
}
}
for d in 0..3 {
let data: Vec<u8> = splats.iter().map(|splat| splat.scale[d]).collect();
buffer.write_all(&data)?;
}
for d in 0..3 {
let data: Vec<u8> = splats
.iter()
.map(|splat| splat.quaternion[d] as u8)
.collect();
buffer.write_all(&data)?;
}
{
let data: Vec<u8> = splats.iter().map(|splat| splat.opacity).collect();
buffer.write_all(&first_differences(settings, data))?;
}
for d in 0..3 {
let data: Vec<u8> = splats.iter().map(|splat| splat.color[d]).collect();
buffer.write_all(&first_differences(settings, data))?;
}
} else {
// !enable_split_dims, write all dimensions interleaved (array of structs).
if settings.enable_split_center_bytes {
if settings.enable_24bit_center {
let data: Vec<u8> = splats
.iter()
.flat_map(|splat| splat.center.iter().map(|&v| (v >> 16) as u8))
.collect();
buffer.write_all(&first_differences(settings, data))?;
}
let data: Vec<u8> = splats
.iter()
.flat_map(|splat| splat.center.iter().map(|&v| (v >> 8) as u8))
.collect();
buffer.write_all(&first_differences(settings, data))?;
let data: Vec<u8> = splats
.iter()
.flat_map(|splat| splat.center.iter().map(|&v| v as u8))
.collect();
buffer.write_all(&first_differences(settings, data))?;
} else {
if settings.enable_24bit_center {
let data: Vec<u8> = splats
.iter()
.flat_map(|splat| {
splat
.center
.iter()
.flat_map(|&d| [d as u8, (d >> 8) as u8, (d >> 16) as u8])
})
.collect();
buffer.write_all(&data)?;
} else {
let data_u16: Vec<u16> = splats
.iter()
.flat_map(|splat| splat.center.iter().map(|&v| v as u16))
.collect();
let data: Vec<u8> = data_u16
.iter()
.flat_map(|&v| v.to_le_bytes().to_vec())
.collect();
buffer.write_all(&data)?;
}
}
let data: Vec<u8> = splats
.iter()
.flat_map(|splat| splat.scale.iter().copied())
.collect();
buffer.write_all(&data)?;
let data: Vec<u8> = splats
.iter()
.flat_map(|splat| splat.quaternion.iter().map(|&v| v as u8))
.collect();
buffer.write_all(&data)?;
let data: Vec<u8> = splats.iter().map(|splat| splat.opacity).collect();
buffer.write_all(&first_differences(settings, data))?;
let data: Vec<u8> = splats
.iter()
.flat_map(|splat| splat.color.iter().copied())
.collect();
buffer.write_all(&first_differences(settings, data))?;
}
let mut payload = if settings.enable_compression {
compress_data(&buffer, settings.zstd_compression_level)?
} else {
buffer
};
if settings.enable_obfuscation {
obfuscate(&mut payload);
}
Ok(payload)
}
rust/wlg/src/wlg0/mod.rs
-274
View File
@@ -1,274 +0,0 @@
use serde::{Deserialize, Serialize};
use std::io::{Read, Write};
use anyhow::anyhow;
mod decode;
mod encode;
#[derive(Debug, Default, Clone)]
pub struct Wlg0Gaussian {
pub center: [f32; 3],
pub ln_scale: [f32; 3],
pub quaternion: [f32; 4],
pub opacity: f32,
pub color: [f32; 3],
}
pub const LN_SCALE_MIN: f32 = -9.0;
pub const LN_SCALE_MAX: f32 = 9.0;
pub const LN_RESCALE: f32 = (LN_SCALE_MAX - LN_SCALE_MIN) / 254.0; // 1..=255
pub fn encode_scale(scale: f32) -> u8 {
if scale == 0.0 {
0
} else {
// Allow scales below LN_SCALE_MIN to be encoded as 0, which signifies a 2DGS
((scale.ln() - LN_SCALE_MIN) / LN_RESCALE + 1.0).clamp(0.0, 255.0).round() as u8
}
}
pub fn decode_scale(scale: u8) -> f32 {
if scale == 0 {
0.0
} else {
(LN_SCALE_MIN + (scale - 1) as f32 * LN_RESCALE).exp()
}
}
pub struct PackedSplats(pub Vec<u32>);
#[derive(Debug, Serialize, Deserialize)]
struct Wlg0Signature {
magic: [u8; 4], // "WLG0"
version: u32, // 1 or 2
}
impl Wlg0Signature {
const MAGIC: [u8; 4] = *b"WLG0";
fn new(version: u32) -> Self {
Self {
magic: Self::MAGIC,
version,
}
}
fn check_version(&self) -> Option<u32> {
if self.magic == Self::MAGIC {
Some(self.version)
} else {
None
}
}
fn write<W: Write>(&self, writer: &mut W) -> std::io::Result<()> {
writer.write_all(&bincode::serialize(self).unwrap())
}
fn read<R: Read>(reader: &mut R) -> std::io::Result<Self> {
let mut buf = [0; std::mem::size_of::<Self>()];
reader.read_exact(&mut buf)?;
Ok(bincode::deserialize(&buf).unwrap())
}
}
#[derive(Debug, Clone)]
pub struct Wlg0Settings {
pub version: u32,
pub enable_compression: bool,
pub enable_first_differences: bool,
pub enable_morton_reordering: bool,
pub enable_hilbert_reordering: bool,
pub enable_split_center_bytes: bool,
pub enable_obfuscation: bool,
pub enable_split_dims: bool,
pub enable_24bit_center: bool,
pub zstd_compression_level: i32, // 0 selects default level (3), max is 22
}
const WLG012_SETTINGS: [Wlg0Settings; 3] = [
// 0: WLGv0: Not a valid version, for testing purposes only
Wlg0Settings {
version: 0,
enable_compression: true,
enable_first_differences: true,
enable_morton_reordering: false,
enable_hilbert_reordering: true,
enable_split_center_bytes: true,
enable_obfuscation: true,
enable_split_dims: true,
enable_24bit_center: true,
zstd_compression_level: 0,
},
// 1: WLGv1: Default settings for WLG v1 files
Wlg0Settings {
version: 1,
enable_compression: true,
enable_first_differences: true,
enable_morton_reordering: true,
enable_hilbert_reordering: false,
enable_split_center_bytes: true,
enable_obfuscation: true,
enable_split_dims: true,
enable_24bit_center: false,
zstd_compression_level: 0,
},
// 2: WLGv2: Default settings for WLG v2 files
Wlg0Settings {
version: 2,
enable_compression: true,
enable_first_differences: true,
enable_morton_reordering: true,
enable_hilbert_reordering: false,
enable_split_center_bytes: true,
enable_obfuscation: true,
enable_split_dims: true,
enable_24bit_center: true,
zstd_compression_level: 0,
},
];
#[derive(Debug, Serialize, Deserialize)]
pub struct Wlg0Header {
// center[d] = center_int[d] * center_scale + center_offset
pub center_scale: f32,
pub center_offset: [f32; 3],
// scale[d] = exp(scale_int[d] / 255 * (scale_max - scale_min) + scale_min)
pub ln_scale_min: f32,
pub ln_scale_max: f32,
// Total # splats, regardless of SH order
pub num_splats: u32,
// 0 for "no spherical harmonics"
pub max_sh_order: u32,
// Supports maximum of SH=3
pub num_sh_splats: [u32; 4],
}
impl Wlg0Header {
pub fn write<W: Write>(&self, writer: &mut W) -> std::io::Result<()> {
writer.write_all(&bincode::serialize(self).unwrap())
}
pub fn read<R: Read>(reader: &mut R) -> std::io::Result<Self> {
let mut buf = [0; std::mem::size_of::<Self>()];
reader.read_exact(&mut buf)?;
Ok(bincode::deserialize(&buf).unwrap())
}
}
fn obfuscate(data: &mut [u8]) {
let mut prev: u8 = 0;
let mut state: u32 = 0x1AB51AB5;
for byte in data {
state = state.wrapping_mul(1664525).wrapping_add(1013904223);
let rand_byte = (state >> 8) as u8;
let obf_byte = (*byte ^ prev).rotate_left(3) ^ rand_byte;
prev = obf_byte;
*byte = obf_byte;
}
}
fn deobfuscate(data: &mut [u8]) {
let mut prev: u8 = 0;
let mut state: u32 = 0x1AB51AB5;
for obf_byte in data {
state = state.wrapping_mul(1664525).wrapping_add(1013904223);
let rand_byte = (state >> 8) as u8;
let tmp = (*obf_byte ^ rand_byte).rotate_right(3);
let byte = tmp ^ prev;
prev = *obf_byte;
*obf_byte = byte;
}
}
fn compress_data(data: &[u8], _zstd_compression_level: i32) -> anyhow::Result<Vec<u8>> {
use ruzstd::encoding::{compress_to_vec, CompressionLevel};
let compressed_data = compress_to_vec(data, CompressionLevel::Fastest);
Ok(compressed_data)
}
fn decompress_data(data: &[u8]) -> anyhow::Result<Vec<u8>> {
let mut decoder = ruzstd::decoding::StreamingDecoder::new(data)?;
let mut decompressed_data = Vec::new();
decoder.read_to_end(&mut decompressed_data)?;
Ok(decompressed_data)
}
fn first_differences(data: &mut [u8]) {
// Compute first differences for array to aid compression
for i in (1..data.len()).rev() {
data[i] = data[i].wrapping_sub(data[i - 1]);
}
}
fn first_cumulative(data: &mut [u8]) {
// Invert first differences
for i in 1..data.len() {
data[i] = data[i].wrapping_add(data[i - 1]);
}
}
pub fn decode12(bytes: &mut [u8]) -> anyhow::Result<(Wlg0Settings, Vec<Wlg0Gaussian>)> {
let (offset, version) = {
let mut reader = std::io::Cursor::new(&bytes);
let signature = Wlg0Signature::read(&mut reader)?;
let Some(version) = signature.check_version() else {
return Err(anyhow!("Invalid WLG signature"));
};
if (version != 1) && (version != 2) {
return Err(anyhow!("Unsupported WLG version"));
}
let offset = reader.position() as usize;
(offset, version)
};
let settings = WLG012_SETTINGS[version as usize].clone();
let mut gaussians = Vec::new();
decode::decode_internal(&mut bytes[offset..], &settings, &mut gaussians)?;
Ok((settings, gaussians))
}
pub fn decode12_packed(bytes: &mut [u8]) -> anyhow::Result<(Wlg0Settings, PackedSplats)> {
let (offset, version) = {
let mut reader = std::io::Cursor::new(&bytes);
let signature = Wlg0Signature::read(&mut reader)?;
let Some(version) = signature.check_version() else {
return Err(anyhow!("Invalid WLG signature"));
};
if (version != 1) && (version != 2) {
return Err(anyhow!("Unsupported WLG version"));
}
let offset = reader.position() as usize;
(offset, version)
};
let settings = WLG012_SETTINGS[version as usize].clone();
let mut packed_splats = PackedSplats(Vec::new());
decode::decode_internal(&mut bytes[offset..], &settings, &mut packed_splats)?;
Ok((settings, packed_splats))
}
pub enum Wlg0EncodeSettings {
Wlg1,
Wlg2,
Custom(Wlg0Settings),
}
pub fn encode12(
settings: &Wlg0EncodeSettings,
gaussians: &[Wlg0Gaussian],
) -> anyhow::Result<(Wlg0Settings, Vec<u8>)> {
let settings = match settings {
Wlg0EncodeSettings::Wlg1 => &WLG012_SETTINGS[1],
Wlg0EncodeSettings::Wlg2 => &WLG012_SETTINGS[2],
Wlg0EncodeSettings::Custom(custom) => custom,
}
.clone();
let payload = encode::encode_internal(&settings, gaussians)?;
let mut buffer = Vec::new();
Wlg0Signature::new(settings.version).write(&mut buffer)?;
buffer.write_all(&payload)?;
Ok((settings, buffer))
}