mirror of https://github.com/vladmandic/human
40 lines
1.9 KiB
TypeScript
40 lines
1.9 KiB
TypeScript
import { log, join } from '../helpers';
|
|
import * as tf from '../../dist/tfjs.esm.js';
|
|
import * as poses from './poses';
|
|
import * as util from './utils';
|
|
import { Body } from '../result';
|
|
import { Tensor, GraphModel } from '../tfjs/types';
|
|
|
|
let model: GraphModel;
|
|
const poseNetOutputs = ['MobilenetV1/offset_2/BiasAdd'/* offsets */, 'MobilenetV1/heatmap_2/BiasAdd'/* heatmapScores */, 'MobilenetV1/displacement_fwd_2/BiasAdd'/* displacementFwd */, 'MobilenetV1/displacement_bwd_2/BiasAdd'/* displacementBwd */];
|
|
|
|
export async function predict(input, config): Promise<Body[]> {
|
|
const res = tf.tidy(() => {
|
|
if (!model.inputs[0].shape) return [];
|
|
const resized = input.resizeBilinear([model.inputs[0].shape[2], model.inputs[0].shape[1]]);
|
|
const normalized = resized.toFloat().div(127.5).sub(1.0);
|
|
const results: Array<Tensor> = model.execute(normalized, poseNetOutputs) as Array<Tensor>;
|
|
const results3d = results.map((y) => y.squeeze([0]));
|
|
results3d[1] = results3d[1].sigmoid(); // apply sigmoid on scores
|
|
return results3d;
|
|
});
|
|
|
|
const buffers = await Promise.all(res.map((tensor) => tensor.buffer()));
|
|
for (const t of res) t.dispose();
|
|
|
|
const decoded = await poses.decode(buffers[0], buffers[1], buffers[2], buffers[3], config.body.maxDetected, config.body.minConfidence);
|
|
if (!model.inputs[0].shape) return [];
|
|
const scaled = util.scalePoses(decoded, [input.shape[1], input.shape[2]], [model.inputs[0].shape[2], model.inputs[0].shape[1]]) as Body[];
|
|
return scaled;
|
|
}
|
|
|
|
export async function load(config) {
|
|
if (!model) {
|
|
// @ts-ignore type mismatch for GraphModel
|
|
model = await tf.loadGraphModel(join(config.modelBasePath, config.body.modelPath));
|
|
if (!model || !model['modelUrl']) log('load model failed:', config.body.modelPath);
|
|
else if (config.debug) log('load model:', model['modelUrl']);
|
|
} else if (config.debug) log('cached model:', model['modelUrl']);
|
|
return model;
|
|
}
|