implemented faceLandmarkNet

justadudewhohacks · justadudewhohacks · commit 72d280cf59f4 · 2018-06-16T02:03:29.000+02:00
diff --git a/src/commons/convLayer.ts b/src/commons/convLayer.ts
@@ -0,0 +1,16 @@
+import * as tf from '@tensorflow/tfjs-core';
+
+import { ConvParams } from './types';
+
+export function convLayer(
+  x: tf.Tensor4D,
+  params: ConvParams,
+  padding: 'valid' | 'same' = 'same'
+): tf.Tensor4D {
+  return tf.tidy(() =>
+    tf.add(
+      tf.conv2d(x, params.filters, [1, 1], padding),
+      params.bias
+    )
+  )
+}
diff --git a/src/faceDetectionNet/boxPredictionLayer.ts b/src/faceDetectionNet/boxPredictionLayer.ts
@@ -1,19 +1,8 @@
 import * as tf from '@tensorflow/tfjs-core';
 
-import { ConvParams } from '../commons/types';
+import { convLayer } from '../commons/convLayer';
 import { FaceDetectionNet } from './types';
 
-function convWithBias(
-  x: tf.Tensor4D,
-  params: ConvParams
-) {
-  return tf.tidy(() =>
-    tf.add(
-      tf.conv2d(x, params.filters, [1, 1], 'same'),
-      params.bias
-    )
-  )
-}
 
 export function boxPredictionLayer(
   x: tf.Tensor4D,
@@ -24,11 +13,11 @@ export function boxPredictionLayer(
     const batchSize = x.shape[0]
 
     const boxPredictionEncoding = tf.reshape(
-      convWithBias(x, params.box_encoding_predictor_params),
+      convLayer(x, params.box_encoding_predictor_params),
       [batchSize, -1, 1, 4]
     )
     const classPrediction = tf.reshape(
-      convWithBias(x, params.class_predictor_params),
+      convLayer(x, params.class_predictor_params),
       [batchSize, -1, 3]
     )
 
diff --git a/src/faceLandmarkNet/fullyConnectedLayer.ts b/src/faceLandmarkNet/fullyConnectedLayer.ts
@@ -0,0 +1,15 @@
+import * as tf from '@tensorflow/tfjs-core';
+
+import { FaceLandmarkNet } from './types';
+
+export function fullyConnectedLayer(
+  x: tf.Tensor2D,
+  params: FaceLandmarkNet.FCParams
+): tf.Tensor2D {
+  return tf.tidy(() =>
+    tf.add(
+      tf.matMul(x, params.weights),
+      params.bias
+    )
+  )
+}
diff --git a/src/faceLandmarkNet/index.ts b/src/faceLandmarkNet/index.ts
@@ -5,6 +5,8 @@ import { NetInput } from '../NetInput';
 import { padToSquare } from '../padToSquare';
 import { TNetInput } from '../types';
 import { extractParams } from './extractParams';
+import { convLayer } from '../commons/convLayer';
+import { fullyConnectedLayer } from './fullyConnectedLayer';
 
 export function faceLandmarkNet(weights: Float32Array) {
   const params = extractParams(weights)
@@ -18,17 +20,22 @@ export function faceLandmarkNet(weights: Float32Array) {
         x = tf.image.resizeBilinear(x, [128, 128])
       }
 
-      // pool 1
-      tf.maxPool(x, [2, 2], [2, 2], 'valid')
-      // pool 2
-      tf.maxPool(x, [2, 2], [2, 2], 'valid')
-      // pool 3
-      tf.maxPool(x, [2, 2], [2, 2], 'valid')
-      // pool 4
-      tf.maxPool(x, [2, 2], [1, 1], 'valid')
-      // TODO
+      let out = convLayer(x, params.conv0_params, 'valid')
+      out = tf.maxPool(out, [2, 2], [2, 2], 'valid')
+      out = convLayer(out, params.conv1_params, 'valid')
+      out = convLayer(out, params.conv2_params, 'valid')
+      out = tf.maxPool(out, [2, 2], [2, 2], 'valid')
+      out = convLayer(out, params.conv3_params, 'valid')
+      out = convLayer(out, params.conv4_params, 'valid')
+      out = tf.maxPool(out, [2, 2], [2, 2], 'valid')
+      out = convLayer(out, params.conv5_params, 'valid')
+      out = convLayer(out, params.conv6_params, 'valid')
+      out = tf.maxPool(out, [2, 2], [1, 1], 'valid')
+      out = convLayer(out, params.conv7_params, 'valid')
+      const fc0 = fullyConnectedLayer(out.as2D(out.shape[0], -1), params.fc0_params)
+      const fc1 = fullyConnectedLayer(fc0, params.fc1_params)
 
-      return x
+      return fc1
     })
   }
 
diff --git a/src/faceRecognitionNet/convLayer.ts b/src/faceRecognitionNet/convLayer.ts
@@ -7,26 +7,26 @@ import { FaceRecognitionNet } from './types';
 function convLayer(
   x: tf.Tensor4D,
   params: FaceRecognitionNet.ConvLayerParams,
-  stride: number,
+  strides: [number, number],
   withRelu: boolean,
   padding: 'valid' | 'same' = 'same'
 ): tf.Tensor4D {
   const { filters, bias } = params.conv
 
-  let out = tf.conv2d(x, filters, [stride, stride], padding)
+  let out = tf.conv2d(x, filters, strides, padding)
   out = tf.add(out, bias)
   out = scale(out, params.scale)
   return withRelu ? tf.relu(out) : out
 }
 
 export function conv(x: tf.Tensor4D, params: FaceRecognitionNet.ConvLayerParams) {
-  return convLayer(x, params, 1, true)
+  return convLayer(x, params, [1, 1], true)
 }
 
 export function convNoRelu(x: tf.Tensor4D, params: FaceRecognitionNet.ConvLayerParams) {
-  return convLayer(x, params, 1, false)
+  return convLayer(x, params, [1, 1], false)
 }
 
 export function convDown(x: tf.Tensor4D, params: FaceRecognitionNet.ConvLayerParams) {
-  return convLayer(x, params, 2, true, 'valid')
+  return convLayer(x, params, [2, 2], true, 'valid')
 }
