import { WorkerResponse , WorkerCommand , CartoonConfig , CartoonOperatipnParams , CartoonFunctionType , } from "./const" ; import { getBrowserType , BrowserType } from "./BrowserUtil" ; import * as tf from "@tensorflow/tfjs" ; import { setWasmPath } from "@tensorflow/tfjs-backend-wasm" ; export { CartoonConfig , CartoonOperatipnParams } from "./const" ; export { BrowserType , getBrowserType } from "./BrowserUtil" ; import workerJs from "worker-loader?inline=no-fallback!./white-box-cartoonization-worker-worker.ts" ; import modelJson from "../resources/white-box-cartoonization/model.json" ; import modelWeight from "../resources/white-box-cartoonization/group1-shard1of1.bin" ; export const generateCartoonDefaultConfig = ( ) : CartoonConfig => { const defaultConf = { browserType : getBrowserType ( ) , processOnLocal : false , useTFWasmBackend : false , wasmPath : "/tfjs-backend-wasm.wasm" , pageUrl : window . location . href , modelJson : modelJson , modelWeight : modelWeight , } ; return defaultConf ; } ; export const generateDefaultCartoonParams = ( ) : CartoonOperatipnParams => { const defaultParams = { type : CartoonFunctionType . Cartoon , processWidth : 320 , processHeight : 320 , } ; return defaultParams ; } ; 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 ) { 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 class LocalCT { model : tf . GraphModel | null = null ; canvas = document . createElement ( "canvas" ) ; init = async ( config ) => { await load_module ( config ) ; await tf . ready ( ) ; tf . env ( ) . set ( "WEBGL_CPU_FORWARD" , false ) ; const modelJson2 = new File ( [ config . modelJson ] , "model.json" , { type : "application/json" , } ) ; const b = Buffer . from ( config . modelWeight . split ( "," ) [ 1 ] , "base64" ) ; const modelWeights = new File ( [ b ] , "group1-shard1of1.bin" ) ; this . model = await tf . loadGraphModel ( tf . io . browserFiles ( [ modelJson2 , modelWeights ] ) ) ; } ; predict = async ( targetCanvas , config , params ) => { this . canvas . width = params . processWidth ; this . canvas . height = params . processHeight ; const ctx = this . canvas . getContext ( "2d" ) ! ; ctx . drawImage ( targetCanvas , 0 , 0 , this . canvas . width , this . canvas . height ) ; tf . tidy ( ( ) => { let tensor = tf . browser . fromPixels ( this . canvas ) ; tensor = tf . sub ( tensor . expandDims ( 0 ) . div ( 127.5 ) , 1 ) ; let prediction = this . model ! . predict ( tensor ) as tf . Tensor ; const alpha = tf . ones ( [ 1 , params . processWidth , params . processHeight , 1 , ] ) ; prediction = tf . concat ( [ prediction , alpha ] , 3 ) ; prediction = tf . add ( prediction , 1 ) ; prediction = tf . mul ( prediction , 127.5 ) ; prediction = prediction . flatten ( ) ; prediction = tf . cast ( prediction , "int32" ) ; prediction = tf . squeeze ( prediction as tf . Tensor ) ; let imgArray = prediction . arraySync ( ) as number [ ] ; let imgArray2 = new Uint8ClampedArray ( imgArray . length ) ; imgArray2 . set ( imgArray ) ; const outputImage = new ImageData ( imgArray2 , this . canvas . width , this . canvas . height ) ; ctx . putImageData ( outputImage , 0 , 0 ) ; } ) ; return this . canvas ; } ; } export class CartoonWorkerManager { private workerCT : Worker | null = null ; private canvasOut = document . createElement ( "canvas" ) ; private canvas = document . createElement ( "canvas" ) ; private config = generateCartoonDefaultConfig ( ) ; private localCT = new LocalCT ( ) ; init = async ( config ) => { if ( config != null ) { this . config = config ; } if ( this . workerCT ) { this . workerCT . terminate ( ) ; } this . workerCT = null ; if ( this . config . processOnLocal == true ) { await this . localCT . init ( this . config ! ) ; return ; } const workerCT = new workerJs ( ) ; workerCT ! . postMessage ( { message : WorkerCommand . INITIALIZE , config : this . config , } ) ; const p = new Promise < void > ( ( onResolve , onFail ) => { workerCT ! . onmessage = ( event ) => { if ( event . data . message === WorkerResponse . INITIALIZED ) { console . log ( "WORKERSS INITIALIZED" ) ; this . workerCT = workerCT ; onResolve ( ) ; } else { console . log ( "cartoon Initialization something wrong.." ) ; onFail ( event ) ; } } ; } ) ; return p ; } ; predict = async ( targetCanvas , params = generateDefaultCartoonParams ( ) ) => { if ( this . config . processOnLocal == true ) { const convertedCanvas = await this . localCT . predict ( targetCanvas , this . config , params ) ; this . canvasOut . width = targetCanvas . width ; this . canvasOut . height = targetCanvas . height ; const ctx = this . canvasOut . getContext ( "2d" ) ! ; ctx . drawImage ( convertedCanvas , 0 , 0 , this . canvasOut . width , this . canvasOut . height ) ; return this . canvasOut ; } if ( ! this . workerCT ) { return null ; } if ( this . config . browserType === BrowserType . SAFARI ) { this . canvas . width = params . processWidth ; this . canvas . height = params . processHeight ; const ctx = this . canvas . getContext ( "2d" ) ! ; ctx . drawImage ( targetCanvas , 0 , 0 , this . canvas . width , this . canvas . height ) ; const imageData = ctx . getImageData ( 0 , 0 , this . canvas . width , this . canvas . height ) ; const dataArray = imageData . data ; const width = imageData . width ; const height = imageData . height ; const uid = performance . now ( ) ; this . workerCT ! . postMessage ( { message : WorkerCommand . PREDICT , uid : uid , config : this . config , params : params , data : dataArray , } , [ dataArray . buffer ] ) ; const p = new Promise ( ( onResolve , onFail ) => { this . workerCT ! . onmessage = ( event ) => { if ( event . data . message === WorkerResponse . PREDICTED && event . data . uid === uid ) { const outDataArray = event . data . converted ; const outImageData = new ImageData ( outDataArray , imageData . width , imageData . height ) ; this . canvas . width = imageData . width ; this . canvas . height = imageData . height ; this . canvas . getContext ( "2d" ) ! . putImageData ( outImageData , 0 , 0 ) ; this . canvasOut . width = targetCanvas . width ; this . canvasOut . height = targetCanvas . height ; const ctx = this . canvasOut . getContext ( "2d" ) ! ; ctx . drawImage ( this . canvas , 0 , 0 , this . canvasOut . width , this . canvasOut . height ) ; onResolve ( this . canvasOut ) ; } else { console . log ( "cartoon Prediction something wrong.." ) ; } } ; } ) ; return p ; } else { const off = new OffscreenCanvas ( targetCanvas . width , targetCanvas . height ) ; off . getContext ( "2d" ) ! . drawImage ( targetCanvas , 0 , 0 , targetCanvas . width , targetCanvas . height ) ; const imageBitmap = off . transferToImageBitmap ( ) ; const uid = performance . now ( ) ; this . workerCT ! . postMessage ( { message : WorkerCommand . PREDICT , uid : uid , config : this . config , params : params , image : imageBitmap , } , [ imageBitmap ] ) ; const p = new Promise ( ( onResolve , onFail ) => { this . workerCT ! . onmessage = ( event ) => { if ( event . data . message === WorkerResponse . PREDICTED && event . data . uid === uid ) { const imagedata = event . data . converted ; this . canvasOut . width = targetCanvas . width ; this . canvasOut . height = targetCanvas . height ; const ctx = this . canvasOut . getContext ( "2d" ) ! ; ctx . drawImage ( imagedata , 0 , 0 , targetCanvas . width , targetCanvas . height ) ; onResolve ( this . canvasOut ) ; } else { console . log ( "cartoon Prediction something wrong.." ) ; } } ; } ) ; return p ; } } ; }