@codait/max-human-pose-estimator
Version:
Detect humans in an image and estimate the pose for each person.
552 lines (478 loc) • 13.1 kB
JavaScript
;Object.defineProperty(exports,'__esModule',{value:true});/* globals tf, Image */
const IMAGESIZE = 432;
const computeTargetSize = function (width, height) {
const resizeRatio = IMAGESIZE / Math.max(width, height);
return {
width: Math.round(resizeRatio * width),
height: Math.round(resizeRatio * height)
}
};
const getImageData = function (imageInput) {
{
return Promise.resolve(imageInput)
}
};
const imageToTensor = function (imageData, mirrorImage = false) {
return tf.tidy(() => {
let imgTensor = tf.browser.fromPixels(imageData);
if (mirrorImage) {
imgTensor = imgTensor.reverse(1);
}
const targetSize = computeTargetSize(imgTensor.shape[0], imgTensor.shape[1]);
return imgTensor
.resizeBilinear([targetSize.width, targetSize.height])
.toFloat()
.expandDims()
})
};
/**
* convert image to Tensor input required by the model
*
* @param {HTMLImageElement} imageInput - the image element
* @param {boolean} mirrorImage - horizontally flip image (default: false)
*/
const preprocess = function (imageInput, mirrorImage = false) {
return getImageData(imageInput)
.then(imageData => {
return imageToTensor(imageData, mirrorImage)
})
.then(inputTensor => {
return Promise.resolve(inputTensor)
})
.catch(err => {
console.error(err);
return Promise.reject(err)
})
};/* globals tf */
let modelPath = null;
{
modelPath = `file://${__dirname}/../model/model.json`;
}
let model = null;
let warmed = false;
/**
* load the human pose estimator model
*/
const load = function (initialize) {
if (!model) {
// console.log('loading model...')
// console.time('model load')
return tf.loadGraphModel(modelPath)
.then(m => {
// console.timeEnd('model load')
model = m;
if (istrue(initialize)) {
warmup();
}
return Promise.resolve(model)
})
.catch(err => {
// console.timeEnd('model load')
console.error(err);
return Promise.reject(err)
})
} else if (istrue(initialize) && !warmed) {
warmup();
return Promise.resolve(model)
} else {
return Promise.resolve(model)
}
};
/**
* run the model to get a prediction
*/
const run = function (imageTensor) {
if (!imageTensor) {
console.error('no image provided');
throw new Error('no image provided')
} else if (!model) {
console.error('model not available');
throw new Error('model not available')
} else {
// console.log('running model...')
// console.time('model inference')
const results = model.predict(imageTensor);
// console.timeEnd('model inference')
warmed = true;
return results
}
};
/**
* run inference on the TensorFlow.js model
*/
const inference = function (imageTensor) {
return load(false).then(() => {
try {
const results = run(imageTensor);
return Promise.resolve(results)
.then((result) => {
tf.dispose(imageTensor);
return result
})
} catch (err) {
return Promise.reject(err)
}
})
};
const warmup = function () {
try {
run(tf.zeros([1, 512, 512, 3]));
} catch (err) { }
};
const istrue = function (param) {
return param === null ||
typeof param === 'undefined' ||
(typeof param === 'string' && param.toLowerCase() === 'true') ||
(typeof param === 'boolean' && param)
};const cocoParts = [
'Nose',
'Neck',
'RShoulder',
'RElbow',
'RWrist',
'LShoulder',
'LElbow',
'LWrist',
'RHip',
'RKnee',
'RAnkle',
'LHip',
'LKnee',
'LAnkle',
'REye',
'LEye',
'REar',
'LEar'
];
const cocoPairs = [
[1, 2],
[1, 5],
[2, 3],
[3, 4],
[5, 6],
[6, 7],
[1, 8],
[8, 9],
[9, 10],
[1, 11],
[11, 12],
[12, 13],
[1, 0],
[0, 14],
[14, 16],
[0, 15],
[15, 17],
[2, 16],
[5, 17]
];
const cocoPairsNetwork = [
[12, 13],
[20, 21],
[14, 15],
[16, 17],
[22, 23],
[24, 25],
[0, 1],
[2, 3],
[4, 5],
[6, 7],
[8, 9],
[10, 11],
[28, 29],
[30, 31],
[34, 35],
[32, 33],
[36, 37],
[18, 19],
[26, 27]
];
const cocoColors = [
[255, 0, 0],
[255, 85, 0],
[255, 170, 0],
[255, 255, 0],
[170, 255, 0],
[85, 255, 0],
[0, 255, 0],
[0, 255, 85],
[0, 255, 170],
[0, 255, 255],
[0, 170, 255],
[0, 85, 255],
[0, 0, 255],
[85, 0, 255],
[170, 0, 255],
[255, 0, 255],
[255, 0, 170],
[255, 0, 85]
];/* globals tf */
const HeatMapCount = 19;
const PafMapCount = 38;
const MaxPairCount = 17;
const DIMFACTOR = 8;
const DEFAULTCONFIG = {
nmsWindowSize: 6,
nmsThreshold: 0.001,
localPAFThreshold: 0.141,
partScoreThreshold: 0.247,
pafCountThreshold: 4,
partCountThreshold: 4
};
let cfg = Object.assign({}, DEFAULTCONFIG);
const configuration = function (config) {
cfg = Object.assign({}, DEFAULTCONFIG, (typeof config === 'object' ? config : {}));
};
const estimatePoses = function (heatmapTensor, pafmapTensor) {
return tf.tidy(() => {
const heatmap = heatmapTensor.bufferSync();
const pafmap = pafmapTensor.bufferSync();
// compute possible parts candidates
const partCandidates = computeParts(heatmap);
// compute possible pairs candidates
const pairCandidates = computePairs(pafmap, partCandidates);
// compute possible poses
const poseCandidates = computePoses(partCandidates, pairCandidates);
tf.dispose(heatmapTensor);
tf.dispose(pafmapTensor);
// create the JSON response (with bodyParts, poseLines, etc)
return formatResponse(poseCandidates)
})
};
const computeParts = function (heatmap) {
const height = heatmap.shape[0];
const width = heatmap.shape[1];
const depth = heatmap.shape[2] - 1;
const parts = new Array(depth);
// extract peak parts from heatmap
for (var y = 0; y < height; y++) {
for (var x = 0; x < width; x++) {
for (var d = 0; d < depth; d++) {
if (!parts[d]) {
parts[d] = [];
}
const score = heatmap.get(y, x, d);
if (score > cfg.nmsThreshold && isMaximum(score, y, x, d, heatmap)) {
parts[d].push([y, x, score]);
}
}
}
}
return parts
};
const isMaximum = function (score, y, x, d, heatmap) {
let isMax = true;
const height = heatmap.shape[0];
const width = heatmap.shape[1];
const h1 = Math.max(0, y - cfg.nmsWindowSize);
const h2 = Math.min(height - 1, y + cfg.nmsWindowSize);
const w1 = Math.max(0, x - cfg.nmsWindowSize);
const w2 = Math.min(width - 1, x + cfg.nmsWindowSize);
for (var h = h1; h <= h2; h++) {
for (var w = w1; w <= w2; w++) {
if (score < heatmap.get(h, w, d)) {
isMax = false;
break
}
}
if (!isMax) {
break
}
}
return isMax
};
const computePairs = function (pafmap, parts) {
const pairsFinal = new Array(MaxPairCount);
const pairs = new Array(MaxPairCount);
cocoPairs.forEach((cocopair, i) => {
const part1 = parts[cocopair[0]];
const part2 = parts[cocopair[1]];
pairs[i] = [];
pairsFinal[i] = [];
// connect the parts, score the connection, and find best matching connections
for (var p1 = 0; p1 < part1.length; p1++) {
for (var p2 = 0; p2 < part2.length; p2++) {
const val = getPairScore(part1[p1][1], part1[p1][0], part2[p2][1], part2[p2][0], pafmap, cocoPairsNetwork[i]);
const score = val.score;
const count = val.count;
if (score > cfg.partScoreThreshold && count >= cfg.pafCountThreshold) {
let inserted = false;
for (var l = 0; l < MaxPairCount; l++) {
if (pairs[i][l] && score > pairs[i][l][2]) {
pairs[i].splice(l, 0, [p1, p2, score]);
inserted = true;
break
}
}
if (!inserted) {
pairs[i].push([p1, p2, score]);
}
}
}
}
const added = {};
for (var m = 0; m < pairs[i].length; m++) {
const p = pairs[i][m];
if (!added[`${p[0]}`] && !added[`${p[1]}`]) {
pairsFinal[i].push(p);
added[`${p[0]}`] = 1;
added[`${p[1]}`] = 1;
}
}
});
return pairsFinal
};
const getPairScore = function (x1, y1, x2, y2, pafmap, cpnetwork) {
let count = 0;
let score = 0;
const dx = x2 - x1;
const dy = y2 - y1;
const normVec = Math.sqrt(Math.pow(dx, 2) + Math.pow(dy, 2));
if (normVec >= 0.0001) {
const shape = pafmap.shape;
const vx = dx / normVec;
const vy = dy / normVec;
for (var t = 0; t < 10; t++) {
const tx = Math.round(x1 + (t * dx / 9) + 0.5);
const ty = Math.round(y1 + (t * dy / 9) + 0.5);
if (shape[0] > ty && shape[1] > tx) {
const s = vy * pafmap.get(ty, tx, cpnetwork[1]) +
vx * pafmap.get(ty, tx, cpnetwork[0]);
if (s > cfg.localPAFThreshold) {
count++;
score += s;
}
}
}
}
return {
score: score,
count: count
}
};
const computePoses = function (parts, pairs) {
const humans = [];
cocoPairs.forEach((cocopair, i) => {
const p1 = cocopair[0];
const p2 = cocopair[1];
pairs[i].forEach((pair, j) => {
const ip1 = pair[0];
const ip2 = pair[1];
let merged = false;
// calculate possible bodies from all pairs found
for (var k = 0; k < humans.length; k++) {
const human = humans[k];
if (ip1 === human.coordsIndexSet[p1] || ip2 === human.coordsIndexSet[p2]) {
human.coordsIndexSet[p1] = ip1;
human.coordsIndexSet[p2] = ip2;
human.partsList[p1] = partsJSON(p1, parts[p1][ip1]);
human.partsList[p2] = partsJSON(p2, parts[p2][ip2]);
merged = true;
break
}
}
if (!merged) {
const human = {
partsList: new Array(18),
coordsIndexSet: new Array(18)
};
human.coordsIndexSet[p1] = ip1;
human.coordsIndexSet[p2] = ip2;
human.partsList[p1] = partsJSON(p1, parts[p1][ip1]);
human.partsList[p2] = partsJSON(p2, parts[p2][ip2]);
humans.push(human);
}
});
});
return humans
};
const partsJSON = function (id, coords) {
return {
x: coords[1] ? coords[1] * DIMFACTOR : coords[1],
y: coords[0] ? coords[0] * DIMFACTOR : coords[0],
partName: cocoParts[id],
partId: id,
score: coords[2]
}
};
const formatResponse = function (humans) {
const humansFinal = [];
for (var i = 0; i < humans.length; i++) {
let bodyPartCount = 0;
for (let j = 0; j < HeatMapCount - 1; j++) {
if (humans[i].coordsIndexSet[j]) {
bodyPartCount += 1;
}
}
// only include poses with enough parts
if (bodyPartCount > cfg.partCountThreshold) {
const pList = humans[i].partsList;
const poseLines = [];
const cocoPairsRender = cocoPairs.slice(0, cocoPairs.length - 2);
cocoPairsRender.forEach((pair, idx) => {
if (pList[pair[0]] && pList[pair[1]]) {
poseLines.push([pList[pair[0]].x, pList[pair[0]].y, pList[pair[1]].x, pList[pair[1]].y]);
}
});
humansFinal.push({
humanId: i,
bodyParts: pList,
poseLines: poseLines
});
}
}
return humansFinal
};
/**
* convert model Tensor output to JSON containing body parts and poses lines data
*
* @param {Tensor} inferenceResults - the output from running the model
*/
const postprocess = function (inferenceResults) {
const [heatmapTensor, pafmapTensor] = tf.tidy(() => {
return inferenceResults.unstack()[0].split([HeatMapCount, PafMapCount], 2)
});
return Promise.all([heatmapTensor.array(), pafmapTensor.array()])
.then(maps => {
tf.dispose(inferenceResults);
return Promise.resolve({
heatMap: maps[0],
pafMap: maps[1],
posesDetected: estimatePoses(heatmapTensor, pafmapTensor),
imageSize: {
width: heatmapTensor.shape[1] * DIMFACTOR,
height: heatmapTensor.shape[0] * DIMFACTOR
}
})
})
};const version="0.3.0";{
global.tf = require('@tensorflow/tfjs-node');
}
const processInput = function (inputImage, mirrorImage) {
return preprocess(inputImage, mirrorImage)
};
const loadModel = function (init) {
return load(init)
};
const runInference = function (inputTensor) {
return inference(inputTensor)
};
const processOutput = function (inferenceResults) {
return postprocess(inferenceResults)
};
const predict = function (inputImage, mirrorImage) {
return processInput(inputImage, mirrorImage)
.then(runInference)
.then(processOutput)
.catch(err => {
console.error(err);
})
};
const config = function (config) {
return configuration(config)
};
const cocoUtil = {
parts: cocoParts,
pairs: cocoPairs,
pairsNetwork: cocoPairsNetwork,
colors: cocoColors
};exports.cocoUtil=cocoUtil;exports.config=config;exports.loadModel=loadModel;exports.predict=predict;exports.processInput=processInput;exports.processOutput=processOutput;exports.runInference=runInference;exports.version=version;