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/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)