apply non max suppression

src/faceDetectionNet/index.ts

@@ -6,6 +6,7 @@ import { mobileNetV1 } from './mobileNetV1';
 import { resizeLayer } from './resizeLayer';
 import { predictionLayer } from './predictionLayer';
 import { outputLayer } from './outputLayer';
+import { nonMaxSuppression } from './nonMaxSuppression';
 
 function fromData(input: number[]): tf.Tensor4D {
   const pxPerChannel = input.length / 3
@@ -75,7 +76,8 @@ export function faceDetectionNet(weights: Float32Array) {
 
   async function locateFaces(
     input: ImageData|ImageData[]|number[],
-    minConfidence: number = 0.8
+    minConfidence: number = 0.8,
+    maxResults: number = 100,
   ) {
     const imgTensor = getImgTensor(input)
 
@@ -91,13 +93,20 @@ export function faceDetectionNet(weights: Float32Array) {
     const scores = _scores[0]
 
     // TODO find a better way to filter by minConfidence
-    const data = await scores.data()
+    const scoresData = Array.from(await scores.data())
+
+    const iouThreshold = 0.5
+    const indices = nonMaxSuppression(
+      boxes,
+      scoresData,
+      maxResults,
+      iouThreshold,
+      minConfidence
+    )
 
-    return Array.from(data)
-      .map((score, idx) => ({ score, idx }))
-      .filter(({ score }) => minConfidence < score)
-      .map(({ score, idx }) => ({
-        score,
+    return indices
+      .map(idx => ({
+        score: scoresData[idx],
         box: {
           top: Math.max(0, height * boxes.get(idx, 0)),
           left: Math.max(0, width * boxes.get(idx, 1)),

src/faceDetectionNet/nonMaxSuppression.ts

+import * as tf from '@tensorflow/tfjs-core'
+
+export function nonMaxSuppression(
+  boxes: tf.Tensor2D,
+  scores: number[],
+  maxOutputSize: number,
+  iouThreshold: number,
+  scoreThreshold: number
+): number[] {
+
+  const numBoxes = boxes.shape[0]
+  const outputSize = Math.min(
+    maxOutputSize,
+    numBoxes
+  )
+
+  const candidates = scores
+    .map((score, boxIndex) => ({ score, boxIndex }))
+    .filter(c => c.score > scoreThreshold)
+    .sort((c1, c2) => c2.score - c1.score)
+
+  const suppressFunc = (x: number) => x <= iouThreshold ? 1 : 0
+
+  const selected: number[] = []
+
+  candidates.forEach(c => {
+    if (selected.length >= outputSize) {
+      return
+    }
+    const originalScore = c.score
+
+    for (let j = selected.length - 1; j >= 0; --j) {
+      const iou = IOU(boxes, c.boxIndex, selected[j])
+      if (iou === 0.0) {
+        continue
+      }
+      c.score *= suppressFunc(iou)
+      if (c.score <= scoreThreshold) {
+        break
+      }
+    }
+
+    if (originalScore === c.score) {
+      selected.push(c.boxIndex)
+    }
+  })
+
+  return selected
+}
+
+function IOU(boxes: tf.Tensor2D, i: number, j: number) {
+  const yminI = Math.min(boxes.get(i, 0), boxes.get(i, 2))
+  const xminI = Math.min(boxes.get(i, 1), boxes.get(i, 3))
+  const ymaxI = Math.max(boxes.get(i, 0), boxes.get(i, 2))
+  const xmaxI = Math.max(boxes.get(i, 1), boxes.get(i, 3))
+  const yminJ = Math.min(boxes.get(j, 0), boxes.get(j, 2))
+  const xminJ = Math.min(boxes.get(j, 1), boxes.get(j, 3))
+  const ymaxJ = Math.max(boxes.get(j, 0), boxes.get(j, 2))
+  const xmaxJ = Math.max(boxes.get(j, 1), boxes.get(j, 3))
+  const areaI = (ymaxI - yminI) * (xmaxI - xminI)
+  const areaJ = (ymaxJ - yminJ) * (xmaxJ - xminJ)
+  if (areaI <= 0 || areaJ <= 0) {
+    return 0.0
+  }
+  const intersectionYmin = Math.max(yminI, yminJ)
+  const intersectionXmin = Math.max(xminI, xminJ)
+  const intersectionYmax = Math.min(ymaxI, ymaxJ)
+  const intersectionXmax = Math.min(xmaxI, xmaxJ)
+  const intersectionArea =
+      Math.max(intersectionYmax - intersectionYmin, 0.0) *
+      Math.max(intersectionXmax - intersectionXmin, 0.0)
+  return intersectionArea / (areaI + areaJ - intersectionArea)
+}
\ No newline at end of file