weight loading of quantized tiny-yolov2-seperable-conv2d + use separable conv…

src/tinyYolov2/TinyYolov2.ts

@@ -20,17 +20,17 @@ import { NetParams, PostProcessingParams, TinyYolov2ForwardParams } from './type
 
 export class TinyYolov2 extends NeuralNetwork<NetParams> {
 
-  private _hasSeparableConvs: boolean
+  private _withSeparableConvs: boolean
   private _anchors: Point[]
 
-  constructor(hasSeparableConvs: boolean = false) {
+  constructor(withSeparableConvs: boolean = true) {
     super('TinyYolov2')
-    this._hasSeparableConvs = hasSeparableConvs
-    this._anchors = hasSeparableConvs ? BOX_ANCHORS_SEPARABLE : BOX_ANCHORS
+    this._withSeparableConvs = withSeparableConvs
+    this._anchors = withSeparableConvs ? BOX_ANCHORS_SEPARABLE : BOX_ANCHORS
   }
 
-  public get hasSeparableConvs(): boolean {
-    return this._hasSeparableConvs
+  public get withSeparableConvs(): boolean {
+    return this._withSeparableConvs
   }
 
   public get anchors(): Point[] {
@@ -48,7 +48,7 @@ export class TinyYolov2 extends NeuralNetwork<NetParams> {
     const out = tf.tidy(() => {
 
       let batchTensor = input.toBatchTensor(inputSize, false)
-      batchTensor = this.hasSeparableConvs
+      batchTensor = this.withSeparableConvs
         ? normalize(batchTensor, MEAN_RGB)
         : batchTensor
       batchTensor = batchTensor.div(tf.scalar(256)) as tf.Tensor4D
@@ -132,7 +132,7 @@ export class TinyYolov2 extends NeuralNetwork<NetParams> {
     const numCells = outputTensor.shape[1]
 
     const [boxesTensor, scoresTensor] = tf.tidy(() => {
-      const reshaped = outputTensor.reshape([numCells, numCells, NUM_BOXES, this.hasSeparableConvs ? 5 : 6])
+      const reshaped = outputTensor.reshape([numCells, numCells, NUM_BOXES, this.withSeparableConvs ? 5 : 6])
 
       const boxes = reshaped.slice([0, 0, 0, 0], [numCells, numCells, NUM_BOXES, 4])
       const scores = reshaped.slice([0, 0, 0, 4], [numCells, numCells, NUM_BOXES, 1])
@@ -172,10 +172,10 @@ export class TinyYolov2 extends NeuralNetwork<NetParams> {
   }
 
   protected loadQuantizedParams(uri: string | undefined) {
-    return loadQuantizedParams(uri)
+    return loadQuantizedParams(uri, this.withSeparableConvs)
   }
 
   protected extractParams(weights: Float32Array) {
-    return extractParams(weights, this.hasSeparableConvs)
+    return extractParams(weights, this.withSeparableConvs)
   }
 }
\ No newline at end of file

src/tinyYolov2/extractParams.ts

@@ -56,7 +56,7 @@ function extractorsFactory(extractWeights: ExtractWeightsFunction, paramMappings
 
 }
 
-export function extractParams(weights: Float32Array, hasSeparableConvs: boolean): { params: NetParams, paramMappings: ParamMapping[] } {
+export function extractParams(weights: Float32Array, withSeparableConvs: boolean): { params: NetParams, paramMappings: ParamMapping[] } {
 
   const {
     extractWeights,
@@ -71,8 +71,8 @@ export function extractParams(weights: Float32Array, hasSeparableConvs: boolean)
     extractSeparableConvParams
   } = extractorsFactory(extractWeights, paramMappings)
 
-  const extractConvFn = hasSeparableConvs ? extractSeparableConvParams : extractConvWithBatchNormParams
-  const numAnchorEncodings = hasSeparableConvs ? 5 : 6
+  const extractConvFn = withSeparableConvs ? extractSeparableConvParams : extractConvWithBatchNormParams
+  const numAnchorEncodings = withSeparableConvs ? 5 : 6
 
   const conv0 = extractConvFn(3, 16, 'conv0',)
   const conv1 = extractConvFn(16, 32, 'conv1')

src/tinyYolov2/index.ts

@@ -2,8 +2,8 @@ import { TinyYolov2 } from './TinyYolov2';
 
 export * from './TinyYolov2';
 
-export function createTinyYolov2(weights: Float32Array) {
-  const net = new TinyYolov2()
+export function createTinyYolov2(weights: Float32Array, withSeparableConvs: boolean = true) {
+  const net = new TinyYolov2(withSeparableConvs)
   net.extractWeights(weights)
   return net
 }
\ No newline at end of file

src/tinyYolov2/loadQuantizedParams.ts

@@ -4,9 +4,10 @@ import { disposeUnusedWeightTensors } from '../commons/disposeUnusedWeightTensor
 import { extractWeightEntryFactory } from '../commons/extractWeightEntryFactory';
 import { loadWeightMap } from '../commons/loadWeightMap';
 import { ConvParams, ParamMapping } from '../commons/types';
-import { BatchNorm, ConvWithBatchNorm, NetParams } from './types';
+import { BatchNorm, ConvWithBatchNorm, NetParams, SeparableConvParams } from './types';
 
 const DEFAULT_MODEL_NAME = 'tiny_yolov2_model'
+const DEFAULT_MODEL_NAME_SEPARABLE_CONV = 'tiny_yolov2_separable_conv_model'
 
 function extractorsFactory(weightMap: any, paramMappings: ParamMapping[]) {
 
@@ -30,35 +31,51 @@ function extractorsFactory(weightMap: any, paramMappings: ParamMapping[]) {
     return { conv, bn }
   }
 
+  function extractSeparableConvParams(prefix: string): SeparableConvParams {
+    const depthwise_filter = extractWeightEntry<tf.Tensor4D>(`${prefix}/depthwise_filter`, 4)
+    const pointwise_filter = extractWeightEntry<tf.Tensor4D>(`${prefix}/pointwise_filter`, 4)
+    const bias = extractWeightEntry<tf.Tensor1D>(`${prefix}/bias`, 1)
+
+    return new SeparableConvParams(
+      depthwise_filter,
+      pointwise_filter,
+      bias
+    )
+  }
+
   return {
     extractConvParams,
-    extractConvWithBatchNormParams
+    extractConvWithBatchNormParams,
+    extractSeparableConvParams
   }
 
 }
 
 export async function loadQuantizedParams(
-  uri: string | undefined
+  uri: string | undefined,
+  withSeparableConvs: boolean
 ): Promise<{ params: NetParams, paramMappings: ParamMapping[] }> {
 
-  const weightMap = await loadWeightMap(uri, DEFAULT_MODEL_NAME)
+  const weightMap = await loadWeightMap(uri,  withSeparableConvs ? DEFAULT_MODEL_NAME_SEPARABLE_CONV : DEFAULT_MODEL_NAME)
   const paramMappings: ParamMapping[] = []
 
   const {
     extractConvParams,
-    extractConvWithBatchNormParams
+    extractConvWithBatchNormParams,
+    extractSeparableConvParams
   } = extractorsFactory(weightMap, paramMappings)
 
+  const extractConvFn = withSeparableConvs ? extractSeparableConvParams : extractConvWithBatchNormParams
 
   const params = {
-    conv0: extractConvWithBatchNormParams('conv0'),
-    conv1: extractConvWithBatchNormParams('conv1'),
-    conv2: extractConvWithBatchNormParams('conv2'),
-    conv3: extractConvWithBatchNormParams('conv3'),
-    conv4: extractConvWithBatchNormParams('conv4'),
-    conv5: extractConvWithBatchNormParams('conv5'),
-    conv6: extractConvWithBatchNormParams('conv6'),
-    conv7: extractConvWithBatchNormParams('conv7'),
+    conv0: extractConvFn('conv0'),
+    conv1: extractConvFn('conv1'),
+    conv2: extractConvFn('conv2'),
+    conv3: extractConvFn('conv3'),
+    conv4: extractConvFn('conv4'),
+    conv5: extractConvFn('conv5'),
+    conv6: extractConvFn('conv6'),
+    conv7: extractConvFn('conv7'),
     conv8: extractConvParams('conv8')
   }
 

test/tests/e2e/expectedResults.ts

@@ -29,6 +29,15 @@ export const expectedTinyYolov2Boxes = [
   { x: 87, y: 30, width: 92, height: 93 }
 ]
 
+export const expectedTinyYolov2SeparableConvBoxes = [
+  { x: 42, y: 257, width: 111, height: 121 },
+  { x: 454, y: 175, width: 104, height: 121 },
+  { x: 230, y: 45, width: 94, height: 104 },
+  { x: 574, y: 62, width: 88, height: 113 },
+  { x: 260, y: 233, width: 82, height: 104 },
+  { x: 83, y: 24, width: 85, height: 111 }
+]
+
 export const expectedMtcnnFaceLandmarks = [
   [new Point(117, 58), new Point(156, 63), new Point(141, 86), new Point(109, 98), new Point(147, 104)],
   [new Point(82, 292), new Point(134, 304), new Point(104, 330), new Point(72, 342), new Point(120, 353)],
@@ -38,7 +47,6 @@ export const expectedMtcnnFaceLandmarks = [
   [new Point(489, 224), new Point(534, 223), new Point(507, 250), new Point(493, 271), new Point(530, 270)]
 ]
 
-
 export function expectMtcnnResults(
   results: { faceDetection: faceapi.FaceDetection, faceLandmarks: faceapi.FaceLandmarks5 }[],
   boxOrder: number[],

test/tests/e2e/tinyYolov2.test.ts

 import * as faceapi from '../../../src';
 import { SizeType } from '../../../src/tinyYolov2/types';
 import { describeWithNets, expectAllTensorsReleased, expectRectClose } from '../../utils';
-import { expectedTinyYolov2Boxes } from './expectedResults';
-
+import { expectedTinyYolov2Boxes, expectedTinyYolov2SeparableConvBoxes } from './expectedResults';
 
 describe('tinyYolov2', () => {
 
@@ -13,11 +12,63 @@ describe('tinyYolov2', () => {
     imgEl = await faceapi.bufferToImage(img)
   })
 
+  describe('with separable convolutions', () => {
+
     describeWithNets('quantized weights', { withTinyYolov2: { quantized: true } }, ({ tinyYolov2 }) => {
 
       it('inputSize lg, finds all faces', async () => {
         const detections = await tinyYolov2.locateFaces(imgEl, { inputSize: SizeType.LG })
 
+        const expectedScores = [0.9, 0.9, 0.89, 0.85, 0.85, 0.85]
+        const maxBoxDelta = 1
+        const boxOrder = [0, 1, 2, 3, 4, 5]
+
+        expect(detections.length).toEqual(6)
+        detections.forEach((det, i) => {
+          expect(det.getScore()).toBeCloseTo(expectedScores[i], 2)
+          expectRectClose(det.getBox(), expectedTinyYolov2SeparableConvBoxes[boxOrder[i]], maxBoxDelta)
+        })
+      })
+
+      it('inputSize md, finds all faces', async () => {
+        const detections = await tinyYolov2.locateFaces(imgEl, { inputSize: SizeType.MD })
+
+        const expectedScores = [0.85, 0.85, 0.84, 0.83, 0.8, 0.8]
+        const maxBoxDelta = 17
+        const boxOrder = [5, 1, 4, 3, 2, 0]
+
+        expect(detections.length).toEqual(6)
+        detections.forEach((det, i) => {
+          expect(det.getScore()).toBeCloseTo(expectedScores[i], 2)
+          expectRectClose(det.getBox(), expectedTinyYolov2SeparableConvBoxes[boxOrder[i]], maxBoxDelta)
+        })
+      })
+
+      it('inputSize custom, finds all faces', async () => {
+        const detections = await tinyYolov2.locateFaces(imgEl, { inputSize: 416 })
+
+        const expectedScores = [0.85, 0.85, 0.84, 0.83, 0.8, 0.8]
+        const maxBoxDelta = 17
+        const boxOrder = [5, 1, 4, 3, 2, 0]
+
+        expect(detections.length).toEqual(6)
+        detections.forEach((det, i) => {
+          expect(det.getScore()).toBeCloseTo(expectedScores[i], 2)
+          expectRectClose(det.getBox(), expectedTinyYolov2SeparableConvBoxes[boxOrder[i]], maxBoxDelta)
+        })
+      })
+
+    })
+
+  })
+
+  describe('without separable convolutions', () => {
+
+    describeWithNets('quantized weights', { withTinyYolov2: { quantized: true, withSeparableConv: false } }, ({ tinyYolov2 }) => {
+
+      it('inputSize lg, finds all faces', async () => {
+        const detections = await tinyYolov2.locateFaces(imgEl, { inputSize: SizeType.LG })
+
         const expectedScores = [0.86, 0.86, 0.85, 0.83, 0.81, 0.81]
         const maxBoxDelta = 3
         const boxOrder = [0, 1, 2, 3, 4, 5]
@@ -59,7 +110,7 @@ describe('tinyYolov2', () => {
 
     })
 
-  describeWithNets('uncompressed weights', { withTinyYolov2: { quantized: false } }, ({ tinyYolov2 }) => {
+    describeWithNets('uncompressed weights', { withTinyYolov2: { quantized: false, withSeparableConv: false } }, ({ tinyYolov2 }) => {
 
       it('inputSize lg, finds all faces', async () => {
         const detections = await tinyYolov2.locateFaces(imgEl, { inputSize: SizeType.LG })
@@ -105,13 +156,17 @@ describe('tinyYolov2', () => {
 
     })
 
+  })
+
   describe('no memory leaks', () => {
 
+    describe('with separable convolutions', () => {
+
       describe('NeuralNetwork, uncompressed model', () => {
 
         it('disposes all param tensors', async () => {
           await expectAllTensorsReleased(async () => {
-          const res = await fetch('base/weights_uncompressed/tiny_yolov2_model.weights')
+            const res = await fetch('base/weights_uncompressed/tiny_yolov2_separable_conv_model.weights')
             const weights = new Float32Array(await res.arrayBuffer())
             const net = faceapi.createTinyYolov2(weights)
             net.dispose()
@@ -134,4 +189,35 @@ describe('tinyYolov2', () => {
 
     })
 
+    describe('without separable convolutions', () => {
+
+      describe('NeuralNetwork, uncompressed model', () => {
+
+        it('disposes all param tensors', async () => {
+          await expectAllTensorsReleased(async () => {
+            const res = await fetch('base/weights_uncompressed/tiny_yolov2_model.weights')
+            const weights = new Float32Array(await res.arrayBuffer())
+            const net = faceapi.createTinyYolov2(weights, false)
+            net.dispose()
+          })
+        })
+
+      })
+
+      describe('NeuralNetwork, quantized model', () => {
+
+        it('disposes all param tensors', async () => {
+          await expectAllTensorsReleased(async () => {
+            const net = new faceapi.TinyYolov2(false)
+            await net.load('base/weights')
+            net.dispose()
+          })
+        })
+
+      })
+
+    })
+
+  })
+
 })
\ No newline at end of file

test/utils.ts

@@ -5,8 +5,7 @@ import * as faceapi from '../src/';
 import { NeuralNetwork } from '../src/commons/NeuralNetwork';
 import { IPoint } from '../src/';
 import { allFacesFactory, allFacesMtcnnFactory } from '../src/allFacesFactory';
-import { allFacesMtcnnFunction, allFacesFunction, tinyYolov2 } from '../src/globalApi';
-import { TinyYolov2 } from '../src/tinyYolov2/TinyYolov2';
+import { allFacesMtcnnFunction, allFacesFunction } from '../src/globalApi';
 
 export function zeros(length: number): Float32Array {
   return new Float32Array(length)
@@ -55,6 +54,10 @@ export type WithNetOptions = {
   quantized?: boolean
 }
 
+export type WithTinyYolov2Options = WithNetOptions & {
+  withSeparableConv?: boolean
+}
+
 export type InjectNetArgs = {
   allFaces: allFacesFunction
   allFacesMtcnn: allFacesMtcnnFunction
@@ -73,7 +76,7 @@ export type DescribeWithNetsOptions = {
   withFaceLandmarkNet?: WithNetOptions
   withFaceRecognitionNet?: WithNetOptions
   withMtcnn?: WithNetOptions
-  withTinyYolov2?: WithNetOptions
+  withTinyYolov2?: WithTinyYolov2Options
 }
 
 async function loadNetWeights(uri: string): Promise<Float32Array> {
@@ -102,7 +105,7 @@ export function describeWithNets(
     let faceLandmarkNet: faceapi.FaceLandmarkNet = new faceapi.FaceLandmarkNet()
     let faceRecognitionNet: faceapi.FaceRecognitionNet = new faceapi.FaceRecognitionNet()
     let mtcnn: faceapi.Mtcnn = new faceapi.Mtcnn()
-    let tinyYolov2: faceapi.TinyYolov2 = new faceapi.TinyYolov2()
+    let tinyYolov2: faceapi.TinyYolov2 = new faceapi.TinyYolov2(options.withTinyYolov2 && options.withTinyYolov2.withSeparableConv)
     let allFaces = allFacesFactory(faceDetectionNet, faceLandmarkNet, faceRecognitionNet)
     let allFacesMtcnn = allFacesMtcnnFactory(mtcnn, faceRecognitionNet)