import { MODNetConfig , MODNetOperationParams , WorkerCommand , WorkerResponse } from "./const" ; import * as tf from "@tensorflow/tfjs" ; import { BrowserType } from "./BrowserUtil" ; import { setWasmPath } from "@tensorflow/tfjs-backend-wasm" ; const ctx = self as any ; let model ; const load_module = async ( config ) => { const dirname = config . pageUrl . substr ( 0 , config . pageUrl . lastIndexOf ( "/" ) ) ; const wasmPath = `${ dirname }${ config . wasmPath }` ; console . log ( `use wasm backend ${ wasmPath }` ) ; setWasmPath ( wasmPath ) ; if ( config . useTFWasmBackend || config . browserType === BrowserType . SAFARI ) { require ( "@tensorflow/tfjs-backend-wasm" ) ; await tf . setBackend ( "wasm" ) ; } else { console . log ( "use webgl backend" ) ; require ( "@tensorflow/tfjs-backend-webgl" ) ; await tf . setBackend ( "webgl" ) ; } } ; const predict = async ( image , config , params ) => { const off = new OffscreenCanvas ( params . processWidth , params . processHeight ) ; const ctx = off . getContext ( "2d" ) ! ; ctx . drawImage ( image , 0 , 0 , off . width , off . height ) ; const imageData = ctx . getImageData ( 0 , 0 , off . width , off . height ) ; let bm = null ; tf . tidy ( ( ) => { let tensor = tf . browser . fromPixels ( imageData ) ; tensor = tensor . expandDims ( 0 ) ; tensor = tf . cast ( tensor , "float32" ) ; tensor = tensor . div ( tf . max ( tensor ) ) ; tensor = tensor . sub ( 0.485 ) . div ( 0.229 ) ; let prediction = model ! . predict ( tensor ) as tf . Tensor ; bm = prediction . reshape ( [ params . processWidth , params . processHeight ] ) . arraySync ( ) as number [ ] [ ] ; } ) ; return bm ! ; } ; const predictWithData = async ( data , config , params ) => { const imageData = new ImageData ( data , params . processWidth , params . processHeight ) ; let bm = null ; tf . tidy ( ( ) => { let tensor = tf . browser . fromPixels ( imageData ) ; tensor = tensor . expandDims ( 0 ) ; tensor = tf . cast ( tensor , "float32" ) ; tensor = tensor . div ( tf . max ( tensor ) ) ; tensor = tensor . sub ( 0.485 ) . div ( 0.229 ) ; let prediction = model ! . predict ( tensor ) as tf . Tensor ; bm = prediction . reshape ( [ params . processWidth , params . processHeight ] ) . arraySync ( ) as number [ ] [ ] ; } ) ; return bm ! ; } ; onmessage = async ( event ) => { if ( event . data . message === WorkerCommand . INITIALIZE ) { const config = event . data . config as MODNetConfig ; await load_module ( config ) ; tf . ready ( ) . then ( async ( ) => { tf . env ( ) . set ( "WEBGL_CPU_FORWARD" , false ) ; if ( config . modelInputSize === 192 ) { const modelJson = new File ( [ config . modelJson_192 ] , "model.json" , { type : "application/json" } ) ; const weight = Buffer . from ( config . modelWeight_192 . split ( "," ) [ 1 ] , "base64" ) ; const modelWeights = new File ( [ weight ] , "group1-shard1of1.bin" ) ; model = await tf . loadGraphModel ( tf . io . browserFiles ( [ modelJson , modelWeights ] ) ) ; } else if ( config . modelInputSize === 256 ) { const modelJson = new File ( [ config . modelJson_256 ] , "model.json" , { type : "application/json" } ) ; const weight = Buffer . from ( config . modelWeight_256 . split ( "," ) [ 1 ] , "base64" ) ; const modelWeights = new File ( [ weight ] , "group1-shard1of1.bin" ) ; model = await tf . loadGraphModel ( tf . io . browserFiles ( [ modelJson , modelWeights ] ) ) ; } else if ( config . modelInputSize === 512 ) { const modelJson = new File ( [ config . modelJson_512 ] , "model.json" , { type : "application/json" } ) ; const weight = Buffer . from ( config . modelWeight_512 . split ( "," ) [ 1 ] , "base64" ) ; const modelWeights = new File ( [ weight ] , "group1-shard1of1.bin" ) ; model = await tf . loadGraphModel ( tf . io . browserFiles ( [ modelJson , modelWeights ] ) ) ; } else { console . log ( "unknwon process size!" , config . modelInputSize ) ; } console . log ( model ! . inputs ) ; console . log ( model ! . inputNodes ) ; console . log ( model ! . outputs ) ; console . log ( model ! . outputNodes ) ; ctx . postMessage ( { message : WorkerResponse . INITIALIZED } ) ; } ) ; } else if ( event . data . message === WorkerCommand . PREDICT ) { const image = event . data . image ; const data = event . data . data ; const config = event . data . config ; const params = event . data . params ; const uid = event . data . uid ; console . log ( "current backend[worker thread]:" , tf . getBackend ( ) ) ; if ( data ) { console . log ( "Browser SAFARI" ) ; const prediction = await predictWithData ( data , config , params ) ; ctx . postMessage ( { message : WorkerResponse . PREDICTED , uid : uid , prediction : prediction } ) ; } else { const prediction = await predict ( image , config , params ) ; ctx . postMessage ( { message : WorkerResponse . PREDICTED , uid : uid , prediction : prediction } ) ; } } } ;