init mtcnn weights + weight loading + PNet

src/commons/extractConvParamsFactory.ts

+import * as tf from '@tensorflow/tfjs-core';
+
+import { ConvParams, ExtractWeightsFunction, ParamMapping } from './types';
+
+export function extractConvParamsFactory(
+  extractWeights: ExtractWeightsFunction,
+  paramMappings: ParamMapping[]
+) {
+
+  return function(
+    channelsIn: number,
+    channelsOut: number,
+    filterSize: number,
+    mappedPrefix: string
+  ): ConvParams {
+
+    const filters = tf.tensor4d(
+      extractWeights(channelsIn * channelsOut * filterSize * filterSize),
+      [filterSize, filterSize, channelsIn, channelsOut]
+    )
+    const bias = tf.tensor1d(extractWeights(channelsOut))
+
+    paramMappings.push(
+      { paramPath: `${mappedPrefix}/filters` },
+      { paramPath: `${mappedPrefix}/bias` }
+    )
+
+    return { filters, bias }
+  }
+
+}

src/commons/extractFCParamsFactory.ts

+import * as tf from '@tensorflow/tfjs-core';
+
+import { ExtractWeightsFunction, FCParams, ParamMapping } from './types';
+
+export function extractFCParamsFactory(
+  extractWeights: ExtractWeightsFunction,
+  paramMappings: ParamMapping[]
+) {
+
+  return function(
+    channelsIn: number,
+    channelsOut: number,
+    mappedPrefix: string
+  ): FCParams {
+
+    const fc_weights = tf.tensor2d(extractWeights(channelsIn * channelsOut), [channelsIn, channelsOut])
+    const fc_bias = tf.tensor1d(extractWeights(channelsOut))
+
+    paramMappings.push(
+      { paramPath: `${mappedPrefix}/weights` },
+      { paramPath: `${mappedPrefix}/bias` }
+    )
+
+    return {
+      weights: fc_weights,
+      bias: fc_bias
+    }
+  }
+
+}

src/commons/types.ts

@@ -5,6 +5,11 @@ export type ConvParams = {
   bias: tf.Tensor1D
 }
 
+export type FCParams = {
+  weights: tf.Tensor2D
+  bias: tf.Tensor1D
+}
+
 export type ExtractWeightsFunction = (numWeights: number) => Float32Array
 
 export type BatchReshapeInfo = {

src/faceLandmarkNet/extractParams.ts

-import * as tf from '@tensorflow/tfjs-core';
-
+import { extractConvParamsFactory } from '../commons/extractConvParamsFactory';
+import { extractFCParamsFactory } from '../commons/extractFCParamsFactory';
 import { extractWeightsFactory } from '../commons/extractWeightsFactory';
-import { ConvParams, ParamMapping } from '../commons/types';
-import { FCParams, NetParams } from './types';
+import { ParamMapping } from '../commons/types';
+import { NetParams } from './types';
 
 export function extractParams(weights: Float32Array): { params: NetParams, paramMappings: ParamMapping[] } {
 
@@ -13,42 +13,8 @@ export function extractParams(weights: Float32Array): { params: NetParams, param
     getRemainingWeights
   } = extractWeightsFactory(weights)
 
-  function extractConvParams(
-    channelsIn: number,
-    channelsOut: number,
-    filterSize: number,
-    mappedPrefix: string
-  ): ConvParams {
-
-    const filters = tf.tensor4d(
-      extractWeights(channelsIn * channelsOut * filterSize * filterSize),
-      [filterSize, filterSize, channelsIn, channelsOut]
-    )
-    const bias = tf.tensor1d(extractWeights(channelsOut))
-
-    paramMappings.push(
-      { paramPath: `${mappedPrefix}/filters` },
-      { paramPath: `${mappedPrefix}/bias` }
-    )
-
-    return { filters, bias }
-  }
-
-  function extractFcParams(channelsIn: number, channelsOut: number, mappedPrefix: string): FCParams {
-
-    const fc_weights = tf.tensor2d(extractWeights(channelsIn * channelsOut), [channelsIn, channelsOut])
-    const fc_bias = tf.tensor1d(extractWeights(channelsOut))
-
-    paramMappings.push(
-      { paramPath: `${mappedPrefix}/weights` },
-      { paramPath: `${mappedPrefix}/bias` }
-    )
-
-    return {
-      weights: fc_weights,
-      bias: fc_bias
-    }
-  }
+  const extractConvParams = extractConvParamsFactory(extractWeights, paramMappings)
+  const extractFCParams = extractFCParamsFactory(extractWeights, paramMappings)
 
   const conv0 = extractConvParams(3, 32, 3, 'conv0')
   const conv1 = extractConvParams(32, 64, 3, 'conv1')
@@ -58,8 +24,8 @@ export function extractParams(weights: Float32Array): { params: NetParams, param
   const conv5 = extractConvParams(64, 128, 3, 'conv5')
   const conv6 = extractConvParams(128, 128, 3, 'conv6')
   const conv7 = extractConvParams(128, 256, 3, 'conv7')
-  const fc0 = extractFcParams(6400, 1024, 'fc0')
-  const fc1 = extractFcParams(1024, 136, 'fc1')
+  const fc0 = extractFCParams(6400, 1024, 'fc0')
+  const fc1 = extractFCParams(1024, 136, 'fc1')
 
   if (getRemainingWeights().length !== 0) {
     throw new Error(`weights remaing after extract: ${getRemainingWeights().length}`)

src/faceLandmarkNet/fullyConnectedLayer.ts

 import * as tf from '@tensorflow/tfjs-core';
 
-import { FCParams } from './types';
+import { FCParams } from '../commons/types';
 
 export function fullyConnectedLayer(
   x: tf.Tensor2D,

src/faceLandmarkNet/loadQuantizedParams.ts

@@ -3,8 +3,8 @@ import * as tf from '@tensorflow/tfjs-core';
 import { disposeUnusedWeightTensors } from '../commons/disposeUnusedWeightTensors';
 import { extractWeightEntryFactory } from '../commons/extractWeightEntryFactory';
 import { loadWeightMap } from '../commons/loadWeightMap';
-import { ConvParams, ParamMapping } from '../commons/types';
-import { FCParams, NetParams } from './types';
+import { ConvParams, FCParams, ParamMapping } from '../commons/types';
+import { NetParams } from './types';
 
 const DEFAULT_MODEL_NAME = 'face_landmark_68_model'
 

src/faceLandmarkNet/types.ts

-import * as tf from '@tensorflow/tfjs-core';
-
-import { ConvParams } from '../commons/types';
-
-export type FCParams = {
-  weights: tf.Tensor2D
-  bias: tf.Tensor1D
-}
+import { ConvParams, FCParams } from '../commons/types';
 
 export type NetParams = {
   conv0: ConvParams

src/mtcnn/Mtcnn.ts

+import * as tf from '@tensorflow/tfjs-core';
+
+import { NeuralNetwork } from '../commons/NeuralNetwork';
+import { NetInput } from '../NetInput';
+import { toNetInput } from '../toNetInput';
+import { TNetInput } from '../types';
+import { extractParams } from './extractParams';
+import { stage1 } from './stage1';
+import { NetParams } from './types';
+
+export class Mtcnn extends NeuralNetwork<NetParams> {
+
+  constructor() {
+    super('Mtcnn')
+  }
+
+  public forwardInput(input: NetInput, minFaceSize: number = 20, scaleFactor: number = 0.709): tf.Tensor2D {
+
+    const { params } = this
+
+    if (!params) {
+      throw new Error('Mtcnn - load model before inference')
+    }
+
+    return tf.tidy(() => {
+      const imgTensor = tf.expandDims(input.inputs[0]).toFloat() as tf.Tensor4D
+
+      function pyramidDown(minFaceSize: number, scaleFactor: number, dims: number[]): number[] {
+
+        const [height, width] = dims
+        const m = 12 / minFaceSize
+
+        const scales = []
+
+        let minLayer = Math.min(height, width) * m
+        let exp = 0
+        while (minLayer >= 12) {
+          scales.push(m * Math.pow(scaleFactor, exp))
+          minLayer = minLayer * scaleFactor
+          exp += 1
+        }
+
+        return scales
+      }
+
+      const scales = pyramidDown(minFaceSize, scaleFactor, imgTensor.shape)
+      const out1 = stage1(imgTensor, scales, params.pnet)
+
+      return tf.tensor2d([0], [1, 1])
+    })
+  }
+
+  public async forward(input: TNetInput): Promise<tf.Tensor2D> {
+    return this.forwardInput(await toNetInput(input, true))
+  }
+
+  protected extractParams(weights: Float32Array) {
+    return extractParams(weights)
+  }
+}
\ No newline at end of file

src/mtcnn/PNet.ts

+import * as tf from '@tensorflow/tfjs-core';
+
+import { convLayer } from '../commons/convLayer';
+import { sharedLayer } from './sharedLayers';
+import { PNetParams } from './types';
+
+export function PNet(x: tf.Tensor4D, params: PNetParams): { prob: tf.Tensor3D, convOut: tf.Tensor4D } {
+  return tf.tidy(() => {
+
+    let out = sharedLayer(x, params)
+    const conv = convLayer(out, params.conv4_1, 'valid')
+    // TODO: tf.reduce_max <=> tf.max ?
+    const logits = tf.sub(conv, tf.max(conv, 3))
+    const prob = tf.softmax(logits, 3) as tf.Tensor3D
+    const convOut = convLayer(out, params.conv4_2, 'valid')
+
+    return { prob, convOut }
+  })
+}
\ No newline at end of file

src/mtcnn/extractParams.ts

+import * as tf from '@tensorflow/tfjs-core';
+
+import { extractConvParamsFactory } from '../commons/extractConvParamsFactory';
+import { extractFCParamsFactory } from '../commons/extractFCParamsFactory';
+import { extractWeightsFactory } from '../commons/extractWeightsFactory';
+import { ExtractWeightsFunction, ParamMapping } from '../commons/types';
+import { NetParams, PNetParams, RNetParams, SharedParams, ONetParams } from './types';
+
+function extractorsFactory(extractWeights: ExtractWeightsFunction, paramMappings: ParamMapping[]) {
+
+  const extractConvParams = extractConvParamsFactory(extractWeights, paramMappings)
+  const extractFCParams = extractFCParamsFactory(extractWeights, paramMappings)
+
+  function extractPReluParams(size: number, paramPath: string): tf.Tensor1D {
+    const alpha = tf.tensor1d(extractWeights(size))
+    paramMappings.push({ paramPath })
+    return alpha
+  }
+
+  function extractSharedParams(numFilters: number[], mappedPrefix: string, isRnet: boolean = false): SharedParams {
+
+    const conv1 = extractConvParams(numFilters[0], numFilters[1], 3, `${mappedPrefix}/conv1`)
+    const prelu1_alpha = extractPReluParams(numFilters[1], `${mappedPrefix}/prelu1_alpha`)
+    const conv2 = extractConvParams(numFilters[1], numFilters[2], 3, `${mappedPrefix}/conv2`)
+    const prelu2_alpha = extractPReluParams(numFilters[2], `${mappedPrefix}/prelu2_alpha`)
+    const conv3 = extractConvParams(numFilters[2], numFilters[3], isRnet ? 2 : 3, `${mappedPrefix}/conv3`)
+    const prelu3_alpha = extractPReluParams(numFilters[3], `${mappedPrefix}/prelu3_alpha`)
+
+    return { conv1, prelu1_alpha, conv2, prelu2_alpha, conv3, prelu3_alpha }
+  }
+
+  function extractPNetParams(): PNetParams {
+
+    const sharedParams = extractSharedParams([3, 10, 16, 32], 'pnet')
+    const conv4_1 = extractConvParams(32, 2, 1, 'pnet/conv4_1')
+    const conv4_2 = extractConvParams(32, 4, 1, 'pnet/conv4_2')
+
+    return { ...sharedParams, conv4_1, conv4_2 }
+  }
+
+  function extractRNetParams(): RNetParams {
+
+    const sharedParams = extractSharedParams([3, 28, 48, 64], 'rnet')
+    const fc1 = extractFCParams(576, 128, 'rnet/fc1')
+    const prelu4_alpha = extractPReluParams(128, 'rnet/prelu4_alpha')
+    const fc2_1 = extractFCParams(128, 2, 'rnet/fc2_1')
+    const fc2_2 = extractFCParams(128, 4, 'rnet/fc2_2')
+
+    return { ...sharedParams, fc1, prelu4_alpha, fc2_1, fc2_2 }
+  }
+
+  function extractONetParams(): ONetParams {
+
+    const sharedParams = extractSharedParams([3, 32, 64, 64], 'onet')
+    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 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')
+
+    return { ...sharedParams, conv4, prelu4_alpha, fc1, prelu5_alpha, fc2_1, fc2_2, fc2_3 }
+  }
+
+  return {
+    extractPNetParams,
+    extractRNetParams,
+    extractONetParams
+  }
+
+}
+
+export function extractParams(weights: Float32Array): { params: NetParams, paramMappings: ParamMapping[] } {
+
+  const {
+    extractWeights,
+    getRemainingWeights
+  } = extractWeightsFactory(weights)
+
+  const paramMappings: ParamMapping[] = []
+
+  const {
+    extractPNetParams,
+    extractRNetParams,
+    extractONetParams
+  } = extractorsFactory(extractWeights, paramMappings)
+
+  const pnet = extractPNetParams()
+  const rnet = extractRNetParams()
+  const onet = extractONetParams()
+
+  return { params: { pnet, rnet, onet }, paramMappings }
+}
\ No newline at end of file

src/mtcnn/index.ts

+import { Mtcnn } from './Mtcnn';
+
+export * from './Mtcnn';
+
+export function mtcnn(weights: Float32Array) {
+  const net = new Mtcnn()
+  net.extractWeights(weights)
+  return net
+}
\ No newline at end of file

src/mtcnn/prelu.ts

+import * as tf from '@tensorflow/tfjs-core';
+
+export function prelu(x: tf.Tensor4D, alpha: tf.Tensor1D): tf.Tensor4D {
+  return tf.tidy(() =>
+    tf.add(
+      tf.relu(x),
+      tf.mul(alpha, tf.neg(tf.relu(tf.neg(x))))
+    )
+  )
+}

src/mtcnn/sharedLayers.ts

+import * as tf from '@tensorflow/tfjs-core';
+
+import { convLayer } from '../commons/convLayer';
+import { prelu } from './prelu';
+import { SharedParams } from './types';
+
+export function sharedLayer(x: tf.Tensor4D, params: SharedParams, isPnet: boolean = false) {
+  return tf.tidy(() => {
+
+    let out = convLayer(x, params.conv1, 'valid')
+    out = prelu(out, params.prelu1_alpha)
+    out = tf.maxPool(out, isPnet ? [2, 2]: [3, 3], [2, 2], 'same')
+    out = convLayer(out, params.conv2, 'valid')
+    out = prelu(out, params.prelu2_alpha)
+    out = isPnet ? out : tf.maxPool(out, [3, 3], [2, 2], 'valid')
+    out = convLayer(out, params.conv3, 'valid')
+    out = prelu(out, params.prelu3_alpha)
+
+    return out
+  })
+}
\ No newline at end of file

src/mtcnn/stage1.ts

+import * as tf from '@tensorflow/tfjs-core';
+
+import { PNet } from './PNet';
+import { PNetParams } from './types';
+
+function rescaleAndNormalize(x: tf.Tensor4D, scale: number): tf.Tensor4D {
+  return tf.tidy(() => {
+    const [height, width] = x.shape
+    const resized = tf.image.resizeBilinear(x, [height * scale, width * scale])
+
+    return tf.mul(tf.sub(resized, tf.scalar(127.5)), tf.scalar(0.0078125))
+
+    // TODO: ?
+    // img_x = np.expand_dims(scaled_image, 0)
+    // img_y = np.transpose(img_x, (0, 2, 1, 3))
+
+  })
+}
+
+export function stage1(x: tf.Tensor4D, scales: number[], params: PNetParams) {
+  return tf.tidy(() => {
+
+    const boxes = scales.map((scale) => {
+      const resized = rescaleAndNormalize(x, scale)
+      const { prob, convOut } = PNet(resized, params)
+    })
+
+  })
+}
+
+/*
+
+  for scale in scales:
+      scaled_image = self.__scale_image(image, scale)
+
+      img_x = np.expand_dims(scaled_image, 0)
+      img_y = np.transpose(img_x, (0, 2, 1, 3))
+
+      out = self.__pnet.feed(img_y)
+
+      out0 = np.transpose(out[0], (0, 2, 1, 3))
+      out1 = np.transpose(out[1], (0, 2, 1, 3))
+
+      boxes, _ = self.__generate_bounding_box(out1[0, :, :, 1].copy(),
+                                              out0[0, :, :, :].copy(), scale, self.__steps_threshold[0])
+
+      # inter-scale nms
+      pick = self.__nms(boxes.copy(), 0.5, 'Union')
+      if boxes.size > 0 and pick.size > 0:
+          boxes = boxes[pick, :]
+          total_boxes = np.append(total_boxes, boxes, axis=0)
+
+
+
+
+  numboxes = total_boxes.shape[0]
+
+  if numboxes > 0:
+      pick = self.__nms(total_boxes.copy(), 0.7, 'Union')
+      total_boxes = total_boxes[pick, :]
+
+      regw = total_boxes[:, 2] - total_boxes[:, 0]
+      regh = total_boxes[:, 3] - total_boxes[:, 1]
+
+      qq1 = total_boxes[:, 0] + total_boxes[:, 5] * regw
+      qq2 = total_boxes[:, 1] + total_boxes[:, 6] * regh
+      qq3 = total_boxes[:, 2] + total_boxes[:, 7] * regw
+      qq4 = total_boxes[:, 3] + total_boxes[:, 8] * regh
+
+      total_boxes = np.transpose(np.vstack([qq1, qq2, qq3, qq4, total_boxes[:, 4]]))
+      total_boxes = self.__rerec(total_boxes.copy())
+
+      total_boxes[:, 0:4] = np.fix(total_boxes[:, 0:4]).astype(np.int32)
+      status = StageStatus(self.__pad(total_boxes.copy(), stage_status.width, stage_status.height),
+                            width=stage_status.width, height=stage_status.height)
+
+  return total_boxes, status
+  */
\ No newline at end of file

src/mtcnn/types.ts

+import { ConvParams, FCParams } from '../commons/types';
+import { tf } from '..';
+
+export type SharedParams = {
+  conv1: ConvParams
+  prelu1_alpha: tf.Tensor1D
+  conv2: ConvParams
+  prelu2_alpha: tf.Tensor1D
+  conv3: ConvParams
+  prelu3_alpha: tf.Tensor1D
+}
+
+export type PNetParams = SharedParams & {
+  conv4_1: ConvParams
+  conv4_2: ConvParams
+}
+
+export type RNetParams = SharedParams & {
+  fc1: FCParams
+  prelu4_alpha: tf.Tensor1D
+  fc2_1: FCParams
+  fc2_2: FCParams
+}
+
+export type ONetParams = SharedParams & {
+  conv4: ConvParams
+  prelu4_alpha: tf.Tensor1D
+  fc1: FCParams
+  prelu5_alpha: tf.Tensor1D
+  fc2_1: FCParams
+  fc2_2: FCParams
+  fc2_3: FCParams
+}
+
+export type NetParams = {
+  pnet: PNetParams
+  rnet: RNetParams
+  onet: ONetParams
+}
\ No newline at end of file

weights/mtcnn_model-weights_manifest.json

+[{"paths":["mtcnn_model-shard1"],"weights":[{"dtype":"float32","name":"pnet/conv1/weights","shape":[3,3,3,10]},{"dtype":"float32","name":"pnet/conv1/biases","shape":[10]},{"dtype":"float32","name":"pnet/prelu1_alpha","shape":[10]},{"dtype":"float32","name":"pnet/conv2/weights","shape":[3,3,10,16]},{"dtype":"float32","name":"pnet/conv2/biases","shape":[16]},{"dtype":"float32","name":"pnet/prelu2_alpha","shape":[16]},{"dtype":"float32","name":"pnet/conv3/weights","shape":[3,3,16,32]},{"dtype":"float32","name":"pnet/conv3/biases","shape":[32]},{"dtype":"float32","name":"pnet/prelu3_alpha","shape":[32]},{"dtype":"float32","name":"pnet/conv4-1/weights","shape":[1,1,32,2]},{"dtype":"float32","name":"pnet/conv4-1/biases","shape":[2]},{"dtype":"float32","name":"pnet/conv4-2/weights","shape":[1,1,32,4]},{"dtype":"float32","name":"pnet/conv4-2/biases","shape":[4]},{"dtype":"float32","name":"rnet/conv1/weights","shape":[3,3,3,28]},{"dtype":"float32","name":"rnet/conv1/biases","shape":[28]},{"dtype":"float32","name":"rnet/prelu1_alpha","shape":[28]},{"dtype":"float32","name":"rnet/conv2/weights","shape":[3,3,28,48]},{"dtype":"float32","name":"rnet/conv2/biases","shape":[48]},{"dtype":"float32","name":"rnet/prelu2_alpha","shape":[48]},{"dtype":"float32","name":"rnet/conv3/weights","shape":[2,2,48,64]},{"dtype":"float32","name":"rnet/conv3/biases","shape":[64]},{"dtype":"float32","name":"rnet/prelu3_alpha","shape":[64]},{"dtype":"float32","name":"rnet/fc1/weights","shape":[576,128]},{"dtype":"float32","name":"rnet/fc1/biases","shape":[128]},{"dtype":"float32","name":"rnet/prelu4_alpha","shape":[128]},{"dtype":"float32","name":"rnet/fc2-1/weights","shape":[128,2]},{"dtype":"float32","name":"rnet/fc2-1/biases","shape":[2]},{"dtype":"float32","name":"rnet/fc2-2/weights","shape":[128,4]},{"dtype":"float32","name":"rnet/fc2-2/biases","shape":[4]},{"dtype":"float32","name":"onet/conv1/weights","shape":[3,3,3,32]},{"dtype":"float32","name":"onet/conv1/biases","shape":[32]},{"dtype":"float32","name":"onet/prelu1_alpha","shape":[32]},{"dtype":"float32","name":"onet/conv2/weights","shape":[3,3,32,64]},{"dtype":"float32","name":"onet/conv2/biases","shape":[64]},{"dtype":"float32","name":"onet/prelu2_alpha","shape":[64]},{"dtype":"float32","name":"onet/conv3/weights","shape":[3,3,64,64]},{"dtype":"float32","name":"onet/conv3/biases","shape":[64]},{"dtype":"float32","name":"onet/prelu3_alpha","shape":[64]},{"dtype":"float32","name":"onet/conv4/weights","shape":[2,2,64,128]},{"dtype":"float32","name":"onet/conv4/biases","shape":[128]},{"dtype":"float32","name":"onet/prelu4_alpha","shape":[128]},{"dtype":"float32","name":"onet/fc1/weights","shape":[1152,256]},{"dtype":"float32","name":"onet/fc1/biases","shape":[256]},{"dtype":"float32","name":"onet/prelu5_alpha","shape":[256]},{"dtype":"float32","name":"onet/fc2-1/weights","shape":[256,2]},{"dtype":"float32","name":"onet/fc2-1/biases","shape":[2]},{"dtype":"float32","name":"onet/fc2-2/weights","shape":[256,4]},{"dtype":"float32","name":"onet/fc2-2/biases","shape":[4]},{"dtype":"float32","name":"onet/fc2-3/weights","shape":[256,10]},{"dtype":"float32","name":"onet/fc2-3/biases","shape":[10]}]}]
\ No newline at end of file