gaussian-splat-lite 1.4.2 → 1.5.0

package/bin/rad/parallel.js

@@ -9,13 +9,12 @@ import {
 import { Lookup, NONE } from "./grid.js";
 import { HALF, halfBits } from "./math.js";
 import {
-  EMPTY,
   INVALID,
-  KEEP,
   LATER,
   Merger,
   SH_DIMS,
   Similarities,
+  filterRecord,
   refresh,
   validRecord,
 } from "./splats.js";
@@ -60,6 +59,10 @@ const WORKER_BYTES = 40 * 1024 ** 2;
 const ENCODER_BYTES = 52 * 1024 ** 2;
 const MAIN_BYTES = 64 * 1024 ** 2;
 const CHUNKS_BYTES = 24 * 1024 ** 2;
+// What an SH clustering worker thread takes besides its kernel and jobs,
+// and the cache that clustering keeps: it reads rows in order.
+const CLUSTER_BYTES = 32 * 1024 ** 2;
+const STREAM_BYTES = 64 * 1024 ** 2;
 export const YOUNG_MB = 8;
 
 // Worker threads with `bytes` each: at most 8, using at most `share` of the
@@ -88,6 +91,79 @@ export function encoders(workspace, chunks) {
   return count;
 }
 
+// Makes room for SH clustering worker threads that take `bytes` each besides
+// the thread itself, and `extra` bytes on this thread, and returns how many
+// there are, at most `jobs`. They only compute, so they may use every other
+// core and all the room but a small cache.
+export function clusterers(workspace, bytes, extra, jobs) {
+  workspace.settle();
+  const each = CLUSTER_BYTES + bytes;
+  const count = Math.max(
+    0,
+    Math.min(
+      availableParallelism() - 1,
+      jobs,
+      Math.floor((workspace.room - extra - STREAM_BYTES) / each),
+    ),
+  );
+  workspace.expect(count * each + extra);
+  return count;
+}
+
+// Ordered request/reply workers for clustering and encoding. Their callers
+// schedule the jobs and close the workers; without workers, requests use the
+// local handler. Keep workers referenced so exit events can reject requests.
+export class WorkerGroup {
+  constructor(url, count, workerData, local, name) {
+    this.local = local;
+    this.workers = Array.from({ length: count }, () => {
+      const worker = new Worker(url, {
+        workerData,
+        resourceLimits: {
+          maxOldGenerationSizeMb: 64,
+          maxYoungGenerationSizeMb: YOUNG_MB,
+        },
+      });
+      worker.waiting = [];
+      worker.on("message", (result) => worker.waiting.shift().resolve(result));
+      // A failed worker fails its pending and later requests.
+      const fail = (error) => {
+        worker.failed ??= error;
+        for (const job of worker.waiting.splice(0)) job.reject(worker.failed);
+      };
+      worker.on("error", fail);
+      worker.on("exit", (code) =>
+        fail(new Error(`${name} worker exited (${code})`)),
+      );
+      return worker;
+    });
+  }
+
+  get threads() {
+    return this.workers.length ? this.workers : [undefined];
+  }
+
+  request(worker, message, transfer = []) {
+    if (!worker) return Promise.resolve(this.local.handle(message));
+    if (worker.failed) return Promise.reject(worker.failed);
+    return new Promise((resolve, reject) => {
+      worker.waiting.push({ resolve, reject });
+      worker.postMessage(message, transfer);
+    });
+  }
+
+  broadcast(message) {
+    return Promise.all(
+      this.threads.map((worker) => this.request(worker, message)),
+    );
+  }
+
+  // Resolves once the workers have stopped and their memory is free.
+  close() {
+    return Promise.all(this.workers.map((worker) => worker.terminate()));
+  }
+}
+
 const shared = (Type, size) =>
   new Type(new SharedArrayBuffer(size * Type.BYTES_PER_ELEMENT));
 
@@ -231,14 +307,15 @@ export class Pool {
     this.run(COPY, count, 1024);
   }
 
-  // Refreshes and checks `count` splats from `base` in place, setting KEEP,
-  // INVALID, EMPTY or LATER in `results`; see Splats.retain().
-  check(splats, base, count) {
+  // Refreshes, checks and optionally crops `count` splats from `base` in
+  // place, storing their filter states in `results`; see Splats.retain().
+  check(splats, base, count, withinDist) {
     this.share({
       splatRecords: splats.records.directory,
       splatSh: splats.sh
         .slice(0, splats.maxSh)
         .map((column) => column.directory),
+      withinDist,
     });
     Atomics.store(this.control, BASE, base);
     this.run(CHECK, count, 1024);
@@ -370,6 +447,7 @@ class Task {
     if (state.column)
       this.column = new SharedColumn(views, state.column, Uint16Array);
     if (state.splatRecords) {
+      this.withinDist = state.withinDist;
       this.splatRecords = new SharedColumn(
         views,
         state.splatRecords,
@@ -450,7 +528,7 @@ class Task {
       for (let d = 0; d < dims; d++)
         if (!Number.isFinite(HALF[row[at + d]])) return INVALID;
     }
-    return record[3] > 0 && record[14] > 0 ? KEEP : EMPTY;
+    return filterRecord(record, this.withinDist);
   }
 
   copy(from, to) {

package/bin/rad/rad.js

@@ -10,20 +10,22 @@ import {
   statSync,
 } from "node:fs";
 import path from "node:path";
-import {
-  Worker,
-  isMainThread,
-  parentPort,
-  workerData,
-} from "node:worker_threads";
+import { isMainThread, parentPort, workerData } from "node:worker_threads";
 import { deflateRawSync } from "node:zlib";
 import { HALF, half, halfBits, opacityRadius } from "./math.js";
-import { YOUNG_MB, encoders } from "./parallel.js";
+import { WorkerGroup, encoders } from "./parallel.js";
 import { Progress } from "./progress.js";
 import { SH_DIMS } from "./splats.js";
 import { readAll, removePath, writeAll } from "./storage.js";
 
 const CHUNK = 65536;
+// Fixed s8 range with SPZ's 1/128 SH step; -1 clips to -127/128.
+const SH_MAX = 127 / 128;
+function encodeSh(value) {
+  const magnitude = Math.min(127, Math.round(Math.abs(value) * 128));
+  return (value < 0 ? -magnitude : magnitude) & 255;
+}
+const SH_CODES = Uint8Array.from(HALF, encodeSh);
 function header(magic, meta) {
   const json = Buffer.from(JSON.stringify(meta));
   const result = Buffer.alloc(8 + Math.ceil(json.length / 8) * 8);
@@ -99,18 +101,15 @@ function alpha(value) {
     ? half(Math.min(2, 1 + 0.25 * (opacityRadius(value) - 1)))
     : value;
 }
-// What the encoding ranges depend on, counted from gathered chunks.
+// The base attributes that determine the encoding ranges.
 class Counts {
   constructor() {
     this.colors = new Histogram();
     this.scales = new Histogram();
-    this.sh = SH_DIMS.map(() => new Histogram());
     this.alphaF16 = false;
   }
-  add(chunk) {
-    const { records, count } = chunk;
-    for (let i = 0; i < count; i++) {
-      const offset = i * 16;
+  add(records) {
+    for (let offset = 0; offset < records.length; offset += 16) {
       this.alphaF16 ||= alpha(records[offset + 3]) > 1;
       for (let d = 4; d <= 6; d++)
         this.colors.add(halfBits(records[offset + d]));
@@ -122,19 +121,15 @@ class Counts {
       );
       this.scales.add(halfBits(Math.max(a, b, c)));
     }
-    for (let band = 0; band < chunk.sh.length; band++)
-      for (const bits of chunk.sh[band]) this.sh[band].add(bits);
   }
   merge(other) {
     this.colors.merge(other.colors);
     this.scales.merge(other.scales);
-    for (let band = 0; band < SH_DIMS.length; band++)
-      this.sh[band].merge(other.sh[band]);
     this.alphaF16 ||= other.alphaF16;
   }
 }
 
-function encoding(counts, maxSh) {
+function encoding(counts) {
   const { colors, scales } = counts;
   const rgbMin = Math.min(0, colors.percentile(0.01));
   const rgbMax = Math.max(1, colors.percentile(0.99));
@@ -143,69 +138,67 @@ function encoding(counts, maxSh) {
     Math.max(Math.log(1e-30), scales.percentile(0.01)),
   );
   const lnScaleMax = Math.max(9, scales.percentile(0.99));
-  const limits = [1, 1, 1];
-  for (let band = 0; band < maxSh; band++) {
-    const histogram = counts.sh[band];
-    limits[band] = Math.max(
-      1,
-      Math.abs(histogram.percentile(0.05)),
-      Math.abs(histogram.percentile(0.95)),
-    );
-  }
-  // The s8 code of each half value, per band.
-  const sh = limits.map((limit) => {
-    const codes = new Uint8Array(65536);
-    for (let bits = 0; bits < 65536; bits++)
-      codes[bits] =
-        signedRound(Math.max(-127, Math.min(127, (HALF[bits] / limit) * 127))) &
-        255;
-    return codes;
-  });
   return {
     alphaF16: counts.alphaF16,
     rgbF16: rgbMin < -1 || rgbMax > 2,
-    sh,
     scaleF16: lnScaleMax - lnScaleMin > 25,
-    limits,
     meta: {
       rgbMin,
       rgbMax,
       lnScaleMin,
       lnScaleMax,
-      sh1Max: limits[0],
-      sh2Max: limits[1],
-      sh3Max: limits[2],
+      sh1Max: SH_MAX,
+      sh2Max: SH_MAX,
+      sh3Max: SH_MAX,
       lodOpacity: true,
     },
   };
 }
 const unorm = (value, min, max) =>
   Math.round(Math.max(0, Math.min(255, ((value - min) / (max - min)) * 255)));
-const signedRound = (value) =>
-  value < 0 ? -Math.round(-value) : Math.round(value);
 
-// The columns of the splats from `base` on that encodeChunk() reads, and
-// Counts without the children.
-function gather(splats, base, count, children = true) {
+// The columns encodeChunk() reads. With a codebook, labels replace the SH.
+function gather(splats, base, count, codebook) {
   return {
     base,
     count,
+    maxSh: splats.maxSh,
     records: splats.records.read(base * 16, new Float32Array(count * 16)),
-    sh: splats.sh
-      .slice(0, splats.maxSh)
-      .map((column, band) =>
-        column.read(
-          base * SH_DIMS[band],
-          new Uint16Array(count * SH_DIMS[band]),
-        ),
-      ),
-    ...(children && {
-      counts: splats.childCount.read(base, new Uint32Array(count)),
-      starts: splats.childStart.read(base, new Float64Array(count)),
+    sh: codebook
+      ? []
+      : splats.sh
+          .slice(0, splats.maxSh)
+          .map((column, band) =>
+            column.read(
+              base * SH_DIMS[band],
+              new Uint16Array(count * SH_DIMS[band]),
+            ),
+          ),
+    ...(codebook && {
+      labels: codebook.labels.read(base, new Uint16Array(count)),
     }),
+    counts: splats.childCount.read(base, new Uint32Array(count)),
+    starts: splats.childStart.read(base, new Float64Array(count)),
   };
 }
 
+// The codebook's entries as s8 codes of each band, coefficient by coefficient.
+function encodeCodes(codebook, maxSh) {
+  const { count, dims, codes } = codebook;
+  const bands = [];
+  let offset = 0;
+  for (let band = 0; band < maxSh; band++) {
+    const size = SH_DIMS[band];
+    const bytes = new Uint8Array(count * size);
+    for (let d = 0; d < size; d++)
+      for (let c = 0; c < count; c++)
+        bytes[d * count + c] = encodeSh(codes[c * dims + offset + d]);
+    bands.push(bytes);
+    offset += size;
+  }
+  return bands;
+}
+
 // Raw property bytes, reused: each property is compressed before the next.
 const scratch = [Buffer.alloc(0), Buffer.alloc(0)];
 function raw(size, slot = 0) {
@@ -326,18 +319,25 @@ function encodeChunk(chunk, encoding) {
     bytes[i * 3 + 2] = unorm(theta, 0, Math.PI);
   }
   property("orientation", "oct88r8", bytes);
+  if (chunk.labels) {
+    // The first chunk holds the codebook for all.
+    for (let band = 0; chunk.codes && band < chunk.maxSh; band++)
+      property(`sh${band + 1}_code`, "s8", chunk.codes[band], [
+        -SH_MAX,
+        SH_MAX,
+      ]);
+    bytes = raw(count * 2);
+    for (let i = 0; i < count; i++) bytes.writeUInt16LE(chunk.labels[i], i * 2);
+    property("sh_label", "u16", bytes);
+  }
   for (let band = 0; band < chunk.sh.length; band++) {
     const dims = SH_DIMS[band];
     const values = chunk.sh[band];
-    const codes = encoding.sh[band];
     bytes = raw(count * dims);
     for (let d = 0; d < dims; d++)
       for (let i = 0; i < count; i++)
-        bytes[d * count + i] = codes[values[i * dims + d]];
-    property(`sh${band + 1}`, "s8", bytes, [
-      -encoding.limits[band],
-      encoding.limits[band],
-    ]);
+        bytes[d * count + i] = SH_CODES[values[i * dims + d]];
+    property(`sh${band + 1}`, "s8", bytes, [-SH_MAX, SH_MAX]);
   }
   const counts = raw(count * 2);
   const starts = raw(count * 4, 1);
@@ -352,7 +352,7 @@ function encodeChunk(chunk, encoding) {
     version: 1,
     base,
     count,
-    maxSh: chunk.sh.length,
+    maxSh: chunk.maxSh,
     lodTree: true,
     splatEncoding: encoding.meta,
     payloadBytes,
@@ -369,8 +369,8 @@ class Coder {
     this.counts = new Counts();
     this.format = undefined;
   }
-  handle({ type, chunk, format }) {
-    if (type === "count") this.counts.add(chunk);
+  handle({ type, records, chunk, format }) {
+    if (type === "count") this.counts.add(records);
     else if (type === "counts") return this.counts;
     else if (type === "format") this.format = format;
     else return encodeChunk(chunk, this.format);
@@ -379,62 +379,33 @@ class Coder {
 }
 
 // Spreads counting and encoding over worker threads.
-class Encoder {
+class Encoder extends WorkerGroup {
   // Encodes `chunks` with workers when there is more than one and the
   // workspace has room for them.
   constructor(workspace, chunks) {
+    const count = encoders(workspace, chunks);
+    super(
+      new URL(import.meta.url),
+      count,
+      { coder: true },
+      count ? undefined : new Coder(),
+      "Encoding",
+    );
     this.next = 0;
-    this.coder = new Coder();
-    this.workers = Array.from({ length: encoders(workspace, chunks) }, () => {
-      const worker = new Worker(new URL(import.meta.url), {
-        workerData: { coder: true },
-        resourceLimits: {
-          maxOldGenerationSizeMb: 64,
-          maxYoungGenerationSizeMb: YOUNG_MB,
-        },
-      });
-      worker.unref();
-      worker.waiting = [];
-      worker.on("message", (result) => worker.waiting.shift().resolve(result));
-      // A failed worker fails its pending and later requests.
-      const fail = (error) => {
-        worker.failed ??= error;
-        for (const job of worker.waiting.splice(0)) job.reject(worker.failed);
-      };
-      worker.on("error", fail);
-      worker.on("exit", (code) =>
-        fail(new Error(`Encoding worker exited (${code})`)),
-      );
-      return worker;
-    });
   }
   get size() {
     return this.workers.length;
   }
-  request(worker, message, transfer = []) {
-    if (!worker) return Promise.resolve(this.coder.handle(message));
-    if (worker.failed) return Promise.reject(worker.failed);
-    return new Promise((resolve, reject) => {
-      worker.waiting.push({ resolve, reject });
-      worker.postMessage(message, transfer);
-    });
-  }
-  // Sends a gathered chunk to the next worker.
-  send(type, chunk) {
+  // Sends the next job to the next worker.
+  send(message, transfer) {
     return this.request(
       this.workers[this.next++ % this.workers.length],
-      { type, chunk },
-      [
-        chunk.records.buffer,
-        ...chunk.sh.map((values) => values.buffer),
-        ...(chunk.counts ? [chunk.counts.buffer, chunk.starts.buffer] : []),
-      ],
+      message,
+      transfer,
     );
   }
-  // Sends `message` to every worker, or to this thread's coder without any.
-  broadcast(message) {
-    const all = this.workers.length ? this.workers : [undefined];
-    return Promise.all(all.map((worker) => this.request(worker, message)));
+  count(records) {
+    return this.send({ type: "count", records }, [records.buffer]);
   }
   async counts() {
     const counts = new Counts();
@@ -446,13 +417,16 @@ class Encoder {
     return this.broadcast({ type: "format", format });
   }
   async encode(chunk) {
-    const bytes = await this.send("encode", chunk);
+    const bytes = await this.send({ type: "encode", chunk }, [
+      chunk.records.buffer,
+      ...chunk.sh.map((values) => values.buffer),
+      ...(chunk.labels ? [chunk.labels.buffer] : []),
+      ...(chunk.codes ?? []).map((values) => values.buffer),
+      chunk.counts.buffer,
+      chunk.starts.buffer,
+    ]);
     return Buffer.from(bytes.buffer, bytes.byteOffset, bytes.length);
   }
-  // Resolves once the workers have stopped and their memory is free.
-  close() {
-    return Promise.all(this.workers.map((worker) => worker.terminate()));
-  }
 }
 
 if (!isMainThread && workerData?.coder) {
@@ -471,7 +445,8 @@ if (!isMainThread && workerData?.coder) {
   });
 }
 
-export async function writeRad(splats, output, chunked, lodBase) {
+// With a codebook from clusterSh(), splats store SH as its labels.
+export async function writeRad(splats, output, chunked, lodBase, codebook) {
   const chunks = [];
   const stem = chunkStem(output);
   const spool = chunked ? undefined : splats.workspace.file();
@@ -529,26 +504,34 @@ export async function writeRad(splats, output, chunked, lodBase) {
     let countedChunks = 0;
     for (let base = 0; base < splats.length; base += CHUNK) {
       splats.workspace.check();
+      const count = Math.min(CHUNK, splats.length - base);
+      const records = splats.records.read(
+        base * 16,
+        new Float32Array(count * 16),
+      );
       const counted = encoder
-        .send(
-          "count",
-          gather(splats, base, Math.min(CHUNK, splats.length - base), false),
-        )
+        .count(records)
         .then(() => progress.update(++countedChunks));
       counted.catch(() => {});
       pending.push(counted);
       if (pending.length > encoder.size + 1) await pending.shift();
     }
     while (pending.length) await pending.shift();
-    const format = encoding(await encoder.counts(), splats.maxSh);
+    const format = encoding(await encoder.counts());
     await encoder.use(format);
     progress.finish();
     progress = new Progress("Encoding RAD chunks", chunkCount);
     for (let base = 0; base < splats.length; base += CHUNK) {
       splats.workspace.check();
-      const bytes = encoder.encode(
-        gather(splats, base, Math.min(CHUNK, splats.length - base)),
+      const chunk = gather(
+        splats,
+        base,
+        Math.min(CHUNK, splats.length - base),
+        codebook,
       );
+      if (codebook && base === 0)
+        chunk.codes = encodeCodes(codebook, splats.maxSh);
+      const bytes = encoder.encode(chunk);
       bytes.catch(() => {});
       pending.push(bytes);
       if (pending.length > encoder.size + 1) store(await pending.shift());
@@ -565,6 +548,7 @@ export async function writeRad(splats, output, chunked, lodBase) {
       allChunkBytes: total,
       chunks,
       splatEncoding: format.meta,
+      ...(codebook && { shCodeCount: codebook.count }),
       comment: `Gaussian Splat Lite quality LOD (base ${lodBase})`,
     });
     const writing = new Progress(

package/bin/rad/splats.js

@@ -20,8 +20,9 @@ export const SH_DIMS = [9, 15, 21];
 export const KEEP = 0;
 export const INVALID = 1;
 export const EMPTY = 2;
+export const CROPPED = 3;
 // A worker found a page that is not cached.
-export const LATER = 3;
+export const LATER = 4;
 
 // Sets a record's feature size and area from its scales and opacity.
 export function refresh(record) {
@@ -45,6 +46,16 @@ export function validRecord(record) {
   return true;
 }
 
+// Classifies a valid record after refresh(), including optional cropping.
+export function filterRecord(record, withinDist) {
+  if (record[3] <= 0 || record[14] <= 0) return EMPTY;
+  if (!withinDist) return KEEP;
+  const [x, y, z, radius] = withinDist;
+  return Math.hypot(record[0] - x, record[1] - y, record[2] - z) > radius
+    ? CROPPED
+    : KEEP;
+}
+
 export class Splats {
   constructor(workspace, maxSh) {
     this.workspace = workspace;
@@ -52,6 +63,7 @@ export class Splats {
     this.inputSh = 0;
     this.length = 0;
     this.invalid = 0;
+    this.cropped = 0;
     this.records = new PagedArray(workspace, Float32Array);
     this.sh = SH_DIMS.map(() => new PagedArray(workspace, Uint16Array));
     // The two children of the i-th merged splat are entries 2i and 2i + 1 of
@@ -127,9 +139,11 @@ export class Splats {
 
   // Skips empty splats and counts invalid ones in `invalid`: NaN or infinite
   // values, including a zero rotation, which the decoder normalizes to NaN.
-  // Workers refresh and check blocks of splats in place; the main thread
+  // Optionally keeps only centers within [x, y, z, radius], counting the
+  // otherwise usable splats outside this sphere in `cropped`.
+  // Workers refresh, check and crop blocks of splats; the main thread
   // checks those they leave, and moves kept splats after skipped ones.
-  retain(pool) {
+  retain(pool, withinDist) {
     const progress = new Progress("Filtering splats", this.length);
     const { results } = pool;
     const block = results.length;
@@ -143,6 +157,7 @@ export class Splats {
     ];
     let count = 0;
     this.invalid = 0;
+    this.cropped = 0;
     for (let base = 0; base < this.length; base += block) {
       const size = Math.min(block, this.length - base);
       for (const [column, stride] of columns)
@@ -151,24 +166,26 @@ export class Splats {
           (base + size) * stride,
           column === this.records,
         );
-      pool.check(this, base, size);
+      pool.check(this, base, size, withinDist);
       for (let i = 0; i < size; i++) {
         const index = base + i;
         let state = results[i];
-        if (state === LATER) {
+        const later = state === LATER;
+        if (later) {
           this.read(index, record);
           refresh(record);
           state = !this.valid(index, record, rows)
             ? INVALID
-            : record[3] > 0 && record[14] > 0
-              ? KEEP
-              : EMPTY;
-          if (state === KEEP) this.records.write(index * 16, record);
+            : filterRecord(record, withinDist);
         }
         if (state === INVALID) this.invalid++;
+        if (state === CROPPED) this.cropped++;
         if (state !== KEEP) continue;
+        if (count !== index || later) {
+          if (!later) this.read(index, record);
+          this.records.write(count * 16, record);
+        }
         if (count !== index) {
-          this.records.write(count * 16, this.read(index, record));
           for (let band = 0; band < this.maxSh; band++) {
             const dims = SH_DIMS[band];
             this.sh[band].read(index * dims, rows[band]);

package/bin/rad/storage.js

@@ -129,7 +129,10 @@ export class Workspace {
   shrink() {
     if (this.slots.length * PAGE_BYTES <= this.capacity) return;
     const slots = Math.floor(this.capacity / this.arenaBytes) * this.arenaPages;
-    for (const page of this.slots.slice(slots)) if (page) this.drop(page);
+    // Without a copy of the slots: a copy kept the released pages alive
+    // through the collection below, and their memory returned only later.
+    for (let slot = slots; slot < this.slots.length; slot++)
+      if (this.slots[slot]) this.drop(this.slots[slot]);
     this.slots.length = slots;
     this.reuse.length = slots;
     this.free = this.free.filter((slot) => slot < slots);