import { WorkerResponse , WorkerCommand , BodypixFunctionType , BodyPixConfig , ModelConfigMobileNetV1_05 , BodyPixOperatipnParams } from "./const" ; import * as bodyPix from "@tensorflow-models/body-pix" ; import { BrowserType , getBrowserType } from "./BrowserUtil" ; import { SemanticPersonSegmentation } from "@tensorflow-models/body-pix" ; export { ModelConfigResNet , ModelConfigMobileNetV1 , ModelConfigMobileNetV1_05 , BodypixFunctionType , BodyPixConfig , BodyPixOperatipnParams } from "./const" ; export { BrowserType , getBrowserType } from "./BrowserUtil" ; export { SemanticPersonSegmentation , SemanticPartSegmentation , PersonSegmentation , PartSegmentation } from "@tensorflow-models/body-pix" ; export { BodyPixInternalResolution } from "@tensorflow-models/body-pix/dist/types" ; import * as tf from "@tensorflow/tfjs" ; import workerJs from "worker-loader?inline=no-fallback!./bodypix-worker-worker.ts" ; const load_module = async ( config ) => { if ( config . useTFWasmBackend ) { console . log ( "use wasm backend" ) ; require ( "@tensorflow/tfjs-backend-wasm" ) ; await tf . setBackend ( "wasm" ) ; } else { console . log ( "use webgl backend" ) ; require ( "@tensorflow/tfjs-backend-webgl" ) ; await tf . setBackend ( "webgl" ) ; } } ; export const generateBodyPixDefaultConfig = ( ) : BodyPixConfig => { const defaultConf = { browserType : getBrowserType ( ) , model : ModelConfigMobileNetV1_05 , processOnLocal : false , useTFWasmBackend : false , } ; return defaultConf ; } ; export const generateDefaultBodyPixParams = ( ) => { const defaultParams = { type : BodypixFunctionType . SegmentPerson , segmentPersonParams : { flipHorizontal : false , internalResolution : "medium" , segmentationThreshold : 0.7 , maxDetections : 10 , scoreThreshold : 0.3 , nmsRadius : 20 , } , segmentPersonPartsParams : { flipHorizontal : false , internalResolution : "medium" , segmentationThreshold : 0.7 , maxDetections : 10 , scoreThreshold : 0.3 , nmsRadius : 20 , } , segmentMultiPersonParams : { flipHorizontal : false , internalResolution : "medium" , segmentationThreshold : 0.7 , maxDetections : 10 , scoreThreshold : 0.3 , nmsRadius : 20 , minKeypointScore : 0.3 , refineSteps : 10 , } , segmentMultiPersonPartsParams : { flipHorizontal : false , internalResolution : "medium" , segmentationThreshold : 0.7 , maxDetections : 10 , scoreThreshold : 0.3 , nmsRadius : 20 , minKeypointScore : 0.3 , refineSteps : 10 , } , processWidth : 300 , processHeight : 300 , } ; return defaultParams ; } ; class LocalBP { model : bodyPix . BodyPix | null = null ; canvas : HTMLCanvasElement = ( ( ) => { const newCanvas = document . createElement ( "canvas" ) ; newCanvas . style . display = "none" ; return newCanvas ; } ) ( ) ; init = async ( config ) => { await load_module ( config ) ; this . model = await bodyPix . load ( config . model ) ; } ; perf = Array . from ( new Array ( 50 ) , ( ) => 0 ) ; predict = async ( targetCanvas , config , params ) => { const processWidth = params . processWidth <= 0 || params . processHeight <= 0 ? targetCanvas . width : params . processWidth ; const processHeight = params . processWidth <= 0 || params . processHeight <= 0 ? targetCanvas . height : params . processHeight ; this . canvas . width = processWidth ; this . canvas . height = processHeight ; const ctx = this . canvas . getContext ( "2d" ) ! ; ctx . drawImage ( targetCanvas , 0 , 0 , processWidth , processHeight ) ; const newImg = ctx . getImageData ( 0 , 0 , processWidth , processHeight ) ; let prediction ; if ( params . type === BodypixFunctionType . SegmentPerson ) { const start = performance . now ( ) ; prediction = await this . model ! . segmentPerson ( newImg , params . segmentPersonParams ) ; const end = performance . now ( ) ; this . perf . shift ( ) ; this . perf . push ( end - start ) ; const average = this . perf . reduce ( ( p , c ) => { return p + c ; } ) / this . perf . length ; } else if ( params . type === BodypixFunctionType . SegmentPersonParts ) { prediction = await this . model ! . segmentPersonParts ( newImg , params . segmentPersonPartsParams ) ; } else if ( params . type === BodypixFunctionType . SegmentMultiPerson ) { prediction = await this . model ! . segmentMultiPerson ( newImg , params . segmentMultiPersonParams ) ; } else if ( params . type === BodypixFunctionType . SegmentMultiPersonParts ) { prediction = await this . model ! . segmentMultiPersonParts ( newImg , params . segmentMultiPersonPartsParams ) ; } else { prediction = await this . model ! . segmentPerson ( newImg , params . segmentPersonParams ) ; } return prediction ; } ; } export class BodypixWorkerManager { private workerBP : Worker | null = null ; config : BodyPixConfig = generateBodyPixDefaultConfig ( ) ; private localBP = new LocalBP ( ) ; init = async ( config = null ) => { if ( config != null ) { this . config = config ; } if ( this . workerBP ) { this . workerBP . terminate ( ) ; } this . workerBP = null ; if ( this . config . browserType === BrowserType . SAFARI || this . config . processOnLocal == true ) { return new Promise < void > ( ( onResolve , onFail ) => { this . localBP . init ( this . config ! ) . then ( ( ) => { onResolve ( ) ; } ) ; } ) ; } const workerBP = new workerJs ( ) ; workerBP ! . postMessage ( { message : WorkerCommand . INITIALIZE , config : this . config , } ) ; const p = new Promise < void > ( ( onResolve , onFail ) => { workerBP ! . onmessage = ( event ) => { if ( event . data . message === WorkerResponse . INITIALIZED ) { console . log ( "WORKERSS INITIALIZED" ) ; this . workerBP = workerBP ; onResolve ( ) ; } else { console . log ( "Bodypix Initialization something wrong.." ) ; onFail ( event ) ; } } ; } ) ; await p ; return ; } ; predict = async ( targetCanvas , params ) => { if ( this . config . browserType === BrowserType . SAFARI || this . config . processOnLocal == true ) { const prediction = await this . localBP . predict ( targetCanvas , this . config , params ) ; return prediction ; } else { if ( ! this . workerBP ) { return Array . from ( new Array ( params . processWidth ) , ( ) => new Array ( params . processHeight ) . fill ( 1 ) ) ; } 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 ( ) ; const p = new Promise ( ( onResolve , onFail ) => { this . workerBP ! . postMessage ( { message : WorkerCommand . PREDICT , uid : uid , params : params , image : imageBitmap , config : this . config , } , [ imageBitmap ] ) ; this . workerBP ! . onmessage = ( event ) => { if ( event . data . message === WorkerResponse . PREDICTED && event . data . uid === uid ) { onResolve ( event . data . prediction ) ; } else { console . log ( "Bodypix Prediction something wrong.." , event . data . message , WorkerResponse . PREDICTED , event . data . uid , uid ) ; } } ; } ) ; return p ; } } ; } export const createForegroundImage = ( srcCanvas , prediction ) => { const tmpCanvas = document . createElement ( "canvas" ) ; tmpCanvas . width = prediction . width ; tmpCanvas . height = prediction . height ; const imageData = tmpCanvas . getContext ( "2d" ) ! . getImageData ( 0 , 0 , tmpCanvas . width , tmpCanvas . height ) ; const data = imageData . data ; for ( let rowIndex = 0 ; rowIndex < prediction . height ; rowIndex ++ ) { for ( let colIndex = 0 ; colIndex < prediction . width ; colIndex ++ ) { const seg_offset = rowIndex * prediction . width + colIndex ; const pix_offset = ( rowIndex * prediction . width + colIndex ) * 4 ; if ( prediction . data [ seg_offset ] === 0 ) { data [ pix_offset ] = 0 ; data [ pix_offset + 1 ] = 0 ; data [ pix_offset + 2 ] = 0 ; data [ pix_offset + 3 ] = 0 ; } else { data [ pix_offset ] = 255 ; data [ pix_offset + 1 ] = 255 ; data [ pix_offset + 2 ] = 255 ; data [ pix_offset + 3 ] = 255 ; } } } const imageDataTransparent = new ImageData ( data , prediction . width , prediction . height ) ; tmpCanvas . getContext ( "2d" ) ! . putImageData ( imageDataTransparent , 0 , 0 ) ; const outputCanvas = document . createElement ( "canvas" ) ; outputCanvas . width = srcCanvas . width ; outputCanvas . height = srcCanvas . height ; const ctx = outputCanvas . getContext ( "2d" ) ! ; ctx . drawImage ( tmpCanvas , 0 , 0 , outputCanvas . width , outputCanvas . height ) ; ctx . globalCompositeOperation = "source-in" ; ctx . drawImage ( srcCanvas , 0 , 0 , outputCanvas . width , outputCanvas . height ) ; const outputImage = outputCanvas . getContext ( "2d" ) ! . getImageData ( 0 , 0 , outputCanvas . width , outputCanvas . height ) ; tmpCanvas . remove ( ) ; outputCanvas . remove ( ) ; return outputImage ; } ;