spark/src/antisplat.ts at main · sparkjsdev/spark · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
import type { SplatEncoding } from "./PackedSplats";
import { computeMaxSplats, setPackedSplat } from "./utils";

export function decodeAntiSplat(
  fileBytes: Uint8Array,
  initNumSplats: (numSplats: number) => void,
  splatCallback: (
    index: number,
    x: number,
    y: number,
    z: number,
    scaleX: number,
    scaleY: number,
    scaleZ: number,
    quatX: number,
    quatY: number,
    quatZ: number,
    quatW: number,
    opacity: number,
    r: number,
    g: number,
    b: number,
  ) => void,
) {
  const numSplats = Math.floor(fileBytes.length / 32); // 32 bytes per splat
  if (numSplats * 32 !== fileBytes.length) {
    throw new Error("Invalid .splat file size");
  }
  initNumSplats(numSplats);

  const f32 = new Float32Array(fileBytes.buffer);
  for (let i = 0; i < numSplats; ++i) {
    const i32 = i * 32;
    const i8 = i * 8;
    const x = f32[i8 + 0];
    const y = f32[i8 + 1];
    const z = f32[i8 + 2];
    const scaleX = f32[i8 + 3];
    const scaleY = f32[i8 + 4];
    const scaleZ = f32[i8 + 5];
    const r = fileBytes[i32 + 24] / 255;
    const g = fileBytes[i32 + 25] / 255;
    const b = fileBytes[i32 + 26] / 255;
    const opacity = fileBytes[i32 + 27] / 255;
    const quatW = (fileBytes[i32 + 28] - 128) / 128;
    const quatX = (fileBytes[i32 + 29] - 128) / 128;
    const quatY = (fileBytes[i32 + 30] - 128) / 128;
    const quatZ = (fileBytes[i32 + 31] - 128) / 128;
    splatCallback(
      i,
      x,
      y,
      z,
      scaleX,
      scaleY,
      scaleZ,
      quatX,
      quatY,
      quatZ,
      quatW,
      opacity,
      r,
      g,
      b,
    );
  }
}

export function unpackAntiSplat(
  fileBytes: Uint8Array,
  splatEncoding: SplatEncoding,
): {
  packedArray: Uint32Array;
  numSplats: number;
} {
  let numSplats = 0;
  let maxSplats = 0;
  let packedArray = new Uint32Array(0);
  decodeAntiSplat(
    fileBytes,
    (cbNumSplats) => {
      numSplats = cbNumSplats;
      maxSplats = computeMaxSplats(numSplats);
      packedArray = new Uint32Array(maxSplats * 4);
    },
    (
      index,
      x,
      y,
      z,
      scaleX,
      scaleY,
      scaleZ,
      quatX,
      quatY,
      quatZ,
      quatW,
      opacity,
      r,
      g,
      b,
    ) => {
      setPackedSplat(
        packedArray,
        index,
        x,
        y,
        z,
        scaleX,
        scaleY,
        scaleZ,
        quatX,
        quatY,
        quatZ,
        quatW,
        opacity,
        r,
        g,
        b,
        splatEncoding,
      );
    },
  );
  return { packedArray, numSplats };
}