1- import { WorkerResponse , WorkerCommand , BodypixFunctionType , BodyPixConfig , ModelConfigMobileNetV1_05 , BodyPixOperatipnParams } from "./const" ;
1+ import { BodyPixConfig , BodyPixOperatipnParams , BodypixFunctionTypes , ModelConfigs } from "./const" ;
22import * as bodyPix from "@tensorflow-models/body-pix" ;
3- import { BrowserType , getBrowserType } from "./BrowserUtil" ;
43import { SemanticPersonSegmentation } from "@tensorflow-models/body-pix" ;
4+ import { WorkerManagerBase , LocalWorker , getBrowserType } from "worker-base" ;
55
6- export { ModelConfigResNet , ModelConfigMobileNetV1 , ModelConfigMobileNetV1_05 , BodypixFunctionType , BodyPixConfig , BodyPixOperatipnParams } from "./const" ;
7- export { BrowserType , getBrowserType } from "./BrowserUtil" ;
6+ export { BodypixFunctionTypes , BodyPixConfig , BodyPixOperatipnParams } from "./const" ;
87export { SemanticPersonSegmentation , SemanticPartSegmentation , PersonSegmentation , PartSegmentation } from "@tensorflow-models/body-pix" ;
9- export { BodyPixInternalResolution } from "@tensorflow-models/body-pix/dist/types" ;
8+ export { BodyPixInternalResolution , BodyPixArchitecture , BodyPixMultiplier , BodyPixOutputStride , BodyPixQuantBytes } from "@tensorflow-models/body-pix/dist/types" ;
109
1110import * as tf from "@tensorflow/tfjs" ;
1211
@@ -28,7 +27,7 @@ const load_module = async (config: BodyPixConfig) => {
2827export const generateBodyPixDefaultConfig = ( ) : BodyPixConfig => {
2928 const defaultConf : BodyPixConfig = {
3029 browserType : getBrowserType ( ) ,
31- model : ModelConfigMobileNetV1_05 ,
30+ model : ModelConfigs . ModelConfigMobileNetV1_05 ,
3231 processOnLocal : false ,
3332 useTFWasmBackend : false ,
3433 } ;
@@ -37,7 +36,7 @@ export const generateBodyPixDefaultConfig = (): BodyPixConfig => {
3736
3837export const generateDefaultBodyPixParams = ( ) => {
3938 const defaultParams : BodyPixOperatipnParams = {
40- type : BodypixFunctionType . SegmentPerson ,
39+ type : BodypixFunctionTypes . SegmentPerson ,
4140 segmentPersonParams : {
4241 flipHorizontal : false ,
4342 internalResolution : "medium" ,
@@ -80,7 +79,7 @@ export const generateDefaultBodyPixParams = () => {
8079 return defaultParams ;
8180} ;
8281
83- class LocalBP {
82+ class LocalBP extends LocalWorker {
8483 model : bodyPix . BodyPix | null = null ;
8584 canvas : HTMLCanvasElement = ( ( ) => {
8685 const newCanvas = document . createElement ( "canvas" ) ;
@@ -93,36 +92,29 @@ class LocalBP {
9392 this . model = await bodyPix . load ( config . model ) ;
9493 } ;
9594
96- perf = Array . from ( new Array ( 50 ) , ( ) => 0 ) ;
97- predict = async ( targetCanvas : HTMLCanvasElement , config : BodyPixConfig , params : BodyPixOperatipnParams ) => {
98- // ImageData作成
99- const processWidth = params . processWidth <= 0 || params . processHeight <= 0 ? targetCanvas . width : params . processWidth ;
100- const processHeight = params . processWidth <= 0 || params . processHeight <= 0 ? targetCanvas . height : params . processHeight ;
95+ predict = async ( config : BodyPixConfig , params : BodyPixOperatipnParams , targetCanvas : HTMLCanvasElement ) => {
96+ // // ImageData作成
97+ // const processWidth = params.processWidth <= 0 || params.processHeight <= 0 ? targetCanvas.width : params.processWidth;
98+ // const processHeight = params.processWidth <= 0 || params.processHeight <= 0 ? targetCanvas.height : params.processHeight;
10199
102- //console.log("process image size:", processWidth, processHeight)
103- this . canvas . width = processWidth ;
104- this . canvas . height = processHeight ;
105- const ctx = this . canvas . getContext ( "2d" ) ! ;
106- ctx . drawImage ( targetCanvas , 0 , 0 , processWidth , processHeight ) ;
107- const newImg = ctx . getImageData ( 0 , 0 , processWidth , processHeight ) ;
100+ // //console.log("process image size:", processWidth, processHeight)
101+ // this.canvas.width = processWidth;
102+ // this.canvas.height = processHeight;
103+ // const ctx = this.canvas.getContext("2d")!;
104+ // ctx.drawImage(targetCanvas, 0, 0, processWidth, processHeight);
105+ // const newImg = ctx.getImageData(0, 0, processWidth, processHeight);
106+
107+ const ctx = targetCanvas . getContext ( "2d" ) ! ;
108+ const newImg = ctx . getImageData ( 0 , 0 , targetCanvas . width , targetCanvas . height ) ;
108109
109110 let prediction ;
110- if ( params . type === BodypixFunctionType . SegmentPerson ) {
111- const start = performance . now ( ) ;
111+ if ( params . type === BodypixFunctionTypes . SegmentPerson ) {
112112 prediction = await this . model ! . segmentPerson ( newImg , params . segmentPersonParams ) ;
113- const end = performance . now ( ) ;
114- this . perf . shift ( ) ;
115- this . perf . push ( end - start ) ;
116- const average =
117- this . perf . reduce ( ( p , c ) => {
118- return p + c ;
119- } ) / this . perf . length ;
120- // console.log("inference average:", average, this.perf.length)
121- } else if ( params . type === BodypixFunctionType . SegmentPersonParts ) {
113+ } else if ( params . type === BodypixFunctionTypes . SegmentPersonParts ) {
122114 prediction = await this . model ! . segmentPersonParts ( newImg , params . segmentPersonPartsParams ) ;
123- } else if ( params . type === BodypixFunctionType . SegmentMultiPerson ) {
115+ } else if ( params . type === BodypixFunctionTypes . SegmentMultiPerson ) {
124116 prediction = await this . model ! . segmentMultiPerson ( newImg , params . segmentMultiPersonParams ) ;
125- } else if ( params . type === BodypixFunctionType . SegmentMultiPersonParts ) {
117+ } else if ( params . type === BodypixFunctionTypes . SegmentMultiPersonParts ) {
126118 prediction = await this . model ! . segmentMultiPersonParts ( newImg , params . segmentMultiPersonPartsParams ) ;
127119 } else {
128120 // segmentPersonに倒す
@@ -132,92 +124,42 @@ class LocalBP {
132124 } ;
133125}
134126
135- export class BodypixWorkerManager {
136- private workerBP : Worker | null = null ;
137-
127+ export class BodypixWorkerManager extends WorkerManagerBase {
138128 config : BodyPixConfig = generateBodyPixDefaultConfig ( ) ;
139- private localBP = new LocalBP ( ) ;
140-
129+ localWorker = new LocalBP ( ) ;
141130 init = async ( config : BodyPixConfig | null = null ) => {
142- if ( config != null ) {
143- this . config = config ;
144- }
145- if ( this . workerBP ) {
146- this . workerBP . terminate ( ) ;
147- }
148- this . workerBP = null ;
149-
150- if ( this . config . browserType === BrowserType . SAFARI || this . config . processOnLocal == true ) {
151- // safariはwebworkerでWebGLが使えないのでworkerは使わない。
152- return new Promise < void > ( ( onResolve , onFail ) => {
153- this . localBP . init ( this . config ! ) . then ( ( ) => {
154- onResolve ( ) ;
155- } ) ;
156- } ) ;
157- }
158-
159- const workerBP : Worker = new workerJs ( ) ;
160-
161- workerBP ! . postMessage ( {
162- message : WorkerCommand . INITIALIZE ,
163- config : this . config ,
164- } ) ;
165- const p = new Promise < void > ( ( onResolve , onFail ) => {
166- workerBP ! . onmessage = ( event ) => {
167- if ( event . data . message === WorkerResponse . INITIALIZED ) {
168- console . log ( "WORKERSS INITIALIZED" ) ;
169- this . workerBP = workerBP ;
170- onResolve ( ) ;
171- } else {
172- console . log ( "Bodypix Initialization something wrong.." ) ;
173- onFail ( event ) ;
174- }
175- } ;
176- } ) ;
177- await p ;
131+ this . config = config || generateBodyPixDefaultConfig ( ) ;
132+ await this . initCommon (
133+ {
134+ useWorkerForSafari : false ,
135+ processOnLocal : this . config . processOnLocal ,
136+ localWorker : ( ) => {
137+ return new workerJs ! ( ) ;
138+ } ,
139+ workerJs : ( ) => {
140+ return new workerJs ( ) ;
141+ } ,
142+ } ,
143+ config
144+ ) ;
178145 return ;
179146 } ;
180147
181- predict = async ( targetCanvas : HTMLCanvasElement , params : BodyPixOperatipnParams ) => {
182- if ( this . config . browserType === BrowserType . SAFARI || this . config . processOnLocal == true ) {
183- const prediction = await this . localBP . predict ( targetCanvas , this . config , params ) ;
148+ predict = async ( params : BodyPixOperatipnParams , targetCanvas : HTMLCanvasElement ) => {
149+ if ( ! this . worker ) {
150+ const resizedCanvas = this . generateTargetCanvas ( targetCanvas , params . processWidth , params . processHeight ) ;
151+ const prediction = await this . localWorker . predict ( this . config , params , resizedCanvas ) ;
184152 return prediction ;
185- } else {
186- if ( ! this . workerBP ) {
187- return Array . from ( new Array ( params . processWidth ) , ( ) => new Array ( params . processHeight ) . fill ( 1 ) ) ;
188- }
189- const offscreen = new OffscreenCanvas ( targetCanvas . width , targetCanvas . height ) ;
190- const offctx = offscreen . getContext ( "2d" ) ! ;
191- offctx . drawImage ( targetCanvas , 0 , 0 , targetCanvas . width , targetCanvas . height ) ;
192- const imageBitmap = offscreen . transferToImageBitmap ( ) ;
193- const uid = performance . now ( ) ;
194- const p = new Promise ( ( onResolve , onFail ) => {
195- this . workerBP ! . postMessage (
196- {
197- message : WorkerCommand . PREDICT ,
198- uid : uid ,
199- params : params ,
200- image : imageBitmap ,
201- config : this . config ,
202- } ,
203- [ imageBitmap ]
204- ) ;
205-
206- this . workerBP ! . onmessage = ( event ) => {
207- if ( event . data . message === WorkerResponse . PREDICTED && event . data . uid === uid ) {
208- onResolve ( event . data . prediction ) ;
209- } else {
210- console . log ( "Bodypix Prediction something wrong.." , event . data . message , WorkerResponse . PREDICTED , event . data . uid , uid ) ;
211- // onFail(event)
212- }
213- } ;
214- } ) ;
215- return p ;
216153 }
154+ const imageBitmap = this . generateImageBitmap ( targetCanvas , params . processWidth , params . processHeight ) ;
155+ const prediction = ( await this . sendToWorker ( this . config , params , imageBitmap ) ) as bodyPix . SemanticPersonSegmentation | bodyPix . SemanticPartSegmentation | bodyPix . PersonSegmentation [ ] | bodyPix . PartSegmentation [ ] ;
156+ return prediction ;
217157 } ;
218158}
219159
220- //// Utility for Demo
160+ /////////////////////////
161+ //// Utility for Demo ///
162+ /////////////////////////
221163export const createForegroundImage = ( srcCanvas : HTMLCanvasElement , prediction : SemanticPersonSegmentation ) => {
222164 const tmpCanvas = document . createElement ( "canvas" ) ;
223165 tmpCanvas . width = prediction . width ;
0 commit comments