add centering option to padToSquare and center input face of face recognition net

justadudewhohacks · justadudewhohacks · commit 2b127e43e1fd · 2018-06-10T15:20:25.000+02:00
diff --git a/src/extractFaceTensors.ts b/src/extractFaceTensors.ts
@@ -2,7 +2,7 @@ import * as tf from '@tensorflow/tfjs-core';
 
 import { FaceDetectionResult } from './faceDetectionNet/FaceDetectionResult';
 import { NetInput } from './NetInput';
-import { getImageTensor } from './transformInputs';
+import { getImageTensor } from './getImageTensor';
 import { TNetInput } from './types';
 
 /**
diff --git a/src/faceDetectionNet/index.ts b/src/faceDetectionNet/index.ts
@@ -1,7 +1,8 @@
 import * as tf from '@tensorflow/tfjs-core';
 
+import { getImageTensor } from '../getImageTensor';
 import { NetInput } from '../NetInput';
-import { getImageTensor, padToSquare } from '../transformInputs';
+import { padToSquare } from '../padToSquare';
 import { TNetInput } from '../types';
 import { extractParams } from './extractParams';
 import { FaceDetectionResult } from './FaceDetectionResult';
@@ -49,7 +50,7 @@ export function faceDetectionNet(weights: Float32Array) {
     } = tf.tidy(() => {
 
       let imgTensor = getImageTensor(input)
-      const [_, height, width] = imgTensor.shape
+      const [height, width] = imgTensor.shape.slice(1)
 
       imgTensor = padToSquare(imgTensor)
       paddedHeightRelative = imgTensor.shape[1] / height
diff --git a/src/faceRecognitionNet/index.ts b/src/faceRecognitionNet/index.ts
@@ -1,7 +1,8 @@
 import * as tf from '@tensorflow/tfjs-core';
 
+import { getImageTensor } from '../getImageTensor';
 import { NetInput } from '../NetInput';
-import { getImageTensor, padToSquare } from '../transformInputs';
+import { padToSquare } from '../padToSquare';
 import { TNetInput } from '../types';
 import { convDown } from './convLayer';
 import { extractParams } from './extractParams';
@@ -14,8 +15,7 @@ export function faceRecognitionNet(weights: Float32Array) {
   function forward(input: tf.Tensor | NetInput | TNetInput) {
     return tf.tidy(() => {
 
-      // TODO pad on both sides, to keep face centered
-      let x = padToSquare(getImageTensor(input))
+      let x = padToSquare(getImageTensor(input), true)
       // work with 150 x 150 sized face images
       if (x.shape[1] !== 150 || x.shape[2] !== 150) {
         x = tf.image.resizeBilinear(x, [150, 150])
diff --git a/src/getImageTensor.ts b/src/getImageTensor.ts
@@ -3,23 +3,6 @@ import * as tf from '@tensorflow/tfjs-core';
 import { NetInput } from './NetInput';
 import { TNetInput } from './types';
 
-export function padToSquare(imgTensor: tf.Tensor4D): tf.Tensor4D {
-  return tf.tidy(() => {
-
-    const [_, height, width] = imgTensor.shape
-    if (height === width) {
-      return imgTensor
-    }
-
-    if (height > width) {
-      const pad = tf.fill([1, height, height - width, 3], 0) as tf.Tensor4D
-      return tf.concat([imgTensor, pad], 2)
-    }
-    const pad = tf.fill([1, width - height, width, 3], 0) as tf.Tensor4D
-    return tf.concat([imgTensor, pad], 1)
-  })
-}
-
 export function getImageTensor(input: tf.Tensor | NetInput | TNetInput): tf.Tensor4D {
   return tf.tidy(() => {
     if (input instanceof tf.Tensor) {
diff --git a/src/index.ts b/src/index.ts
@@ -1,15 +1,18 @@
+import * as tf from '@tensorflow/tfjs-core';
+
 import { euclideanDistance } from './euclideanDistance';
 import { faceDetectionNet } from './faceDetectionNet';
 import { faceRecognitionNet } from './faceRecognitionNet';
 import { NetInput } from './NetInput';
-import * as tf from '@tensorflow/tfjs-core';
+import { padToSquare } from './padToSquare';
 
 export {
   euclideanDistance,
   faceDetectionNet,
   faceRecognitionNet,
   NetInput,
-  tf
+  tf,
+  padToSquare
 }
 
 export * from './extractFaces'
diff --git a/src/padToSquare.ts b/src/padToSquare.ts
@@ -0,0 +1,30 @@
+import * as tf from '@tensorflow/tfjs-core';
+
+/**
+ * Pads the smaller dimension of an image tensor with zeros, such that width === height.
+ *
+ * @param imgTensor The image tensor.
+ * @param isCenterImage (optional, default: false) If true, add padding on both sides of the image, such that the image
+ * @returns The padded tensor with width === height.
+ */
+export function padToSquare(
+  imgTensor: tf.Tensor4D,
+  isCenterImage: boolean = false
+): tf.Tensor4D {
+  return tf.tidy(() => {
+
+    const [height, width] = imgTensor.shape.slice(1)
+    if (height === width) {
+      return imgTensor
+    }
+
+    const paddingAmount = Math.floor(Math.abs(height - width) * (isCenterImage ? 0.5 : 1))
+    const paddingAxis = height > width ? 2 : 1
+    const paddingTensorShape = imgTensor.shape.slice() as [number, number, number, number]
+    paddingTensorShape[paddingAxis] = paddingAmount
+
+    const tensorsToStack = (isCenterImage ? [tf.fill(paddingTensorShape, 0)] : [])
+      .concat([imgTensor,  tf.fill(paddingTensorShape, 0)]) as tf.Tensor4D[]
+    return tf.concat(tensorsToStack, paddingAxis)
+  })
+}
