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
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
import * as tf from '@tensorflow/tfjs-core';
import { BoundingBox, Point } from '../classes';
import { nonMaxSuppression } from '../ops';
import { CELL_SIZE, CELL_STRIDE } from './config';
import { getSizesForScale } from './getSizesForScale';
import { MtcnnBox } from './MtcnnBox';
import { normalize } from './normalize';
import { PNet } from './PNet';
function rescaleAndNormalize(x, scale) {
    return tf.tidy(function () {
        var _a = getSizesForScale(scale, x.shape.slice(1)), height = _a.height, width = _a.width;
        var resized = tf.image.resizeBilinear(x, [height, width]);
        var normalized = normalize(resized);
        return tf.transpose(normalized, [0, 2, 1, 3]);
    });
}
function extractBoundingBoxes(scoresTensor, regionsTensor, scale, scoreThreshold) {
    // TODO: fix this!, maybe better to use tf.gather here
    var indices = [];
    var scoresData = scoresTensor.arraySync();
    for (var y = 0; y < scoresTensor.shape[0]; y++) {
        for (var x = 0; x < scoresTensor.shape[1]; x++) {
            if (scoresData[y][x] >= scoreThreshold) {
                indices.push(new Point(x, y));
            }
        }
    }
    var boundingBoxes = indices.map(function (idx) {
        var cell = new BoundingBox(Math.round((idx.y * CELL_STRIDE + 1) / scale), Math.round((idx.x * CELL_STRIDE + 1) / scale), Math.round((idx.y * CELL_STRIDE + CELL_SIZE) / scale), Math.round((idx.x * CELL_STRIDE + CELL_SIZE) / scale));
        var score = scoresData[idx.y][idx.x];
        var regionsData = regionsTensor.arraySync();
        var region = new MtcnnBox(regionsData[idx.y][idx.x][0], regionsData[idx.y][idx.x][1], regionsData[idx.y][idx.x][2], regionsData[idx.y][idx.x][3]);
        return {
            cell: cell,
            score: score,
            region: region
        };
    });
    return boundingBoxes;
}
export function stage1(imgTensor, scales, scoreThreshold, params, stats) {
    stats.stage1 = [];
    var pnetOutputs = scales.map(function (scale) { return tf.tidy(function () {
        var statsForScale = { scale: scale };
        var resized = rescaleAndNormalize(imgTensor, scale);
        var ts = Date.now();
        var _a = PNet(resized, params), prob = _a.prob, regions = _a.regions;
        statsForScale.pnet = Date.now() - ts;
        var scoresTensor = tf.unstack(tf.unstack(prob, 3)[1])[0];
        var regionsTensor = tf.unstack(regions)[0];
        return {
            scoresTensor: scoresTensor,
            regionsTensor: regionsTensor,
            scale: scale,
            statsForScale: statsForScale
        };
    }); });
    var boxesForScale = pnetOutputs.map(function (_a) {
        var scoresTensor = _a.scoresTensor, regionsTensor = _a.regionsTensor, scale = _a.scale, statsForScale = _a.statsForScale;
        var boundingBoxes = extractBoundingBoxes(scoresTensor, regionsTensor, scale, scoreThreshold);
        scoresTensor.dispose();
        regionsTensor.dispose();
        if (!boundingBoxes.length) {
            stats.stage1.push(statsForScale);
            return [];
        }
        var ts = Date.now();
        var indices = nonMaxSuppression(boundingBoxes.map(function (bbox) { return bbox.cell; }), boundingBoxes.map(function (bbox) { return bbox.score; }), 0.5);
        statsForScale.nms = Date.now() - ts;
        statsForScale.numBoxes = indices.length;
        stats.stage1.push(statsForScale);
        return indices.map(function (boxIdx) { return boundingBoxes[boxIdx]; });
    });
    var allBoxes = boxesForScale.reduce(function (all, boxes) { return all.concat(boxes); }, []);
    var finalBoxes = [];
    var finalScores = [];
    if (allBoxes.length > 0) {
        var ts = Date.now();
        var indices = nonMaxSuppression(allBoxes.map(function (bbox) { return bbox.cell; }), allBoxes.map(function (bbox) { return bbox.score; }), 0.7);
        stats.stage1_nms = Date.now() - ts;
        finalScores = indices.map(function (idx) { return allBoxes[idx].score; });
        finalBoxes = indices
            .map(function (idx) { return allBoxes[idx]; })
            .map(function (_a) {
            var cell = _a.cell, region = _a.region;
            return new BoundingBox(cell.left + (region.left * cell.width), cell.top + (region.top * cell.height), cell.right + (region.right * cell.width), cell.bottom + (region.bottom * cell.height)).toSquare().round();
        });
    }
    return {
        boxes: finalBoxes,
        scores: finalScores
    };
}
//# sourceMappingURL=stage1.js.map