fix grayscale segmentation regression + RGB masks recoloring issue

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

@@ -0,0 +1,43 @@
+import { describe, expect, it } from "vitest";
+import { isGrayscale } from "./canvas-decoder";
+
+const createData = (
+  pixels: Array<[number, number, number, number]>
+): Uint8ClampedArray => {
+  return new Uint8ClampedArray(pixels.flat());
+};
+
+describe("isGrayscale", () => {
+  it("should return true for a perfectly grayscale image", () => {
+    const data = createData(Array(100).fill([100, 100, 100, 255]));
+    expect(isGrayscale(data)).toBe(true);
+  });
+
+  it("should return false if alpha is not 255", () => {
+    const data = createData([
+      [100, 100, 100, 255],
+      [100, 100, 100, 254],
+      ...Array(98).fill([100, 100, 100, 255]),
+    ]);
+    expect(isGrayscale(data)).toBe(false);
+  });
+
+  it("should return false if any pixel is not grayscale", () => {
+    const data = createData([
+      [100, 100, 100, 255],
+      [100, 101, 100, 255],
+      ...Array(98).fill([100, 100, 100, 255]),
+    ]);
+    expect(isGrayscale(data)).toBe(false);
+  });
+
+  it("should detect a non-grayscale pixel placed deep enough to ensure at least 1% of pixels are checked", () => {
+    // large image: 100,000 pixels. 1% of 100,000 is 1,000.
+    // the function will check at least 1,000 pixels.
+    // place a non-grayscale pixel after 800 pixels.
+    const pixels = Array(100000).fill([50, 50, 50, 255]);
+    pixels[800] = [50, 51, 50, 255]; // this is within the first 1% of pixels
+    const data = createData(pixels);
+    expect(isGrayscale(data)).toBe(false);
+  });
+});

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

@@ -1,5 +1,26 @@
 import { OverlayMask } from "../numpy";
 
+/**
+ * Checks if the given pixel data is grayscale by sampling a subset of pixels.
+ * The function will check at least 500 pixels or 1% of all pixels, whichever is larger.
+ * If the image is grayscale, the R, G, and B channels will be equal for all sampled pixels,
+ * and the alpha channel will always be 255.
+ */
+export const isGrayscale = (data: Uint8ClampedArray): boolean => {
+  const totalPixels = data.length / 4;
+  const checks = Math.max(500, Math.floor(totalPixels * 0.01));
+  const step = Math.max(1, Math.floor(totalPixels / checks));
+
+  for (let p = 0; p < totalPixels; p += step) {
+    const i = p * 4;
+    const [r, g, b, a] = [data[i], data[i + 1], data[i + 2], data[i + 3]];
+    if (a !== 255 || r !== g || g !== b) {
+      return false;
+    }
+  }
+  return true;
+};
+
 /**
  * Decodes a given image source into an OverlayMask using an OffscreenCanvas
  */
@@ -12,25 +33,26 @@ export const decodeWithCanvas = async (blob: ImageBitmapSource) => {
   const ctx = canvas.getContext("2d");
 
   ctx.drawImage(imageBitmap, 0, 0);
+  imageBitmap.close();
 
   const imageData = ctx.getImageData(0, 0, width, height);
 
-  const numChannels = imageData.data.length / (width * height);
-
-  const overlayData = {
-    width,
-    height,
-    data: imageData.data,
-    channels: numChannels,
-  };
+  // for nongrayscale images, channel is guaranteed to be 4 (RGBA)
+  const channels = isGrayscale(imageData.data) ? 1 : 4;
 
-  // dispose
-  imageBitmap.close();
+  if (channels === 1) {
+    // get rid of the G, B, and A channels, new buffer will be 1/4 the size
+    const data = new Uint8ClampedArray(width * height);
+    for (let i = 0; i < data.length; i++) {
+      data[i] = imageData.data[i * 4];
+    }
+    imageData.data.set(data);
+  }
 
   return {
-    buffer: overlayData.data.buffer,
-    channels: numChannels,
-    arrayType: overlayData.data.constructor.name as OverlayMask["arrayType"],
-    shape: [overlayData.height, overlayData.width],
+    buffer: imageData.data.buffer,
+    channels,
+    arrayType: "Uint8ClampedArray",
+    shape: [height, width],
   } as OverlayMask;
 };