added some tooling to retrain the face landmark model

tools/train/package.json

+{
+  "scripts": {
+    "start": "node server.js"
+  },
+  "author": "justadudewhohacks",
+  "license": "MIT",
+  "dependencies": {
+    "express": "^4.16.3",
+    "file-saver": "^1.3.8"
+  }
+}

tools/train/public/faceLandmarksTrain.js

+async function promiseSequential(promises) {
+  const curr = promises[0]
+  if (!curr) {
+    return
+  }
+
+  await curr()
+  return promiseSequential(promises.slice(1))
+}
+
+async function trainStep(batchCreators) {
+  await promiseSequential(batchCreators.map((batchCreator, dataIdx) => async () => {
+
+    const { batchInput, landmarksBatchTensor } = await batchCreator()
+    let ts = Date.now()
+    const cost = optimizer.minimize(() => {
+      const out = window.trainNet.forwardInput(batchInput.managed())
+      const loss = lossFunction(
+        landmarksBatchTensor,
+        out
+      )
+      return loss
+    }, true)
+
+    ts = Date.now() - ts
+    console.log(`loss[${dataIdx}]: ${await cost.data()}, ${ts} ms (${ts / batchInput.batchSize} ms / batch element)`)
+
+    landmarksBatchTensor.dispose()
+    cost.dispose()
+
+    await tf.nextFrame()
+  }))
+}
+
+function createBatchCreators(data, batchSize) {
+  if (batchSize < 1) {
+    throw new Error('invalid batch size: ' + batchSize)
+  }
+
+  const batches = []
+  const pushToBatch = (remaining) => {
+    if (remaining.length) {
+      batches.push(remaining.slice(0, batchSize))
+      pushToBatch(remaining.slice(batchSize))
+    }
+    return batches
+  }
+
+  pushToBatch(data)
+
+  const batchCreators = batches.map(dataForBatch => async () => {
+    const imgs = dataForBatch.map(d => d.img)
+    const allLandmarks = dataForBatch.map(d => landmarkPositionsToArray(d.landmarks))
+
+    const batchInput = await faceapi.toNetInput(imgs)
+    const landmarksBatchTensor = tf.tidy(() => {
+      const landmarkTensors = allLandmarks.map(arr => tf.tensor2d(arr, [1, 136]))
+      return tf.stack(landmarkTensors, 0).as2D(-1, 136)
+    })
+
+    return {
+      batchInput,
+      landmarksBatchTensor
+    }
+  })
+
+  return batchCreators
+}
+
+function landmarkPositionsToArray(landmarks) {
+  return landmarks.getRelativePositions().map(pt => [pt.x, pt.y])
+    .reduce((flat, arr) => flat.concat(arr))
+}
+
+function toFaceLandmarks(landmarks, { naturalWidth, naturalHeight }) {
+  return new faceapi.FaceLandmarks(
+    landmarks.map(l => new faceapi.Point(l.x / naturalWidth, l.y / naturalHeight)),
+    { width: naturalWidth, height: naturalHeight }
+  )
+}
+
+async function loadImagesInBatch(allLandmarks, offset = 0) {
+  return Promise.all(allLandmarks.map(async (landmarks, i) => {
+    const imgUri = `/train_images/${i + offset}.png`
+    const img = await faceapi.bufferToImage(await fetchImage(imgUri))
+    return {
+      imgUri,
+      img,
+      landmarks: toFaceLandmarks(landmarks, img)
+    }
+  }))
+}
+
+async function getTrainData() {
+  const landmarksJson = (await (await fetch('/train_landmarks.json')).json())
+
+  const numLandmarks = Object.keys(landmarksJson).length
+  const allLandmarks = Array.from(
+    new Array(numLandmarks),
+    (_, i) => landmarksJson[i]
+  )
+
+  const batch1 = await loadImagesInBatch(allLandmarks.slice(0, 4000))
+  const batch2 = await loadImagesInBatch(allLandmarks.slice(4000), 4000)
+
+  return batch1.concat(batch2)
+}
+
+// https://stackoverflow.com/questions/6274339/how-can-i-shuffle-an-array
+function shuffle(a) {
+  var j, x, i;
+  for (i = a.length - 1; i > 0; i--) {
+      j = Math.floor(Math.random() * (i + 1));
+      x = a[i];
+      a[i] = a[j];
+      a[j] = x;
+  }
+  return a;
+}
\ No newline at end of file

tools/train/public/faceLandmarksUi.js

+function onKeyDown(e) {
+  e.target.value = (
+    parseInt(e.target.value) + (e.keyCode === 38 ? 1 : (e.keyCode === 40 ? -1 : 0))
+  ) || e.target.value || 0
+
+  const uri = '/train_images/' + e.target.value + '.png'
+
+  console.log(uri)
+  onSelectionChanged(uri)
+}
+
+function onChangeDrawLines(e) {
+  window.drawLines = $(e.target).prop('checked')
+  redraw()
+}
+
+async function onSelectionChanged(uri) {
+  const imgBuf = await fetchImage(uri)
+  window.currentImg = await faceapi.bufferToImage(imgBuf)
+  window.landmarks = await Promise.all(window.nets.map(
+    async net => net.detectLandmarks(window.currentImg)
+  ))
+  redraw(uri)
+}
+
+function drawLandmarkCanvas(img, landmarks, drawOpts) {
+  const canvas = faceapi.createCanvasFromMedia(img)
+  $('#faceContainer').append(canvas)
+
+  faceapi.drawLandmarks(
+    canvas,
+    landmarks,
+    drawOpts
+  )
+}
+
+function redraw(uri) {
+  $('#faceContainer').empty()
+
+  const drawOpts = { lineWidth: window.drawLines ? 2 : 4, drawLines: window.drawLines }
+  window.landmarks.forEach(landmarks => {
+    drawLandmarkCanvas(window.currentImg, landmarks, drawOpts)
+  })
+
+  const trainImgAndLandmarks = window.trainData.find(t => t.imgUri === uri)
+  if (trainImgAndLandmarks) {
+    drawLandmarkCanvas(trainImgAndLandmarks.img, trainImgAndLandmarks.landmarks, drawOpts)
+  }
+}
+
+async function loadNet(file) {
+  const res = await fetch(file)
+  const weights = new Float32Array(await res.arrayBuffer())
+  return faceapi.faceLandmarkNet(weights)
+}
+
+async function init() {
+  //await faceapi.loadFaceLandmarkModel('/')
+  //window.nets.push(faceapi.landmarkNet)
+  //window.nets.push(await loadNet('retrained/landmarks_v0.weights'))
+  //window.nets.push(await loadNet('retrained/landmarks_v2.weights'))
+
+  window.trainNet = await loadNet('retrained/landmarks_v6.weights')
+  window.nets.push(trainNet)
+
+  $('#loader').hide()
+  await run()
+  await onSelectionChanged($('#selectList select').val())
+}
+
+$(document).ready(function() {
+  $('#imgByNr').keydown(onKeyDown);
+  renderFaceImageSelectList(
+    '#selectList',
+    onSelectionChanged,
+    { className: 'sheldon', imageIdx: 1 }
+  )
+  init()
+})
\ No newline at end of file

tools/train/server.js

+require('./.env')
+
+const express = require('express')
+const path = require('path')
+
+const app = express()
+
+const viewsDir = path.join(__dirname, 'views')
+app.use(express.static(viewsDir))
+app.use(express.static(path.join(__dirname, './public')))
+app.use(express.static(path.join(__dirname, './tmp')))
+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')))
+
+const trainDataPath = path.resolve(process.env.TRAIN_DATA_PATH)
+app.use(express.static(trainDataPath))
+
+app.get('/', (req, res) => res.redirect('/face_landmarks'))
+app.get('/face_landmarks', (req, res) => res.sendFile(path.join(viewsDir, 'faceLandmarks.html')))
+
+app.listen(3000, () => console.log('Listening on port 3000!'))
\ No newline at end of file

tools/train/views/faceLandmarks.html

+<!DOCTYPE html>
+<html>
+<head>
+  <link rel="stylesheet" href="styles.css">
+  <link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/materialize/0.100.2/css/materialize.css">
+  <script type="text/javascript" src="https://code.jquery.com/jquery-2.1.1.min.js"></script>
+  <script src="https://cdnjs.cloudflare.com/ajax/libs/materialize/0.100.2/js/materialize.min.js"></script>
+  <script src="face-api.js"></script>
+  <script src="commons.js"></script>
+  <script src="faceLandmarksUi.js"></script>
+  <script src="faceLandmarksTrain.js"></script>
+  <script src="FileSaver.js"></script>
+</head>
+<body>
+  <div id="navbar"></div>
+  <div class="center-content page-container">
+    <div>
+      <div class="progress" id="loader">
+        <div class="indeterminate"></div>
+      </div>
+      <div class="row side-by-side">
+        <div class="center-content">
+          <div id="faceContainer"></div>
+          <div id="selectList"></div>
+        </div>
+      </div>
+      <p>
+        <input type="checkbox" id="drawLinesCheckbox" checked="checked" onchange="onChangeDrawLines(event)" />
+        <label for="drawLinesCheckbox">Draw Lines</label>
+      </p>
+      <div class="row">
+        <label for="imgByNr">Enter image NR: </label>
+        <input id="imgByNr" type="text" class="bold">
+      </div>
+    </div>
+  </div>
+
+  <script>
+    tf = faceapi.tf
+    window.saveEveryNthIteration = 2
+    window.drawLines = true
+    window.nets = []
+    window.landmarks = []
+
+    window.trainSteps = 100
+    window.learningRate = 0.01
+    //window.optimizer = tf.train.sgd(learningRate)
+    window.optimizer = tf.train.adam(0.001, 0.9, 0.999, 1e-8)
+
+    function lossFunction(labels, out) {
+      return tf.losses.meanSquaredError(labels, out)
+    }
+
+    async function run() {
+      window.trainData = await getTrainData()
+      window.trainNet.variable()
+
+      return
+
+      await train()
+    }
+
+    async function train(batchSize = 10) {
+      for (let i = 0; i < trainSteps; i++) {
+        console.log('step', i)
+        const batchCreators = createBatchCreators(shuffle(window.trainData), batchSize)
+        let ts = Date.now()
+        await trainStep(batchCreators)
+        ts = Date.now() - ts
+        console.log('step %s done (%s ms)', i, ts)
+        if (((i + 1) % saveEveryNthIteration) === 0) {
+          saveWeights(i)
+        }
+      }
+    }
+
+    function saveWeights(idx = 0) {
+      const binaryWeights = new Float32Array(
+        window.trainNet.getParamList()
+          .map(({ tensor }) => Array.from(tensor.dataSync()))
+          .reduce((flat, arr) => flat.concat(arr))
+      )
+      saveAs(new Blob([binaryWeights]), 'landmark_trained_weights_' + idx + '.weights')
+    }
+
+  </script>
+
+</body>
+</html>
\ No newline at end of file