allow arbitrary input size

justadudewhohacks · justadudewhohacks · commit e95dd35db029 · 2018-07-30T18:07:51.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, IOU_THRESHOLD, NUM_BOXES, NUM_CELLS } from './config';
+import { BOX_ANCHORS, INPUT_SIZES, IOU_THRESHOLD, NUM_BOXES } from './config';
 import { convWithBatchNorm } from './convWithBatchNorm';
 import { extractParams } from './extractParams';
 import { getDefaultParams } from './getDefaultParams';
@@ -59,18 +59,19 @@ export class TinyYolov2 extends NeuralNetwork<NetParams> {
 
   public async locateFaces(input: TNetInput, forwardParams: TinyYolov2ForwardParams = {}): Promise<FaceDetection[]> {
 
-    const { sizeType, scoreThreshold } = getDefaultParams(forwardParams)
+    const { inputSize: _inputSize, scoreThreshold } = getDefaultParams(forwardParams)
 
+    const inputSize = typeof _inputSize === 'string'
+      ? INPUT_SIZES[_inputSize]
+      : _inputSize
 
-    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`)
+    if (typeof inputSize !== 'number') {
+      throw new Error(`TinyYolov2 - unkown inputSize: ${inputSize}, expected number or one of xs | sm | md | lg`)
     }
 
     const netInput = await toNetInput(input, true)
     const out = await this.forwardInput(netInput, inputSize)
+    const numCells = out.shape[1]
 
     const [boxesTensor, scoresTensor] = tf.tidy(() => {
       const reshaped = out.reshape([numCells, numCells, NUM_BOXES, 6])
diff --git a/src/tinyYolov2/config.ts b/src/tinyYolov2/config.ts
@@ -1,7 +1,6 @@
 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
 
diff --git a/src/tinyYolov2/types.ts b/src/tinyYolov2/types.ts
@@ -32,6 +32,6 @@ export enum SizeType {
 }
 
 export type TinyYolov2ForwardParams = {
-  sizeType?: SizeType
+  inputSize?: SizeType | number
   scoreThreshold?: number
 }

Original file line number	Diff line number	Diff line change
`@@ -32,6 +32,6 @@ export enum SizeType {`
`32`	`32`	`}`
`33`	`33`
`34`	`34`	`export type TinyYolov2ForwardParams = {`
`35`		`- sizeType?: SizeType`
	`35`	`+ inputSize?: SizeType \| number`
`36`	`36`	`scoreThreshold?: number`
`37`	`37`	`}`