Stream during save (#222)

src/asset-loader.ts

@@ -1,8 +1,7 @@
 import { Asset, AssetRegistry, GraphicsDevice, GSplatData, GSplatResource, TEXTURETYPE_RGBP } from 'playcanvas';
 import { Splat } from './splat';
 import { Env } from './env';
-
-import { startSpinner, stopSpinner } from './ui/spinner';
+import { Events } from './events';
 
 interface ModelLoadRequest {
     url?: string;
@@ -94,17 +93,19 @@ const deserializeFromSSplat = (data: ArrayBufferLike) => {
 class AssetLoader {
     device: GraphicsDevice;
     registry: AssetRegistry;
+    events: Events;
     defaultAnisotropy: number;
     loadAllData = true;
 
-    constructor(device: GraphicsDevice, registry: AssetRegistry, defaultAnisotropy?: number) {
+    constructor(device: GraphicsDevice, registry: AssetRegistry, events: Events, defaultAnisotropy?: number) {
         this.device = device;
         this.registry = registry;
+        this.events = events;
         this.defaultAnisotropy = defaultAnisotropy || 1;
     }
 
     loadPly(loadRequest: ModelLoadRequest) {
-        startSpinner();
+        this.events.fire('startSpinner');
 
         return new Promise<Splat>((resolve, reject) => {
             const asset = new Asset(
@@ -154,12 +155,12 @@ class AssetLoader {
             this.registry.add(asset);
             this.registry.load(asset);
         }).finally(() => {
-            stopSpinner();
+            this.events.fire('stopSpinner');
         });
     }
 
     loadSplat(loadRequest: ModelLoadRequest) {
-        startSpinner();
+        this.events.fire('startSpinner');
 
         return new Promise<Splat>((resolve, reject) => {
             fetch(loadRequest.url || loadRequest.filename)
@@ -184,7 +185,7 @@ class AssetLoader {
                 reject('Failed to load splat data');
             });
         }).finally(() => {
-            stopSpinner();
+            this.events.fire('stopSpinner');
         });
     }
 

src/file-handler.ts

@@ -2,8 +2,7 @@ import { path, Vec3 } from 'playcanvas';
 import { Scene } from './scene';
 import { Events } from './events';
 import { CreateDropHandler } from './drop-handler';
-import { convertPly, convertPlyCompressed, convertSplat } from './splat-convert';
-import { startSpinner, stopSpinner } from './ui/spinner';
+import { WriteFunc, serializePly, serializePlyCompressed, serializeSplat } from './splat-serialize';
 import { ElementType } from './element';
 import { Splat } from './splat';
 import { localize } from './ui/localization';
@@ -47,7 +46,7 @@ let fileHandle: FileSystemFileHandle = null;
 const vec = new Vec3();
 
 // download the data to the given filename
-const download = (filename: string, data: ArrayBuffer) => {
+const download = (filename: string, data: Uint8Array) => {
     const blob = new Blob([data], { type: "octet/stream" });
     const url = window.URL.createObjectURL(blob);
 
@@ -81,14 +80,6 @@ const sendToRemoteStorage = async (filename: string, data: ArrayBuffer, remoteSt
     });
 };
 
-// write the data to file
-const writeToFile = async (stream: FileSystemWritableFileStream, data: ArrayBuffer) => {
-    await stream.seek(0);
-    await stream.write(data);
-    await stream.truncate(data.byteLength);
-    await stream.close();
-};
-
 const loadCameraPoses = async (url: string, filename: string, events: Events) => {
     const response = await fetch(url);
     const json = await response.json();
@@ -331,39 +322,61 @@ const initFileHandler = async (scene: Scene, events: Events, dropTarget: HTMLEle
         }
     });
 
+    const writeScene = async (type: ExportType, writeFunc: WriteFunc) => {
+        const splats = getSplats();
+
+        switch (type) {
+            case 'ply':
+                await serializePly(splats, writeFunc);
+                break;
+            case 'compressed-ply':
+                await serializePlyCompressed(splats, writeFunc);
+                return;
+            case 'splat':
+                await serializeSplat(splats, writeFunc);
+                return;
+        }
+    };
+
     events.function('scene.write', async (options: SceneWriteOptions) => {
-        startSpinner();
+        events.fire('startSpinner');
 
         try {
-            const splats = getSplats();
-
             // setTimeout so spinner has a chance to activate
             await new Promise<void>((resolve) => {
                 setTimeout(resolve);
             });
 
-            const data = (() => {
-                switch (options.type) {
-                    case 'ply':
-                        return convertPly(splats);
-                    case 'compressed-ply':
-                        return convertPlyCompressed(splats);
-                    case 'splat':
-                        return convertSplat(splats);
-                    default:
-                        return null;
-                }
-            })();
-
-            if (options.stream) {
-                // write to stream
-                await writeToFile(options.stream, data);
-            } else if (remoteStorageDetails) {
-                // write data to remote storage
-                await sendToRemoteStorage(options.filename, data, remoteStorageDetails);
+            const { stream } = options;
+
+            if (stream) {
+                // writer must keep track of written bytes because JS streams don't
+                let cursor = 0;
+                const writeFunc = (data: Uint8Array) => {
+                    cursor += data.byteLength;
+                    return stream.write(data);
+                };
+
+                await stream.seek(0);
+                await writeScene(options.type, writeFunc);
+                await stream.truncate(cursor);
+                await stream.close();
             } else if (options.filename) {
-                // download file to local machine
-                download(options.filename, data);
+                // (safari): concatenate data into single buffer
+                let cursor = 0;
+                let data = new Uint8Array(1024 * 1024);
+
+                const writeFunc = (chunk: Uint8Array) => {
+                    if (cursor + chunk.byteLength > data.byteLength) {
+                        const newData = new Uint8Array(data.byteLength * 2);
+                        newData.set(data);
+                        data = newData;
+                    }
+                    data.set(chunk, cursor);
+                    cursor += chunk.byteLength;
+                };
+                await writeScene(options.type, writeFunc);
+                download(options.filename, new Uint8Array(data.buffer, 0, cursor));
             }
         } catch (err) {
             events.invoke('showPopup', {
@@ -372,7 +385,7 @@ const initFileHandler = async (scene: Scene, events: Events, dropTarget: HTMLEle
                 message: `${err.message ?? err} while saving file`
             });
         } finally {
-            stopSpinner();
+            events.fire('stopSpinner');
         }
     });
 };

src/scene.ts

@@ -175,7 +175,7 @@ class Scene {
         layers.push(this.gizmoLayer);
 
         this.dataProcessor = new DataProcessor(this.app.graphicsDevice);
-        this.assetLoader = new AssetLoader(graphicsDevice, this.app.assets, this.app.graphicsDevice.maxAnisotropy);
+        this.assetLoader = new AssetLoader(graphicsDevice, this.app.assets, events, this.app.graphicsDevice.maxAnisotropy);
 
         // create root entities
         this.contentRoot = new Entity('contentRoot');

src/shaders/splat-shader.ts

@@ -87,7 +87,7 @@ void main(void)
 
     texCoord = vertex_position.xy * 0.5;
 
-    #if FORWARD_PASS
+    #ifdef FORWARD_PASS
         // get color
         color = texelFetch(splatColor, splatUV, 0);
 

src/splat-convert.ts → src/splat-serialize.ts

@@ -10,6 +10,9 @@ import { Splat } from './splat';
 import { SHRotation } from './sh-utils';
 import { version } from '../package.json';
 
+// async function for writing data
+type WriteFunc = (data: Uint8Array) => void;
+
 const generatedByString =  `Generated by SuperSplat ${version}`;
 
 const countTotalSplats = (splats: Splat[]) => {
@@ -45,17 +48,18 @@ class SplatTransformCache {
     getSHRot: (index: number) => SHRotation;
 
     constructor(splat: Splat) {
-        const transforms = new Map<number, { mat: Mat4, rot: Quat, scale: Vec3, shRot: SHRotation }>();
+        const transforms = new Map<number, { transformIndex: number, mat: Mat4, rot: Quat, scale: Vec3, shRot: SHRotation }>();
         const indices = splat.transformTexture.getSource() as unknown as Uint32Array;
         const tmpMat = new Mat4();
         const tmpMat3 = new Mat3();
         const tmpQuat = new Quat();
 
         const getTransform = (index: number) => {
-            let result = transforms.get(index);
+            const transformIndex = indices?.[index] ?? 0;
+            let result = transforms.get(transformIndex);
             if (!result) {
-                result = { mat: null, rot: null, scale: null, shRot: null };
-                transforms.set(index, result);
+                result = { transformIndex, mat: null, rot: null, scale: null, shRot: null };
+                transforms.set(transformIndex, result);
             }
             return result;
         };
@@ -71,9 +75,8 @@ class SplatTransformCache {
                 mat.mul2(mat, splat.entity.getWorldTransform());
 
                 // combine with transform palette matrix
-                const transformIndex = indices?.[index] ?? 0;
-                if (transformIndex > 0) {
-                    splat.transformPalette.getTransform(transformIndex, tmpMat);
+                if (transform.transformIndex > 0) {
+                    splat.transformPalette.getTransform(transform.transformIndex, tmpMat);
                     mat.mul2(mat, tmpMat);
                 }
 
@@ -122,7 +125,7 @@ class SplatTransformCache {
 const v = new Vec3();
 const q = new Quat();
 
-const convertPly = (splats: Splat[]) => {
+const serializePly = async (splats: Splat[], write: WriteFunc) => {
     // count the number of non-deleted splats
     const totalSplats = countTotalSplats(splats);
 
@@ -150,19 +153,18 @@ const convertPly = (splats: Splat[]) => {
          ``
     ].flat().join('\n');
 
-    const header = (new TextEncoder()).encode(headerText);
-    const result = new Uint8Array(header.byteLength + totalSplats * propNames.length * 4);
-    const dataView = new DataView(result.buffer);
+    // write encoded header
+    await write((new TextEncoder()).encode(headerText));
 
     // construct an object for holding a single splat's properties
     const splat = propNames.reduce((acc: any, name) => {
         acc[name] = 0;
         return acc;
     }, {});
 
-    result.set(header);
-
-    let offset = header.byteLength;
+    const buf = new Uint8Array(1024 * propNames.length * 4);
+    const dataView = new DataView(buf.buffer);
+    let offset = 0;
 
     for (let e = 0; e < splats.length; ++e) {
         const splatData = splats[e].splatData;
@@ -231,10 +233,19 @@ const convertPly = (splats: Splat[]) => {
                 dataView.setFloat32(offset, splat[propNames[j]], true);
                 offset += 4;
             }
+
+            // buffer is full, write it to the output stream
+            if (offset === buf.byteLength) {
+                await write(buf);
+                offset = 0;
+            }
         }
     }
 
-    return result;
+    // write the last (most likely partially filled) buf
+    if (offset > 0) {
+        await write(new Uint8Array(buf.buffer, 0, offset));
+    }
 };
 
 interface CompressedIndex {
@@ -517,7 +528,7 @@ const sortSplats = (splats: Splat[], indices: CompressedIndex[]) => {
     indices.sort((a, b) => morton[a.globalIndex] - morton[b.globalIndex]);
 };
 
-const convertPlyCompressed = (splats: Splat[]) => {
+const serializePlyCompressed = async (splats: Splat[], write: WriteFunc) => {
     const chunkProps = [
         'min_x', 'min_y', 'min_z',
         'max_x', 'max_y', 'max_z',
@@ -634,10 +645,10 @@ const convertPlyCompressed = (splats: Splat[]) => {
         }
     }
 
-    return result;
+    await write(result);
 };
 
-const convertSplat = (splats: Splat[]) => {
+const serializeSplat = async (splats: Splat[], write: WriteFunc) => {
     const totalSplats = countTotalSplats(splats);
 
     // position.xyz: float32, scale.xyz: float32, color.rgba: uint8, quaternion.ijkl: uint8
@@ -706,7 +717,12 @@ const convertSplat = (splats: Splat[]) => {
         }
     }
 
-    return result;
+    await write(result);
 };
 
-export { convertPly, convertPlyCompressed, convertSplat };
+export {
+    WriteFunc,
+    serializePly,
+    serializePlyCompressed,
+    serializeSplat
+};

src/ui/editor.ts

@@ -12,8 +12,9 @@ import { RightToolbar } from './right-toolbar';
 import { ModeToggle } from './mode-toggle';
 import { Tooltips } from './tooltips';
 import { ShortcutsPopup } from './shortcuts-popup';
-import { localizeInit } from './localization';
+import { Spinner } from './spinner';
 
+import { localizeInit } from './localization';
 import { version } from '../../package.json';
 import logo from './playcanvas-logo.png';
 
@@ -153,6 +154,20 @@ class EditorUI {
             return this.popup.show(options);
         });
 
+        // spinner
+
+        const spinner = new Spinner();
+
+        topContainer.append(spinner);
+
+        events.on('startSpinner', () => {
+            spinner.hidden = false;
+        });
+
+        events.on('stopSpinner', () => {
+            spinner.hidden = true;
+        });
+
         // initialize canvas to correct size before creating graphics device etc
         const pixelRatio = window.devicePixelRatio;
         canvas.width = Math.ceil(canvasContainer.dom.offsetWidth * pixelRatio);