Add frontend support for loading overlay masks from disk

app/packages/app/package.json

@@ -71,6 +71,7 @@
     "typescript": "^4.7.4",
     "typescript-plugin-css-modules": "^3.4.0",
     "vite": "^3.0.0",
-    "vite-plugin-relay": "^1.0.7"
+    "vite-plugin-relay": "^1.0.7",
+    "vite-plugin-rewrite-all": "^1.0.0"
   }
 }

app/packages/app/vite.config.ts

@@ -2,6 +2,7 @@ import reactRefresh from "@vitejs/plugin-react-refresh";
 import nodePolyfills from "rollup-plugin-polyfill-node";
 import { defineConfig } from "vite";
 import relay from "vite-plugin-relay";
+import pluginRewriteAll from "vite-plugin-rewrite-all";
 
 export default defineConfig(({ mode }) => {
   return {
@@ -12,6 +13,8 @@ export default defineConfig(({ mode }) => {
       }),
       relay,
       nodePolyfills(),
+      // pluginRewriteAll to address this vite bug: https://github.com/vitejs/vite/issues/2415
+      pluginRewriteAll(),
     ],
     server: {
       proxy: {

app/packages/looker/package.json

@@ -20,16 +20,19 @@
   "dependencies": {
     "@ungap/event-target": "^0.2.2",
     "copy-to-clipboard": "^3.3.1",
+    "fast-png": "^6.1.0",
     "immutable": "^4.0.0-rc.12",
     "json-format-highlight": "^1.0.4",
     "lru-cache": "^6.0.0",
     "mime": "^2.5.2",
     "uuid": "^8.3.2"
   },
   "devDependencies": {
+    "@rollup/plugin-inject": "^5.0.2",
     "@types/color-string": "^1.5.0",
     "@types/lru-cache": "^5.1.0",
     "@types/uuid": "^8.3.0",
+    "buffer": "^6.0.3",
     "prettier": "^2.7.1",
     "typescript": "^4.7.4",
     "typescript-plugin-css-modules": "^3.4.0",

app/packages/looker/src/index.ts

@@ -1,9 +1,6 @@
 /**
  * Copyright 2017-2022, Voxel51, Inc.
  */
-import LRU from "lru-cache";
-import { v4 as uuid } from "uuid";
-
 import {
   AppError,
   DATE_FIELD,
@@ -16,16 +13,18 @@ import {
   Schema,
   withPath,
 } from "@fiftyone/utilities";
+import LRU from "lru-cache";
+import { v4 as uuid } from "uuid";
 
 import {
+  BASE_ALPHA,
+  CHUNK_SIZE,
+  DASH_LENGTH,
   FONT_SIZE,
-  STROKE_WIDTH,
+  MAX_FRAME_CACHE_SIZE_BYTES,
   PAD,
   POINT_RADIUS,
-  MAX_FRAME_CACHE_SIZE_BYTES,
-  CHUNK_SIZE,
-  DASH_LENGTH,
-  BASE_ALPHA,
+  STROKE_WIDTH,
 } from "./constants";
 import {
   getFrameElements,
@@ -37,29 +36,29 @@ import {
   LookerElement,
   VIDEO_SHORTCUTS,
 } from "./elements/common";
-import processOverlays from "./processOverlays";
 import { ClassificationsOverlay, loadOverlays } from "./overlays";
 import { CONTAINS, Overlay } from "./overlays/base";
+import processOverlays from "./processOverlays";
 import {
-  FrameState,
-  ImageState,
-  VideoState,
-  StateUpdate,
   BaseState,
+  BufferRange,
+  Buffers,
+  Coloring,
+  Coordinates,
   DEFAULT_FRAME_OPTIONS,
   DEFAULT_IMAGE_OPTIONS,
   DEFAULT_VIDEO_OPTIONS,
-  Coordinates,
-  Optional,
+  Dimensions,
   FrameChunkResponse,
-  VideoSample,
   FrameSample,
-  Buffers,
+  FrameState,
+  ImageState,
   LabelData,
-  BufferRange,
-  Dimensions,
+  Optional,
   Sample,
-  Coloring,
+  StateUpdate,
+  VideoSample,
+  VideoState,
 } from "./state";
 import {
   addToBuffers,
@@ -75,9 +74,9 @@ import {
 
 import { zoomToContent } from "./zoom";
 
-import { getFrameNumber } from "./elements/util";
 import { getColor } from "@fiftyone/utilities";
 import { Events } from "./elements/base";
+import { getFrameNumber } from "./elements/util";
 
 export { createColorGenerator, getRGB } from "@fiftyone/utilities";
 export { freeVideos } from "./elements/util";
@@ -91,7 +90,6 @@ export type {
   VideoConfig,
   VideoOptions,
 } from "./state";
-
 export { zoomAspectRatio } from "./zoom";
 
 export type RGB = [number, number, number];
@@ -646,7 +644,7 @@ export abstract class Looker<
 
   private loadSample(sample: Sample) {
     const messageUUID = uuid();
-    const worker = getLabelsWorker((event, detail) =>
+    const labelsWorker = getLabelsWorker((event, detail) =>
       this.dispatchEvent(event, detail)
     );
     const listener = ({ data: { sample, uuid } }) => {
@@ -658,12 +656,12 @@ export abstract class Looker<
           disabled: false,
           reloading: false,
         });
-        worker.removeEventListener("message", listener);
+        labelsWorker.removeEventListener("message", listener);
       }
     };
-    worker.addEventListener("message", listener);
+    labelsWorker.addEventListener("message", listener);
 
-    worker.postMessage({
+    labelsWorker.postMessage({
       method: "processSample",
       coloring: this.state.options.coloring,
       sample,

app/packages/looker/src/numpy.ts

@@ -2,6 +2,7 @@
  * Copyright 2017-2022, Voxel51, Inc.
  */
 
+import { Buffer } from "buffer";
 import pako from "./pako.js";
 
 export { deserialize };
@@ -164,8 +165,8 @@ function parse(array: Uint8Array): NumpyResult {
 }
 
 /**
- * Deserializes and parses a saved numpy array
+ * Deserializes and parses a base64 encoded numpy array
  */
-function deserialize(str: string): NumpyResult {
-  return parse(pako.inflate(atob(str)));
+function deserialize(compressedBase64Array: string): NumpyResult {
+  return parse(pako.inflate(Buffer.from(compressedBase64Array, "base64")));
 }

app/packages/looker/src/util.ts

@@ -1,9 +1,16 @@
 /**
  * Copyright 2017-2022, Voxel51, Inc.
  */
-import mime from "mime";
 import { mergeWith } from "immutable";
+import mime from "mime";
 
+import {
+  AppError,
+  getFetchParameters,
+  GraphQLError,
+  NetworkError,
+  ServerError,
+} from "@fiftyone/utilities";
 import { MIN_PIXELS } from "./constants";
 import {
   BaseState,
@@ -15,15 +22,7 @@ import {
   DispatchEvent,
   Optional,
 } from "./state";
-
 import LookerWorker from "./worker.ts?worker&inline";
-import {
-  AppError,
-  getFetchParameters,
-  GraphQLError,
-  NetworkError,
-  ServerError,
-} from "@fiftyone/utilities";
 
 /**
  * Shallow data-object comparison for equality

app/packages/looker/src/worker.ts

@@ -2,16 +2,22 @@
  * Copyright 2017-2022, Voxel51, Inc.
  */
 
+import { getSampleSrc } from "@fiftyone/state/src/recoil/utils";
 import {
-  LABEL_LIST,
-  VALID_LABEL_TYPES,
+  DETECTION,
+  DETECTIONS,
+  get32BitColor,
   getFetchFunction,
+  HEATMAP,
+  LABEL_LIST,
+  SEGMENTATION,
   setFetchFunction,
   Stage,
-  get32BitColor,
+  VALID_LABEL_TYPES,
 } from "@fiftyone/utilities";
+import { decode as decodePng } from "fast-png";
 import { CHUNK_SIZE } from "./constants";
-import { ARRAY_TYPES, deserialize } from "./numpy";
+import { ARRAY_TYPES, deserialize, NumpyResult } from "./numpy";
 import { Coloring, FrameChunk } from "./state";
 
 interface ResolveColor {
@@ -126,9 +132,76 @@ const mapId = (obj) => {
   return obj;
 };
 
-const LABELS = new Set(VALID_LABEL_TYPES);
+const ALL_VALID_LABELS = new Set(VALID_LABEL_TYPES);
+
+const LABELS_THAT_CAN_HAVE_OVERLAYS_ON_DISK = new Set([
+  SEGMENTATION,
+  HEATMAP,
+  DETECTION,
+  DETECTIONS,
+]);
 
-const processLabels = (
+/**
+ * Some label types (example: segmentation, heatmap) can have their overlay data stored on-disk,
+ * we want to impute the relevant mask property of these labels from what's stored in the disk
+ */
+const imputeOverlayFromPath = async (
+  label: Record<string, any>,
+  buffers: ArrayBuffer[]
+) => {
+  // handle all list types here
+  if (label._cls === DETECTIONS) {
+    label.detections.forEach((detection) =>
+      imputeOverlayFromPath(detection, buffers)
+    );
+    return;
+  }
+
+  // overlay path is in `map_path` property for heatmap, or else, it's in `mask_path` property (for segmentation or detection)
+  const overlayPathField = label._cls === HEATMAP ? "map_path" : "mask_path";
+  const overlayField = overlayPathField === "map_path" ? "map" : "mask";
+
+  if (
+    Object.hasOwn(label, overlayField) ||
+    !Object.hasOwn(label, overlayPathField)
+  ) {
+    // nothing to be done
+    return;
+  }
+
+  // convert absolute file path to a URL that we can "fetch" from
+  const overlayPngImageUrl = getSampleSrc(label[overlayPathField] as string);
+
+  const pngArrayBuffer: ArrayBuffer = await getFetchFunction()(
+    "GET",
+    overlayPngImageUrl,
+    null,
+    "arrayBuffer"
+  );
+  const overlayData = decodePng(pngArrayBuffer);
+
+  const width = overlayData.width;
+  const height = overlayData.height;
+
+  // frame what we have as a specialized type `NumpyResult` that's used in downstream while processing overlays
+  const data: NumpyResult = {
+    shape: [height, width],
+    buffer: overlayData.data.buffer,
+    arrayType: overlayData.data.constructor.name as NumpyResult["arrayType"],
+  };
+
+  // set the `mask` property for this label
+  label[overlayField] = {
+    data,
+    image: new ArrayBuffer(width * height * 4),
+  };
+
+  // transfer buffers
+  buffers.push(data.buffer);
+  buffers.push(label[overlayField].image);
+};
+
+const processLabels = async (
   sample: { [key: string]: any },
   coloring: Coloring,
   prefix: string = ""
@@ -142,11 +215,15 @@ const processLabels = (
       continue;
     }
 
+    if (LABELS_THAT_CAN_HAVE_OVERLAYS_ON_DISK.has(label._cls)) {
+      await imputeOverlayFromPath(label, buffers);
+    }
+
     if (label._cls in DESERIALIZE) {
       DESERIALIZE[label._cls](label, buffers);
     }
 
-    if (LABELS.has(label._cls)) {
+    if (ALL_VALID_LABELS.has(label._cls)) {
       if (label._cls in LABEL_LIST) {
         const list = label[LABEL_LIST[label._cls]];
         if (Array.isArray(list)) {

app/packages/looker/vite.config.ts

@@ -1,4 +1,5 @@
 import { UserConfig } from "vite";
+import inject from "@rollup/plugin-inject";
 
 export default <UserConfig>{
   esbuild: true,
@@ -7,6 +8,9 @@ export default <UserConfig>{
       entry: "src/index.ts",
       formats: ["es"],
     },
+    rollupOptions: {
+      plugins: [inject({ Buffer: ["Buffer", "Buffer"] })],
+    },
     target: "es2015",
     minify: false,
   },

app/packages/utilities/src/fetch.ts

@@ -17,7 +17,7 @@ export interface FetchFunction {
     method: string,
     path: string,
     body?: A,
-    result?: "json" | "blob",
+    result?: "json" | "blob" | "arrayBuffer",
     retries?: number,
     retryDelay?: number
   ): Promise<R>;
@@ -32,13 +32,18 @@ export const getFetchHeaders = () => {
 };
 
 export const getFetchOrigin = () => {
-  if (window.FIFTYONE_SERVER_ADDRESS) {
+  // window is not defined in the web worker
+  if (typeof window !== "undefined" && window.FIFTYONE_SERVER_ADDRESS) {
     return window.FIFTYONE_SERVER_ADDRESS;
   }
   return fetchOrigin;
 };
 export function getFetchPathPrefix(): string {
-  if (typeof window.FIFTYONE_SERVER_PATH_PREFIX === "string") {
+  // window is not defined in the web worker
+  if (
+    typeof window !== "undefined" &&
+    typeof window.FIFTYONE_SERVER_PATH_PREFIX === "string"
+  ) {
     return window.FIFTYONE_SERVER_PATH_PREFIX;
   }
   return "";