diff --git a/rust/spark-lib/Cargo.toml b/rust/spark-lib/Cargo.toml index be19442b..93b970df 100644 --- a/rust/spark-lib/Cargo.toml +++ b/rust/spark-lib/Cargo.toml @@ -7,6 +7,9 @@ license.workspace = true authors.workspace = true repository.workspace = true +[features] +sogs_web_decode = [] + [dependencies] ahash.workspace = true anyhow.workspace = true diff --git a/rust/spark-lib/src/decoder.rs b/rust/spark-lib/src/decoder.rs index 74c4af30..d9fce7dd 100644 --- a/rust/spark-lib/src/decoder.rs +++ b/rust/spark-lib/src/decoder.rs @@ -8,7 +8,7 @@ use crate::{ ksplat::KsplatDecoder, ply::{PLY_MAGIC, PlyDecoder}, rad::{RAD_CHUNK_MAGIC, RAD_MAGIC, RadDecoder}, - sogs::SogsDecoder, + sogs::{SogsDecoder, PK_MAGIC, CUSTOM_SOGS_MAGIC}, spz::{SPZ_MAGIC, SpzDecoder} }; @@ -355,6 +355,7 @@ impl SplatFileType { "spz" => Ok(Self::SPZ), "splat" => Ok(Self::ANTISPLAT), "ksplat" => Ok(Self::KSPLAT), + "pcsogs" => Ok(Self::SOGS), "pcsogszip" => Ok(Self::SOGS), "rad" => Ok(Self::RAD), _ => Err(anyhow::anyhow!("Invalid file type: {}", enum_str)), @@ -488,13 +489,15 @@ impl ChunkReceiver for MultiDecoder { } } } - } else if magic == 0x04034b50 { + } else if magic == PK_MAGIC { detection_complete = true; if let Some(pathname) = &self.pathname { if let Some(SplatFileType::SOGS) = SplatFileType::from_pathname(pathname) { return self.init_file_type(SplatFileType::SOGS); } } + } else if magic == CUSTOM_SOGS_MAGIC { + return self.init_file_type(SplatFileType::SOGS); } else if magic == RAD_MAGIC || magic == RAD_CHUNK_MAGIC { return self.init_file_type(SplatFileType::RAD); } else { diff --git a/rust/spark-lib/src/sogs.rs b/rust/spark-lib/src/sogs.rs index 59a39785..7707b343 100644 --- a/rust/spark-lib/src/sogs.rs +++ b/rust/spark-lib/src/sogs.rs @@ -1,14 +1,19 @@ -use std::{collections::HashMap, io::Cursor}; +use std::collections::HashMap; +#[cfg(not(feature = "sogs_web_decode"))] +use std::io::Cursor; use anyhow::{anyhow, Context}; +#[cfg(not(feature = "sogs_web_decode"))] use image::{DynamicImage, GenericImageView, ImageReader}; use serde_json; use serde::Deserialize; +#[cfg(not(feature = "sogs_web_decode"))] use zip::ZipArchive; use crate::decoder::{ChunkReceiver, SplatInit, SplatProps, SplatReceiver}; -const PK_MAGIC: u32 = 0x04034b50; +pub const PK_MAGIC: u32 = 0x04034b50; +pub const CUSTOM_SOGS_MAGIC: u32 = 0x53474F53; const SH_C0: f32 = 0.28209479177387814; const MAX_SPLAT_CHUNK: usize = 65536; @@ -133,7 +138,7 @@ impl ChunkReceiver for SogsDecoder { 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]]); - if magic != PK_MAGIC { + if magic != PK_MAGIC && magic != CUSTOM_SOGS_MAGIC { return Err(anyhow!("Not a ZIP/SOGS file")); } decode_sogs(&self.buffer, &mut self.splats, None)?; @@ -141,6 +146,53 @@ impl ChunkReceiver for SogsDecoder { } } +#[cfg(feature = "sogs_web_decode")] +fn decode_sogs(bytes: &[u8], splats: &mut T, _pathname: Option<&str>) -> anyhow::Result<()> { + let mut file_cache: HashMap> = HashMap::new(); + let mut offset: usize = 4; // Skip magic number + + while offset < bytes.len() { + let name_len = u16::from_le_bytes(bytes[offset..offset + 2].try_into().unwrap()) as usize; + offset += 2; + let name_bytes = &bytes[offset..offset + name_len]; + offset += name_len; + let name = std::str::from_utf8(name_bytes).unwrap(); + + let data_size = u32::from_le_bytes(bytes[offset..offset + 4].try_into().unwrap()) as usize; + offset += 4; + let data = &bytes[offset..offset + data_size]; + offset += data_size; + + file_cache.insert(name.to_string(), data.to_vec()); + } + + let mut get_image_data = |name: &str| -> anyhow::Result { + let mut file = file_cache.get(name).cloned().ok_or_else(|| anyhow!("Missing file {name} in cache"))?; + + let width = u32::from_le_bytes( + file[0..4].try_into().unwrap() + ) as usize; + + let height = u32::from_le_bytes( + file[4..8].try_into().unwrap() + ) as usize; + + let rgba = file.split_off(8); + + Ok(ImageData { width, height, rgba }) + }; + + let meta_bytes = file_cache.get("meta.json").cloned().ok_or_else(|| anyhow!("Missing meta.json in cache"))?; + let meta: PcSogsRoot = serde_json::from_slice(&meta_bytes) + .context("Failed to parse meta.json for SOGS")?; + + match meta { + PcSogsRoot::V2(v2) => decode_v2(v2, splats, &mut get_image_data), + PcSogsRoot::V1(v1) => decode_v1(v1, splats, &mut get_image_data), + } +} + +#[cfg(not(feature = "sogs_web_decode"))] fn decode_sogs(bytes: &[u8], splats: &mut T, _pathname: Option<&str>) -> anyhow::Result<()> { let cursor = Cursor::new(bytes); let mut zip = ZipArchive::new(cursor)?; @@ -169,20 +221,21 @@ fn decode_sogs(bytes: &[u8], splats: &mut T, _pathname: Option let mut file_cache: HashMap> = HashMap::new(); preload_all(&meta, &prefix, &mut zip, &mut file_cache)?; - let mut get_file = |name: &str| -> anyhow::Result> { - file_cache.get(name).cloned().ok_or_else(|| anyhow!("Missing file {name} in cache")) + let mut get_image_data = |name: &str| -> anyhow::Result { + let file = file_cache.get(name).cloned().ok_or_else(|| anyhow!("Missing file {name} in cache")); + decode_image(&file?) }; match meta { - PcSogsRoot::V2(v2) => decode_v2(v2, splats, &mut get_file), - PcSogsRoot::V1(v1) => decode_v1(v1, splats, &mut get_file), + PcSogsRoot::V2(v2) => decode_v2(v2, splats, &mut get_image_data), + PcSogsRoot::V1(v1) => decode_v1(v1, splats, &mut get_image_data), } } fn decode_v2( meta: PcSogsV2, splats: &mut T, - get_file: &mut dyn FnMut(&str) -> anyhow::Result>, + get_image_data: &mut dyn FnMut(&str) -> anyhow::Result, ) -> anyhow::Result<()> { let _ = meta.version; let num_splats = meta.count; @@ -191,16 +244,11 @@ fn decode_v2( }).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 means0 = get_image_data(&meta.means.files[0])?; + let means1 = get_image_data(&meta.means.files[1])?; + let scales_img = get_image_data(&meta.scales.files[0])?; + let quats_img = get_image_data(&meta.quats.files[0])?; + let sh0_img = get_image_data(&meta.sh0.files[0])?; let mut center = vec![0.0f32; num_splats * 3]; let mut scale = vec![0.0f32; num_splats * 3]; @@ -222,7 +270,7 @@ fn decode_v2( if let Some(shn) = meta.shn { decode_shn_v2( shn, - get_file, + get_image_data, num_splats, &mut sh1, &mut sh2, @@ -248,7 +296,7 @@ fn decode_v2( fn decode_v1( meta: PcSogsV1, splats: &mut T, - get_file: &mut dyn FnMut(&str) -> anyhow::Result>, + get_image_data: &mut dyn FnMut(&str) -> anyhow::Result, ) -> anyhow::Result<()> { let num_splats = meta.means.shape[0]; if meta.quats.encoding.as_deref() != Some("quaternion_packed") { @@ -270,16 +318,11 @@ 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 means0 = get_image_data(&meta.means.files[0])?; + let means1 = get_image_data(&meta.means.files[1])?; + let scales_img = get_image_data(&meta.scales.files[0])?; + let quats_img = get_image_data(&meta.quats.files[0])?; + let sh0_img = get_image_data(&meta.sh0.files[0])?; let mut center = vec![0.0f32; num_splats * 3]; let mut scale = vec![0.0f32; num_splats * 3]; @@ -298,7 +341,7 @@ fn decode_v1( if let Some(shn) = meta.shn { decode_shn_v1( shn, - get_file, + get_image_data, num_splats, max_sh_degree, &mut sh1, @@ -447,14 +490,14 @@ fn decode_sh0_v1(mins: &[f32; 4], maxs: &[f32; 4], img: &ImageData, out_rgb: &mu fn decode_shn_v2( shn: ShNV2, - get_file: &mut dyn FnMut(&str) -> anyhow::Result>, + get_image_data: &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 centroids = get_image_data(&shn.files[0])?; + let labels = get_image_data(&shn.files[1])?; let lookup = shn.codebook; let use_sh1 = shn.bands >= 1; let use_sh2 = shn.bands >= 2; @@ -490,15 +533,15 @@ fn decode_shn_v2( fn decode_shn_v1( shn: ShNV1, - get_file: &mut dyn FnMut(&str) -> anyhow::Result>, + get_image_data: &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 centroids = get_image_data(&shn.files[0])?; + let labels = get_image_data(&shn.files[1])?; let lookup: Vec = (0..256) .map(|i| shn.mins + (shn.maxs - shn.mins) * (i as f32 / 255.0)) .collect(); @@ -569,6 +612,7 @@ fn emit_to_receiver( splats.finish() } +#[cfg(not(feature = "sogs_web_decode"))] fn preload_all( meta: &PcSogsRoot, prefix: &str, @@ -598,6 +642,7 @@ fn preload_all( Ok(()) } +#[cfg(not(feature = "sogs_web_decode"))] fn preload_file( zip: &mut ZipArchive>, prefix: &str, @@ -629,11 +674,7 @@ struct ImageData { height: usize, } -fn decode_rgba(bytes: &[u8]) -> anyhow::Result { - let img = decode_image(bytes)?; - Ok(img) -} - +#[cfg(not(feature = "sogs_web_decode"))] fn decode_image(bytes: &[u8]) -> anyhow::Result { let img = ImageReader::new(Cursor::new(bytes)) .with_guessed_format()? diff --git a/rust/spark-rs/Cargo.toml b/rust/spark-rs/Cargo.toml index 7c4f3e13..09fffdb8 100644 --- a/rust/spark-rs/Cargo.toml +++ b/rust/spark-rs/Cargo.toml @@ -23,7 +23,7 @@ ordered-float.workspace = true smallvec.workspace = true wasm-bindgen.workspace = true web-sys = { workspace = true, features = ["Window", "Performance"] } -spark-lib = { path = "../spark-lib" } +spark-lib = { path = "../spark-lib", features = ["sogs_web_decode"] } serde-wasm-bindgen.workspace = true serde_json.workspace = true itertools.workspace = true diff --git a/src/sogs.ts b/src/sogs.ts new file mode 100644 index 00000000..726471fd --- /dev/null +++ b/src/sogs.ts @@ -0,0 +1,181 @@ +import { unzipSync } from "fflate"; + +// Custom magic number for (unzipped) and decoded SOG files +const HEADER = new Uint8Array([0x53, 0x4f, 0x47, 0x53]); + +const temp = new Uint8Array(4); +const dataView = new DataView(temp.buffer); +const textDecoder = new TextDecoder(); +const textEncoder = new TextEncoder(); + +type ImageData = { width: number; height: number; rgba: Uint8Array }; + +export function unzipAndDecodeImages(zipSize: number) { + const data = new Uint8Array(zipSize); + let processed = 0; + + return new TransformStream({ + start(controller) { + controller.enqueue(HEADER); + }, + async transform(chunk, controller) { + const chunkData = await chunk; + data.set(chunkData, processed); + processed += chunkData.length; + }, + async flush(controller) { + const unzipped = unzipSync(data); + + const promises: Array> = []; + for (const fileName in unzipped) { + if (fileName.endsWith(".webp")) { + promises.push( + decodeImage(unzipped[fileName].buffer as ArrayBuffer).then( + (imageData) => enqueueImage(controller, fileName, imageData), + ), + ); + } else { + enqueueFile(controller, fileName, unzipped[fileName]); + } + } + + await Promise.allSettled(promises); + }, + }); +} + +export function fetchAndDecodeImages(url: string) { + // Strip the file name + const baseUrl = url.substring(0, url.lastIndexOf("/")); + + return new ReadableStream({ + async start(controller) { + // Fetch the meta.json file + const arrayBuffer = await (await fetch(url)).arrayBuffer(); + const json = JSON.parse(textDecoder.decode(arrayBuffer)); + const refFiles = [ + ...json.means.files, + ...json.scales.files, + ...json.quats.files, + ...json.sh0.files, + ...(json.shN?.files ?? []), + ]; + + // Start outputting + controller.enqueue(HEADER); + enqueueFile(controller, "meta.json", new Uint8Array(arrayBuffer)); + + const promises = refFiles.map(async (imageFile) => { + const response = await fetch(`${baseUrl}/${imageFile}`); + const arrayBuffer = await response.arrayBuffer(); + const imageData = await decodeImage(arrayBuffer); + + enqueueImage(controller, imageFile, imageData); + }); + + await Promise.allSettled(promises); + controller.close(); + }, + }); +} + +function enqueueFileName( + controller: + | ReadableStreamDefaultController + | TransformStreamDefaultController, + fileName: string, +) { + const encodedFileName = textEncoder.encode(fileName); + dataView.setUint16(0, encodedFileName.byteLength, true); + controller.enqueue(temp.slice(0, 2)); + controller.enqueue(encodedFileName); +} + +function enqueueFile( + controller: + | ReadableStreamDefaultController + | TransformStreamDefaultController, + fileName: string, + data: Uint8Array, +) { + enqueueFileName(controller, fileName); + + dataView.setUint32(0, data.byteLength, true); + controller.enqueue(temp.slice()); + controller.enqueue(data); +} + +function enqueueImage( + controller: + | ReadableStreamDefaultController + | TransformStreamDefaultController, + fileName: string, + imageData: ImageData, +) { + enqueueFileName(controller, fileName); + + // byte size + dataView.setUint32(0, imageData.rgba.byteLength + 8, true); + controller.enqueue(temp.slice()); + + // width + dataView.setUint32(0, imageData.width, true); + controller.enqueue(temp.slice()); + // height + dataView.setUint32(0, imageData.height, true); + controller.enqueue(temp.slice()); + // rgba + controller.enqueue(imageData.rgba); +} + +// WebGL context for reading raw pixel data of WebP images +let offscreenGlContext: WebGL2RenderingContext | null = null; + +export async function decodeImage(fileBytes: ArrayBuffer) { + if (!offscreenGlContext) { + const canvas = new OffscreenCanvas(1, 1); + offscreenGlContext = canvas.getContext("webgl2"); + if (!offscreenGlContext) { + throw new Error("Failed to create WebGL2 context"); + } + } + + const imageBlob = new Blob([fileBytes]); + const bitmap = await createImageBitmap(imageBlob, { + premultiplyAlpha: "none", + }); + + const gl = offscreenGlContext; + const texture = gl.createTexture(); + gl.bindTexture(gl.TEXTURE_2D, texture); + gl.pixelStorei(gl.UNPACK_FLIP_Y_WEBGL, true); + gl.texImage2D(gl.TEXTURE_2D, 0, gl.RGBA, gl.RGBA, gl.UNSIGNED_BYTE, bitmap); + gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_MAG_FILTER, gl.NEAREST); + gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_MIN_FILTER, gl.NEAREST); + + const framebuffer = gl.createFramebuffer(); + gl.bindFramebuffer(gl.FRAMEBUFFER, framebuffer); + gl.framebufferTexture2D( + gl.FRAMEBUFFER, + gl.COLOR_ATTACHMENT0, + gl.TEXTURE_2D, + texture, + 0, + ); + + const data = new Uint8Array(bitmap.width * bitmap.height * 4); + gl.readPixels( + 0, + 0, + bitmap.width, + bitmap.height, + gl.RGBA, + gl.UNSIGNED_BYTE, + data, + ); + + gl.deleteTexture(texture); + gl.deleteFramebuffer(framebuffer); + + return { rgba: data, width: bitmap.width, height: bitmap.height }; +} diff --git a/src/worker.ts b/src/worker.ts index 1801a758..f0864c65 100644 --- a/src/worker.ts +++ b/src/worker.ts @@ -20,8 +20,9 @@ import init_wasm, { get_lod_tree_level, } from "spark-rs"; import type { ExtResult, PackedResult, SplatEncoding } from "./defines"; +import { fetchAndDecodeImages, unzipAndDecodeImages } from "./sogs"; -const rpcHandlers = { +export const rpcHandlers = { sortSplats16, sortSplats32, loadPackedSplats, @@ -96,37 +97,9 @@ function sortSplats32({ return { activeSplats, readback, ordering }; } -async function fetchRange({ - url, - requestHeader, - withCredentials, - offset, - bytes, -}: { - url: string; - requestHeader?: Record; - withCredentials?: string; - offset?: number; - bytes?: number; -}): Promise { - const request = new Request(url, { - headers: requestHeader ? new Headers(requestHeader) : undefined, - credentials: withCredentials ? "include" : "same-origin", - }); - if (offset !== undefined && bytes !== undefined) { - request.headers.set("Range", `bytes=${offset}-${offset + bytes - 1}`); - } - const response = await fetch(request); - if (!response.ok || !response.body) { - throw new Error( - `Failed to fetch "${url}": ${response.status} ${response.statusText}`, - ); - } - return new Uint8Array(await response.arrayBuffer()); -} - async function decodeBytesUrl({ decoder, + fileType, fileBytes, url, requestHeader, @@ -136,6 +109,7 @@ async function decodeBytesUrl({ sendStatus, }: { decoder: ChunkDecoder; + fileType?: string; fileBytes?: Uint8Array; url?: string; requestHeader?: Record; @@ -144,13 +118,20 @@ async function decodeBytesUrl({ chunkedLength?: number; sendStatus: (data: unknown) => void; }) { + let readStream: ReadableStream; + let streamLength = 0; + if (fileBytes) { - const CHUNK_SIZE = 1048576; // 1 MB - for (let i = 0; i < fileBytes.length; i += CHUNK_SIZE) { - decoder.push( - fileBytes.subarray(i, Math.min(i + CHUNK_SIZE, fileBytes.length)), - ); - } + readStream = new ReadableStream({ + start(controller) { + controller.enqueue(fileBytes); + controller.close(); + }, + }); + streamLength = fileBytes.length; + } else if (url && fileType === "pcsogs") { + // Unbundled SOG files require fetching and decoding + readStream = fetchAndDecodeImages(url); } else if (url) { const request = new Request(url, { headers: requestHeader ? new Headers(requestHeader) : undefined, @@ -163,49 +144,65 @@ async function decodeBytesUrl({ `Failed to fetch "${url}": ${response.status} ${response.statusText}`, ); } - const readStream = response.body.getReader(); + readStream = response.body; const contentLength = Number.parseInt( response.headers.get("Content-Length") || "0", ); - const total = Number.isNaN(contentLength) ? 0 : contentLength; - let loaded = 0; - - while (true) { - const { done, value } = await readStream.read(); - if (done) { - readStream.releaseLock(); - break; - } - loaded += value.length; - sendStatus({ loaded, total }); - - decoder.push(value); - } + streamLength = Number.isNaN(contentLength) ? 0 : contentLength; } else if (chunked) { - let loaded = 0; - const total = chunkedLength ?? 0; - while (true) { - const readNextChunk: Promise = new Promise((resolve) => { - nextChunkWaiter = resolve; - }); - sendStatus({ nextChunk: true }); - const nextChunk = await readNextChunk; - - if (nextChunk.length === 0) { - break; - } - - decoder.push(nextChunk); - loaded += nextChunk.length; - sendStatus({ progress: { loaded, total } }); - } - if (total === 0) { - sendStatus({ progress: { loaded, total: loaded } }); - } + readStream = new ReadableStream({ + async start(controller) { + async function readNext() { + const readNextChunk: Promise = new Promise((resolve) => { + nextChunkWaiter = resolve; + }); + sendStatus({ nextChunk: true }); + const nextChunk = await readNextChunk; + + if (nextChunk.length === 0) { + controller.close(); + return true; + } + + controller.enqueue(nextChunk); + return false; + } + + let final: boolean; + do { + final = await readNext(); + } while (!final); + }, + }); + streamLength = chunkedLength ?? 0; } else { throw new Error("No url or fileBytes provided"); } + // Handle SOG files + if ( + fileType === "pcsogszip" || + url?.endsWith(".sog") || + url?.endsWith(".sogs") || + url?.endsWith(".zip") + ) { + readStream = readStream.pipeThrough(unzipAndDecodeImages(streamLength)); + } + + const reader = readStream.getReader(); + let loaded = 0; + while (true) { + const { done, value } = await reader.read(); + if (done) { + reader.releaseLock(); + break; + } + loaded += value.length; + sendStatus({ loaded, total: streamLength }); + + decoder.push(value); + } + const decoded = decoder.finish(); return decoded; } @@ -294,6 +291,7 @@ async function loadPackedSplats( ); const decoded = await decodeBytesUrl({ decoder, + fileType, fileBytes, url, requestHeader, @@ -312,6 +310,7 @@ async function loadPackedSplats( const decoder = decode_to_csplatarray(fileType, pathName ?? url, encoding); const decoded = await decodeBytesUrl({ decoder, + fileType, fileBytes, url, requestHeader, @@ -455,6 +454,7 @@ async function loadExtSplats( ); const decoded = await decodeBytesUrl({ decoder, + fileType, fileBytes, url, requestHeader, @@ -473,6 +473,7 @@ async function loadExtSplats( const decoder = decode_to_gsplatarray(fileType, pathName ?? url); const decoded = await decodeBytesUrl({ decoder, + fileType, fileBytes, url, requestHeader,