Merge pull request #5406 from voxel51/fix/segmentation-jpg

"fix" jpg segmentations
voxel51 · Jan 22, 2025 · 67bdf7e · 67bdf7e
2 parents 7166c5b + 8e0e047
commit 67bdf7e
Show file tree

Hide file tree

Showing 3 changed files with 54 additions and 12 deletions.
diff --git a/app/packages/looker/src/worker/canvas-decoder.ts b/app/packages/looker/src/worker/canvas-decoder.ts
@@ -1,5 +1,7 @@
-import { HEATMAP } from "@fiftyone/utilities";
+import { HEATMAP, SEGMENTATION } from "@fiftyone/utilities";
+import { Coloring } from "..";
 import { OverlayMask } from "../numpy";
+import { isRgbMaskTargets } from "../overlays/util";
 
 const canvasAndCtx = (() => {
   if (typeof OffscreenCanvas !== "undefined") {
@@ -67,7 +69,12 @@ export const recastBufferToMonoChannel = (
   return uint8Array.slice(0, totalPixels).buffer;
 };
 
-export const decodeWithCanvas = async (blob: Blob, cls: string) => {
+export const decodeWithCanvas = async (
+  blob: Blob,
+  cls: string,
+  field: string,
+  coloring: Coloring
+) => {
   let channels: number = 4;
 
   if (blob.type === "image/png") {
@@ -121,6 +128,31 @@ export const decodeWithCanvas = async (blob: Blob, cls: string) => {
     channels = 1;
   }
 
+  // if it's segmentation, we need to recast according to whether or not this field is mapped to RGB targets
+  if (cls === SEGMENTATION) {
+    let maskTargets = coloring.maskTargets?.[field];
+    if (maskTargets === undefined) {
+      maskTargets = coloring.defaultMaskTargets;
+    }
+    const isRgbMaskTargets_ = isRgbMaskTargets(maskTargets);
+
+    if (!isRgbMaskTargets_ && channels > 1) {
+      // recast to mono channel because we don't need the other channels
+      targetsBuffer = recastBufferToMonoChannel(
+        imageData.data,
+        width,
+        height,
+        channels
+      );
+      channels = 1;
+    }
+
+    // note: for JPG segmentations with RGB mask targets, we don't need to recast
+    // although depending on the JPG compression, we might have some artifacts.
+    // even the slightest change in color can cause the mask to be rendered as
+    // background color (transparent) instead of the actual mask color
+  }
+
   return {
     buffer: targetsBuffer,
     channels,

diff --git a/app/packages/looker/src/worker/disk-overlay-decoder.test.ts b/app/packages/looker/src/worker/disk-overlay-decoder.test.ts
@@ -92,7 +92,12 @@ describe("decodeOverlayOnDisk", () => {
       url: sampleSrcUrl,
       options: { priority: "low" },
     });
-    expect(decodeWithCanvas).toHaveBeenCalledWith(mockBlob, SEGMENTATION);
+    expect(decodeWithCanvas).toHaveBeenCalledWith(
+      mockBlob,
+      SEGMENTATION,
+      field,
+      COLORING
+    );
     expect(label.mask).toBeDefined();
     expect(label.mask.data).toBe(overlayMask);
     expect(label.mask.image).toBeInstanceOf(ArrayBuffer);
@@ -128,7 +133,12 @@ describe("decodeOverlayOnDisk", () => {
       url: sampleSrcUrl,
       options: { priority: "low" },
     });
-    expect(decodeWithCanvas).toHaveBeenCalledWith(mockBlob, HEATMAP);
+    expect(decodeWithCanvas).toHaveBeenCalledWith(
+      mockBlob,
+      HEATMAP,
+      field,
+      COLORING
+    );
     expect(label.map).toBeDefined();
     expect(label.map.data).toBe(overlayMask);
     expect(label.map.image).toBeInstanceOf(ArrayBuffer);

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