fixed memory leaks in extractFaces, imageTensorToCanvas, toNetInput with batch…

src/NetInput.ts

@@ -29,6 +29,16 @@ export class NetInput {
           return tf.clone(input as tf.Tensor3D)
         }
 
+        if (isTensor4D(input)) {
+          const shape = (input as tf.Tensor4D).shape
+          const batchSize = shape[0]
+          if (batchSize !== 1) {
+            throw new Error(`NetInput - tf.Tensor4D with batchSize ${batchSize} passed, but not supported in input array`)
+          }
+
+          return (input as tf.Tensor4D).reshape(shape.slice(1) as [number, number, number]) as tf.Tensor3D
+        }
+
         return tf.fromPixels(
           input instanceof HTMLCanvasElement ? input : createCanvasFromMedia(input as HTMLImageElement | HTMLVideoElement)
         )

src/extractFaces.ts

@@ -3,6 +3,7 @@ import { Rect } from './Rect';
 import { toNetInput } from './toNetInput';
 import { TNetInput } from './types';
 import { createCanvas, getContext2dOrThrow, imageTensorToCanvas } from './utils';
+import * as tf from '@tensorflow/tfjs-core';
 
 /**
  * Extracts the image regions containing the detected faces.
@@ -29,6 +30,10 @@ export async function extractFaces(
     }
 
     canvas = await imageTensorToCanvas(netInput.inputs[0])
+
+    if (netInput.isManaged) {
+      netInput.dispose()
+    }
   }
 
   const ctx = getContext2dOrThrow(canvas)

src/toNetInput.ts

@@ -41,23 +41,10 @@ export async function toNetInput(
 
   const getIdxHint = (idx: number) => Array.isArray(inputs) ? ` at input index ${idx}:` : ''
 
-  const inputArray = inputArgArray
-    .map(resolveInput)
-    .map((input, i) => {
-      if (isTensor4D(input)) {
-        // if tf.Tensor4D is passed in the input array, the batch size has to be 1
-        const batchSize = input.shape[0]
-        if (batchSize !== 1) {
-          throw new Error(`toNetInput -${getIdxHint(i)} tf.Tensor4D with batchSize ${batchSize} passed, but not supported in input array`)
-        }
-        // to tf.Tensor3D
-        return input.reshape(input.shape.slice(1))
-      }
-      return input
-    })
+  const inputArray = inputArgArray.map(resolveInput)
 
   inputArray.forEach((input, i) => {
-    if (!isMediaElement(input) && !isTensor3D(input)) {
+    if (!isMediaElement(input) && !isTensor3D(input) && !isTensor4D(input)) {
 
       if (typeof inputArgArray[i] === 'string') {
         throw new Error(`toNetInput -${getIdxHint(i)} string passed, but could not resolve HTMLElement for element id ${inputArgArray[i]}`)
@@ -65,6 +52,14 @@ export async function toNetInput(
 
       throw new Error(`toNetInput -${getIdxHint(i)} expected media to be of type HTMLImageElement | HTMLVideoElement | HTMLCanvasElement | tf.Tensor3D, or to be an element id`)
     }
+
+    if (isTensor4D(input)) {
+      // if tf.Tensor4D is passed in the input array, the batch size has to be 1
+      const batchSize = input.shape[0]
+      if (batchSize !== 1) {
+        throw new Error(`toNetInput -${getIdxHint(i)} tf.Tensor4D with batchSize ${batchSize} passed, but not supported in input array`)
+      }
+    }
   })
 
   // wait for all media elements being loaded

src/utils.ts

@@ -113,7 +113,10 @@ export async function imageTensorToCanvas(
   const targetCanvas = canvas ||  document.createElement('canvas')
 
   const [height, width, numChannels] = imgTensor.shape.slice(isTensor4D(imgTensor) ? 1 : 0)
-  await tf.toPixels(imgTensor.as3D(height, width, numChannels).toInt(), targetCanvas)
+  const imgTensor3D = tf.tidy(() => imgTensor.as3D(height, width, numChannels).toInt())
+  await tf.toPixels(imgTensor3D, targetCanvas)
+
+  imgTensor3D.dispose()
 
   return targetCanvas
 }
\ No newline at end of file

test/tests/NetInput.test.ts

@@ -7,12 +7,11 @@ import { expectAllTensorsReleased, tensor3D } from '../utils';
 
 describe('NetInput', () => {
 
-  let imgEl: HTMLImageElement, canvasEl: HTMLCanvasElement
+  let imgEl: HTMLImageElement
 
   beforeAll(async () => {
     const img = await (await fetch('base/test/images/face1.png')).blob()
     imgEl = await bufferToImage(img)
-    canvasEl = createCanvasFromMedia(imgEl)
   })
 
   describe('no memory leaks', () => {

test/tests/e2e/faceLandmarkNet.test.ts

@@ -5,8 +5,9 @@ import { isTensor3D } from '../../../src/commons/isTensor';
 import { FaceLandmarks } from '../../../src/faceLandmarkNet/FaceLandmarks';
 import { Point } from '../../../src/Point';
 import { Dimensions, TMediaElement } from '../../../src/types';
-import { expectMaxDelta, expectAllTensorsReleased } from '../../utils';
+import { expectMaxDelta, expectAllTensorsReleased, tensor3D } from '../../utils';
 import { NetInput } from '../../../src/NetInput';
+import { toNetInput } from '../../../src';
 
 function getInputDims (input: tf.Tensor | TMediaElement): Dimensions {
   if (input instanceof tf.Tensor) {
@@ -255,6 +256,53 @@ describe('faceLandmarkNet', () => {
         })
       })
 
+      it('single tf.Tensor3D', async () => {
+        const tensor = tf.fromPixels(imgEl1)
+
+        await expectAllTensorsReleased(async () => {
+          const netInput = (new NetInput([tensor])).managed()
+          const outTensor = await faceLandmarkNet.forwardInput(netInput)
+          outTensor.dispose()
+        })
+
+        tensor.dispose()
+      })
+
+      it('multiple tf.Tensor3Ds', async () => {
+        const tensors = [imgEl1, imgEl1, imgEl1].map(el => tf.fromPixels(el))
+
+        await expectAllTensorsReleased(async () => {
+          const netInput = (new NetInput(tensors)).managed()
+          const outTensor = await faceLandmarkNet.forwardInput(netInput)
+          outTensor.dispose()
+        })
+
+        tensors.forEach(t => t.dispose())
+      })
+
+      it('single batch size 1 tf.Tensor4Ds', async () => {
+        const tensor = tf.tidy(() => tf.fromPixels(imgEl1).expandDims()) as tf.Tensor4D
+
+        await expectAllTensorsReleased(async () => {
+          const outTensor = await faceLandmarkNet.forwardInput(await toNetInput(tensor, true))
+          outTensor.dispose()
+        })
+
+        tensor.dispose()
+      })
+
+      it('multiple batch size 1 tf.Tensor4Ds', async () => {
+        const tensors = [imgEl1, imgEl1, imgEl1]
+          .map(el => tf.tidy(() => tf.fromPixels(el).expandDims())) as tf.Tensor4D[]
+
+        await expectAllTensorsReleased(async () => {
+          const outTensor = await faceLandmarkNet.forwardInput(await toNetInput(tensors, true))
+          outTensor.dispose()
+        })
+
+        tensors.forEach(t => t.dispose())
+      })
+
     })
 
     describe('detectLandmarks', () => {
@@ -271,6 +319,48 @@ describe('faceLandmarkNet', () => {
         })
       })
 
+      it('single tf.Tensor3D', async () => {
+        const tensor = tf.fromPixels(imgEl1)
+
+        await expectAllTensorsReleased(async () => {
+          await faceLandmarkNet.detectLandmarks(tensor)
+        })
+
+        tensor.dispose()
+      })
+
+      it('multiple tf.Tensor3Ds', async () => {
+        const tensors = [imgEl1, imgEl1, imgEl1].map(el => tf.fromPixels(el))
+
+
+        await expectAllTensorsReleased(async () => {
+          await faceLandmarkNet.detectLandmarks(tensors)
+        })
+
+        tensors.forEach(t => t.dispose())
+      })
+
+      it('single batch size 1 tf.Tensor4Ds', async () => {
+        const tensor = tf.tidy(() => tf.fromPixels(imgEl1).expandDims()) as tf.Tensor4D
+
+        await expectAllTensorsReleased(async () => {
+          await faceLandmarkNet.detectLandmarks(tensor)
+        })
+
+        tensor.dispose()
+      })
+
+      it('multiple batch size 1 tf.Tensor4Ds', async () => {
+        const tensors = [imgEl1, imgEl1, imgEl1]
+          .map(el => tf.tidy(() => tf.fromPixels(el).expandDims())) as tf.Tensor4D[]
+
+        await expectAllTensorsReleased(async () => {
+          await faceLandmarkNet.detectLandmarks(tensors)
+        })
+
+        tensors.forEach(t => t.dispose())
+      })
+
     })
   })
 

test/tests/toNetInput.test.ts

+import { tf } from '../../src';
 import { NetInput } from '../../src/NetInput';
 import { toNetInput } from '../../src/toNetInput';
 import { bufferToImage, createCanvasFromMedia } from '../../src/utils';
-import { createFakeHTMLVideoElement } from '../utils';
+import { expectAllTensorsReleased } from '../utils';
 
 describe('toNetInput', () => {
 
@@ -21,13 +22,6 @@ describe('toNetInput', () => {
       expect(netInput.batchSize).toEqual(1)
     })
 
-    it('from HTMLVideoElement', async () => {
-      const videoEl = await createFakeHTMLVideoElement()
-      const netInput = await toNetInput(videoEl, true)
-      expect(netInput instanceof NetInput).toBe(true)
-      expect(netInput.batchSize).toEqual(1)
-    })
-
     it('from HTMLCanvasElement', async () => {
       const netInput = await toNetInput(canvasEl, true)
       expect(netInput instanceof NetInput).toBe(true)
@@ -43,20 +37,6 @@ describe('toNetInput', () => {
       expect(netInput.batchSize).toEqual(2)
     })
 
-    it('from HTMLVideoElement array', async () => {
-      const videoElements = [
-        await createFakeHTMLVideoElement(),
-        await createFakeHTMLVideoElement()
-      ]
-      videoElements.forEach(videoEl =>
-        spyOnProperty(videoEl, 'readyState', 'get').and.returnValue(4)
-      )
-
-      const netInput = await toNetInput(videoElements, true)
-      expect(netInput instanceof NetInput).toBe(true)
-      expect(netInput.batchSize).toEqual(2)
-    })
-
     it('from HTMLCanvasElement array', async () => {
       const netInput = await toNetInput([
         canvasEl,
@@ -70,7 +50,7 @@ describe('toNetInput', () => {
       const netInput = await toNetInput([
         imgEl,
         canvasEl,
-        await createFakeHTMLVideoElement()
+        canvasEl
       ], true)
       expect(netInput instanceof NetInput).toBe(true)
       expect(netInput.batchSize).toEqual(3)
@@ -111,4 +91,67 @@ describe('toNetInput', () => {
 
   })
 
+  describe('no memory leaks', () => {
+
+    it('single image element', async () => {
+      await expectAllTensorsReleased(async () => {
+        const netInput = await toNetInput(imgEl)
+        netInput.dispose()
+      })
+    })
+
+    it('multiple image elements', async () => {
+      await expectAllTensorsReleased(async () => {
+        const netInput = await toNetInput([imgEl, imgEl, imgEl])
+        netInput.dispose()
+      })
+    })
+
+    it('single tf.Tensor3D', async () => {
+      const tensor = tf.fromPixels(imgEl)
+
+      await expectAllTensorsReleased(async () => {
+        const netInput = await toNetInput(tensor)
+        netInput.dispose()
+      })
+
+      tensor.dispose()
+    })
+
+    it('multiple tf.Tensor3Ds', async () => {
+      const tensors = [imgEl, imgEl, imgEl].map(el => tf.fromPixels(el))
+
+      await expectAllTensorsReleased(async () => {
+        const netInput = await toNetInput(tensors)
+        netInput.dispose()
+      })
+
+      tensors.forEach(t => t.dispose())
+    })
+
+    it('single batch size 1 tf.Tensor4Ds', async () => {
+      const tensor = tf.tidy(() => tf.fromPixels(imgEl).expandDims()) as tf.Tensor4D
+
+      await expectAllTensorsReleased(async () => {
+        const netInput = await toNetInput(tensor)
+        netInput.dispose()
+      })
+
+      tensor.dispose()
+    })
+
+    it('multiple batch size 1 tf.Tensor4Ds', async () => {
+      const tensors = [imgEl, imgEl, imgEl]
+        .map(el => tf.tidy(() => tf.fromPixels(el).expandDims())) as tf.Tensor4D[]
+
+      await expectAllTensorsReleased(async () => {
+        const netInput = await toNetInput(tensors)
+        netInput.dispose()
+      })
+
+      tensors.forEach(t => t.dispose())
+    })
+
+  })
+
 })

test/utils.ts

@@ -12,15 +12,6 @@ export function expectMaxDelta(val1: number, val2: number, maxDelta: number) {
   expect(Math.abs(val1 - val2)).toBeLessThan(maxDelta)
 }
 
-export async function createFakeHTMLVideoElement() {
-  const videoEl = document.createElement('video')
-  videoEl.muted = true
-  videoEl.src = 'base/test/media/video.mp4'
-  await videoEl.pause()
-  await videoEl.play()
-  return videoEl
-}
-
 export async function expectAllTensorsReleased(fn: () => any) {
   const numTensorsBefore = tf.memory().numTensors
   await fn()