cleanup overlays in the worker listener callback instead

voxel51 · Dec 12, 2024 · 4f76488 · 4f76488
1 parent d3bf174
commit 4f76488
Show file tree

Hide file tree

Showing 6 changed files with 23 additions and 38 deletions.
diff --git a/app/packages/looker/src/lookers/abstract.ts b/app/packages/looker/src/lookers/abstract.ts
@@ -520,10 +520,6 @@ export abstract class AbstractLooker<
     const arrayBuffers: ArrayBuffer[] = [];
 
     for (const overlay of this.pluckedOverlays ?? []) {
-      // we paint overlays again, so cleanup the old ones
-      // this helps prevent memory leaks from, for instance, dangling ImageBitmaps
-      overlay.cleanup();
-
       let overlayData: LabelMask = null;
 
       if ("mask" in overlay.label) {
@@ -739,13 +735,22 @@ export abstract class AbstractLooker<
     );
   }
 
+  protected cleanOverlays() {
+    for (const overlay of this.sampleOverlays ?? []) {
+      overlay.cleanup();
+    }
+  }
+
   private loadSample(sample: Sample, transfer: Transferable[] = []) {
     const messageUUID = uuid();
 
     const labelsWorker = getLabelsWorker();
 
     const listener = ({ data: { sample, coloring, uuid } }) => {
       if (uuid === messageUUID) {
+        // we paint overlays again, so cleanup the old ones
+        // this helps prevent memory leaks from, for instance, dangling ImageBitmaps
+        this.cleanOverlays();
         this.sample = sample;
         this.state.options.coloring = coloring;
         this.loadOverlays(sample);

diff --git a/app/packages/looker/src/overlays/base.ts b/app/packages/looker/src/overlays/base.ts
@@ -42,7 +42,6 @@ export interface SelectData {
 
 export type LabelMask = {
   bitmap?: ImageBitmap;
-  closedBitmapDims?: { width: number; height: number };
   data?: OverlayMask;
 };
 

diff --git a/app/packages/looker/src/overlays/detection.ts b/app/packages/looker/src/overlays/detection.ts
@@ -262,16 +262,8 @@ export default class DetectionOverlay<
   }
 
   public cleanup(): void {
-    if (this.label.mask?.bitmap) {
-      // store height and width in bitmap object since it might be used again
-      const height = this.label.mask.bitmap.height;
-      const width = this.label.mask.bitmap.width;
-
-      this.label.mask?.bitmap.close();
-      this.label.mask.bitmap = null;
-
-      this.label.mask.closedBitmapDims = { width, height };
-    }
+    this.label.mask?.bitmap?.close();
+    this.label.mask.bitmap = null;
   }
 }
 

diff --git a/app/packages/looker/src/overlays/heatmap.ts b/app/packages/looker/src/overlays/heatmap.ts
@@ -207,16 +207,8 @@ export default class HeatmapOverlay<State extends BaseState>
   }
 
   public cleanup(): void {
-    if (this.label.map?.bitmap) {
-      // store height and width in bitmap object since it might be used again
-      const height = this.label.map.bitmap.height;
-      const width = this.label.map.bitmap.width;
-
-      this.label.map?.bitmap.close();
-      this.label.map.bitmap = null;
-
-      this.label.map.closedBitmapDims = { width, height };
-    }
+    this.label.map?.bitmap?.close();
+    this.label.map.bitmap = null;
   }
 }
 

diff --git a/app/packages/looker/src/overlays/segmentation.ts b/app/packages/looker/src/overlays/segmentation.ts
@@ -262,16 +262,8 @@ export default class SegmentationOverlay<State extends BaseState>
   }
 
   public cleanup(): void {
-    if (this.label.mask?.bitmap) {
-      // store height and width in bitmap object since it might be used again
-      const height = this.label.mask.bitmap.height;
-      const width = this.label.mask.bitmap.width;
-
-      this.label.mask?.bitmap.close();
-      this.label.mask.bitmap = null;
-
-      this.label.mask.closedBitmapDims = { width, height };
-    }
+    this.label.mask?.bitmap?.close();
+    this.label.mask.bitmap = null;
   }
 }
 

diff --git a/app/packages/looker/src/worker/disk-overlay-decoder.ts b/app/packages/looker/src/worker/disk-overlay-decoder.ts
@@ -56,11 +56,16 @@ export const decodeOverlayOnDisk = async (
     // it's possible we're just re-coloring, in which case re-init mask image and set bitmap to null
     if (
       label[overlayField] &&
-      label[overlayField].closedBitmapDims &&
+      label[overlayField].bitmap &&
       !label[overlayField].image
     ) {
-      const height = label[overlayField].closedBitmapDims.height;
-      const width = label[overlayField].closedBitmapDims.width;
+      const height = label[overlayField].bitmap.height;
+      const width = label[overlayField].bitmap.width;
+
+      // close the copied bitmap
+      label[overlayField].bitmap.close();
+      label[overlayField].bitmap = null;
+
       label[overlayField].image = new ArrayBuffer(height * width * 4);
       label[overlayField].bitmap.close();
       label[overlayField].bitmap = null;