init face expression model training script

src/faceExpressionNet/FaceExpressionNet.ts

+import * as tf from '@tensorflow/tfjs-core';
+import { NetInput, TNetInput, toNetInput } from 'tfjs-image-recognition-base';
+
 import { FaceFeatureExtractor } from '../faceFeatureExtractor/FaceFeatureExtractor';
 import { FaceFeatureExtractorParams } from '../faceFeatureExtractor/types';
 import { FaceProcessor } from '../faceProcessor/FaceProcessor';
+import { EmotionLabels } from './types';
 
 export class FaceExpressionNet extends FaceProcessor<FaceFeatureExtractorParams> {
 
-  constructor(faceFeatureExtractor: FaceFeatureExtractor) {
+  public static getEmotionLabel(emotion: string) {
+    const label = EmotionLabels[emotion.toUpperCase()]
+
+    if (typeof label !== 'number') {
+      throw new Error(`getEmotionLabel - no label for emotion: ${emotion}`)
+    }
+
+    return label
+  }
+
+  public static decodeEmotions(probabilities: number[] | Float32Array) {
+    if (probabilities.length !== 7) {
+      throw new Error(`decodeEmotions - expected probabilities.length to be 7, have: ${probabilities.length}`)
+    }
+    return Object.keys(EmotionLabels).map(label => ({ label, probability: probabilities[EmotionLabels[label]] }))
+  }
+
+  constructor(faceFeatureExtractor: FaceFeatureExtractor = new FaceFeatureExtractor()) {
     super('FaceExpressionNet', faceFeatureExtractor)
   }
 
+  public runNet(input: NetInput | tf.Tensor4D): tf.Tensor2D {
+    return tf.tidy(() => {
+      const out = super.runNet(input)
+      return tf.softmax(out)
+    })
+  }
+
+  public forwardInput(input: NetInput | tf.Tensor4D): tf.Tensor2D {
+    return tf.tidy(() => this.runNet(input))
+  }
+
+  public async forward(input: TNetInput): Promise<tf.Tensor2D> {
+    return this.forwardInput(await toNetInput(input))
+  }
+
+  public async predictExpressions(input: TNetInput) {
+    const out = await this.forward(input)
+    const probabilitesByBatch = await Promise.all(tf.unstack(out).map(t => t.data()))
+    out.dispose()
+
+    return probabilitesByBatch.map(propablities => FaceExpressionNet.decodeEmotions(propablities as Float32Array))
+  }
+
   public dispose(throwOnRedispose: boolean = true) {
     this.faceFeatureExtractor.dispose(throwOnRedispose)
     super.dispose(throwOnRedispose)

src/faceExpressionNet/index.ts

-export * from './FaceExpressionNet';
\ No newline at end of file
+export * from './FaceExpressionNet';
+export * from './types';
\ No newline at end of file

src/faceExpressionNet/types.ts

+export enum EmotionLabels {
+  NEUTRAL = 0,
+  HAPPY = 1,
+  SAD = 2,
+  ANGRY = 3,
+  FEARFUL = 4,
+  DISGUSTED = 5,
+  SURPRISED = 6
+}
\ No newline at end of file

src/faceFeatureExtractor/FaceFeatureExtractor.ts

 import * as tf from '@tensorflow/tfjs-core';
-import { NetInput, NeuralNetwork, normalize } from 'tfjs-image-recognition-base';
+import { NetInput, NeuralNetwork, normalize, TNetInput, toNetInput } 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, IFaceFeatureExtractor, FaceFeatureExtractorParams } from './types';
+import { DenseBlock4Params, FaceFeatureExtractorParams, IFaceFeatureExtractor } from './types';
 
 function denseBlock(
   x: tf.Tensor4D,
@@ -39,7 +39,7 @@ export class FaceFeatureExtractor extends NeuralNetwork<FaceFeatureExtractorPara
     super('FaceFeatureExtractor')
   }
 
-  public forward(input: NetInput): tf.Tensor4D {
+  public forwardInput(input: NetInput): tf.Tensor4D {
 
     const { params } = this
 
@@ -62,6 +62,10 @@ export class FaceFeatureExtractor extends NeuralNetwork<FaceFeatureExtractorPara
     })
   }
 
+  public async forward(input: TNetInput): Promise<tf.Tensor4D> {
+    return this.forwardInput(await toNetInput(input))
+  }
+
   protected getDefaultModelName(): string {
     return 'face_feature_extractor_model'
   }

src/faceFeatureExtractor/TinyFaceFeatureExtractor.ts

 import * as tf from '@tensorflow/tfjs-core';
-import { NetInput, NeuralNetwork, normalize } from 'tfjs-image-recognition-base';
+import { NetInput, NeuralNetwork, normalize, TNetInput, toNetInput } from 'tfjs-image-recognition-base';
 import { ConvParams, SeparableConvParams } from 'tfjs-tiny-yolov2';
 
 import { depthwiseSeparableConv } from './depthwiseSeparableConv';
@@ -36,7 +36,7 @@ export class TinyFaceFeatureExtractor extends NeuralNetwork<TinyFaceFeatureExtra
     super('TinyFaceFeatureExtractor')
   }
 
-  public forward(input: NetInput): tf.Tensor4D {
+  public forwardInput(input: NetInput): tf.Tensor4D {
 
     const { params } = this
 
@@ -58,6 +58,10 @@ export class TinyFaceFeatureExtractor extends NeuralNetwork<TinyFaceFeatureExtra
     })
   }
 
+  public async forward(input: TNetInput): Promise<tf.Tensor4D> {
+    return this.forwardInput(await toNetInput(input))
+  }
+
   protected getDefaultModelName(): string {
     return 'face_landmark_68_tiny_model'
   }

src/faceFeatureExtractor/types.ts

 import * as tf from '@tensorflow/tfjs-core';
-import { NetInput, NeuralNetwork } from 'tfjs-image-recognition-base';
+import { NetInput, NeuralNetwork, TNetInput } from 'tfjs-image-recognition-base';
 import { ConvParams, SeparableConvParams } from 'tfjs-tiny-yolov2';
 
 export type ConvWithBatchNormParams = BatchNormParams & {
@@ -42,5 +42,6 @@ export type FaceFeatureExtractorParams = {
 }
 
 export interface IFaceFeatureExtractor<TNetParams extends TinyFaceFeatureExtractorParams | FaceFeatureExtractorParams> extends NeuralNetwork<TNetParams> {
-  forward(input: NetInput): tf.Tensor4D
+  forwardInput(input: NetInput): tf.Tensor4D
+  forward(input: TNetInput): Promise<tf.Tensor4D>
 }
\ No newline at end of file

src/faceProcessor/FaceProcessor.ts

@@ -42,7 +42,7 @@ export abstract class FaceProcessor<
 
     return tf.tidy(() => {
       const bottleneckFeatures = input instanceof NetInput
-        ? this.faceFeatureExtractor.forward(input)
+        ? this.faceFeatureExtractor.forwardInput(input)
         : input
       return fullyConnectedLayer(bottleneckFeatures.as2D(bottleneckFeatures.shape[0], -1), params.fc)
     })
@@ -53,6 +53,16 @@ export abstract class FaceProcessor<
     super.dispose(throwOnRedispose)
   }
 
+  public loadClassifierParams(weights: Float32Array) {
+    const { params, paramMappings } = this.extractClassifierParams(weights)
+    this._params = params
+    this._paramMappings = paramMappings
+  }
+
+  public extractClassifierParams(weights: Float32Array) {
+    return extractParams(weights, this.getClassifierChannelsIn(), this.getClassifierChannelsOut())
+  }
+
   protected extractParamsFromWeigthMap(weightMap: tf.NamedTensorMap) {
 
     const { featureExtractorMap, classifierMap } = seperateWeightMaps(weightMap)
@@ -72,6 +82,6 @@ export abstract class FaceProcessor<
     const classifierWeights = weights.slice(weights.length - classifierWeightSize)
 
     this.faceFeatureExtractor.extractWeights(featureExtractorWeights)
-    return extractParams(classifierWeights, cIn, cOut)
+    return this.extractClassifierParams(classifierWeights)
   }
 }
\ No newline at end of file

tools/train/faceExpressions/public/initWeights.html

+<!DOCTYPE html>
+<html>
+<head>
+  <script src="https://cdn.jsdelivr.net/npm/@tensorflow/tfjs@0.12.0"> </script>
+  <script src="../../node_modules/file-saver/FileSaver.js"></script>
+</head>
+<body>
+
+  <div class="row side-by-side">
+    <button
+      class="waves-effect waves-light btn"
+      onclick="save()"
+    >
+      Save
+    </button>
+  </div>
+
+  <script>
+    function toDataArray(tensor) {
+      return Array.from(tensor.dataSync())
+    }
+
+    function flatten(arrs) {
+      return arrs.reduce((flat, arr) => flat.concat(arr))
+    }
+
+    function initWeights(initializer) {
+      const numOutputFilters = 256
+      const outputSize = 7
+
+      const weights = toDataArray(initializer.apply([1, 1, numOutputFilters, 7]))
+      const bias = toDataArray(tf.zeros([outputSize]))
+
+      return new Float32Array(weights.concat(bias))
+    }
+
+    function save() {
+      const initialWeights = initWeights(
+        tf.initializers.glorotNormal()
+      )
+      saveAs(new Blob([initialWeights]), `initial_glorot.weights`)
+    }
+
+  </script>
+</body>
+</html>
\ No newline at end of file

tools/train/faceExpressions/public/js/commons.js

+function getImageUrl({ db, label, img }) {
+  if (db === 'kaggle') {
+    return `kaggle-face-expressions-db/${label}/${img}`
+  }
+
+  const id = parseInt(img.replace('.png'))
+  const dirNr = Math.floor(id / 5000)
+  return `cropped-faces/jpgs${dirNr + 1}/${img}`
+}
+

tools/train/faceExpressions/public/js/fetchData.js

-export function getImageUrl({ db, img }) {
-  if (db === 'kaggle') {
-    return `kaggle-face-expressions-db/${label}/${id}.png`
-  }
-
-  const dirNr = Math.floor(id / NUM_IMAGES_PER_DIR)
-  return `cropped-faces/jpgs${dirNr + 1}/${id}.jpg`
-}
\ No newline at end of file

tools/train/faceExpressions/public/trainClassifier.html

+<!DOCTYPE html>
+<html>
+<head>
+  <script src="face-api.js"></script>
+  <script src="FileSaver.js"></script>
+  <script src="js/commons.js"></script>
+</head>
+<body>
+  <div id="container"></div>
+
+  <script>
+    tf = faceapi.tf
+
+    // load the FaceLandmark68Net and use it's feature extractor since we only
+    // train the output layer of the FaceExpressionNet
+    const dummyLandmarkNet = new faceapi.FaceLandmark68Net()
+    window.net = new faceapi.FaceExpressionNet(dummyLandmarkNet.faceFeatureExtractor)
+
+    // uri to weights file of last checkpoint
+    const modelCheckpoint = 'tmp/initial_classifier.weights'
+    const startEpoch = 0
+
+    const learningRate = 0.001 // 0.001
+    window.optimizer = tf.train.adam(learningRate, 0.9, 0.999, 1e-8)
+
+    window.saveEveryNthSample = Infinity
+
+    window.batchSize = 16
+    //window.batchSize = 32
+
+    window.lossValues = []
+
+    window.iterDelay = 0
+    window.withLogging = true
+
+    const log = (str, ...args) => console.log(`[${[(new Date()).toTimeString().substr(0, 8)]}] ${str || ''}`, ...args)
+
+    function saveWeights(net, filename = 'train_tmp') {
+      saveAs(new Blob([net.serializeParams()]), filename)
+    }
+
+    async function delay(ms) {
+      return new Promise(res => setTimeout(res, ms))
+    }
+
+    async function load() {
+      window.trainData = await faceapi.fetchJson('trainData.json')
+      await dummyLandmarkNet.load('/')
+
+      // fetch the actual output layer weights
+      const classifierWeights = await faceapi.fetchNetWeights(modelCheckpoint)
+      await window.net.loadClassifierParams(classifierWeights)
+      window.net.variable()
+    }
+
+    async function onEpochDone(epoch) {
+      saveWeights(window.net, `face_expression_model_${epoch}.weights`)
+
+      const loss = window.lossValues[epoch]
+      saveAs(new Blob([JSON.stringify({ loss, avgLoss: loss / window.trainIds.length })]), `face_expression_model_${epoch}.json`)
+
+    }
+
+    function prepareDataForEpoch() {
+      return faceapi.shuffleArray(
+        Object.keys(window.trainData).map(label => {
+          let dataForLabel = window.trainData[label].map(data => ({ ...data, label }))
+          // since train data for "disgusted" have less than 2000 samples
+          // use some data twice to ensure an even distribution
+          dataForLabel = label === 'disgusted'
+            ? faceapi.shuffleArray(dataForLabel.concat(dataForLabel).concat(dataForLabel)).slice(0, 2000)
+            : dataForLabel
+          return dataForLabel
+        }).reduce((flat, arr) => arr.concat(flat))
+      )
+    }
+
+    function getLabelOneHotVector(emotion) {
+      const label = faceapi.FaceExpressionNet.getEmotionLabel(emotion)
+      return Array(7).fill(0).map((_, i) => i === label ? 1 : 0)
+    }
+
+    async function train() {
+      await load()
+
+      for (let epoch = startEpoch; epoch < Infinity; epoch++) {
+
+        if (epoch !== startEpoch) {
+          // ugly hack to wait for loss datas for that epoch to be resolved
+          setTimeout(() => onEpochDone(epoch - 1), 10000)
+        }
+        window.lossValues[epoch] = 0
+
+        const shuffledInputs = prepareDataForEpoch()
+          console.log(shuffledInputs)
+
+        for (let dataIdx = 0; dataIdx < shuffledInputs.length; dataIdx += window.batchSize) {
+          const tsIter = Date.now()
+
+          const batchData = shuffledInputs.slice(dataIdx, dataIdx + window.batchSize)
+          const bImages = await Promise.all(
+            batchData
+              .map(data => getImageUrl(data))
+              .map(imgUrl => faceapi.fetchImage(imgUrl))
+          )
+          const bOneHotVectors = batchData
+            .map(data => getLabelOneHotVector(data.label))
+
+          let tsBackward = Date.now()
+          let tsForward = Date.now()
+          const bottleneckFeatures = await window.net.faceFeatureExtractor.forward(bImages)
+          tsForward = Date.now() - tsForward
+
+          const loss = optimizer.minimize(() => {
+            tsBackward = Date.now()
+            const labels = tf.tensor2d(bOneHotVectors)
+            const out = window.net.forwardInput(bottleneckFeatures)
+
+            const loss = tf.losses.softmaxCrossEntropy(
+              labels,
+              out,
+              tf.Reduction.MEAN
+            )
+
+            return loss
+          }, true)
+          tsBackward = Date.now() - tsBackward
+
+          bottleneckFeatures.dispose()
+
+          // start next iteration without waiting for loss data
+
+          loss.data().then(data => {
+            const lossValue = data[0]
+            window.lossValues[epoch] += lossValue
+            window.withLogging && log(`epoch ${epoch}, dataIdx ${dataIdx} - loss: ${lossValue}, ( ${window.lossValues[epoch]})`)
+            loss.dispose()
+          })
+
+          window.withLogging && log(`epoch ${epoch}, dataIdx ${dataIdx} - forward: ${tsForward} ms, backprop: ${tsBackward} ms, iter: ${Date.now() - tsIter} ms`)
+          if (window.logsilly)  {
+            log(`fetch: ${tsFetch} ms, pts: ${tsFetchPts} ms, jpgs: ${tsFetchJpgs} ms, bufferToImage: ${tsBufferToImage} ms`)
+          }
+          if (window.iterDelay) {
+            await delay(window.iterDelay)
+          } else {
+            await tf.nextFrame()
+          }
+        }
+      }
+    }
+
+  </script>
+</body>
+</html>
\ No newline at end of file

tools/train/faceExpressions/server.js

@@ -10,7 +10,11 @@ const publicDir = path.join(__dirname, './public')
 app.use(express.static(publicDir))
 app.use(express.static(path.join(__dirname, '../shared')))
 app.use(express.static(path.join(__dirname, '../node_modules/file-saver')))
-app.use(express.static(path.join(__dirname, '../../../examples/public')))
 app.use(express.static(path.join(__dirname, '../../../weights')))
 app.use(express.static(path.join(__dirname, '../../../dist')))
-app.use(express.static(path.resolve(process.env.DATA_PATH)))
\ No newline at end of file
+app.use(express.static(path.resolve(process.env.DATA_PATH)))
+
+app.get('/', (req, res) => res.redirect('/train'))
+app.get('/train', (req, res) => res.sendFile(path.join(publicDir, 'trainClassifier.html')))
+
+app.listen(8000, () => console.log('Listening on port 8000!'))
\ No newline at end of file

tools/train/faceLandmarks/public/initWeightsDensenet.html

@@ -2,7 +2,7 @@
 <html>
 <head>
   <script src="https://cdn.jsdelivr.net/npm/@tensorflow/tfjs@0.12.0"> </script>
-  <script src="FileSaver.js"></script>
+  <script src="../../node_modules/file-saver/FileSaver.js"></script>
 </head>
 <body>
 

tools/train/faceLandmarks/server.js

@@ -8,9 +8,9 @@ const app = express()
 
 const publicDir = path.join(__dirname, './public')
 app.use(express.static(publicDir))
+
 app.use(express.static(path.join(__dirname, '../shared')))
 app.use(express.static(path.join(__dirname, '../node_modules/file-saver')))
-app.use(express.static(path.join(__dirname, '../../../examples/public')))
 app.use(express.static(path.join(__dirname, '../../../weights')))
 app.use(express.static(path.join(__dirname, '../../../dist')))