fixed tinyYolov2 implementation

justadudewhohacks · justadudewhohacks · commit e82cc5ab71a9 · 2018-07-28T20:10:53.000+02:00
diff --git a/src/tinyYolov2/TinyYolov2.ts b/src/tinyYolov2/TinyYolov2.ts
@@ -25,22 +25,23 @@ export class TinyYolov2 extends NeuralNetwork<NetParams> {
     }
 
     const out = tf.tidy(() => {
-      const batchTensor = input.toBatchTensor(416).div(tf.scalar(255)).toFloat()
+      //const batchTensor = input.toBatchTensor(416).div(tf.scalar(255)).toFloat() as tf.Tensor4D
 
-      let out = tf.pad(batchTensor, [[0, 0], [1, 1], [1, 1], [0, 0]]) as tf.Tensor4D
+      // TODO: fix boxes after padding
+      const batchTensor = tf.image.resizeBilinear(input.inputs[0], [416, 416]).toFloat().div(tf.scalar(255)).expandDims() as tf.Tensor4D
 
-      out = convWithBatchNorm(out, params.conv0)
-      out = tf.maxPool(out, [2, 2], [2, 2], 'valid')
+      let out = convWithBatchNorm(batchTensor, params.conv0)
+      out = tf.maxPool(out, [2, 2], [2, 2], 'same')
       out = convWithBatchNorm(out, params.conv1)
-      out = tf.maxPool(out, [2, 2], [2, 2], 'valid')
+      out = tf.maxPool(out, [2, 2], [2, 2], 'same')
       out = convWithBatchNorm(out, params.conv2)
-      out = tf.maxPool(out, [2, 2], [2, 2], 'valid')
+      out = tf.maxPool(out, [2, 2], [2, 2], 'same')
       out = convWithBatchNorm(out, params.conv3)
-      out = tf.maxPool(out, [2, 2], [2, 2], 'valid')
+      out = tf.maxPool(out, [2, 2], [2, 2], 'same')
       out = convWithBatchNorm(out, params.conv4)
-      out = tf.maxPool(out, [2, 2], [2, 2], 'valid')
+      out = tf.maxPool(out, [2, 2], [2, 2], 'same')
       out = convWithBatchNorm(out, params.conv5)
-      out = tf.maxPool(out, [2, 2], [1, 1], 'valid')
+      out = tf.maxPool(out, [2, 2], [1, 1], 'same')
       out = convWithBatchNorm(out, params.conv6)
       out = convWithBatchNorm(out, params.conv7)
       out = convLayer(out, params.conv8, 'valid', false)
diff --git a/src/tinyYolov2/convWithBatchNorm.ts b/src/tinyYolov2/convWithBatchNorm.ts
@@ -5,7 +5,8 @@ import { ConvWithBatchNorm } from './types';
 
 export function convWithBatchNorm(x: tf.Tensor4D, params: ConvWithBatchNorm): tf.Tensor4D  {
   return tf.tidy(() => {
-    let out = tf.conv2d(x, params.conv.filters, [1, 1], 'valid')
+    let out = tf.pad(x, [[0, 0], [1, 1], [1, 1], [0, 0]]) as tf.Tensor4D
+    out = tf.conv2d(out, params.conv.filters, [1, 1], 'valid')
     out = tf.sub(out, params.bn.sub)
     out = tf.mul(out, params.bn.truediv)
     out = tf.add(out, params.conv.bias)
