fixed nms issue

justadudewhohacks · justadudewhohacks · commit 28a774538d16 · 2018-07-30T17:06:36.000+02:00
diff --git a/src/tinyYolov2/TinyYolov2.ts b/src/tinyYolov2/TinyYolov2.ts
@@ -8,7 +8,7 @@ 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 { BOX_ANCHORS, INPUT_SIZES, IOU_THRESHOLD, NUM_BOXES, NUM_CELLS } from './config';
 import { convWithBatchNorm } from './convWithBatchNorm';
 import { extractParams } from './extractParams';
 import { getDefaultParams } from './getDefaultParams';
@@ -111,7 +111,17 @@ export class TinyYolov2 extends NeuralNetwork<NetParams> {
     boxesTensor.dispose()
     scoresTensor.dispose()
 
-    const indices = nonMaxSuppression(boxes, scores, 0.4, true)
+    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(
diff --git a/src/tinyYolov2/config.ts b/src/tinyYolov2/config.ts
@@ -3,6 +3,7 @@ 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 IOU_THRESHOLD = 0.4
 
 export const BOX_ANCHORS = [
   new Point(0.738768, 0.874946),
