import { PoseNetConfig , ModelConfigResNet50 , WorkerCommand , WorkerResponse , PoseNetFunctionType , PoseNetOperatipnParams , } from "./const" ; import { getBrowserType , BrowserType } from "./BrowserUtil" ; import * as poseNet from "@tensorflow-models/posenet" ; import * as tf from "@tensorflow/tfjs" ; export { Pose , getAdjacentKeyPoints } from "@tensorflow-models/posenet" ; export { ModelConfigResNet50 , ModelConfigMobileNetV1 , PoseNetOperatipnParams , PoseNetConfig , PoseNetFunctionType , } from "./const" ; export const generatePoseNetDefaultConfig = ( ) : PoseNetConfig => { const defaultConf = { browserType : getBrowserType ( ) , model : ModelConfigResNet50 , processOnLocal : false , useTFWasmBackend : false , wasmPath : "/tfjs-backend-wasm.wasm" , workerPath : "./posenet-worker-worker.js" , } ; if ( defaultConf . browserType == BrowserType . SAFARI ) { defaultConf . processOnLocal = true ; } return defaultConf ; } ; export const generateDefaultPoseNetParams = ( ) => { const defaultParams = { type : PoseNetFunctionType . SinglePerson , singlePersonParams : { flipHorizontal : false , } , multiPersonParams : { flipHorizontal : false , maxDetections : 5 , scoreThreshold : 0.5 , nmsRadius : 20 , } , } ; return defaultParams ; } ; import workerJs from "worker-loader?inline=no-fallback!./posenet-worker-worker.ts" ; class LocalPN { model : poseNet . PoseNet | null = null ; canvas : HTMLCanvasElement = document . createElement ( "canvas" ) ; init = ( config ) => { return poseNet . load ( config . model ) . then ( ( res ) => { console . log ( "posenet loaded locally" , config ) ; this . model = res ; return ; } ) ; } ; predict = async ( canvas , config , params ) : Promise < poseNet . Pose [ ] > => { console . log ( "current backend[main thread]:" , tf . getBackend ( ) ) ; const processWidth = canvas . width ; const processHeight = canvas . height ; this . canvas . width = processWidth ; this . canvas . height = processHeight ; const ctx = canvas . getContext ( "2d" ) ! ; ctx . drawImage ( canvas , 0 , 0 , processWidth , processHeight ) ; const newImg = ctx . getImageData ( 0 , 0 , processWidth , processHeight ) ; if ( params . type === PoseNetFunctionType . SinglePerson ) { const prediction = await this . model ! . estimateSinglePose ( newImg , params . singlePersonParams ! ) ; return [ prediction ] ; } else if ( params . type === PoseNetFunctionType . MultiPerson ) { const prediction = await this . model ! . estimateMultiplePoses ( newImg , params . multiPersonParams ! ) ; return prediction ; } else { const prediction = await this . model ! . estimateMultiplePoses ( newImg , params . multiPersonParams ! ) ; return prediction ; } } ; } export class PoseNetWorkerManager { private workerPN : Worker | null = null ; private config : PoseNetConfig = generatePoseNetDefaultConfig ( ) ; private localPN = new LocalPN ( ) ; private canvas = document . createElement ( "canvas" ) ; init = async ( config = null ) => { if ( config != null ) { this . config = config ; } if ( this . workerPN ) { this . workerPN . terminate ( ) ; } this . workerPN = null ; if ( this . config . processOnLocal === true ) { await this . localPN . init ( this . config ! ) ; return ; } const workerPN = new workerJs ( ) ; workerPN ! . postMessage ( { message : WorkerCommand . INITIALIZE , config : this . config , } ) ; const p = new Promise < void > ( ( onResolve , onFail ) => { workerPN ! . onmessage = ( event ) => { if ( event . data . message === WorkerResponse . INITIALIZED ) { console . log ( "WORKERSS INITIALIZED" ) ; this . workerPN = workerPN ; onResolve ( ) ; } else { console . log ( "Bodypix Initialization something wrong.." ) ; onFail ( event ) ; } } ; } ) ; return p ; } ; predict = async ( targetCanvas , params = generateDefaultPoseNetParams ( ) ) => { if ( this . config . processOnLocal === true ) { const prediction = await this . localPN . predict ( targetCanvas , this . config , params ) ; return prediction ; } if ( ! this . workerPN ) { return null ; } if ( this . config . browserType === BrowserType . SAFARI ) { const data = targetCanvas . getContext ( "2d" ) ! . getImageData ( 0 , 0 , targetCanvas . width , targetCanvas . height ) . data ; const uid = performance . now ( ) ; const p = new Promise ( async ( onResolve , onFail ) => { this . workerPN ! . postMessage ( { message : WorkerCommand . PREDICT , uid : uid , config : this . config , params : params , data : data , width : targetCanvas . width , height : targetCanvas . height , } , [ data . buffer ] ) ; this . workerPN ! . onmessage = ( event ) => { if ( event . data . message === WorkerResponse . PREDICTED && event . data . uid === uid ) { onResolve ( event . data . prediction ) ; } else { console . log ( `Posenet Prediction something wrong.. ${ event . data . message }, ${ event . data . uid }, ${ uid }` ) ; } } ; } ) ; return p ; } else { const offscreen = new OffscreenCanvas ( targetCanvas . width , targetCanvas . height ) ; const offctx = offscreen . getContext ( "2d" ) ! ; offctx . drawImage ( targetCanvas , 0 , 0 , targetCanvas . width , targetCanvas . height ) ; const imageBitmap = offscreen . transferToImageBitmap ( ) ; const uid = performance . now ( ) ; this . workerPN ! . postMessage ( { message : WorkerCommand . PREDICT , uid : uid , image : imageBitmap , config : this . config , params : params , } , [ imageBitmap ] ) ; const p = new Promise ( ( onResolve , onFail ) => { this . workerPN ! . onmessage = ( event ) => { if ( event . data . message === WorkerResponse . PREDICTED && event . data . uid === uid ) { onResolve ( event . data . prediction ) ; } else { console . log ( `Posenet Prediction something wrong.. ${ event . data . message }, ${ event . data . uid }, ${ uid }` ) ; } } ; } ) ; return p ; } } ; } const drawPoints = ( canvas , prediction ) => { const keypoints = prediction . keypoints ; for ( let i = 0 ; i < keypoints . length ; i ++ ) { const keypoint = keypoints [ i ] ; const scaleX = 1 ; const scaleY = 1 ; const x = keypoint . position . x ; const y = keypoint . position . y ; const ctx = canvas . getContext ( "2d" ) ! ; ctx . fillStyle = "rgba(0,0,255,0.3)" ; ctx . fillRect ( x * scaleX , y * scaleY , 6 , 6 ) ; } } ; const drawSkeleton = ( canvas , prediction ) => { const adjacentKeyPoints = poseNet . getAdjacentKeyPoints ( prediction . keypoints , 0.0 ) ; const scaleX = 1 ; const scaleY = 1 ; const ctx = canvas . getContext ( "2d" ) ! ; adjacentKeyPoints . forEach ( ( keypoints ) => { ctx . beginPath ( ) ; ctx . moveTo ( keypoints [ 0 ] . position . x * scaleX , keypoints [ 0 ] . position . y * scaleY ) ; ctx . lineTo ( keypoints [ 1 ] . position . x * scaleX , keypoints [ 1 ] . position . y * scaleY ) ; ctx . lineWidth = 6 ; ctx . strokeStyle = "rgba(255,0,0,0.3)" ; ctx . stroke ( ) ; } ) ; } ; export const drawSkeltonAndPoint = ( srcCanvas , prediction ) => { const canvas = document . createElement ( "canvas" ) ; canvas . width = srcCanvas . width ; canvas . height = srcCanvas . height ; prediction . forEach ( ( x ) => { drawPoints ( canvas , x ) ; drawSkeleton ( canvas , x ) ; } ) ; const imageData = canvas . getContext ( "2d" ) ! . getImageData ( 0 , 0 , canvas . width , canvas . height ) ; canvas . remove ( ) ; return imageData ; } ;