@d5techs/3dgs-lib 1.4.22 → 1.4.23

package/dist/3dgs-lib.js

@@ -6301,91 +6301,231 @@ function findNearest(sortedCB, value) {
   }
   return lo;
 }
-function kmeansVectors(vectors, dim, k, iterations) {
-  const n = vectors.length;
-  const labels = new Uint16Array(n);
-  if (n <= k) {
-    const centroids2 = [];
-    for (let i = 0; i < n; i++) {
-      centroids2.push(new Float32Array(vectors[i]));
-      labels[i] = i;
+function getGpuClusterShader(dim) {
+  return (
+    /* wgsl */
+    `
+struct Uniforms {
+  numPoints: u32,
+  numCentroids: u32
+};
+
+@group(0) @binding(0) var<uniform> uniforms: Uniforms;
+@group(0) @binding(1) var<storage, read> points: array<f32>;
+@group(0) @binding(2) var<storage, read> centroids: array<f32>;
+@group(0) @binding(3) var<storage, read_write> results: array<u32>;
+
+const DIM = ${dim}u;
+const CHUNK_SIZE = 128u;
+var<workgroup> sharedChunk: array<f32, ${dim * 128}>;
+
+fn calcDistSqr(point: array<f32, ${dim}>, cIdx: u32) -> f32 {
+  var result = 0.0;
+  var ci = cIdx * DIM;
+  for (var i = 0u; i < DIM; i++) {
+    let v = point[i] - sharedChunk[ci + i];
+    result += v * v;
+  }
+  return result;
+}
+
+@compute @workgroup_size(64)
+fn main(
+  @builtin(local_invocation_index) localId: u32,
+  @builtin(global_invocation_id) globalId: vec3u,
+  @builtin(num_workgroups) numWgs: vec3u
+) {
+  let pointIdx = globalId.x + globalId.y * numWgs.x * 64u;
+
+  var point: array<f32, ${dim}>;
+  if (pointIdx < uniforms.numPoints) {
+    for (var i = 0u; i < DIM; i++) {
+      point[i] = points[pointIdx * DIM + i];
     }
-    while (centroids2.length < k) centroids2.push(new Float32Array(dim));
-    return { centroids: centroids2, labels };
   }
-  const centroids = [];
-  centroids.push(new Float32Array(vectors[Math.floor(Math.random() * n)]));
-  const dist = new Float64Array(n);
-  dist.fill(Infinity);
-  for (let c = 1; c < k; c++) {
-    const last = centroids[c - 1];
-    let totalDist = 0;
-    for (let i = 0; i < n; i++) {
-      let d2 = 0;
-      for (let d = 0; d < dim; d++) {
-        const diff = vectors[i][d] - last[d];
-        d2 += diff * diff;
-      }
-      if (d2 < dist[i]) dist[i] = d2;
-      totalDist += dist[i];
+
+  var mind = 1e20;
+  var mini = 0u;
+
+  let numChunks = (uniforms.numCentroids + CHUNK_SIZE - 1u) / CHUNK_SIZE;
+  for (var chunk = 0u; chunk < numChunks; chunk++) {
+
+    let rowsPerThread = CHUNK_SIZE / 64u;
+    let dstRow = localId * rowsPerThread;
+    let srcRow = min(uniforms.numCentroids, chunk * CHUNK_SIZE + dstRow);
+    let numRows = min(uniforms.numCentroids, srcRow + rowsPerThread) - srcRow;
+
+    var dst = dstRow * DIM;
+    var src = srcRow * DIM;
+    for (var c = 0u; c < numRows * DIM; c++) {
+      sharedChunk[dst + c] = centroids[src + c];
     }
-    let target = Math.random() * totalDist;
-    let chosen = 0;
-    for (let i = 0; i < n; i++) {
-      target -= dist[i];
-      if (target <= 0) {
-        chosen = i;
-        break;
+
+    workgroupBarrier();
+
+    if (pointIdx < uniforms.numPoints) {
+      let thisChunkSize = min(CHUNK_SIZE, uniforms.numCentroids - chunk * CHUNK_SIZE);
+      for (var c = 0u; c < thisChunkSize; c++) {
+        let d = calcDistSqr(point, c);
+        if (d < mind) {
+          mind = d;
+          mini = chunk * CHUNK_SIZE + c;
+        }
       }
     }
-    centroids.push(new Float32Array(vectors[chosen]));
+
+    workgroupBarrier();
   }
-  for (let iter = 0; iter < iterations; iter++) {
-    for (let i = 0; i < n; i++) {
-      let bestK = 0;
-      let bestD = Infinity;
-      for (let c = 0; c < k; c++) {
-        let d2 = 0;
-        for (let d = 0; d < dim; d++) {
-          const diff = vectors[i][d] - centroids[c][d];
-          d2 += diff * diff;
-        }
-        if (d2 < bestD) {
-          bestD = d2;
-          bestK = c;
-        }
-      }
-      labels[i] = bestK;
+
+  if (pointIdx < uniforms.numPoints) {
+    results[pointIdx] = mini;
+  }
+}`
+  );
+}
+let _cachedDevice = null;
+let _deviceFailed = false;
+async function getGpuDevice() {
+  if (_cachedDevice) return _cachedDevice;
+  if (_deviceFailed) return null;
+  try {
+    if (typeof navigator === "undefined" || !navigator.gpu) {
+      _deviceFailed = true;
+      return null;
     }
-    const sums = [];
-    for (let c = 0; c < k; c++) sums.push(new Float64Array(dim));
-    const counts = new Uint32Array(k);
-    for (let i = 0; i < n; i++) {
-      const l = labels[i];
-      counts[l]++;
-      for (let d = 0; d < dim; d++) sums[l][d] += vectors[i][d];
+    const adapter = await navigator.gpu.requestAdapter();
+    if (!adapter) {
+      _deviceFailed = true;
+      return null;
     }
-    let changed = false;
-    for (let c = 0; c < k; c++) {
-      if (counts[c] > 0) {
-        for (let d = 0; d < dim; d++) {
-          const v = sums[c][d] / counts[c];
-          if (v !== centroids[c][d]) {
-            centroids[c][d] = v;
-            changed = true;
-          }
-        }
+    _cachedDevice = await adapter.requestDevice({
+      requiredLimits: {
+        maxStorageBufferBindingSize: adapter.limits.maxStorageBufferBindingSize,
+        maxBufferSize: adapter.limits.maxBufferSize
       }
+    });
+    _cachedDevice.lost.then(() => {
+      _cachedDevice = null;
+    });
+    return _cachedDevice;
+  } catch {
+    _deviceFailed = true;
+    return null;
+  }
+}
+async function gpuAssignLabels(device, pointsFlat, centroidsFlat, numPoints, numCentroids, dim) {
+  const WORKGROUP_SIZE2 = 64;
+  const BATCH_SIZE = 1024 * WORKGROUP_SIZE2;
+  const numBatches = Math.ceil(numPoints / BATCH_SIZE);
+  const shaderModule = device.createShaderModule({ code: getGpuClusterShader(dim) });
+  const bindGroupLayout = device.createBindGroupLayout({
+    entries: [
+      { binding: 0, visibility: GPUShaderStage.COMPUTE, buffer: { type: "uniform" } },
+      { binding: 1, visibility: GPUShaderStage.COMPUTE, buffer: { type: "read-only-storage" } },
+      { binding: 2, visibility: GPUShaderStage.COMPUTE, buffer: { type: "read-only-storage" } },
+      { binding: 3, visibility: GPUShaderStage.COMPUTE, buffer: { type: "storage" } }
+    ]
+  });
+  const pipeline = device.createComputePipeline({
+    layout: device.createPipelineLayout({ bindGroupLayouts: [bindGroupLayout] }),
+    compute: { module: shaderModule, entryPoint: "main" }
+  });
+  const uniformBuf = device.createBuffer({ size: 8, usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST });
+  const centroidsBuf = device.createBuffer({ size: centroidsFlat.byteLength, usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST });
+  device.queue.writeBuffer(centroidsBuf, 0, centroidsFlat.buffer);
+  const pointsBufSize = BATCH_SIZE * dim * 4;
+  const pointsBuf = device.createBuffer({ size: pointsBufSize, usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST });
+  const resultsBufSize = BATCH_SIZE * 4;
+  const resultsBuf = device.createBuffer({ size: resultsBufSize, usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC });
+  const readbackBuf = device.createBuffer({ size: resultsBufSize, usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST });
+  const labels = new Uint32Array(numPoints);
+  const bindGroup = device.createBindGroup({
+    layout: bindGroupLayout,
+    entries: [
+      { binding: 0, resource: { buffer: uniformBuf } },
+      { binding: 1, resource: { buffer: pointsBuf } },
+      { binding: 2, resource: { buffer: centroidsBuf } },
+      { binding: 3, resource: { buffer: resultsBuf } }
+    ]
+  });
+  for (let batch = 0; batch < numBatches; batch++) {
+    const offset = batch * BATCH_SIZE;
+    const currentBatchSize = Math.min(numPoints - offset, BATCH_SIZE);
+    const groups = Math.ceil(currentBatchSize / WORKGROUP_SIZE2);
+    device.queue.writeBuffer(uniformBuf, 0, new Uint32Array([currentBatchSize, numCentroids]));
+    device.queue.writeBuffer(pointsBuf, 0, pointsFlat.buffer, offset * dim * 4, currentBatchSize * dim * 4);
+    const encoder = device.createCommandEncoder();
+    const pass = encoder.beginComputePass();
+    pass.setPipeline(pipeline);
+    pass.setBindGroup(0, bindGroup);
+    pass.dispatchWorkgroups(groups);
+    pass.end();
+    encoder.copyBufferToBuffer(resultsBuf, 0, readbackBuf, 0, currentBatchSize * 4);
+    device.queue.submit([encoder.finish()]);
+    await readbackBuf.mapAsync(GPUMapMode.READ);
+    const mapped = new Uint32Array(readbackBuf.getMappedRange(0, currentBatchSize * 4));
+    labels.set(mapped, offset);
+    readbackBuf.unmap();
+  }
+  uniformBuf.destroy();
+  centroidsBuf.destroy();
+  pointsBuf.destroy();
+  resultsBuf.destroy();
+  readbackBuf.destroy();
+  return labels;
+}
+async function kmeansGpu(pointsFlat, n, dim, k, iterations, onProgress) {
+  const labels = new Uint32Array(n);
+  if (n <= k) {
+    const centroids2 = new Float32Array(k * dim);
+    for (let i = 0; i < n; i++) {
+      centroids2.set(pointsFlat.subarray(i * dim, i * dim + dim), i * dim);
+      labels[i] = i;
     }
-    if (!changed) break;
+    return { centroids: centroids2, labels };
   }
+  const centroids = new Float32Array(k * dim);
+  const indices = pickRandomIndices(n, k);
+  for (let i = 0; i < k; i++) {
+    centroids.set(pointsFlat.subarray(indices[i] * dim, indices[i] * dim + dim), i * dim);
+  }
+  const device = await getGpuDevice();
+  for (let iter = 0; iter < iterations; iter++) {
+    onProgress == null ? void 0 : onProgress(iter, iterations);
+    if (device) {
+      const newLabels = await gpuAssignLabels(device, pointsFlat, centroids, n, k, dim);
+      labels.set(newLabels);
+    } else {
+      assignLabelsCpu(pointsFlat, centroids, labels, n, k, dim);
+    }
+    updateCentroids(pointsFlat, centroids, labels, n, k, dim);
+  }
+  onProgress == null ? void 0 : onProgress(iterations, iterations);
+  if (device) {
+    const finalLabels = await gpuAssignLabels(device, pointsFlat, centroids, n, k, dim);
+    labels.set(finalLabels);
+  } else {
+    assignLabelsCpu(pointsFlat, centroids, labels, n, k, dim);
+  }
+  return { centroids, labels };
+}
+function pickRandomIndices(n, m) {
+  const chosen = /* @__PURE__ */ new Set();
+  for (let j = n - m; j < n; j++) {
+    const t = Math.floor(Math.random() * (j + 1));
+    chosen.add(chosen.has(t) ? j : t);
+  }
+  return [...chosen];
+}
+function assignLabelsCpu(points, centroids, labels, n, k, dim) {
   for (let i = 0; i < n; i++) {
-    let bestK = 0;
-    let bestD = Infinity;
+    const pOff = i * dim;
+    let bestK = 0, bestD = Infinity;
     for (let c = 0; c < k; c++) {
+      const cOff = c * dim;
       let d2 = 0;
       for (let d = 0; d < dim; d++) {
-        const diff = vectors[i][d] - centroids[c][d];
+        const diff = points[pOff + d] - centroids[cOff + d];
         d2 += diff * diff;
       }
       if (d2 < bestD) {
@@ -6395,7 +6535,25 @@ function kmeansVectors(vectors, dim, k, iterations) {
     }
     labels[i] = bestK;
   }
-  return { centroids: centroids.map((c) => new Float32Array(c)), labels };
+}
+function updateCentroids(points, centroids, labels, n, k, dim) {
+  const sums = new Float64Array(k * dim);
+  const counts = new Uint32Array(k);
+  for (let i = 0; i < n; i++) {
+    const l = labels[i];
+    counts[l]++;
+    const pOff = i * dim, sOff = l * dim;
+    for (let d = 0; d < dim; d++) sums[sOff + d] += points[pOff + d];
+  }
+  for (let c = 0; c < k; c++) {
+    if (counts[c] > 0) {
+      const off = c * dim;
+      for (let d = 0; d < dim; d++) centroids[off + d] = sums[off + d] / counts[c];
+    } else {
+      const idx = Math.floor(Math.random() * n);
+      centroids.set(points.subarray(idx * dim, idx * dim + dim), c * dim);
+    }
+  }
 }
 function mortonSort(positions, count) {
   let minx = Infinity, miny = Infinity, minz = Infinity;
@@ -6466,11 +6624,12 @@ async function encodeWebP(width, height, pixelData) {
   }
   return new Uint8Array(await blob.arrayBuffer());
 }
-async function serializeSOG(data, coordinateSystem = "blender", options = {}) {
+async function serializeSOG(data, coordinateSystem = "blender", options = {}, onProgress) {
   const { maxSHBands = 3, iterations = 10, paletteSize = 1024 } = options;
   const { count, colors, opacities } = data;
   let { positions, scales, rotations } = data;
   let shCoeffs = data.shCoeffs ? new Float32Array(data.shCoeffs) : void 0;
+  onProgress == null ? void 0 : onProgress("坐标系变换", 0);
   if (coordinateSystem === "blender") {
     positions = new Float32Array(positions);
     scales = new Float32Array(scales);
@@ -6493,6 +6652,7 @@ async function serializeSOG(data, coordinateSystem = "blender", options = {}) {
       }
     }
   }
+  onProgress == null ? void 0 : onProgress("Morton 排序", 0.05);
   const sortOrder = mortonSort(positions, count);
   const sPos = new Float32Array(count * 3);
   const sSca = new Float32Array(count * 3);
@@ -6525,6 +6685,7 @@ async function serializeSOG(data, coordinateSystem = "blender", options = {}) {
   }
   const { width, height } = calcImageDimensions(count);
   const totalPixels = width * height;
+  onProgress == null ? void 0 : onProgress("编码位置", 0.1);
   const logPos = new Float32Array(count * 3);
   for (let i = 0; i < count * 3; i++) logPos[i] = symLog(sPos[i]);
   const mins = [Infinity, Infinity, Infinity];
@@ -6552,6 +6713,7 @@ async function serializeSOG(data, coordinateSystem = "blender", options = {}) {
     meansLData[off + 3] = 255;
     meansUData[off + 3] = 255;
   }
+  onProgress == null ? void 0 : onProgress("编码缩放", 0.2);
   const logScales = new Float32Array(count * 3);
   for (let i = 0; i < count * 3; i++) logScales[i] = Math.log(Math.max(1e-8, sSca[i]));
   const scaleCodebook = buildCodebook(logScales, 256);
@@ -6564,6 +6726,7 @@ async function serializeSOG(data, coordinateSystem = "blender", options = {}) {
     scalesData[off + 2] = findNearest(scaleCodebook, logScales[i * 3 + 2]);
     scalesData[off + 3] = 255;
   }
+  onProgress == null ? void 0 : onProgress("编码四元数", 0.3);
   const quatsData = new Uint8ClampedArray(totalPixels * 4);
   quatsData.fill(255);
   const SQRT2 = Math.SQRT2;
@@ -6616,6 +6779,7 @@ async function serializeSOG(data, coordinateSystem = "blender", options = {}) {
     quatsData[off + 2] = Math.round(Math.max(0, Math.min(255, (r2 / SQRT2 + 0.5) * 255)));
     quatsData[off + 3] = largest + 252;
   }
+  onProgress == null ? void 0 : onProgress("编码颜色", 0.4);
   const dcValues = new Float32Array(count * 3);
   for (let i = 0; i < count; i++) {
     dcValues[i * 3] = (sCol[i * 3] - 0.5) / SH_C0;
@@ -6644,18 +6808,28 @@ async function serializeSOG(data, coordinateSystem = "blender", options = {}) {
     encodeWebP(width, height, sh0Data)
   ];
   if (actualBands > 0 && sSH) {
-    const vectors = new Array(count);
+    onProgress == null ? void 0 : onProgress("压缩球谐系数 (GPU)", 0.5);
+    const pointsFlat = new Float32Array(count * shDim);
     for (let i = 0; i < count; i++) {
-      const v = new Float32Array(shDim);
       const b = i * 45;
-      for (let j = 0; j < shDim; j++) v[j] = sSH[b + j];
-      vectors[i] = v;
+      for (let j = 0; j < shDim; j++) pointsFlat[i * shDim + j] = sSH[b + j];
     }
     const effectivePaletteSize = Math.min(paletteSize, count);
-    const { centroids, labels } = kmeansVectors(vectors, shDim, effectivePaletteSize, iterations);
-    const allCoeffValues = new Float32Array(centroids.length * shDim);
-    for (let c = 0; c < centroids.length; c++) {
-      for (let d = 0; d < shDim; d++) allCoeffValues[c * shDim + d] = centroids[c][d];
+    const { centroids, labels } = await kmeansGpu(
+      pointsFlat,
+      count,
+      shDim,
+      effectivePaletteSize,
+      iterations,
+      (iter, total) => {
+        const p = 0.5 + 0.35 * (iter / total);
+        onProgress == null ? void 0 : onProgress(`球谐聚类 ${iter}/${total}`, p);
+      }
+    );
+    onProgress == null ? void 0 : onProgress("量化球谐 codebook", 0.88);
+    const allCoeffValues = new Float32Array(effectivePaletteSize * shDim);
+    for (let c = 0; c < effectivePaletteSize; c++) {
+      for (let d = 0; d < shDim; d++) allCoeffValues[c * shDim + d] = centroids[c * shDim + d];
     }
     const shCodebook = buildCodebook(allCoeffValues, 256);
     const centW = 64 * numCoeffs;
@@ -6663,14 +6837,14 @@ async function serializeSOG(data, coordinateSystem = "blender", options = {}) {
     const centData = new Uint8ClampedArray(centW * centH * 4);
     centData.fill(255);
     for (let n = 0; n < effectivePaletteSize; n++) {
-      const entry = centroids[n];
       for (let c = 0; c < numCoeffs; c++) {
         const u = n % 64 * numCoeffs + c;
         const v = Math.floor(n / 64);
         const off = (v * centW + u) * 4;
-        centData[off] = findNearest(shCodebook, entry[c * 3]);
-        centData[off + 1] = findNearest(shCodebook, entry[c * 3 + 1]);
-        centData[off + 2] = findNearest(shCodebook, entry[c * 3 + 2]);
+        const cOff = n * shDim;
+        centData[off] = findNearest(shCodebook, centroids[cOff + c * 3]);
+        centData[off + 1] = findNearest(shCodebook, centroids[cOff + c * 3 + 1]);
+        centData[off + 2] = findNearest(shCodebook, centroids[cOff + c * 3 + 2]);
         centData[off + 3] = 255;
       }
     }
@@ -6694,6 +6868,7 @@ async function serializeSOG(data, coordinateSystem = "blender", options = {}) {
       files: ["sh_centroids.webp", "sh_labels.webp"]
     };
   }
+  onProgress == null ? void 0 : onProgress("编码 WebP 图像", 0.92);
   const webps = await Promise.all(webpPromises);
   const [meansLWebP, meansUWebP, scalesWebP, quatsWebP, sh0WebP] = webps;
   const meta = {
@@ -6719,7 +6894,9 @@ async function serializeSOG(data, coordinateSystem = "blender", options = {}) {
     zipEntries["sh_centroids.webp"] = webps[5];
     zipEntries["sh_labels.webp"] = webps[6];
   }
+  onProgress == null ? void 0 : onProgress("打包 ZIP", 0.98);
   const zipData = zipSync(zipEntries, { level: 0 });
+  onProgress == null ? void 0 : onProgress("完成", 1);
   return zipData.buffer;
 }
 const WORKGROUP_SIZE$1 = 256;