Skip to content

Commit 2defaf3

Browse files
mrxzmrxz
authored
Add support for loading SOGSv2/SOG (#179)
1 parent 6f327a7 commit 2defaf3

3 files changed

Lines changed: 165 additions & 74 deletions

File tree

‎examples/editor/index.html‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -753,7 +753,8 @@
753753
file.name.toLowerCase().endsWith('.spz') ||
754754
file.name.toLowerCase().endsWith('.splat') ||
755755
file.name.toLowerCase().endsWith('.ksplat') ||
756-
file.name.toLowerCase().endsWith('.zip')
756+
file.name.toLowerCase().endsWith('.zip') ||
757+
file.name.toLowerCase().endsWith('.sog')
757758
);
758759

759760
if (splatFiles.length > 0) {

‎src/SplatLoader.ts‎

Lines changed: 58 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -282,6 +282,9 @@ export function getSplatFileTypeFromPath(
282282
if (extension === "ksplat") {
283283
return SplatFileType.KSPLAT;
284284
}
285+
if (extension === "sog") {
286+
return SplatFileType.PCSOGSZIP;
287+
}
285288
return undefined;
286289
}
287290

@@ -318,14 +321,40 @@ export type PcSogsJson = {
318321
};
319322
};
320323

324+
export type PcSogsV2Json = {
325+
version: 2;
326+
count: number;
327+
antialias?: boolean;
328+
means: {
329+
mins: number[];
330+
maxs: number[];
331+
files: string[];
332+
};
333+
scales: {
334+
codebook: number[];
335+
files: string[];
336+
};
337+
quats: { files: string[] };
338+
sh0: {
339+
codebook: number[];
340+
files: string[];
341+
};
342+
shN?: {
343+
count: number;
344+
bands: number;
345+
codebook: number[];
346+
files: string[];
347+
};
348+
};
349+
321350
export function isPcSogs(input: ArrayBuffer | Uint8Array | string): boolean {
322351
// Returns true if the input seems to be a valid PC SOGS file
323352
return tryPcSogs(input) !== undefined;
324353
}
325354

326355
export function tryPcSogs(
327356
input: ArrayBuffer | Uint8Array | string,
328-
): PcSogsJson | undefined {
357+
): PcSogsJson | PcSogsV2Json | undefined {
329358
// Try to parse input as SOGS JSON and see if it's valid
330359
try {
331360
let text: string;
@@ -345,6 +374,8 @@ export function tryPcSogs(
345374
if (!json || typeof json !== "object" || Array.isArray(json)) {
346375
return undefined;
347376
}
377+
const isVersion2 = json.version === 2;
378+
348379
for (const key of ["means", "scales", "quats", "sh0"]) {
349380
if (
350381
!json[key] ||
@@ -353,15 +384,33 @@ export function tryPcSogs(
353384
) {
354385
return undefined;
355386
}
356-
if (!json[key].shape || !json[key].files) {
357-
return undefined;
358-
}
359-
if (key !== "quats" && (!json[key].mins || !json[key].maxs)) {
360-
return undefined;
387+
if (isVersion2) {
388+
// Expect files
389+
if (!json[key].files) {
390+
return undefined;
391+
}
392+
393+
// Scales and sh0 should have codebooks
394+
if ((key === "scales" || key === "sh0") && !json[key].codebook) {
395+
return undefined;
396+
}
397+
// Means should have mins and maxs defined
398+
if (key === "means" && (!json[key].mins || !json[key].maxs)) {
399+
return undefined;
400+
}
401+
} else {
402+
// Expect shape and files
403+
if (!json[key].shape || !json[key].files) {
404+
return undefined;
405+
}
406+
// Besides 'quats' all other properties have mins and maxs
407+
if (key !== "quats" && (!json[key].mins || !json[key].maxs)) {
408+
return undefined;
409+
}
361410
}
362411
}
363412
// This is probably a PC SOGS file
364-
return json as PcSogsJson;
413+
return json as PcSogsJson | PcSogsV2Json;
365414
} catch {
366415
return undefined;
367416
}
@@ -388,6 +437,8 @@ export function tryPcSogsZip(
388437
if (!metaFilename) {
389438
return undefined;
390439
}
440+
441+
// Check for PC SOGS V1 and V2 (aka SOG)
391442
const json = tryPcSogs(unzipped[metaFilename]);
392443
if (!json) {
393444
return undefined;

‎src/pcsogs.ts‎

Lines changed: 105 additions & 66 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,10 @@
11
import { unzip } from "fflate";
22
import type { SplatEncoding } from "./PackedSplats";
3-
import { type PcSogsJson, tryPcSogsZip } from "./SplatLoader";
3+
import {
4+
type PcSogsJson,
5+
type PcSogsV2Json,
6+
tryPcSogsZip,
7+
} from "./SplatLoader";
48
import {
59
computeMaxSplats,
610
encodeSh1Rgb,
@@ -13,19 +17,21 @@ import {
1317
} from "./utils";
1418

1519
export async function unpackPcSogs(
16-
json: PcSogsJson,
20+
json: PcSogsJson | PcSogsV2Json,
1721
extraFiles: Record<string, ArrayBuffer>,
1822
splatEncoding: SplatEncoding,
1923
): Promise<{
2024
packedArray: Uint32Array;
2125
numSplats: number;
2226
extra: Record<string, unknown>;
2327
}> {
24-
if (json.quats.encoding !== "quaternion_packed") {
28+
const isVersion2 = "version" in json;
29+
30+
if (!isVersion2 && json.quats.encoding !== "quaternion_packed") {
2531
throw new Error("Unsupported quaternion encoding");
2632
}
2733

28-
const numSplats = json.means.shape[0];
34+
const numSplats = isVersion2 ? json.count : json.means.shape[0];
2935
const maxSplats = computeMaxSplats(numSplats);
3036
const packedArray = new Uint32Array(maxSplats * 4);
3137
const extra: Record<string, unknown> = {};
@@ -54,30 +60,41 @@ export async function unpackPcSogs(
5460

5561
const scalesPromise = decodeImageRgba(extraFiles[json.scales.files[0]]).then(
5662
(scales) => {
57-
const xLookup = new Array(256)
58-
.fill(0)
59-
.map(
60-
(_, i) =>
61-
json.scales.mins[0] +
62-
(json.scales.maxs[0] - json.scales.mins[0]) * (i / 255),
63-
)
64-
.map((x) => Math.exp(x));
65-
const yLookup = new Array(256)
66-
.fill(0)
67-
.map(
68-
(_, i) =>
69-
json.scales.mins[1] +
70-
(json.scales.maxs[1] - json.scales.mins[1]) * (i / 255),
71-
)
72-
.map((x) => Math.exp(x));
73-
const zLookup = new Array(256)
74-
.fill(0)
75-
.map(
76-
(_, i) =>
77-
json.scales.mins[2] +
78-
(json.scales.maxs[2] - json.scales.mins[2]) * (i / 255),
79-
)
80-
.map((x) => Math.exp(x));
63+
let xLookup: number[];
64+
let yLookup: number[];
65+
let zLookup: number[];
66+
67+
if (isVersion2) {
68+
xLookup =
69+
yLookup =
70+
zLookup =
71+
json.scales.codebook.map((x) => Math.exp(x));
72+
} else {
73+
xLookup = new Array(256)
74+
.fill(0)
75+
.map(
76+
(_, i) =>
77+
json.scales.mins[0] +
78+
(json.scales.maxs[0] - json.scales.mins[0]) * (i / 255),
79+
)
80+
.map((x) => Math.exp(x));
81+
yLookup = new Array(256)
82+
.fill(0)
83+
.map(
84+
(_, i) =>
85+
json.scales.mins[1] +
86+
(json.scales.maxs[1] - json.scales.mins[1]) * (i / 255),
87+
)
88+
.map((x) => Math.exp(x));
89+
zLookup = new Array(256)
90+
.fill(0)
91+
.map(
92+
(_, i) =>
93+
json.scales.mins[2] +
94+
(json.scales.maxs[2] - json.scales.mins[2]) * (i / 255),
95+
)
96+
.map((x) => Math.exp(x));
97+
}
8198

8299
for (let i = 0; i < numSplats; ++i) {
83100
const i4 = i * 4;
@@ -118,38 +135,51 @@ export async function unpackPcSogs(
118135
const sh0Promise = decodeImageRgba(extraFiles[json.sh0.files[0]]).then(
119136
(sh0) => {
120137
const SH_C0 = 0.28209479177387814;
121-
const rLookup = new Array(256)
122-
.fill(0)
123-
.map(
124-
(_, i) =>
125-
json.sh0.mins[0] +
126-
(json.sh0.maxs[0] - json.sh0.mins[0]) * (i / 255),
127-
)
128-
.map((x) => SH_C0 * x + 0.5);
129-
const gLookup = new Array(256)
130-
.fill(0)
131-
.map(
132-
(_, i) =>
133-
json.sh0.mins[1] +
134-
(json.sh0.maxs[1] - json.sh0.mins[1]) * (i / 255),
135-
)
136-
.map((x) => SH_C0 * x + 0.5);
137-
const bLookup = new Array(256)
138-
.fill(0)
139-
.map(
140-
(_, i) =>
141-
json.sh0.mins[2] +
142-
(json.sh0.maxs[2] - json.sh0.mins[2]) * (i / 255),
143-
)
144-
.map((x) => SH_C0 * x + 0.5);
145-
const aLookup = new Array(256)
146-
.fill(0)
147-
.map(
148-
(_, i) =>
149-
json.sh0.mins[3] +
150-
(json.sh0.maxs[3] - json.sh0.mins[3]) * (i / 255),
151-
)
152-
.map((x) => 1.0 / (1.0 + Math.exp(-x)));
138+
let rLookup: number[];
139+
let gLookup: number[];
140+
let bLookup: number[];
141+
let aLookup: number[];
142+
143+
if (isVersion2) {
144+
rLookup =
145+
gLookup =
146+
bLookup =
147+
json.sh0.codebook.map((x) => SH_C0 * x + 0.5);
148+
aLookup = new Array(256).fill(0).map((_, i) => i / 255);
149+
} else {
150+
rLookup = new Array(256)
151+
.fill(0)
152+
.map(
153+
(_, i) =>
154+
json.sh0.mins[0] +
155+
(json.sh0.maxs[0] - json.sh0.mins[0]) * (i / 255),
156+
)
157+
.map((x) => SH_C0 * x + 0.5);
158+
gLookup = new Array(256)
159+
.fill(0)
160+
.map(
161+
(_, i) =>
162+
json.sh0.mins[1] +
163+
(json.sh0.maxs[1] - json.sh0.mins[1]) * (i / 255),
164+
)
165+
.map((x) => SH_C0 * x + 0.5);
166+
bLookup = new Array(256)
167+
.fill(0)
168+
.map(
169+
(_, i) =>
170+
json.sh0.mins[2] +
171+
(json.sh0.maxs[2] - json.sh0.mins[2]) * (i / 255),
172+
)
173+
.map((x) => SH_C0 * x + 0.5);
174+
aLookup = new Array(256)
175+
.fill(0)
176+
.map(
177+
(_, i) =>
178+
json.sh0.mins[3] +
179+
(json.sh0.maxs[3] - json.sh0.mins[3]) * (i / 255),
180+
)
181+
.map((x) => 1.0 / (1.0 + Math.exp(-x)));
182+
}
153183

154184
for (let i = 0; i < numSplats; ++i) {
155185
const i4 = i * 4;
@@ -168,9 +198,15 @@ export async function unpackPcSogs(
168198

169199
const promises = [meansPromise, scalesPromise, quatsPromise, sh0Promise];
170200
if (json.shN) {
171-
const useSH3 = json.shN.shape[1] >= 48 - 3;
172-
const useSH2 = json.shN.shape[1] >= 27 - 3;
173-
const useSH1 = json.shN.shape[1] >= 12 - 3;
201+
const useSH3 = isVersion2
202+
? json.shN.bands >= 3
203+
: json.shN.shape[1] >= 48 - 3;
204+
const useSH2 = isVersion2
205+
? json.shN.bands >= 2
206+
: json.shN.shape[1] >= 27 - 3;
207+
const useSH1 = isVersion2
208+
? json.shN.bands >= 1
209+
: json.shN.shape[1] >= 12 - 3;
174210

175211
if (useSH1) extra.sh1 = new Uint32Array(numSplats * 2);
176212
if (useSH2) extra.sh2 = new Uint32Array(numSplats * 4);
@@ -185,9 +221,12 @@ export async function unpackPcSogs(
185221
decodeImage(extraFiles[json.shN.files[0]]),
186222
decodeImage(extraFiles[json.shN.files[1]]),
187223
]).then(([centroids, labels]) => {
188-
const lookup = new Array(256)
189-
.fill(0)
190-
.map((_, i) => shN.mins + (shN.maxs - shN.mins) * (i / 255));
224+
const lookup =
225+
"codebook" in shN
226+
? shN.codebook
227+
: new Array(256)
228+
.fill(0)
229+
.map((_, i) => shN.mins + (shN.maxs - shN.mins) * (i / 255));
191230

192231
for (let i = 0; i < numSplats; ++i) {
193232
const i4 = i * 4;

0 commit comments

Comments
 (0)