From 6e4775ab933f09e24421af10f5e4eef985236e00 Mon Sep 17 00:00:00 2001 From: Sashank Aryal Date: Wed, 20 Nov 2024 14:42:44 -0600 Subject: [PATCH 01/31] change typedef: export ImageBitmap from worker --- app/packages/looker/src/overlays/detection.ts | 3 +-- app/packages/looker/src/overlays/heatmap.ts | 5 ++--- app/packages/looker/src/overlays/segmentation.ts | 5 ++--- 3 files changed, 5 insertions(+), 8 deletions(-) diff --git a/app/packages/looker/src/overlays/detection.ts b/app/packages/looker/src/overlays/detection.ts index 4930771692..40168159f4 100644 --- a/app/packages/looker/src/overlays/detection.ts +++ b/app/packages/looker/src/overlays/detection.ts @@ -12,8 +12,7 @@ import { t } from "./util"; export interface DetectionLabel extends RegularLabel { mask?: { - data: OverlayMask; - image: ArrayBuffer; + bitmap?: ImageBitmap; }; bounding_box: BoundingBox; diff --git a/app/packages/looker/src/overlays/heatmap.ts b/app/packages/looker/src/overlays/heatmap.ts index c53a3ad971..94ef7163f5 100644 --- a/app/packages/looker/src/overlays/heatmap.ts +++ b/app/packages/looker/src/overlays/heatmap.ts @@ -7,7 +7,7 @@ import { getRGBA, getRGBAColor, } from "@fiftyone/utilities"; -import { ARRAY_TYPES, OverlayMask, TypedArray } from "../numpy"; +import { ARRAY_TYPES, TypedArray } from "../numpy"; import { BaseState, Coordinates } from "../state"; import { isFloatArray } from "../util"; import { clampedIndex } from "../worker/painter"; @@ -22,8 +22,7 @@ import { import { strokeCanvasRect, t } from "./util"; interface HeatMap { - data: OverlayMask; - image: ArrayBuffer; + bitmap?: ImageBitmap; } interface HeatmapLabel extends BaseLabel { diff --git a/app/packages/looker/src/overlays/segmentation.ts b/app/packages/looker/src/overlays/segmentation.ts index c55a8b5ef5..3a35651632 100644 --- a/app/packages/looker/src/overlays/segmentation.ts +++ b/app/packages/looker/src/overlays/segmentation.ts @@ -3,7 +3,7 @@ */ import { getColor } from "@fiftyone/utilities"; -import { ARRAY_TYPES, OverlayMask, TypedArray } from "../numpy"; +import { ARRAY_TYPES, TypedArray } from "../numpy"; import { BaseState, Coordinates, MaskTargets } from "../state"; import { BaseLabel, @@ -17,8 +17,7 @@ import { isRgbMaskTargets, strokeCanvasRect, t } from "./util"; interface SegmentationLabel extends BaseLabel { mask?: { - data: OverlayMask; - image: ArrayBuffer; + bitmap?: ImageBitmap; }; } From 53efcf3a38937d0ef2eb77a338635b746c4e561f Mon Sep 17 00:00:00 2001 From: Sashank Aryal Date: Wed, 20 Nov 2024 18:25:00 -0600 Subject: [PATCH 02/31] implement getSizeBytes in heatmap and segmentation, too --- app/packages/looker/src/overlays/detection.ts | 1 - app/packages/looker/src/overlays/heatmap.ts | 5 +++++ app/packages/looker/src/overlays/segmentation.ts | 6 +++++- 3 files changed, 10 insertions(+), 2 deletions(-) diff --git a/app/packages/looker/src/overlays/detection.ts b/app/packages/looker/src/overlays/detection.ts index 40168159f4..2dce852d53 100644 --- a/app/packages/looker/src/overlays/detection.ts +++ b/app/packages/looker/src/overlays/detection.ts @@ -4,7 +4,6 @@ import { NONFINITES } from "@fiftyone/utilities"; import { INFO_COLOR } from "../constants"; -import { OverlayMask } from "../numpy"; import { BaseState, BoundingBox, Coordinates, NONFINITE } from "../state"; import { distanceFromLineSegment } from "../util"; import { CONTAINS, CoordinateOverlay, PointInfo, RegularLabel } from "./base"; diff --git a/app/packages/looker/src/overlays/heatmap.ts b/app/packages/looker/src/overlays/heatmap.ts index 94ef7163f5..5e0a9f7e19 100644 --- a/app/packages/looker/src/overlays/heatmap.ts +++ b/app/packages/looker/src/overlays/heatmap.ts @@ -6,6 +6,7 @@ import { getColor, getRGBA, getRGBAColor, + sizeBytesEstimate, } from "@fiftyone/utilities"; import { ARRAY_TYPES, TypedArray } from "../numpy"; import { BaseState, Coordinates } from "../state"; @@ -234,6 +235,10 @@ export default class HeatmapOverlay return this.targets[index]; } + + getSizeBytes(): number { + return sizeBytesEstimate(this.label); + } } export const getHeatmapPoints = (labels: HeatmapLabel[]): Coordinates[] => { diff --git a/app/packages/looker/src/overlays/segmentation.ts b/app/packages/looker/src/overlays/segmentation.ts index 3a35651632..ec0bd89e39 100644 --- a/app/packages/looker/src/overlays/segmentation.ts +++ b/app/packages/looker/src/overlays/segmentation.ts @@ -2,7 +2,7 @@ * Copyright 2017-2024, Voxel51, Inc. */ -import { getColor } from "@fiftyone/utilities"; +import { getColor, sizeBytesEstimate } from "@fiftyone/utilities"; import { ARRAY_TYPES, TypedArray } from "../numpy"; import { BaseState, Coordinates, MaskTargets } from "../state"; import { @@ -277,6 +277,10 @@ export default class SegmentationOverlay } return this.targets[index]; } + + getSizeBytes(): number { + return sizeBytesEstimate(this.label); + } } export const getSegmentationPoints = ( From 7e00596d24d3bf83f5272babd07b7c050cb75ca0 Mon Sep 17 00:00:00 2001 From: Sashank Aryal Date: Wed, 20 Nov 2024 18:29:41 -0600 Subject: [PATCH 03/31] remove canvas creation in detection constructor --- app/packages/looker/src/overlays/detection.ts | 34 +++---------------- 1 file changed, 4 insertions(+), 30 deletions(-) diff --git a/app/packages/looker/src/overlays/detection.ts b/app/packages/looker/src/overlays/detection.ts index 2dce852d53..e4527016a5 100644 --- a/app/packages/looker/src/overlays/detection.ts +++ b/app/packages/looker/src/overlays/detection.ts @@ -25,10 +25,9 @@ export interface DetectionLabel extends RegularLabel { export default class DetectionOverlay< State extends BaseState > extends CoordinateOverlay { - private imageData: ImageData; private is3D: boolean; private labelBoundingBox: BoundingBox; - private canvas: HTMLCanvasElement; + private imageBitmap: ImageBitmap | null = null; constructor(field, label) { super(field, label); @@ -38,32 +37,6 @@ export default class DetectionOverlay< } else { this.is3D = false; } - - if (this.label.mask) { - const [height, width] = this.label.mask.data.shape; - - if (!height || !width) { - return; - } - - this.canvas = document.createElement("canvas"); - this.canvas.width = width; - this.canvas.height = height; - this.imageData = new ImageData( - new Uint8ClampedArray(this.label.mask.image), - width, - height - ); - const maskCtx = this.canvas.getContext("2d"); - maskCtx.imageSmoothingEnabled = false; - maskCtx.clearRect( - 0, - 0, - this.label.mask.data.shape[1], - this.label.mask.data.shape[0] - ); - maskCtx.putImageData(this.imageData, 0, 0); - } } containsPoint(state: Readonly): CONTAINS { @@ -167,7 +140,7 @@ export default class DetectionOverlay< } private drawMask(ctx: CanvasRenderingContext2D, state: Readonly) { - if (!this.canvas) { + if (!this.label.mask?.bitmap) { return; } @@ -175,8 +148,9 @@ export default class DetectionOverlay< const [x, y] = t(state, tlx, tly); const tmp = ctx.globalAlpha; ctx.globalAlpha = state.options.alpha; + ctx.imageSmoothingEnabled = false; ctx.drawImage( - this.canvas, + this.label.mask.bitmap, x, y, w * state.canvasBBox[2], From 3f305216162838b00a7fa32ff4c84e9d2a69fe4d Mon Sep 17 00:00:00 2001 From: Sashank Aryal Date: Wed, 20 Nov 2024 18:53:36 -0600 Subject: [PATCH 04/31] remove buffers arg from mask deserializer --- app/packages/looker/src/worker/deserializer.ts | 18 +++++------------- 1 file changed, 5 insertions(+), 13 deletions(-) diff --git a/app/packages/looker/src/worker/deserializer.ts b/app/packages/looker/src/worker/deserializer.ts index 02a7b03867..7e240a8842 100644 --- a/app/packages/looker/src/worker/deserializer.ts +++ b/app/packages/looker/src/worker/deserializer.ts @@ -14,7 +14,7 @@ const extractSerializedMask = ( }; export const DeserializerFactory = { - Detection: (label, buffers) => { + Detection: (label) => { const serializedMask = extractSerializedMask(label, "mask"); if (serializedMask) { @@ -24,17 +24,15 @@ export const DeserializerFactory = { data, image: new ArrayBuffer(width * height * 4), }; - buffers.push(data.buffer); - buffers.push(label.mask.image); } }, - Detections: (labels, buffers) => { + Detections: (labels) => { const list = labels?.detections || []; for (const label of list) { - DeserializerFactory.Detection(label, buffers); + DeserializerFactory.Detection(label); } }, - Heatmap: (label, buffers) => { + Heatmap: (label) => { const serializedMask = extractSerializedMask(label, "map"); if (serializedMask) { @@ -45,12 +43,9 @@ export const DeserializerFactory = { data, image: new ArrayBuffer(width * height * 4), }; - - buffers.push(data.buffer); - buffers.push(label.map.image); } }, - Segmentation: (label, buffers) => { + Segmentation: (label) => { const serializedMask = extractSerializedMask(label, "mask"); if (serializedMask) { @@ -61,9 +56,6 @@ export const DeserializerFactory = { data, image: new ArrayBuffer(width * height * 4), }; - - buffers.push(data.buffer); - buffers.push(label.mask.image); } }, }; From e29f14079b2cc0b56d77d1f378089eeccc04131c Mon Sep 17 00:00:00 2001 From: Sashank Aryal Date: Wed, 20 Nov 2024 18:54:02 -0600 Subject: [PATCH 05/31] make worker send bitmaps back to main thread --- app/packages/looker/src/worker/index.ts | 152 +++++++++++++++++------- 1 file changed, 108 insertions(+), 44 deletions(-) diff --git a/app/packages/looker/src/worker/index.ts b/app/packages/looker/src/worker/index.ts index 21859407e2..403e71b2cd 100644 --- a/app/packages/looker/src/worker/index.ts +++ b/app/packages/looker/src/worker/index.ts @@ -106,7 +106,6 @@ const imputeOverlayFromPath = async ( coloring: Coloring, customizeColorSetting: CustomizeColor[], colorscale: Colorscale, - buffers: ArrayBuffer[], sources: { [path: string]: string }, cls: string, maskPathDecodingPromises: Promise[] = [] @@ -122,7 +121,6 @@ const imputeOverlayFromPath = async ( coloring, customizeColorSetting, colorscale, - buffers, {}, DETECTION ) @@ -170,16 +168,26 @@ const imputeOverlayFromPath = async ( const [overlayHeight, overlayWidth] = overlayMask.shape; // set the `mask` property for this label + // we need to do this because we need raw image pixel data + // to iterate through and paint it with the color + // defined by the user for this particular label label[overlayField] = { data: overlayMask, image: new ArrayBuffer(overlayWidth * overlayHeight * 4), }; - - // transfer buffers - buffers.push(overlayMask.buffer); - buffers.push(label[overlayField].image); }; +/** + * 1. Start deserializing on-disk masks. Accumulate promises. + * 2. Await mask path decoding to finish. + * 3. Start painting overlays. Accumulate promises. + * 4. Await overlay painting to finish. + * 5. Start bitmap generation. Accumulate promises. + * 6. Await bitmap generation to finish. + * 7. Transfer bitmaps back to the main thread. + * + * Note that on-disk masks support async deserialization, which means they are more performant. + */ const processLabels = async ( sample: ProcessSample["sample"], coloring: ProcessSample["coloring"], @@ -190,11 +198,10 @@ const processLabels = async ( labelTagColors: ProcessSample["labelTagColors"], selectedLabelTags: ProcessSample["selectedLabelTags"], schema: Schema -): Promise => { - const buffers: ArrayBuffer[] = []; - const painterPromises = []; - +): Promise[]> => { const maskPathDecodingPromises = []; + const painterPromises = []; + const bitmapPromises = []; // mask deserialization / mask_path decoding loop for (const field in sample) { @@ -217,7 +224,6 @@ const processLabels = async ( coloring, customizeColorSetting, colorscale, - buffers, sources, cls, maskPathDecodingPromises @@ -226,11 +232,11 @@ const processLabels = async ( } if (cls in DeserializerFactory) { - DeserializerFactory[cls](label, buffers); + DeserializerFactory[cls](label); } if ([EMBEDDED_DOCUMENT, DYNAMIC_EMBEDDED_DOCUMENT].includes(cls)) { - const moreBuffers = await processLabels( + const moreBitmapPromises = await processLabels( label, coloring, `${prefix ? prefix : ""}${field}.`, @@ -241,7 +247,7 @@ const processLabels = async ( selectedLabelTags, schema ); - buffers.push(...moreBuffers); + bitmapPromises.push(...moreBitmapPromises); } if (ALL_VALID_LABELS.has(cls)) { @@ -286,7 +292,64 @@ const processLabels = async ( } } - return Promise.all(painterPromises).then(() => buffers); + await Promise.allSettled(painterPromises); + + // bitmap generation loop + for (const field in sample) { + let labels = sample[field]; + if (!Array.isArray(labels)) { + labels = [labels]; + } + const cls = getCls(`${prefix ? prefix : ""}${field}`, schema); + + for (const label of labels) { + if (!label) { + continue; + } + + collectBitmapPromises(label, cls, bitmapPromises); + } + } + + return bitmapPromises; +}; + +const collectBitmapPromises = (label, cls, bitmapPromises) => { + if (cls === DETECTIONS) { + label?.detections?.forEach((detection) => + collectBitmapPromises(detection, DETECTION, bitmapPromises) + ); + return; + } + + // we are detection now + if (cls !== DETECTION) { + return; + } + + if (label.mask) { + const [height, width] = label.mask.data.shape; + + const imageData = new ImageData( + new Uint8ClampedArray(label.mask.image), + width, + height + ); + + bitmapPromises.push( + new Promise((resolve) => { + createImageBitmap(imageData).then((imageBitmap) => { + // release buffers (will be garbage collected) + label.mask.data = null; + label.mask.image = null; + + label.mask.bitmap = imageBitmap; + + resolve(imageBitmap); + }); + }) + ); + } }; /** GLOBALS */ @@ -316,7 +379,7 @@ export interface ProcessSample { type ProcessSampleMethod = ReaderMethod & ProcessSample; -const processSample = ({ +const processSample = async ({ sample, uuid, coloring, @@ -329,13 +392,13 @@ const processSample = ({ }: ProcessSample) => { mapId(sample); - let bufferPromises = []; + let imageBitmapPromises: Promise[] = []; if (sample?._media_type === "point-cloud" || sample?._media_type === "3d") { process3DLabels(schema, sample); } else { - bufferPromises = [ - processLabels( + imageBitmapPromises.push( + ...(await processLabels( sample, coloring, null, @@ -345,32 +408,33 @@ const processSample = ({ labelTagColors, selectedLabelTags, schema - ), - ]; - } - - if (sample.frames && sample.frames.length) { - bufferPromises = [ - ...bufferPromises, - ...sample.frames - .map((frame) => - processLabels( - frame, - coloring, - "frames.", - sources, - customizeColorSetting, - colorscale, - labelTagColors, - selectedLabelTags, - schema - ) - ) - .flat(), - ]; + )) + ); } - Promise.all(bufferPromises).then((buffers) => { + // todo: address frames + // if (sample.frames && sample.frames.length) { + // bufferPromises = [ + // ...bufferPromises, + // ...sample.frames + // .map((frame) => + // processLabels( + // frame, + // coloring, + // "frames.", + // sources, + // customizeColorSetting, + // colorscale, + // labelTagColors, + // selectedLabelTags, + // schema + // ) + // ) + // .flat(), + // ]; + // } + + Promise.all(imageBitmapPromises).then((bitmaps) => { postMessage( { method: "processSample", @@ -383,7 +447,7 @@ const processSample = ({ selectedLabelTags, }, // @ts-ignore - buffers.flat() + bitmaps.flat() ); }); }; From 6cfce47eaccc88a9961a0bdf77441f38caad9257 Mon Sep 17 00:00:00 2001 From: Sashank Aryal Date: Wed, 20 Nov 2024 19:16:09 -0600 Subject: [PATCH 06/31] move path decoder to another file --- .../looker/src/worker/disk-overlay-decoder.ts | 87 ++++++++++++++++++ app/packages/looker/src/worker/index.ts | 88 +------------------ 2 files changed, 89 insertions(+), 86 deletions(-) create mode 100644 app/packages/looker/src/worker/disk-overlay-decoder.ts diff --git a/app/packages/looker/src/worker/disk-overlay-decoder.ts b/app/packages/looker/src/worker/disk-overlay-decoder.ts new file mode 100644 index 0000000000..01b725c703 --- /dev/null +++ b/app/packages/looker/src/worker/disk-overlay-decoder.ts @@ -0,0 +1,87 @@ +import { getSampleSrc } from "@fiftyone/state"; +import { DETECTION, DETECTIONS, HEATMAP } from "@fiftyone/utilities"; +import { Coloring, CustomizeColor } from ".."; +import { Colorscale } from "../state"; +import { decodeWithCanvas } from "./canvas-decoder"; +import { fetchWithLinearBackoff } from "./decorated-fetch"; + +/** + * 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 + */ +export const decodeOverlayOnDisk = async ( + field: string, + label: Record, + coloring: Coloring, + customizeColorSetting: CustomizeColor[], + colorscale: Colorscale, + sources: { [path: string]: string }, + cls: string, + maskPathDecodingPromises: Promise[] = [] +) => { + // handle all list types here + if (cls === DETECTIONS) { + const promises: Promise[] = []; + for (const detection of label.detections) { + promises.push( + decodeOverlayOnDisk( + field, + detection, + coloring, + customizeColorSetting, + colorscale, + {}, + DETECTION + ) + ); + } + maskPathDecodingPromises.push(...promises); + } + + // overlay path is in `map_path` property for heatmap, or else, it's in `mask_path` property (for segmentation or detection) + const overlayPathField = 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 overlayImageUrl = getSampleSrc( + sources[`${field}.${overlayPathField}`] || label[overlayPathField] + ); + const urlTokens = overlayImageUrl.split("?"); + + let baseUrl = overlayImageUrl; + + // remove query params if not local URL + if (!urlTokens.at(1)?.startsWith("filepath=")) { + baseUrl = overlayImageUrl.split("?")[0]; + } + + let overlayImageBlob: Blob; + try { + const overlayImageFetchResponse = await fetchWithLinearBackoff(baseUrl); + overlayImageBlob = await overlayImageFetchResponse.blob(); + } catch (e) { + console.error(e); + // skip decoding if fetch fails altogether + return; + } + + const overlayMask = await decodeWithCanvas(overlayImageBlob); + const [overlayHeight, overlayWidth] = overlayMask.shape; + + // set the `mask` property for this label + // we need to do this because we need raw image pixel data + // to iterate through and paint it with the color + // defined by the user for this particular label + label[overlayField] = { + data: overlayMask, + image: new ArrayBuffer(overlayWidth * overlayHeight * 4), + }; +}; diff --git a/app/packages/looker/src/worker/index.ts b/app/packages/looker/src/worker/index.ts index 403e71b2cd..af414cae37 100644 --- a/app/packages/looker/src/worker/index.ts +++ b/app/packages/looker/src/worker/index.ts @@ -2,14 +2,12 @@ * Copyright 2017-2024, Voxel51, Inc. */ -import { getSampleSrc } from "@fiftyone/state/src/recoil/utils"; import { DENSE_LABELS, DETECTION, DETECTIONS, DYNAMIC_EMBEDDED_DOCUMENT, EMBEDDED_DOCUMENT, - HEATMAP, LABEL_LIST, Schema, Stage, @@ -29,9 +27,8 @@ import { LabelTagColor, Sample, } from "../state"; -import { decodeWithCanvas } from "./canvas-decoder"; -import { fetchWithLinearBackoff } from "./decorated-fetch"; import { DeserializerFactory } from "./deserializer"; +import { decodeOverlayOnDisk } from "./disk-overlay-decoder"; import { PainterFactory } from "./painter"; import { mapId } from "./shared"; import { process3DLabels } from "./threed-label-processor"; @@ -96,87 +93,6 @@ const painterFactory = PainterFactory(requestColor); const ALL_VALID_LABELS = new Set(VALID_LABEL_TYPES); -/** - * 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 ( - field: string, - label: Record, - coloring: Coloring, - customizeColorSetting: CustomizeColor[], - colorscale: Colorscale, - sources: { [path: string]: string }, - cls: string, - maskPathDecodingPromises: Promise[] = [] -) => { - // handle all list types here - if (cls === DETECTIONS) { - const promises: Promise[] = []; - for (const detection of label.detections) { - promises.push( - imputeOverlayFromPath( - field, - detection, - coloring, - customizeColorSetting, - colorscale, - {}, - DETECTION - ) - ); - } - maskPathDecodingPromises.push(...promises); - } - - // overlay path is in `map_path` property for heatmap, or else, it's in `mask_path` property (for segmentation or detection) - const overlayPathField = 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 overlayImageUrl = getSampleSrc( - sources[`${field}.${overlayPathField}`] || label[overlayPathField] - ); - const urlTokens = overlayImageUrl.split("?"); - - let baseUrl = overlayImageUrl; - - // remove query params if not local URL - if (!urlTokens.at(1)?.startsWith("filepath=")) { - baseUrl = overlayImageUrl.split("?")[0]; - } - - let overlayImageBlob: Blob; - try { - const overlayImageFetchResponse = await fetchWithLinearBackoff(baseUrl); - overlayImageBlob = await overlayImageFetchResponse.blob(); - } catch (e) { - console.error(e); - // skip decoding if fetch fails altogether - return; - } - - const overlayMask = await decodeWithCanvas(overlayImageBlob); - const [overlayHeight, overlayWidth] = overlayMask.shape; - - // set the `mask` property for this label - // we need to do this because we need raw image pixel data - // to iterate through and paint it with the color - // defined by the user for this particular label - label[overlayField] = { - data: overlayMask, - image: new ArrayBuffer(overlayWidth * overlayHeight * 4), - }; -}; - /** * 1. Start deserializing on-disk masks. Accumulate promises. * 2. Await mask path decoding to finish. @@ -218,7 +134,7 @@ const processLabels = async ( if (DENSE_LABELS.has(cls)) { maskPathDecodingPromises.push( - imputeOverlayFromPath( + decodeOverlayOnDisk( `${prefix || ""}${field}`, label, coloring, From 6cb954e5a61445f8405eea30cc9ffcf1fc5b9172 Mon Sep 17 00:00:00 2001 From: Sashank Aryal Date: Wed, 20 Nov 2024 19:42:15 -0600 Subject: [PATCH 07/31] add unit tests for on disk decoder --- .../src/worker/disk-overlay-decoder.test.ts | 216 ++++++++++++++++++ .../looker/src/worker/disk-overlay-decoder.ts | 13 +- 2 files changed, 224 insertions(+), 5 deletions(-) create mode 100644 app/packages/looker/src/worker/disk-overlay-decoder.test.ts diff --git a/app/packages/looker/src/worker/disk-overlay-decoder.test.ts b/app/packages/looker/src/worker/disk-overlay-decoder.test.ts new file mode 100644 index 0000000000..5d20f454dd --- /dev/null +++ b/app/packages/looker/src/worker/disk-overlay-decoder.test.ts @@ -0,0 +1,216 @@ +import { getSampleSrc } from "@fiftyone/state"; +import { DETECTIONS, HEATMAP } from "@fiftyone/utilities"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import type { Coloring, CustomizeColor } from ".."; +import { LabelMask } from "../overlays/base"; +import type { Colorscale } from "../state"; +import { decodeWithCanvas } from "./canvas-decoder"; +import { fetchWithLinearBackoff } from "./decorated-fetch"; +import { decodeOverlayOnDisk, IntermediateMask } from "./disk-overlay-decoder"; + +vi.mock("@fiftyone/state", () => ({ + getSampleSrc: vi.fn(), +})); + +vi.mock("@fiftyone/utilities", () => ({ + DETECTION: "Detection", + DETECTIONS: "Detections", + HEATMAP: "Heatmap", +})); + +vi.mock("./canvas-decoder", () => ({ + decodeWithCanvas: vi.fn(), +})); + +vi.mock("./decorated-fetch", () => ({ + fetchWithLinearBackoff: vi.fn(), +})); + +const COLORING = {} as Coloring; +const COLOR_SCALE = {} as Colorscale; +const CUSTOMIZE_COLOR_SETTING: CustomizeColor[] = []; +const SOURCES = {}; + +type MaskUnion = (IntermediateMask & LabelMask) | null; + +describe("decodeOverlayOnDisk", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + it("should return early if label already has overlay field (not on disk)", async () => { + const field = "testField"; + const label = { mask: {}, mask_path: "shouldBeIgnored" }; + const cls = "Segmentation"; + const maskPathDecodingPromises: Promise[] = []; + + await decodeOverlayOnDisk( + field, + label, + COLORING, + CUSTOMIZE_COLOR_SETTING, + COLOR_SCALE, + SOURCES, + cls, + maskPathDecodingPromises + ); + + expect(label.mask).toBeDefined(); + expect(fetchWithLinearBackoff).not.toHaveBeenCalled(); + }); + + it("should fetch and decode overlay when label has overlay path field", async () => { + const field = "testField"; + const label = { mask_path: "/path/to/mask", mask: null as MaskUnion }; + const cls = "Segmentation"; + const maskPathDecodingPromises: Promise[] = []; + + const sampleSrcUrl = "http://example.com/path/to/mask"; + const mockBlob = new Blob(["mock data"], { type: "image/png" }); + const overlayMask = { shape: [100, 200] }; + + vi.mocked(getSampleSrc).mockReturnValue(sampleSrcUrl); + vi.mocked(fetchWithLinearBackoff).mockResolvedValue({ + blob: () => Promise.resolve(mockBlob), + } as Response); + vi.mocked(decodeWithCanvas).mockResolvedValue(overlayMask); + + await decodeOverlayOnDisk( + field, + label, + COLORING, + CUSTOMIZE_COLOR_SETTING, + COLOR_SCALE, + SOURCES, + cls, + maskPathDecodingPromises + ); + + expect(getSampleSrc).toHaveBeenCalledWith("/path/to/mask"); + expect(fetchWithLinearBackoff).toHaveBeenCalledWith(sampleSrcUrl); + expect(decodeWithCanvas).toHaveBeenCalledWith(mockBlob); + expect(label.mask).toBeDefined(); + expect(label.mask.data).toBe(overlayMask); + expect(label.mask.image).toBeInstanceOf(ArrayBuffer); + expect(label.mask.image.byteLength).toBe(100 * 200 * 4); + }); + + it("should handle HEATMAP class", async () => { + const field = "testField"; + const label = { map_path: "/path/to/map", map: null as MaskUnion }; + const cls = HEATMAP; + const maskPathDecodingPromises: Promise[] = []; + + const sampleSrcUrl = "http://example.com/path/to/map"; + const mockBlob = new Blob(["mock data"], { type: "image/png" }); + const overlayMask = { shape: [100, 200] }; + + vi.mocked(getSampleSrc).mockReturnValue(sampleSrcUrl); + vi.mocked(fetchWithLinearBackoff).mockResolvedValue({ + blob: () => Promise.resolve(mockBlob), + } as Response); + vi.mocked(decodeWithCanvas).mockResolvedValue(overlayMask); + + await decodeOverlayOnDisk( + field, + label, + COLORING, + CUSTOMIZE_COLOR_SETTING, + COLOR_SCALE, + SOURCES, + cls, + maskPathDecodingPromises + ); + + expect(getSampleSrc).toHaveBeenCalledWith("/path/to/map"); + expect(fetchWithLinearBackoff).toHaveBeenCalledWith(sampleSrcUrl); + expect(decodeWithCanvas).toHaveBeenCalledWith(mockBlob); + expect(label.map).toBeDefined(); + expect(label.map.data).toBe(overlayMask); + expect(label.map.image).toBeInstanceOf(ArrayBuffer); + expect(label.map.image.byteLength).toBe(100 * 200 * 4); + }); + + it("should handle DETECTIONS class and process detections recursively", async () => { + const field = "testField"; + const label = { + detections: [ + { mask_path: "/path/to/mask1", mask: null as MaskUnion }, + { mask_path: "/path/to/mask2", mask: null as MaskUnion }, + ], + }; + const cls = DETECTIONS; + const maskPathDecodingPromises: Promise[] = []; + + const sampleSrcUrl1 = "http://example.com/path/to/mask1"; + const sampleSrcUrl2 = "http://example.com/path/to/mask2"; + const mockBlob1 = new Blob(["mock data 1"], { type: "image/png" }); + const mockBlob2 = new Blob(["mock data 2"], { type: "image/png" }); + const overlayMask1 = { shape: [50, 50] }; + const overlayMask2 = { shape: [60, 60] }; + + vi.mocked(getSampleSrc) + .mockReturnValueOnce(sampleSrcUrl1) + .mockReturnValueOnce(sampleSrcUrl2); + vi.mocked(fetchWithLinearBackoff) + .mockResolvedValueOnce({ + blob: () => Promise.resolve(mockBlob1), + } as Response) + .mockResolvedValueOnce({ + blob: () => Promise.resolve(mockBlob2), + } as Response); + vi.mocked(decodeWithCanvas) + .mockResolvedValueOnce(overlayMask1) + .mockResolvedValueOnce(overlayMask2); + + await decodeOverlayOnDisk( + field, + label, + COLORING, + CUSTOMIZE_COLOR_SETTING, + COLOR_SCALE, + SOURCES, + cls, + maskPathDecodingPromises + ); + + await Promise.all(maskPathDecodingPromises); + + expect(getSampleSrc).toHaveBeenNthCalledWith(1, "/path/to/mask1"); + expect(getSampleSrc).toHaveBeenNthCalledWith(2, "/path/to/mask2"); + expect(label.detections[0].mask).toBeDefined(); + expect(label.detections[0].mask.data).toBe(overlayMask1); + expect(label.detections[1].mask).toBeDefined(); + expect(label.detections[1].mask.data).toBe(overlayMask2); + }); + + it("should return early if fetch (with retry) fails", async () => { + const field = "testField"; + const label = { mask_path: "/path/to/mask", mask: null as MaskUnion }; + const cls = "Segmentation"; + const maskPathDecodingPromises: Promise[] = []; + + const sampleSrcUrl = "http://example.com/path/to/mask"; + + vi.mocked(getSampleSrc).mockReturnValue(sampleSrcUrl); + vi.mocked(fetchWithLinearBackoff).mockRejectedValue( + new Error("Fetch failed") + ); + + await decodeOverlayOnDisk( + field, + label, + COLORING, + CUSTOMIZE_COLOR_SETTING, + COLOR_SCALE, + SOURCES, + cls, + maskPathDecodingPromises + ); + + expect(getSampleSrc).toHaveBeenCalledWith("/path/to/mask"); + expect(fetchWithLinearBackoff).toHaveBeenCalledWith(sampleSrcUrl); + expect(decodeWithCanvas).not.toHaveBeenCalled(); + expect(label.mask).toBeNull(); + }); +}); diff --git a/app/packages/looker/src/worker/disk-overlay-decoder.ts b/app/packages/looker/src/worker/disk-overlay-decoder.ts index 01b725c703..418ff981b9 100644 --- a/app/packages/looker/src/worker/disk-overlay-decoder.ts +++ b/app/packages/looker/src/worker/disk-overlay-decoder.ts @@ -1,10 +1,16 @@ import { getSampleSrc } from "@fiftyone/state"; import { DETECTION, DETECTIONS, HEATMAP } from "@fiftyone/utilities"; import { Coloring, CustomizeColor } from ".."; +import { OverlayMask } from "../numpy"; import { Colorscale } from "../state"; import { decodeWithCanvas } from "./canvas-decoder"; import { fetchWithLinearBackoff } from "./decorated-fetch"; +export type IntermediateMask = { + data: OverlayMask; + image: ArrayBuffer; +}; + /** * 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 @@ -42,10 +48,7 @@ export const decodeOverlayOnDisk = async ( const overlayPathField = cls === HEATMAP ? "map_path" : "mask_path"; const overlayField = overlayPathField === "map_path" ? "map" : "mask"; - if ( - Object.hasOwn(label, overlayField) || - !Object.hasOwn(label, overlayPathField) - ) { + if (Boolean(label[overlayField]) || !Object.hasOwn(label, overlayPathField)) { // nothing to be done return; } @@ -83,5 +86,5 @@ export const decodeOverlayOnDisk = async ( label[overlayField] = { data: overlayMask, image: new ArrayBuffer(overlayWidth * overlayHeight * 4), - }; + } as IntermediateMask; }; From fcc480ea51bdabadf2f0bb3b4be02a9baffcb679 Mon Sep 17 00:00:00 2001 From: Sashank Aryal Date: Wed, 20 Nov 2024 21:36:31 -0600 Subject: [PATCH 08/31] add LabelMask type --- app/packages/looker/src/overlays/base.ts | 4 ++++ app/packages/looker/src/overlays/detection.ts | 12 ++++++++---- app/packages/looker/src/overlays/heatmap.ts | 7 ++----- app/packages/looker/src/overlays/segmentation.ts | 5 ++--- 4 files changed, 16 insertions(+), 12 deletions(-) diff --git a/app/packages/looker/src/overlays/base.ts b/app/packages/looker/src/overlays/base.ts index fd817ecf9d..995c6349cc 100644 --- a/app/packages/looker/src/overlays/base.ts +++ b/app/packages/looker/src/overlays/base.ts @@ -39,6 +39,10 @@ export interface SelectData { frameNumber?: number; } +export type LabelMask = { + bitmap?: ImageBitmap; +}; + export interface RegularLabel extends BaseLabel { _id?: string; label?: string; diff --git a/app/packages/looker/src/overlays/detection.ts b/app/packages/looker/src/overlays/detection.ts index e4527016a5..4b9a2f9a0b 100644 --- a/app/packages/looker/src/overlays/detection.ts +++ b/app/packages/looker/src/overlays/detection.ts @@ -6,13 +6,17 @@ import { NONFINITES } from "@fiftyone/utilities"; import { INFO_COLOR } from "../constants"; import { BaseState, BoundingBox, Coordinates, NONFINITE } from "../state"; import { distanceFromLineSegment } from "../util"; -import { CONTAINS, CoordinateOverlay, PointInfo, RegularLabel } from "./base"; +import { + CONTAINS, + CoordinateOverlay, + LabelMask, + PointInfo, + RegularLabel, +} from "./base"; import { t } from "./util"; export interface DetectionLabel extends RegularLabel { - mask?: { - bitmap?: ImageBitmap; - }; + mask?: LabelMask; bounding_box: BoundingBox; // valid for 3D bounding boxes diff --git a/app/packages/looker/src/overlays/heatmap.ts b/app/packages/looker/src/overlays/heatmap.ts index 5e0a9f7e19..4695096b94 100644 --- a/app/packages/looker/src/overlays/heatmap.ts +++ b/app/packages/looker/src/overlays/heatmap.ts @@ -15,6 +15,7 @@ import { clampedIndex } from "../worker/painter"; import { BaseLabel, CONTAINS, + LabelMask, Overlay, PointInfo, SelectData, @@ -22,12 +23,8 @@ import { } from "./base"; import { strokeCanvasRect, t } from "./util"; -interface HeatMap { - bitmap?: ImageBitmap; -} - interface HeatmapLabel extends BaseLabel { - map?: HeatMap; + map?: LabelMask; range?: [number, number]; } diff --git a/app/packages/looker/src/overlays/segmentation.ts b/app/packages/looker/src/overlays/segmentation.ts index ec0bd89e39..01aa2409df 100644 --- a/app/packages/looker/src/overlays/segmentation.ts +++ b/app/packages/looker/src/overlays/segmentation.ts @@ -8,6 +8,7 @@ import { BaseState, Coordinates, MaskTargets } from "../state"; import { BaseLabel, CONTAINS, + LabelMask, Overlay, PointInfo, SelectData, @@ -16,9 +17,7 @@ import { import { isRgbMaskTargets, strokeCanvasRect, t } from "./util"; interface SegmentationLabel extends BaseLabel { - mask?: { - bitmap?: ImageBitmap; - }; + mask?: LabelMask; } interface SegmentationInfo extends BaseLabel { From e428200b87adda1c812c6e9c034d2272638c0be6 Mon Sep 17 00:00:00 2001 From: Sashank Aryal Date: Wed, 20 Nov 2024 21:54:02 -0600 Subject: [PATCH 09/31] use modified import for getSampleSrc --- app/packages/looker/src/worker/disk-overlay-decoder.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/app/packages/looker/src/worker/disk-overlay-decoder.ts b/app/packages/looker/src/worker/disk-overlay-decoder.ts index 418ff981b9..b72d466528 100644 --- a/app/packages/looker/src/worker/disk-overlay-decoder.ts +++ b/app/packages/looker/src/worker/disk-overlay-decoder.ts @@ -1,4 +1,4 @@ -import { getSampleSrc } from "@fiftyone/state"; +import { getSampleSrc } from "@fiftyone/state/src/recoil/utils"; import { DETECTION, DETECTIONS, HEATMAP } from "@fiftyone/utilities"; import { Coloring, CustomizeColor } from ".."; import { OverlayMask } from "../numpy"; From 997828d7804e582bef71469b5281dfd3d9c49b6f Mon Sep 17 00:00:00 2001 From: Sashank Aryal Date: Wed, 20 Nov 2024 22:06:25 -0600 Subject: [PATCH 10/31] remove detections specific logic in bitmap collection function --- .../looker/src/worker/disk-overlay-decoder.ts | 9 ++++--- app/packages/looker/src/worker/index.ts | 25 +++++++++---------- app/packages/looker/src/worker/shared.ts | 11 ++++++++ 3 files changed, 28 insertions(+), 17 deletions(-) diff --git a/app/packages/looker/src/worker/disk-overlay-decoder.ts b/app/packages/looker/src/worker/disk-overlay-decoder.ts index b72d466528..e072b9da95 100644 --- a/app/packages/looker/src/worker/disk-overlay-decoder.ts +++ b/app/packages/looker/src/worker/disk-overlay-decoder.ts @@ -1,10 +1,11 @@ import { getSampleSrc } from "@fiftyone/state/src/recoil/utils"; -import { DETECTION, DETECTIONS, HEATMAP } from "@fiftyone/utilities"; +import { DETECTION, DETECTIONS } from "@fiftyone/utilities"; import { Coloring, CustomizeColor } from ".."; import { OverlayMask } from "../numpy"; import { Colorscale } from "../state"; import { decodeWithCanvas } from "./canvas-decoder"; import { fetchWithLinearBackoff } from "./decorated-fetch"; +import { getOverlayFieldFromCls } from "./shared"; export type IntermediateMask = { data: OverlayMask; @@ -44,9 +45,9 @@ export const decodeOverlayOnDisk = async ( maskPathDecodingPromises.push(...promises); } - // overlay path is in `map_path` property for heatmap, or else, it's in `mask_path` property (for segmentation or detection) - const overlayPathField = cls === HEATMAP ? "map_path" : "mask_path"; - const overlayField = overlayPathField === "map_path" ? "map" : "mask"; + const overlayFields = getOverlayFieldFromCls(cls); + const overlayPathField = overlayFields.disk; + const overlayField = overlayFields.canonical; if (Boolean(label[overlayField]) || !Object.hasOwn(label, overlayPathField)) { // nothing to be done diff --git a/app/packages/looker/src/worker/index.ts b/app/packages/looker/src/worker/index.ts index af414cae37..446995418c 100644 --- a/app/packages/looker/src/worker/index.ts +++ b/app/packages/looker/src/worker/index.ts @@ -30,7 +30,7 @@ import { import { DeserializerFactory } from "./deserializer"; import { decodeOverlayOnDisk } from "./disk-overlay-decoder"; import { PainterFactory } from "./painter"; -import { mapId } from "./shared"; +import { getOverlayFieldFromCls, mapId } from "./shared"; import { process3DLabels } from "./threed-label-processor"; interface ResolveColor { @@ -238,28 +238,27 @@ const collectBitmapPromises = (label, cls, bitmapPromises) => { return; } - // we are detection now - if (cls !== DETECTION) { - return; - } + const overlayFields = getOverlayFieldFromCls(cls); + const overlayField = overlayFields.canonical; - if (label.mask) { - const [height, width] = label.mask.data.shape; + if (label[overlayField]) { + const [height, width] = label[overlayField].data.shape; const imageData = new ImageData( - new Uint8ClampedArray(label.mask.image), + new Uint8ClampedArray(label[overlayField].image), width, height ); + // release buffers (will be garbage collected) + // we created ImageData and don't need the raw data anymore + label[overlayField].data = null; + label[overlayField].image = null; + bitmapPromises.push( new Promise((resolve) => { createImageBitmap(imageData).then((imageBitmap) => { - // release buffers (will be garbage collected) - label.mask.data = null; - label.mask.image = null; - - label.mask.bitmap = imageBitmap; + label[overlayField].bitmap = imageBitmap; resolve(imageBitmap); }); diff --git a/app/packages/looker/src/worker/shared.ts b/app/packages/looker/src/worker/shared.ts index adfda58d29..ec383b7536 100644 --- a/app/packages/looker/src/worker/shared.ts +++ b/app/packages/looker/src/worker/shared.ts @@ -1,3 +1,5 @@ +import { HEATMAP } from "@fiftyone/utilities"; + /** * Map the _id field to id */ @@ -8,3 +10,12 @@ export const mapId = (obj) => { } return obj; }; + +export const getOverlayFieldFromCls = (cls: string) => { + switch (cls) { + case HEATMAP: + return { canonical: "map", disk: "map_path" }; + default: + return { canonical: "mask", disk: "mask_path" }; + } +}; From 4ea1df56b9c62c5cbe156cd20f0ca48dd42ba9c2 Mon Sep 17 00:00:00 2001 From: Sashank Aryal Date: Wed, 20 Nov 2024 22:38:23 -0600 Subject: [PATCH 11/31] return targets buffers to main thread but not image buffers --- app/packages/looker/src/overlays/base.ts | 2 + app/packages/looker/src/overlays/heatmap.ts | 44 +++-------- .../looker/src/worker/disk-overlay-decoder.ts | 7 +- app/packages/looker/src/worker/index.ts | 77 ++++++++++--------- 4 files changed, 60 insertions(+), 70 deletions(-) diff --git a/app/packages/looker/src/overlays/base.ts b/app/packages/looker/src/overlays/base.ts index 995c6349cc..179b3aec18 100644 --- a/app/packages/looker/src/overlays/base.ts +++ b/app/packages/looker/src/overlays/base.ts @@ -3,6 +3,7 @@ */ import { getCls, sizeBytesEstimate } from "@fiftyone/utilities"; +import { OverlayMask } from "../numpy"; import type { BaseState, Coordinates, NONFINITE } from "../state"; import { getLabelColor, shouldShowLabelTag } from "./util"; @@ -41,6 +42,7 @@ export interface SelectData { export type LabelMask = { bitmap?: ImageBitmap; + data?: OverlayMask; }; export interface RegularLabel extends BaseLabel { diff --git a/app/packages/looker/src/overlays/heatmap.ts b/app/packages/looker/src/overlays/heatmap.ts index 4695096b94..b5bb143853 100644 --- a/app/packages/looker/src/overlays/heatmap.ts +++ b/app/packages/looker/src/overlays/heatmap.ts @@ -42,8 +42,6 @@ export default class HeatmapOverlay private label: HeatmapLabel; private targets?: TypedArray; private readonly range: [number, number]; - private canvas: HTMLCanvasElement; - private imageData: ImageData; constructor(field: string, label: HeatmapLabel) { this.field = field; @@ -65,25 +63,6 @@ export default class HeatmapOverlay if (!width || !height) { return; } - - this.canvas = document.createElement("canvas"); - this.canvas.width = width; - this.canvas.height = height; - - this.imageData = new ImageData( - new Uint8ClampedArray(this.label.map.image), - width, - height - ); - const maskCtx = this.canvas.getContext("2d"); - maskCtx.imageSmoothingEnabled = false; - maskCtx.clearRect( - 0, - 0, - this.label.map.data.shape[1], - this.label.map.data.shape[0] - ); - maskCtx.putImageData(this.imageData, 0, 0); } containsPoint(state: Readonly): CONTAINS { @@ -98,22 +77,23 @@ export default class HeatmapOverlay } draw(ctx: CanvasRenderingContext2D, state: Readonly): void { - if (this.imageData) { - const maskCtx = this.canvas.getContext("2d"); - maskCtx.imageSmoothingEnabled = false; - maskCtx.clearRect( - 0, - 0, - this.label.map.data.shape[1], - this.label.map.data.shape[0] - ); - maskCtx.putImageData(this.imageData, 0, 0); + if (this.label.map.bitmap) { + // const maskCtx = this.canvas.getContext("2d"); + // maskCtx.imageSmoothingEnabled = false; + // maskCtx.clearRect( + // 0, + // 0, + // this.label.map.data.shape[1], + // this.label.map.data.shape[0] + // ); + // maskCtx.putImageData(this.imageData, 0, 0); const [tlx, tly] = t(state, 0, 0); const [brx, bry] = t(state, 1, 1); const tmp = ctx.globalAlpha; ctx.globalAlpha = state.options.alpha; - ctx.drawImage(this.canvas, tlx, tly, brx - tlx, bry - tly); + // ctx.drawImage(this.canvas, tlx, tly, brx - tlx, bry - tly); + ctx.drawImage(this.label.map.bitmap, tlx, tly, brx - tlx, bry - tly); ctx.globalAlpha = tmp; } diff --git a/app/packages/looker/src/worker/disk-overlay-decoder.ts b/app/packages/looker/src/worker/disk-overlay-decoder.ts index e072b9da95..a53de4f50e 100644 --- a/app/packages/looker/src/worker/disk-overlay-decoder.ts +++ b/app/packages/looker/src/worker/disk-overlay-decoder.ts @@ -24,7 +24,8 @@ export const decodeOverlayOnDisk = async ( colorscale: Colorscale, sources: { [path: string]: string }, cls: string, - maskPathDecodingPromises: Promise[] = [] + maskPathDecodingPromises: Promise[] = [], + maskTargetsBuffers: ArrayBuffer[] = [] ) => { // handle all list types here if (cls === DETECTIONS) { @@ -88,4 +89,8 @@ export const decodeOverlayOnDisk = async ( data: overlayMask, image: new ArrayBuffer(overlayWidth * overlayHeight * 4), } as IntermediateMask; + + // no need to transfer image's buffer + //since we'll be constructing ImageBitmap and transfering that + maskTargetsBuffers.push(overlayMask.buffer); }; diff --git a/app/packages/looker/src/worker/index.ts b/app/packages/looker/src/worker/index.ts index 446995418c..32fe080528 100644 --- a/app/packages/looker/src/worker/index.ts +++ b/app/packages/looker/src/worker/index.ts @@ -114,12 +114,13 @@ const processLabels = async ( labelTagColors: ProcessSample["labelTagColors"], selectedLabelTags: ProcessSample["selectedLabelTags"], schema: Schema -): Promise[]> => { - const maskPathDecodingPromises = []; - const painterPromises = []; - const bitmapPromises = []; +): Promise<[Promise[], ArrayBuffer[]]> => { + const maskPathDecodingPromises: Promise[] = []; + const painterPromises: Promise[] = []; + const bitmapPromises: Promise[] = []; + const maskTargetsBuffers: ArrayBuffer[] = []; - // mask deserialization / mask_path decoding loop + // mask deserialization / on-disk overlay decoding loop for (const field in sample) { let labels = sample[field]; if (!Array.isArray(labels)) { @@ -142,28 +143,31 @@ const processLabels = async ( colorscale, sources, cls, - maskPathDecodingPromises + maskPathDecodingPromises, + maskTargetsBuffers ) ); } if (cls in DeserializerFactory) { - DeserializerFactory[cls](label); + DeserializerFactory[cls](label, maskTargetsBuffers); } if ([EMBEDDED_DOCUMENT, DYNAMIC_EMBEDDED_DOCUMENT].includes(cls)) { - const moreBitmapPromises = await processLabels( - label, - coloring, - `${prefix ? prefix : ""}${field}.`, - sources, - customizeColorSetting, - colorscale, - labelTagColors, - selectedLabelTags, - schema - ); + const [moreBitmapPromises, moreMaskTargetsBuffers] = + await processLabels( + label, + coloring, + `${prefix ? prefix : ""}${field}.`, + sources, + customizeColorSetting, + colorscale, + labelTagColors, + selectedLabelTags, + schema + ); bitmapPromises.push(...moreBitmapPromises); + maskTargetsBuffers.push(...moreMaskTargetsBuffers); } if (ALL_VALID_LABELS.has(cls)) { @@ -227,7 +231,7 @@ const processLabels = async ( } } - return bitmapPromises; + return [bitmapPromises, maskTargetsBuffers]; }; const collectBitmapPromises = (label, cls, bitmapPromises) => { @@ -250,16 +254,14 @@ const collectBitmapPromises = (label, cls, bitmapPromises) => { height ); - // release buffers (will be garbage collected) - // we created ImageData and don't need the raw data anymore - label[overlayField].data = null; + // set raw image to null - will be garbage collected + // we don't need it anymore since we copied to ImageData label[overlayField].image = null; bitmapPromises.push( new Promise((resolve) => { createImageBitmap(imageData).then((imageBitmap) => { label[overlayField].bitmap = imageBitmap; - resolve(imageBitmap); }); }) @@ -307,24 +309,25 @@ const processSample = async ({ }: ProcessSample) => { mapId(sample); - let imageBitmapPromises: Promise[] = []; + const imageBitmapPromises: Promise[] = []; + const maskTargetsBuffers: ArrayBuffer[] = []; if (sample?._media_type === "point-cloud" || sample?._media_type === "3d") { process3DLabels(schema, sample); } else { - imageBitmapPromises.push( - ...(await processLabels( - sample, - coloring, - null, - sources, - customizeColorSetting, - colorscale, - labelTagColors, - selectedLabelTags, - schema - )) + const [bitmapPromises, moreMaskTargetsBuffers] = await processLabels( + sample, + coloring, + null, + sources, + customizeColorSetting, + colorscale, + labelTagColors, + selectedLabelTags, + schema ); + imageBitmapPromises.push(...bitmapPromises); + maskTargetsBuffers.push(...moreMaskTargetsBuffers); } // todo: address frames @@ -362,7 +365,7 @@ const processSample = async ({ selectedLabelTags, }, // @ts-ignore - bitmaps.flat() + bitmaps.flat().concat(maskTargetsBuffers.flat()) ); }); }; From 4a0cf0654390ac9b237e6673b5bafcfb35b212dd Mon Sep 17 00:00:00 2001 From: Sashank Aryal Date: Wed, 20 Nov 2024 22:50:05 -0600 Subject: [PATCH 12/31] fix frames --- app/packages/looker/src/overlays/heatmap.ts | 11 ------ app/packages/looker/src/worker/index.ts | 38 +++++++++------------ 2 files changed, 17 insertions(+), 32 deletions(-) diff --git a/app/packages/looker/src/overlays/heatmap.ts b/app/packages/looker/src/overlays/heatmap.ts index b5bb143853..5e67504b29 100644 --- a/app/packages/looker/src/overlays/heatmap.ts +++ b/app/packages/looker/src/overlays/heatmap.ts @@ -78,21 +78,10 @@ export default class HeatmapOverlay draw(ctx: CanvasRenderingContext2D, state: Readonly): void { if (this.label.map.bitmap) { - // const maskCtx = this.canvas.getContext("2d"); - // maskCtx.imageSmoothingEnabled = false; - // maskCtx.clearRect( - // 0, - // 0, - // this.label.map.data.shape[1], - // this.label.map.data.shape[0] - // ); - // maskCtx.putImageData(this.imageData, 0, 0); - const [tlx, tly] = t(state, 0, 0); const [brx, bry] = t(state, 1, 1); const tmp = ctx.globalAlpha; ctx.globalAlpha = state.options.alpha; - // ctx.drawImage(this.canvas, tlx, tly, brx - tlx, bry - tly); ctx.drawImage(this.label.map.bitmap, tlx, tly, brx - tlx, bry - tly); ctx.globalAlpha = tmp; } diff --git a/app/packages/looker/src/worker/index.ts b/app/packages/looker/src/worker/index.ts index 32fe080528..bd0bd484e6 100644 --- a/app/packages/looker/src/worker/index.ts +++ b/app/packages/looker/src/worker/index.ts @@ -330,27 +330,23 @@ const processSample = async ({ maskTargetsBuffers.push(...moreMaskTargetsBuffers); } - // todo: address frames - // if (sample.frames && sample.frames.length) { - // bufferPromises = [ - // ...bufferPromises, - // ...sample.frames - // .map((frame) => - // processLabels( - // frame, - // coloring, - // "frames.", - // sources, - // customizeColorSetting, - // colorscale, - // labelTagColors, - // selectedLabelTags, - // schema - // ) - // ) - // .flat(), - // ]; - // } + if (sample.frames && sample.frames.length) { + for (const frame of sample.frames) { + const [moreBitmapPromises, moreMaskTargetsBuffers] = await processLabels( + frame, + coloring, + "frames.", + sources, + customizeColorSetting, + colorscale, + labelTagColors, + selectedLabelTags, + schema + ); + imageBitmapPromises.push(...moreBitmapPromises); + maskTargetsBuffers.push(...moreMaskTargetsBuffers); + } + } Promise.all(imageBitmapPromises).then((bitmaps) => { postMessage( From e28dce47917b5d73f84769e2f8997bc26f2d7f16 Mon Sep 17 00:00:00 2001 From: Sashank Aryal Date: Wed, 20 Nov 2024 22:52:03 -0600 Subject: [PATCH 13/31] push buffers in deserializer --- app/packages/looker/src/worker/deserializer.ts | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/app/packages/looker/src/worker/deserializer.ts b/app/packages/looker/src/worker/deserializer.ts index 7e240a8842..363522b01f 100644 --- a/app/packages/looker/src/worker/deserializer.ts +++ b/app/packages/looker/src/worker/deserializer.ts @@ -14,7 +14,7 @@ const extractSerializedMask = ( }; export const DeserializerFactory = { - Detection: (label) => { + Detection: (label, buffers) => { const serializedMask = extractSerializedMask(label, "mask"); if (serializedMask) { @@ -24,15 +24,16 @@ export const DeserializerFactory = { data, image: new ArrayBuffer(width * height * 4), }; + buffers.push(data.buffer); } }, - Detections: (labels) => { + Detections: (labels, buffers) => { const list = labels?.detections || []; for (const label of list) { - DeserializerFactory.Detection(label); + DeserializerFactory.Detection(label, buffers); } }, - Heatmap: (label) => { + Heatmap: (label, buffers) => { const serializedMask = extractSerializedMask(label, "map"); if (serializedMask) { @@ -43,9 +44,11 @@ export const DeserializerFactory = { data, image: new ArrayBuffer(width * height * 4), }; + + buffers.push(data.buffer); } }, - Segmentation: (label) => { + Segmentation: (label, buffers) => { const serializedMask = extractSerializedMask(label, "mask"); if (serializedMask) { @@ -56,6 +59,8 @@ export const DeserializerFactory = { data, image: new ArrayBuffer(width * height * 4), }; + + buffers.push(data.buffer); } }, }; From c55387a4fe7ec3664dc7f1e90ea43e25326b8411 Mon Sep 17 00:00:00 2001 From: Sashank Aryal Date: Wed, 20 Nov 2024 22:54:32 -0600 Subject: [PATCH 14/31] draw bitmap in segmentation, too --- app/packages/looker/src/overlays/heatmap.ts | 2 +- .../looker/src/overlays/segmentation.ts | 26 +++---------------- 2 files changed, 4 insertions(+), 24 deletions(-) diff --git a/app/packages/looker/src/overlays/heatmap.ts b/app/packages/looker/src/overlays/heatmap.ts index 5e67504b29..5609550a06 100644 --- a/app/packages/looker/src/overlays/heatmap.ts +++ b/app/packages/looker/src/overlays/heatmap.ts @@ -77,7 +77,7 @@ export default class HeatmapOverlay } draw(ctx: CanvasRenderingContext2D, state: Readonly): void { - if (this.label.map.bitmap) { + if (this.label.map?.bitmap) { const [tlx, tly] = t(state, 0, 0); const [brx, bry] = t(state, 1, 1); const tmp = ctx.globalAlpha; diff --git a/app/packages/looker/src/overlays/segmentation.ts b/app/packages/looker/src/overlays/segmentation.ts index 01aa2409df..e8e5359823 100644 --- a/app/packages/looker/src/overlays/segmentation.ts +++ b/app/packages/looker/src/overlays/segmentation.ts @@ -32,8 +32,6 @@ export default class SegmentationOverlay readonly field: string; private label: SegmentationLabel; private targets?: TypedArray; - private canvas: HTMLCanvasElement; - private imageData: ImageData; private isRgbMaskTargets = false; @@ -51,6 +49,7 @@ export default class SegmentationOverlay if (!this.label.mask) { return; } + const [height, width] = this.label.mask.data.shape; if (!height || !width) { @@ -60,25 +59,6 @@ export default class SegmentationOverlay this.targets = new ARRAY_TYPES[this.label.mask.data.arrayType]( this.label.mask.data.buffer ); - - this.canvas = document.createElement("canvas"); - this.canvas.width = width; - this.canvas.height = height; - - this.imageData = new ImageData( - new Uint8ClampedArray(this.label.mask.image), - width, - height - ); - const maskCtx = this.canvas.getContext("2d"); - maskCtx.imageSmoothingEnabled = false; - maskCtx.clearRect( - 0, - 0, - this.label.mask.data.shape[1], - this.label.mask.data.shape[0] - ); - maskCtx.putImageData(this.imageData, 0, 0); } containsPoint(state: Readonly): CONTAINS { @@ -97,12 +77,12 @@ export default class SegmentationOverlay return; } - if (this.imageData) { + if (this.label.mask?.bitmap) { const [tlx, tly] = t(state, 0, 0); const [brx, bry] = t(state, 1, 1); const tmp = ctx.globalAlpha; ctx.globalAlpha = state.options.alpha; - ctx.drawImage(this.canvas, tlx, tly, brx - tlx, bry - tly); + ctx.drawImage(this.label.mask.bitmap, tlx, tly, brx - tlx, bry - tly); ctx.globalAlpha = tmp; } From 35cceeecf70a5b99a485d56da2a696f16c5f670c Mon Sep 17 00:00:00 2001 From: Sashank Aryal Date: Wed, 20 Nov 2024 23:08:36 -0600 Subject: [PATCH 15/31] clarify comments --- app/packages/looker/src/worker/index.ts | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/app/packages/looker/src/worker/index.ts b/app/packages/looker/src/worker/index.ts index bd0bd484e6..e11906b51a 100644 --- a/app/packages/looker/src/worker/index.ts +++ b/app/packages/looker/src/worker/index.ts @@ -94,15 +94,14 @@ const painterFactory = PainterFactory(requestColor); const ALL_VALID_LABELS = new Set(VALID_LABEL_TYPES); /** - * 1. Start deserializing on-disk masks. Accumulate promises. + * This function processes labels in a recursive manner. It follows the following steps: + * 1. Deserialize masks. Accumulate promises. * 2. Await mask path decoding to finish. * 3. Start painting overlays. Accumulate promises. * 4. Await overlay painting to finish. * 5. Start bitmap generation. Accumulate promises. * 6. Await bitmap generation to finish. - * 7. Transfer bitmaps back to the main thread. - * - * Note that on-disk masks support async deserialization, which means they are more performant. + * 7. Transfer bitmaps and mask targets array buffers back to the main thread. */ const processLabels = async ( sample: ProcessSample["sample"], From 00d7381a4d1ce0d58b2f1491304323a9915460aa Mon Sep 17 00:00:00 2001 From: Sashank Aryal Date: Fri, 22 Nov 2024 16:41:34 -0600 Subject: [PATCH 16/31] take fetch opts --- .../looker/src/worker/decorated-fetch.test.ts | 15 ++++++++++----- app/packages/looker/src/worker/decorated-fetch.ts | 3 ++- 2 files changed, 12 insertions(+), 6 deletions(-) diff --git a/app/packages/looker/src/worker/decorated-fetch.test.ts b/app/packages/looker/src/worker/decorated-fetch.test.ts index 67ed853200..3a9a15e1e7 100644 --- a/app/packages/looker/src/worker/decorated-fetch.test.ts +++ b/app/packages/looker/src/worker/decorated-fetch.test.ts @@ -15,7 +15,7 @@ describe("fetchWithLinearBackoff", () => { expect(response).toBe(mockResponse); expect(global.fetch).toHaveBeenCalledTimes(1); - expect(global.fetch).toHaveBeenCalledWith("http://fiftyone.ai"); + expect(global.fetch).toHaveBeenCalledWith("http://fiftyone.ai", {}); }); it("should retry when fetch fails and eventually succeed", async () => { @@ -35,7 +35,7 @@ describe("fetchWithLinearBackoff", () => { global.fetch = vi.fn().mockRejectedValue(new Error("Network Error")); await expect( - fetchWithLinearBackoff("http://fiftyone.ai", 3, 10) + fetchWithLinearBackoff("http://fiftyone.ai", {}, 3, 10) ).rejects.toThrowError(new RegExp("Max retries for fetch reached")); expect(global.fetch).toHaveBeenCalledTimes(3); @@ -46,7 +46,7 @@ describe("fetchWithLinearBackoff", () => { global.fetch = vi.fn().mockResolvedValue(mockResponse); await expect( - fetchWithLinearBackoff("http://fiftyone.ai", 5, 10) + fetchWithLinearBackoff("http://fiftyone.ai", {}, 5, 10) ).rejects.toThrow("HTTP error: 500"); expect(global.fetch).toHaveBeenCalledTimes(5); @@ -57,7 +57,7 @@ describe("fetchWithLinearBackoff", () => { global.fetch = vi.fn().mockResolvedValue(mockResponse); await expect( - fetchWithLinearBackoff("http://fiftyone.ai", 5, 10) + fetchWithLinearBackoff("http://fiftyone.ai", {}, 5, 10) ).rejects.toThrow("Non-retryable HTTP error: 404"); expect(global.fetch).toHaveBeenCalledTimes(1); @@ -73,7 +73,12 @@ describe("fetchWithLinearBackoff", () => { vi.useFakeTimers(); - const fetchPromise = fetchWithLinearBackoff("http://fiftyone.ai", 5, 10); + const fetchPromise = fetchWithLinearBackoff( + "http://fiftyone.ai", + {}, + 5, + 10 + ); // advance timers to simulate delays // after first delay diff --git a/app/packages/looker/src/worker/decorated-fetch.ts b/app/packages/looker/src/worker/decorated-fetch.ts index c77059d551..9f0a910ea2 100644 --- a/app/packages/looker/src/worker/decorated-fetch.ts +++ b/app/packages/looker/src/worker/decorated-fetch.ts @@ -12,12 +12,13 @@ class NonRetryableError extends Error { export const fetchWithLinearBackoff = async ( url: string, + opts: RequestInit = {}, retries = DEFAULT_MAX_RETRIES, delay = DEFAULT_BASE_DELAY ) => { for (let i = 0; i < retries; i++) { try { - const response = await fetch(url); + const response = await fetch(url, opts); if (response.ok) { return response; } else { From bc9ba4150b7fe995bd7f4afc8ba96cea9c6d2a33 Mon Sep 17 00:00:00 2001 From: Sashank Aryal Date: Fri, 22 Nov 2024 16:50:24 -0600 Subject: [PATCH 17/31] pool fetch --- .../looker/src/worker/disk-overlay-decoder.ts | 7 +++- .../looker/src/worker/pooled-fetch.ts | 38 +++++++++++++++++++ 2 files changed, 43 insertions(+), 2 deletions(-) create mode 100644 app/packages/looker/src/worker/pooled-fetch.ts diff --git a/app/packages/looker/src/worker/disk-overlay-decoder.ts b/app/packages/looker/src/worker/disk-overlay-decoder.ts index a53de4f50e..989f002032 100644 --- a/app/packages/looker/src/worker/disk-overlay-decoder.ts +++ b/app/packages/looker/src/worker/disk-overlay-decoder.ts @@ -4,7 +4,7 @@ import { Coloring, CustomizeColor } from ".."; import { OverlayMask } from "../numpy"; import { Colorscale } from "../state"; import { decodeWithCanvas } from "./canvas-decoder"; -import { fetchWithLinearBackoff } from "./decorated-fetch"; +import { enqueueFetch } from "./pooled-fetch"; import { getOverlayFieldFromCls } from "./shared"; export type IntermediateMask = { @@ -70,7 +70,10 @@ export const decodeOverlayOnDisk = async ( let overlayImageBlob: Blob; try { - const overlayImageFetchResponse = await fetchWithLinearBackoff(baseUrl); + const overlayImageFetchResponse = await enqueueFetch({ + url: baseUrl, + options: { priority: "low" }, + }); overlayImageBlob = await overlayImageFetchResponse.blob(); } catch (e) { console.error(e); diff --git a/app/packages/looker/src/worker/pooled-fetch.ts b/app/packages/looker/src/worker/pooled-fetch.ts new file mode 100644 index 0000000000..434efa177f --- /dev/null +++ b/app/packages/looker/src/worker/pooled-fetch.ts @@ -0,0 +1,38 @@ +import { fetchWithLinearBackoff } from "./decorated-fetch"; + +// note: arbitrary number that seems to work well +const MAX_CONCURRENT_REQUESTS = 100; + +let activeRequests = 0; +const requestQueue = []; + +export const enqueueFetch = (request: { + url: string; + options?: RequestInit; +}): Promise => { + return new Promise((resolve, reject) => { + requestQueue.push({ request, resolve, reject }); + processFetchQueue(); + }); +}; + +const processFetchQueue = () => { + if (activeRequests >= MAX_CONCURRENT_REQUESTS || requestQueue.length === 0) { + return; + } + + const { request, resolve, reject } = requestQueue.shift(); + activeRequests++; + + fetchWithLinearBackoff(request.url, request.options) + .then((response) => { + activeRequests--; + resolve(response); + processFetchQueue(); + }) + .catch((error) => { + activeRequests--; + reject(error); + processFetchQueue(); + }); +}; From 0fcedd130f9fefc5cdf2af6f96dd52d3a3ab132a Mon Sep 17 00:00:00 2001 From: Sashank Aryal Date: Fri, 22 Nov 2024 17:02:03 -0600 Subject: [PATCH 18/31] pass missing params in recursive calls --- app/packages/looker/src/worker/disk-overlay-decoder.ts | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/app/packages/looker/src/worker/disk-overlay-decoder.ts b/app/packages/looker/src/worker/disk-overlay-decoder.ts index 989f002032..764666ee25 100644 --- a/app/packages/looker/src/worker/disk-overlay-decoder.ts +++ b/app/packages/looker/src/worker/disk-overlay-decoder.ts @@ -39,7 +39,9 @@ export const decodeOverlayOnDisk = async ( customizeColorSetting, colorscale, {}, - DETECTION + DETECTION, + maskPathDecodingPromises, + maskTargetsBuffers ) ); } From b732e20a60b83557fa1c509221435ba591f9dfb2 Mon Sep 17 00:00:00 2001 From: Sashank Aryal Date: Fri, 22 Nov 2024 17:10:21 -0600 Subject: [PATCH 19/31] add types --- app/packages/looker/src/worker/pooled-fetch.ts | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/app/packages/looker/src/worker/pooled-fetch.ts b/app/packages/looker/src/worker/pooled-fetch.ts index 434efa177f..3a61c8b0ce 100644 --- a/app/packages/looker/src/worker/pooled-fetch.ts +++ b/app/packages/looker/src/worker/pooled-fetch.ts @@ -1,10 +1,19 @@ import { fetchWithLinearBackoff } from "./decorated-fetch"; +interface QueueItem { + request: { + url: string; + options?: RequestInit; + }; + resolve: (value: Response | PromiseLike) => void; + reject: (reason?: any) => void; +} + // note: arbitrary number that seems to work well const MAX_CONCURRENT_REQUESTS = 100; let activeRequests = 0; -const requestQueue = []; +const requestQueue: QueueItem[] = []; export const enqueueFetch = (request: { url: string; @@ -33,6 +42,5 @@ const processFetchQueue = () => { .catch((error) => { activeRequests--; reject(error); - processFetchQueue(); }); }; From 1a0331ab6fd08946e8fac0c132c3deb8ce064038 Mon Sep 17 00:00:00 2001 From: Sashank Aryal Date: Fri, 22 Nov 2024 17:20:41 -0600 Subject: [PATCH 20/31] close bitmap when looker destroys --- app/packages/looker/src/lookers/abstract.ts | 5 +++++ app/packages/looker/src/overlays/base.ts | 1 + app/packages/looker/src/overlays/detection.ts | 8 +++++++- app/packages/looker/src/overlays/heatmap.ts | 6 ++++++ app/packages/looker/src/overlays/segmentation.ts | 6 ++++++ 5 files changed, 25 insertions(+), 1 deletion(-) diff --git a/app/packages/looker/src/lookers/abstract.ts b/app/packages/looker/src/lookers/abstract.ts index 7d18e8bddb..3eca774de8 100644 --- a/app/packages/looker/src/lookers/abstract.ts +++ b/app/packages/looker/src/lookers/abstract.ts @@ -293,6 +293,11 @@ export abstract class AbstractLooker< return; } + if (this.state.destroyed && this.sampleOverlays) { + // close all current overlays + this.pluckedOverlays.forEach((overlay) => overlay.cleanup?.()); + } + if ( !this.state.windowBBox || this.state.destroyed || diff --git a/app/packages/looker/src/overlays/base.ts b/app/packages/looker/src/overlays/base.ts index 179b3aec18..a3ec867766 100644 --- a/app/packages/looker/src/overlays/base.ts +++ b/app/packages/looker/src/overlays/base.ts @@ -73,6 +73,7 @@ export interface Overlay> { getPoints(state: Readonly): Coordinates[]; getSelectData(state: Readonly): SelectData; getSizeBytes(): number; + cleanup?(): void; } export abstract class CoordinateOverlay< diff --git a/app/packages/looker/src/overlays/detection.ts b/app/packages/looker/src/overlays/detection.ts index 4b9a2f9a0b..d5c62f3f1b 100644 --- a/app/packages/looker/src/overlays/detection.ts +++ b/app/packages/looker/src/overlays/detection.ts @@ -31,7 +31,6 @@ export default class DetectionOverlay< > extends CoordinateOverlay { private is3D: boolean; private labelBoundingBox: BoundingBox; - private imageBitmap: ImageBitmap | null = null; constructor(field, label) { super(field, label); @@ -261,6 +260,13 @@ export default class DetectionOverlay< const oh = state.strokeWidth / state.canvasBBox[3]; return [(bx - ow) * w, (by - oh) * h, (bw + ow * 2) * w, (bh + oh * 2) * h]; } + + public cleanup(): void { + if (this.label.mask?.bitmap) { + this.label.mask?.bitmap.close(); + console.log(">>>cleanup"); + } + } } export const getDetectionPoints = (labels: DetectionLabel[]): Coordinates[] => { diff --git a/app/packages/looker/src/overlays/heatmap.ts b/app/packages/looker/src/overlays/heatmap.ts index 5609550a06..e8e8817643 100644 --- a/app/packages/looker/src/overlays/heatmap.ts +++ b/app/packages/looker/src/overlays/heatmap.ts @@ -205,6 +205,12 @@ export default class HeatmapOverlay getSizeBytes(): number { return sizeBytesEstimate(this.label); } + + public cleanup(): void { + if (this.label.map?.bitmap) { + this.label.map?.bitmap.close(); + } + } } export const getHeatmapPoints = (labels: HeatmapLabel[]): Coordinates[] => { diff --git a/app/packages/looker/src/overlays/segmentation.ts b/app/packages/looker/src/overlays/segmentation.ts index e8e5359823..a4cb098254 100644 --- a/app/packages/looker/src/overlays/segmentation.ts +++ b/app/packages/looker/src/overlays/segmentation.ts @@ -260,6 +260,12 @@ export default class SegmentationOverlay getSizeBytes(): number { return sizeBytesEstimate(this.label); } + + public cleanup(): void { + if (this.label.mask?.bitmap) { + this.label.mask?.bitmap.close(); + } + } } export const getSegmentationPoints = ( From a56869c64e8898822ab214db93c0b5b4c57fa489 Mon Sep 17 00:00:00 2001 From: Sashank Aryal Date: Fri, 22 Nov 2024 17:22:57 -0600 Subject: [PATCH 21/31] add error handling --- app/packages/looker/src/worker/disk-overlay-decoder.ts | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/app/packages/looker/src/worker/disk-overlay-decoder.ts b/app/packages/looker/src/worker/disk-overlay-decoder.ts index 764666ee25..9573d4f49a 100644 --- a/app/packages/looker/src/worker/disk-overlay-decoder.ts +++ b/app/packages/looker/src/worker/disk-overlay-decoder.ts @@ -83,7 +83,15 @@ export const decodeOverlayOnDisk = async ( return; } - const overlayMask = await decodeWithCanvas(overlayImageBlob); + let overlayMask: OverlayMask; + + try { + overlayMask = await decodeWithCanvas(overlayImageBlob); + } catch (e) { + console.error(e); + return; + } + const [overlayHeight, overlayWidth] = overlayMask.shape; // set the `mask` property for this label From 80039ddc5e7dc1e339b1c8b1cf2a654333865abf Mon Sep 17 00:00:00 2001 From: Sashank Aryal Date: Sun, 24 Nov 2024 23:23:32 -0600 Subject: [PATCH 22/31] remove log --- app/packages/looker/src/overlays/detection.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/app/packages/looker/src/overlays/detection.ts b/app/packages/looker/src/overlays/detection.ts index d5c62f3f1b..ec6d45086f 100644 --- a/app/packages/looker/src/overlays/detection.ts +++ b/app/packages/looker/src/overlays/detection.ts @@ -264,7 +264,7 @@ export default class DetectionOverlay< public cleanup(): void { if (this.label.mask?.bitmap) { this.label.mask?.bitmap.close(); - console.log(">>>cleanup"); + this.label.mask.bitmap = null; } } } From 781a4f18dbedc4f1bca792f2c269f76fdb74c8a1 Mon Sep 17 00:00:00 2001 From: Sashank Aryal Date: Sun, 24 Nov 2024 23:33:06 -0600 Subject: [PATCH 23/31] address edge case of empty mask --- app/packages/looker/src/worker/index.ts | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/app/packages/looker/src/worker/index.ts b/app/packages/looker/src/worker/index.ts index e11906b51a..61daa9a173 100644 --- a/app/packages/looker/src/worker/index.ts +++ b/app/packages/looker/src/worker/index.ts @@ -247,6 +247,11 @@ const collectBitmapPromises = (label, cls, bitmapPromises) => { if (label[overlayField]) { const [height, width] = label[overlayField].data.shape; + if (!height || !width) { + label[overlayField].image = null; + return; + } + const imageData = new ImageData( new Uint8ClampedArray(label[overlayField].image), width, From 0e347ffbde23fd089e77b23fa98fd37637f6db6e Mon Sep 17 00:00:00 2001 From: Sashank Aryal Date: Mon, 25 Nov 2024 00:46:29 -0600 Subject: [PATCH 24/31] Fix app test --- .../src/worker/disk-overlay-decoder.test.ts | 51 ++++++++----------- 1 file changed, 20 insertions(+), 31 deletions(-) diff --git a/app/packages/looker/src/worker/disk-overlay-decoder.test.ts b/app/packages/looker/src/worker/disk-overlay-decoder.test.ts index 5d20f454dd..dc7ea31fdf 100644 --- a/app/packages/looker/src/worker/disk-overlay-decoder.test.ts +++ b/app/packages/looker/src/worker/disk-overlay-decoder.test.ts @@ -1,31 +1,25 @@ -import { getSampleSrc } from "@fiftyone/state"; +import { getSampleSrc } from "@fiftyone/state/src/recoil/utils"; import { DETECTIONS, HEATMAP } from "@fiftyone/utilities"; import { beforeEach, describe, expect, it, vi } from "vitest"; import type { Coloring, CustomizeColor } from ".."; import { LabelMask } from "../overlays/base"; import type { Colorscale } from "../state"; import { decodeWithCanvas } from "./canvas-decoder"; -import { fetchWithLinearBackoff } from "./decorated-fetch"; import { decodeOverlayOnDisk, IntermediateMask } from "./disk-overlay-decoder"; +import { enqueueFetch } from "./pooled-fetch"; -vi.mock("@fiftyone/state", () => ({ +vi.mock("@fiftyone/state/src/recoil/utils", () => ({ getSampleSrc: vi.fn(), })); -vi.mock("@fiftyone/utilities", () => ({ - DETECTION: "Detection", - DETECTIONS: "Detections", - HEATMAP: "Heatmap", +vi.mock("./pooled-fetch", () => ({ + enqueueFetch: vi.fn(), })); vi.mock("./canvas-decoder", () => ({ decodeWithCanvas: vi.fn(), })); -vi.mock("./decorated-fetch", () => ({ - fetchWithLinearBackoff: vi.fn(), -})); - const COLORING = {} as Coloring; const COLOR_SCALE = {} as Colorscale; const CUSTOMIZE_COLOR_SETTING: CustomizeColor[] = []; @@ -56,7 +50,7 @@ describe("decodeOverlayOnDisk", () => { ); expect(label.mask).toBeDefined(); - expect(fetchWithLinearBackoff).not.toHaveBeenCalled(); + expect(enqueueFetch).not.toHaveBeenCalled(); }); it("should fetch and decode overlay when label has overlay path field", async () => { @@ -70,7 +64,7 @@ describe("decodeOverlayOnDisk", () => { const overlayMask = { shape: [100, 200] }; vi.mocked(getSampleSrc).mockReturnValue(sampleSrcUrl); - vi.mocked(fetchWithLinearBackoff).mockResolvedValue({ + vi.mocked(enqueueFetch).mockResolvedValue({ blob: () => Promise.resolve(mockBlob), } as Response); vi.mocked(decodeWithCanvas).mockResolvedValue(overlayMask); @@ -87,7 +81,10 @@ describe("decodeOverlayOnDisk", () => { ); expect(getSampleSrc).toHaveBeenCalledWith("/path/to/mask"); - expect(fetchWithLinearBackoff).toHaveBeenCalledWith(sampleSrcUrl); + expect(enqueueFetch).toHaveBeenCalledWith({ + url: sampleSrcUrl, + options: { priority: "low" }, + }); expect(decodeWithCanvas).toHaveBeenCalledWith(mockBlob); expect(label.mask).toBeDefined(); expect(label.mask.data).toBe(overlayMask); @@ -106,9 +103,6 @@ describe("decodeOverlayOnDisk", () => { const overlayMask = { shape: [100, 200] }; vi.mocked(getSampleSrc).mockReturnValue(sampleSrcUrl); - vi.mocked(fetchWithLinearBackoff).mockResolvedValue({ - blob: () => Promise.resolve(mockBlob), - } as Response); vi.mocked(decodeWithCanvas).mockResolvedValue(overlayMask); await decodeOverlayOnDisk( @@ -123,7 +117,10 @@ describe("decodeOverlayOnDisk", () => { ); expect(getSampleSrc).toHaveBeenCalledWith("/path/to/map"); - expect(fetchWithLinearBackoff).toHaveBeenCalledWith(sampleSrcUrl); + expect(enqueueFetch).toHaveBeenCalledWith({ + url: sampleSrcUrl, + options: { priority: "low" }, + }); expect(decodeWithCanvas).toHaveBeenCalledWith(mockBlob); expect(label.map).toBeDefined(); expect(label.map.data).toBe(overlayMask); @@ -144,21 +141,12 @@ describe("decodeOverlayOnDisk", () => { const sampleSrcUrl1 = "http://example.com/path/to/mask1"; const sampleSrcUrl2 = "http://example.com/path/to/mask2"; - const mockBlob1 = new Blob(["mock data 1"], { type: "image/png" }); - const mockBlob2 = new Blob(["mock data 2"], { type: "image/png" }); const overlayMask1 = { shape: [50, 50] }; const overlayMask2 = { shape: [60, 60] }; vi.mocked(getSampleSrc) .mockReturnValueOnce(sampleSrcUrl1) .mockReturnValueOnce(sampleSrcUrl2); - vi.mocked(fetchWithLinearBackoff) - .mockResolvedValueOnce({ - blob: () => Promise.resolve(mockBlob1), - } as Response) - .mockResolvedValueOnce({ - blob: () => Promise.resolve(mockBlob2), - } as Response); vi.mocked(decodeWithCanvas) .mockResolvedValueOnce(overlayMask1) .mockResolvedValueOnce(overlayMask2); @@ -193,9 +181,7 @@ describe("decodeOverlayOnDisk", () => { const sampleSrcUrl = "http://example.com/path/to/mask"; vi.mocked(getSampleSrc).mockReturnValue(sampleSrcUrl); - vi.mocked(fetchWithLinearBackoff).mockRejectedValue( - new Error("Fetch failed") - ); + vi.mocked(enqueueFetch).mockRejectedValue(new Error("Fetch failed")); await decodeOverlayOnDisk( field, @@ -209,7 +195,10 @@ describe("decodeOverlayOnDisk", () => { ); expect(getSampleSrc).toHaveBeenCalledWith("/path/to/mask"); - expect(fetchWithLinearBackoff).toHaveBeenCalledWith(sampleSrcUrl); + expect(enqueueFetch).toHaveBeenCalledWith({ + url: sampleSrcUrl, + options: { priority: "low" }, + }); expect(decodeWithCanvas).not.toHaveBeenCalled(); expect(label.mask).toBeNull(); }); From 882f5aa7d8db3991f03a92fcd3a637e80ab6b253 Mon Sep 17 00:00:00 2001 From: Sashank Aryal Date: Mon, 25 Nov 2024 16:19:33 -0600 Subject: [PATCH 25/31] fix video --- app/packages/looker/src/worker/index.ts | 77 +++++++++++++++++++------ 1 file changed, 58 insertions(+), 19 deletions(-) diff --git a/app/packages/looker/src/worker/index.ts b/app/packages/looker/src/worker/index.ts index 61daa9a173..e5550ce278 100644 --- a/app/packages/looker/src/worker/index.ts +++ b/app/packages/looker/src/worker/index.ts @@ -127,6 +127,10 @@ const processLabels = async ( } const cls = getCls(`${prefix ? prefix : ""}${field}`, schema); + if (!cls) { + continue; + } + for (const label of labels) { if (!label) { continue; @@ -186,11 +190,17 @@ const processLabels = async ( // overlay painting loop for (const field in sample) { let labels = sample[field]; + if (!Array.isArray(labels)) { labels = [labels]; } + const cls = getCls(`${prefix ? prefix : ""}${field}`, schema); + if (!cls) { + continue; + } + for (const label of labels) { if (!label) { continue; @@ -216,11 +226,17 @@ const processLabels = async ( // bitmap generation loop for (const field in sample) { let labels = sample[field]; + if (!Array.isArray(labels)) { labels = [labels]; } + const cls = getCls(`${prefix ? prefix : ""}${field}`, schema); + if (!cls) { + continue; + } + for (const label of labels) { if (!label) { continue; @@ -314,7 +330,7 @@ const processSample = async ({ mapId(sample); const imageBitmapPromises: Promise[] = []; - const maskTargetsBuffers: ArrayBuffer[] = []; + let maskTargetsBuffers: ArrayBuffer[] = []; if (sample?._media_type === "point-cloud" || sample?._media_type === "3d") { process3DLabels(schema, sample); @@ -330,29 +346,51 @@ const processSample = async ({ selectedLabelTags, schema ); - imageBitmapPromises.push(...bitmapPromises); - maskTargetsBuffers.push(...moreMaskTargetsBuffers); + + if (bitmapPromises.length !== 0) { + imageBitmapPromises.push(...bitmapPromises); + } + + if (moreMaskTargetsBuffers.length !== 0) { + maskTargetsBuffers.push(...moreMaskTargetsBuffers); + } } - if (sample.frames && sample.frames.length) { + // this usually only applies to thumbnail frame + // other frames are processed in the stream (see `getSendChunk`) + if (sample.frames?.length) { + const allFramePromises: ReturnType[] = []; for (const frame of sample.frames) { - const [moreBitmapPromises, moreMaskTargetsBuffers] = await processLabels( - frame, - coloring, - "frames.", - sources, - customizeColorSetting, - colorscale, - labelTagColors, - selectedLabelTags, - schema + allFramePromises.push( + processLabels( + frame, + coloring, + "frames.", + sources, + customizeColorSetting, + colorscale, + labelTagColors, + selectedLabelTags, + schema + ) ); - imageBitmapPromises.push(...moreBitmapPromises); - maskTargetsBuffers.push(...moreMaskTargetsBuffers); + } + const framePromisesResolved = await Promise.all(allFramePromises); + for (const [bitmapPromises, buffers] of framePromisesResolved) { + if (bitmapPromises.length !== 0) { + imageBitmapPromises.push(...bitmapPromises); + } + + if (buffers.length !== 0) { + maskTargetsBuffers.push(...buffers); + } } } Promise.all(imageBitmapPromises).then((bitmaps) => { + const flatBitmaps = bitmaps.flat() ?? []; + const flatMaskTargetsBuffers = maskTargetsBuffers.flat() ?? []; + const transferables = [...flatBitmaps, ...flatMaskTargetsBuffers]; postMessage( { method: "processSample", @@ -365,7 +403,7 @@ const processSample = async ({ selectedLabelTags, }, // @ts-ignore - bitmaps.flat().concat(maskTargetsBuffers.flat()) + transferables ); }); }; @@ -501,7 +539,8 @@ const getSendChunk = value.schema ) ) - ).then((buffers) => { + ).then(([bitmaps, buffers]) => { + const transferables = [...bitmaps.flat(), ...buffers.flat()]; postMessage( { method: "frameChunk", @@ -510,7 +549,7 @@ const getSendChunk = uuid, }, // @ts-ignore - buffers.flat() + transferables ); }); } From 532fe7b824b5a1c2e93d5707a3d37a6f2ce27949 Mon Sep 17 00:00:00 2001 From: Sashank Aryal Date: Mon, 25 Nov 2024 16:31:11 -0600 Subject: [PATCH 26/31] fix edge case where bitmap is already set but we need to reinit image buffer --- app/packages/looker/src/worker/disk-overlay-decoder.ts | 8 ++++++++ app/packages/looker/src/worker/index.ts | 3 +++ 2 files changed, 11 insertions(+) diff --git a/app/packages/looker/src/worker/disk-overlay-decoder.ts b/app/packages/looker/src/worker/disk-overlay-decoder.ts index 9573d4f49a..4fdb3d823d 100644 --- a/app/packages/looker/src/worker/disk-overlay-decoder.ts +++ b/app/packages/looker/src/worker/disk-overlay-decoder.ts @@ -53,6 +53,14 @@ export const decodeOverlayOnDisk = async ( const overlayField = overlayFields.canonical; if (Boolean(label[overlayField]) || !Object.hasOwn(label, overlayPathField)) { + // it's possible we're just re-coloring, in which case re-init mask image and set bitmap to null + if (!label[overlayField].image && label[overlayField].bitmap) { + const height = label[overlayField].bitmap.height; + const width = label[overlayField].bitmap.width; + label[overlayField].image = new ArrayBuffer(height * width * 4); + label[overlayField].bitmap.close(); + label[overlayField].bitmap = null; + } // nothing to be done return; } diff --git a/app/packages/looker/src/worker/index.ts b/app/packages/looker/src/worker/index.ts index e5550ce278..1bd503c00f 100644 --- a/app/packages/looker/src/worker/index.ts +++ b/app/packages/looker/src/worker/index.ts @@ -188,6 +188,8 @@ const processLabels = async ( await Promise.allSettled(maskPathDecodingPromises); // overlay painting loop + if (sample.id.endsWith("50a4")) console.log(">>>Painting overlays for hen"); + for (const field in sample) { let labels = sample[field]; @@ -206,6 +208,7 @@ const processLabels = async ( continue; } if (painterFactory[cls]) { + if (sample.id.endsWith("50a4")) debugger; painterPromises.push( painterFactory[cls]( prefix ? prefix + field : field, From f65aa4d8f6a2b43f92ddc7e19b99607d5a2a7424 Mon Sep 17 00:00:00 2001 From: Sashank Aryal Date: Mon, 25 Nov 2024 16:36:46 -0600 Subject: [PATCH 27/31] fix null check and remove logs --- app/packages/looker/src/worker/disk-overlay-decoder.ts | 6 +++++- app/packages/looker/src/worker/index.ts | 3 --- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/app/packages/looker/src/worker/disk-overlay-decoder.ts b/app/packages/looker/src/worker/disk-overlay-decoder.ts index 4fdb3d823d..8730f74bf0 100644 --- a/app/packages/looker/src/worker/disk-overlay-decoder.ts +++ b/app/packages/looker/src/worker/disk-overlay-decoder.ts @@ -54,7 +54,11 @@ export const decodeOverlayOnDisk = async ( if (Boolean(label[overlayField]) || !Object.hasOwn(label, overlayPathField)) { // it's possible we're just re-coloring, in which case re-init mask image and set bitmap to null - if (!label[overlayField].image && label[overlayField].bitmap) { + if ( + label[overlayField] && + label[overlayField].bitmap && + !label[overlayField].image + ) { const height = label[overlayField].bitmap.height; const width = label[overlayField].bitmap.width; label[overlayField].image = new ArrayBuffer(height * width * 4); diff --git a/app/packages/looker/src/worker/index.ts b/app/packages/looker/src/worker/index.ts index 1bd503c00f..e5550ce278 100644 --- a/app/packages/looker/src/worker/index.ts +++ b/app/packages/looker/src/worker/index.ts @@ -188,8 +188,6 @@ const processLabels = async ( await Promise.allSettled(maskPathDecodingPromises); // overlay painting loop - if (sample.id.endsWith("50a4")) console.log(">>>Painting overlays for hen"); - for (const field in sample) { let labels = sample[field]; @@ -208,7 +206,6 @@ const processLabels = async ( continue; } if (painterFactory[cls]) { - if (sample.id.endsWith("50a4")) debugger; painterPromises.push( painterFactory[cls]( prefix ? prefix + field : field, From 3b0ccf6ac5b2d3adb02f92904e5b4002d01d0490 Mon Sep 17 00:00:00 2001 From: Sashank Aryal Date: Tue, 26 Nov 2024 14:02:15 -0600 Subject: [PATCH 28/31] fix frames --- app/packages/looker/src/worker/index.ts | 46 +++++++++++++++++-------- 1 file changed, 31 insertions(+), 15 deletions(-) diff --git a/app/packages/looker/src/worker/index.ts b/app/packages/looker/src/worker/index.ts index e5550ce278..dcf0b2e79b 100644 --- a/app/packages/looker/src/worker/index.ts +++ b/app/packages/looker/src/worker/index.ts @@ -523,9 +523,9 @@ const createReader = ({ const getSendChunk = (uuid: string) => - ({ value }: { done: boolean; value?: FrameChunkResponse }) => { + async ({ value }: { done: boolean; value?: FrameChunkResponse }) => { if (value) { - Promise.all( + const allLabelsPromiseResults = await Promise.allSettled( value.frames.map((frame) => processLabels( frame, @@ -539,19 +539,35 @@ const getSendChunk = value.schema ) ) - ).then(([bitmaps, buffers]) => { - const transferables = [...bitmaps.flat(), ...buffers.flat()]; - postMessage( - { - method: "frameChunk", - frames: value.frames, - range: value.range, - uuid, - }, - // @ts-ignore - transferables - ); - }); + ); + + const allLabelsResults = allLabelsPromiseResults + .filter((result) => result.status === "fulfilled") + .map((result) => result.value); + + const allBuffers = allLabelsResults.map((result) => result[1]).flat(); + + const allBitmapsPromises = allLabelsResults + .map((result) => result[0]) + .flat(); + + const bitmapPromiseResults = ( + await Promise.allSettled(allBitmapsPromises) + ) + .map((result) => (result.status === "fulfilled" ? result.value : [])) + .flat(); + + const transferables = [...bitmapPromiseResults, ...allBuffers]; + postMessage( + { + method: "frameChunk", + frames: value.frames, + range: value.range, + uuid, + }, + // @ts-ignore + transferables + ); } }; From 7d5d7c6bab1e4c7ac30af191c513b41d50abc7bc Mon Sep 17 00:00:00 2001 From: Sashank Aryal Date: Tue, 26 Nov 2024 14:15:21 -0600 Subject: [PATCH 29/31] fix frames memory leak --- app/packages/looker/src/lookers/frame-reader.ts | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/app/packages/looker/src/lookers/frame-reader.ts b/app/packages/looker/src/lookers/frame-reader.ts index a85472e7e4..ce0489dc43 100644 --- a/app/packages/looker/src/lookers/frame-reader.ts +++ b/app/packages/looker/src/lookers/frame-reader.ts @@ -52,7 +52,15 @@ interface AcquireReaderOptions { export const { acquireReader, clearReader } = (() => { const createCache = (removeFrame: RemoveFrame) => { return new LRUCache({ - dispose: (_, key) => removeFrame(key), + dispose: (frame, key) => { + const overlays = frame.overlays; + + for (let i = 0; i < overlays.length; i++) { + overlays[i].cleanup?.(); + } + + removeFrame(key); + }, max: MAX_FRAME_STREAM_SIZE, maxSize: MAX_FRAME_STREAM_SIZE_BYTES, noDisposeOnSet: true, From abc6a13a879846aa89cf1a9a2fa378c02a21db64 Mon Sep 17 00:00:00 2001 From: Sashank Aryal Date: Tue, 26 Nov 2024 14:33:32 -0600 Subject: [PATCH 30/31] add unit tests for pooled fetch --- .../looker/src/worker/decorated-fetch.ts | 18 ++- .../looker/src/worker/pooled-fetch.test.ts | 110 ++++++++++++++++++ .../looker/src/worker/pooled-fetch.ts | 12 +- 3 files changed, 129 insertions(+), 11 deletions(-) create mode 100644 app/packages/looker/src/worker/pooled-fetch.test.ts diff --git a/app/packages/looker/src/worker/decorated-fetch.ts b/app/packages/looker/src/worker/decorated-fetch.ts index 9f0a910ea2..d01f0b48b2 100644 --- a/app/packages/looker/src/worker/decorated-fetch.ts +++ b/app/packages/looker/src/worker/decorated-fetch.ts @@ -3,6 +3,10 @@ const DEFAULT_BASE_DELAY = 200; // list of HTTP status codes that are client errors (4xx) and should not be retried const NON_RETRYABLE_STATUS_CODES = [400, 401, 403, 404, 405, 422]; +export interface RetryOptions { + retries: number; + delay: number; +} class NonRetryableError extends Error { constructor(message: string) { super(message); @@ -13,10 +17,12 @@ class NonRetryableError extends Error { export const fetchWithLinearBackoff = async ( url: string, opts: RequestInit = {}, - retries = DEFAULT_MAX_RETRIES, - delay = DEFAULT_BASE_DELAY + retry: RetryOptions = { + retries: DEFAULT_MAX_RETRIES, + delay: DEFAULT_BASE_DELAY, + } ) => { - for (let i = 0; i < retries; i++) { + for (let i = 0; i < retry.retries; i++) { try { const response = await fetch(url, opts); if (response.ok) { @@ -36,8 +42,10 @@ export const fetchWithLinearBackoff = async ( // immediately throw throw e; } - if (i < retries - 1) { - await new Promise((resolve) => setTimeout(resolve, delay * (i + 1))); + if (i < retry.retries - 1) { + await new Promise((resolve) => + setTimeout(resolve, retry.delay * (i + 1)) + ); } else { // max retries reached throw new Error( diff --git a/app/packages/looker/src/worker/pooled-fetch.test.ts b/app/packages/looker/src/worker/pooled-fetch.test.ts new file mode 100644 index 0000000000..6804f3260a --- /dev/null +++ b/app/packages/looker/src/worker/pooled-fetch.test.ts @@ -0,0 +1,110 @@ +import { beforeEach, describe, expect, it, Mock, vi } from "vitest"; +import { enqueueFetch } from "./pooled-fetch"; + +const MAX_CONCURRENT_REQUESTS = 100; + +// helper function to create a deferred promise +function createDeferredPromise() { + let resolve: (value: T | PromiseLike) => void; + let reject: (reason?: any) => void; + const promise = new Promise((res, rej) => { + resolve = res; + reject = rej; + }); + return { promise, resolve: resolve!, reject: reject! }; +} + +describe("enqueueFetch", () => { + let mockedFetch: Mock; + + beforeEach(() => { + vi.resetAllMocks(); + mockedFetch = vi.fn(); + global.fetch = mockedFetch; + }); + + it("should return response when fetch succeeds", async () => { + const mockResponse = new Response("OK", { status: 200 }); + mockedFetch.mockResolvedValue(mockResponse); + + const response = await enqueueFetch({ url: "https://fiftyone.ai" }); + expect(response).toBe(mockResponse); + }); + + it("should process multiple requests in order", async () => { + const mockResponse1 = new Response("First", { status: 200 }); + const mockResponse2 = new Response("Second", { status: 200 }); + + const deferred1 = createDeferredPromise(); + const deferred2 = createDeferredPromise(); + + mockedFetch + .mockImplementationOnce(() => deferred1.promise) + .mockImplementationOnce(() => deferred2.promise); + + const promise1 = enqueueFetch({ url: "https://fiftyone.ai/1" }); + const promise2 = enqueueFetch({ url: "https://fiftyone.ai/2" }); + + deferred1.resolve(mockResponse1); + + const response1 = await promise1; + expect(response1).toBe(mockResponse1); + + deferred2.resolve(mockResponse2); + + const response2 = await promise2; + expect(response2).toBe(mockResponse2); + }); + + it("should not exceed MAX_CONCURRENT_REQUESTS", async () => { + const numRequests = MAX_CONCURRENT_REQUESTS + 50; + const deferredPromises = []; + + for (let i = 0; i < numRequests; i++) { + const deferred = createDeferredPromise(); + deferredPromises.push(deferred); + mockedFetch.mockImplementationOnce(() => deferred.promise); + enqueueFetch({ url: `https://fiftyone.ai/${i}` }); + } + + // at this point, fetch should have been called MAX_CONCURRENT_REQUESTS times + expect(mockedFetch).toHaveBeenCalledTimes(MAX_CONCURRENT_REQUESTS); + + // resolve all deferred promises + deferredPromises.forEach((deferred, index) => { + deferred.resolve(new Response(`Response ${index}`, { status: 200 })); + }); + + // wait for all promises to resolve + await Promise.all(deferredPromises.map((dp) => dp.promise)); + + // all requests should have been processed + expect(mockedFetch).toHaveBeenCalledTimes(numRequests); + }); + + it("should reject immediately on non-retryable error", async () => { + const mockResponse = new Response("Not Found", { status: 404 }); + mockedFetch.mockResolvedValue(mockResponse); + + await expect(enqueueFetch({ url: "https://fiftyone.ai" })).rejects.toThrow( + "Non-retryable HTTP error: 404" + ); + }); + + it("should retry on retryable errors up to MAX_RETRIES times", async () => { + const MAX_RETRIES = 3; + mockedFetch.mockRejectedValue(new Error("Network Error")); + + await expect( + enqueueFetch({ + url: "https://fiftyone.ai", + retryOptions: { + retries: MAX_RETRIES, + delay: 50, + }, + }) + ).rejects.toThrow("Max retries for fetch reached"); + + expect(mockedFetch).toHaveBeenCalledTimes(MAX_RETRIES); + }); +}); diff --git a/app/packages/looker/src/worker/pooled-fetch.ts b/app/packages/looker/src/worker/pooled-fetch.ts index 3a61c8b0ce..a23e1cb739 100644 --- a/app/packages/looker/src/worker/pooled-fetch.ts +++ b/app/packages/looker/src/worker/pooled-fetch.ts @@ -1,9 +1,10 @@ -import { fetchWithLinearBackoff } from "./decorated-fetch"; +import { fetchWithLinearBackoff, RetryOptions } from "./decorated-fetch"; interface QueueItem { request: { url: string; options?: RequestInit; + retryOptions?: RetryOptions; }; resolve: (value: Response | PromiseLike) => void; reject: (reason?: any) => void; @@ -15,10 +16,9 @@ const MAX_CONCURRENT_REQUESTS = 100; let activeRequests = 0; const requestQueue: QueueItem[] = []; -export const enqueueFetch = (request: { - url: string; - options?: RequestInit; -}): Promise => { +export const enqueueFetch = ( + request: QueueItem["request"] +): Promise => { return new Promise((resolve, reject) => { requestQueue.push({ request, resolve, reject }); processFetchQueue(); @@ -33,7 +33,7 @@ const processFetchQueue = () => { const { request, resolve, reject } = requestQueue.shift(); activeRequests++; - fetchWithLinearBackoff(request.url, request.options) + fetchWithLinearBackoff(request.url, request.options, request.retryOptions) .then((response) => { activeRequests--; resolve(response); From 28e2ab543ff0b23aefae8786e3ef2d74778b87e0 Mon Sep 17 00:00:00 2001 From: Sashank Aryal Date: Tue, 26 Nov 2024 14:38:46 -0600 Subject: [PATCH 31/31] fix fetch unit tests --- .../looker/src/worker/decorated-fetch.test.ts | 30 +++++++++++++++---- 1 file changed, 25 insertions(+), 5 deletions(-) diff --git a/app/packages/looker/src/worker/decorated-fetch.test.ts b/app/packages/looker/src/worker/decorated-fetch.test.ts index 3a9a15e1e7..52fa49d21b 100644 --- a/app/packages/looker/src/worker/decorated-fetch.test.ts +++ b/app/packages/looker/src/worker/decorated-fetch.test.ts @@ -35,7 +35,14 @@ describe("fetchWithLinearBackoff", () => { global.fetch = vi.fn().mockRejectedValue(new Error("Network Error")); await expect( - fetchWithLinearBackoff("http://fiftyone.ai", {}, 3, 10) + fetchWithLinearBackoff( + "http://fiftyone.ai", + {}, + { + retries: 3, + delay: 10, + } + ) ).rejects.toThrowError(new RegExp("Max retries for fetch reached")); expect(global.fetch).toHaveBeenCalledTimes(3); @@ -46,7 +53,14 @@ describe("fetchWithLinearBackoff", () => { global.fetch = vi.fn().mockResolvedValue(mockResponse); await expect( - fetchWithLinearBackoff("http://fiftyone.ai", {}, 5, 10) + fetchWithLinearBackoff( + "http://fiftyone.ai", + {}, + { + retries: 5, + delay: 10, + } + ) ).rejects.toThrow("HTTP error: 500"); expect(global.fetch).toHaveBeenCalledTimes(5); @@ -57,7 +71,14 @@ describe("fetchWithLinearBackoff", () => { global.fetch = vi.fn().mockResolvedValue(mockResponse); await expect( - fetchWithLinearBackoff("http://fiftyone.ai", {}, 5, 10) + fetchWithLinearBackoff( + "http://fiftyone.ai", + {}, + { + retries: 5, + delay: 10, + } + ) ).rejects.toThrow("Non-retryable HTTP error: 404"); expect(global.fetch).toHaveBeenCalledTimes(1); @@ -76,8 +97,7 @@ describe("fetchWithLinearBackoff", () => { const fetchPromise = fetchWithLinearBackoff( "http://fiftyone.ai", {}, - 5, - 10 + { retries: 5, delay: 10 } ); // advance timers to simulate delays