1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
| import * as tf from '../../dist/tfjs.esm';
|
| import { convLayer } from '../common/index';
| 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 };
| });
| }
|
|
