license import { NumberTuple , partIds , partNames , StringTuple } from '../keypoints' ; import { Keypoint , PartWithScore , TensorBuffer3D , Vector2D } from '../types' ; import { clamp , getOffsetPoint } from './util' ; import { addVectors , getImageCoords } from './util' ; const poseChain = [ [ 'nose' , 'leftEye' ] , [ 'leftEye' , 'leftEar' ] , [ 'nose' , 'rightEye' ] , [ 'rightEye' , 'rightEar' ] , [ 'nose' , 'leftShoulder' ] , [ 'leftShoulder' , 'leftElbow' ] , [ 'leftElbow' , 'leftWrist' ] , [ 'leftShoulder' , 'leftHip' ] , [ 'leftHip' , 'leftKnee' ] , [ 'leftKnee' , 'leftAnkle' ] , [ 'nose' , 'rightShoulder' ] , [ 'rightShoulder' , 'rightElbow' ] , [ 'rightElbow' , 'rightWrist' ] , [ 'rightShoulder' , 'rightHip' ] , [ 'rightHip' , 'rightKnee' ] , [ 'rightKnee' , 'rightAnkle' ] ] ; const parentChildrenTuples = poseChain . map ( ( [ parentJoinName , childJoinName ] ) : NumberTuple => ( [ partIds [ parentJoinName ] , partIds [ childJoinName ] ] ) ) ; const parentToChildEdges = parentChildrenTuples . map ( ( [ , childJointId ] ) => childJointId ) ; const childToParentEdges = parentChildrenTuples . map ( ( [ parentJointId , ] ) => parentJointId ) ; function getDisplacement ( i , point , displacements ) { const numEdges = displacements . shape [ 2 ] / 2 ; return { y : displacements . get ( point . y , point . x , i ) , x : displacements . get ( point . y , point . x , numEdges + i ) } ; } function decode ( point , outputStride , height , width ) { return { y : clamp ( Math . round ( point . y / outputStride ) , 0 , height - 1 ) , x : clamp ( Math . round ( point . x / outputStride ) , 0 , width - 1 ) } ; } /**
 * We get a new keypoint along the `edgeId` for the pose instance, assuming
 * that the position of the `idSource` part is already known. For this, we
 * follow the displacement vector from the source to target part (stored in
 * the `i`-t channel of the displacement tensor).
 */ function traverseToTargetKeypoint ( edgeId , sourceKeypoint , targetKeypointId , scoresBuffer , offsets , outputStride , displacements ) { const [ height , width ] = scoresBuffer . shape ; const sourceKeypointIndeces = decode ( sourceKeypoint . position , outputStride , height , width ) ; const displacement = getDisplacement ( edgeId , sourceKeypointIndeces , displacements ) ; const displacedPoint = addVectors ( sourceKeypoint . position , displacement ) ; const displacedPointIndeces = decode ( displacedPoint , outputStride , height , width ) ; const offsetPoint = getOffsetPoint ( displacedPointIndeces . y , displacedPointIndeces . x , targetKeypointId , offsets ) ; const targetKeypoint = addVectors ( displacedPoint , { x : offsetPoint . x , y : offsetPoint . y } ) ; const targetKeypointIndeces = decode ( targetKeypoint , outputStride , height , width ) ; const score = scoresBuffer . get ( targetKeypointIndeces . y , targetKeypointIndeces . x , targetKeypointId ) ; return { position : targetKeypoint , part : partNames [ targetKeypointId ] , score } ; } return export function decodePose ( root , scores , offsets , outputStride , displacementsFwd , displacementsBwd ) { const numParts = scores . shape [ 2 ] ; const numEdges = parentToChildEdges . length ; const instanceKeypoints = new Array ( numParts ) ; const { part : rootPart , score : rootScore } = root ; const rootPoint = getImageCoords ( rootPart , outputStride , offsets ) ; instanceKeypoints [ rootPart . id ] = { score : rootScore , part : partNames [ rootPart . id ] , position : rootPoint } ; for ( let edge = numEdges - 1 ; edge >= 0 ; -- edge ) { const sourceKeypointId = parentToChildEdges [ edge ] ; const targetKeypointId = childToParentEdges [ edge ] ; if ( instanceKeypoints [ sourceKeypointId ] && ! instanceKeypoints [ targetKeypointId ] ) { instanceKeypoints [ targetKeypointId ] = traverseToTargetKeypoint ( edge , instanceKeypoints [ sourceKeypointId ] , targetKeypointId , scores , offsets , outputStride , displacementsBwd ) ; } } for ( let edge = 0 ; edge < numEdges ; ++ edge ) { const sourceKeypointId = childToParentEdges [ edge ] ; const targetKeypointId = parentToChildEdges [ edge ] ; if ( instanceKeypoints [ sourceKeypointId ] && ! instanceKeypoints [ targetKeypointId ] ) { instanceKeypoints [ targetKeypointId ] = traverseToTargetKeypoint ( edge , instanceKeypoints [ sourceKeypointId ] , targetKeypointId , scores , offsets , outputStride , displacementsFwd ) ; } } return instanceKeypoints ; }