Merge pull request #177 from allen-cell-animated/feature/loader-worker

Feature/loader worker

public/index.ts

@@ -1,5 +1,5 @@
 import "regenerator-runtime/runtime";
-import { Vector2, Vector3, Vector4 } from "three";
+import { Vector2, Vector3 } from "three";
 import * as dat from "dat.gui";
 
 import {
@@ -9,21 +9,19 @@ import {
   JsonImageInfoLoader,
   View3d,
   Volume,
-  VolumeCache,
   VolumeMaker,
   Light,
   AREA_LIGHT,
   RENDERMODE_PATHTRACE,
   RENDERMODE_RAYMARCH,
   SKY_LIGHT,
   VolumeFileFormat,
-  createVolumeLoader,
-  SubscribableRequestQueue,
 } from "../src";
 // special loader really just for this demo app but lives with the other loaders
 import { OpenCellLoader } from "../src/loaders/OpenCellLoader";
 import { State, TestDataSpec } from "./types";
 import { getDefaultImageInfo } from "../src/Volume";
+import VolumeLoaderContext from "../src/workers/LoadWorkerHandle";
 
 const CACHE_MAX_SIZE = 1_000_000_000;
 const CONCURRENCY_LIMIT = 8;
@@ -88,8 +86,7 @@ const TEST_DATA: Record<string, TestDataSpec> = {
 
 let view3D: View3d;
 
-const volumeCache = new VolumeCache(CACHE_MAX_SIZE);
-const requestQueue = new SubscribableRequestQueue(CONCURRENCY_LIMIT, PREFETCH_CONCURRENCY_LIMIT);
+const loaderContext = new VolumeLoaderContext(CACHE_MAX_SIZE, CONCURRENCY_LIMIT, PREFETCH_CONCURRENCY_LIMIT);
 
 const myState: State = {
   file: "",
@@ -1015,6 +1012,8 @@ async function createLoader(data: TestDataSpec): Promise<IVolumeLoader> {
     return new OpenCellLoader();
   }
 
+  await loaderContext.onOpen();
+
   let path: string | string[] = data.url;
   if (data.type === VolumeFileFormat.JSON) {
     path = [];
@@ -1024,10 +1023,8 @@ async function createLoader(data: TestDataSpec): Promise<IVolumeLoader> {
     }
   }
 
-  return await createVolumeLoader(path, {
-    cache: volumeCache,
-    queue: requestQueue,
-    fetchOptions: { maxPrefetchChunks: MAX_PREFETCH_CHUNKS, maxPrefetchDistance: PREFETCH_DISTANCE },
+  return await loaderContext.createLoader(path, {
+    fetchOptions: { maxPrefetchDistance: PREFETCH_DISTANCE, maxPrefetchChunks: MAX_PREFETCH_CHUNKS },
   });
 }
 

src/index.ts

@@ -12,6 +12,7 @@ import { LoadSpec } from "./loaders/IVolumeLoader";
 import { OMEZarrLoader } from "./loaders/OmeZarrLoader";
 import { JsonImageInfoLoader } from "./loaders/JsonImageInfoLoader";
 import { TiffLoader } from "./loaders/TiffLoader";
+import VolumeLoaderContext from "./workers/LoadWorkerHandle";
 
 import { Light, AREA_LIGHT, SKY_LIGHT } from "./Light";
 
@@ -20,6 +21,7 @@ export type { ControlPoint, Lut } from "./Histogram";
 export type { CreateLoaderOptions } from "./loaders";
 export type { IVolumeLoader, PerChannelCallback } from "./loaders/IVolumeLoader";
 export type { ZarrLoaderFetchOptions } from "./loaders/OmeZarrLoader";
+export type { WorkerLoader } from "./workers/LoadWorkerHandle";
 export {
   Histogram,
   View3d,
@@ -32,6 +34,7 @@ export {
   OMEZarrLoader,
   JsonImageInfoLoader,
   TiffLoader,
+  VolumeLoaderContext,
   VolumeFileFormat,
   createVolumeLoader,
   Channel,

src/loaders/IVolumeLoader.ts

@@ -1,6 +1,7 @@
 import { Box3, Vector3 } from "three";
 
-import Volume from "../Volume";
+import Volume, { ImageInfo } from "../Volume";
+import { buildDefaultMetadata } from "./VolumeLoaderUtils";
 
 export class LoadSpec {
   time = 0;
@@ -26,6 +27,11 @@ export class VolumeDims {
   dataType = "uint8";
 }
 
+export type LoadedVolumeInfo = {
+  imageInfo: ImageInfo;
+  loadSpec: LoadSpec;
+};
+
 /**
  * @callback PerChannelCallback
  * @param {string} imageurl
@@ -34,6 +40,8 @@ export class VolumeDims {
  */
 export type PerChannelCallback = (volume: Volume, channelIndex: number) => void;
 
+export type RawChannelDataCallback = (channelIndex: number, data: Uint8Array, atlasDims?: [number, number]) => void;
+
 /**
  * Loads volume data from a source specified by a `LoadSpec`.
  *
@@ -47,10 +55,6 @@ export interface IVolumeLoader {
   /**
    * Create an empty `Volume` from a `LoadSpec`, which must be passed to `loadVolumeData` to begin loading.
    * Optionally pass a callback to respond whenever new channel data is loaded into the volume.
-   *
-   * Loaders are allowed to assume that they will only be called on a single data source, in order to cache
-   * information about that source. Once this method has been called, every subsequent call to it or
-   * `loadVolumeData` should reference the same source.
    */
   createVolume(loadSpec: LoadSpec, onChannelLoaded?: PerChannelCallback): Promise<Volume>;
 
@@ -65,3 +69,64 @@ export interface IVolumeLoader {
   // TODO explicitly passing a `LoadSpec` is now rarely useful. Remove?
   loadVolumeData(volume: Volume, loadSpec?: LoadSpec, onChannelLoaded?: PerChannelCallback): void;
 }
+
+/** Abstract class which allows loaders to accept and return types that are easier to transfer to/from a worker. */
+export abstract class ThreadableVolumeLoader implements IVolumeLoader {
+  /** Unchanged from `IVolumeLoader`. See that interface for details. */
+  abstract loadDims(loadSpec: LoadSpec): Promise<VolumeDims[]>;
+
+  /**
+   * Creates an `ImageInfo` object from a `LoadSpec`, which may be passed to the `Volume` constructor to create an
+   * empty volume that can accept data loaded with the given `LoadSpec`.
+   *
+   * Also returns a new `LoadSpec` that may have been modified from the input `LoadSpec` to reflect the constraints or
+   * abilities of the loader. This new `LoadSpec` should be used when constructing the `Volume`, _not_ the original.
+   */
+  abstract createImageInfo(loadSpec: LoadSpec): Promise<LoadedVolumeInfo>;
+
+  /**
+   * Begins loading per-channel data for the volume specified by `imageInfo` and `loadSpec`.
+   *
+   * Returns a promise that resolves to reflect any modifications to `imageInfo` and/or `loadSpec` that need to be made
+   * based on this load. Actual loaded channel data is passed to `onData` as it is loaded. Depending on the format,
+   * the returned array may be in simple 3d dimension order or reflect a 2d atlas. If the latter, the dimensions of the
+   * atlas are passed as the third argument to `onData`.
+   */
+  abstract loadRawChannelData(
+    imageInfo: ImageInfo,
+    loadSpec: LoadSpec,
+    onData: RawChannelDataCallback
+  ): Promise<Partial<LoadedVolumeInfo>>;
+
+  async createVolume(loadSpec: LoadSpec, onChannelLoaded?: PerChannelCallback): Promise<Volume> {
+    const { imageInfo, loadSpec: adjustedLoadSpec } = await this.createImageInfo(loadSpec);
+    const vol = new Volume(imageInfo, adjustedLoadSpec, this);
+    vol.channelLoadCallback = onChannelLoaded;
+    vol.imageMetadata = buildDefaultMetadata(imageInfo);
+    return vol;
+  }
+
+  async loadVolumeData(
+    volume: Volume,
+    loadSpecOverride?: LoadSpec,
+    onChannelLoaded?: PerChannelCallback
+  ): Promise<void> {
+    const onChannelData: RawChannelDataCallback = (channelIndex, data, atlasDims) => {
+      if (atlasDims) {
+        volume.setChannelDataFromAtlas(channelIndex, data, atlasDims[0], atlasDims[1]);
+      } else {
+        volume.setChannelDataFromVolume(channelIndex, data);
+      }
+      onChannelLoaded?.(volume, channelIndex);
+    };
+
+    const spec = { ...loadSpecOverride, ...volume.loadSpec };
+    const { imageInfo, loadSpec } = await this.loadRawChannelData(volume.imageInfo, spec, onChannelData);
+
+    if (imageInfo) {
+      volume.imageInfo = imageInfo;
+      volume.updateDimensions();
+    }
+    volume.loadSpec = { ...loadSpec, ...spec };
+  }
+}