Skip to content

Commit 623ab30

Browse files
justadudewhohacksjustadudewhohacks
committed
fixed memory leaks in extractFaces, imageTensorToCanvas, toNetInput with batch size 1 tensor4d array + removed tests using createFakeHTMLVideoElement as it seems to not load the example video properly some times
1 parent d2e2d36 commit 623ab30

File tree

8 files changed

+187
-51
lines changed

8 files changed

+187
-51
lines changed

‎src/NetInput.ts

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,16 @@ export class NetInput {
2929
return tf.clone(input as tf.Tensor3D)
3030
}
3131

32+
if (isTensor4D(input)) {
33+
const shape = (input as tf.Tensor4D).shape
34+
const batchSize = shape[0]
35+
if (batchSize !== 1) {
36+
throw new Error(`NetInput - tf.Tensor4D with batchSize ${batchSize} passed, but not supported in input array`)
37+
}
38+
39+
return (input as tf.Tensor4D).reshape(shape.slice(1) as [number, number, number]) as tf.Tensor3D
40+
}
41+
3242
return tf.fromPixels(
3343
input instanceof HTMLCanvasElement ? input : createCanvasFromMedia(input as HTMLImageElement | HTMLVideoElement)
3444
)

‎src/extractFaces.ts

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ import { Rect } from './Rect';
33
import { toNetInput } from './toNetInput';
44
import { TNetInput } from './types';
55
import { createCanvas, getContext2dOrThrow, imageTensorToCanvas } from './utils';
6+
import * as tf from '@tensorflow/tfjs-core';
67

78
/**
89
* Extracts the image regions containing the detected faces.
@@ -29,6 +30,10 @@ export async function extractFaces(
2930
}
3031

3132
canvas = await imageTensorToCanvas(netInput.inputs[0])
33+
34+
if (netInput.isManaged) {
35+
netInput.dispose()
36+
}
3237
}
3338

3439
const ctx = getContext2dOrThrow(canvas)

‎src/toNetInput.ts

Lines changed: 10 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -41,30 +41,25 @@ export async function toNetInput(
4141

4242
const getIdxHint = (idx: number) => Array.isArray(inputs) ? ` at input index ${idx}:` : ''
4343

44-
const inputArray = inputArgArray
45-
.map(resolveInput)
46-
.map((input, i) => {
47-
if (isTensor4D(input)) {
48-
// if tf.Tensor4D is passed in the input array, the batch size has to be 1
49-
const batchSize = input.shape[0]
50-
if (batchSize !== 1) {
51-
throw new Error(`toNetInput -${getIdxHint(i)} tf.Tensor4D with batchSize ${batchSize} passed, but not supported in input array`)
52-
}
53-
// to tf.Tensor3D
54-
return input.reshape(input.shape.slice(1))
55-
}
56-
return input
57-
})
44+
const inputArray = inputArgArray.map(resolveInput)
5845

5946
inputArray.forEach((input, i) => {
60-
if (!isMediaElement(input) && !isTensor3D(input)) {
47+
if (!isMediaElement(input) && !isTensor3D(input) && !isTensor4D(input)) {
6148

6249
if (typeof inputArgArray[i] === 'string') {
6350
throw new Error(`toNetInput -${getIdxHint(i)} string passed, but could not resolve HTMLElement for element id ${inputArgArray[i]}`)
6451
}
6552

6653
throw new Error(`toNetInput -${getIdxHint(i)} expected media to be of type HTMLImageElement | HTMLVideoElement | HTMLCanvasElement | tf.Tensor3D, or to be an element id`)
6754
}
55+
56+
if (isTensor4D(input)) {
57+
// if tf.Tensor4D is passed in the input array, the batch size has to be 1
58+
const batchSize = input.shape[0]
59+
if (batchSize !== 1) {
60+
throw new Error(`toNetInput -${getIdxHint(i)} tf.Tensor4D with batchSize ${batchSize} passed, but not supported in input array`)
61+
}
62+
}
6863
})
6964

7065
// wait for all media elements being loaded

‎src/utils.ts

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -113,7 +113,10 @@ export async function imageTensorToCanvas(
113113
const targetCanvas = canvas || document.createElement('canvas')
114114

115115
const [height, width, numChannels] = imgTensor.shape.slice(isTensor4D(imgTensor) ? 1 : 0)
116-
await tf.toPixels(imgTensor.as3D(height, width, numChannels).toInt(), targetCanvas)
116+
const imgTensor3D = tf.tidy(() => imgTensor.as3D(height, width, numChannels).toInt())
117+
await tf.toPixels(imgTensor3D, targetCanvas)
118+
119+
imgTensor3D.dispose()
117120

118121
return targetCanvas
119122
}

‎test/tests/NetInput.test.ts

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,12 +7,11 @@ import { expectAllTensorsReleased, tensor3D } from '../utils';
77

88
describe('NetInput', () => {
99

10-
let imgEl: HTMLImageElement, canvasEl: HTMLCanvasElement
10+
let imgEl: HTMLImageElement
1111

1212
beforeAll(async () => {
1313
const img = await (await fetch('base/test/images/face1.png')).blob()
1414
imgEl = await bufferToImage(img)
15-
canvasEl = createCanvasFromMedia(imgEl)
1615
})
1716

1817
describe('no memory leaks', () => {

‎test/tests/e2e/faceLandmarkNet.test.ts

Lines changed: 91 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,8 +5,9 @@ import { isTensor3D } from '../../../src/commons/isTensor';
55
import { FaceLandmarks } from '../../../src/faceLandmarkNet/FaceLandmarks';
66
import { Point } from '../../../src/Point';
77
import { Dimensions, TMediaElement } from '../../../src/types';
8-
import { expectMaxDelta, expectAllTensorsReleased } from '../../utils';
8+
import { expectMaxDelta, expectAllTensorsReleased, tensor3D } from '../../utils';
99
import { NetInput } from '../../../src/NetInput';
10+
import { toNetInput } from '../../../src';
1011

1112
function getInputDims (input: tf.Tensor | TMediaElement): Dimensions {
1213
if (input instanceof tf.Tensor) {
@@ -255,6 +256,53 @@ describe('faceLandmarkNet', () => {
255256
})
256257
})
257258

259+
it('single tf.Tensor3D', async () => {
260+
const tensor = tf.fromPixels(imgEl1)
261+
262+
await expectAllTensorsReleased(async () => {
263+
const netInput = (new NetInput([tensor])).managed()
264+
const outTensor = await faceLandmarkNet.forwardInput(netInput)
265+
outTensor.dispose()
266+
})
267+
268+
tensor.dispose()
269+
})
270+
271+
it('multiple tf.Tensor3Ds', async () => {
272+
const tensors = [imgEl1, imgEl1, imgEl1].map(el => tf.fromPixels(el))
273+
274+
await expectAllTensorsReleased(async () => {
275+
const netInput = (new NetInput(tensors)).managed()
276+
const outTensor = await faceLandmarkNet.forwardInput(netInput)
277+
outTensor.dispose()
278+
})
279+
280+
tensors.forEach(t => t.dispose())
281+
})
282+
283+
it('single batch size 1 tf.Tensor4Ds', async () => {
284+
const tensor = tf.tidy(() => tf.fromPixels(imgEl1).expandDims()) as tf.Tensor4D
285+
286+
await expectAllTensorsReleased(async () => {
287+
const outTensor = await faceLandmarkNet.forwardInput(await toNetInput(tensor, true))
288+
outTensor.dispose()
289+
})
290+
291+
tensor.dispose()
292+
})
293+
294+
it('multiple batch size 1 tf.Tensor4Ds', async () => {
295+
const tensors = [imgEl1, imgEl1, imgEl1]
296+
.map(el => tf.tidy(() => tf.fromPixels(el).expandDims())) as tf.Tensor4D[]
297+
298+
await expectAllTensorsReleased(async () => {
299+
const outTensor = await faceLandmarkNet.forwardInput(await toNetInput(tensors, true))
300+
outTensor.dispose()
301+
})
302+
303+
tensors.forEach(t => t.dispose())
304+
})
305+
258306
})
259307

260308
describe('detectLandmarks', () => {
@@ -271,6 +319,48 @@ describe('faceLandmarkNet', () => {
271319
})
272320
})
273321

322+
it('single tf.Tensor3D', async () => {
323+
const tensor = tf.fromPixels(imgEl1)
324+
325+
await expectAllTensorsReleased(async () => {
326+
await faceLandmarkNet.detectLandmarks(tensor)
327+
})
328+
329+
tensor.dispose()
330+
})
331+
332+
it('multiple tf.Tensor3Ds', async () => {
333+
const tensors = [imgEl1, imgEl1, imgEl1].map(el => tf.fromPixels(el))
334+
335+
336+
await expectAllTensorsReleased(async () => {
337+
await faceLandmarkNet.detectLandmarks(tensors)
338+
})
339+
340+
tensors.forEach(t => t.dispose())
341+
})
342+
343+
it('single batch size 1 tf.Tensor4Ds', async () => {
344+
const tensor = tf.tidy(() => tf.fromPixels(imgEl1).expandDims()) as tf.Tensor4D
345+
346+
await expectAllTensorsReleased(async () => {
347+
await faceLandmarkNet.detectLandmarks(tensor)
348+
})
349+
350+
tensor.dispose()
351+
})
352+
353+
it('multiple batch size 1 tf.Tensor4Ds', async () => {
354+
const tensors = [imgEl1, imgEl1, imgEl1]
355+
.map(el => tf.tidy(() => tf.fromPixels(el).expandDims())) as tf.Tensor4D[]
356+
357+
await expectAllTensorsReleased(async () => {
358+
await faceLandmarkNet.detectLandmarks(tensors)
359+
})
360+
361+
tensors.forEach(t => t.dispose())
362+
})
363+
274364
})
275365
})
276366

‎test/tests/toNetInput.test.ts

Lines changed: 66 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,8 @@
1+
import { tf } from '../../src';
12
import { NetInput } from '../../src/NetInput';
23
import { toNetInput } from '../../src/toNetInput';
34
import { bufferToImage, createCanvasFromMedia } from '../../src/utils';
4-
import { createFakeHTMLVideoElement } from '../utils';
5+
import { expectAllTensorsReleased } from '../utils';
56

67
describe('toNetInput', () => {
78

@@ -21,13 +22,6 @@ describe('toNetInput', () => {
2122
expect(netInput.batchSize).toEqual(1)
2223
})
2324

24-
it('from HTMLVideoElement', async () => {
25-
const videoEl = await createFakeHTMLVideoElement()
26-
const netInput = await toNetInput(videoEl, true)
27-
expect(netInput instanceof NetInput).toBe(true)
28-
expect(netInput.batchSize).toEqual(1)
29-
})
30-
3125
it('from HTMLCanvasElement', async () => {
3226
const netInput = await toNetInput(canvasEl, true)
3327
expect(netInput instanceof NetInput).toBe(true)
@@ -43,20 +37,6 @@ describe('toNetInput', () => {
4337
expect(netInput.batchSize).toEqual(2)
4438
})
4539

46-
it('from HTMLVideoElement array', async () => {
47-
const videoElements = [
48-
await createFakeHTMLVideoElement(),
49-
await createFakeHTMLVideoElement()
50-
]
51-
videoElements.forEach(videoEl =>
52-
spyOnProperty(videoEl, 'readyState', 'get').and.returnValue(4)
53-
)
54-
55-
const netInput = await toNetInput(videoElements, true)
56-
expect(netInput instanceof NetInput).toBe(true)
57-
expect(netInput.batchSize).toEqual(2)
58-
})
59-
6040
it('from HTMLCanvasElement array', async () => {
6141
const netInput = await toNetInput([
6242
canvasEl,
@@ -70,7 +50,7 @@ describe('toNetInput', () => {
7050
const netInput = await toNetInput([
7151
imgEl,
7252
canvasEl,
73-
await createFakeHTMLVideoElement()
53+
canvasEl
7454
], true)
7555
expect(netInput instanceof NetInput).toBe(true)
7656
expect(netInput.batchSize).toEqual(3)
@@ -111,4 +91,67 @@ describe('toNetInput', () => {
11191

11292
})
11393

94+
describe('no memory leaks', () => {
95+
96+
it('single image element', async () => {
97+
await expectAllTensorsReleased(async () => {
98+
const netInput = await toNetInput(imgEl)
99+
netInput.dispose()
100+
})
101+
})
102+
103+
it('multiple image elements', async () => {
104+
await expectAllTensorsReleased(async () => {
105+
const netInput = await toNetInput([imgEl, imgEl, imgEl])
106+
netInput.dispose()
107+
})
108+
})
109+
110+
it('single tf.Tensor3D', async () => {
111+
const tensor = tf.fromPixels(imgEl)
112+
113+
await expectAllTensorsReleased(async () => {
114+
const netInput = await toNetInput(tensor)
115+
netInput.dispose()
116+
})
117+
118+
tensor.dispose()
119+
})
120+
121+
it('multiple tf.Tensor3Ds', async () => {
122+
const tensors = [imgEl, imgEl, imgEl].map(el => tf.fromPixels(el))
123+
124+
await expectAllTensorsReleased(async () => {
125+
const netInput = await toNetInput(tensors)
126+
netInput.dispose()
127+
})
128+
129+
tensors.forEach(t => t.dispose())
130+
})
131+
132+
it('single batch size 1 tf.Tensor4Ds', async () => {
133+
const tensor = tf.tidy(() => tf.fromPixels(imgEl).expandDims()) as tf.Tensor4D
134+
135+
await expectAllTensorsReleased(async () => {
136+
const netInput = await toNetInput(tensor)
137+
netInput.dispose()
138+
})
139+
140+
tensor.dispose()
141+
})
142+
143+
it('multiple batch size 1 tf.Tensor4Ds', async () => {
144+
const tensors = [imgEl, imgEl, imgEl]
145+
.map(el => tf.tidy(() => tf.fromPixels(el).expandDims())) as tf.Tensor4D[]
146+
147+
await expectAllTensorsReleased(async () => {
148+
const netInput = await toNetInput(tensors)
149+
netInput.dispose()
150+
})
151+
152+
tensors.forEach(t => t.dispose())
153+
})
154+
155+
})
156+
114157
})

‎test/utils.ts

Lines changed: 0 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -12,15 +12,6 @@ export function expectMaxDelta(val1: number, val2: number, maxDelta: number) {
1212
expect(Math.abs(val1 - val2)).toBeLessThan(maxDelta)
1313
}
1414

15-
export async function createFakeHTMLVideoElement() {
16-
const videoEl = document.createElement('video')
17-
videoEl.muted = true
18-
videoEl.src = 'base/test/media/video.mp4'
19-
await videoEl.pause()
20-
await videoEl.play()
21-
return videoEl
22-
}
23-
2415
export async function expectAllTensorsReleased(fn: () => any) {
2516
const numTensorsBefore = tf.memory().numTensors
2617
await fn()

0 commit comments

Comments
 (0)