prepare tinyYolov2 for using depthwise seperable conv2ds + retraining

src/BoundingBox.ts

 import { Rect } from './Rect';
+import { Dimensions } from './types';
+import { isDimensions } from './utils';
 
 export class BoundingBox {
   constructor(
@@ -33,6 +35,10 @@ export class BoundingBox {
     return this.bottom - this.top
   }
 
+  public get area() : number {
+    return this.width * this.height
+  }
+
   public toSquare(): BoundingBox {
     let { left, top, right, bottom } = this
 
@@ -100,6 +106,12 @@ export class BoundingBox {
     ).toSquare().round()
   }
 
+  public rescale(s: Dimensions | number) {
+    const scaleX = isDimensions(s) ? (s as Dimensions).width : s as number
+    const scaleY = isDimensions(s) ? (s as Dimensions).height : s as number
+    return new BoundingBox(this.left * scaleX, this.top * scaleY, this.right * scaleX, this.bottom * scaleY)
+  }
+
   public toRect(): Rect {
     return new Rect(this.left, this.top, this.width, this.height)
   }

src/commons/extractWeightsFactory.ts

@@ -2,6 +2,7 @@ export function extractWeightsFactory(weights: Float32Array) {
   let remainingWeights = weights
 
   function extractWeights(numWeights: number): Float32Array {
+    console.log(numWeights)
     const ret = remainingWeights.slice(0, numWeights)
     remainingWeights = remainingWeights.slice(numWeights)
     return ret

src/commons/nonMaxSuppression.ts

 import { BoundingBox } from '../BoundingBox';
+import { iou } from '../iou';
 
 export function nonMaxSuppression(
   boxes: BoundingBox[],
@@ -7,10 +8,6 @@ export function nonMaxSuppression(
   isIOU: boolean = true
 ): number[] {
 
-  const areas = boxes.map(
-    box => (box.width + 1) * (box.height + 1)
-  )
-
   let indicesSortedByScore = scores
     .map((score, boxIndex) => ({ score, boxIndex }))
     .sort((c1, c2) => c1.score - c2.score)
@@ -31,15 +28,7 @@ export function nonMaxSuppression(
       const currBox = boxes[curr]
       const idxBox = boxes[idx]
 
-      const width = Math.max(0.0, Math.min(currBox.right, idxBox.right) - Math.max(currBox.left, idxBox.left) + 1)
-      const height = Math.max(0.0, Math.min(currBox.bottom, idxBox.bottom) - Math.max(currBox.top, idxBox.top) + 1)
-      const interSection = width * height
-
-      const out = isIOU
-        ? interSection / (areas[curr] + areas[idx] - interSection)
-        : interSection / Math.min(areas[curr], areas[idx])
-
-      outputs.push(out)
+      outputs.push(iou(currBox, idxBox, isIOU))
     }
 
     indicesSortedByScore = indicesSortedByScore.filter(

src/index.ts

@@ -19,6 +19,7 @@ export * from './faceDetectionNet';
 export * from './faceLandmarkNet';
 export * from './faceRecognitionNet';
 export * from './globalApi';
+export * from './iou';
 export * from './mtcnn';
 export * from './padToSquare';
 export * from './tinyYolov2';

src/iou.ts

+import { BoundingBox } from './BoundingBox';
+
+export function iou(box1: BoundingBox, box2: BoundingBox, isIOU: boolean = true) {
+  const width = Math.max(0.0, Math.min(box1.right, box2.right) - Math.max(box1.left, box2.left) + 1)
+  const height = Math.max(0.0, Math.min(box1.bottom, box2.bottom) - Math.max(box1.top, box2.top) + 1)
+  const interSection = width * height
+
+  return isIOU
+    ? interSection / (box1.area + box2.area - interSection)
+    : interSection / Math.min(box1.area, box2.area)
+}
\ No newline at end of file

src/tinyYolov2/TinyYolov2.ts

@@ -6,19 +6,34 @@ import { NeuralNetwork } from '../commons/NeuralNetwork';
 import { nonMaxSuppression } from '../commons/nonMaxSuppression';
 import { FaceDetection } from '../FaceDetection';
 import { NetInput } from '../NetInput';
+import { Point } from '../Point';
 import { toNetInput } from '../toNetInput';
 import { TNetInput } from '../types';
-import { BOX_ANCHORS, INPUT_SIZES, IOU_THRESHOLD, NUM_BOXES } from './config';
+import { sigmoid } from '../utils';
+import { BOX_ANCHORS, BOX_ANCHORS_SEPARABLE, INPUT_SIZES, IOU_THRESHOLD, NUM_BOXES } from './config';
 import { convWithBatchNorm } from './convWithBatchNorm';
 import { extractParams } from './extractParams';
 import { getDefaultParams } from './getDefaultParams';
 import { loadQuantizedParams } from './loadQuantizedParams';
-import { NetParams, TinyYolov2ForwardParams } from './types';
+import { NetParams, PostProcessingParams, TinyYolov2ForwardParams } from './types';
 
 export class TinyYolov2 extends NeuralNetwork<NetParams> {
 
-  constructor() {
+  private _hasSeparableConvs: boolean
+  private _anchors: Point[]
+
+  constructor(hasSeparableConvs: boolean = false) {
     super('TinyYolov2')
+    this._hasSeparableConvs = hasSeparableConvs
+    this._anchors = hasSeparableConvs ? BOX_ANCHORS_SEPARABLE : BOX_ANCHORS
+  }
+
+  public get hasSeparableConvs(): boolean {
+    return this._hasSeparableConvs
+  }
+
+  public get anchors(): Point[] {
+    return this._anchors
   }
 
   public forwardInput(input: NetInput, inputSize: number): tf.Tensor4D {
@@ -30,7 +45,7 @@ export class TinyYolov2 extends NeuralNetwork<NetParams> {
     }
 
     const out = tf.tidy(() => {
-      const batchTensor = input.toBatchTensor(inputSize, false).div(tf.scalar(255)).toFloat() as tf.Tensor4D
+      const batchTensor = input.toBatchTensor(inputSize, false).div(tf.scalar(255)) as tf.Tensor4D
 
       let out = convWithBatchNorm(batchTensor, params.conv0)
       out = tf.maxPool(out, [2, 2], [2, 2], 'same')
@@ -72,39 +87,74 @@ export class TinyYolov2 extends NeuralNetwork<NetParams> {
 
     const netInput = await toNetInput(input, true)
     const out = await this.forwardInput(netInput, inputSize)
-    const numCells = out.shape[1]
+
+    const inputDimensions = {
+      width: netInput.getInputWidth(0),
+      height: netInput.getInputHeight(0)
+    }
+
+    const paddings = new Point(
+      (netInput.getPaddings(0).x + netInput.getInputWidth(0)) / netInput.getInputWidth(0),
+      (netInput.getPaddings(0).y + netInput.getInputHeight(0)) / netInput.getInputHeight(0)
+    )
+
+    const results = this.postProcess(out, { scoreThreshold, paddings })
+    const boxes = results.map(res => res.box)
+    const scores = results.map(res => res.score)
+
+    out.dispose()
+
+    const indices = nonMaxSuppression(
+      boxes.map(box => box.rescale(inputSize)),
+      scores,
+      IOU_THRESHOLD,
+      true
+    )
+
+    const detections = indices.map(idx =>
+      new FaceDetection(
+        scores[idx],
+        boxes[idx].toRect(),
+        inputDimensions
+      )
+    )
+
+    return detections
+  }
+
+  public postProcess(outputTensor: tf.Tensor4D, { scoreThreshold, paddings }: PostProcessingParams) {
+
+    const numCells = outputTensor.shape[1]
 
     const [boxesTensor, scoresTensor] = tf.tidy(() => {
-      const reshaped = out.reshape([numCells, numCells, NUM_BOXES, 6])
-      out.dispose()
+      const reshaped = outputTensor.reshape([numCells, numCells, NUM_BOXES, this.hasSeparableConvs ? 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])
       return [boxes, scores]
     })
 
-    const expit = (x: number): number => 1 / (1 + Math.exp(-x))
-
-    const paddedHeightRelative = (netInput.getPaddings(0).y + netInput.getInputHeight(0)) / netInput.getInputHeight(0)
-    const paddedWidthRelative = (netInput.getPaddings(0).x + netInput.getInputWidth(0)) / netInput.getInputWidth(0)
-
-    const boxes: BoundingBox[] = []
-    const scores: number[] = []
+    const results = []
 
     for (let row = 0; row < numCells; row ++) {
       for (let col = 0; col < numCells; col ++) {
-        for (let box = 0; box < NUM_BOXES; box ++) {
-          const score = expit(scoresTensor.get(row, col, box, 0))
+        for (let anchor = 0; anchor < NUM_BOXES; anchor ++) {
+          const score = sigmoid(scoresTensor.get(row, col, anchor, 0))
           if (score > scoreThreshold) {
-            const ctX = ((col + expit(boxesTensor.get(row, col, box, 0))) / numCells) * paddedWidthRelative
-            const ctY = ((row + expit(boxesTensor.get(row, col, box, 1))) / numCells) * paddedHeightRelative
-            const width = ((Math.exp(boxesTensor.get(row, col, box, 2)) * BOX_ANCHORS[box].x) / numCells) * paddedWidthRelative
-            const height = ((Math.exp(boxesTensor.get(row, col, box, 3)) * BOX_ANCHORS[box].y) / numCells) * paddedHeightRelative
+            const ctX = ((col + sigmoid(boxesTensor.get(row, col, anchor, 0))) / numCells) * paddings.x
+            const ctY = ((row + sigmoid(boxesTensor.get(row, col, anchor, 1))) / numCells) * paddings.y
+            const width = ((Math.exp(boxesTensor.get(row, col, anchor, 2)) * this.anchors[anchor].x) / numCells) * paddings.x
+            const height = ((Math.exp(boxesTensor.get(row, col, anchor, 3)) * this.anchors[anchor].y) / numCells) * paddings.y
 
             const x = (ctX - (width / 2))
             const y = (ctY - (height / 2))
-            boxes.push(new BoundingBox(x, y, x + width, y + height))
-            scores.push(score)
+            results.push({
+              box: new BoundingBox(x, y, x + width, y + height),
+              score,
+              row,
+              col,
+              anchor
+            })
           }
         }
       }
@@ -113,27 +163,7 @@ export class TinyYolov2 extends NeuralNetwork<NetParams> {
     boxesTensor.dispose()
     scoresTensor.dispose()
 
-    const indices = nonMaxSuppression(
-      boxes.map(box => new BoundingBox(
-        box.left * inputSize,
-        box.top * inputSize,
-        box.right * inputSize,
-        box.bottom * inputSize
-      )),
-      scores,
-      IOU_THRESHOLD,
-      true
-    )
-
-    const detections = indices.map(idx =>
-      new FaceDetection(
-        scores[idx],
-        boxes[idx].toRect(),
-        { width: netInput.getInputWidth(0), height: netInput.getInputHeight(0) }
-      )
-    )
-
-    return detections
+    return results
   }
 
   protected loadQuantizedParams(uri: string | undefined) {
@@ -141,6 +171,6 @@ export class TinyYolov2 extends NeuralNetwork<NetParams> {
   }
 
   protected extractParams(weights: Float32Array) {
-    return extractParams(weights)
+    return extractParams(weights, this.hasSeparableConvs)
   }
 }
\ No newline at end of file

src/tinyYolov2/config.ts

@@ -10,4 +10,12 @@ export const BOX_ANCHORS = [
   new Point(4.30971, 7.04493),
   new Point(10.246, 4.59428),
   new Point(12.6868, 11.8741)
+]
+
+export const BOX_ANCHORS_SEPARABLE = [
+  new Point(1.603231, 2.094468),
+  new Point(6.041143, 7.080126),
+  new Point(2.882459, 3.518061),
+  new Point(4.266906, 5.178857),
+  new Point(9.041765, 10.66308)
 ]
\ No newline at end of file

src/tinyYolov2/convWithBatchNorm.ts

 import * as tf from '@tensorflow/tfjs-core';
 
 import { leaky } from './leaky';
-import { ConvWithBatchNorm } from './types';
+import { ConvWithBatchNorm, SeparableConvParams } from './types';
 
-export function convWithBatchNorm(x: tf.Tensor4D, params: ConvWithBatchNorm): tf.Tensor4D  {
+export function convWithBatchNorm(x: tf.Tensor4D, params: ConvWithBatchNorm | SeparableConvParams): tf.Tensor4D  {
   return tf.tidy(() => {
     let out = tf.pad(x, [[0, 0], [1, 1], [1, 1], [0, 0]]) as tf.Tensor4D
-    out = tf.conv2d(out, params.conv.filters, [1, 1], 'valid')
-    out = tf.sub(out, params.bn.sub)
-    out = tf.mul(out, params.bn.truediv)
-    out = tf.add(out, params.conv.bias)
+
+    if (params instanceof SeparableConvParams) {
+      out = tf.separableConv2d(out, params.depthwise_filter, params.pointwise_filter, [1, 1], 'same')
+      out = tf.add(out, params.bias)
+    } else {
+      out = tf.conv2d(out, params.conv.filters, [1, 1], 'valid')
+      out = tf.sub(out, params.bn.sub)
+      out = tf.mul(out, params.bn.truediv)
+      out = tf.add(out, params.conv.bias)
+    }
+
     return leaky(out)
   })
 }
\ No newline at end of file

src/tinyYolov2/extractParams.ts

@@ -3,7 +3,7 @@ import * as tf from '@tensorflow/tfjs-core';
 import { extractConvParamsFactory } from '../commons/extractConvParamsFactory';
 import { extractWeightsFactory } from '../commons/extractWeightsFactory';
 import { ExtractWeightsFunction, ParamMapping } from '../commons/types';
-import { BatchNorm, ConvWithBatchNorm, NetParams } from './types';
+import { BatchNorm, ConvWithBatchNorm, NetParams, SeparableConvParams } from './types';
 
 function extractorsFactory(extractWeights: ExtractWeightsFunction, paramMappings: ParamMapping[]) {
 
@@ -30,14 +30,34 @@ function extractorsFactory(extractWeights: ExtractWeightsFunction, paramMappings
     return { conv, bn }
   }
 
+  function extractSeparableConvParams(channelsIn: number, channelsOut: number, mappedPrefix: string): SeparableConvParams {
+    console.log(mappedPrefix)
+    const depthwise_filter = tf.tensor4d(extractWeights(3 * 3 * channelsIn), [3, 3, channelsIn, 1])
+    const pointwise_filter = tf.tensor4d(extractWeights(channelsIn * channelsOut), [1, 1, channelsIn, channelsOut])
+    const bias = tf.tensor1d(extractWeights(channelsOut))
+    console.log('done')
+    paramMappings.push(
+      { paramPath: `${mappedPrefix}/depthwise_filter` },
+      { paramPath: `${mappedPrefix}/pointwise_filter` },
+      { paramPath: `${mappedPrefix}/bias` }
+    )
+
+    return new SeparableConvParams(
+      depthwise_filter,
+      pointwise_filter,
+      bias
+    )
+  }
+
   return {
     extractConvParams,
-    extractConvWithBatchNormParams
+    extractConvWithBatchNormParams,
+    extractSeparableConvParams
   }
 
 }
 
-export function extractParams(weights: Float32Array): { params: NetParams, paramMappings: ParamMapping[] } {
+export function extractParams(weights: Float32Array, hasSeparableConvs: boolean): { params: NetParams, paramMappings: ParamMapping[] } {
 
   const {
     extractWeights,
@@ -48,18 +68,22 @@ export function extractParams(weights: Float32Array): { params: NetParams, param
 
   const {
     extractConvParams,
-    extractConvWithBatchNormParams
+    extractConvWithBatchNormParams,
+    extractSeparableConvParams
   } = extractorsFactory(extractWeights, paramMappings)
 
-  const conv0 = extractConvWithBatchNormParams(3, 16, 'conv0')
-  const conv1 = extractConvWithBatchNormParams(16, 32, 'conv1')
-  const conv2 = extractConvWithBatchNormParams(32, 64, 'conv2')
-  const conv3 = extractConvWithBatchNormParams(64, 128, 'conv3')
-  const conv4 = extractConvWithBatchNormParams(128, 256, 'conv4')
-  const conv5 = extractConvWithBatchNormParams(256, 512, 'conv5')
-  const conv6 = extractConvWithBatchNormParams(512, 1024, 'conv6')
-  const conv7 = extractConvWithBatchNormParams(1024, 1024, 'conv7')
-  const conv8 = extractConvParams(1024, 30, 1, 'conv8')
+  const extractConvFn = hasSeparableConvs ? extractSeparableConvParams : extractConvWithBatchNormParams
+  const numAnchorEncodings = hasSeparableConvs ? 5 : 6
+
+  const conv0 = extractConvFn(3, 16, 'conv0',)
+  const conv1 = extractConvFn(16, 32, 'conv1')
+  const conv2 = extractConvFn(32, 64, 'conv2')
+  const conv3 = extractConvFn(64, 128, 'conv3')
+  const conv4 = extractConvFn(128, 256, 'conv4')
+  const conv5 = extractConvFn(256, 512, 'conv5')
+  const conv6 = extractConvFn(512, 1024, 'conv6')
+  const conv7 = extractConvFn(1024, 1024, 'conv7')
+  const conv8 = extractConvParams(1024, 5 * numAnchorEncodings, 1, 'conv8')
 
   if (getRemainingWeights().length !== 0) {
     throw new Error(`weights remaing after extract: ${getRemainingWeights().length}`)

src/tinyYolov2/types.ts

 import * as tf from '@tensorflow/tfjs-core';
 
 import { ConvParams } from '../commons/types';
+import { Point } from '../Point';
 
 export type BatchNorm = {
   sub: tf.Tensor1D
   truediv: tf.Tensor1D
 }
 
+export class SeparableConvParams {
+  constructor(
+    public depthwise_filter: tf.Tensor4D,
+    public pointwise_filter: tf.Tensor4D,
+    public bias: tf.Tensor1D
+  ) {}
+}
+
 export type ConvWithBatchNorm = {
   conv: ConvParams
   bn: BatchNorm
 }
 
 export type NetParams = {
-  conv0: ConvWithBatchNorm
-  conv1: ConvWithBatchNorm
-  conv2: ConvWithBatchNorm
-  conv3: ConvWithBatchNorm
-  conv4: ConvWithBatchNorm
-  conv5: ConvWithBatchNorm
-  conv6: ConvWithBatchNorm
-  conv7: ConvWithBatchNorm
+  conv0: ConvWithBatchNorm | SeparableConvParams
+  conv1: ConvWithBatchNorm | SeparableConvParams
+  conv2: ConvWithBatchNorm | SeparableConvParams
+  conv3: ConvWithBatchNorm | SeparableConvParams
+  conv4: ConvWithBatchNorm | SeparableConvParams
+  conv5: ConvWithBatchNorm | SeparableConvParams
+  conv6: ConvWithBatchNorm | SeparableConvParams
+  conv7: ConvWithBatchNorm | SeparableConvParams
   conv8: ConvParams
 }
 
@@ -34,4 +43,9 @@ export enum SizeType {
 export type TinyYolov2ForwardParams = {
   inputSize?: SizeType | number
   scoreThreshold?: number
+}
+
+export type PostProcessingParams = {
+  scoreThreshold: number
+  paddings: Point
 }
\ No newline at end of file

src/utils.ts

@@ -15,6 +15,14 @@ export function round(num: number) {
   return Math.floor(num * 100) / 100
 }
 
+export function sigmoid(x: number) {
+  return 1 / (1 + Math.exp(-x))
+}
+
+export function isDimensions(obj: any): boolean {
+  return obj && obj.width && obj.height
+}
+
 export function resolveInput(arg: string | any) {
   if (typeof arg === 'string') {
     return document.getElementById(arg)