From 8e9347b4a0b7cb9235d6503d94457b35968466dd Mon Sep 17 00:00:00 2001
From: wukko <me@wukko.me>
Date: Fri, 17 Jan 2025 01:03:59 +0600
Subject: [PATCH] web/removebg: fix functionality after build, improve pipeline

- no longer killing the worker if it has done its job correctly and is expected to shut itself down
- no longer reading messages not intended for the worker handler and also made the cobalt messaging distnict
---
 web/src/lib/workers/removebg.ts    | 23 +++++------
 web/src/routes/cutout/+page.svelte | 62 ++++++++++++++++--------------
 2 files changed, 46 insertions(+), 39 deletions(-)

diff --git a/web/src/lib/workers/removebg.ts b/web/src/lib/workers/removebg.ts
index 811e5f13..29f62ff6 100644
--- a/web/src/lib/workers/removebg.ts
+++ b/web/src/lib/workers/removebg.ts
@@ -11,12 +11,8 @@ const models = {
     }
 }
 
-export const maskImage = async (source: Blob, mask: RawImage) => {
-    const image = await RawImage.fromBlob(source);
-
-    const canvas = document.createElement('canvas');
-    canvas.width = image.width;
-    canvas.height = image.height;
+export const maskImage = (image: RawImage, mask: RawImage) => {
+    const canvas = new OffscreenCanvas(image.width, image.height);
     const ctx = canvas.getContext('2d');
 
     if (!ctx) return;
@@ -29,12 +25,11 @@ export const maskImage = async (source: Blob, mask: RawImage) => {
     }
     ctx.putImageData(pixelData, 0, 0);
 
-    return canvas;
+    return canvas.transferToImageBitmap();
 }
 
 const removeImageBackground = async (file: File) => {
-    const originalImageBlob = new Blob([file]);
-    const image = await RawImage.fromBlob(originalImageBlob);
+    const image = await RawImage.fromBlob(file);
 
     const model_type = "light";
     const model = await AutoModel.from_pretrained(models[model_type].id, {
@@ -52,11 +47,17 @@ const removeImageBackground = async (file: File) => {
         const { output } = await model({ [models[model_type].input]: pixel_values });
         const mask = await RawImage.fromTensor(output[0].mul(255).to('uint8')).resize(image.width, image.height);
 
-        self.postMessage({ source: originalImageBlob, mask });
+        self.postMessage({
+            cobaltRemoveBgWorker: {
+                result: maskImage(image, mask),
+            }
+         });
     }
 }
 
 self.onmessage = async (event: MessageEvent) => {
-    await removeImageBackground(event.data.file);
+    if (event.data.file) {
+        await removeImageBackground(event.data.file);
+    }
     self.close();
 }
diff --git a/web/src/routes/cutout/+page.svelte b/web/src/routes/cutout/+page.svelte
index fb51f235..3d923819 100644
--- a/web/src/routes/cutout/+page.svelte
+++ b/web/src/routes/cutout/+page.svelte
@@ -4,64 +4,69 @@
 
     import { onMount } from "svelte";
     import { goto } from "$app/navigation";
-
     import { downloadFile } from "$lib/download";
-    import { maskImage } from "$lib/workers/removebg";
 
+    import Skeleton from "$components/misc/Skeleton.svelte";
     import DropReceiver from "$components/misc/DropReceiver.svelte";
     import FileReceiver from "$components/misc/FileReceiver.svelte";
-    import Skeleton from "$components/misc/Skeleton.svelte";
 
     let draggedOver = false;
     let file: File | undefined;
-    let result: HTMLCanvasElement;
+    let result: ImageBitmap;
     let imageContainer: HTMLElement;
+    let canvas: HTMLCanvasElement;
 
     let state: "empty" | "busy" | "done" = "empty";
 
     let worker: Worker;
 
+    const renderImageToCanvas = () => {
+        if (canvas && result) {
+            canvas.width = result.width;
+            canvas.height = result.height;
+            canvas.getContext('bitmaprenderer')?.transferFromImageBitmap(result);
+        }
+    }
+
     const processImage = async () => {
         if (!file) return;
 
         state = "busy";
-
         worker = new RemoveBgWorker();
 
         worker.postMessage({ file });
         worker.onmessage = async (event) => {
-            const maskedCanvas = await maskImage(event.data.source, event.data.mask);
-
-            if (!maskedCanvas) {
-                state = "empty";
-                return;
-            };
-
-            state = "done";
-
-            result = maskedCanvas;
-            imageContainer.append(maskedCanvas);
-
-            worker.terminate();
+            console.log("event received by removebg page:", event)
+            const eventData = event.data.cobaltRemoveBgWorker;
+            if (eventData.result) {
+                state = "done";
+                result = eventData.result;
+                renderImageToCanvas();
+            }
         };
 
         worker.onerror = (e) => {
-            console.error("bg removal worker exploded:", e);
-
             state = "empty";
+            console.error("bg removal worker exploded:", e);
             worker.terminate();
         }
     };
 
     const exportImage = async () => {
-        result.toBlob(async (blob) => {
-            if (!blob || !file) return;
-            return await downloadFile({
-                file: new File([blob], `${file.name} (cutout).png`, {
-                    type: "image/png",
-                }),
-            });
-        }, "image/png");
+        if (!result || !file) return;
+
+        const resultBlob = await new Promise<Blob>((resolve, reject) => {
+            canvas.toBlob(blob => {
+                if (blob) resolve(blob);
+                else reject();
+            }, "image/png")
+        })
+
+        return await downloadFile({
+            file: new File([resultBlob], `${file.name} (cutout).png`, {
+                type: "image/png",
+            }),
+        });
     };
 
     onMount(() => {
@@ -109,6 +114,7 @@
             {#if state === "busy"}
                 <Skeleton width="100%" height="100%" class="big" />
             {/if}
+            <canvas bind:this={canvas}></canvas>
         </div>
     {/if}