import { WorkerCommand , WorkerResponse } from "./const" ; import { PoseNetConfig , PoseNetFunctionType , PoseNetOperatipnParams , } from "./const" ; import * as poseNet from "@tensorflow-models/posenet" ; import * as tf from "@tensorflow/tfjs" ; const ctx = self as any ; let model ; const predict = async ( image , config , params ) : Promise < poseNet . Pose [ ] > => { const processWidth = image . width ; const processHeight = image . height ; const offscreen = new OffscreenCanvas ( processWidth , processHeight ) ; const ctx = offscreen . getContext ( "2d" ) ! ; ctx . drawImage ( image , 0 , 0 , processWidth , processHeight ) ; const newImg = ctx . getImageData ( 0 , 0 , processWidth , processHeight ) ; if ( params . type === PoseNetFunctionType . SinglePerson ) { const prediction = await model ! . estimateSinglePose ( newImg , params . singlePersonParams ! ) ; return [ prediction ] ; } else if ( params . type === PoseNetFunctionType . MultiPerson ) { const prediction = await model ! . estimateMultiplePoses ( newImg , params . multiPersonParams ! ) ; return prediction ; } else { const prediction = await model ! . estimateMultiplePoses ( newImg , params . multiPersonParams ! ) ; return prediction ; } } ; const predictForSafari = async ( data , width , height , config , params ) => { console . log ( "for safari" ) ; const newImg = new ImageData ( new Uint8ClampedArray ( data ) , width , height ) ; if ( params . type === PoseNetFunctionType . SinglePerson ) { const prediction = await model ! . estimateSinglePose ( newImg , params . singlePersonParams ! ) ; return [ prediction ] ; } else if ( params . type === PoseNetFunctionType . MultiPerson ) { const prediction = await model ! . estimateMultiplePoses ( newImg , params . multiPersonParams ! ) ; return prediction ; } else { const prediction = await model ! . estimateMultiplePoses ( newImg , params . multiPersonParams ! ) ; return prediction ; } } ; onmessage = async ( event ) => { if ( event . data . message === WorkerCommand . INITIALIZE ) { console . log ( "Initialize model!." , event ) ; await tf . ready ( ) ; tf . env ( ) . set ( "WEBGL_CPU_FORWARD" , false ) ; model = await poseNet . load ( event . data . config . model ) ; ctx . postMessage ( { message : WorkerResponse . INITIALIZED } ) ; } else if ( event . data . message === WorkerCommand . PREDICT ) { const config = event . data . config ; const params = event . data . params ; const image = event . data . image ; const data = event . data . data ; const width = event . data . width ; const height = event . data . height ; const uid = event . data . uid ; const functionType = event . data . functionType ; if ( params . type === PoseNetFunctionType . SinglePerson ) { let prediction ; if ( data ) { prediction = await predictForSafari ( data , width , height , config , params ) ; ctx . postMessage ( { message : WorkerResponse . PREDICTED , uid : uid , prediction : prediction , } ) ; } else { prediction = await predict ( image , config , params ) ; ctx . postMessage ( { message : WorkerResponse . PREDICTED , uid : uid , prediction : prediction , } ) ; image . close ( ) ; } } else if ( params . type === PoseNetFunctionType . MultiPerson ) { let prediction ; if ( data ) { prediction = await predictForSafari ( data , width , height , config , params ) ; ctx . postMessage ( { message : WorkerResponse . PREDICTED , uid : uid , prediction : prediction , } ) ; } else { prediction = await predict ( image , config , params ) ; ctx . postMessage ( { message : WorkerResponse . PREDICTED , uid : uid , prediction : prediction , } ) ; image . close ( ) ; } } else { console . log ( "not implemented" , functionType ) ; } } } ;