perf: Only create the texture on the parent layer (not sub-layers)

Author: manzt

src/io.ts

@@ -227,7 +227,7 @@ export function initLayerStateFromSource(source: SourceData & { id: string }): L
         loader: label.loader,
         modelMatrix: label.modelMatrix,
         opacity: 1,
-        lut: label.lut,
+        colors: label.colors,
       })),
       transformSourceSelection: getTransformSourceSelectionFromLabels(
         source.labels.map((label) => label.loader[0]),

src/layers/label-layer.ts

@@ -1,21 +1,21 @@
-import { getImageSize } from "@hms-dbmi/viv";
-import { BitmapLayer, TileLayer, type UpdateParameters } from "deck.gl";
+import { BitmapLayer, TileLayer } from "deck.gl";
 import * as utils from "../utils";
 
+import type { Layer, UpdateParameters } from "deck.gl";
 import { type Matrix4, clamp } from "math.gl";
 import type * as zarr from "zarrita";
 import type { ZarrPixelSource } from "../ZarrPixelSource";
 
 type Texture = ReturnType<BitmapLayer["context"]["device"]["createTexture"]>;
 
-export type LabelLayerLut = Readonly<Record<number, readonly [number, number, number, number]>>;
+export type OmeColors = Readonly<Record<number, readonly [number, number, number, number]>>;
 export interface LabelLayerProps {
   id: string;
   loader: Array<ZarrPixelSource>;
   selection: Array<number>;
   opacity: number;
   modelMatrix: Matrix4;
-  lut?: LabelLayerLut;
+  colors?: OmeColors;
 }
 
 /**
@@ -35,12 +35,18 @@ type LabelPixelData = {
   height: number;
 };
 
-export class LabelLayer extends TileLayer<LabelPixelData> {
+export class LabelLayer extends TileLayer<LabelPixelData, LabelLayerProps> {
   static layerName = "VizarrLabelLayer";
+  // @ts-expect-error - only way to extend the base state type
+  state!: { colorTexture: Texture } & TileLayer["state"];
 
   constructor(props: LabelLayerProps) {
     const resolutions = props.loader;
-    const dimensions = getImageSize(resolutions[0]);
+    const dimensions = {
+      height: resolutions[0].shape.at(-2),
+      width: resolutions[0].shape.at(-1),
+    };
+    utils.assert(dimensions.width && dimensions.height);
     const tileSize = getTileSizeForResolutions(resolutions);
     super({
       id: `labels-${props.id}`,
@@ -50,6 +56,7 @@ export class LabelLayer extends TileLayer<LabelPixelData> {
       opacity: props.opacity,
       maxZoom: 0,
       modelMatrix: props.modelMatrix,
+      colors: props.colors,
       zoomOffset: Math.round(Math.log2(props.modelMatrix ? props.modelMatrix.getScale()[0] : 1)),
       updateTriggers: {
         getTileData: [props.loader, props.selection],
@@ -65,29 +72,68 @@ export class LabelLayer extends TileLayer<LabelPixelData> {
         );
         return { data, width, height };
       },
-      renderSubLayers({ tile, data }) {
-        const [[left, bottom], [right, top]] = tile.boundingBox;
-        const { width, height } = dimensions;
-        return new GrayscaleBitmapLayer({
-          id: `tile-${tile.index.x}.${tile.index.y}.${tile.index.z}-${props.id}`,
-          pixelData: data,
-          opacity: props.opacity,
-          modelMatrix: props.modelMatrix,
-          lut: props.lut,
-          bounds: [clamp(left, 0, width), clamp(top, 0, height), clamp(right, 0, width), clamp(bottom, 0, height)],
-          // For underlying class
-          image: new ImageData(data.width, data.height),
-          pickable: false,
-        });
-      },
     });
   }
+
+  renderSubLayers(
+    params: TileLayer["props"] & {
+      data: LabelPixelData;
+      tile: {
+        index: { x: number; y: number; z: number };
+        boundingBox: [min: Array<number>, max: Array<number>];
+      };
+    },
+  ): Layer {
+    const { tile, data, ...props } = params;
+    const [[left, bottom], [right, top]] = tile.boundingBox;
+    utils.assert(props.extent, "missing extent");
+    const [_x0, _y0, width, height] = props.extent;
+    return new GrayscaleBitmapLayer({
+      id: `tile-${tile.index.x}.${tile.index.y}.${tile.index.z}-${props.id}`,
+      pixelData: data,
+      opacity: props.opacity,
+      modelMatrix: props.modelMatrix,
+      colorTexture: this.state.colorTexture,
+      bounds: [clamp(left, 0, width), clamp(top, 0, height), clamp(right, 0, width), clamp(bottom, 0, height)],
+      // For underlying class
+      image: new ImageData(data.width, data.height),
+      pickable: false,
+    });
+  }
+
+  updateState({ props, oldProps, changeFlags, ...rest }: UpdateParameters<this>): void {
+    super.updateState({ props, oldProps, changeFlags, ...rest });
+    // we make the colorTexture on this layer so we can share it amoung all the sublayers
+    if (props.colors !== oldProps.colors || !this.state.colorTexture) {
+      this.state.colorTexture?.destroy();
+      const colorTexture = createColorTexture({
+        source: props.colors,
+        maxTextureDimension2D: this.context.device.limits.maxTextureDimension2D,
+      });
+      this.setState({
+        colorTexture: this.context.device.createTexture({
+          width: colorTexture.width,
+          height: colorTexture.height,
+          data: colorTexture.data,
+          dimension: "2d",
+          mipmaps: false,
+          sampler: {
+            minFilter: "nearest",
+            magFilter: "nearest",
+            addressModeU: "clamp-to-edge",
+            addressModeV: "clamp-to-edge",
+          },
+          format: "rgba8unorm",
+        }),
+      });
+    }
+  }
 }
 
-export class GrayscaleBitmapLayer extends BitmapLayer<{ pixelData: LabelPixelData; lut?: LabelLayerLut }> {
+export class GrayscaleBitmapLayer extends BitmapLayer<{ pixelData: LabelPixelData; colorTexture: Texture }> {
   static layerName = "VizarrGrayscaleBitmapLayer";
   // @ts-expect-error - only way to extend the base state type
-  state!: { texture: Texture; colorTexture: Texture } & BitmapLayer["state"];
+  state!: { texture: Texture } & BitmapLayer["state"];
 
   getShaders() {
     const sampler = (
@@ -162,40 +208,21 @@ void main() {
         }),
       });
     }
-
-    if (props.lut !== oldProps.lut || !this.state.colorTexture) {
-      this.state.colorTexture?.destroy();
-      const colorTexture = createColorTexture({
-        source: props.lut,
-        maxTextureDimension2D: this.context.device.limits.maxTextureDimension2D,
-      });
-      this.setState({
-        colorTexture: this.context.device.createTexture({
-          width: colorTexture.width,
-          height: colorTexture.height,
-          data: colorTexture.data,
-          dimension: "2d",
-          mipmaps: false,
-          sampler: {
-            minFilter: "nearest",
-            magFilter: "nearest",
-            addressModeU: "clamp-to-edge",
-            addressModeV: "clamp-to-edge",
-          },
-          format: "rgba8unorm",
-        }),
-      });
-    }
   }
 
   draw(opts: unknown) {
-    const { model, texture, colorTexture } = this.state;
+    const { model, texture } = this.state;
+    const { colorTexture } = this.props;
+
     if (model && texture && colorTexture) {
       model.setUniforms({
         colorTextureWidth: colorTexture.width,
         colorTextureHeight: colorTexture.height,
       });
-      model.setBindings({ grayscaleTexture: texture, colorTexture: colorTexture });
+      model.setBindings({
+        grayscaleTexture: texture,
+        colorTexture: colorTexture,
+      });
     }
     super.draw(opts);
   }
@@ -210,7 +237,7 @@ function getTileSizeForResolutions(resolutions: Array<ZarrPixelSource>): number
   return tileSize;
 }
 
-const SEEN_LUTS = new WeakSet<LabelLayerLut>();
+const SEEN_LUTS = new WeakSet<OmeColors>();
 
 /**
  * Creates a color lookup table (LUT) as a 2D texture.
@@ -219,7 +246,7 @@ const SEEN_LUTS = new WeakSet<LabelLayerLut>();
  * @param options.maxTextureDimension2D - The maximum texture dimension size.
  */
 function createColorTexture(options: {
-  source?: LabelLayerLut;
+  source?: OmeColors;
   maxTextureDimension2D: number;
 }) {
   const { source, maxTextureDimension2D } = options;

src/ome.ts

@@ -3,7 +3,7 @@ import * as zarr from "zarrita";
 import type { ImageLabels, ImageLayerConfig, OnClickData, SourceData } from "./state";
 
 import { ZarrPixelSource } from "./ZarrPixelSource";
-import type { LabelLayerLut } from "./layers/label-layer";
+import type { OmeColors } from "./layers/label-layer";
 import * as utils from "./utils";
 
 export async function loadWell(
@@ -283,9 +283,9 @@ async function loadOmeImageLabel(root: zarr.Location<zarr.Readable>, name: strin
   const labels = utils.getNgffAxisLabels(axes);
   return {
     name,
-    lut: resolveImageLabelsLut(attrs["image-label"]),
     modelMatrix: utils.coordinateTransformationsToMatrix(attrs.multiscales),
     loader: data.map((arr) => new ZarrPixelSource(arr, { labels, tileSize })),
+    colors: resolveImageLabelsLut(attrs["image-label"]),
   };
 }
 
@@ -360,7 +360,7 @@ function parseOmeroMeta({ rdefs, channels, name }: Ome.Omero, axes: Ome.Axis[]):
   };
 }
 
-function resolveImageLabelsLut(attrs: Ome.ImageLabel): LabelLayerLut | undefined {
+function resolveImageLabelsLut(attrs: Ome.ImageLabel): OmeColors | undefined {
   if (!attrs.colors || attrs.colors.length === 0) {
     return undefined;
   }

src/state.ts

@@ -11,7 +11,7 @@ import type { ZarrPixelSource } from "./ZarrPixelSource";
 import { initLayerStateFromSource } from "./io";
 
 import { GridLayer, type GridLayerProps, type GridLoader } from "./layers/grid-layer";
-import { LabelLayer, type LabelLayerLut, type LabelLayerProps } from "./layers/label-layer";
+import { LabelLayer, type LabelLayerProps, type OmeColors } from "./layers/label-layer";
 import {
   ImageLayer,
   type ImageLayerProps,
@@ -59,7 +59,7 @@ export type ImageLabels = Array<{
   name: string;
   loader: ZarrPixelSource[];
   modelMatrix: Matrix4;
-  lut?: LabelLayerLut;
+  colors?: OmeColors;
 }>;
 
 export type SourceData = {
@@ -187,7 +187,6 @@ export const layerAtoms = atom((get) => {
       const { layerProps, transformSourceSelection } = layer.labels;
       const selection = transformSourceSelection(layer.layerProps.selections[0]);
       const imageLabelLayers = layerProps.map((props) => new LabelLayer({ ...props, selection }));
-      // @ts-expect-error - TS can't resolve that Array<ImageLabelLayer> == Array<Layer<ImageLabelLayerProps>>
       layers.push(...imageLabelLayers);
     }
     return layers;