From 0a14b78d67d10717919d97064488107148656d28 Mon Sep 17 00:00:00 2001
From: Juliusz Chroboczek <jch@irif.fr>
Date: Sat, 14 Dec 2024 14:48:50 +0100
Subject: [PATCH] Move MediaPipe initialisation to filter init.

We used to initialise MediaPipe at worker start, which prevented
us from handling errors.  We now do it at filter init, and stop
the stream with an error message if initialisation fails.
---
 static/background-blur-worker.js | 33 +++++++++++++------------
 static/galene.js                 | 42 +++++++++++++++++++++-----------
 2 files changed, 46 insertions(+), 29 deletions(-)

diff --git a/static/background-blur-worker.js b/static/background-blur-worker.js
index 5c7ff75..2076b60 100644
--- a/static/background-blur-worker.js
+++ b/static/background-blur-worker.js
@@ -2,17 +2,15 @@
 
 let imageSegmenter;
 
-async function loadImageSegmenter() {
+async function loadImageSegmenter(model) {
     let module = await import('/third-party/tasks-vision/vision_bundle.mjs');
-
     let vision = await module.FilesetResolver.forVisionTasks(
         "/third-party/tasks-vision/wasm"
     );
 
-    imageSegmenter =
-        await module.ImageSegmenter.createFromOptions(vision, {
+    return await module.ImageSegmenter.createFromOptions(vision, {
             baseOptions: {
-                modelAssetPath: '/third-party/tasks-vision/models/selfie_segmenter.tflite',
+                modelAssetPath: model,
             },
             outputCategoryMask: true,
             outputConfidenceMasks: false,
@@ -20,19 +18,24 @@ async function loadImageSegmenter() {
         });
 }
 
-loadImageSegmenter();
-
-onmessage = e => {
-    let bitmap = e.data.bitmap;
-    if(!(bitmap instanceof ImageBitmap)) {
-        postMessage(new Error('Bad type for worker data'));
+onmessage = async e => {
+    let data = e.data;
+    if(imageSegmenter == null) {
+        try {
+            imageSegmenter = await loadImageSegmenter(data.model);
+            if(imageSegmenter == null)
+                throw new Error("loadImageSegmenter returned null");
+        } catch(e) {
+            postMessage(e);
+            return;
+        }
+        postMessage(null);
         return;
     }
 
-    if(!imageSegmenter) {
-        // not ready yet
-        bitmap.close();
-        postMessage(null);
+    let bitmap = e.data.bitmap;
+    if(!(bitmap instanceof ImageBitmap)) {
+        postMessage(new Error('Bad type for worker data'));
         return;
     }
 
diff --git a/static/galene.js b/static/galene.js
index 36a56a1..5a1e00e 100644
--- a/static/galene.js
+++ b/static/galene.js
@@ -1188,6 +1188,30 @@ async function setFilter(c) {
     c.userdata.filter = filter;
 }
 
+/**
+ * Sends a message to a worker, then waits for a reply.
+ *
+ * @param {Worker} worker
+ * @param {any} message
+ * @param {any[]} [transfer]
+ */
+async function workerSendReceive(worker, message, transfer) {
+    let p = new Promise((resolve, reject) => {
+        worker.onmessage = e => {
+            if(e && e.data) {
+                if(e.data instanceof Error)
+                    reject(e.data);
+                else
+                    resolve(e.data);
+            } else {
+                resolve(null);
+            }
+        };
+    });
+    worker.postMessage(message, transfer);
+    return await p
+}
+
 /**
  * @type {Object.<string,filterDefinition>}
  */
@@ -1242,6 +1266,9 @@ let filters = {
             if(this.userdata.worker)
                 throw new Error("Worker already running (this shouldn't happen)")
             this.userdata.worker = new Worker('/background-blur-worker.js');
+            await workerSendReceive(this.userdata.worker, {
+                model: '/third-party/tasks-vision/models/selfie_segmenter.tflite',
+            });
         },
         cleanup: async function() {
             if(this.userdata.worker.onmessage) {
@@ -1253,23 +1280,10 @@ let filters = {
         draw: async function(src, ctx) {
             let bitmap = await createImageBitmap(src);
             try {
-                let p = new Promise((resolve, reject) => {
-                    this.userdata.worker.onmessage = e => {
-                        if(e && e.data) {
-                            if(e.data instanceof Error)
-                                reject(e.data);
-                            else
-                                resolve(e.data);
-                        } else {
-                            resolve(null);
-                        }
-                    };
-                });
-                this.userdata.worker.postMessage({
+                let result = await workerSendReceive(this.userdata.worker, {
                     bitmap: bitmap,
                     timestamp: performance.now(),
                 }, [bitmap]);
-                let result = await p;
 
                 if(!result)
                     return false;