seperate feature extraction nets from landmark nets to make them reusable

src/faceFeatureExtractor/FaceFeatureExtractor.ts

+import * as tf from '@tensorflow/tfjs-core';
+import { NetInput, NeuralNetwork, normalize } from 'tfjs-image-recognition-base';
+import { ConvParams, SeparableConvParams } from 'tfjs-tiny-yolov2';
+
+import { depthwiseSeparableConv } from './depthwiseSeparableConv';
+import { extractParams } from './extractParams';
+import { extractParamsFromWeigthMap } from './extractParamsFromWeigthMap';
+import { DenseBlock4Params, NetParams } from './types';
+
+function denseBlock(
+  x: tf.Tensor4D,
+  denseBlockParams: DenseBlock4Params,
+  isFirstLayer: boolean = false
+): tf.Tensor4D {
+  return tf.tidy(() => {
+    const out1 = tf.relu(
+      isFirstLayer
+        ? tf.add(
+          tf.conv2d(x, (denseBlockParams.conv0 as ConvParams).filters, [2, 2], 'same'),
+          denseBlockParams.conv0.bias
+        )
+        : depthwiseSeparableConv(x, denseBlockParams.conv0 as SeparableConvParams, [2, 2])
+    ) as tf.Tensor4D
+    const out2 = depthwiseSeparableConv(out1, denseBlockParams.conv1, [1, 1])
+
+    const in3 = tf.relu(tf.add(out1, out2)) as tf.Tensor4D
+    const out3 = depthwiseSeparableConv(in3, denseBlockParams.conv2, [1, 1])
+
+    const in4 = tf.relu(tf.add(out1, tf.add(out2, out3))) as tf.Tensor4D
+    const out4 = depthwiseSeparableConv(in4, denseBlockParams.conv3, [1, 1])
+
+    return tf.relu(tf.add(out1, tf.add(out2, tf.add(out3, out4)))) as tf.Tensor4D
+  })
+}
+
+export class FaceFeatureExtractor extends NeuralNetwork<NetParams> {
+
+  constructor() {
+    super('FaceFeatureExtractor')
+  }
+
+  public forward(input: NetInput): tf.Tensor4D {
+
+    const { params } = this
+
+    if (!params) {
+      throw new Error('FaceFeatureExtractor - load model before inference')
+    }
+
+    return tf.tidy(() => {
+      const batchTensor = input.toBatchTensor(112, true)
+      const meanRgb = [122.782, 117.001, 104.298]
+      const normalized = normalize(batchTensor, meanRgb).div(tf.scalar(255)) as tf.Tensor4D
+
+      let out = denseBlock(normalized, params.dense0, true)
+      out = denseBlock(out, params.dense1)
+      out = denseBlock(out, params.dense2)
+      out = denseBlock(out, params.dense3)
+      out = tf.avgPool(out, [7, 7], [2, 2], 'valid')
+
+      return out
+    })
+  }
+
+  protected getDefaultModelName(): string {
+    return 'face_feature_extractor_model'
+  }
+
+  protected extractParamsFromWeigthMap(weightMap: tf.NamedTensorMap) {
+    return extractParamsFromWeigthMap(weightMap)
+  }
+
+  protected extractParams(weights: Float32Array) {
+    return extractParams(weights)
+  }
+}
\ No newline at end of file

src/faceFeatureExtractor/TinyFaceFeatureExtractor.ts

+import * as tf from '@tensorflow/tfjs-core';
+import { NetInput, NeuralNetwork, normalize } from 'tfjs-image-recognition-base';
+import { ConvParams, SeparableConvParams } from 'tfjs-tiny-yolov2';
+
+import { depthwiseSeparableConv } from './depthwiseSeparableConv';
+import { extractParamsFromWeigthMapTiny } from './extractParamsFromWeigthMapTiny';
+import { extractParamsTiny } from './extractParamsTiny';
+import { DenseBlock3Params, TinyNetParams } from './types';
+
+function denseBlock(
+  x: tf.Tensor4D,
+  denseBlockParams: DenseBlock3Params,
+  isFirstLayer: boolean = false
+): tf.Tensor4D {
+  return tf.tidy(() => {
+    const out1 = tf.relu(
+      isFirstLayer
+        ? tf.add(
+          tf.conv2d(x, (denseBlockParams.conv0 as ConvParams).filters, [2, 2], 'same'),
+          denseBlockParams.conv0.bias
+        )
+        : depthwiseSeparableConv(x, denseBlockParams.conv0 as SeparableConvParams, [2, 2])
+    ) as tf.Tensor4D
+    const out2 = depthwiseSeparableConv(out1, denseBlockParams.conv1, [1, 1])
+
+    const in3 = tf.relu(tf.add(out1, out2)) as tf.Tensor4D
+    const out3 = depthwiseSeparableConv(in3, denseBlockParams.conv2, [1, 1])
+
+    return tf.relu(tf.add(out1, tf.add(out2, out3))) as tf.Tensor4D
+  })
+}
+
+export class TinyFaceFeatureExtractor extends NeuralNetwork<TinyNetParams> {
+
+  constructor() {
+    super('TinyFaceFeatureExtractor')
+  }
+
+  public forward(input: NetInput): tf.Tensor4D {
+
+    const { params } = this
+
+    if (!params) {
+      throw new Error('TinyFaceFeatureExtractor - load model before inference')
+    }
+
+    return tf.tidy(() => {
+      const batchTensor = input.toBatchTensor(112, true)
+      const meanRgb = [122.782, 117.001, 104.298]
+      const normalized = normalize(batchTensor, meanRgb).div(tf.scalar(255)) as tf.Tensor4D
+
+      let out = denseBlock(normalized, params.dense0, true)
+      out = denseBlock(out, params.dense1)
+      out = denseBlock(out, params.dense2)
+      out = tf.avgPool(out, [14, 14], [2, 2], 'valid')
+
+      return out
+    })
+  }
+
+  protected getDefaultModelName(): string {
+    return 'face_landmark_68_tiny_model'
+  }
+
+  protected extractParamsFromWeigthMap(weightMap: tf.NamedTensorMap) {
+    return extractParamsFromWeigthMapTiny(weightMap)
+  }
+
+  protected extractParams(weights: Float32Array) {
+    return extractParamsTiny(weights)
+  }
+}
\ No newline at end of file

src/faceFeatureExtractor/extractParams.ts

+import { extractWeightsFactory, ParamMapping } from 'tfjs-image-recognition-base';
+
+import { extractorsFactory } from './extractorsFactory';
+import { NetParams } from './types';
+
+export function extractParams(weights: Float32Array): { params: NetParams, paramMappings: ParamMapping[] } {
+
+  const paramMappings: ParamMapping[] = []
+
+  const {
+    extractWeights,
+    getRemainingWeights
+  } = extractWeightsFactory(weights)
+
+  const {
+    extractDenseBlock4Params
+  } = extractorsFactory(extractWeights, paramMappings)
+
+  const dense0 = extractDenseBlock4Params(3, 32, 'dense0', true)
+  const dense1 = extractDenseBlock4Params(32, 64, 'dense1')
+  const dense2 = extractDenseBlock4Params(64, 128, 'dense2')
+  const dense3 = extractDenseBlock4Params(128, 256, 'dense3')
+
+  if (getRemainingWeights().length !== 0) {
+    throw new Error(`weights remaing after extract: ${getRemainingWeights().length}`)
+  }
+
+  return {
+    paramMappings,
+    params: { dense0, dense1, dense2, dense3 }
+  }
+}
\ No newline at end of file

src/faceFeatureExtractor/extractParamsFromWeigthMap.ts

+import * as tf from '@tensorflow/tfjs-core';
+import { disposeUnusedWeightTensors, ParamMapping } from 'tfjs-image-recognition-base';
+
+import { loadParamsFactory } from './loadParamsFactory';
+import { NetParams } from './types';
+
+export function extractParamsFromWeigthMap(
+  weightMap: tf.NamedTensorMap
+): { params: NetParams, paramMappings: ParamMapping[] } {
+
+  const paramMappings: ParamMapping[] = []
+
+  const {
+    extractDenseBlock4Params
+  } = loadParamsFactory(weightMap, paramMappings)
+
+  const params = {
+    dense0: extractDenseBlock4Params('dense0', true),
+    dense1: extractDenseBlock4Params('dense1'),
+    dense2: extractDenseBlock4Params('dense2'),
+    dense3: extractDenseBlock4Params('dense3')
+  }
+
+  disposeUnusedWeightTensors(weightMap, paramMappings)
+
+  return { params, paramMappings }
+}
\ No newline at end of file

src/faceLandmarkNet/extractParamsFromWeigthMapTiny.ts → src/faceFeatureExtractor/extractParamsFromWeigthMapTiny.ts

@@ -11,15 +11,13 @@ export function extractParamsFromWeigthMapTiny(
   const paramMappings: ParamMapping[] = []
 
   const {
-    extractDenseBlock3Params,
-    extractFcParams
+    extractDenseBlock3Params
   } = loadParamsFactory(weightMap, paramMappings)
 
   const params = {
     dense0: extractDenseBlock3Params('dense0', true),
     dense1: extractDenseBlock3Params('dense1'),
-    dense2: extractDenseBlock3Params('dense2'),
-    fc: extractFcParams('fc')
+    dense2: extractDenseBlock3Params('dense2')
   }
 
   disposeUnusedWeightTensors(weightMap, paramMappings)

src/faceLandmarkNet/extractParamsTiny.ts → src/faceFeatureExtractor/extractParamsTiny.ts

@@ -13,14 +13,12 @@ export function extractParamsTiny(weights: Float32Array): { params: TinyNetParam
   } = extractWeightsFactory(weights)
 
   const {
-    extractDenseBlock3Params,
-    extractFCParams
+    extractDenseBlock3Params
   } = extractorsFactory(extractWeights, paramMappings)
 
   const dense0 = extractDenseBlock3Params(3, 32, 'dense0', true)
   const dense1 = extractDenseBlock3Params(32, 64, 'dense1')
   const dense2 = extractDenseBlock3Params(64, 128, 'dense2')
-  const fc = extractFCParams(128, 136, 'fc')
 
   if (getRemainingWeights().length !== 0) {
     throw new Error(`weights remaing after extract: ${getRemainingWeights().length}`)
@@ -28,6 +26,6 @@ export function extractParamsTiny(weights: Float32Array): { params: TinyNetParam
 
   return {
     paramMappings,
-    params: { dense0, dense1, dense2, fc }
+    params: { dense0, dense1, dense2 }
   }
 }
\ No newline at end of file

src/faceLandmarkNet/extractorsFactory.ts → src/faceFeatureExtractor/extractorsFactory.ts

@@ -24,21 +24,6 @@ export function extractorsFactory(extractWeights: ExtractWeightsFunction, paramM
     )
   }
 
-  function extractFCParams(channelsIn: number, channelsOut: number, mappedPrefix: string): FCParams {
-    const weights = tf.tensor2d(extractWeights(channelsIn * channelsOut), [channelsIn, channelsOut])
-    const bias = tf.tensor1d(extractWeights(channelsOut))
-
-    paramMappings.push(
-      { paramPath: `${mappedPrefix}/weights` },
-      { paramPath: `${mappedPrefix}/bias` }
-    )
-
-    return {
-      weights,
-      bias
-    }
-  }
-
   const extractConvParams = extractConvParamsFactory(extractWeights, paramMappings)
 
   function extractDenseBlock3Params(channelsIn: number, channelsOut: number, mappedPrefix: string, isFirstLayer: boolean = false): DenseBlock3Params {
@@ -62,8 +47,7 @@ export function extractorsFactory(extractWeights: ExtractWeightsFunction, paramM
 
   return {
     extractDenseBlock3Params,
-    extractDenseBlock4Params,
-    extractFCParams
+    extractDenseBlock4Params
   }
 
 }
\ No newline at end of file

src/faceFeatureExtractor/index.ts

+export * from './FaceFeatureExtractor';
+export * from './TinyFaceFeatureExtractor';
\ No newline at end of file

src/faceLandmarkNet/loadParamsFactory.ts → src/faceFeatureExtractor/loadParamsFactory.ts

 import * as tf from '@tensorflow/tfjs-core';
 import { extractWeightEntryFactory, ParamMapping } from 'tfjs-image-recognition-base';
-import { ConvParams, FCParams, SeparableConvParams } from 'tfjs-tiny-yolov2';
+import { ConvParams, SeparableConvParams } from 'tfjs-tiny-yolov2';
 
 import { DenseBlock3Params, DenseBlock4Params } from './types';
 
@@ -48,16 +48,8 @@ export function loadParamsFactory(weightMap: any, paramMappings: ParamMapping[])
     return { conv0, conv1, conv2, conv3 }
   }
 
-  function extractFcParams(prefix: string): FCParams {
-    const weights = extractWeightEntry<tf.Tensor2D>(`${prefix}/weights`, 2)
-    const bias = extractWeightEntry<tf.Tensor1D>(`${prefix}/bias`, 1)
-
-    return { weights, bias }
-  }
-
   return {
     extractDenseBlock3Params,
-    extractDenseBlock4Params,
-    extractFcParams
+    extractDenseBlock4Params
   }
 }
\ No newline at end of file

src/faceFeatureExtractor/types.ts

+import * as tf from '@tensorflow/tfjs-core';
+import { ConvParams, SeparableConvParams } from 'tfjs-tiny-yolov2';
+
+export type ConvWithBatchNormParams = BatchNormParams & {
+  filter: tf.Tensor4D
+}
+
+export type BatchNormParams = {
+  mean: tf.Tensor1D
+  variance: tf.Tensor1D
+  scale: tf.Tensor1D
+  offset: tf.Tensor1D
+}
+
+export type SeparableConvWithBatchNormParams = {
+  depthwise: ConvWithBatchNormParams
+  pointwise: ConvWithBatchNormParams
+}
+
+export type DenseBlock3Params = {
+  conv0: SeparableConvParams | ConvParams
+  conv1: SeparableConvParams
+  conv2: SeparableConvParams
+}
+
+export type DenseBlock4Params = DenseBlock3Params & {
+  conv3: SeparableConvParams
+}
+
+export type TinyNetParams = {
+  dense0: DenseBlock3Params
+  dense1: DenseBlock3Params
+  dense2: DenseBlock3Params
+}
+
+export type NetParams = {
+  dense0: DenseBlock4Params
+  dense1: DenseBlock4Params
+  dense2: DenseBlock4Params
+  dense3: DenseBlock4Params
+}
+

src/faceLandmarkNet/FaceLandmark68Net.ts

 import * as tf from '@tensorflow/tfjs-core';
-import { NetInput, normalize } from 'tfjs-image-recognition-base';
-import { ConvParams, SeparableConvParams } from 'tfjs-tiny-yolov2';
+import { NetInput } from 'tfjs-image-recognition-base';
 
-import { depthwiseSeparableConv } from './depthwiseSeparableConv';
+import { FaceFeatureExtractor } from '../faceFeatureExtractor/FaceFeatureExtractor';
 import { extractParams } from './extractParams';
 import { extractParamsFromWeigthMap } from './extractParamsFromWeigthMap';
 import { FaceLandmark68NetBase } from './FaceLandmark68NetBase';
 import { fullyConnectedLayer } from './fullyConnectedLayer';
-import { DenseBlock4Params, NetParams } from './types';
-
-function denseBlock(
-  x: tf.Tensor4D,
-  denseBlockParams: DenseBlock4Params,
-  isFirstLayer: boolean = false
-): tf.Tensor4D {
-  return tf.tidy(() => {
-    const out1 = tf.relu(
-      isFirstLayer
-        ? tf.add(
-          tf.conv2d(x, (denseBlockParams.conv0 as ConvParams).filters, [2, 2], 'same'),
-          denseBlockParams.conv0.bias
-        )
-        : depthwiseSeparableConv(x, denseBlockParams.conv0 as SeparableConvParams, [2, 2])
-    ) as tf.Tensor4D
-    const out2 = depthwiseSeparableConv(out1, denseBlockParams.conv1, [1, 1])
-
-    const in3 = tf.relu(tf.add(out1, out2)) as tf.Tensor4D
-    const out3 = depthwiseSeparableConv(in3, denseBlockParams.conv2, [1, 1])
-
-    const in4 = tf.relu(tf.add(out1, tf.add(out2, out3))) as tf.Tensor4D
-    const out4 = depthwiseSeparableConv(in4, denseBlockParams.conv3, [1, 1])
-
-    return tf.relu(tf.add(out1, tf.add(out2, tf.add(out3, out4)))) as tf.Tensor4D
-  })
-}
+import { NetParams } from './types';
+import { seperateWeightMaps } from './util';
 
 export class FaceLandmark68Net extends FaceLandmark68NetBase<NetParams> {
 
-  constructor() {
+  private static classifierNumFilters: number = 256
+
+  private _faceFeatureExtractor: FaceFeatureExtractor
+
+
+  constructor(faceFeatureExtractor: FaceFeatureExtractor) {
     super('FaceLandmark68Net')
+    this._faceFeatureExtractor = faceFeatureExtractor
+  }
+
+  public get faceFeatureExtractor(): FaceFeatureExtractor {
+    return this._faceFeatureExtractor
   }
 
-  public runNet(input: NetInput): tf.Tensor2D {
+  public runNet(input: NetInput | tf.Tensor4D): tf.Tensor2D {
 
     const { params } = this
 
@@ -49,30 +33,44 @@ export class FaceLandmark68Net extends FaceLandmark68NetBase<NetParams> {
       throw new Error('FaceLandmark68Net - load model before inference')
     }
 
-    return tf.tidy(() => {
-      const batchTensor = input.toBatchTensor(112, true)
-      const meanRgb = [122.782, 117.001, 104.298]
-      const normalized = normalize(batchTensor, meanRgb).div(tf.scalar(255)) as tf.Tensor4D
-
-      let out = denseBlock(normalized, params.dense0, true)
-      out = denseBlock(out, params.dense1)
-      out = denseBlock(out, params.dense2)
-      out = denseBlock(out, params.dense3)
-      out = tf.avgPool(out, [7, 7], [2, 2], 'valid')
+    if (!this.faceFeatureExtractor.isLoaded) {
+      throw new Error('FaceLandmark68Net - load face feature extractor model before inference')
+    }
 
-      return fullyConnectedLayer(out.as2D(out.shape[0], -1), params.fc)
+    return tf.tidy(() => {
+      const bottleneckFeatures = input instanceof NetInput
+        ? this.faceFeatureExtractor.forward(input)
+        : input
+      return fullyConnectedLayer(bottleneckFeatures.as2D(bottleneckFeatures.shape[0], -1), params.fc)
     })
   }
 
+  public dispose(throwOnRedispose: boolean = true) {
+    this.faceFeatureExtractor.dispose(throwOnRedispose)
+    super.dispose(throwOnRedispose)
+  }
+
   protected getDefaultModelName(): string {
     return 'face_landmark_68_model'
   }
 
   protected extractParamsFromWeigthMap(weightMap: tf.NamedTensorMap) {
-    return extractParamsFromWeigthMap(weightMap)
+
+    const { featureExtractorMap, classifierMap } = seperateWeightMaps(weightMap)
+
+    this.faceFeatureExtractor.loadFromWeightMap(featureExtractorMap)
+
+    return extractParamsFromWeigthMap(classifierMap)
   }
 
   protected extractParams(weights: Float32Array) {
-    return extractParams(weights)
+
+    const classifierWeightSize = 136 * FaceLandmark68Net.classifierNumFilters + 136
+
+    const featureExtractorWeights = weights.slice(0, weights.length - classifierWeightSize)
+    const classifierWeights = weights.slice(weights.length - classifierWeightSize)
+
+    this.faceFeatureExtractor.extractWeights(featureExtractorWeights)
+    return extractParams(classifierWeights, FaceLandmark68Net.classifierNumFilters)
   }
 }
\ No newline at end of file

src/faceLandmarkNet/FaceLandmark68TinyNet.ts

 import * as tf from '@tensorflow/tfjs-core';
-import { NetInput, normalize } from 'tfjs-image-recognition-base';
-import { ConvParams, SeparableConvParams } from 'tfjs-tiny-yolov2';
+import { NetInput } from 'tfjs-image-recognition-base';
 
-import { depthwiseSeparableConv } from './depthwiseSeparableConv';
-import { extractParamsTiny } from './extractParamsTiny';
+import { TinyFaceFeatureExtractor } from '../faceFeatureExtractor/TinyFaceFeatureExtractor';
+import { extractParams } from './extractParams';
+import { extractParamsFromWeigthMap } from './extractParamsFromWeigthMap';
 import { FaceLandmark68NetBase } from './FaceLandmark68NetBase';
 import { fullyConnectedLayer } from './fullyConnectedLayer';
-import { extractParamsFromWeigthMapTiny } from './extractParamsFromWeigthMapTiny';
-import { DenseBlock3Params, TinyNetParams } from './types';
-
-function denseBlock(
-  x: tf.Tensor4D,
-  denseBlockParams: DenseBlock3Params,
-  isFirstLayer: boolean = false
-): tf.Tensor4D {
-  return tf.tidy(() => {
-    const out1 = tf.relu(
-      isFirstLayer
-        ? tf.add(
-          tf.conv2d(x, (denseBlockParams.conv0 as ConvParams).filters, [2, 2], 'same'),
-          denseBlockParams.conv0.bias
-        )
-        : depthwiseSeparableConv(x, denseBlockParams.conv0 as SeparableConvParams, [2, 2])
-    ) as tf.Tensor4D
-    const out2 = depthwiseSeparableConv(out1, denseBlockParams.conv1, [1, 1])
-
-    const in3 = tf.relu(tf.add(out1, out2)) as tf.Tensor4D
-    const out3 = depthwiseSeparableConv(in3, denseBlockParams.conv2, [1, 1])
-
-    return tf.relu(tf.add(out1, tf.add(out2, out3))) as tf.Tensor4D
-  })
-}
-
-export class FaceLandmark68TinyNet extends FaceLandmark68NetBase<TinyNetParams> {
-
-  constructor() {
+import { NetParams } from './types';
+import { seperateWeightMaps } from './util';
+
+export class FaceLandmark68TinyNet extends FaceLandmark68NetBase<NetParams> {
+
+  private static classifierNumFilters: number = 128
+
+  private _faceFeatureExtractor: TinyFaceFeatureExtractor
+
+  constructor(faceFeatureExtractor: TinyFaceFeatureExtractor) {
     super('FaceLandmark68TinyNet')
+    this._faceFeatureExtractor = faceFeatureExtractor
   }
 
-  public runNet(input: NetInput): tf.Tensor2D {
+  public get faceFeatureExtractor(): TinyFaceFeatureExtractor {
+    return this._faceFeatureExtractor
+  }
+
+  public runNet(input: NetInput | tf.Tensor4D): tf.Tensor2D {
 
     const { params } = this
 
@@ -46,29 +32,44 @@ export class FaceLandmark68TinyNet extends FaceLandmark68NetBase<TinyNetParams>
       throw new Error('FaceLandmark68TinyNet - load model before inference')
     }
 
-    return tf.tidy(() => {
-      const batchTensor = input.toBatchTensor(112, true)
-      const meanRgb = [122.782, 117.001, 104.298]
-      const normalized = normalize(batchTensor, meanRgb).div(tf.scalar(255)) as tf.Tensor4D
-
-      let out = denseBlock(normalized, params.dense0, true)
-      out = denseBlock(out, params.dense1)
-      out = denseBlock(out, params.dense2)
-      out = tf.avgPool(out, [14, 14], [2, 2], 'valid')
+    if (!this.faceFeatureExtractor.isLoaded) {
+      throw new Error('FaceLandmark68TinyNet - load face feature extractor model before inference')
+    }
 
-      return fullyConnectedLayer(out.as2D(out.shape[0], -1), params.fc)
+    return tf.tidy(() => {
+      const bottleneckFeatures = input instanceof NetInput
+        ? this.faceFeatureExtractor.forward(input)
+        : input
+      return fullyConnectedLayer(bottleneckFeatures.as2D(bottleneckFeatures.shape[0], -1), params.fc)
     })
   }
 
+  public dispose(throwOnRedispose: boolean = true) {
+    this.faceFeatureExtractor.dispose(throwOnRedispose)
+    super.dispose(throwOnRedispose)
+  }
+
   protected getDefaultModelName(): string {
     return 'face_landmark_68_tiny_model'
   }
 
   protected extractParamsFromWeigthMap(weightMap: tf.NamedTensorMap) {
-    return extractParamsFromWeigthMapTiny(weightMap)
+
+    const { featureExtractorMap, classifierMap } = seperateWeightMaps(weightMap)
+
+    this.faceFeatureExtractor.loadFromWeightMap(featureExtractorMap)
+
+    return extractParamsFromWeigthMap(classifierMap)
   }
 
   protected extractParams(weights: Float32Array) {
-    return extractParamsTiny(weights)
+
+    const classifierWeightSize = 136 * FaceLandmark68TinyNet.classifierNumFilters + 136
+
+    const featureExtractorWeights = weights.slice(0, weights.length - classifierWeightSize)
+    const classifierWeights = weights.slice(weights.length - classifierWeightSize)
+
+    this.faceFeatureExtractor.extractWeights(featureExtractorWeights)
+    return extractParams(classifierWeights, FaceLandmark68TinyNet.classifierNumFilters)
   }
 }
\ No newline at end of file

src/faceLandmarkNet/extractParams.ts

+import * as tf from '@tensorflow/tfjs-core';
 import { extractWeightsFactory, ParamMapping } from 'tfjs-image-recognition-base';
+import { FCParams } from 'tfjs-tiny-yolov2';
 
-import { extractorsFactory } from './extractorsFactory';
 import { NetParams } from './types';
 
-export function extractParams(weights: Float32Array): { params: NetParams, paramMappings: ParamMapping[] } {
+export function extractParams(weights: Float32Array, numFilters: number): { params: NetParams, paramMappings: ParamMapping[] } {
 
   const paramMappings: ParamMapping[] = []
 
@@ -12,16 +13,22 @@ export function extractParams(weights: Float32Array): { params: NetParams, param
     getRemainingWeights
   } = extractWeightsFactory(weights)
 
-  const {
-    extractDenseBlock4Params,
-    extractFCParams
-  } = extractorsFactory(extractWeights, paramMappings)
+  function extractFCParams(channelsIn: number, channelsOut: number, mappedPrefix: string): FCParams {
+    const weights = tf.tensor2d(extractWeights(channelsIn * channelsOut), [channelsIn, channelsOut])
+    const bias = tf.tensor1d(extractWeights(channelsOut))
+
+    paramMappings.push(
+      { paramPath: `${mappedPrefix}/weights` },
+      { paramPath: `${mappedPrefix}/bias` }
+    )
+
+    return {
+      weights,
+      bias
+    }
+  }
 
-  const dense0 = extractDenseBlock4Params(3, 32, 'dense0', true)
-  const dense1 = extractDenseBlock4Params(32, 64, 'dense1')
-  const dense2 = extractDenseBlock4Params(64, 128, 'dense2')
-  const dense3 = extractDenseBlock4Params(128, 256, 'dense3')
-  const fc = extractFCParams(256, 136, 'fc')
+  const fc = extractFCParams(numFilters, 136, 'fc')
 
   if (getRemainingWeights().length !== 0) {
     throw new Error(`weights remaing after extract: ${getRemainingWeights().length}`)
@@ -29,6 +36,6 @@ export function extractParams(weights: Float32Array): { params: NetParams, param
 
   return {
     paramMappings,
-    params: { dense0, dense1, dense2, dense3, fc }
+    params: { fc }
   }
 }
\ No newline at end of file

src/faceLandmarkNet/extractParamsFromWeigthMap.ts

 import * as tf from '@tensorflow/tfjs-core';
-import { disposeUnusedWeightTensors, ParamMapping } from 'tfjs-image-recognition-base';
+import { disposeUnusedWeightTensors, extractWeightEntryFactory, ParamMapping } from 'tfjs-image-recognition-base';
+import { FCParams } from 'tfjs-tiny-yolov2';
 
-import { loadParamsFactory } from './loadParamsFactory';
 import { NetParams } from './types';
 
 export function extractParamsFromWeigthMap(
@@ -10,16 +10,15 @@ export function extractParamsFromWeigthMap(
 
   const paramMappings: ParamMapping[] = []
 
-  const {
-    extractDenseBlock4Params,
-    extractFcParams
-  } = loadParamsFactory(weightMap, paramMappings)
+  const extractWeightEntry = extractWeightEntryFactory(weightMap, paramMappings)
+
+  function extractFcParams(prefix: string): FCParams {
+    const weights = extractWeightEntry<tf.Tensor2D>(`${prefix}/weights`, 2)
+    const bias = extractWeightEntry<tf.Tensor1D>(`${prefix}/bias`, 1)
+    return { weights, bias }
+  }
 
   const params = {
-    dense0: extractDenseBlock4Params('dense0', true),
-    dense1: extractDenseBlock4Params('dense1'),
-    dense2: extractDenseBlock4Params('dense2'),
-    dense3: extractDenseBlock4Params('dense3'),
     fc: extractFcParams('fc')
   }
 

src/faceLandmarkNet/index.ts

@@ -3,10 +3,4 @@ import { FaceLandmark68Net } from './FaceLandmark68Net';
 export * from './FaceLandmark68Net';
 export * from './FaceLandmark68TinyNet';
 
-export class FaceLandmarkNet extends FaceLandmark68Net {}
-
-export function createFaceLandmarkNet(weights: Float32Array) {
-  const net = new FaceLandmarkNet()
-  net.extractWeights(weights)
-  return net
-}
\ No newline at end of file
+export class FaceLandmarkNet extends FaceLandmark68Net {}
\ No newline at end of file

src/faceLandmarkNet/types.ts

-import * as tf from '@tensorflow/tfjs-core';
-import { ConvParams, FCParams, SeparableConvParams } from 'tfjs-tiny-yolov2';
-
-export type ConvWithBatchNormParams = BatchNormParams & {
-  filter: tf.Tensor4D
-}
-
-export type BatchNormParams = {
-  mean: tf.Tensor1D
-  variance: tf.Tensor1D
-  scale: tf.Tensor1D
-  offset: tf.Tensor1D
-}
-
-export type SeparableConvWithBatchNormParams = {
-  depthwise: ConvWithBatchNormParams
-  pointwise: ConvWithBatchNormParams
-}
-
-export declare type FCWithBatchNormParams = BatchNormParams & {
-    weights: tf.Tensor2D
-}
-
-export type DenseBlock3Params = {
-  conv0: SeparableConvParams | ConvParams
-  conv1: SeparableConvParams
-  conv2: SeparableConvParams
-}
-
-export type DenseBlock4Params = DenseBlock3Params & {
-  conv3: SeparableConvParams
-}
-
-export type TinyNetParams = {
-  dense0: DenseBlock3Params
-  dense1: DenseBlock3Params
-  dense2: DenseBlock3Params
-  fc: FCParams
-}
+import { FCParams } from 'tfjs-tiny-yolov2';
 
 export type NetParams = {
-  dense0: DenseBlock4Params
-  dense1: DenseBlock4Params
-  dense2: DenseBlock4Params
-  dense3: DenseBlock4Params
   fc: FCParams
 }
 

src/faceLandmarkNet/util.ts

+import * as tf from '@tensorflow/tfjs-core';
+
+export function seperateWeightMaps(weightMap: tf.NamedTensorMap) {
+
+  const featureExtractorMap: tf.NamedTensorMap = {}
+  const classifierMap: tf.NamedTensorMap = {}
+
+  Object.keys(weightMap).forEach(key => {
+    const map = key.startsWith('fc') ? classifierMap : featureExtractorMap
+    map[key] = weightMap[key]
+  })
+
+  return { featureExtractorMap, classifierMap }
+
+}
\ No newline at end of file

src/globalApi/nets.ts

@@ -4,6 +4,7 @@ import { ITinyYolov2Options } from 'tfjs-tiny-yolov2';
 import { FaceDetection } from '../classes/FaceDetection';
 import { FaceLandmarks5 } from '../classes/FaceLandmarks5';
 import { FaceLandmarks68 } from '../classes/FaceLandmarks68';
+import { FaceFeatureExtractor, TinyFaceFeatureExtractor } from '../faceFeatureExtractor';
 import { FaceLandmark68Net } from '../faceLandmarkNet/FaceLandmark68Net';
 import { FaceLandmark68TinyNet } from '../faceLandmarkNet/FaceLandmark68TinyNet';
 import { FaceRecognitionNet } from '../faceRecognitionNet/FaceRecognitionNet';
@@ -17,13 +18,16 @@ import { TinyFaceDetector } from '../tinyFaceDetector/TinyFaceDetector';
 import { TinyFaceDetectorOptions } from '../tinyFaceDetector/TinyFaceDetectorOptions';
 import { TinyYolov2 } from '../tinyYolov2/TinyYolov2';
 
+const faceFeatureExtractor = new FaceFeatureExtractor()
+const tinyFaceFeatureExtractor = new TinyFaceFeatureExtractor()
+
 export const nets = {
   ssdMobilenetv1: new SsdMobilenetv1(),
   tinyFaceDetector: new TinyFaceDetector(),
   tinyYolov2: new TinyYolov2(),
   mtcnn: new Mtcnn(),
-  faceLandmark68Net: new FaceLandmark68Net(),
-  faceLandmark68TinyNet: new FaceLandmark68TinyNet(),
+  faceLandmark68Net: new FaceLandmark68Net(faceFeatureExtractor),
+  faceLandmark68TinyNet: new FaceLandmark68TinyNet(tinyFaceFeatureExtractor),
   faceRecognitionNet: new FaceRecognitionNet()
 }