@vladmandic/face-api/src/ssdMobilenetv1/boxPredictionLayer.ts

JavaScript module for Face Detection and Face Recognition Using Tensorflow/JS

import * as tf from '@tensorflow/tfjs/dist/tf.es2017.js';

import { convLayer } from '../common';
import { BoxPredictionParams } from './types';


export function boxPredictionLayer(
  x: tf.Tensor4D,
  params: BoxPredictionParams
) {
  return tf.tidy(() => {

    const batchSize = x.shape[0]

    const boxPredictionEncoding = tf.reshape(
      convLayer(x, params.box_encoding_predictor),
      [batchSize, -1, 1, 4]
    )
    const classPrediction = tf.reshape(
      convLayer(x, params.class_predictor),
      [batchSize, -1, 3]
    )

    return {
      boxPredictionEncoding,
      classPrediction
    }
  })
}