Merge pull request #5120 from voxel51/feat/on-disk-detection-mask-path

add support for detection.mask_path

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

@@ -5,6 +5,7 @@
 import { getSampleSrc } from "@fiftyone/state/src/recoil/utils";
 import {
   DENSE_LABELS,
+  DETECTION,
   DETECTIONS,
   DYNAMIC_EMBEDDED_DOCUMENT,
   EMBEDDED_DOCUMENT,
@@ -110,18 +111,24 @@ const imputeOverlayFromPath = async (
 ) => {
   // handle all list types here
   if (cls === DETECTIONS) {
-    label?.detections?.forEach((detection) =>
-      imputeOverlayFromPath(
-        field,
-        detection,
-        coloring,
-        customizeColorSetting,
-        colorscale,
-        buffers,
-        {},
-        cls
-      )
-    );
+    const promises = [];
+    for (const detection of label.detections) {
+      promises.push(
+        imputeOverlayFromPath(
+          field,
+          detection,
+          coloring,
+          customizeColorSetting,
+          colorscale,
+          buffers,
+          {},
+          DETECTION
+        )
+      );
+    }
+    // if some paths fail to load, it's okay, we can still proceed
+    // hence we use `allSettled` instead of `all`
+    await Promise.allSettled(promises);
     return;
   }
 
@@ -150,27 +157,14 @@ const imputeOverlayFromPath = async (
     baseUrl = overlayImageUrl.split("?")[0];
   }
 
-  const fileExtension = baseUrl.split(".").pop();
-
-  const overlayImageBuffer: ArrayBuffer = await getFetchFunction()(
+  const overlayImageBuffer: Blob = await getFetchFunction()(
     "GET",
     overlayImageUrl,
     null,
-    "arrayBuffer"
+    "blob"
   );
 
-  const mimeTypes = {
-    png: "image/png",
-    jpg: "image/jpeg",
-    jpeg: "image/jpeg",
-    gif: "image/gif",
-    bmp: "image/bmp",
-  };
-  const blobType =
-    mimeTypes[fileExtension.toLowerCase()] || "application/octet-stream";
-  const blob = new Blob([overlayImageBuffer], { type: blobType });
-
-  const overlayMask = await decodeWithCanvas(blob);
+  const overlayMask = await decodeWithCanvas(overlayImageBuffer);
   const [overlayHeight, overlayWidth] = overlayMask.shape;
 
   // set the `mask` property for this label
@@ -211,16 +205,20 @@ const processLabels = async (
       }
 
       if (DENSE_LABELS.has(cls)) {
-        await imputeOverlayFromPath(
-          `${prefix || ""}${field}`,
-          label,
-          coloring,
-          customizeColorSetting,
-          colorscale,
-          buffers,
-          sources,
-          cls
-        );
+        try {
+          await imputeOverlayFromPath(
+            `${prefix || ""}${field}`,
+            label,
+            coloring,
+            customizeColorSetting,
+            colorscale,
+            buffers,
+            sources,
+            cls
+          );
+        } catch (e) {
+          console.error("Couldn't decode overlay image from disk: ", e);
+        }
       }
 
       if (cls in DeserializerFactory) {

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

@@ -119,10 +119,25 @@ export const PainterFactory = (requestColor) => ({
     );
     const bitColor = get32BitColor(color);
 
-    // these for loops must be fast. no "in" or "of" syntax
-    for (let i = 0; i < overlay.length; i++) {
-      if (targets[i]) {
-        overlay[i] = bitColor;
+    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) {
+        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;
+          overlay[overlayIndex] = bitColor;
+        }
+      }
+    } else {
+      for (let i = 0; i < overlay.length; i++) {
+        if (targets[i]) {
+          overlay[i] = bitColor;
+        }
       }
     }
   },

docs/source/user_guide/using_datasets.rst

@@ -2556,25 +2556,41 @@ Instance segmentations
 ----------------------
 
 Object detections stored in |Detections| may also have instance segmentation
-masks, which should be stored in the
-:attr:`mask <fiftyone.core.labels.Detection.mask>` attribute of each
-|Detection|.
+masks.
 
-The mask must be a 2D numpy array containing either booleans or 0/1 integers
-encoding the extent of the instance mask within the
+These masks can be stored in one of two ways: either directly in the database
+via the :attr:`mask<fiftyone.core.labels.Detection.mask>` attribute, or on
+disk referenced by the
+:attr:`mask_path <fiftyone.core.labels.Detection.mask_path>` attribute.
+
+Masks stored directly in the database must be 2D numpy arrays
+containing either booleans or 0/1 integers that encode the extent of the
+instance mask within the
 :attr:`bounding_box <fiftyone.core.labels.Detection.bounding_box>` of the
-object. The array can be of any size; it is stretched as necessary to fill the
+object.
+
+For masks stored on disk, the
+:attr:`mask_path <fiftyone.core.labels.Detection.mask_path>` attribute should
+contain the file path to the mask image. We recommend storing masks as
+single-channel PNG images, where a pixel value of 0 indicates the
+background (rendered as transparent in the App), and any other 
+value indicates the object.
+
+Masks can be of any size; they are stretched as necessary to fill the
 object's bounding box when visualizing in the App.
 
 .. code-block:: python
     :linenos:
 
     import numpy as np
+    from PIL import Image
 
     import fiftyone as fo
 
     # Example instance mask
-    mask = (np.random.randn(32, 32) > 0)
+    mask = ((np.random.randn(32, 32) > 0) * 255).astype(np.uint8)
+    mask_path = "/path/to/mask.png"
+    Image.fromarray(mask).save(mask_path)
 
     sample = fo.Sample(filepath="/path/to/image.png")
 
@@ -2583,7 +2599,7 @@ object's bounding box when visualizing in the App.
             fo.Detection(
                 label="cat",
                 bounding_box=[0.480, 0.513, 0.397, 0.288],
-                mask=mask,
+                mask_path=mask_path,
                 confidence=0.96,
             ),
         ]
@@ -2608,13 +2624,7 @@ object's bounding box when visualizing in the App.
                     'attributes': {},
                     'label': 'cat',
                     'bounding_box': [0.48, 0.513, 0.397, 0.288],
-                    'mask': array([[False,  True, False, ...,  True,  True, False],
-                           [ True, False,  True, ..., False,  True,  True],
-                           [False,  True, False, ..., False,  True, False],
-                           ...,
-                           [ True,  True, False, ..., False, False,  True],
-                           [ True,  True,  True, ...,  True,  True, False],
-                           [False,  True,  True, ..., False,  True,  True]]),
+                    'mask_path': '/path/to/mask.png',
                     'confidence': 0.96,
                     'index': None,
                 }>,
@@ -2634,7 +2644,7 @@ by dynamically adding new fields to each |Detection| instance:
     detection = fo.Detection(
         label="cat",
         bounding_box=[0.5, 0.5, 0.4, 0.3],
-        mask=np.random.randn(32, 32) > 0,
+        mask_path="/path/to/mask.png",
         age=51,  # custom attribute
         mood="salty",  # custom attribute
     )
@@ -2649,13 +2659,7 @@ by dynamically adding new fields to each |Detection| instance:
         'tags': [],
         'label': 'cat',
         'bounding_box': [0.5, 0.5, 0.4, 0.3],
-        'mask': array([[False, False,  True, ...,  True,  True, False],
-               [ True,  True, False, ...,  True, False,  True],
-               [False, False,  True, ..., False, False, False],
-               ...,
-               [False, False,  True, ...,  True,  True, False],
-               [ True, False,  True, ...,  True, False,  True],
-               [False,  True, False, ...,  True,  True,  True]]),
+        'mask_path': '/path/to/mask.png',
         'confidence': None,
         'index': None,
         'age': 51,

e2e-pw/src/oss/specs/overlays/detection-mask.spec.ts

@@ -1,10 +1,17 @@
-import { test as base } from "src/oss/fixtures";
+import { test as base, expect } from "src/oss/fixtures";
 import { GridPom } from "src/oss/poms/grid";
 import { ModalPom } from "src/oss/poms/modal";
 import { getUniqueDatasetNameWithPrefix } from "src/oss/utils";
 
 const datasetName = getUniqueDatasetNameWithPrefix("detection-mask");
-const testImgPath = "/tmp/detection-mask-img.png";
+
+const colors = ["#ff0000", "#00ff00", "#0000ff"];
+
+const badDetectionMaskSampleImage = "/tmp/detection-bad-mask-img.png";
+const goodDetectionMaskSampleImage = "/tmp/detection-good-mask-img.png";
+const goodDetectionMaskPathSampleImage = "/tmp/detection-mask-path-img.png";
+
+const goodDetectionMaskOnDisk = "/tmp/detection-mask-on-disk.png";
 
 const test = base.extend<{ modal: ModalPom; grid: GridPom }>({
   modal: async ({ page, eventUtils }, use) => {
@@ -16,22 +23,37 @@ const test = base.extend<{ modal: ModalPom; grid: GridPom }>({
 });
 
 test.beforeAll(async ({ fiftyoneLoader, mediaFactory }) => {
-  await mediaFactory.createBlankImage({
-    outputPath: testImgPath,
-    width: 25,
-    height: 25,
-  });
+  await Promise.all(
+    [
+      badDetectionMaskSampleImage,
+      goodDetectionMaskSampleImage,
+      goodDetectionMaskPathSampleImage,
+    ].map((img, index) => {
+      const fillColor = colors[index];
+      mediaFactory.createBlankImage({
+        outputPath: img,
+        width: 25,
+        height: 25,
+        fillColor: fillColor,
+      });
+    })
+  );
 
   await fiftyoneLoader.executePythonCode(
     `
         import fiftyone as fo
+        import numpy as np
+
+        from PIL import Image
 
         dataset = fo.Dataset("${datasetName}")
         dataset.persistent = True
-        
-        dataset.add_sample(fo.Sample(filepath="${testImgPath}"))
-        sample = dataset.first()
-        sample["ground_truth"] = fo.Detections(
+
+        samples = []
+
+        # sample with bad detection mask
+        badDetectionMaskSample = fo.Sample(filepath="${badDetectionMaskSampleImage}")
+        badDetectionMaskSample["ground_truth"] = fo.Detections(
             detections=[
                 fo.Detection(
                     label="bad_mask_detection",
@@ -40,7 +62,34 @@ test.beforeAll(async ({ fiftyoneLoader, mediaFactory }) => {
                 ),
             ]
         )
-        sample.save()
+        samples.append(badDetectionMaskSample)
+
+        # sample with good detection mask
+        goodDetectionMaskSample = fo.Sample(filepath="${goodDetectionMaskSampleImage}")
+        goodDetectionMaskSample["ground_truth"] = fo.Detections(
+            detections=[
+                fo.Detection(
+                    label="good_mask_detection",
+                    bounding_box=[0.0, 0.0, 0.5, 0.5],
+                    mask=np.ones((15, 15)),
+                ),
+            ]
+        )
+        samples.append(goodDetectionMaskSample)
+
+        # sample with good detection mask _path_
+        img = Image.fromarray(np.ones((15, 15), dtype=np.uint8))
+        img.save("${goodDetectionMaskOnDisk}")
+
+        goodDetectionMaskPathSample = fo.Sample(filepath="${goodDetectionMaskPathSampleImage}")
+        goodDetectionMaskPathSample["prediction"] = fo.Detection(
+                    label="good_mask_detection_path",
+                    bounding_box=[0.0, 0.0, 0.5, 0.5],
+                    mask_path="${goodDetectionMaskOnDisk}",
+                )
+        samples.append(goodDetectionMaskPathSample)
+        
+        dataset.add_samples(samples)
         `
   );
 });
@@ -50,9 +99,21 @@ test.beforeEach(async ({ page, fiftyoneLoader }) => {
 });
 
 test.describe("detection-mask", () => {
-  test("should load empty mask fine", async ({ grid, modal }) => {
-    await grid.assert.isEntryCountTextEqualTo("1 sample");
+  test("should load all masks fine", async ({ grid, modal }) => {
+    await grid.assert.isEntryCountTextEqualTo("3 samples");
+
+    // bad sample, assert it loads in the modal fine, too
     await grid.openFirstSample();
     await modal.waitForSampleLoadDomAttribute();
+
+    // close modal and assert grid screenshot (compares all detections)
+    await modal.close();
+
+    await expect(grid.getForwardSection()).toHaveScreenshot(
+      "grid-detections.png",
+      {
+        animations: "allow",
+      }
+    );
   });
 });

e2e-pw/src/oss/specs/plugins/histograms.spec.ts

@@ -59,6 +59,7 @@ test("histograms panel", async ({ histogram, panel }) => {
     "detections.detections.confidence",
     "detections.detections.index",
     "detections.detections.label",
+    "detections.detections.mask_path",
     "detections.detections.tags",
     "float",
     "int",