license import * as tf from '@tensorflow/tfjs' ; import { connectedPartIndeces } from './keypoints' ; import { OutputStride } from './mobilenet' ; import { Keypoint , Pose , TensorBuffer3D , Vector2D } from './types' ; function eitherPointDoesntMeetConfidence ( a , b , minConfidence ) { return ( a < minConfidence || b < minConfidence ) ; } export function getAdjacentKeyPoints ( keypoints , minConfidence ) { return connectedPartIndeces . reduce ( ( result , [ leftJoint , rightJoint ] ) : Keypoint [ ] [ ] => { if ( eitherPointDoesntMeetConfidence ( keypoints [ leftJoint ] . score , keypoints [ rightJoint ] . score , minConfidence ) ) { return result ; } result . push ( [ keypoints [ leftJoint ] , keypoints [ rightJoint ] ] ) ; return result ; } , [ ] ) ; } const { NEGATIVE_INFINITY , POSITIVE_INFINITY } = Number ; export function getBoundingBox ( keypoints ) { return keypoints . reduce ( ( { maxX , maxY , minX , minY } , { position : { x , y } } ) => { return { maxX : Math . max ( maxX , x ) , maxY : Math . max ( maxY , y ) , minX : Math . min ( minX , x ) , minY : Math . min ( minY , y ) } ; } , { maxX : NEGATIVE_INFINITY , maxY : NEGATIVE_INFINITY , minX : POSITIVE_INFINITY , minY : POSITIVE_INFINITY } ) ; } export function getBoundingBoxPoints ( keypoints ) { const { minX , minY , maxX , maxY } = getBoundingBox ( keypoints ) ; return [ { x : minX , y : minY } , { x : maxX , y : minY } , { x : maxX , y : maxY } , { x : minX , y : maxY } ] ; } export async function toTensorBuffer < rank extends tf . Rank > ( tensor , type = 'float32' ) { const tensorData = await tensor . data ( ) ; return new tf . TensorBuffer < rank > ( tensor . shape , type , tensorData ) ; } export async function toTensorBuffers3D ( tensors ) { return Promise . all ( tensors . map ( tensor => toTensorBuffer ( tensor , 'float32' ) ) ) ; } export function scalePose ( pose , scaleX , scaleY ) { return { score : pose . score , keypoints : pose . keypoints . map ( ( { score , part , position } ) => ( { score , part , position : { x : position . x * scaleX , y : position . y * scaleY } } ) ) } ; } export function scalePoses ( poses , scaleY , scaleX ) { if ( scaleX === 1 && scaleY === 1 ) { return poses ; } return poses . map ( pose => scalePose ( pose , scaleX , scaleY ) ) ; } export function getValidResolution ( imageScaleFactor , inputDimension , outputStride ) { const evenResolution = inputDimension * imageScaleFactor - 1 ; return evenResolution - ( evenResolution % outputStride ) + 1 ; }