extract image patch data in bgr format to avoid channel swapping on gpu

justadudewhohacks · justadudewhohacks · commit 0e899bd29fe5 · 2018-07-12T17:08:07.000+02:00
diff --git a/src/mtcnn/extractImagePatches.ts b/src/mtcnn/extractImagePatches.ts
@@ -2,7 +2,6 @@ import * as tf from '@tensorflow/tfjs-core';
 
 import { Dimensions } from '../types';
 import { createCanvas, getContext2dOrThrow } from '../utils';
-import { bgrToRgbTensor } from './bgrToRgbTensor';
 import { BoundingBox } from './BoundingBox';
 import { normalize } from './normalize';
 
@@ -35,8 +34,10 @@ export async function extractImagePatches(
     const { data } = patchCtx.getImageData(0, 0, width, height)
 
     const currData = []
-    for(let i = 0; i < data.length; i++) {
-      if ((i + 1) % 4 === 0) continue
+    // RGBA -> BGR
+    for(let i = 0; i < data.length; i+=4) {
+      currData.push(data[i + 2])
+      currData.push(data[i + 1])
       currData.push(data[i])
     }
     imagePatchesDatas.push(currData)
@@ -45,10 +46,10 @@ export async function extractImagePatches(
 
   return imagePatchesDatas.map(data => {
     const t = tf.tidy(() => {
-      const imagePatchTensor = bgrToRgbTensor(tf.transpose(
+      const imagePatchTensor = tf.transpose(
         tf.tensor4d(data, [1, width, height, 3]),
         [0, 2, 1, 3]
-      ).toFloat()) as tf.Tensor4D
+      ).toFloat() as tf.Tensor4D
 
       return normalize(imagePatchTensor)
     })
diff --git a/src/mtcnn/stage1.ts b/src/mtcnn/stage1.ts
@@ -73,25 +73,26 @@ export function stage1(
 ) {
   stats.stage1 = []
 
-  const boxesForScale = scales.map((scale) => {
+  const pnetOutputs = scales.map((scale) => tf.tidy(() => {
     const statsForScale: any = { scale }
+    const resized = rescaleAndNormalize(imgTensor, scale)
 
-    const { scoresTensor, regionsTensor } = tf.tidy(() => {
-      const resized = rescaleAndNormalize(imgTensor, scale)
-
-      let ts = Date.now()
-      const { prob, regions } = PNet(resized, params)
-      statsForScale.pnet = Date.now() - ts
+    let ts = Date.now()
+    const { prob, regions } = PNet(resized, params)
+    statsForScale.pnet = Date.now() - ts
 
-      const scoresTensor = tf.unstack(tf.unstack(prob, 3)[1])[0] as tf.Tensor2D
-      const regionsTensor = tf.unstack(regions)[0] as tf.Tensor3D
+    const scoresTensor = tf.unstack(tf.unstack(prob, 3)[1])[0] as tf.Tensor2D
+    const regionsTensor = tf.unstack(regions)[0] as tf.Tensor3D
 
-      return {
-        scoresTensor,
-        regionsTensor
-      }
-    })
+    return {
+      scoresTensor,
+      regionsTensor,
+      scale,
+      statsForScale
+    }
+  }))
 
+  const boxesForScale = pnetOutputs.map(({ scoresTensor, regionsTensor, scale, statsForScale }) => {
     const boundingBoxes = extractBoundingBoxes(
       scoresTensor,
       regionsTensor,
@@ -121,7 +122,7 @@ export function stage1(
   })
 
   const allBoxes = boxesForScale.reduce(
-    (all, boxes) => all.concat(boxes)
+    (all, boxes) => all.concat(boxes), []
   )
 
   let finalBoxes: BoundingBox[] = []
