playcanvas
Version:
Open-source WebGL/WebGPU 3D engine for the web
113 lines (112 loc) • 4.17 kB
JavaScript
import { GSplatData } from "../../../../scene/gsplat/gsplat-data.js";
import { GSplatResource } from "../../../../scene/gsplat/gsplat-resource.js";
import { GltfAccessor } from "../gltf-accessor.js";
const extensionName = "KHR_gaussian_splatting";
const shDegreeCoefCounts = [3, 5, 7];
const shBandCoefCounts = { 1: 3, 2: 8, 3: 15 };
const hasGSplatExtension = (primitive) => {
return !!primitive?.extensions?.[extensionName];
};
const createGSplatData = (primitive, accessors, bufferViews) => {
const extensionData = primitive.extensions[extensionName];
const attributes = primitive.attributes;
const numSplats = accessors[attributes.POSITION]?.count ?? 0;
const readAttribute = (name) => {
const accessor = accessors[attributes[name]];
if (!accessor || accessor.count !== numSplats) {
return null;
}
return GltfAccessor.getDataFloat32(accessor, bufferViews);
};
const positions = readAttribute("POSITION");
const rotations = readAttribute(`${extensionName}:ROTATION`);
const scales = readAttribute(`${extensionName}:SCALE`);
const opacities = readAttribute(`${extensionName}:OPACITY`);
const sh0 = readAttribute(`${extensionName}:SH_DEGREE_0_COEF_0`);
if (!numSplats || !positions || !rotations || !scales || !opacities || !sh0) {
return null;
}
const properties = [];
const addProp = (name, storage) => {
properties.push({ type: "float", name, storage, byteSize: 4 });
};
addProp("x", GltfAccessor.extractComponent(positions, 3, 0, numSplats));
addProp("y", GltfAccessor.extractComponent(positions, 3, 1, numSplats));
addProp("z", GltfAccessor.extractComponent(positions, 3, 2, numSplats));
addProp("rot_0", GltfAccessor.extractComponent(rotations, 4, 3, numSplats));
addProp("rot_1", GltfAccessor.extractComponent(rotations, 4, 0, numSplats));
addProp("rot_2", GltfAccessor.extractComponent(rotations, 4, 1, numSplats));
addProp("rot_3", GltfAccessor.extractComponent(rotations, 4, 2, numSplats));
addProp("scale_0", GltfAccessor.extractComponent(scales, 3, 0, numSplats));
addProp("scale_1", GltfAccessor.extractComponent(scales, 3, 1, numSplats));
addProp("scale_2", GltfAccessor.extractComponent(scales, 3, 2, numSplats));
addProp("opacity", GltfAccessor.extractComponent(opacities, 1, 0, numSplats));
addProp("f_dc_0", GltfAccessor.extractComponent(sh0, 3, 0, numSplats));
addProp("f_dc_1", GltfAccessor.extractComponent(sh0, 3, 1, numSplats));
addProp("f_dc_2", GltfAccessor.extractComponent(sh0, 3, 2, numSplats));
let bands = 0;
for (let d = 1; d <= 3; d++) {
let complete = true;
for (let c = 0; c < shDegreeCoefCounts[d - 1]; c++) {
if (attributes[`${extensionName}:SH_DEGREE_${d}_COEF_${c}`] === void 0) {
complete = false;
break;
}
}
if (!complete) {
break;
}
bands = d;
}
if (bands > 0) {
const numCoefs = shBandCoefCounts[bands];
const rest = new Array(numCoefs * 3);
let k = 0;
for (let d = 1; d <= bands; d++) {
for (let c = 0; c < shDegreeCoefCounts[d - 1]; c++) {
const name = `${extensionName}:SH_DEGREE_${d}_COEF_${c}`;
const source = readAttribute(name);
if (!source) {
return null;
}
rest[k] = GltfAccessor.extractComponent(source, 3, 0, numSplats);
rest[numCoefs + k] = GltfAccessor.extractComponent(source, 3, 1, numSplats);
rest[numCoefs * 2 + k] = GltfAccessor.extractComponent(source, 3, 2, numSplats);
k++;
}
}
for (let i = 0; i < numCoefs * 3; i++) {
addProp(`f_rest_${i}`, rest[i]);
}
}
const data = new GSplatData([{
name: "vertex",
count: numSplats,
properties
}]);
data.activated = true;
return data;
};
const createGSplats = (device, gltf, bufferViews) => {
if (!gltf.hasOwnProperty("meshes")) {
return [];
}
return gltf.meshes.map((gltfMesh) => {
let resources = null;
gltfMesh.primitives.forEach((primitive) => {
if (hasGSplatExtension(primitive)) {
const gsplatData = createGSplatData(primitive, gltf.accessors, bufferViews);
if (gsplatData) {
gsplatData.reorderData();
if (!resources) resources = [];
resources.push(new GSplatResource(device, gsplatData));
}
}
});
return resources;
});
};
export {
createGSplats,
hasGSplatExtension
};