implemented mtcnn model loading from url + expose mtcnn to global api + fixed some minor issues

justadudewhohacks · justadudewhohacks · commit 5d87492262c2 · 2018-07-14T12:42:10.000+02:00
diff --git a/src/FaceDetection.ts b/src/FaceDetection.ts
@@ -1,5 +1,5 @@
-import { Rect } from '../Rect';
-import { Dimensions } from '../types';
+import { Rect } from './Rect';
+import { Dimensions } from './types';
 
 export class FaceDetection {
   private _score: number
diff --git a/src/FullFaceDescription.ts b/src/FullFaceDescription.ts
@@ -1,4 +1,4 @@
-import { FaceDetection } from './faceDetectionNet/FaceDetection';
+import { FaceDetection } from './FaceDetection';
 import { FaceLandmarks68 } from './faceLandmarkNet/FaceLandmarks68';
 
 export class FullFaceDescription {
diff --git a/src/drawing/index.ts b/src/drawing/index.ts
@@ -1,4 +1,4 @@
-import { FaceDetection } from '../faceDetectionNet/FaceDetection';
+import { FaceDetection } from '../FaceDetection';
 import { FaceLandmarks68 } from '../faceLandmarkNet';
 import { FaceLandmarks } from '../FaceLandmarks';
 import { Point } from '../Point';
diff --git a/src/extractFaceTensors.ts b/src/extractFaceTensors.ts
@@ -1,6 +1,6 @@
 import * as tf from '@tensorflow/tfjs-core';
 
-import { FaceDetection } from './faceDetectionNet/FaceDetection';
+import { FaceDetection } from './FaceDetection';
 import { Rect } from './Rect';
 import { toNetInput } from './toNetInput';
 import { TNetInput } from './types';
diff --git a/src/extractFaces.ts b/src/extractFaces.ts
@@ -1,9 +1,8 @@
-import { FaceDetection } from './faceDetectionNet/FaceDetection';
+import { FaceDetection } from './FaceDetection';
 import { Rect } from './Rect';
 import { toNetInput } from './toNetInput';
 import { TNetInput } from './types';
 import { createCanvas, getContext2dOrThrow, imageTensorToCanvas } from './utils';
-import * as tf from '@tensorflow/tfjs-core';
 
 /**
  * Extracts the image regions containing the detected faces.
diff --git a/src/faceDetectionNet/FaceDetectionNet.ts b/src/faceDetectionNet/FaceDetectionNet.ts
@@ -1,12 +1,12 @@
 import * as tf from '@tensorflow/tfjs-core';
 
 import { NeuralNetwork } from '../commons/NeuralNetwork';
+import { FaceDetection } from '../FaceDetection';
 import { NetInput } from '../NetInput';
 import { Rect } from '../Rect';
 import { toNetInput } from '../toNetInput';
 import { TNetInput } from '../types';
 import { extractParams } from './extractParams';
-import { FaceDetection } from './FaceDetection';
 import { loadQuantizedParams } from './loadQuantizedParams';
 import { mobileNetV1 } from './mobileNetV1';
 import { nonMaxSuppression } from './nonMaxSuppression';
diff --git a/src/faceDetectionNet/index.ts b/src/faceDetectionNet/index.ts
@@ -1,10 +1,14 @@
 import { FaceDetectionNet } from './FaceDetectionNet';
 
 export * from './FaceDetectionNet';
-export * from './FaceDetection';
 
-export function faceDetectionNet(weights: Float32Array) {
+export function createFaceDetectionNet(weights: Float32Array) {
   const net = new FaceDetectionNet()
   net.extractWeights(weights)
   return net
+}
+
+export function faceDetectionNet(weights: Float32Array) {
+  console.warn('faceDetectionNet(weights: Float32Array) will be deprecated in future, use createFaceDetectionNet instead')
+  return createFaceDetectionNet(weights)
 }
diff --git a/src/faceDetectionNet/loadQuantizedParams.ts b/src/faceDetectionNet/loadQuantizedParams.ts
@@ -1,7 +1,8 @@
-import { tf } from '..';
+import * as tf from '@tensorflow/tfjs-core';
+
 import { disposeUnusedWeightTensors } from '../commons/disposeUnusedWeightTensors';
 import { extractWeightEntryFactory } from '../commons/extractWeightEntryFactory';
-import { isTensor1D, isTensor3D, isTensor4D } from '../commons/isTensor';
+import { isTensor3D } from '../commons/isTensor';
 import { loadWeightMap } from '../commons/loadWeightMap';
 import { ConvParams, ParamMapping } from '../commons/types';
 import { BoxPredictionParams, MobileNetV1, NetParams, PointwiseConvParams, PredictionLayerParams } from './types';
diff --git a/src/faceLandmarkNet/FaceLandmarks68.ts b/src/faceLandmarkNet/FaceLandmarks68.ts
@@ -1,9 +1,8 @@
 import { getCenterPoint } from '../commons/getCenterPoint';
-import { FaceDetection } from '../faceDetectionNet/FaceDetection';
+import { FaceDetection } from '../FaceDetection';
 import { FaceLandmarks } from '../FaceLandmarks';
 import { IPoint, Point } from '../Point';
 import { Rect } from '../Rect';
-import { Dimensions } from '../types';
 
 // face alignment constants
 const relX = 0.5
@@ -70,7 +69,7 @@ export class FaceLandmarks68 extends FaceLandmarks {
    * @returns The bounding box of the aligned face.
    */
   public align(
-    detection?: Rect
+    detection?: FaceDetection | Rect
   ): Rect {
     if (detection) {
       const box = detection instanceof FaceDetection
diff --git a/src/faceLandmarkNet/index.ts b/src/faceLandmarkNet/index.ts
@@ -3,8 +3,13 @@ import { FaceLandmarkNet } from './FaceLandmarkNet';
 export * from './FaceLandmarkNet';
 export * from './FaceLandmarks68';
 
-export function faceLandmarkNet(weights: Float32Array) {
+export function createFaceLandmarkNet(weights: Float32Array) {
   const net = new FaceLandmarkNet()
   net.extractWeights(weights)
   return net
+}
+
+export function faceLandmarkNet(weights: Float32Array) {
+  console.warn('faceLandmarkNet(weights: Float32Array) will be deprecated in future, use createFaceLandmarkNet instead')
+  return createFaceLandmarkNet(weights)
 }
diff --git a/src/faceRecognitionNet/index.ts b/src/faceRecognitionNet/index.ts
@@ -2,8 +2,13 @@ import { FaceRecognitionNet } from './FaceRecognitionNet';
 
 export * from './FaceRecognitionNet';
 
-export function faceRecognitionNet(weights: Float32Array) {
+export function createFaceRecognitionNet(weights: Float32Array) {
   const net = new FaceRecognitionNet()
   net.extractWeights(weights)
   return net
+}
+
+export function faceRecognitionNet(weights: Float32Array) {
+  console.warn('faceRecognitionNet(weights: Float32Array) will be deprecated in future, use createFaceRecognitionNet instead')
+  return createFaceRecognitionNet(weights)
 }
diff --git a/src/globalApi.ts b/src/globalApi.ts
@@ -1,36 +1,53 @@
 import * as tf from '@tensorflow/tfjs-core';
 
 import { allFacesFactory } from './allFacesFactory';
-import { FaceDetection } from './faceDetectionNet/FaceDetection';
+import { FaceDetection } from './FaceDetection';
 import { FaceDetectionNet } from './faceDetectionNet/FaceDetectionNet';
 import { FaceLandmarkNet } from './faceLandmarkNet/FaceLandmarkNet';
 import { FaceLandmarks68 } from './faceLandmarkNet/FaceLandmarks68';
 import { FaceRecognitionNet } from './faceRecognitionNet/FaceRecognitionNet';
 import { FullFaceDescription } from './FullFaceDescription';
+import { getDefaultMtcnnForwardParams } from './mtcnn/getDefaultMtcnnForwardParams';
+import { Mtcnn } from './mtcnn/Mtcnn';
+import { MtcnnForwardParams, MtcnnResult } from './mtcnn/types';
 import { NetInput } from './NetInput';
 import { TNetInput } from './types';
 
 export const detectionNet = new FaceDetectionNet()
 export const landmarkNet = new FaceLandmarkNet()
 export const recognitionNet = new FaceRecognitionNet()
 
+// nets need more specific names, to avoid ambiguity in future
+// when alternative net implementations are provided
+export const nets = {
+  ssdMobilenet: detectionNet,
+  faceLandmark68Net: landmarkNet,
+  faceNet: recognitionNet,
+  mtcnn: new Mtcnn()
+}
+
 export function loadFaceDetectionModel(url: string) {
-  return detectionNet.load(url)
+  return nets.ssdMobilenet.load(url)
 }
 
 export function loadFaceLandmarkModel(url: string) {
-  return landmarkNet.load(url)
+  return nets.faceLandmark68Net.load(url)
 }
 
 export function loadFaceRecognitionModel(url: string) {
-  return recognitionNet.load(url)
+  return nets.faceNet.load(url)
+}
+
+export function loadMtcnnModel(url: string) {
+  return nets.mtcnn.load(url)
 }
 
 export function loadModels(url: string) {
   return Promise.all([
     loadFaceDetectionModel(url),
     loadFaceLandmarkModel(url),
-    loadFaceRecognitionModel(url)
+    loadFaceRecognitionModel(url),
+    loadMtcnnModel(url)
   ])
 }
 
@@ -39,19 +56,26 @@ export function locateFaces(
   minConfidence?: number,
   maxResults?: number
 ): Promise<FaceDetection[]> {
-  return detectionNet.locateFaces(input, minConfidence, maxResults)
+  return nets.ssdMobilenet.locateFaces(input, minConfidence, maxResults)
 }
 
 export function detectLandmarks(
   input: TNetInput
 ): Promise<FaceLandmarks68 | FaceLandmarks68[]> {
-  return landmarkNet.detectLandmarks(input)
+  return nets.faceLandmark68Net.detectLandmarks(input)
 }
 
 export function computeFaceDescriptor(
   input: TNetInput
 ): Promise<Float32Array | Float32Array[]>  {
-  return recognitionNet.computeFaceDescriptor(input)
+  return nets.faceNet.computeFaceDescriptor(input)
+}
+
+export function mtcnn(
+  input: TNetInput,
+  forwardParameters: MtcnnForwardParams = getDefaultMtcnnForwardParams()
+): Promise<MtcnnResult[]> {
+  return nets.mtcnn.forward(input, forwardParameters)
 }
 
 export const allFaces: (
diff --git a/src/index.ts b/src/index.ts
@@ -4,6 +4,8 @@ export {
   tf
 }
 
+
+export * from './FaceDetection';
 export * from './FullFaceDescription';
 export * from './NetInput';
 export * from './Point';
diff --git a/src/mtcnn/Mtcnn.ts b/src/mtcnn/Mtcnn.ts
@@ -1,21 +1,24 @@
 import * as tf from '@tensorflow/tfjs-core';
 
 import { NeuralNetwork } from '../commons/NeuralNetwork';
-import { FaceDetection } from '../faceDetectionNet/FaceDetection';
+import { FaceDetection } from '../FaceDetection';
 import { NetInput } from '../NetInput';
 import { Point } from '../Point';
 import { Rect } from '../Rect';
 import { toNetInput } from '../toNetInput';
 import { TNetInput } from '../types';
 import { bgrToRgbTensor } from './bgrToRgbTensor';
+import { CELL_SIZE } from './config';
 import { extractParams } from './extractParams';
 import { FaceLandmarks5 } from './FaceLandmarks5';
+import { getDefaultMtcnnForwardParams } from './getDefaultMtcnnForwardParams';
 import { getSizesForScale } from './getSizesForScale';
+import { loadQuantizedParams } from './loadQuantizedParams';
 import { pyramidDown } from './pyramidDown';
 import { stage1 } from './stage1';
 import { stage2 } from './stage2';
 import { stage3 } from './stage3';
-import { MtcnnResult, NetParams } from './types';
+import { MtcnnForwardParams, MtcnnResult, NetParams } from './types';
 
 export class Mtcnn extends NeuralNetwork<NetParams> {
 
@@ -25,10 +28,7 @@ export class Mtcnn extends NeuralNetwork<NetParams> {
 
   public async forwardInput(
     input: NetInput,
-    minFaceSize: number = 20,
-    scaleFactor: number = 0.709,
-    maxNumScales: number = 10,
-    scoreThresholds: number[] = [0.6, 0.7, 0.7]
+    { minFaceSize, scaleFactor, maxNumScales, scoreThresholds, scaleSteps } = getDefaultMtcnnForwardParams()
   ): Promise<{ results: MtcnnResult[], stats: any }> {
 
     const { params } = this
@@ -64,10 +64,10 @@ export class Mtcnn extends NeuralNetwork<NetParams> {
 
     const [height, width] = imgTensor.shape.slice(1)
 
-    const scales = pyramidDown(minFaceSize, scaleFactor, [height, width])
+    const scales = scaleSteps || pyramidDown(minFaceSize, scaleFactor, [height, width])
       .filter(scale => {
         const sizes = getSizesForScale(scale, [height, width])
-        return Math.min(sizes.width, sizes.height) > 48
+        return Math.min(sizes.width, sizes.height) > CELL_SIZE
       })
       .slice(0, maxNumScales)
 
@@ -124,38 +124,31 @@ export class Mtcnn extends NeuralNetwork<NetParams> {
 
   public async forward(
     input: TNetInput,
-    minFaceSize: number = 20,
-    scaleFactor: number = 0.709,
-    maxNumScales: number = 10,
-    scoreThresholds: number[] = [0.6, 0.7, 0.7]
+    forwardParameters: MtcnnForwardParams = getDefaultMtcnnForwardParams()
   ): Promise<MtcnnResult[]> {
     return (
       await this.forwardInput(
         await toNetInput(input, true, true),
-        minFaceSize,
-        scaleFactor,
-        maxNumScales,
-        scoreThresholds
+        forwardParameters
       )
     ).results
   }
 
   public async forwardWithStats(
     input: TNetInput,
-    minFaceSize: number = 20,
-    scaleFactor: number = 0.709,
-    maxNumScales: number = 10,
-    scoreThresholds: number[] = [0.6, 0.7, 0.7]
+    forwardParameters: MtcnnForwardParams = getDefaultMtcnnForwardParams()
   ): Promise<{ results: MtcnnResult[], stats: any }> {
     return this.forwardInput(
       await toNetInput(input, true, true),
-      minFaceSize,
-      scaleFactor,
-      maxNumScales,
-      scoreThresholds
+      forwardParameters
     )
   }
 
+  // none of the param tensors are quantized yet
+  protected loadQuantizedParams(uri: string | undefined) {
+    return loadQuantizedParams(uri)
+  }
+
   protected extractParams(weights: Float32Array) {
     return extractParams(weights)
   }
diff --git a/src/mtcnn/extractParams.ts b/src/mtcnn/extractParams.ts
@@ -55,10 +55,10 @@ function extractorsFactory(extractWeights: ExtractWeightsFunction, paramMappings
     const conv4 = extractConvParams(64, 128, 2, 'onet/conv4')
     const prelu4_alpha = extractPReluParams(128, 'onet/prelu4_alpha')
     const fc1 = extractFCParams(1152, 256, 'onet/fc1')
-    const prelu5_alpha = extractPReluParams(256, 'onet/prelu4_alpha')
+    const prelu5_alpha = extractPReluParams(256, 'onet/prelu5_alpha')
     const fc2_1 = extractFCParams(256, 2, 'onet/fc2_1')
     const fc2_2 = extractFCParams(256, 4, 'onet/fc2_2')
-    const fc2_3 = extractFCParams(256, 10, 'onet/fc2_2')
+    const fc2_3 = extractFCParams(256, 10, 'onet/fc2_3')
 
     return { ...sharedParams, conv4, prelu4_alpha, fc1, prelu5_alpha, fc2_1, fc2_2, fc2_3 }
   }
diff --git a/src/mtcnn/getDefaultMtcnnForwardParams.ts b/src/mtcnn/getDefaultMtcnnForwardParams.ts
@@ -0,0 +1,10 @@
+import { MtcnnForwardParams } from './types';
+
+export function getDefaultMtcnnForwardParams(): MtcnnForwardParams {
+  return {
+    minFaceSize: 20,
+    scaleFactor: 0.709,
+    maxNumScales: 10,
+    scoreThresholds: [0.6, 0.7, 0.7]
+  }
+}
diff --git a/src/mtcnn/index.ts b/src/mtcnn/index.ts
@@ -1,8 +1,9 @@
 import { Mtcnn } from './Mtcnn';
 
 export * from './Mtcnn';
+export * from './FaceLandmarks5';
 
-export function mtcnn(weights: Float32Array) {
+export function createMtcnn(weights: Float32Array) {
   const net = new Mtcnn()
   net.extractWeights(weights)
   return net
diff --git a/src/mtcnn/loadQuantizedParams.ts b/src/mtcnn/loadQuantizedParams.ts
diff --git a/src/mtcnn/stage2.ts b/src/mtcnn/stage2.ts
diff --git a/src/mtcnn/stage3.ts b/src/mtcnn/stage3.ts
diff --git a/src/mtcnn/types.ts b/src/mtcnn/types.ts

Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,4 @@`
`1`		`-import { FaceDetection } from './faceDetectionNet/FaceDetection';`
	`1`	`+import { FaceDetection } from './FaceDetection';`
`2`	`2`	`import { FaceLandmarks68 } from './faceLandmarkNet/FaceLandmarks68';`
`3`	`3`
`4`	`4`	`export class FullFaceDescription {`
Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,4 @@`
`1`		`-import { FaceDetection } from '../faceDetectionNet/FaceDetection';`
	`1`	`+import { FaceDetection } from '../FaceDetection';`
`2`	`2`	`import { FaceLandmarks68 } from '../faceLandmarkNet';`
`3`	`3`	`import { FaceLandmarks } from '../FaceLandmarks';`
`4`	`4`	`import { Point } from '../Point';`