mirror of https://github.com/vladmandic/human
55 lines
1.9 KiB
TypeScript
55 lines
1.9 KiB
TypeScript
import * as tf from '../../dist/tfjs.esm.js';
|
|
import * as kpt from './keypoints';
|
|
|
|
export function getPointsConfidence(heatmapScores, heatMapCoords) {
|
|
const numKeypoints = heatMapCoords.shape[0];
|
|
const result = new Float32Array(numKeypoints);
|
|
for (let keypoint = 0; keypoint < numKeypoints; keypoint++) {
|
|
const y = heatMapCoords.get(keypoint, 0);
|
|
const x = heatMapCoords.get(keypoint, 1);
|
|
result[keypoint] = heatmapScores.get(y, x, keypoint);
|
|
}
|
|
return result;
|
|
}
|
|
|
|
function getOffsetPoint(y, x, keypoint, offsetsBuffer) {
|
|
return {
|
|
y: offsetsBuffer.get(y, x, keypoint),
|
|
x: offsetsBuffer.get(y, x, keypoint + kpt.NUM_KEYPOINTS),
|
|
};
|
|
}
|
|
|
|
export function getOffsetVectors(heatMapCoordsBuffer, offsetsBuffer) {
|
|
const result: Array<number> = [];
|
|
for (let keypoint = 0; keypoint < kpt.NUM_KEYPOINTS; keypoint++) {
|
|
const heatmapY = heatMapCoordsBuffer.get(keypoint, 0).valueOf();
|
|
const heatmapX = heatMapCoordsBuffer.get(keypoint, 1).valueOf();
|
|
const { x, y } = getOffsetPoint(heatmapY, heatmapX, keypoint, offsetsBuffer);
|
|
result.push(y);
|
|
result.push(x);
|
|
}
|
|
return tf.tensor2d(result, [kpt.NUM_KEYPOINTS, 2]);
|
|
}
|
|
|
|
export function getOffsetPoints(heatMapCoordsBuffer, outputStride, offsetsBuffer) {
|
|
return tf.tidy(() => heatMapCoordsBuffer.toTensor().mul(tf.scalar(outputStride, 'int32')).toFloat().add(getOffsetVectors(heatMapCoordsBuffer, offsetsBuffer)));
|
|
}
|
|
|
|
function mod(a, b) {
|
|
return tf.tidy(() => {
|
|
const floored = a.div(tf.scalar(b, 'int32'));
|
|
return a.sub(floored.mul(tf.scalar(b, 'int32')));
|
|
});
|
|
}
|
|
|
|
export function argmax2d(inputs) {
|
|
const [height, width, depth] = inputs.shape;
|
|
return tf.tidy(() => {
|
|
const reshaped = inputs.reshape([height * width, depth]);
|
|
const coords = reshaped.argMax(0);
|
|
const yCoords = coords.div(tf.scalar(width, 'int32')).expandDims(1);
|
|
const xCoords = mod(coords, width).expandDims(1);
|
|
return tf.concat([yCoords, xCoords], 1);
|
|
});
|
|
}
|