license import * as tf from '@tensorflow/tfjs' ; import { Keypoint , Pose } from '../types' ; import { toTensorBuffers3D } from '../util' ; import { buildPartWithScoreQueue } from './buildPartWithScoreQueue' ; import { decodePose } from './decodePose' ; import { getImageCoords , squaredDistance } from './util' ; function withinNmsRadiusOfCorrespondingPoint ( poses , squaredNmsRadius , { x , y } , keypointId ) { return poses . some ( ( { keypoints } ) => { const correspondingKeypoint = keypoints [ keypointId ] . position ; return squaredDistance ( y , x , correspondingKeypoint . y , correspondingKeypoint . x ) <= squaredNmsRadius ; } ) ; } function getInstanceScore ( existingPoses , squaredNmsRadius , instanceKeypoints ) { let notOverlappedKeypointScores = instanceKeypoints . reduce ( ( result , { position , score } , keypointId ) : number => { if ( ! withinNmsRadiusOfCorrespondingPoint ( existingPoses , squaredNmsRadius , position , keypointId ) ) { result += score ; } return result ; } , 0.0 ) ; return notOverlappedKeypointScores /= instanceKeypoints . length ; } const kLocalMaximumRadius = 1 ; param heatmapScores param offsets param displacementsFwd param displacementsBwd param outputStride param maxPoseDetections param scoreThreshold param nmsRadius return export default async function decodeMultiplePoses ( heatmapScores , offsets , displacementsFwd , displacementsBwd , outputStride , maxPoseDetections , scoreThreshold = 0.5 , nmsRadius = 20 ) { const poses = [ ] ; const [ scoresBuffer , offsetsBuffer , displacementsFwdBuffer , displacementsBwdBuffer ] = await toTensorBuffers3D ( [ heatmapScores , offsets , displacementsFwd , displacementsBwd ] ) ; const queue = buildPartWithScoreQueue ( scoreThreshold , kLocalMaximumRadius , scoresBuffer ) ; const squaredNmsRadius = nmsRadius * nmsRadius ; while ( poses . length < maxPoseDetections && ! queue . empty ( ) ) { const root = queue . dequeue ( ) ; const rootImageCoords = getImageCoords ( root . part , outputStride , offsetsBuffer ) ; if ( withinNmsRadiusOfCorrespondingPoint ( poses , squaredNmsRadius , rootImageCoords , root . part . id ) ) { continue ; } const keypoints = decodePose ( root , scoresBuffer , offsetsBuffer , outputStride , displacementsFwdBuffer , displacementsBwdBuffer ) ; const score = getInstanceScore ( poses , squaredNmsRadius , keypoints ) ; poses . push ( { keypoints , score } ) ; } return poses ; }