postprocessing for tiny-yolov2

src/mtcnn/BoundingBox.ts → src/BoundingBox.ts

+import { Rect } from './Rect';
+
 export class BoundingBox {
   constructor(
     private _left: number,
@@ -31,7 +33,6 @@ export class BoundingBox {
     return this.bottom - this.top
   }
 
-
   public toSquare(): BoundingBox {
     let { left, top, right, bottom } = this
 
@@ -98,4 +99,8 @@ export class BoundingBox {
       this.bottom + (region.bottom * this.height)
     ).toSquare().round()
   }
+
+  public toRect(): Rect {
+    return new Rect(this.left, this.top, this.width, this.height)
+  }
 }
\ No newline at end of file

src/mtcnn/nms.ts → src/commons/nonMaxSuppression.ts

-import { BoundingBox } from './BoundingBox';
+import { BoundingBox } from '../BoundingBox';
 
-export function nms(
+export function nonMaxSuppression(
   boxes: BoundingBox[],
   scores: number[],
   iouThreshold: number,

src/mtcnn/extractImagePatches.ts

 import * as tf from '@tensorflow/tfjs-core';
 
+import { BoundingBox } from '../BoundingBox';
 import { Dimensions } from '../types';
 import { createCanvas, getContext2dOrThrow } from '../utils';
-import { BoundingBox } from './BoundingBox';
 import { normalize } from './normalize';
 
 export async function extractImagePatches(

src/mtcnn/stage1.ts

 import * as tf from '@tensorflow/tfjs-core';
 
+import { BoundingBox } from '../BoundingBox';
+import { nonMaxSuppression } from '../commons/nonMaxSuppression';
 import { Point } from '../Point';
-import { BoundingBox } from './BoundingBox';
 import { CELL_SIZE, CELL_STRIDE } from './config';
-import { nms } from './nms';
+import { getSizesForScale } from './getSizesForScale';
 import { normalize } from './normalize';
 import { PNet } from './PNet';
 import { PNetParams } from './types';
-import { getSizesForScale } from './getSizesForScale';
 
 function rescaleAndNormalize(x: tf.Tensor4D, scale: number): tf.Tensor4D {
   return tf.tidy(() => {
@@ -109,7 +109,7 @@ export function stage1(
     }
 
     let ts = Date.now()
-    const indices = nms(
+    const indices = nonMaxSuppression(
       boundingBoxes.map(bbox => bbox.cell),
       boundingBoxes.map(bbox => bbox.score),
       0.5
@@ -130,7 +130,7 @@ export function stage1(
 
   if (allBoxes.length > 0) {
     let ts = Date.now()
-    const indices = nms(
+    const indices = nonMaxSuppression(
       allBoxes.map(bbox => bbox.cell),
       allBoxes.map(bbox => bbox.score),
       0.7

src/mtcnn/stage2.ts

 import * as tf from '@tensorflow/tfjs-core';
 
-import { BoundingBox } from './BoundingBox';
+import { BoundingBox } from '../BoundingBox';
+import { nonMaxSuppression } from '../commons/nonMaxSuppression';
 import { extractImagePatches } from './extractImagePatches';
-import { nms } from './nms';
 import { RNet } from './RNet';
 import { RNetParams } from './types';
 
@@ -47,7 +47,7 @@ export async function stage2(
 
   if (filteredBoxes.length > 0) {
     ts = Date.now()
-    const indicesNms = nms(
+    const indicesNms = nonMaxSuppression(
       filteredBoxes,
       filteredScores,
       0.7

src/mtcnn/stage3.ts

 import * as tf from '@tensorflow/tfjs-core';
 
+import { BoundingBox } from '../BoundingBox';
+import { nonMaxSuppression } from '../commons/nonMaxSuppression';
 import { Point } from '../Point';
-import { BoundingBox } from './BoundingBox';
 import { extractImagePatches } from './extractImagePatches';
-import { nms } from './nms';
 import { ONet } from './ONet';
 import { ONetParams } from './types';
 
@@ -57,7 +57,7 @@ export async function stage3(
   if (filteredBoxes.length > 0) {
 
     ts = Date.now()
-    const indicesNms = nms(
+    const indicesNms = nonMaxSuppression(
       filteredBoxes,
       filteredScores,
       0.7,

src/tinyYolov2/TinyYolov2.ts

 import * as tf from '@tensorflow/tfjs-core';
 
+import { BoundingBox } from '../BoundingBox';
 import { convLayer } from '../commons/convLayer';
 import { NeuralNetwork } from '../commons/NeuralNetwork';
+import { nonMaxSuppression } from '../commons/nonMaxSuppression';
+import { FaceDetection } from '../FaceDetection';
 import { NetInput } from '../NetInput';
 import { toNetInput } from '../toNetInput';
 import { TNetInput } from '../types';
+import { BOX_ANCHORS, INPUT_SIZES, NUM_BOXES, NUM_CELLS } from './config';
 import { convWithBatchNorm } from './convWithBatchNorm';
 import { extractParams } from './extractParams';
-import { NetParams } from './types';
-
+import { getDefaultParams } from './getDefaultParams';
+import { NetParams, TinyYolov2ForwardParams } from './types';
 
 export class TinyYolov2 extends NeuralNetwork<NetParams> {
 
@@ -16,7 +20,7 @@ export class TinyYolov2 extends NeuralNetwork<NetParams> {
     super('TinyYolov2')
   }
 
-  public async forwardInput(input: NetInput): Promise<any> {
+  public forwardInput(input: NetInput, inputSize: number): tf.Tensor4D {
 
     const { params } = this
 
@@ -25,10 +29,7 @@ export class TinyYolov2 extends NeuralNetwork<NetParams> {
     }
 
     const out = tf.tidy(() => {
-      //const batchTensor = input.toBatchTensor(416).div(tf.scalar(255)).toFloat() as tf.Tensor4D
-
-      // TODO: fix boxes after padding
-      const batchTensor = tf.image.resizeBilinear(input.inputs[0], [416, 416]).toFloat().div(tf.scalar(255)).expandDims() as tf.Tensor4D
+      const batchTensor = input.toBatchTensor(inputSize, false).div(tf.scalar(255)).toFloat() as tf.Tensor4D
 
       let out = convWithBatchNorm(batchTensor, params.conv0)
       out = tf.maxPool(out, [2, 2], [2, 2], 'same')
@@ -52,8 +53,75 @@ export class TinyYolov2 extends NeuralNetwork<NetParams> {
     return out
   }
 
-  public async forward(input: TNetInput): Promise<any> {
-    return await this.forwardInput(await toNetInput(input, true, true))
+  public async forward(input: TNetInput, inputSize: number): Promise<tf.Tensor4D> {
+    return await this.forwardInput(await toNetInput(input, true, true), inputSize)
+  }
+
+  public async locateFaces(input: TNetInput, forwardParams: TinyYolov2ForwardParams = {}): Promise<FaceDetection[]> {
+
+    const { sizeType, scoreThreshold } = getDefaultParams(forwardParams)
+
+
+    const inputSize = INPUT_SIZES[sizeType]
+    const numCells = NUM_CELLS[sizeType]
+
+    if (!inputSize) {
+      throw new Error(`TinyYolov2 - unkown sizeType: ${sizeType}, expected one of: xs | sm | md | lg`)
+    }
+
+    const netInput = await toNetInput(input, true)
+    const out = await this.forwardInput(netInput, inputSize)
+
+    const [boxesTensor, scoresTensor] = tf.tidy(() => {
+      const reshaped = out.reshape([numCells, numCells, NUM_BOXES, 6])
+      out.dispose()
+
+      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[] = []
+
+    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))
+          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 x = (ctX - (width / 2))
+            const y = (ctY - (height / 2))
+            boxes.push(new BoundingBox(x, y, x + width, y + height))
+            scores.push(score)
+          }
+        }
+      }
+    }
+
+    boxesTensor.dispose()
+    scoresTensor.dispose()
+
+    const indices = nonMaxSuppression(boxes, scores, 0.4, true)
+
+    const detections = indices.map(idx =>
+      new FaceDetection(
+        scores[idx],
+        boxes[idx].toRect(),
+        { width: netInput.getInputWidth(0), height: netInput.getInputHeight(0) }
+      )
+    )
+
+    return detections
   }
 
   /* TODO

src/tinyYolov2/config.ts

+import { Point } from '../Point';
+
+export const INPUT_SIZES = { xs: 224, sm: 320, md: 416, lg: 608 }
+export const NUM_CELLS = { xs: 7, sm: 10, md: 13, lg: 19 }
+export const NUM_BOXES = 5
+
+export const BOX_ANCHORS = [
+  new Point(0.738768, 0.874946),
+  new Point(2.42204, 2.65704),
+  new Point(4.30971, 7.04493),
+  new Point(10.246, 4.59428),
+  new Point(12.6868, 11.8741)
+]
\ No newline at end of file

src/tinyYolov2/getDefaultParams.ts

+import { SizeType, TinyYolov2ForwardParams } from './types';
+
+export function getDefaultParams(params: TinyYolov2ForwardParams) {
+  return Object.assign(
+    {},
+    {
+      sizeType: SizeType.MD,
+      scoreThreshold: 0.5
+    },
+    params
+  )
+}
\ No newline at end of file

src/tinyYolov2/types.ts

@@ -22,4 +22,16 @@ export type NetParams = {
   conv6: ConvWithBatchNorm
   conv7: ConvWithBatchNorm
   conv8: ConvParams
+}
+
+export enum SizeType {
+  XS = 'xs',
+  SM = 'sm',
+  MD = 'md',
+  LG = 'lg'
+}
+
+export type TinyYolov2ForwardParams = {
+  sizeType?: SizeType
+  scoreThreshold?: number
 }
\ No newline at end of file