@tensorflow-models/body-pix
Version:
Pretrained BodyPix model in TensorFlow.js
58 lines • 2.59 kB
JavaScript
;
/**
* @license
* Copyright 2019 Google LLC. All Rights Reserved.
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
* =============================================================================
*/
Object.defineProperty(exports, "__esModule", { value: true });
exports.getOffsetPoints = exports.getOffsetVectors = exports.getPointsConfidence = void 0;
var tf = require("@tensorflow/tfjs-core");
var keypoints_1 = require("../keypoints");
function getPointsConfidence(heatmapScores, heatMapCoords) {
var numKeypoints = heatMapCoords.shape[0];
var result = new Float32Array(numKeypoints);
for (var keypoint = 0; keypoint < numKeypoints; keypoint++) {
var y = heatMapCoords.get(keypoint, 0);
var x = heatMapCoords.get(keypoint, 1);
result[keypoint] = heatmapScores.get(y, x, keypoint);
}
return result;
}
exports.getPointsConfidence = getPointsConfidence;
function getOffsetPoint(y, x, keypoint, offsetsBuffer) {
return {
y: offsetsBuffer.get(y, x, keypoint),
x: offsetsBuffer.get(y, x, keypoint + keypoints_1.NUM_KEYPOINTS)
};
}
function getOffsetVectors(heatMapCoordsBuffer, offsetsBuffer) {
var result = [];
for (var keypoint = 0; keypoint < keypoints_1.NUM_KEYPOINTS; keypoint++) {
var heatmapY = heatMapCoordsBuffer.get(keypoint, 0).valueOf();
var heatmapX = heatMapCoordsBuffer.get(keypoint, 1).valueOf();
var _a = getOffsetPoint(heatmapY, heatmapX, keypoint, offsetsBuffer), x = _a.x, y = _a.y;
result.push(y);
result.push(x);
}
return tf.tensor2d(result, [keypoints_1.NUM_KEYPOINTS, 2]);
}
exports.getOffsetVectors = getOffsetVectors;
function getOffsetPoints(heatMapCoordsBuffer, outputStride, offsetsBuffer) {
return tf.tidy(function () {
var offsetVectors = getOffsetVectors(heatMapCoordsBuffer, offsetsBuffer);
return tf.add(tf.cast(tf.mul(heatMapCoordsBuffer.toTensor(), tf.scalar(outputStride, 'int32')), 'float32'), offsetVectors);
});
}
exports.getOffsetPoints = getOffsetPoints;
//# sourceMappingURL=util.js.map