Merge pull request #163 from allen-cell-animated/fix/load-hires-slice

Load single slices at higher resolution

public/index.ts

@@ -978,6 +978,9 @@ function createTestVolume() {
     timeScale: 1,
     timeUnit: "",
 
+    numMultiscaleLevels: 1,
+    multiscaleLevel: 0,
+
     transform: { translation: new Vector3(0, 0, 0), rotation: new Vector3(0, 0, 0) },
   };
 

src/Volume.ts

@@ -4,12 +4,12 @@ import Channel from "./Channel";
 import Histogram from "./Histogram";
 import { getColorByChannelIndex } from "./constants/colors";
 import { IVolumeLoader, LoadSpec, PerChannelCallback } from "./loaders/IVolumeLoader";
+import { estimateLevelForAtlas } from "./loaders/VolumeLoaderUtils";
 
 export type ImageInfo = Readonly<{
   name: string;
 
   /** XY size of the *original* (not downsampled) volume, in pixels */
-  // If we ever allow downsampling in z, replace with Vector3
   originalSize: Vector3;
   /**
    * XY dimensions of the texture atlas used by `RayMarchedAtlasVolume` and `Atlas2DSlice`, in number of z-slice
@@ -41,6 +41,11 @@ export type ImageInfo = Readonly<{
   /** Symbol of temporal unit used by `timeScale`, e.g. "hr" */
   timeUnit: string;
 
+  /** Number of scale levels available for this volume */
+  numMultiscaleLevels: number;
+  /** The scale level from which this image was loaded, between `0` and `numMultiscaleLevels-1` */
+  multiscaleLevel: number;
+
   transform: {
     /** Translation of the volume from the center of space, in volume voxels */
     translation: Vector3;
@@ -67,6 +72,8 @@ export const getDefaultImageInfo = (): ImageInfo => ({
   times: 1,
   timeScale: 1,
   timeUnit: "",
+  numMultiscaleLevels: 1,
+  multiscaleLevel: 0,
   transform: {
     translation: new Vector3(0, 0, 0),
     rotation: new Vector3(0, 0, 0),
@@ -229,14 +236,29 @@ export default class Volume {
     this.normRegionOffset = subregionOffset.clone().divide(volumeSize);
   }
 
-  updateRequiredData(required: Partial<LoadSpec>, onChannelLoaded?: PerChannelCallback) {
+  /** Call on any state update that may require new data to be loaded (subregion, enabled channels, time, etc.) */
+  async updateRequiredData(required: Partial<LoadSpec>, onChannelLoaded?: PerChannelCallback): Promise<void> {
     this.loadSpecRequired = { ...this.loadSpecRequired, ...required };
+    let noReload =
+      this.loadSpec.time === this.loadSpecRequired.time &&
+      this.loadSpec.subregion.containsBox(this.loadSpecRequired.subregion);
+
+    // An update to `subregion` should trigger a reload when the new subregion is not contained in the old one
+    // OR when the new subregion is smaller than the old one by enough that we can load a higher scale level.
+    if (noReload && !this.loadSpec.subregion.equals(this.loadSpecRequired.subregion)) {
+      const currentScale = this.imageInfo.multiscaleLevel;
+      // `LoadSpec.multiscaleLevel`, if specified, forces a cap on the scale level we can load.
+      const minLevel = this.loadSpec.multiscaleLevel ?? 0;
+      // Loaders should cache loaded dimensions so that this call blocks no more than once per valid `LoadSpec`.
+      const dims = await this.loader?.loadDims(this.loadSpecRequired);
+      if (dims) {
+        const loadableLevel = estimateLevelForAtlas(dims.map(({ shape }) => [shape[2], shape[3], shape[4]]));
+        noReload = currentScale <= Math.max(loadableLevel, minLevel);
+      }
+    }
+
     // if newly required data is not currently contained in this volume...
-    if (
-      this.loadSpecRequired.time !== this.loadSpec.time ||
-      this.loadSpecRequired.scene !== this.loadSpec.scene ||
-      !this.loadSpec.subregion.containsBox(this.loadSpecRequired.subregion)
-    ) {
+    if (!noReload) {
       // ...clone `loadSpecRequired` into `loadSpec` and load
       this.setUnloaded();
       this.loadSpec = {

src/loaders/IVolumeLoader.ts

@@ -3,7 +3,6 @@ import { Box3, Vector3 } from "three";
 import Volume from "../Volume";
 
 export class LoadSpec {
-  scene = 0;
   time = 0;
   multiscaleLevel?: number;
   /** Subregion of volume to load. If not specified, the entire volume is loaded. Specify as floats between 0-1. */
@@ -13,11 +12,10 @@ export class LoadSpec {
 
 export function loadSpecToString(spec: LoadSpec): string {
   const { min, max } = spec.subregion;
-  return `${spec.scene}:${spec.multiscaleLevel}:${spec.time}:x(${min.x},${max.x}):y(${min.y},${max.y}):z(${min.z},${max.z})`;
+  return `${spec.multiscaleLevel}:${spec.time}:x(${min.x},${max.x}):y(${min.y},${max.y}):z(${min.z},${max.z})`;
 }
 
 export class VolumeDims {
-  subpath = "";
   // shape: [t, c, z, y, x]
   shape: number[] = [0, 0, 0, 0, 0];
   // spacing: [t, c, z, y, x]; generally expect 1 for non-spatial dimensions

src/loaders/JsonImageInfoLoader.ts

@@ -80,6 +80,9 @@ const convertImageInfo = (json: JsonImageInfo): ImageInfo => ({
   timeScale: json.time_scale || 1,
   timeUnit: json.time_unit || "s",
 
+  numMultiscaleLevels: 1,
+  multiscaleLevel: 0,
+
   transform: {
     translation: json.transform?.translation
       ? new Vector3().fromArray(json.transform.translation)
@@ -92,8 +95,7 @@ const convertImageInfo = (json: JsonImageInfo): ImageInfo => ({
 
 class JsonImageInfoLoader implements IVolumeLoader {
   urls: string[];
-  time: number;
-  jsonInfo: JsonImageInfo | null = null;
+  jsonInfo: (JsonImageInfo | undefined)[];
 
   cache?: VolumeCache;
 
@@ -104,24 +106,29 @@ class JsonImageInfoLoader implements IVolumeLoader {
       this.urls = [urls];
     }
 
-    this.time = 0;
+    this.jsonInfo = new Array(this.urls.length);
     this.cache = cache;
   }
 
-  async getJsonImageInfo(loadSpec: LoadSpec): Promise<JsonImageInfo> {
-    const response = await fetch(this.urls[loadSpec.time]);
+  private async getJsonImageInfo(time: number): Promise<JsonImageInfo> {
+    const cachedInfo = this.jsonInfo[time];
+    if (cachedInfo) {
+      return cachedInfo;
+    }
+
+    const response = await fetch(this.urls[time]);
     const imageInfo = (await response.json()) as JsonImageInfo;
 
     imageInfo.pixel_size_unit = imageInfo.pixel_size_unit || "μm";
     imageInfo.times = imageInfo.times || this.urls.length;
+    this.jsonInfo[time] = imageInfo;
     return imageInfo;
   }
 
   async loadDims(loadSpec: LoadSpec): Promise<VolumeDims[]> {
-    const jsonInfo = await this.getJsonImageInfo(loadSpec);
+    const jsonInfo = await this.getJsonImageInfo(loadSpec.time);
 
     const d = new VolumeDims();
-    d.subpath = "";
     d.shape = [jsonInfo.times || 1, jsonInfo.channels, jsonInfo.tiles, jsonInfo.tile_height, jsonInfo.tile_width];
     d.spacing = [1, 1, jsonInfo.pixel_size_z, jsonInfo.pixel_size_y, jsonInfo.pixel_size_x];
     d.spaceUnit = jsonInfo.pixel_size_unit || "μm";
@@ -130,9 +137,8 @@ class JsonImageInfoLoader implements IVolumeLoader {
   }
 
   async createVolume(loadSpec: LoadSpec, onChannelLoaded?: PerChannelCallback): Promise<Volume> {
-    this.time = loadSpec.time;
-    this.jsonInfo = await this.getJsonImageInfo(loadSpec);
-    const imageInfo = convertImageInfo(this.jsonInfo);
+    const jsonInfo = await this.getJsonImageInfo(loadSpec.time);
+    const imageInfo = convertImageInfo(jsonInfo);
 
     const vol = new Volume(imageInfo, loadSpec, this);
     vol.channelLoadCallback = onChannelLoaded;
@@ -146,12 +152,9 @@ class JsonImageInfoLoader implements IVolumeLoader {
     // Try to figure out the urlPrefix from the LoadSpec.
     // For this format we assume the image data is in the same directory as the json file.
     const loadSpec = explicitLoadSpec || vol.loadSpec;
-    if (this.time !== loadSpec.time) {
-      this.time = loadSpec.time;
-      this.jsonInfo = await this.getJsonImageInfo(loadSpec);
-    }
+    const jsonInfo = await this.getJsonImageInfo(loadSpec.time);
 
-    let images = this.jsonInfo?.images;
+    let images = jsonInfo?.images;
     if (!images) {
       return;
     }
@@ -163,7 +166,7 @@ class JsonImageInfoLoader implements IVolumeLoader {
     }
 
     // This regex removes everything after the last slash, so the url had better be simple.
-    const urlPrefix = this.urls[this.time].replace(/[^/]*$/, "");
+    const urlPrefix = this.urls[loadSpec.time].replace(/[^/]*$/, "");
     images = images.map((element) => ({ ...element, name: urlPrefix + element.name }));
 
     vol.loadSpec = {

src/loaders/OmeZarrLoader.ts

@@ -17,7 +17,6 @@ import {
   unitNameToSymbol,
 } from "./VolumeLoaderUtils";
 
-const MAX_ATLAS_DIMENSION = 2048;
 const CHUNK_REQUEST_CANCEL_REASON = "chunk request cancelled";
 
 type CoordinateTransformation =
@@ -194,22 +193,21 @@ function remapAxesToTCZYX(axes: Axis[]): [number, number, number, number, number
   return axisTCZYX;
 }
 
-function pickLevelToLoad(loadSpec: LoadSpec, zarrMultiscales: ZarrArray[], zi: number, yi: number, xi: number): number {
+/**
+ * Picks the best scale level to load based on scale level dimensions, a max atlas size, and a `LoadSpec`.
+ * This works like `estimateLevelForAtlas` but factors in `LoadSpec`'s `subregion` property (shrinks the size of the
+ * data, maybe enough to allow loading a higher level) and its `multiscaleLevel` property (sets a max scale level).
+ */
+function pickLevelToLoad(loadSpec: LoadSpec, spatialDimsZYX: [number, number, number][]): number {
   const size = loadSpec.subregion.getSize(new Vector3());
-
-  const spatialDims = zarrMultiscales.map(({ shape }) => [
-    Math.max(shape[zi] * size.z, 1),
-    Math.max(shape[yi] * size.y, 1),
-    Math.max(shape[xi] * size.x, 1),
+  const dims = spatialDimsZYX.map(([z, y, x]): [number, number, number] => [
+    Math.max(z * size.z, 1),
+    Math.max(y * size.y, 1),
+    Math.max(x * size.x, 1),
   ]);
 
-  const optimalLevel = estimateLevelForAtlas(spatialDims, MAX_ATLAS_DIMENSION);
-
-  if (loadSpec.multiscaleLevel) {
-    return Math.max(optimalLevel, loadSpec.multiscaleLevel);
-  } else {
-    return optimalLevel;
-  }
+  const optimalLevel = estimateLevelForAtlas(dims);
+  return Math.max(optimalLevel, loadSpec.multiscaleLevel ?? 0);
 }
 
 function convertChannel(channelData: TypedArray, dtype: string): Uint8Array {
@@ -338,7 +336,12 @@ class OMEZarrLoader implements IVolumeLoader {
     return scale;
   }
 
-  async loadDims(loadSpec: LoadSpec): Promise<VolumeDims[]> {
+  private getLevelShapesZYX(): [number, number, number][] {
+    const [z, y, x] = this.axesTCZYX.slice(-3);
+    return this.scaleLevels.map(({ shape }) => [shape[z], shape[y], shape[x]]);
+  }
+
+  loadDims(loadSpec: LoadSpec): Promise<VolumeDims[]> {
     const [spaceUnit, timeUnit] = this.getUnitSymbols();
     // Compute subregion size so we can factor that in
     const maxExtent = this.maxExtent ?? new Box3(new Vector3(0, 0, 0), new Vector3(1, 1, 1));
@@ -352,7 +355,6 @@ class OMEZarrLoader implements IVolumeLoader {
 
       dims.spaceUnit = spaceUnit;
       dims.timeUnit = timeUnit;
-      dims.subpath = level.path ?? "";
       dims.shape = [-1, -1, -1, -1, -1];
       dims.spacing = [1, 1, 1, 1, 1];
 
@@ -366,7 +368,7 @@ class OMEZarrLoader implements IVolumeLoader {
       return dims;
     });
 
-    return result;
+    return Promise.resolve(result);
   }
 
   async createVolume(loadSpec: LoadSpec, onChannelLoaded?: PerChannelCallback): Promise<Volume> {
@@ -375,7 +377,7 @@ class OMEZarrLoader implements IVolumeLoader {
     const hasC = c > -1;
 
     const shape0 = this.scaleLevels[0].shape;
-    const levelToLoad = pickLevelToLoad(loadSpec, this.scaleLevels, z, y, x);
+    const levelToLoad = pickLevelToLoad(loadSpec, this.getLevelShapesZYX());
     const shapeLv = this.scaleLevels[levelToLoad].shape;
 
     const [spatialUnit, timeUnit] = this.getUnitSymbols();
@@ -413,6 +415,8 @@ class OMEZarrLoader implements IVolumeLoader {
       times,
       timeScale,
       timeUnit,
+      numMultiscaleLevels: this.scaleLevels.length,
+      multiscaleLevel: levelToLoad,
 
       transform: {
         translation: new Vector3(0, 0, 0),
@@ -442,7 +446,7 @@ class OMEZarrLoader implements IVolumeLoader {
     const [z, y, x] = this.axesTCZYX.slice(2);
     const subregion = composeSubregion(vol.loadSpec.subregion, maxExtent);
 
-    const levelIdx = pickLevelToLoad({ ...vol.loadSpec, subregion }, this.scaleLevels, z, y, x);
+    const levelIdx = pickLevelToLoad({ ...vol.loadSpec, subregion }, this.getLevelShapesZYX());
     const level = this.scaleLevels[levelIdx];
     const levelShape = level.shape;
 
@@ -458,6 +462,7 @@ class OMEZarrLoader implements IVolumeLoader {
       volumeSize: volSizePx,
       subregionSize: regionSizePx,
       subregionOffset: regionPx.min,
+      multiscaleLevel: levelIdx,
     };
     vol.updateDimensions();
 

src/loaders/OpenCellLoader.ts

@@ -9,7 +9,6 @@ import { JsonImageInfoLoader } from "./JsonImageInfoLoader";
 class OpenCellLoader implements IVolumeLoader {
   async loadDims(_: LoadSpec): Promise<VolumeDims[]> {
     const d = new VolumeDims();
-    d.subpath = "";
     d.shape = [1, 2, 27, 600, 600];
     d.spacing = [1, 1, 2, 1, 1];
     d.spaceUnit = ""; // unknown unit.
@@ -41,6 +40,9 @@ class OpenCellLoader implements IVolumeLoader {
       timeScale: 1,
       timeUnit: "",
 
+      numMultiscaleLevels: 1,
+      multiscaleLevel: 0,
+
       transform: {
         translation: new Vector3(0, 0, 0),
         rotation: new Vector3(0, 0, 0),