Add support for loading SOGSv2/SOG (#179)

This commit is contained in:
Noeri Huisman
2025-09-25 17:24:41 +02:00
committed by GitHub
parent 6f327a716f
commit 2defaf3f64
3 changed files with 165 additions and 74 deletions
+2 -1
View File
@@ -753,7 +753,8 @@
file.name.toLowerCase().endsWith('.spz') || file.name.toLowerCase().endsWith('.spz') ||
file.name.toLowerCase().endsWith('.splat') || file.name.toLowerCase().endsWith('.splat') ||
file.name.toLowerCase().endsWith('.ksplat') || file.name.toLowerCase().endsWith('.ksplat') ||
file.name.toLowerCase().endsWith('.zip') file.name.toLowerCase().endsWith('.zip') ||
file.name.toLowerCase().endsWith('.sog')
); );
if (splatFiles.length > 0) { if (splatFiles.length > 0) {
+58 -7
View File
@@ -282,6 +282,9 @@ export function getSplatFileTypeFromPath(
if (extension === "ksplat") { if (extension === "ksplat") {
return SplatFileType.KSPLAT; return SplatFileType.KSPLAT;
} }
if (extension === "sog") {
return SplatFileType.PCSOGSZIP;
}
return undefined; return undefined;
} }
@@ -318,6 +321,32 @@ export type PcSogsJson = {
}; };
}; };
export type PcSogsV2Json = {
version: 2;
count: number;
antialias?: boolean;
means: {
mins: number[];
maxs: number[];
files: string[];
};
scales: {
codebook: number[];
files: string[];
};
quats: { files: string[] };
sh0: {
codebook: number[];
files: string[];
};
shN?: {
count: number;
bands: number;
codebook: number[];
files: string[];
};
};
export function isPcSogs(input: ArrayBuffer | Uint8Array | string): boolean { export function isPcSogs(input: ArrayBuffer | Uint8Array | string): boolean {
// Returns true if the input seems to be a valid PC SOGS file // Returns true if the input seems to be a valid PC SOGS file
return tryPcSogs(input) !== undefined; return tryPcSogs(input) !== undefined;
@@ -325,7 +354,7 @@ export function isPcSogs(input: ArrayBuffer | Uint8Array | string): boolean {
export function tryPcSogs( export function tryPcSogs(
input: ArrayBuffer | Uint8Array | string, input: ArrayBuffer | Uint8Array | string,
): PcSogsJson | undefined { ): PcSogsJson | PcSogsV2Json | undefined {
// Try to parse input as SOGS JSON and see if it's valid // Try to parse input as SOGS JSON and see if it's valid
try { try {
let text: string; let text: string;
@@ -345,6 +374,8 @@ export function tryPcSogs(
if (!json || typeof json !== "object" || Array.isArray(json)) { if (!json || typeof json !== "object" || Array.isArray(json)) {
return undefined; return undefined;
} }
const isVersion2 = json.version === 2;
for (const key of ["means", "scales", "quats", "sh0"]) { for (const key of ["means", "scales", "quats", "sh0"]) {
if ( if (
!json[key] || !json[key] ||
@@ -353,15 +384,33 @@ export function tryPcSogs(
) { ) {
return undefined; return undefined;
} }
if (!json[key].shape || !json[key].files) { if (isVersion2) {
return undefined; // Expect files
} if (!json[key].files) {
if (key !== "quats" && (!json[key].mins || !json[key].maxs)) { return undefined;
return undefined; }
// Scales and sh0 should have codebooks
if ((key === "scales" || key === "sh0") && !json[key].codebook) {
return undefined;
}
// Means should have mins and maxs defined
if (key === "means" && (!json[key].mins || !json[key].maxs)) {
return undefined;
}
} else {
// Expect shape and files
if (!json[key].shape || !json[key].files) {
return undefined;
}
// Besides 'quats' all other properties have mins and maxs
if (key !== "quats" && (!json[key].mins || !json[key].maxs)) {
return undefined;
}
} }
} }
// This is probably a PC SOGS file // This is probably a PC SOGS file
return json as PcSogsJson; return json as PcSogsJson | PcSogsV2Json;
} catch { } catch {
return undefined; return undefined;
} }
@@ -388,6 +437,8 @@ export function tryPcSogsZip(
if (!metaFilename) { if (!metaFilename) {
return undefined; return undefined;
} }
// Check for PC SOGS V1 and V2 (aka SOG)
const json = tryPcSogs(unzipped[metaFilename]); const json = tryPcSogs(unzipped[metaFilename]);
if (!json) { if (!json) {
return undefined; return undefined;
+105 -66
View File
@@ -1,6 +1,10 @@
import { unzip } from "fflate"; import { unzip } from "fflate";
import type { SplatEncoding } from "./PackedSplats"; import type { SplatEncoding } from "./PackedSplats";
import { type PcSogsJson, tryPcSogsZip } from "./SplatLoader"; import {
type PcSogsJson,
type PcSogsV2Json,
tryPcSogsZip,
} from "./SplatLoader";
import { import {
computeMaxSplats, computeMaxSplats,
encodeSh1Rgb, encodeSh1Rgb,
@@ -13,7 +17,7 @@ import {
} from "./utils"; } from "./utils";
export async function unpackPcSogs( export async function unpackPcSogs(
json: PcSogsJson, json: PcSogsJson | PcSogsV2Json,
extraFiles: Record<string, ArrayBuffer>, extraFiles: Record<string, ArrayBuffer>,
splatEncoding: SplatEncoding, splatEncoding: SplatEncoding,
): Promise<{ ): Promise<{
@@ -21,11 +25,13 @@ export async function unpackPcSogs(
numSplats: number; numSplats: number;
extra: Record<string, unknown>; extra: Record<string, unknown>;
}> { }> {
if (json.quats.encoding !== "quaternion_packed") { const isVersion2 = "version" in json;
if (!isVersion2 && json.quats.encoding !== "quaternion_packed") {
throw new Error("Unsupported quaternion encoding"); throw new Error("Unsupported quaternion encoding");
} }
const numSplats = json.means.shape[0]; const numSplats = isVersion2 ? json.count : json.means.shape[0];
const maxSplats = computeMaxSplats(numSplats); const maxSplats = computeMaxSplats(numSplats);
const packedArray = new Uint32Array(maxSplats * 4); const packedArray = new Uint32Array(maxSplats * 4);
const extra: Record<string, unknown> = {}; const extra: Record<string, unknown> = {};
@@ -54,30 +60,41 @@ export async function unpackPcSogs(
const scalesPromise = decodeImageRgba(extraFiles[json.scales.files[0]]).then( const scalesPromise = decodeImageRgba(extraFiles[json.scales.files[0]]).then(
(scales) => { (scales) => {
const xLookup = new Array(256) let xLookup: number[];
.fill(0) let yLookup: number[];
.map( let zLookup: number[];
(_, i) =>
json.scales.mins[0] + if (isVersion2) {
(json.scales.maxs[0] - json.scales.mins[0]) * (i / 255), xLookup =
) yLookup =
.map((x) => Math.exp(x)); zLookup =
const yLookup = new Array(256) json.scales.codebook.map((x) => Math.exp(x));
.fill(0) } else {
.map( xLookup = new Array(256)
(_, i) => .fill(0)
json.scales.mins[1] + .map(
(json.scales.maxs[1] - json.scales.mins[1]) * (i / 255), (_, i) =>
) json.scales.mins[0] +
.map((x) => Math.exp(x)); (json.scales.maxs[0] - json.scales.mins[0]) * (i / 255),
const zLookup = new Array(256) )
.fill(0) .map((x) => Math.exp(x));
.map( yLookup = new Array(256)
(_, i) => .fill(0)
json.scales.mins[2] + .map(
(json.scales.maxs[2] - json.scales.mins[2]) * (i / 255), (_, i) =>
) json.scales.mins[1] +
.map((x) => Math.exp(x)); (json.scales.maxs[1] - json.scales.mins[1]) * (i / 255),
)
.map((x) => Math.exp(x));
zLookup = new Array(256)
.fill(0)
.map(
(_, i) =>
json.scales.mins[2] +
(json.scales.maxs[2] - json.scales.mins[2]) * (i / 255),
)
.map((x) => Math.exp(x));
}
for (let i = 0; i < numSplats; ++i) { for (let i = 0; i < numSplats; ++i) {
const i4 = i * 4; const i4 = i * 4;
@@ -118,38 +135,51 @@ export async function unpackPcSogs(
const sh0Promise = decodeImageRgba(extraFiles[json.sh0.files[0]]).then( const sh0Promise = decodeImageRgba(extraFiles[json.sh0.files[0]]).then(
(sh0) => { (sh0) => {
const SH_C0 = 0.28209479177387814; const SH_C0 = 0.28209479177387814;
const rLookup = new Array(256) let rLookup: number[];
.fill(0) let gLookup: number[];
.map( let bLookup: number[];
(_, i) => let aLookup: number[];
json.sh0.mins[0] +
(json.sh0.maxs[0] - json.sh0.mins[0]) * (i / 255), if (isVersion2) {
) rLookup =
.map((x) => SH_C0 * x + 0.5); gLookup =
const gLookup = new Array(256) bLookup =
.fill(0) json.sh0.codebook.map((x) => SH_C0 * x + 0.5);
.map( aLookup = new Array(256).fill(0).map((_, i) => i / 255);
(_, i) => } else {
json.sh0.mins[1] + rLookup = new Array(256)
(json.sh0.maxs[1] - json.sh0.mins[1]) * (i / 255), .fill(0)
) .map(
.map((x) => SH_C0 * x + 0.5); (_, i) =>
const bLookup = new Array(256) json.sh0.mins[0] +
.fill(0) (json.sh0.maxs[0] - json.sh0.mins[0]) * (i / 255),
.map( )
(_, i) => .map((x) => SH_C0 * x + 0.5);
json.sh0.mins[2] + gLookup = new Array(256)
(json.sh0.maxs[2] - json.sh0.mins[2]) * (i / 255), .fill(0)
) .map(
.map((x) => SH_C0 * x + 0.5); (_, i) =>
const aLookup = new Array(256) json.sh0.mins[1] +
.fill(0) (json.sh0.maxs[1] - json.sh0.mins[1]) * (i / 255),
.map( )
(_, i) => .map((x) => SH_C0 * x + 0.5);
json.sh0.mins[3] + bLookup = new Array(256)
(json.sh0.maxs[3] - json.sh0.mins[3]) * (i / 255), .fill(0)
) .map(
.map((x) => 1.0 / (1.0 + Math.exp(-x))); (_, i) =>
json.sh0.mins[2] +
(json.sh0.maxs[2] - json.sh0.mins[2]) * (i / 255),
)
.map((x) => SH_C0 * x + 0.5);
aLookup = new Array(256)
.fill(0)
.map(
(_, i) =>
json.sh0.mins[3] +
(json.sh0.maxs[3] - json.sh0.mins[3]) * (i / 255),
)
.map((x) => 1.0 / (1.0 + Math.exp(-x)));
}
for (let i = 0; i < numSplats; ++i) { for (let i = 0; i < numSplats; ++i) {
const i4 = i * 4; const i4 = i * 4;
@@ -168,9 +198,15 @@ export async function unpackPcSogs(
const promises = [meansPromise, scalesPromise, quatsPromise, sh0Promise]; const promises = [meansPromise, scalesPromise, quatsPromise, sh0Promise];
if (json.shN) { if (json.shN) {
const useSH3 = json.shN.shape[1] >= 48 - 3; const useSH3 = isVersion2
const useSH2 = json.shN.shape[1] >= 27 - 3; ? json.shN.bands >= 3
const useSH1 = json.shN.shape[1] >= 12 - 3; : json.shN.shape[1] >= 48 - 3;
const useSH2 = isVersion2
? json.shN.bands >= 2
: json.shN.shape[1] >= 27 - 3;
const useSH1 = isVersion2
? json.shN.bands >= 1
: json.shN.shape[1] >= 12 - 3;
if (useSH1) extra.sh1 = new Uint32Array(numSplats * 2); if (useSH1) extra.sh1 = new Uint32Array(numSplats * 2);
if (useSH2) extra.sh2 = new Uint32Array(numSplats * 4); if (useSH2) extra.sh2 = new Uint32Array(numSplats * 4);
@@ -185,9 +221,12 @@ export async function unpackPcSogs(
decodeImage(extraFiles[json.shN.files[0]]), decodeImage(extraFiles[json.shN.files[0]]),
decodeImage(extraFiles[json.shN.files[1]]), decodeImage(extraFiles[json.shN.files[1]]),
]).then(([centroids, labels]) => { ]).then(([centroids, labels]) => {
const lookup = new Array(256) const lookup =
.fill(0) "codebook" in shN
.map((_, i) => shN.mins + (shN.maxs - shN.mins) * (i / 255)); ? shN.codebook
: new Array(256)
.fill(0)
.map((_, i) => shN.mins + (shN.maxs - shN.mins) * (i / 255));
for (let i = 0; i < numSplats; ++i) { for (let i = 0; i < numSplats; ++i) {
const i4 = i * 4; const i4 = i * 4;