Last active
September 1, 2020 03:37
Star
You must be signed in to star a gist
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
{ | |
"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" | |
} | |
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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