web/removebg: convert to a proper web worker
no more hanging ui :3
This commit is contained in:
@@ -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();
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user