license import * as tf from '@tensorflow/tfjs' ; import { CheckpointLoader } from './checkpoint_loader' ; import { checkpoints } from './checkpoints' ; import { assertValidOutputStride , assertValidScaleFactor , MobileNet , MobileNetMultiplier , OutputStride } from './mobilenet' ; import decodeMultiplePoses from './multiPose/decodeMultiplePoses' ; import decodeSinglePose from './singlePose/decodeSinglePose' ; import { Pose } from './types' ; import { getValidResolution , scalePose , scalePoses } from './util' ; export type PoseNetResolution = 161 | 193 | 257 | 289 | 321 | 353 | 385 | 417 | 449 | 481 | 513 ; export type InputType = ImageData | HTMLImageElement | HTMLCanvasElement | HTMLVideoElement ; function toInputTensor ( input , resizeHeight , resizeWidth , flipHorizontal ) { const imageTensor = tf . fromPixels ( input ) ; if ( flipHorizontal ) { return imageTensor . reverse ( 1 ) . resizeBilinear ( [ resizeHeight , resizeWidth ] ) ; } else { return imageTensor . resizeBilinear ( [ resizeHeight , resizeWidth ] ) ; } } export class PoseNet { mobileNet : MobileNet ; constructor ( mobileNet ) { this . mobileNet = mobileNet ; } param input param outputStride return predictForSinglePose ( input , outputStride = 16 ) { assertValidOutputStride ( outputStride ) ; return tf . tidy ( ( ) => { const mobileNetOutput = this . mobileNet . predict ( input , outputStride ) ; const heatmaps = this . mobileNet . convToOutput ( mobileNetOutput , 'heatmap_2' ) ; const offsets = this . mobileNet . convToOutput ( mobileNetOutput , 'offset_2' ) ; return { heatmapScores : heatmaps . sigmoid ( ) , offsets } ; } ) ; } param input param outputStride return predictForMultiPose ( input , outputStride = 16 ) { return tf . tidy ( ( ) => { const mobileNetOutput = this . mobileNet . predict ( input , outputStride ) ; const heatmaps = this . mobileNet . convToOutput ( mobileNetOutput , 'heatmap_2' ) ; const offsets = this . mobileNet . convToOutput ( mobileNetOutput , 'offset_2' ) ; const displacementFwd = this . mobileNet . convToOutput ( mobileNetOutput , 'displacement_fwd_2' ) ; const displacementBwd = this . mobileNet . convToOutput ( mobileNetOutput , 'displacement_bwd_2' ) ; return { heatmapScores : heatmaps . sigmoid ( ) , offsets , displacementFwd , displacementBwd } ; } ) ; } param input param imageScaleFactor param flipHorizontal . Defaults param outputStride return async estimateSinglePose ( input , imageScaleFactor = 0.5 , flipHorizontal = false , outputStride = 16 ) { assertValidOutputStride ( outputStride ) ; assertValidScaleFactor ( imageScaleFactor ) ; const resizedHeight = getValidResolution ( imageScaleFactor , input . height , outputStride ) ; const resizedWidth = getValidResolution ( imageScaleFactor , input . width , outputStride ) ; const { heatmapScores , offsets } = tf . tidy ( ( ) => { const inputTensor = toInputTensor ( input , resizedHeight , resizedWidth , flipHorizontal ) ; return this . predictForSinglePose ( inputTensor , outputStride ) ; } ) ; const pose = await decodeSinglePose ( heatmapScores , offsets , outputStride ) ; heatmapScores . dispose ( ) ; offsets . dispose ( ) ; const scaleY = input . height / resizedHeight ; const scaleX = input . width / resizedWidth ; return scalePose ( pose , scaleY , scaleX ) ; } param input param imageScaleFactor param flipHorizontal param outputStride param maxDetections param scoreThreshold param nmsRadius return async estimateMultiplePoses ( input , imageScaleFactor = 0.5 , flipHorizontal = false , outputStride = 16 , maxDetections = 5 , scoreThreshold = .5 , nmsRadius = 20 ) { assertValidOutputStride ( outputStride ) ; assertValidScaleFactor ( imageScaleFactor ) ; const resizedHeight = getValidResolution ( imageScaleFactor , input . height , outputStride ) ; const resizedWidth = getValidResolution ( imageScaleFactor , input . width , outputStride ) ; const { heatmapScores , offsets , displacementFwd , displacementBwd } = tf . tidy ( ( ) => { const inputTensor = toInputTensor ( input , resizedHeight , resizedWidth , flipHorizontal ) ; return this . predictForMultiPose ( inputTensor , outputStride ) ; } ) ; const poses = await decodeMultiplePoses ( heatmapScores , offsets , displacementFwd , displacementBwd , outputStride , maxDetections , scoreThreshold , nmsRadius ) ; heatmapScores . dispose ( ) ; offsets . dispose ( ) ; displacementFwd . dispose ( ) ; displacementBwd . dispose ( ) ; const scaleY = input . height / resizedHeight ; const scaleX = input . width / resizedWidth ; return scalePoses ( poses , scaleY , scaleX ) ; } public dispose ( ) { this . mobileNet . dispose ( ) ; } } param multiplier export async function load ( multiplier = 1.01 ) { if ( tf == null ) { throw new Error ( `Cannot find TensorFlow.js. If you are using a <script> tag, please ` + `also include @tensorflow/tfjs on the page before using this model.` ) ; } const possibleMultipliers = Object . keys ( checkpoints ) ; tf . util . assert ( typeof multiplier === 'number' , `got multiplier type of ${ typeof multiplier } when it should be a ` + `number.` ) ; tf . util . assert ( possibleMultipliers . indexOf ( multiplier . toString ( ) ) >= 0 , `invalid multiplier value of ${ multiplier }.  No checkpoint exists for that ` + `multiplier. Must be one of ${ possibleMultipliers . join ( ',' ) }.` ) ; const checkpoint = checkpoints [ multiplier ] ; const checkpointLoader = new CheckpointLoader ( checkpoint . url ) ; const variables = await checkpointLoader . getAllVariables ( ) ; const mobileNet = new MobileNet ( variables , checkpoint . architecture ) ; return new PoseNet ( mobileNet ) ; }