Skip to content
New issue

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

Sign up for GitHub

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

Already on GitHub? Sign in to your account

Jump to bottom

labels rendering performance improvement: create ImageBitmaps in worker #5169

Open
wants to merge 21 commits into
base: develop
Choose a base branch
from
Open

labels rendering performance improvement: create ImageBitmaps in worker #5169

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
21 commits
Select commit Hold shift + click to select a range
6e4775a
change typedef: export ImageBitmap from worker
sashankaryal Nov 20, 2024
53efcf3
implement getSizeBytes in heatmap and segmentation, too
sashankaryal Nov 21, 2024
7e00596
remove canvas creation in detection constructor
sashankaryal Nov 21, 2024
3f30521
remove buffers arg from mask deserializer
sashankaryal Nov 21, 2024
e29f140
make worker send bitmaps back to main thread
sashankaryal Nov 21, 2024
6cfce47
move path decoder to another file
sashankaryal Nov 21, 2024
6cb954e
add unit tests for on disk decoder
sashankaryal Nov 21, 2024
fcc480e
add LabelMask type
sashankaryal Nov 21, 2024
e428200
use modified import for getSampleSrc
sashankaryal Nov 21, 2024
997828d
remove detections specific logic in bitmap collection function
sashankaryal Nov 21, 2024
4ea1df5
return targets buffers to main thread but not image buffers
sashankaryal Nov 21, 2024
4a0cf06
fix frames
sashankaryal Nov 21, 2024
e28dce4
push buffers in deserializer
sashankaryal Nov 21, 2024
c55387a
draw bitmap in segmentation, too
sashankaryal Nov 21, 2024
35cceee
clarify comments
sashankaryal Nov 21, 2024
00d7381
take fetch opts
sashankaryal Nov 22, 2024
bc9ba41
pool fetch
sashankaryal Nov 22, 2024
0fcedd1
pass missing params in recursive calls
sashankaryal Nov 22, 2024
b732e20
add types
sashankaryal Nov 22, 2024
1a0331a
close bitmap when looker destroys
sashankaryal Nov 22, 2024
a56869c
add error handling
sashankaryal Nov 22, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions app/packages/looker/src/lookers/abstract.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 ||
Expand Down
7 changes: 7 additions & 0 deletions app/packages/looker/src/overlays/base.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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";

Expand Down Expand Up @@ -39,6 +40,11 @@ export interface SelectData {
frameNumber?: number;
}

export type LabelMask = {
bitmap?: ImageBitmap;
data?: OverlayMask;
};

export interface RegularLabel extends BaseLabel {
_id?: string;
label?: string;
Expand Down Expand Up @@ -67,6 +73,7 @@ export interface Overlay<State extends Partial<BaseState>> {
getPoints(state: Readonly<State>): Coordinates[];
getSelectData(state: Readonly<State>): SelectData;
getSizeBytes(): number;
cleanup?(): void;
}

export abstract class CoordinateOverlay<
Expand Down
54 changes: 18 additions & 36 deletions app/packages/looker/src/overlays/detection.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,17 +4,19 @@
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";
import {
CONTAINS,
CoordinateOverlay,
LabelMask,
PointInfo,
RegularLabel,
} from "./base";
import { t } from "./util";

export interface DetectionLabel extends RegularLabel {
mask?: {
data: OverlayMask;
image: ArrayBuffer;
};
mask?: LabelMask;
bounding_box: BoundingBox;

// valid for 3D bounding boxes
Expand All @@ -27,10 +29,8 @@ export interface DetectionLabel extends RegularLabel {
export default class DetectionOverlay<
State extends BaseState
> extends CoordinateOverlay<State, DetectionLabel> {
private imageData: ImageData;
private is3D: boolean;
private labelBoundingBox: BoundingBox;
private canvas: HTMLCanvasElement;

constructor(field, label) {
super(field, label);
Expand All @@ -40,32 +40,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<State>): CONTAINS {
Expand Down Expand Up @@ -169,16 +143,17 @@ export default class DetectionOverlay<
}

private drawMask(ctx: CanvasRenderingContext2D, state: Readonly<State>) {
if (!this.canvas) {
if (!this.label.mask?.bitmap) {
return;
}

const [tlx, tly, w, h] = this.label.bounding_box;
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,
sashankaryal marked this conversation as resolved.
Show resolved Hide resolved
x,
y,
w * state.canvasBBox[2],
Expand Down Expand Up @@ -285,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");
}
}
Comment on lines +264 to +269
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🛠️ Refactor suggestion

Remove debug log and enhance cleanup method

The cleanup method is essential but needs improvements:

  1. Remove the debug console.log
  2. Add error handling for bitmap closure
 public cleanup(): void {
   if (this.label.mask?.bitmap) {
-    this.label.mask?.bitmap.close();
-    console.log(">>>cleanup");
+    try {
+      this.label.mask.bitmap.close();
+    } catch (error) {
+      console.error('Error closing bitmap:', error);
+    }
+    this.label.mask.bitmap = null;
   }
 }
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
public cleanup(): void {
if (this.label.mask?.bitmap) {
this.label.mask?.bitmap.close();
console.log(">>>cleanup");
}
}
public cleanup(): void {
if (this.label.mask?.bitmap) {
try {
this.label.mask.bitmap.close();
} catch (error) {
console.error('Error closing bitmap:', error);
}
this.label.mask.bitmap = null;
}
}

💡 Codebase verification

Memory leak risk: Cleanup method not properly invoked

Based on the codebase analysis, there are potential memory leak concerns:

  • The cleanup() method is defined in multiple overlay classes (Detection, Segmentation, Heatmap) but there's no clear evidence of it being called during overlay disposal
  • While overlays are instantiated in overlays/index.ts, there's no corresponding cleanup logic visible
  • No component unmounting or disposal logic was found that invokes these cleanup methods

Recommendations:

  • Ensure cleanup() is called when overlays are disposed
  • Add cleanup calls in the parent component's unmounting lifecycle
  • Consider implementing a disposal pattern in the overlay management logic
🔗 Analysis chain

Verify cleanup method invocation

Ensure that the cleanup method is called appropriately when the overlay is destroyed or removed.

Let me gather more information about where these cleanup methods are being called.

Let me check one more aspect to see how these overlay classes are used and disposed.

Let me try one final search to understand the overlay lifecycle management.

🏁 Scripts executed

The following scripts were executed for the analysis:

Script:

#!/bin/bash
# Search for cleanup method calls
rg -A 5 "cleanup\(\)" app/packages/looker/src/

# Search for component unmounting or disposal logic
ast-grep --pattern 'class $_ extends $_ {
  $$$
  componentWillUnmount() {
    $$$
  }
  $$$
}'

Length of output: 1354

Script:

#!/bin/bash
# Search for imports and usage of these overlay classes
rg -A 2 "import.*\{.*(Detection|Segmentation|Heatmap).*\}" app/packages/looker/src/

# Search for instances where these overlays are destroyed/disposed
rg -B 2 -A 2 "dispose\(\)|destroy\(\)" app/packages/looker/src/

# Look for any class that might be managing these overlays
ast-grep --pattern 'class $_ {
  $$$
  dispose() {
    $$$
  }
  $$$
}'

Length of output: 3128

Script:

#!/bin/bash
# Search for overlay usage in components and their lifecycle methods
rg -B 3 -A 3 "new (Detection|Segmentation|Heatmap)Overlay" app/packages/looker/src/

# Look for any cleanup calls in component unmounting
ast-grep --pattern 'componentWillUnmount() {
  $$$
}'

# Check if overlays are managed in any manager class
rg -B 2 -A 5 "overlay" app/packages/looker/src/managers/

Length of output: 272

Script:

#!/bin/bash
# Search for overlay instantiation and management
rg -B 3 -A 3 "new.*Overlay" app/packages/looker/src/

# Look for cleanup or disposal patterns in the main components
rg -B 3 -A 3 "\.cleanup\(\)" app/packages/looker/src/

# Check base overlay class for lifecycle management
ast-grep --pattern 'class $_ {
  $$$
  cleanup() {
    $$$
  }
  $$$
}'

Length of output: 723

}

export const getDetectionPoints = (labels: DetectionLabel[]): Coordinates[] => {
Expand Down
56 changes: 16 additions & 40 deletions app/packages/looker/src/overlays/heatmap.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,28 +6,25 @@ import {
getColor,
getRGBA,
getRGBAColor,
sizeBytesEstimate,
} 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";
import {
BaseLabel,
CONTAINS,
LabelMask,
Overlay,
PointInfo,
SelectData,
isShown,
} from "./base";
import { strokeCanvasRect, t } from "./util";

interface HeatMap {
data: OverlayMask;
image: ArrayBuffer;
}

interface HeatmapLabel extends BaseLabel {
map?: HeatMap;
map?: LabelMask;
range?: [number, number];
}

Expand All @@ -45,8 +42,6 @@ export default class HeatmapOverlay<State extends BaseState>
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;
Expand All @@ -68,25 +63,6 @@ export default class HeatmapOverlay<State extends BaseState>
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<State>): CONTAINS {
Expand All @@ -101,22 +77,12 @@ export default class HeatmapOverlay<State extends BaseState>
}

draw(ctx: CanvasRenderingContext2D, state: Readonly<State>): 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 [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;
}

Expand Down Expand Up @@ -235,6 +201,16 @@ export default class HeatmapOverlay<State extends BaseState>

return this.targets[index];
}

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[] => {
Expand Down
46 changes: 17 additions & 29 deletions app/packages/looker/src/overlays/segmentation.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,12 +2,13 @@
* Copyright 2017-2024, Voxel51, Inc.
*/

import { getColor } from "@fiftyone/utilities";
import { ARRAY_TYPES, OverlayMask, TypedArray } from "../numpy";
import { getColor, sizeBytesEstimate } from "@fiftyone/utilities";
import { ARRAY_TYPES, TypedArray } from "../numpy";
import { BaseState, Coordinates, MaskTargets } from "../state";
import {
BaseLabel,
CONTAINS,
LabelMask,
Overlay,
PointInfo,
SelectData,
Expand All @@ -16,10 +17,7 @@ import {
import { isRgbMaskTargets, strokeCanvasRect, t } from "./util";

interface SegmentationLabel extends BaseLabel {
mask?: {
data: OverlayMask;
image: ArrayBuffer;
};
mask?: LabelMask;
}

interface SegmentationInfo extends BaseLabel {
Expand All @@ -34,8 +32,6 @@ export default class SegmentationOverlay<State extends BaseState>
readonly field: string;
private label: SegmentationLabel;
private targets?: TypedArray;
private canvas: HTMLCanvasElement;
private imageData: ImageData;

private isRgbMaskTargets = false;

Expand All @@ -53,6 +49,7 @@ export default class SegmentationOverlay<State extends BaseState>
if (!this.label.mask) {
return;
}

const [height, width] = this.label.mask.data.shape;

if (!height || !width) {
Expand All @@ -62,25 +59,6 @@ export default class SegmentationOverlay<State extends BaseState>
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<State>): CONTAINS {
Expand All @@ -99,12 +77,12 @@ export default class SegmentationOverlay<State extends BaseState>
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;
}

Expand Down Expand Up @@ -278,6 +256,16 @@ export default class SegmentationOverlay<State extends BaseState>
}
return this.targets[index];
}

getSizeBytes(): number {
return sizeBytesEstimate(this.label);
}
sashankaryal marked this conversation as resolved.
Show resolved Hide resolved

public cleanup(): void {
if (this.label.mask?.bitmap) {
this.label.mask?.bitmap.close();
}
}
}

export const getSegmentationPoints = (
Expand Down
Loading
Loading