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
This commit is contained in:
wukko 2025-01-17 01:03:59 +06:00
parent 2812960088
commit 8e9347b4a0
No known key found for this signature in database
GPG key ID: 3E30B3F26C7B4AA2
2 changed files with 46 additions and 39 deletions

View file

@ -11,12 +11,8 @@ const models = {
} }
} }
export const maskImage = async (source: Blob, mask: RawImage) => { export const maskImage = (image: RawImage, mask: RawImage) => {
const image = await RawImage.fromBlob(source); const canvas = new OffscreenCanvas(image.width, image.height);
const canvas = document.createElement('canvas');
canvas.width = image.width;
canvas.height = image.height;
const ctx = canvas.getContext('2d'); const ctx = canvas.getContext('2d');
if (!ctx) return; if (!ctx) return;
@ -29,12 +25,11 @@ export const maskImage = async (source: Blob, mask: RawImage) => {
} }
ctx.putImageData(pixelData, 0, 0); ctx.putImageData(pixelData, 0, 0);
return canvas; return canvas.transferToImageBitmap();
} }
const removeImageBackground = async (file: File) => { const removeImageBackground = async (file: File) => {
const originalImageBlob = new Blob([file]); const image = await RawImage.fromBlob(file);
const image = await RawImage.fromBlob(originalImageBlob);
const model_type = "light"; const model_type = "light";
const model = await AutoModel.from_pretrained(models[model_type].id, { 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 { 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 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) => { self.onmessage = async (event: MessageEvent) => {
if (event.data.file) {
await removeImageBackground(event.data.file); await removeImageBackground(event.data.file);
}
self.close(); self.close();
} }

View file

@ -4,64 +4,69 @@
import { onMount } from "svelte"; import { onMount } from "svelte";
import { goto } from "$app/navigation"; import { goto } from "$app/navigation";
import { downloadFile } from "$lib/download"; 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 DropReceiver from "$components/misc/DropReceiver.svelte";
import FileReceiver from "$components/misc/FileReceiver.svelte"; import FileReceiver from "$components/misc/FileReceiver.svelte";
import Skeleton from "$components/misc/Skeleton.svelte";
let draggedOver = false; let draggedOver = false;
let file: File | undefined; let file: File | undefined;
let result: HTMLCanvasElement; let result: ImageBitmap;
let imageContainer: HTMLElement; let imageContainer: HTMLElement;
let canvas: HTMLCanvasElement;
let state: "empty" | "busy" | "done" = "empty"; let state: "empty" | "busy" | "done" = "empty";
let worker: Worker; let worker: Worker;
const renderImageToCanvas = () => {
if (canvas && result) {
canvas.width = result.width;
canvas.height = result.height;
canvas.getContext('bitmaprenderer')?.transferFromImageBitmap(result);
}
}
const processImage = async () => { const processImage = async () => {
if (!file) return; if (!file) return;
state = "busy"; state = "busy";
worker = new RemoveBgWorker(); worker = new RemoveBgWorker();
worker.postMessage({ file }); worker.postMessage({ file });
worker.onmessage = async (event) => { worker.onmessage = async (event) => {
const maskedCanvas = await maskImage(event.data.source, event.data.mask); console.log("event received by removebg page:", event)
const eventData = event.data.cobaltRemoveBgWorker;
if (!maskedCanvas) { if (eventData.result) {
state = "empty";
return;
};
state = "done"; state = "done";
result = eventData.result;
result = maskedCanvas; renderImageToCanvas();
imageContainer.append(maskedCanvas); }
worker.terminate();
}; };
worker.onerror = (e) => { worker.onerror = (e) => {
console.error("bg removal worker exploded:", e);
state = "empty"; state = "empty";
console.error("bg removal worker exploded:", e);
worker.terminate(); worker.terminate();
} }
}; };
const exportImage = async () => { const exportImage = async () => {
result.toBlob(async (blob) => { if (!result || !file) return;
if (!blob || !file) return;
const resultBlob = await new Promise<Blob>((resolve, reject) => {
canvas.toBlob(blob => {
if (blob) resolve(blob);
else reject();
}, "image/png")
})
return await downloadFile({ return await downloadFile({
file: new File([blob], `${file.name} (cutout).png`, { file: new File([resultBlob], `${file.name} (cutout).png`, {
type: "image/png", type: "image/png",
}), }),
}); });
}, "image/png");
}; };
onMount(() => { onMount(() => {
@ -109,6 +114,7 @@
{#if state === "busy"} {#if state === "busy"}
<Skeleton width="100%" height="100%" class="big" /> <Skeleton width="100%" height="100%" class="big" />
{/if} {/if}
<canvas bind:this={canvas}></canvas>
</div> </div>
{/if} {/if}