Merge pull request #5169 from voxel51/improv/image-bitmaps

labels rendering performance improvement: create ImageBitmaps in worker

app/packages/looker/src/lookers/abstract.ts

@@ -293,6 +293,11 @@ export abstract class AbstractLooker<
           return;
         }
 
+        if (this.state.destroyed && this.sampleOverlays) {
+          // close all current overlays
+          this.pluckedOverlays.forEach((overlay) => overlay.cleanup?.());
+        }
+
         if (
           !this.state.windowBBox ||
           this.state.destroyed ||

app/packages/looker/src/lookers/frame-reader.ts

@@ -52,7 +52,15 @@ interface AcquireReaderOptions {
 export const { acquireReader, clearReader } = (() => {
   const createCache = (removeFrame: RemoveFrame) => {
     return new LRUCache<number, Frame>({
-      dispose: (_, key) => removeFrame(key),
+      dispose: (frame, key) => {
+        const overlays = frame.overlays;
+
+        for (let i = 0; i < overlays.length; i++) {
+          overlays[i].cleanup?.();
+        }
+
+        removeFrame(key);
+      },
       max: MAX_FRAME_STREAM_SIZE,
       maxSize: MAX_FRAME_STREAM_SIZE_BYTES,
       noDisposeOnSet: true,

app/packages/looker/src/overlays/base.ts

@@ -3,6 +3,7 @@
  */
 
 import { getCls, sizeBytesEstimate } from "@fiftyone/utilities";
+import { OverlayMask } from "../numpy";
 import type { BaseState, Coordinates, NONFINITE } from "../state";
 import { getLabelColor, shouldShowLabelTag } from "./util";
 
@@ -39,6 +40,11 @@ export interface SelectData {
   frameNumber?: number;
 }
 
+export type LabelMask = {
+  bitmap?: ImageBitmap;
+  data?: OverlayMask;
+};
+
 export interface RegularLabel extends BaseLabel {
   _id?: string;
   label?: string;
@@ -67,6 +73,7 @@ export interface Overlay<State extends Partial<BaseState>> {
   getPoints(state: Readonly<State>): Coordinates[];
   getSelectData(state: Readonly<State>): SelectData;
   getSizeBytes(): number;
+  cleanup?(): void;
 }
 
 export abstract class CoordinateOverlay<

app/packages/looker/src/overlays/detection.ts

@@ -4,17 +4,19 @@
 import { NONFINITES } from "@fiftyone/utilities";
 
 import { INFO_COLOR } from "../constants";
-import { OverlayMask } from "../numpy";
 import { BaseState, BoundingBox, Coordinates, NONFINITE } from "../state";
 import { distanceFromLineSegment } from "../util";
-import { CONTAINS, CoordinateOverlay, PointInfo, RegularLabel } from "./base";
+import {
+  CONTAINS,
+  CoordinateOverlay,
+  LabelMask,
+  PointInfo,
+  RegularLabel,
+} from "./base";
 import { t } from "./util";
 
 export interface DetectionLabel extends RegularLabel {
-  mask?: {
-    data: OverlayMask;
-    image: ArrayBuffer;
-  };
+  mask?: LabelMask;
   bounding_box: BoundingBox;
 
   // valid for 3D bounding boxes
@@ -27,10 +29,8 @@ export interface DetectionLabel extends RegularLabel {
 export default class DetectionOverlay<
   State extends BaseState
 > extends CoordinateOverlay<State, DetectionLabel> {
-  private imageData: ImageData;
   private is3D: boolean;
   private labelBoundingBox: BoundingBox;
-  private canvas: HTMLCanvasElement;
 
   constructor(field, label) {
     super(field, label);
@@ -40,32 +40,6 @@ export default class DetectionOverlay<
     } else {
       this.is3D = false;
     }
-
-    if (this.label.mask) {
-      const [height, width] = this.label.mask.data.shape;
-
-      if (!height || !width) {
-        return;
-      }
-
-      this.canvas = document.createElement("canvas");
-      this.canvas.width = width;
-      this.canvas.height = height;
-      this.imageData = new ImageData(
-        new Uint8ClampedArray(this.label.mask.image),
-        width,
-        height
-      );
-      const maskCtx = this.canvas.getContext("2d");
-      maskCtx.imageSmoothingEnabled = false;
-      maskCtx.clearRect(
-        0,
-        0,
-        this.label.mask.data.shape[1],
-        this.label.mask.data.shape[0]
-      );
-      maskCtx.putImageData(this.imageData, 0, 0);
-    }
   }
 
   containsPoint(state: Readonly<State>): CONTAINS {
@@ -169,16 +143,17 @@ export default class DetectionOverlay<
   }
 
   private drawMask(ctx: CanvasRenderingContext2D, state: Readonly<State>) {
-    if (!this.canvas) {
+    if (!this.label.mask?.bitmap) {
       return;
     }
 
     const [tlx, tly, w, h] = this.label.bounding_box;
     const [x, y] = t(state, tlx, tly);
     const tmp = ctx.globalAlpha;
     ctx.globalAlpha = state.options.alpha;
+    ctx.imageSmoothingEnabled = false;
     ctx.drawImage(
-      this.canvas,
+      this.label.mask.bitmap,
       x,
       y,
       w * state.canvasBBox[2],
@@ -285,6 +260,13 @@ export default class DetectionOverlay<
     const oh = state.strokeWidth / state.canvasBBox[3];
     return [(bx - ow) * w, (by - oh) * h, (bw + ow * 2) * w, (bh + oh * 2) * h];
   }
+
+  public cleanup(): void {
+    if (this.label.mask?.bitmap) {
+      this.label.mask?.bitmap.close();
+      this.label.mask.bitmap = null;
+    }
+  }
 }
 
 export const getDetectionPoints = (labels: DetectionLabel[]): Coordinates[] => {

app/packages/looker/src/overlays/heatmap.ts

@@ -6,28 +6,25 @@ import {
   getColor,
   getRGBA,
   getRGBAColor,
+  sizeBytesEstimate,
 } from "@fiftyone/utilities";
-import { ARRAY_TYPES, OverlayMask, TypedArray } from "../numpy";
+import { ARRAY_TYPES, TypedArray } from "../numpy";
 import { BaseState, Coordinates } from "../state";
 import { isFloatArray } from "../util";
 import { clampedIndex } from "../worker/painter";
 import {
   BaseLabel,
   CONTAINS,
+  LabelMask,
   Overlay,
   PointInfo,
   SelectData,
   isShown,
 } from "./base";
 import { strokeCanvasRect, t } from "./util";
 
-interface HeatMap {
-  data: OverlayMask;
-  image: ArrayBuffer;
-}
-
 interface HeatmapLabel extends BaseLabel {
-  map?: HeatMap;
+  map?: LabelMask;
   range?: [number, number];
 }
 
@@ -45,8 +42,6 @@ export default class HeatmapOverlay<State extends BaseState>
   private label: HeatmapLabel;
   private targets?: TypedArray;
   private readonly range: [number, number];
-  private canvas: HTMLCanvasElement;
-  private imageData: ImageData;
 
   constructor(field: string, label: HeatmapLabel) {
     this.field = field;
@@ -68,25 +63,6 @@ export default class HeatmapOverlay<State extends BaseState>
     if (!width || !height) {
       return;
     }
-
-    this.canvas = document.createElement("canvas");
-    this.canvas.width = width;
-    this.canvas.height = height;
-
-    this.imageData = new ImageData(
-      new Uint8ClampedArray(this.label.map.image),
-      width,
-      height
-    );
-    const maskCtx = this.canvas.getContext("2d");
-    maskCtx.imageSmoothingEnabled = false;
-    maskCtx.clearRect(
-      0,
-      0,
-      this.label.map.data.shape[1],
-      this.label.map.data.shape[0]
-    );
-    maskCtx.putImageData(this.imageData, 0, 0);
   }
 
   containsPoint(state: Readonly<State>): CONTAINS {
@@ -101,22 +77,12 @@ export default class HeatmapOverlay<State extends BaseState>
   }
 
   draw(ctx: CanvasRenderingContext2D, state: Readonly<State>): void {
-    if (this.imageData) {
-      const maskCtx = this.canvas.getContext("2d");
-      maskCtx.imageSmoothingEnabled = false;
-      maskCtx.clearRect(
-        0,
-        0,
-        this.label.map.data.shape[1],
-        this.label.map.data.shape[0]
-      );
-      maskCtx.putImageData(this.imageData, 0, 0);
-
+    if (this.label.map?.bitmap) {
       const [tlx, tly] = t(state, 0, 0);
       const [brx, bry] = t(state, 1, 1);
       const tmp = ctx.globalAlpha;
       ctx.globalAlpha = state.options.alpha;
-      ctx.drawImage(this.canvas, tlx, tly, brx - tlx, bry - tly);
+      ctx.drawImage(this.label.map.bitmap, tlx, tly, brx - tlx, bry - tly);
       ctx.globalAlpha = tmp;
     }
 
@@ -235,6 +201,16 @@ export default class HeatmapOverlay<State extends BaseState>
 
     return this.targets[index];
   }
+
+  getSizeBytes(): number {
+    return sizeBytesEstimate(this.label);
+  }
+
+  public cleanup(): void {
+    if (this.label.map?.bitmap) {
+      this.label.map?.bitmap.close();
+    }
+  }
 }
 
 export const getHeatmapPoints = (labels: HeatmapLabel[]): Coordinates[] => {

app/packages/looker/src/overlays/segmentation.ts

@@ -2,12 +2,13 @@
  * Copyright 2017-2024, Voxel51, Inc.
  */
 
-import { getColor } from "@fiftyone/utilities";
-import { ARRAY_TYPES, OverlayMask, TypedArray } from "../numpy";
+import { getColor, sizeBytesEstimate } from "@fiftyone/utilities";
+import { ARRAY_TYPES, TypedArray } from "../numpy";
 import { BaseState, Coordinates, MaskTargets } from "../state";
 import {
   BaseLabel,
   CONTAINS,
+  LabelMask,
   Overlay,
   PointInfo,
   SelectData,
@@ -16,10 +17,7 @@ import {
 import { isRgbMaskTargets, strokeCanvasRect, t } from "./util";
 
 interface SegmentationLabel extends BaseLabel {
-  mask?: {
-    data: OverlayMask;
-    image: ArrayBuffer;
-  };
+  mask?: LabelMask;
 }
 
 interface SegmentationInfo extends BaseLabel {
@@ -34,8 +32,6 @@ export default class SegmentationOverlay<State extends BaseState>
   readonly field: string;
   private label: SegmentationLabel;
   private targets?: TypedArray;
-  private canvas: HTMLCanvasElement;
-  private imageData: ImageData;
 
   private isRgbMaskTargets = false;
 
@@ -53,6 +49,7 @@ export default class SegmentationOverlay<State extends BaseState>
     if (!this.label.mask) {
       return;
     }
+
     const [height, width] = this.label.mask.data.shape;
 
     if (!height || !width) {
@@ -62,25 +59,6 @@ export default class SegmentationOverlay<State extends BaseState>
     this.targets = new ARRAY_TYPES[this.label.mask.data.arrayType](
       this.label.mask.data.buffer
     );
-
-    this.canvas = document.createElement("canvas");
-    this.canvas.width = width;
-    this.canvas.height = height;
-
-    this.imageData = new ImageData(
-      new Uint8ClampedArray(this.label.mask.image),
-      width,
-      height
-    );
-    const maskCtx = this.canvas.getContext("2d");
-    maskCtx.imageSmoothingEnabled = false;
-    maskCtx.clearRect(
-      0,
-      0,
-      this.label.mask.data.shape[1],
-      this.label.mask.data.shape[0]
-    );
-    maskCtx.putImageData(this.imageData, 0, 0);
   }
 
   containsPoint(state: Readonly<State>): CONTAINS {
@@ -99,12 +77,12 @@ export default class SegmentationOverlay<State extends BaseState>
       return;
     }
 
-    if (this.imageData) {
+    if (this.label.mask?.bitmap) {
       const [tlx, tly] = t(state, 0, 0);
       const [brx, bry] = t(state, 1, 1);
       const tmp = ctx.globalAlpha;
       ctx.globalAlpha = state.options.alpha;
-      ctx.drawImage(this.canvas, tlx, tly, brx - tlx, bry - tly);
+      ctx.drawImage(this.label.mask.bitmap, tlx, tly, brx - tlx, bry - tly);
       ctx.globalAlpha = tmp;
     }
 
@@ -278,6 +256,16 @@ export default class SegmentationOverlay<State extends BaseState>
     }
     return this.targets[index];
   }
+
+  getSizeBytes(): number {
+    return sizeBytesEstimate(this.label);
+  }
+
+  public cleanup(): void {
+    if (this.label.mask?.bitmap) {
+      this.label.mask?.bitmap.close();
+    }
+  }
 }
 
 export const getSegmentationPoints = (