@tensorflow-models/body-pix
Version:
Pretrained BodyPix model in TensorFlow.js
80 lines • 4.1 kB
JavaScript
;
Object.defineProperty(exports, "__esModule", { value: true });
var keypoints_1 = require("../keypoints");
var util_1 = require("./util");
var util_2 = require("./util");
var parentChildrenTuples = keypoints_1.POSE_CHAIN.map(function (_a) {
var parentJoinName = _a[0], childJoinName = _a[1];
return ([keypoints_1.PART_IDS[parentJoinName], keypoints_1.PART_IDS[childJoinName]]);
});
var parentToChildEdges = parentChildrenTuples.map(function (_a) {
var childJointId = _a[1];
return childJointId;
});
var childToParentEdges = parentChildrenTuples.map(function (_a) {
var parentJointId = _a[0];
return parentJointId;
});
function getDisplacement(edgeId, point, displacements) {
var numEdges = displacements.shape[2] / 2;
return {
y: displacements.get(point.y, point.x, edgeId),
x: displacements.get(point.y, point.x, numEdges + edgeId)
};
}
function getStridedIndexNearPoint(point, outputStride, height, width) {
return {
y: util_1.clamp(Math.round(point.y / outputStride), 0, height - 1),
x: util_1.clamp(Math.round(point.x / outputStride), 0, width - 1)
};
}
function traverseToTargetKeypoint(edgeId, sourceKeypoint, targetKeypointId, scoresBuffer, offsets, outputStride, displacements, offsetRefineStep) {
if (offsetRefineStep === void 0) { offsetRefineStep = 2; }
var _a = scoresBuffer.shape, height = _a[0], width = _a[1];
var sourceKeypointIndices = getStridedIndexNearPoint(sourceKeypoint.position, outputStride, height, width);
var displacement = getDisplacement(edgeId, sourceKeypointIndices, displacements);
var displacedPoint = util_2.addVectors(sourceKeypoint.position, displacement);
var targetKeypoint = displacedPoint;
for (var i = 0; i < offsetRefineStep; i++) {
var targetKeypointIndices = getStridedIndexNearPoint(targetKeypoint, outputStride, height, width);
var offsetPoint = util_1.getOffsetPoint(targetKeypointIndices.y, targetKeypointIndices.x, targetKeypointId, offsets);
targetKeypoint = util_2.addVectors({
x: targetKeypointIndices.x * outputStride,
y: targetKeypointIndices.y * outputStride
}, { x: offsetPoint.x, y: offsetPoint.y });
}
var targetKeyPointIndices = getStridedIndexNearPoint(targetKeypoint, outputStride, height, width);
var score = scoresBuffer.get(targetKeyPointIndices.y, targetKeyPointIndices.x, targetKeypointId);
return { position: targetKeypoint, part: keypoints_1.PART_NAMES[targetKeypointId], score: score };
}
function decodePose(root, scores, offsets, outputStride, displacementsFwd, displacementsBwd) {
var numParts = scores.shape[2];
var numEdges = parentToChildEdges.length;
var instanceKeypoints = new Array(numParts);
var rootPart = root.part, rootScore = root.score;
var rootPoint = util_2.getImageCoords(rootPart, outputStride, offsets);
instanceKeypoints[rootPart.id] = {
score: rootScore,
part: keypoints_1.PART_NAMES[rootPart.id],
position: rootPoint
};
for (var edge = numEdges - 1; edge >= 0; --edge) {
var sourceKeypointId = parentToChildEdges[edge];
var targetKeypointId = childToParentEdges[edge];
if (instanceKeypoints[sourceKeypointId] &&
!instanceKeypoints[targetKeypointId]) {
instanceKeypoints[targetKeypointId] = traverseToTargetKeypoint(edge, instanceKeypoints[sourceKeypointId], targetKeypointId, scores, offsets, outputStride, displacementsBwd);
}
}
for (var edge = 0; edge < numEdges; ++edge) {
var sourceKeypointId = childToParentEdges[edge];
var targetKeypointId = parentToChildEdges[edge];
if (instanceKeypoints[sourceKeypointId] &&
!instanceKeypoints[targetKeypointId]) {
instanceKeypoints[targetKeypointId] = traverseToTargetKeypoint(edge, instanceKeypoints[sourceKeypointId], targetKeypointId, scores, offsets, outputStride, displacementsFwd);
}
}
return instanceKeypoints;
}
exports.decodePose = decodePose;
//# sourceMappingURL=decode_pose.js.map