mirror of
https://github.com/imputnet/cobalt.git
synced 2025-01-15 19:35:31 +00:00
web/removebg: convert to a proper web worker
Some checks are pending
Some checks are pending
no more hanging ui :3
This commit is contained in:
parent
2f2d39dc4c
commit
28d8927c08
|
@ -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();
|
||||
}
|
||||
|
|
|
@ -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;
|
||||
}}
|
||||
>
|
||||
|
|
1
web/src/routes/cutout/+page.ts
Normal file
1
web/src/routes/cutout/+page.ts
Normal file
|
@ -0,0 +1 @@
|
|||
export const ssr = false;
|
Loading…
Reference in a new issue