test cases for AgeGenderNet + fixed memory leaks and batch inputs for AgeGenderNet

src/ageGenderNet/AgeGenderNet.ts

@@ -42,8 +42,10 @@ export class AgeGenderNet extends NeuralNetwork<NetParams> {
   }
 
   public forwardInput(input: NetInput | tf.Tensor4D): NetOutput {
-    const { age, gender } = this.runNet(input)
-    return tf.tidy(() => ({ age, gender: tf.softmax(gender) }))
+    return tf.tidy(() => {
+      const { age, gender } = this.runNet(input)
+      return { age, gender: tf.softmax(gender) }
+    })
   }
 
   public async forward(input: TNetInput): Promise<NetOutput> {
@@ -64,14 +66,18 @@ export class AgeGenderNet extends NeuralNetwork<NetParams> {
     const predictionsByBatch = await Promise.all(
       ageAndGenderTensors.map(async ({ ageTensor, genderTensor }) => {
         const age = (await ageTensor.data())[0]
-        const probMale = (await out.gender.data())[0]
+        const probMale = (await genderTensor.data())[0]
         const isMale = probMale > 0.5
         const gender = isMale ? Gender.MALE : Gender.FEMALE
         const genderProbability = isMale ? probMale : (1 - probMale)
 
+        ageTensor.dispose()
+        genderTensor.dispose()
         return { age, gender, genderProbability }
       })
     )
+    out.age.dispose()
+    out.gender.dispose()
 
     return netInput.isBatchInput
       ? predictionsByBatch

test/tests/ageGenderNet/ageGenderNet.test.ts

+import * as tf from '@tensorflow/tfjs-core';
+
+import { createCanvasFromMedia, NetInput, toNetInput } from '../../../src';
+import { AgeAndGenderPrediction } from '../../../src/ageGenderNet/types';
+import { loadImage } from '../../env';
+import { describeWithBackend, describeWithNets, expectAllTensorsReleased } from '../../utils';
+
+function expectResultsAngry(result: AgeAndGenderPrediction) {
+  expect(result.age).toBeGreaterThanOrEqual(38)
+  expect(result.age).toBeLessThanOrEqual(42)
+  expect(result.gender).toEqual('male')
+  expect(result.genderProbability).toBeGreaterThanOrEqual(0.9)
+}
+
+function expectResultsSurprised(result: AgeAndGenderPrediction) {
+  expect(result.age).toBeGreaterThanOrEqual(24)
+  expect(result.age).toBeLessThanOrEqual(28)
+  expect(result.gender).toEqual('female')
+  expect(result.genderProbability).toBeGreaterThanOrEqual(0.8)
+}
+
+describeWithBackend('ageGenderNet', () => {
+
+  let imgElAngry: HTMLImageElement
+  let imgElSurprised: HTMLImageElement
+
+  beforeAll(async () => {
+    imgElAngry = await loadImage('test/images/angry_cropped.jpg')
+    imgElSurprised = await loadImage('test/images/surprised_cropped.jpg')
+  })
+
+  describeWithNets('quantized weights', { withAgeGenderNet: { quantized: true } }, ({ ageGenderNet }) => {
+
+    it('recognizes age and gender', async () => {
+      const result = await ageGenderNet.predictAgeAndGender(imgElAngry) as AgeAndGenderPrediction
+      expectResultsAngry(result)
+    })
+
+  })
+
+  describeWithNets('batch inputs', { withAgeGenderNet: { quantized: true } }, ({ ageGenderNet }) => {
+
+    it('recognizes age and gender for batch of image elements', async () => {
+      const inputs = [imgElAngry, imgElSurprised]
+
+      const results = await ageGenderNet.predictAgeAndGender(inputs) as AgeAndGenderPrediction[]
+      expect(Array.isArray(results)).toBe(true)
+      expect(results.length).toEqual(2)
+
+      const [resultAngry, resultSurprised] = results
+      expectResultsAngry(resultAngry)
+      expectResultsSurprised(resultSurprised)
+    })
+
+    it('computes age and gender for batch of tf.Tensor3D', async () => {
+      const inputs = [imgElAngry, imgElSurprised].map(el => tf.browser.fromPixels(createCanvasFromMedia(el)))
+
+      const results = await ageGenderNet.predictAgeAndGender(inputs) as AgeAndGenderPrediction[]
+      expect(Array.isArray(results)).toBe(true)
+      expect(results.length).toEqual(2)
+
+      const [resultAngry, resultSurprised] = results
+      expectResultsAngry(resultAngry)
+      expectResultsSurprised(resultSurprised)
+    })
+
+    it('computes age and gender for batch of mixed inputs', async () => {
+      const inputs = [imgElAngry, tf.browser.fromPixels(createCanvasFromMedia(imgElSurprised))]
+
+      const results = await ageGenderNet.predictAgeAndGender(inputs) as AgeAndGenderPrediction[]
+      expect(Array.isArray(results)).toBe(true)
+      expect(results.length).toEqual(2)
+
+      const [resultAngry, resultSurprised] = results
+      expectResultsAngry(resultAngry)
+      expectResultsSurprised(resultSurprised)
+    })
+
+  })
+
+  describeWithNets('no memory leaks', { withAgeGenderNet: { quantized: true } }, ({ ageGenderNet }) => {
+
+    describe('forwardInput', () => {
+
+      it('single image element', async () => {
+        await expectAllTensorsReleased(async () => {
+          const netInput = new NetInput([imgElAngry])
+          const { age, gender } = await ageGenderNet.forwardInput(netInput)
+          age.dispose()
+          gender.dispose()
+        })
+      })
+
+      it('multiple image elements', async () => {
+        await expectAllTensorsReleased(async () => {
+          const netInput = new NetInput([imgElAngry, imgElAngry])
+          const { age, gender } = await ageGenderNet.forwardInput(netInput)
+          age.dispose()
+          gender.dispose()
+        })
+      })
+
+      it('single tf.Tensor3D', async () => {
+        const tensor = tf.browser.fromPixels(createCanvasFromMedia(imgElAngry))
+
+        await expectAllTensorsReleased(async () => {
+          const { age, gender } = await ageGenderNet.forwardInput(await toNetInput(tensor))
+          age.dispose()
+          gender.dispose()
+        })
+
+        tensor.dispose()
+      })
+
+      it('multiple tf.Tensor3Ds', async () => {
+        const tensors = [imgElAngry, imgElAngry, imgElAngry].map(el => tf.browser.fromPixels(createCanvasFromMedia(el)))
+
+        await expectAllTensorsReleased(async () => {
+          const { age, gender } = await ageGenderNet.forwardInput(await toNetInput(tensors))
+          age.dispose()
+          gender.dispose()
+        })
+
+        tensors.forEach(t => t.dispose())
+      })
+
+      it('single batch size 1 tf.Tensor4Ds', async () => {
+        const tensor = tf.tidy(() => tf.browser.fromPixels(createCanvasFromMedia(imgElAngry)).expandDims()) as tf.Tensor4D
+
+        await expectAllTensorsReleased(async () => {
+          const { age, gender } = await ageGenderNet.forwardInput(await toNetInput(tensor))
+          age.dispose()
+          gender.dispose()
+        })
+
+        tensor.dispose()
+      })
+
+      it('multiple batch size 1 tf.Tensor4Ds', async () => {
+        const tensors = [imgElAngry, imgElAngry, imgElAngry]
+          .map(el => tf.tidy(() => tf.browser.fromPixels(createCanvasFromMedia(el)).expandDims())) as tf.Tensor4D[]
+
+        await expectAllTensorsReleased(async () => {
+          const { age, gender } = await ageGenderNet.forwardInput(await toNetInput(tensors))
+          age.dispose()
+          gender.dispose()
+        })
+
+        tensors.forEach(t => t.dispose())
+      })
+
+    })
+
+    describe('predictExpressions', () => {
+
+      it('single image element', async () => {
+        await expectAllTensorsReleased(async () => {
+          await ageGenderNet.predictAgeAndGender(imgElAngry)
+        })
+      })
+
+      it('multiple image elements', async () => {
+        await expectAllTensorsReleased(async () => {
+          await ageGenderNet.predictAgeAndGender([imgElAngry, imgElAngry, imgElAngry])
+        })
+      })
+
+      it('single tf.Tensor3D', async () => {
+        const tensor = tf.browser.fromPixels(createCanvasFromMedia(imgElAngry))
+
+        await expectAllTensorsReleased(async () => {
+          await ageGenderNet.predictAgeAndGender(tensor)
+        })
+
+        tensor.dispose()
+      })
+
+      it('multiple tf.Tensor3Ds', async () => {
+        const tensors = [imgElAngry, imgElAngry, imgElAngry].map(el => tf.browser.fromPixels(createCanvasFromMedia(el)))
+
+
+        await expectAllTensorsReleased(async () => {
+          await ageGenderNet.predictAgeAndGender(tensors)
+        })
+
+        tensors.forEach(t => t.dispose())
+      })
+
+      it('single batch size 1 tf.Tensor4Ds', async () => {
+        const tensor = tf.tidy(() => tf.browser.fromPixels(createCanvasFromMedia(imgElAngry)).expandDims()) as tf.Tensor4D
+
+        await expectAllTensorsReleased(async () => {
+          await ageGenderNet.predictAgeAndGender(tensor)
+        })
+
+        tensor.dispose()
+      })
+
+      it('multiple batch size 1 tf.Tensor4Ds', async () => {
+        const tensors = [imgElAngry, imgElAngry, imgElAngry]
+          .map(el => tf.tidy(() => tf.browser.fromPixels(createCanvasFromMedia(el)).expandDims())) as tf.Tensor4D[]
+
+        await expectAllTensorsReleased(async () => {
+          await ageGenderNet.predictAgeAndGender(tensors)
+        })
+
+        tensors.forEach(t => t.dispose())
+      })
+
+    })
+  })
+
+})
+

test/tests/faceExpressionNet/faceExpressionNet.test.ts

@@ -41,7 +41,7 @@ describeWithBackend('faceExpressionNet', () => {
       expect(resultSurprised.surprised).toBeGreaterThan(0.95)
     })
 
-    it('computes face landmarks for batch of tf.Tensor3D', async () => {
+    it('computes face expressions for batch of tf.Tensor3D', async () => {
       const inputs = [imgElAngry, imgElSurprised].map(el => tf.browser.fromPixels(createCanvasFromMedia(el)))
 
       const results = await faceExpressionNet.predictExpressions(inputs) as FaceExpressions[]
@@ -55,7 +55,7 @@ describeWithBackend('faceExpressionNet', () => {
       expect(resultSurprised.surprised).toBeGreaterThan(0.95)
     })
 
-    it('computes face landmarks for batch of mixed inputs', async () => {
+    it('computes face expressions for batch of mixed inputs', async () => {
       const inputs = [imgElAngry, tf.browser.fromPixels(createCanvasFromMedia(imgElSurprised))]
 
       const results = await faceExpressionNet.predictExpressions(inputs) as FaceExpressions[]

test/utils.ts

@@ -2,6 +2,7 @@ import * as tf from '@tensorflow/tfjs-core';
 
 import * as faceapi from '../src';
 import { FaceRecognitionNet, IPoint, IRect, Mtcnn, TinyYolov2 } from '../src/';
+import { AgeGenderNet } from '../src/ageGenderNet/AgeGenderNet';
 import { FaceDetection } from '../src/classes/FaceDetection';
 import { FaceLandmarks } from '../src/classes/FaceLandmarks';
 import { FaceExpressionNet } from '../src/faceExpressionNet/FaceExpressionNet';
@@ -114,6 +115,7 @@ export type InjectNetArgs = {
   faceRecognitionNet: FaceRecognitionNet
   mtcnn: Mtcnn
   faceExpressionNet: FaceExpressionNet
+  ageGenderNet: AgeGenderNet
   tinyYolov2: TinyYolov2
 }
 
@@ -129,6 +131,7 @@ export type DescribeWithNetsOptions = {
   withFaceRecognitionNet?: WithNetOptions
   withMtcnn?: WithNetOptions
   withFaceExpressionNet?: WithNetOptions
+  withAgeGenderNet?: WithNetOptions
   withTinyYolov2?: WithTinyYolov2Options
 }
 
@@ -176,6 +179,7 @@ export function describeWithNets(
       faceRecognitionNet,
       mtcnn,
       faceExpressionNet,
+      ageGenderNet,
       tinyYolov2
     } = faceapi.nets
 
@@ -192,6 +196,7 @@ export function describeWithNets(
         withFaceRecognitionNet,
         withMtcnn,
         withFaceExpressionNet,
+        withAgeGenderNet,
         withTinyYolov2
       } = options
 
@@ -244,6 +249,13 @@ export function describeWithNets(
         )
       }
 
+      if (withAgeGenderNet) {
+        await initNet<AgeGenderNet>(
+          ageGenderNet,
+          !!withAgeGenderNet && !withAgeGenderNet.quantized && 'age_gender_model.weights'
+        )
+      }
+
       if (withTinyYolov2 || withAllFacesTinyYolov2) {
         await initNet<TinyYolov2>(
           tinyYolov2,
@@ -273,6 +285,7 @@ export function describeWithNets(
       faceRecognitionNet,
       mtcnn,
       faceExpressionNet,
+      ageGenderNet,
       tinyYolov2
     })
   })