From 28d8927c085171449e4942151f0589e83f1f0b63 Mon Sep 17 00:00:00 2001
From: wukko <me@wukko.me>
Date: Wed, 15 Jan 2025 17:22:34 +0600
Subject: [PATCH] web/removebg: convert to a proper web worker

no more hanging ui :3
---
 web/src/lib/workers/removebg.ts    | 48 ++++++++++++++----------
 web/src/routes/cutout/+page.svelte | 59 ++++++++++++++++++------------
 web/src/routes/cutout/+page.ts     |  1 +
 3 files changed, 66 insertions(+), 42 deletions(-)
 create mode 100644 web/src/routes/cutout/+page.ts

diff --git a/web/src/lib/workers/removebg.ts b/web/src/lib/workers/removebg.ts
index e6a59bcc..4f8e7bd8 100644
--- a/web/src/lib/workers/removebg.ts
+++ b/web/src/lib/workers/removebg.ts
@@ -11,8 +11,30 @@ 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;
+    const ctx = canvas.getContext('2d');
+
+    if (!ctx) return;
+
+    ctx.drawImage(image.toCanvas(), 0, 0);
+
+    const pixelData = ctx.getImageData(0, 0, image.width, image.height);
+    for (let i = 0; i < mask.data.length; ++i) {
+        pixelData.data[4 * i + 3] = mask.data[i];
+    }
+    ctx.putImageData(pixelData, 0, 0);
+
+    return canvas;
+}
+
 export const removeImageBackground = async (file: File) => {
-    const image = await RawImage.fromBlob(new Blob([file]));
+    const originalImageBlob = new Blob([file]);
+    const image = await RawImage.fromBlob(originalImageBlob);
 
     const model_type = "light";
     const model = await AutoModel.from_pretrained(models[model_type].id, {
@@ -24,26 +46,14 @@ export const removeImageBackground = async (file: File) => {
 
     if (model && processor) {
         const { pixel_values } = await processor(image);
-
         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);
 
-        const canvas = document.createElement('canvas');
-        canvas.width = image.width;
-        canvas.height = image.height;
-        const ctx = canvas.getContext('2d');
-
-        if (!ctx) return;
-
-        ctx.drawImage(image.toCanvas(), 0, 0);
-
-        const pixelData = ctx.getImageData(0, 0, image.width, image.height);
-        for (let i = 0; i < mask.data.length; ++i) {
-            pixelData.data[4 * i + 3] = mask.data[i];
-        }
-        ctx.putImageData(pixelData, 0, 0);
-
-        return canvas;
+        self.postMessage({ source: originalImageBlob, mask });
     }
 }
+
+self.onmessage = async (event: MessageEvent) => {
+    await removeImageBackground(event.data.file);
+    self.close();
+}
diff --git a/web/src/routes/cutout/+page.svelte b/web/src/routes/cutout/+page.svelte
index dded67bb..eed2ef19 100644
--- a/web/src/routes/cutout/+page.svelte
+++ b/web/src/routes/cutout/+page.svelte
@@ -1,36 +1,49 @@
 <script lang="ts">
     import settings from "$lib/state/settings";
+    import RemoveBgWorker from '$lib/workers/removebg?worker';
 
     import { onMount } from "svelte";
     import { goto } from "$app/navigation";
 
     import { downloadFile } from "$lib/download";
-    import { removeImageBackground } from "$lib/workers/removebg";
+    import { maskImage } from "$lib/workers/removebg";
 
     import DropReceiver from "$components/misc/DropReceiver.svelte";
     import FileReceiver from "$components/misc/FileReceiver.svelte";
     import Skeleton from "$components/misc/Skeleton.svelte";
 
-    let imageContainer: HTMLElement;
-
     let draggedOver = false;
     let file: File | undefined;
-
-    let thinking = false;
-    let done = false;
-
     let result: HTMLCanvasElement;
+    let imageContainer: HTMLElement;
+
+    let state: "empty" | "busy" | "done" = "empty";
+
+    const worker = new RemoveBgWorker();
 
     const processImage = async () => {
-        if (file) {
-            thinking = true;
-            const removedBackground = await removeImageBackground(file);
-            if (removedBackground) {
-                thinking = false;
-                done = true;
-                result = removedBackground;
-                imageContainer.append(removedBackground);
-            }
+        if (!file) return;
+
+        state = "busy";
+
+        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.onerror = (e) => {
+            console.error("bg removal worker exploded:", e)
+            worker.terminate();
         }
     };
 
@@ -53,7 +66,7 @@
 </script>
 
 <DropReceiver bind:file bind:draggedOver id="cutout-container">
-    {#if !thinking && !done}
+    {#if state === "empty"}
         <FileReceiver
             bind:draggedOver
             bind:file
@@ -68,27 +81,27 @@
         {/if}
     {/if}
 
-    {#if thinking}
+    {#if state === "busy"}
         <div>thinking very hard rn...</div>
     {/if}
 
-    {#if done}
+    {#if state === "done"}
         <div>thought a lot, here's what i got:</div>
     {/if}
 
-    {#if thinking || done}
+    {#if ["busy", "done"].includes(state)}
         <div id="image-preview" bind:this={imageContainer}>
-            {#if !done}
+            {#if state === "busy"}
                 <Skeleton width="100%" height="100%" class="big" />
             {/if}
         </div>
     {/if}
 
-    {#if done}
+    {#if state === "done"}
         <div id="finished-actions">
             <button
                 on:click={() => {
-                    done = false;
+                    state = "empty"
                     file = undefined;
                 }}
             >
diff --git a/web/src/routes/cutout/+page.ts b/web/src/routes/cutout/+page.ts
new file mode 100644
index 00000000..a3d15781
--- /dev/null
+++ b/web/src/routes/cutout/+page.ts
@@ -0,0 +1 @@
+export const ssr = false;