mirror of
https://github.com/wukko/cobalt.git
synced 2025-01-22 10:46:19 +01:00
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:
parent
2812960088
commit
8e9347b4a0
2 changed files with 46 additions and 39 deletions
|
@ -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();
|
||||||
}
|
}
|
||||||
|
|
|
@ -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}
|
||||||
|
|
||||||
|
|
Loading…
Reference in a new issue