UNPKG

@tensorflow-models/body-pix

Version:

Pretrained BodyPix model in TensorFlow.js

314 lines 13.4 kB
"use strict"; var __awaiter = (this && this.__awaiter) || function (thisArg, _arguments, P, generator) { return new (P || (P = Promise))(function (resolve, reject) { function fulfilled(value) { try { step(generator.next(value)); } catch (e) { reject(e); } } function rejected(value) { try { step(generator["throw"](value)); } catch (e) { reject(e); } } function step(result) { result.done ? resolve(result.value) : new P(function (resolve) { resolve(result.value); }).then(fulfilled, rejected); } step((generator = generator.apply(thisArg, _arguments || [])).next()); }); }; var __generator = (this && this.__generator) || function (thisArg, body) { var _ = { label: 0, sent: function() { if (t[0] & 1) throw t[1]; return t[1]; }, trys: [], ops: [] }, f, y, t, g; return g = { next: verb(0), "throw": verb(1), "return": verb(2) }, typeof Symbol === "function" && (g[Symbol.iterator] = function() { return this; }), g; function verb(n) { return function (v) { return step([n, v]); }; } function step(op) { if (f) throw new TypeError("Generator is already executing."); while (_) try { if (f = 1, y && (t = op[0] & 2 ? y["return"] : op[0] ? y["throw"] || ((t = y["return"]) && t.call(y), 0) : y.next) && !(t = t.call(y, op[1])).done) return t; if (y = 0, t) op = [op[0] & 2, t.value]; switch (op[0]) { case 0: case 1: t = op; break; case 4: _.label++; return { value: op[1], done: false }; case 5: _.label++; y = op[1]; op = [0]; continue; case 7: op = _.ops.pop(); _.trys.pop(); continue; default: if (!(t = _.trys, t = t.length > 0 && t[t.length - 1]) && (op[0] === 6 || op[0] === 2)) { _ = 0; continue; } if (op[0] === 3 && (!t || (op[1] > t[0] && op[1] < t[3]))) { _.label = op[1]; break; } if (op[0] === 6 && _.label < t[1]) { _.label = t[1]; t = op; break; } if (t && _.label < t[2]) { _.label = t[2]; _.ops.push(op); break; } if (t[2]) _.ops.pop(); _.trys.pop(); continue; } op = body.call(thisArg, _); } catch (e) { op = [6, e]; y = 0; } finally { f = t = 0; } if (op[0] & 5) throw op[1]; return { value: op[0] ? op[1] : void 0, done: true }; } }; var _a; Object.defineProperty(exports, "__esModule", { value: true }); var tf = require("@tensorflow/tfjs-core"); function getSizeFromImageLikeElement(input) { if (input.offsetHeight !== 0 && input.offsetWidth !== 0) { return [input.offsetHeight, input.offsetWidth]; } else if (input.height != null && input.width != null) { return [input.height, input.width]; } else { throw new Error("HTMLImageElement must have height and width attributes set."); } } function getSizeFromVideoElement(input) { if (input.height != null && input.width != null) { return [input.height, input.width]; } else { return [input.videoHeight, input.videoWidth]; } } function getInputSize(input) { if ((typeof (HTMLCanvasElement) !== 'undefined' && input instanceof HTMLCanvasElement) || (typeof (HTMLImageElement) !== 'undefined' && input instanceof HTMLImageElement)) { return getSizeFromImageLikeElement(input); } else if (typeof (ImageData) !== 'undefined' && input instanceof ImageData) { return [input.height, input.width]; } else if (typeof (HTMLVideoElement) !== 'undefined' && input instanceof HTMLVideoElement) { return getSizeFromVideoElement(input); } else if (input instanceof tf.Tensor) { return [input.shape[0], input.shape[1]]; } else { throw new Error("error: Unknown input type: " + input + "."); } } exports.getInputSize = getInputSize; function isValidInputResolution(resolution, outputStride) { return (resolution - 1) % outputStride === 0; } function toValidInputResolution(inputResolution, outputStride) { if (isValidInputResolution(inputResolution, outputStride)) { return inputResolution; } return Math.floor(inputResolution / outputStride) * outputStride + 1; } exports.toValidInputResolution = toValidInputResolution; var INTERNAL_RESOLUTION_STRING_OPTIONS = { low: 'low', medium: 'medium', high: 'high', full: 'full' }; var INTERNAL_RESOLUTION_PERCENTAGES = (_a = {}, _a[INTERNAL_RESOLUTION_STRING_OPTIONS.low] = 0.25, _a[INTERNAL_RESOLUTION_STRING_OPTIONS.medium] = 0.5, _a[INTERNAL_RESOLUTION_STRING_OPTIONS.high] = 0.75, _a[INTERNAL_RESOLUTION_STRING_OPTIONS.full] = 1.0, _a); var MIN_INTERNAL_RESOLUTION = 0.1; var MAX_INTERNAL_RESOLUTION = 2.0; function toInternalResolutionPercentage(internalResolution) { if (typeof internalResolution === 'string') { var result = INTERNAL_RESOLUTION_PERCENTAGES[internalResolution]; tf.util.assert(typeof result === 'number', function () { return "string value of inputResolution must be one of " + Object.values(INTERNAL_RESOLUTION_STRING_OPTIONS) .join(',') + " but was " + internalResolution + "."; }); return result; } else { tf.util.assert(typeof internalResolution === 'number' && internalResolution <= MAX_INTERNAL_RESOLUTION && internalResolution >= MIN_INTERNAL_RESOLUTION, function () { return "inputResolution must be a string or number between 0 and 4, but " + ("was " + internalResolution); }); return internalResolution; } } function toInputResolutionHeightAndWidth(internalResolution, outputStride, _a) { var inputHeight = _a[0], inputWidth = _a[1]; var internalResolutionPercentage = toInternalResolutionPercentage(internalResolution); return [ toValidInputResolution(inputHeight * internalResolutionPercentage, outputStride), toValidInputResolution(inputWidth * internalResolutionPercentage, outputStride) ]; } exports.toInputResolutionHeightAndWidth = toInputResolutionHeightAndWidth; function toInputTensor(input) { return input instanceof tf.Tensor ? input : tf.browser.fromPixels(input); } exports.toInputTensor = toInputTensor; function resizeAndPadTo(imageTensor, _a, flipHorizontal) { var targetH = _a[0], targetW = _a[1]; if (flipHorizontal === void 0) { flipHorizontal = false; } var _b = imageTensor.shape, height = _b[0], width = _b[1]; var targetAspect = targetW / targetH; var aspect = width / height; var resizeW; var resizeH; var padL; var padR; var padT; var padB; if (aspect > targetAspect) { resizeW = targetW; resizeH = Math.ceil(resizeW / aspect); var padHeight = targetH - resizeH; padL = 0; padR = 0; padT = Math.floor(padHeight / 2); padB = targetH - (resizeH + padT); } else { resizeH = targetH; resizeW = Math.ceil(targetH * aspect); var padWidth = targetW - resizeW; padL = Math.floor(padWidth / 2); padR = targetW - (resizeW + padL); padT = 0; padB = 0; } var resizedAndPadded = tf.tidy(function () { var resized; if (flipHorizontal) { resized = imageTensor.reverse(1).resizeBilinear([resizeH, resizeW]); } else { resized = imageTensor.resizeBilinear([resizeH, resizeW]); } var padded = tf.pad3d(resized, [[padT, padB], [padL, padR], [0, 0]]); return padded; }); return { resizedAndPadded: resizedAndPadded, paddedBy: [[padT, padB], [padL, padR]] }; } exports.resizeAndPadTo = resizeAndPadTo; function scaleAndCropToInputTensorShape(tensor, _a, _b, _c, applySigmoidActivation) { var inputTensorHeight = _a[0], inputTensorWidth = _a[1]; var resizedAndPaddedHeight = _b[0], resizedAndPaddedWidth = _b[1]; var _d = _c[0], padT = _d[0], padB = _d[1], _e = _c[1], padL = _e[0], padR = _e[1]; if (applySigmoidActivation === void 0) { applySigmoidActivation = false; } return tf.tidy(function () { var inResizedAndPadded = tensor.resizeBilinear([resizedAndPaddedHeight, resizedAndPaddedWidth], true); if (applySigmoidActivation) { inResizedAndPadded = inResizedAndPadded.sigmoid(); } return removePaddingAndResizeBack(inResizedAndPadded, [inputTensorHeight, inputTensorWidth], [[padT, padB], [padL, padR]]); }); } exports.scaleAndCropToInputTensorShape = scaleAndCropToInputTensorShape; function removePaddingAndResizeBack(resizedAndPadded, _a, _b) { var originalHeight = _a[0], originalWidth = _a[1]; var _c = _b[0], padT = _c[0], padB = _c[1], _d = _b[1], padL = _d[0], padR = _d[1]; return tf.tidy(function () { return tf.image .cropAndResize(resizedAndPadded.expandDims(), [[ padT / (originalHeight + padT + padB - 1.0), padL / (originalWidth + padL + padR - 1.0), (padT + originalHeight - 1.0) / (originalHeight + padT + padB - 1.0), (padL + originalWidth - 1.0) / (originalWidth + padL + padR - 1.0) ]], [0], [originalHeight, originalWidth]) .squeeze([0]); }); } exports.removePaddingAndResizeBack = removePaddingAndResizeBack; function resize2d(tensor, resolution, nearestNeighbor) { return tf.tidy(function () { return tensor.expandDims(2) .resizeBilinear(resolution, nearestNeighbor) .squeeze(); }); } exports.resize2d = resize2d; function padAndResizeTo(input, _a) { var targetH = _a[0], targetW = _a[1]; var _b = getInputSize(input), height = _b[0], width = _b[1]; var targetAspect = targetW / targetH; var aspect = width / height; var _c = [0, 0, 0, 0], padT = _c[0], padB = _c[1], padL = _c[2], padR = _c[3]; if (aspect < targetAspect) { padT = 0; padB = 0; padL = Math.round(0.5 * (targetAspect * height - width)); padR = Math.round(0.5 * (targetAspect * height - width)); } else { padT = Math.round(0.5 * ((1.0 / targetAspect) * width - height)); padB = Math.round(0.5 * ((1.0 / targetAspect) * width - height)); padL = 0; padR = 0; } var resized = tf.tidy(function () { var imageTensor = toInputTensor(input); imageTensor = tf.pad3d(imageTensor, [[padT, padB], [padL, padR], [0, 0]]); return imageTensor.resizeBilinear([targetH, targetW]); }); return { resized: resized, padding: { top: padT, left: padL, right: padR, bottom: padB } }; } exports.padAndResizeTo = padAndResizeTo; function toTensorBuffers3D(tensors) { return __awaiter(this, void 0, void 0, function () { return __generator(this, function (_a) { return [2, Promise.all(tensors.map(function (tensor) { return tensor.buffer(); }))]; }); }); } exports.toTensorBuffers3D = toTensorBuffers3D; function scalePose(pose, scaleY, scaleX, offsetY, offsetX) { if (offsetY === void 0) { offsetY = 0; } if (offsetX === void 0) { offsetX = 0; } return { score: pose.score, keypoints: pose.keypoints.map(function (_a) { var score = _a.score, part = _a.part, position = _a.position; return ({ score: score, part: part, position: { x: position.x * scaleX + offsetX, y: position.y * scaleY + offsetY } }); }) }; } exports.scalePose = scalePose; function scalePoses(poses, scaleY, scaleX, offsetY, offsetX) { if (offsetY === void 0) { offsetY = 0; } if (offsetX === void 0) { offsetX = 0; } if (scaleX === 1 && scaleY === 1 && offsetY === 0 && offsetX === 0) { return poses; } return poses.map(function (pose) { return scalePose(pose, scaleY, scaleX, offsetY, offsetX); }); } exports.scalePoses = scalePoses; function flipPoseHorizontal(pose, imageWidth) { return { score: pose.score, keypoints: pose.keypoints.map(function (_a) { var score = _a.score, part = _a.part, position = _a.position; return ({ score: score, part: part, position: { x: imageWidth - 1 - position.x, y: position.y } }); }) }; } exports.flipPoseHorizontal = flipPoseHorizontal; function flipPosesHorizontal(poses, imageWidth) { if (imageWidth <= 0) { return poses; } return poses.map(function (pose) { return flipPoseHorizontal(pose, imageWidth); }); } exports.flipPosesHorizontal = flipPosesHorizontal; function scaleAndFlipPoses(poses, _a, _b, padding, flipHorizontal) { var height = _a[0], width = _a[1]; var inputResolutionHeight = _b[0], inputResolutionWidth = _b[1]; var scaleY = (height + padding.top + padding.bottom) / (inputResolutionHeight); var scaleX = (width + padding.left + padding.right) / (inputResolutionWidth); var scaledPoses = scalePoses(poses, scaleY, scaleX, -padding.top, -padding.left); if (flipHorizontal) { return flipPosesHorizontal(scaledPoses, width); } else { return scaledPoses; } } exports.scaleAndFlipPoses = scaleAndFlipPoses; //# sourceMappingURL=util.js.map