Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add initial support for 3D Gaussian Splatting #244

Merged
merged 29 commits into from
Oct 13, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
29 commits
Select commit Hold shift + click to select a range
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
52 changes: 52 additions & 0 deletions src/ply-parser.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,52 @@
import {
AppBase,
AssetRegistry,
ContainerHandler,
GraphicsDevice
} from 'playcanvas';
import { SplatResource } from './splat-resource';
import { readPly } from './ply-reader';

// filter out element data we're not going to use
const elements = [
'x', 'y', 'z',
'red', 'green', 'blue',
'opacity',
'f_dc_0', 'f_dc_1', 'f_dc_2',
'scale_0', 'scale_1', 'scale_2',
'rot_0', 'rot_1', 'rot_2', 'rot_3'
];

class PlyParser {
device: GraphicsDevice;
assets: AssetRegistry;
maxRetries: number;

constructor(device: GraphicsDevice, assets: AssetRegistry, maxRetries: number) {
this.device = device;
this.assets = assets;
this.maxRetries = maxRetries;
}

async load(url: any, callback: (err: string, resource: SplatResource) => void) {
const response = await fetch(url.load);
readPly(response.body.getReader(), new Set(elements))
.then((response) => {
callback(null, new SplatResource(this.device, response));
})
.catch((err) => {
callback(err, null);
});
}

open(url: string, data: any) {
return data;
}
}

const registerPlyParser = (app: AppBase) => {
const containerHandler = app.loader.getHandler('container') as ContainerHandler;
containerHandler.parsers.ply = new PlyParser(app.graphicsDevice, app.assets, app.loader.maxRetries);
};

export { registerPlyParser };
211 changes: 211 additions & 0 deletions src/ply-reader.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,211 @@
const magicBytes = new Uint8Array([112, 108, 121, 10]); // ply\n
const endHeaderBytes = new Uint8Array([10, 101, 110, 100, 95, 104, 101, 97, 100, 101, 114, 10]); // \nend_header\n

const dataTypeMap = new Map<string, any>([
['char', Int8Array],
['uchar', Uint8Array],
['short', Int16Array],
['ushort', Uint16Array],
['int', Int32Array],
['uint', Uint32Array],
['float', Float32Array],
['double', Float64Array]
]);

type PlyProperty = {
type: string;
name: string;
storage: any;
byteSize: number;
}

type PlyElement = {
name: string;
count: number;
properties: PlyProperty[];
}

// asynchronously read a ply file data
const readPly = async (reader: ReadableStreamDefaultReader, propertyFilter: Set<string> = null): Promise<PlyElement[]> => {
const concat = (a: Uint8Array, b: Uint8Array) => {
const c = new Uint8Array(a.byteLength + b.byteLength);
c.set(a);
c.set(b, a.byteLength);
return c;
};

const find = (buf: Uint8Array, search: Uint8Array) => {
const endIndex = buf.length - search.length;
let i, j;
for (i = 0; i < endIndex; ++i) {
for (j = 0; j < search.length; ++j) {
if (buf[i + j] !== search[j]) {
break;
}
}
if (j === search.length) {
return i;
}
}
return -1;
};

const startsWith = (a: Uint8Array, b: Uint8Array) => {
if (a.length < b.length) {
return false;
}

for (let i = 0; i < b.length; ++i) {
if (a[i] !== b[i]) {
return false;
}
}

return true;
};

let buf: Uint8Array;
let endHeaderIndex: number;

while (true) {
// get the next chunk
/* eslint-disable no-await-in-loop */
const { value, done } = await reader.read();

if (done) {
throw new Error('Stream finished before end of header');
}

// combine new chunk with the previous
buf = buf ? concat(buf, value) : value;

// check magic bytes
if (buf.length >= magicBytes.length && !startsWith(buf, magicBytes)) {
throw new Error('Invalid ply header');
}

// check if we can find the end-of-header marker
endHeaderIndex = find(buf, endHeaderBytes);

if (endHeaderIndex !== -1) {
break;
}
}

// decode buffer header text
const headerText = new TextDecoder('ascii').decode(buf.slice(0, endHeaderIndex));

// split into lines and remove comments
const headerLines = headerText.split('\n')
.filter(line => !line.startsWith('comment '));

// decode header and allocate data storage
const elements: any[] = [];
for (let i = 1; i < headerLines.length; ++i) {
const words = headerLines[i].split(' ');

switch (words[0]) {
case 'format':
if (words[1] !== 'binary_little_endian') {
throw new Error('Unsupported ply format');
}
break;
case 'element':
elements.push({
name: words[1],
count: parseInt(words[2], 10),
properties: []
});
break;
case 'property': {
if (!dataTypeMap.has(words[1])) {
throw new Error(`Unrecognized property data type '${words[1]}' in ply header`);
}
const element = elements[elements.length - 1];
const storageType = dataTypeMap.get(words[1]);
const storage = (!propertyFilter || propertyFilter.has(words[2])) ? new storageType(element.count) : null;
element.properties.push({
type: words[1],
name: words[2],
storage: storage,
byteSize: storageType.BYTES_PER_ELEMENT
});
break;
}
default:
throw new Error(`Unrecognized header value '${words[0]}' in ply header`);
}
}

// read data
let readIndex = endHeaderIndex + endHeaderBytes.length;
let remaining = buf.length - readIndex;
let dataView = new DataView(buf.buffer);

for (let i = 0; i < elements.length; ++i) {
const element = elements[i];

for (let e = 0; e < element.count; ++e) {
for (let j = 0; j < element.properties.length; ++j) {
const property = element.properties[j];

// if we've run out of data, load the next chunk
while (remaining < property.byteSize) {
const { value, done } = await reader.read();

if (done) {
throw new Error('Stream finished before end of data');
}

// create buffer with left-over data from previous chunk and the new data
const tmp = new Uint8Array(remaining + value.byteLength);
tmp.set(buf.slice(readIndex));
tmp.set(value, remaining);

buf = tmp;
dataView = new DataView(buf.buffer);
readIndex = 0;
remaining = buf.length;
}

if (property.storage) {
switch (property.type) {
case 'char':
property.storage[e] = dataView.getInt8(readIndex);
break;
case 'uchar':
property.storage[e] = dataView.getUint8(readIndex);
break;
case 'short':
property.storage[e] = dataView.getInt16(readIndex, true);
break;
case 'ushort':
property.storage[e] = dataView.getUint16(readIndex, true);
break;
case 'int':
property.storage[e] = dataView.getInt32(readIndex, true);
break;
case 'uint':
property.storage[e] = dataView.getUint32(readIndex, true);
break;
case 'float':
property.storage[e] = dataView.getFloat32(readIndex, true);
break;
case 'double':
property.storage[e] = dataView.getFloat64(readIndex, true);
break;
}
}

readIndex += property.byteSize;
remaining -= property.byteSize;
}
}
}

// console.log(elements);

return elements;
};

export { readPly, PlyProperty, PlyElement };
112 changes: 112 additions & 0 deletions src/sort-worker.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,112 @@
type Vec3 = {
x: number,
y: number,
z: number
};

// sort blind set of data
function SortWorker() {
const epsilon = 0.0001;

let data: Float32Array;
let stride: number;
let cameraPosition: Vec3;
let cameraDirection: Vec3;

const lastCameraPosition = { x: 0, y: 0, z: 0 };
const lastCameraDirection = { x: 0, y: 0, z: 0 };

let distanceBuffer: Float32Array;
let orderBuffer: Uint32Array;
let target: Float32Array;

const update = () => {
if (!data || !stride || !cameraPosition || !cameraDirection) return;

// early out if camera hasn't moved
if (Math.abs(cameraPosition.x - lastCameraPosition.x) < epsilon &&
Math.abs(cameraPosition.y - lastCameraPosition.y) < epsilon &&
Math.abs(cameraPosition.z - lastCameraPosition.z) < epsilon &&
Math.abs(cameraDirection.x - lastCameraDirection.x) < epsilon &&
Math.abs(cameraDirection.y - lastCameraDirection.y) < epsilon &&
Math.abs(cameraDirection.z - lastCameraDirection.z) < epsilon) {
return;
}

const numVertices = data.length / stride;

// create distance buffer
if (!distanceBuffer || distanceBuffer.length !== numVertices) {
distanceBuffer = new Float32Array(numVertices);
orderBuffer = new Uint32Array(numVertices);
target = new Float32Array(numVertices * stride);
}

// store
lastCameraPosition.x = cameraPosition.x;
lastCameraPosition.y = cameraPosition.y;
lastCameraPosition.z = cameraPosition.z;
lastCameraDirection.x = cameraDirection.x;
lastCameraDirection.y = cameraDirection.y;
lastCameraDirection.z = cameraDirection.z;

const px = cameraPosition.x;
const py = cameraPosition.y;
const pz = cameraPosition.z;
const dx = cameraDirection.x;
const dy = cameraDirection.y;
const dz = cameraDirection.z;

// generate per vertex distance to camera
for (let i = 0; i < numVertices; ++i) {
distanceBuffer[i] =
(data[i * stride + 0] - px) * dx +
(data[i * stride + 1] - py) * dy +
(data[i * stride + 2] - pz) * dz;
orderBuffer[i] = i;
}

// sort indices
orderBuffer.sort((a, b) => distanceBuffer[a] - distanceBuffer[b]);

const orderChanged = orderBuffer.some((v, i) => v !== i);

if (orderChanged) {
// order the splat data
for (let i = 0; i < numVertices; ++i) {
const ti = i * stride;
const si = orderBuffer[i] * stride;
for (let j = 0; j < stride; ++j) {
target[ti + j] = data[si + j];
}
}

// swap
const tmp = data;
data = target;
target = tmp;

// send results
self.postMessage({
data: data.buffer
}, [data.buffer]);

data = null;
}
};

self.onmessage = (message: any) => {
if (message.data.data) {
data = new Float32Array(message.data.data);
}
if (message.data.stride) {
stride = message.data.stride;
}
if (message.data.cameraPosition) cameraPosition = message.data.cameraPosition;
if (message.data.cameraDirection) cameraDirection = message.data.cameraDirection;

update();
};
}

export { SortWorker };
Loading