Web worker gpu compilation UI freeze

html/js/owt/owtstream.js

@@ -1,37 +1,106 @@
-const createOWTStream = async () => {
-  stream = await Owt.Base.MediaStreamFactory.createMediaStream(
-    avTrackConstraint
-  )
-  console.log(inputvideo)
-  if ("srcObject" in inputvideo) {
-    inputvideo.srcObject = stream
-  } else {
-    inputvideo.src = URL.createObjectURL(stream)
-  }
+let instanceType = "deeplabtflite";
+let computeTime = 0;
+let computeStart = 0;
+const outputDimensions = [1, 321, 321, 21];
 
-  inputvideo.autoplay = true
-  console.log(inputvideo.srcObject)
-}
+const inputOptions = {
+  mean: [127.5, 127.5, 127.5],
+  std: [127.5, 127.5, 127.5],
+  scaledFlag: false,
+  inputLayout: 'nhwc',
+  inputDimensions: [1, 321, 321, 3], // deeplab
+  inputResolution: [321, 321]
+};
+
+let outputBuffer;
+let continueinputvideo = true;
+console.log("not in worker create context: ", navigator.ml);
 
-let continueinputvideo = true
+const worker = new Worker('../js/tflite/deeplab/webnn/webnnworker.js');
 
 const pipeline2 = buildWebGL2Pipeline(
   inputvideo,
   backgroundImageSource,
   "none",
   [321, 321],
   outputcanvas,
-  null
+  null,
 );
 
+function stopCamera() {
+  stream.getTracks().forEach((track) => {
+    if (track.readyState === 'live' && track.kind === 'video') {
+      track.stop();
+    }
+  });
+}
+
+function renderCamStream() {
+  const inputBuffer = getInputTensor(inputvideo, inputOptions);
+  let obj = {
+    task: 'nncompute',
+    value: inputBuffer,
+  }
+  computeStart = performance.now();
+  worker.postMessage(obj);
+}
+
+async function drawOutput(outputBuffer, srcElement) {
+  if (instanceType.startsWith('deeplab')) {
+    outputBuffer = tf.tidy(() => {
+      const a = tf.tensor(outputBuffer, outputDimensions, 'float32');
+      let axis = 3;
+      if (instanceType === 'deeplabnchw') {
+        axis = 1;
+      }
+      const b = tf.argMax(a, axis);
+      const c = tf.tensor(b.dataSync(), b.shape, 'float32');
+      return c.dataSync();
+    });
+  }
+  console.log('output: ', outputBuffer);
+  outputcanvas.width = srcElement.naturalWidth | srcElement.videoWidth;
+  outputcanvas.height = srcElement.naturalHeight | srcElement.videoHeight;
+  const pipeline = buildWebGL2Pipeline(
+    srcElement,
+    backgroundImageSource,
+    backgroundType,
+    inputOptions.inputResolution,
+    outputcanvas,
+    outputBuffer,
+  );
+  const postProcessingConfig = {
+    smoothSegmentationMask: true,
+    jointBilateralFilter: { sigmaSpace: 1, sigmaColor: 0.1 },
+    coverage: [0.5, 0.75],
+    lightWrapping: 0.3,
+    blendMode: 'screen',
+  }
+  pipeline.updatePostProcessingConfig(postProcessingConfig);
+  await pipeline.render();
+}
+
+const createOWTStream = async () => {
+  stream = await Owt.Base.MediaStreamFactory.createMediaStream(
+    avTrackConstraint
+  );
+  if ("srcObject" in inputvideo) {
+    inputvideo.srcObject = stream;
+  } else {
+    inputvideo.src = URL.createObjectURL(stream);
+  }
+
+  inputvideo.autoplay = true;
+}
+
 const videoCanvasOnFrame = async () => {
-  if(continueinputvideo) {
+  if (continueinputvideo) {
     window.requestAnimationFrame(videoCanvasOnFrame);
     // ctx2d.drawImage(inputvideo, 0, 0, cW, cH);
-    if(stream) {
+    if (stream) {
       const postProcessingConfig2 = {
         smoothSegmentationMask: true,
-        jointBilateralFilter: {sigmaSpace: 1, sigmaColor: 0.1},
+        jointBilateralFilter: { sigmaSpace: 1, sigmaColor: 0.1 },
         coverage: [0.5, 0.75],
         lightWrapping: 0.3,
         blendMode: 'screen',
@@ -44,40 +113,80 @@ const videoCanvasOnFrame = async () => {
 
 
 const oneWebMeetOWT = async () => {
-  await createOWTStream()
-  backgroundType = "blur"
-  continueinputvideo = true
+  await createOWTStream();
+  backgroundType = "blur";
+  continueinputvideo = true;
   await videoCanvasOnFrame();
   getProcessedStream();
   initConference();
   userMarquee();
+  // ssLoad();
 
-  if(parsePathnameBackend().toLowerCase() === 'webnn' && parseSearchParams("ds") === "1") {
-    console.log("^^^^^^^^^^^^^^^^^ WIN ^^^^^^^^^^^^^^^^^^^^")
-    let worker = new Worker('../js/tflite/deeplab/webnn/webnnworker.js');
-    worker.addEventListener('message', (e) => {
-      const { msg } = e.data;
-      $("#worker").html(msg);
-    });
-  } else {
-    ssLoad();
+  let obj = {
+    task: 'nnwarmup',
+    value: null,
   }
+  worker.postMessage(obj);
+
+
+  // // if(parsePathnameBackend().toLowerCase() === 'webnn' && parseSearchParams("ds") === "1") {
+  // if(parsePathnameBackend().toLowerCase() === 'wasm' && parseSearchParams("ds") === "0") {
+  //   console.log("^^^^^^^^^^^^^^^^^ WIN ^^^^^^^^^^^^^^^^^^^^")
+  //   let worker = new Worker('../js/tflite/deeplab/webnn/webnnworker.js');
+  //   worker.postMessage(['loadtask', 'computetask']);
+
+  //   worker.addEventListener('message', (e) => {
+  //     const { msg } = e.data;
+  //     $("#worker").html(msg);
+  //   });
+  // } else {
+  //   ssLoad();
+  // }
 };
 
+worker.addEventListener('message', async (e) => {
+  const { msg } = e.data;
+  if (msg === "modelloaded") {
+    $("#modelloadstatus").html('Model Loaded');
+    $("#tbb").prop('disabled', false);
+    $("#tbr").prop('disabled', false);
+    $("#tbb").removeClass('disabled');
+    $("#tbr").removeClass('disabled');
+  } else {
+    // Start receiving computed outputBuffer
+    outputBuffer = e.data;
+    computeTime = (performance.now() - computeStart).toFixed(2);
+    await drawOutput(outputBuffer, inputvideo);
+    console.log(`  done in ${computeTime} ms.`);
+
+    if (continueAnimating) {
+      renderCamStream();
+      spaninference.html(computeTime);
+      let ct = parseInt(computeTime);
+      $("#fps").html((1000 / ct).toFixed(0));
+    }
+  }
+});
+
 const ssConfig = async (isSS, effect) => {
-  if(isSS && effect) {
+  if (isSS && effect) {
     backgroundImageSource.src = '../../assets/img/ssbg/0.jpg'
     continueinputvideo = false
     console.log(isSS + ' ' + effect)
     backgroundType = effect;
-    continueAnimating = true
-    await ssCompute()
+    continueAnimating = true;
+
+    try {
+      inputvideo.onloadedmediadata = renderCamStream();
+    } catch (error) {
+      console.log(error);
+    }
   } else {
     // gl = outputcanvas.getContext("2d");
     backgroundImageSource.src = '';
     continueAnimating = false;
-    continueinputvideo = true
-    backgroundType = "none"
-    await videoCanvasOnFrame()
+    continueinputvideo = true;
+    backgroundType = "none";
+    await videoCanvasOnFrame();
   }
 }

html/js/tflite/deeplab/webnn/webnnworker.js

@@ -1,27 +1,111 @@
-let netInstance;
-let inputOptions;
+
+importScripts('../tflite_model_runner_cc_simd.js');
+importScripts('../../../webnn-polyfill.js');
+
 let modelRunner;
-let inputBuffer;
-let outputBuffer;
-
-let loadtask = async () => {
-    importScripts("../../utils.js");
-    importScripts("./deeplabv3_mnv2_tflite.js");
-    netInstance = new DeepLabV3MNV2TFLite();
-    inputOptions = netInstance.inputOptions;
-    self.postMessage({msg: "loading model"});
-    modelRunner = await netInstance.load();
-    self.postMessage({msg: "model loaded"});
+
+function sizeOfShape(shape) {
+    return shape.reduce((a, b) => {
+        return a * b;
+    });
+}
+
+let warmup = async () => {
+    const inputDimensions = [1, 321, 321, 3];
+    const inputBuffer = new Float32Array(sizeOfShape(inputDimensions));
+    modelRunner = await load();
+    compute(modelRunner, inputBuffer);
+}
+
+let nncompute = (inputBuffer) => {
+    let outputBuffer = compute(modelRunner, inputBuffer);
+    outputBuffer = new Float32Array(outputBuffer);
+    postMessage(outputBuffer, [outputBuffer.buffer]);
+}
+
+addEventListener('message', async (e) => {
+    const task = e.data.task;
+    const value = e.data.value;
+    if(task === "nnwarmup") {
+       await warmup();
+       postMessage({msg: "modelloaded"});
+    }
+    if(task === "nncompute") {
+       nncompute(value);
+    }
+});
+
+async function load() {
+    // Create the model runner with the model.
+
+    const MODEL_PATH = '../../../../assets/models/deeplab/deeplab_mobilenetv2_321_no_argmax.tflite';
+
+    // Load WASM module and model.
+    const [module, modelArrayBuffer] = await Promise.all([
+        tflite_model_runner_ModuleFactory(),
+        (await fetch(MODEL_PATH)).arrayBuffer(),
+    ]);
+    const modelBytes = new Uint8Array(modelArrayBuffer);
+    const offset = module._malloc(modelBytes.length);
+    module.HEAPU8.set(modelBytes, offset);
+    // Create model runner.
+    const modelRunnerResult =
+        module.TFLiteWebModelRunner.CreateFromBufferAndOptions(
+            offset, modelBytes.length, {
+            numThreads: Math.min(
+                4, Math.max(1, (navigator.hardwareConcurrency || 1) / 2)),
+            enableWebNNDelegate: true,
+            webNNDevicePreference: 1 // 0 - default, 1 - gpu, 2 - cpu
+        });
+    if (!modelRunnerResult.ok()) {
+        throw new Error(
+            'Failed to create TFLiteWebModelRunner: ' + modelRunnerResult.errorMessage());
+    }
+    const modelRunner = modelRunnerResult.value();
+    return modelRunner;
+}
+
+function compute(modelRunner, inputData) {
+    // Get input and output info.
+    const inputs = this.callAndDelete(
+        modelRunner.GetInputs(), results => this.convertCppVectorToArray(results));
+    const input = inputs[0];
+    const outputs = this.callAndDelete(
+        modelRunner.GetOutputs(), results => this.convertCppVectorToArray(results));
+    const output = outputs[0];
+
+    // Set input tensor data from the image
+    const inputBuffer = input.data();
+    // const inputData = new Float32Array(inputBuffer.length);
+    inputBuffer.set(inputData);
+
+    const success = modelRunner.Infer();
+    if (!success) return;
+    return output.data();
 }
 
-let computetask = async () => {
-    inputBuffer = getInputTensor(null, inputOptions);
-    outputBuffer = netInstance.compute(modelRunner, inputBuffer);
-    self.postMessage({msg: "prediction completed"});
+/**
+ * Calls the given function with the given deletable argument, ensuring that
+ * the argument gets deleted afterwards (even if the function throws an error).
+*/
+function callAndDelete(arg, func) {
+    try {
+        return func(arg);
+    } finally {
+        if (arg != null) arg.delete();
+    }
 }
 
-loadtask();
-computetask();
+/** Converts the given c++ vector to a JS array. */
+function convertCppVectorToArray(vector) {
+    if (vector == null) return [];
 
+    const result = [];
+    for (let i = 0; i < vector.size(); i++) {
+        const item = vector.get(i);
+        result.push(item);
+    }
+    return result;
+}