gx

chenyc

2025-02-12 ea42ff3ebee1eeb3fb29423aa848a249441db81c

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
import * as tf from '../../dist/tfjs.esm';
 
import { OutputLayerParams } from './types';
 
function getCenterCoordinatesAndSizesLayer(x: tf.Tensor2D) {
  const vec = tf.unstack(tf.transpose(x, [1, 0]));
 
  const sizes = [
    tf.sub(vec[2], vec[0]),
    tf.sub(vec[3], vec[1]),
  ];
  const centers = [
    tf.add(vec[0], tf.div(sizes[0], 2)),
    tf.add(vec[1], tf.div(sizes[1], 2)),
  ];
  return { sizes, centers };
}
 
function decodeBoxesLayer(x0: tf.Tensor2D, x1: tf.Tensor2D) {
  const { sizes, centers } = getCenterCoordinatesAndSizesLayer(x0);
 
  const vec = tf.unstack(tf.transpose(x1, [1, 0]));
  const div0_out = tf.div(tf.mul(tf.exp(tf.div(vec[2], 5)), sizes[0]), 2);
  const add0_out = tf.add(tf.mul(tf.div(vec[0], 10), sizes[0]), centers[0]);
  const div1_out = tf.div(tf.mul(tf.exp(tf.div(vec[3], 5)), sizes[1]), 2);
  const add1_out = tf.add(tf.mul(tf.div(vec[1], 10), sizes[1]), centers[1]);
 
  return tf.transpose(
    tf.stack([
      tf.sub(add0_out, div0_out),
      tf.sub(add1_out, div1_out),
      tf.add(add0_out, div0_out),
      tf.add(add1_out, div1_out),
    ]),
    [1, 0],
  );
}
 
export function outputLayer(boxPredictions: tf.Tensor4D, classPredictions: tf.Tensor4D, params: OutputLayerParams) {
  return tf.tidy(() => {
    const batchSize = boxPredictions.shape[0];
 
    let boxes = decodeBoxesLayer(
      tf.reshape(tf.tile(params.extra_dim, [batchSize, 1, 1]), [-1, 4]) as tf.Tensor2D,
      tf.reshape(boxPredictions, [-1, 4]) as tf.Tensor2D,
    );
    boxes = tf.reshape(boxes, [batchSize, (boxes.shape[0] / batchSize), 4]);
 
    const scoresAndClasses = tf.sigmoid(tf.slice(classPredictions, [0, 0, 1], [-1, -1, -1]));
    let scores = tf.slice(scoresAndClasses, [0, 0, 0], [-1, -1, 1]) as tf.Tensor;
 
    scores = tf.reshape(scores, [batchSize, scores.shape[1] as number]);
 
    const boxesByBatch = tf.unstack(boxes) as tf.Tensor2D[];
    const scoresByBatch = tf.unstack(scores) as tf.Tensor1D[];
 
    return { boxes: boxesByBatch, scores: scoresByBatch };
  });
}