reduce heatmap targets dims

app/packages/looker/src/worker/canvas-decoder.ts

@@ -1,9 +1,16 @@
+import { HEATMAP } from "@fiftyone/utilities";
 import { OverlayMask } from "../numpy";
 
-const offScreenCanvas = new OffscreenCanvas(1, 1);
-const offScreenCanvasCtx = offScreenCanvas.getContext("2d", {
-  willReadFrequently: true,
-})!;
+const canvasAndCtx = (() => {
+  if ("OffscreenCanvas" in self) {
+    const offScreenCanvas = new OffscreenCanvas(1, 1);
+    const offScreenCanvasCtx = offScreenCanvas.getContext("2d", {
+      willReadFrequently: true,
+    })!;
+
+    return { canvas: offScreenCanvas, ctx: offScreenCanvasCtx };
+  }
+})();
 
 const PNG_SIGNATURE = [0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a];
 /**
@@ -34,7 +41,33 @@ const getPngcolorType = async (blob: Blob): Promise<number | undefined> => {
   return colorType;
 };
 
-export const decodeWithCanvas = async (blob: Blob) => {
+/**
+ * Sets the buffer in place to grayscale by removing the G, B, and A channels.
+ *
+ * This is meant for images that are packed like the following, since the other channels are storing redundant data:
+ *
+ * X, X, X, 255, Y, Y, Y, 255, Z, Z, Z, 255, ...
+ */
+export const recastBufferToMonoChannel = (
+  uint8Array: Uint8ClampedArray,
+  width: number,
+  height: number,
+  stride: number
+) => {
+  const totalPixels = width * height;
+
+  let read = 0;
+  let write = 0;
+
+  while (write < totalPixels) {
+    uint8Array[write++] = uint8Array[read];
+    read += stride;
+  }
+
+  return uint8Array.slice(0, totalPixels).buffer;
+};
+
+export const decodeWithCanvas = async (blob: Blob, cls: string) => {
   let channels: number = 4;
 
   if (blob.type === "image/png") {
@@ -58,6 +91,8 @@ export const decodeWithCanvas = async (blob: Blob) => {
   const imageBitmap = await createImageBitmap(blob);
   const { width, height } = imageBitmap;
 
+  const { canvas: offScreenCanvas, ctx: offScreenCanvasCtx } = canvasAndCtx!;
+
   offScreenCanvas.width = width;
   offScreenCanvas.height = height;
 
@@ -67,26 +102,27 @@ export const decodeWithCanvas = async (blob: Blob) => {
 
   const imageData = offScreenCanvasCtx.getImageData(0, 0, width, height);
 
-  if (channels === 1) {
-    // get rid of the G, B, and A channels, new buffer will be 1/4 the size
-    const rawBuffer = imageData.data;
-    const totalPixels = width * height;
-
-    let read = 0;
-    let write = 0;
+  let targetsBuffer = imageData.data.buffer;
 
-    while (write < totalPixels) {
-      rawBuffer[write++] = rawBuffer[read];
-      // skip "G,B,A"
-      read += 4;
-    }
+  if (channels === 1) {
+    // recasting because we know from png header that it's grayscale,
+    // but when we decoded using canvas, it's RGBA
+    targetsBuffer = recastBufferToMonoChannel(imageData.data, width, height, 4);
+  }
 
-    const grayScaleData = rawBuffer.slice(0, totalPixels);
-    rawBuffer.set(grayScaleData);
+  if (cls === HEATMAP && channels > 1) {
+    // recast to mono channel because we don't need the other channels
+    targetsBuffer = recastBufferToMonoChannel(
+      imageData.data,
+      width,
+      height,
+      channels
+    );
+    channels = 1;
   }
 
   return {
-    buffer: imageData.data.buffer,
+    buffer: targetsBuffer,
     channels,
     arrayType: "Uint8ClampedArray",
     shape: [height, width],

app/packages/looker/src/worker/disk-overlay-decoder.test.ts

@@ -1,10 +1,10 @@
 import { getSampleSrc } from "@fiftyone/state/src/recoil/utils";
-import { DETECTIONS, HEATMAP } from "@fiftyone/utilities";
+import { DETECTIONS, HEATMAP, SEGMENTATION } from "@fiftyone/utilities";
 import { beforeEach, describe, expect, it, vi } from "vitest";
 import type { Coloring, CustomizeColor } from "..";
 import { LabelMask } from "../overlays/base";
 import type { Colorscale } from "../state";
-import { decodeWithCanvas } from "./canvas-decoder";
+import { decodeWithCanvas, recastBufferToMonoChannel } from "./canvas-decoder";
 import { decodeOverlayOnDisk, IntermediateMask } from "./disk-overlay-decoder";
 import { enqueueFetch } from "./pooled-fetch";
 
@@ -16,9 +16,16 @@ vi.mock("./pooled-fetch", () => ({
   enqueueFetch: vi.fn(),
 }));
 
-vi.mock("./canvas-decoder", () => ({
-  decodeWithCanvas: vi.fn(),
-}));
+vi.mock("./canvas-decoder", async () => {
+  const actual = await vi.importActual<typeof import("./canvas-decoder")>(
+    "./canvas-decoder"
+  );
+
+  return {
+    ...actual,
+    decodeWithCanvas: vi.fn(),
+  };
+});
 
 const COLORING = {} as Coloring;
 const COLOR_SCALE = {} as Colorscale;
@@ -85,7 +92,7 @@ describe("decodeOverlayOnDisk", () => {
       url: sampleSrcUrl,
       options: { priority: "low" },
     });
-    expect(decodeWithCanvas).toHaveBeenCalledWith(mockBlob);
+    expect(decodeWithCanvas).toHaveBeenCalledWith(mockBlob, SEGMENTATION);
     expect(label.mask).toBeDefined();
     expect(label.mask.data).toBe(overlayMask);
     expect(label.mask.image).toBeInstanceOf(ArrayBuffer);
@@ -121,7 +128,7 @@ describe("decodeOverlayOnDisk", () => {
       url: sampleSrcUrl,
       options: { priority: "low" },
     });
-    expect(decodeWithCanvas).toHaveBeenCalledWith(mockBlob);
+    expect(decodeWithCanvas).toHaveBeenCalledWith(mockBlob, HEATMAP);
     expect(label.map).toBeDefined();
     expect(label.map.data).toBe(overlayMask);
     expect(label.map.image).toBeInstanceOf(ArrayBuffer);
@@ -203,3 +210,77 @@ describe("decodeOverlayOnDisk", () => {
     expect(label.mask).toBeNull();
   });
 });
+
+describe("recastBufferToMonoChannel", () => {
+  it("should handle a single grayscale pixel without alpha (stride=1)", () => {
+    const input = new Uint8ClampedArray([128]);
+    const width = 1;
+    const height = 1;
+    const stride = 1;
+
+    const resultBuffer = recastBufferToMonoChannel(
+      input,
+      width,
+      height,
+      stride
+    );
+    const resultArray = new Uint8Array(resultBuffer);
+
+    expect(resultArray).toEqual(new Uint8Array([128]));
+  });
+
+  it("should handle stride=4 (e.g., RGBA)", () => {
+    // two pixels, each RGBA For example:
+    //  pixel 1: R=10, G=10, B=10, A=255
+    //  pixel 2: R=40, G=40, B=40, A=255
+    const input = new Uint8ClampedArray([10, 10, 10, 255, 40, 40, 40, 255]);
+    const width = 2;
+    const height = 1;
+    const stride = 4;
+
+    const resultBuffer = recastBufferToMonoChannel(
+      input,
+      width,
+      height,
+      stride
+    );
+    const resultArray = new Uint8Array(resultBuffer);
+    expect(resultArray).toEqual(new Uint8Array([10, 40]));
+  });
+
+  it("should handle stride=3 (e.g., RGB without alpha)", () => {
+    // two pixels, each RGB (no alpha). For example:
+    //  pixel 1: R=10, G=10, B=10
+    //  pixel 2: R=40, G=40, B=40
+    const input = new Uint8ClampedArray([10, 10, 10, 40, 40, 40]);
+    const width = 2;
+    const height = 1;
+    const stride = 3;
+
+    const resultBuffer = recastBufferToMonoChannel(
+      input,
+      width,
+      height,
+      stride
+    );
+    const resultArray = new Uint8Array(resultBuffer);
+    expect(resultArray).toEqual(new Uint8Array([10, 40]));
+  });
+
+  it("should return an empty buffer if width or height is zero", () => {
+    const input = new Uint8ClampedArray([1, 2, 3, 4]);
+    const width = 0;
+    const height = 1;
+    const stride = 4;
+
+    const resultBuffer = recastBufferToMonoChannel(
+      input,
+      width,
+      height,
+      stride
+    );
+    const resultArray = new Uint8Array(resultBuffer);
+
+    expect(resultArray.length).toBe(0);
+  });
+});

app/packages/looker/src/worker/disk-overlay-decoder.ts

@@ -1,9 +1,14 @@
 import { getSampleSrc } from "@fiftyone/state/src/recoil/utils";
-import { DETECTION, DETECTIONS } from "@fiftyone/utilities";
+import {
+  DETECTION,
+  DETECTIONS,
+  HEATMAP,
+  SEGMENTATION,
+} from "@fiftyone/utilities";
 import { Coloring, CustomizeColor } from "..";
 import { OverlayMask } from "../numpy";
 import { Colorscale } from "../state";
-import { decodeWithCanvas } from "./canvas-decoder";
+import { decodeWithCanvas, recastBufferToMonoChannel } from "./canvas-decoder";
 import { enqueueFetch } from "./pooled-fetch";
 import { getOverlayFieldFromCls } from "./shared";
 
@@ -114,7 +119,7 @@ export const decodeOverlayOnDisk = async (
   let overlayMask: OverlayMask;
 
   try {
-    overlayMask = await decodeWithCanvas(overlayImageBlob);
+    overlayMask = await decodeWithCanvas(overlayImageBlob, cls);
   } catch (e) {
     console.error(e);
     return;