license import * as tf from '@tensorflow/tfjs' ; import { NUM_KEYPOINTS } from '../keypoints' ; import { Vector2D } from '../types' ; export function getPointsConfidence ( heatmapScores , heatMapCoords ) { const numKeypoints = heatMapCoords . shape [ 0 ] ; const result = new Float32Array ( numKeypoints ) ; for ( let keypoint = 0 ; keypoint < numKeypoints ; keypoint ++ ) { const y = heatMapCoords . get ( keypoint , 0 ) ; const x = heatMapCoords . get ( keypoint , 1 ) ; result [ keypoint ] = heatmapScores . get ( y , x , keypoint ) ; } return result ; } function getOffsetPoint ( y , x , keypoint , offsetsBuffer ) { return { y : offsetsBuffer . get ( y , x , keypoint ) , x : offsetsBuffer . get ( y , x , keypoint + NUM_KEYPOINTS ) } ; } export function getOffsetVectors ( heatMapCoordsBuffer , offsetsBuffer ) { const result = [ ] ; for ( let keypoint = 0 ; keypoint < NUM_KEYPOINTS ; keypoint ++ ) { const heatmapY = heatMapCoordsBuffer . get ( keypoint , 0 ) . valueOf ( ) ; const heatmapX = heatMapCoordsBuffer . get ( keypoint , 1 ) . valueOf ( ) ; const { x , y } = getOffsetPoint ( heatmapY , heatmapX , keypoint , offsetsBuffer ) ; result . push ( y ) ; result . push ( x ) ; } return tf . tensor2d ( result , [ NUM_KEYPOINTS , 2 ] ) ; } export function getOffsetPoints ( heatMapCoordsBuffer , outputStride , offsetsBuffer ) { return tf . tidy ( ( ) => { const offsetVectors = getOffsetVectors ( heatMapCoordsBuffer , offsetsBuffer ) ; return heatMapCoordsBuffer . toTensor ( ) . mul ( tf . scalar ( outputStride , 'int32' ) ) . toFloat ( ) . add ( offsetVectors ) ; } ) ; }