Skip to content

Instantly share code, notes, and snippets.

@marufeuille
Last active September 1, 2020 03:37
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
Star You must be signed in to star a gist
Save marufeuille/c044c6a6ad4abeeffcb5ab222651d4c6 to your computer and use it in GitHub Desktop.
{
"dependencies": {
"@tensorflow-models/posenet": "^2.2.1",
"@tensorflow/tfjs": "^2.3.0",
"@tensorflow/tfjs-node": "^2.3.0",
"canvas": "^2.6.1",
"mathjs": "^7.2.0",
"typescript": "^4.0.2"
}
}
const tf = require('@tensorflow/tfjs-node');
const posenet = require('@tensorflow-models/posenet');
const fs = require('fs');
const path = require('path');
const { dot, acos, norm } = require('mathjs');
const { createCanvas, loadImage, Canvas, JPEGStream } = require('canvas')
// parameters
const imageScaleFactor = 0.5;
const outputStride = 16;
const flipHorizontal = false;
const image_width = 300;
const image_height = 300;
const image_dir_path = "./img";
const output_dir_path = "./dst/"
/**
* pose estimation for single picture
* @param {PoseNet} net - loadされたPoseNetオブジェクト
* @param {string} path - PoseNetで姿勢推定をしたい画像のパス
* @return {Promise} Promiseオブジェクト。解決されるとdictionaly形式の姿勢情報が返る
*/
const estimate = async (net, path) => {
const img = await loadImage(path)
const canvas = createCanvas(image_width, image_height)
const ctx = canvas.getContext('2d')
ctx.drawImage(img, 0, 0, image_width, image_height)
return net.estimateSinglePose(canvas, imageScaleFactor, flipHorizontal, outputStride)
}
/**
* 不要情報を削除し、必要なもののみに変形する
* @param {Object} poses - PoseNetが出力する姿勢情報のリスト
* @return {Object} 変形された姿勢情報
*/
const filterPoseInfo = async poses => {
const results = []
for (const pose of poses) {
const pose_info = {}
for (const key of pose.keypoints) {
pose_info[key.part] = [key.position.x, key.position.y]
}
results.push(pose_info)
}
return results
}
/**
* jsonファイルを指定したファイル名で保存する
* @param {Object} json - 保存するJSON
* @param {string} filename - 保存するファイル名
*/
const saveFile = (json, filename) => {
fs.writeFileSync(path.join(output_dir_path, filename), JSON.stringify(json));
}
/**
* 2つのベクトルを受け取り、角度を計算します
* @param {Array<Number>} v1 - ベクトル
* @param {Array<Number>} v2 - ベクトル
* @return {Number} 角度
*/
const calcDegree = (v1, v2) => {
return acos(dot(v1, v2) / (norm(v1) * norm(v2))) * 180 / Math.PI
}
/**
* 姿勢情報から左右の肘の角度を算出します
* @param {Object} pose - PoseNetが出力する姿勢情報のリスト
* @return {Array<Number>} 左右の角度の配列
*/
const calcJointAngles = (pose) => {
// left elbow
const l_x0 = pose["leftElbow"][0]
const l_y0 = pose["leftElbow"][1]
const l_v1 = [pose["leftWrist"][0] - l_x0, pose["leftWrist"][1] - l_y0]
const l_v2 = [pose["leftShoulder"][0] - l_x0, pose["leftShoulder"][1] - l_y0]
const leftDeg = calcDegree(l_v1, l_v2)
// right elbow
const r_x0 = pose["rightElbow"][0]
const r_y0 = pose["rightElbow"][1]
const r_v1 = [pose["rightWrist"][0] - r_x0, pose["rightWrist"][1] - r_y0]
const r_v2 = [pose["rightShoulder"][0] - r_x0, pose["rightShoulder"][1] - r_y0]
const rightDeg = calcDegree(r_v1, r_v2)
return [leftDeg, rightDeg]
}
/**
* Entry point
*/
const main = async () => {
net = await posenet.load({
architecture: 'ResNet50',
outputStride: 32,
inputResolution: { width: image_width, height: image_height },
quantBytes: 2
})
fs.readdir(image_dir_path, (err, files) => {
const promisses = []
files = files.filter(filename => filename.endsWith(".jpg"))
for (const file of files) {
const filepath = path.join(image_dir_path, file);
promisses.push(estimate(net, filepath))
}
Promise.all(promisses)
.then(filterPoseInfo)
.then(async (json) => {
saveFile(json, "raw_output.json")
return json
})
.then((json) => {
const angles = []
for (const pose of json) {
angles.push(calcJointAngles(pose))
}
saveFile(angles, "angles_output.json")
})
.catch(err => console.log(err))
})
}
main();
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment