init train data fetching

justadudewhohacks · justadudewhohacks · commit aa7dd87c475a · 2018-08-04T20:05:22.000+02:00
diff --git a/tools/train/faceLandmarks/faceLandmarksTrain.js b/tools/train/faceLandmarks/faceLandmarksTrain.js
@@ -1,13 +1,3 @@
-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 () => {
 
diff --git a/tools/train/serveFaceLandmarks.js b/tools/train/serveFaceLandmarks.js
@@ -7,6 +7,7 @@ const app = express()
 
 const publicDir = path.join(__dirname, './faceLandmarks')
 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')))
diff --git a/tools/train/serveTinyYolov2.js b/tools/train/serveTinyYolov2.js
@@ -2,21 +2,26 @@ require('./tinyYolov2/.env')
 
 const express = require('express')
 const path = require('path')
+const fs = require('fs')
 
 const app = express()
 
 const publicDir = path.join(__dirname, './tinyYolov2')
 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')))
 
 const trainDataPath = path.resolve(process.env.TRAIN_DATA_PATH)
+const imagesPath = path.join(trainDataPath, './final_images')
+const detectionsPath = path.join(trainDataPath, './final_detections')
+const detectionFilenames = fs.readdirSync(detectionsPath)
+
 app.use(express.static(trainDataPath))
 
-//app.get('/', (req, res) => res.sendFile(path.join(publicDir, 'train.html')))
-//app.get('/', (req, res) => res.sendFile(path.join(publicDir, 'tinyYolov2FaceDetectionVideo.html')))
-app.get('/', (req, res) => res.sendFile(path.join(publicDir, 'testLoss.html')))
+app.get('/detection_filenames', (req, res) => res.status(202).send(detectionFilenames))
+app.get('/', (req, res) => res.sendFile(path.join(publicDir, 'train.html')))
 
 app.listen(3000, () => console.log('Listening on port 3000!'))
diff --git a/tools/train/shared/trainUtils.js b/tools/train/shared/trainUtils.js
@@ -1,4 +1,12 @@
+async function promiseSequential(promises) {
+  const curr = promises[0]
+  if (!curr) {
+    return
+  }
 
+  await curr()
+  return promiseSequential(promises.slice(1))
+}
 
 // https://stackoverflow.com/questions/6274339/how-can-i-shuffle-an-array
 function shuffle(a) {
diff --git a/tools/train/tinyYolov2/train.html b/tools/train/tinyYolov2/train.html
@@ -29,16 +29,16 @@
       return new Float32Array(await (await fetch(uri)).arrayBuffer())
     }
 
-    async function getTrainData() {
-      // TBD
+    async function fetchDetectionFilenames() {
+      return fetch('/detection_filenames').then(res => res.json())
     }
 
     async function run() {
       const weights = await loadNetWeights(weightsUrl)
       window.net = new faceapi.TinyYolov2(true)
       window.net.load(weights)
-      window.trainData = await getTrainData()
       window.net.variable()
+      window.detectionFilenames = await fetchDetectionFilenames()
     }
 
 /*
@@ -68,7 +68,7 @@
     async function train(batchSize = 1) {
       for (let i = 0; i < trainSteps; i++) {
         console.log('step', i)
-        const batchCreators = createBatchCreators(shuffle(window.trainData), batchSize)
+        const batchCreators = createBatchCreators(shuffle(detectionFilenames), batchSize)
         let ts = Date.now()
         await trainStep(batchCreators)
         ts = Date.now() - ts
diff --git a/tools/train/tinyYolov2/train.js b/tools/train/tinyYolov2/train.js
@@ -0,0 +1,61 @@
+async function trainStep(batchCreators) {
+  await promiseSequential(batchCreators.map((batchCreator, dataIdx) => async () => {
+
+    const { batchInput, groundTruthBoxes } = 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(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(window.detectionFilenames)
+
+  const batchCreators = batches.map(detectionFilenames => async () => {
+    const imgs = detectionFilenames.map(
+      detectionFilenames.map(file => fetch(file).then(res => res.json()))
+    )
+    const groundTruthBoxes = await Promise.all(
+      detectionFilenames.map(async file => await faceapi.bufferToImage(await fetchImage(file.replace('.json', ''))))
+    )
+
+    const batchInput = await faceapi.toNetInput(imgs)
+
+    return {
+      batchInput,
+      groundTruthBoxes
+    }
+  })
+
+  return batchCreators
+}
