Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

add support for detection.mask_path #5120

Open
wants to merge 10 commits into
base: develop
Choose a base branch
from
54 changes: 23 additions & 31 deletions app/packages/looker/src/worker/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
import { getSampleSrc } from "@fiftyone/state/src/recoil/utils";
import {
DENSE_LABELS,
DETECTION,
DETECTIONS,
DYNAMIC_EMBEDDED_DOCUMENT,
EMBEDDED_DOCUMENT,
Expand Down Expand Up @@ -110,18 +111,18 @@ const imputeOverlayFromPath = async (
) => {
// handle all list types here
if (cls === DETECTIONS) {
label?.detections?.forEach((detection) =>
imputeOverlayFromPath(
for (const detection of label.detections) {
await imputeOverlayFromPath(
field,
detection,
coloring,
customizeColorSetting,
colorscale,
buffers,
{},
cls
)
);
DETECTION
);
}
return;
}

Expand Down Expand Up @@ -150,27 +151,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
Expand Down Expand Up @@ -211,16 +199,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) {
Expand Down
23 changes: 19 additions & 4 deletions app/packages/looker/src/worker/painter.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}
}
}
},
Expand Down
89 changes: 75 additions & 14 deletions e2e-pw/src/oss/specs/overlays/detection-mask.spec.ts
Original file line number Diff line number Diff line change
@@ -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";

sashankaryal marked this conversation as resolved.
Show resolved Hide resolved
const test = base.extend<{ modal: ModalPom; grid: GridPom }>({
modal: async ({ page, eventUtils }, use) => {
Expand All @@ -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,
});
})
);
sashankaryal marked this conversation as resolved.
Show resolved Hide resolved

sashankaryal marked this conversation as resolved.
Show resolved Hide resolved
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",
Expand All @@ -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)
`
);
});
Expand All @@ -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)
sashankaryal marked this conversation as resolved.
Show resolved Hide resolved
await modal.close();

await expect(grid.getForwardSection()).toHaveScreenshot(
"grid-detections.png",
{
animations: "allow",
}
);
});
sashankaryal marked this conversation as resolved.
Show resolved Hide resolved
sashankaryal marked this conversation as resolved.
Show resolved Hide resolved
});
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
1 change: 1 addition & 0 deletions e2e-pw/src/oss/specs/plugins/histograms.spec.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
60 changes: 58 additions & 2 deletions fiftyone/core/labels.py
Original file line number Diff line number Diff line change
Expand Up @@ -401,18 +401,74 @@ class Detection(_HasAttributesDict, _HasID, Label):
mask (None): an instance segmentation mask for the detection within
its bounding box, which should be a 2D binary or 0/1 integer numpy
array
mask_path (None): the absolute path to the instance segmentation image
on disk
confidence (None): a confidence in ``[0, 1]`` for the detection
index (None): an index for the object
attributes ({}): a dict mapping attribute names to :class:`Attribute`
instances
"""

_MEDIA_FIELD = "mask_path"

label = fof.StringField()
bounding_box = fof.ListField(fof.FloatField())
mask = fof.ArrayField()
mask_path = fof.StringField()
confidence = fof.FloatField()
index = fof.IntField()

@property
def has_mask(self):
"""Whether this instance has a mask."""
return self.mask is not None or self.mask_path is not None

def get_mask(self):
"""Returns the detection mask for this instance.

Returns:
a numpy array, or ``None``
"""
if self.mask is not None:
return self.mask

if self.mask_path is not None:
return _read_mask(self.mask_path)

return None

def import_mask(self, update=False):
"""Imports this instance's mask from disk to its :attr:`mask`
attribute.

Args:
update (False): whether to clear this instance's :attr:`mask_path`
attribute after importing
"""
if self.mask_path is not None:
self.mask = _read_mask(self.mask_path)

if update:
self.mask_path = None

def export_mask(self, outpath, update=False):
"""Exports this instance's mask to the given path.

Args:
outpath: the path to write the mask
update (False): whether to clear this instance's :attr:`mask`
attribute and set its :attr:`mask_path` attribute when
exporting in-database segmentations
"""
if self.mask_path is not None:
etau.copy_file(self.mask_path, outpath)
else:
_write_mask(self.mask, outpath)

if update:
self.mask = None
self.mask_path = outpath

def to_polyline(self, tolerance=2, filled=True):
"""Returns a :class:`Polyline` representation of this instance.

Expand Down Expand Up @@ -467,7 +523,8 @@ def to_segmentation(self, mask=None, frame_size=None, target=255):
Returns:
a :class:`Segmentation`
"""
if self.mask is None:
mask = self.get_mask()
if mask is None:
raise ValueError(
"Only detections with their `mask` attributes populated can "
"be converted to segmentations"
Expand Down Expand Up @@ -1044,7 +1101,6 @@ def import_mask(self, update=False):
attribute.

Args:
outpath: the path to write the map
update (False): whether to clear this instance's :attr:`mask_path`
attribute after importing
"""
Expand Down
4 changes: 2 additions & 2 deletions fiftyone/utils/coco.py
Original file line number Diff line number Diff line change
Expand Up @@ -1304,7 +1304,7 @@ def from_label(
x, y, w, h = label.bounding_box
bbox = [x * width, y * height, w * width, h * height]

if label.mask is not None:
if label.has_mask() is not None:
sashankaryal marked this conversation as resolved.
Show resolved Hide resolved
segmentation = _instance_to_coco_segmentation(
label, frame_size, iscrowd=iscrowd, tolerance=tolerance
)
Expand Down Expand Up @@ -2116,7 +2116,7 @@ def _coco_objects_to_detections(
)

if detection is not None and (
not load_segmentations or detection.mask is not None
not load_segmentations or detection.has_mask() is not None
sashankaryal marked this conversation as resolved.
Show resolved Hide resolved
):
detections.append(detection)

Expand Down
2 changes: 1 addition & 1 deletion fiftyone/utils/cvat.py
Original file line number Diff line number Diff line change
Expand Up @@ -6400,7 +6400,7 @@ def _create_detection_shapes(
}
)
elif label_type in ("instance", "instances"):
if det.mask is None:
if det.has_mask() is None:
continue

polygon = det.to_polyline()
Expand Down
2 changes: 1 addition & 1 deletion fiftyone/utils/eta.py
Original file line number Diff line number Diff line change
Expand Up @@ -596,7 +596,7 @@ def to_detected_object(detection, name=None, extra_attrs=True):
bry = tly + h
bounding_box = etag.BoundingBox.from_coords(tlx, tly, brx, bry)

mask = detection.mask
mask = detection.get_mask()
confidence = detection.confidence

attrs = _to_eta_attributes(detection, extra_attrs=extra_attrs)
Expand Down
Loading