Merge pull request #5294 from voxel51/fix/png-decoding

fix png decoding

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

@@ -1,45 +1,65 @@
 import { OverlayMask } from "../numpy";
 
+const PNG_SIGNATURE = [0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a];
 /**
- * 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.
+ * Reads the PNG's image header chunk to determine the color type.
+ * Returns the color type if PNG, otherwise undefined.
  */
-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;
+const getPngcolorType = async (blob: Blob): Promise<number | undefined> => {
+  // https://www.w3.org/TR/2003/REC-PNG-20031110/#11IHDR
+
+  // PNG signature is 8 bytes
+  // IHDR (image header): length(4 bytes), chunk type(4 bytes), then data(13 bytes)
+  // data layout of IHDR: width(4), height(4), bit depth(1), color type(1), ...
+  // color type is at offset: 8(signature) + 4(length) + 4(chunk type) + 8(width+height) + 1(bit depth)
+  // = 8 + 4 + 4 + 8 + 1 = 25 (0-based index)
+
+  const header = new Uint8Array(await blob.slice(0, 26).arrayBuffer());
+
+  // check PNG signature
+  for (let i = 0; i < PNG_SIGNATURE.length; i++) {
+    if (header[i] !== PNG_SIGNATURE[i]) {
+      // not a PNG
+      return undefined;
     }
   }
-  return true;
+
+  // color type at byte 25 (0-based)
+  const colorType = header[25];
+  return colorType;
 };
 
-/**
- * Decodes a given image source into an OverlayMask using an OffscreenCanvas
- */
-export const decodeWithCanvas = async (blob: ImageBitmapSource) => {
+export const decodeWithCanvas = async (blob: Blob) => {
+  let channels: number = 4;
+
+  if (blob.type === "image/png") {
+    const colorType = await getPngcolorType(blob);
+    if (colorType !== undefined) {
+      // according to PNG specs:
+      // 0: Grayscale          => 1 channel
+      // 2: Truecolor (RGB)   => (would be 3 channels, but we can safely use 4)
+      // 3: Indexed-color     => (palette-based, treat as non-grayscale => 4)
+      // 4: Grayscale+Alpha    => Grayscale image (so treat as grayscale => 1)
+      // 6: RGBA               => non-grayscale => 4
+      if (colorType === 0 || colorType === 4) {
+        channels = 1;
+      } else {
+        channels = 4;
+      }
+    }
+  }
+  // if not PNG, use 4 channels
+
   const imageBitmap = await createImageBitmap(blob);
-  const width = imageBitmap.width;
-  const height = imageBitmap.height;
+  const { width, height } = imageBitmap;
 
   const canvas = new OffscreenCanvas(width, height);
-  const ctx = canvas.getContext("2d");
-
+  const ctx = canvas.getContext("2d")!;
   ctx.drawImage(imageBitmap, 0, 0);
   imageBitmap.close();
 
   const imageData = ctx.getImageData(0, 0, width, height);
 
-  // for nongrayscale images, channel is guaranteed to be 4 (RGBA)
-  const channels = isGrayscale(imageData.data) ? 1 : 4;
-
   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);

app/packages/looker/src/worker/painter.ts

@@ -113,23 +113,24 @@ export const PainterFactory = (requestColor) => ({
       }
     }
 
+    const numChannels = label.mask.data.channels;
     const overlay = new Uint32Array(label.mask.image);
     const targets = new ARRAY_TYPES[label.mask.data.arrayType](
       label.mask.data.buffer
     );
     const bitColor = get32BitColor(color);
 
     if (label.mask_path) {
-      // putImageData results in an UInt8ClampedArray (for both grayscale or RGB masks),
-      // where each pixel is represented by 4 bytes (RGBA)
-      // it's packed like: [R, G, B, A, R, G, B, A, ...]
-      // use first channel info to determine if the pixel is in the mask
-      // skip second (G), third (B) and fourth (A) channels
-      for (let i = 0; i < targets.length; i += 4) {
+      // putImageData results in an UInt8ClampedArray,
+      // if image is grayscale, it'll be packed as:
+      // [I, I, I, I, I...], where I is the intensity value
+      // or else it'll be packed as:
+      // [R, G, B, A, R, G, B, A...]
+      // for non-grayscale masks, we can check every nth byte,
+      // where n = numChannels, to check for whether or not the pixel is part of the mask
+      for (let i = 0; i < targets.length; i += numChannels) {
         if (targets[i]) {
-          // overlay image is a Uint32Array, where each pixel is represented by 4 bytes (RGBA)
-          // so we need to divide by 4 to get the correct index to assign 32 bit color
-          const overlayIndex = i / 4;
+          const overlayIndex = i / numChannels;
           overlay[overlayIndex] = bitColor;
         }
       }