web: very early proof-of-concept of on-device image background removal

web/src/lib/workers/removebg.ts

@@ -0,0 +1,56 @@
+import { AutoModel, AutoProcessor, RawImage } from "@huggingface/transformers";
+
+const models = {
+    light: {
+        id: "briaai/RMBG-1.4",
+        input: "input",
+        modelConfig: {
+            device: "wasm",
+            dtype: "fp32",
+        },
+        processorConfig: {},
+    },
+    heavy: {
+        id: "onnx-community/BiRefNet_lite",
+        input: "input_image",
+        modelConfig: {
+            device: "webgpu",
+            dtype: "fp16",
+        },
+        processorConfig: {},
+    }
+}
+
+export const removeImageBackground = async (file: File) => {
+    const image = await RawImage.fromBlob(new Blob([file]));
+
+    const model_type = "light";
+    const model = await AutoModel.from_pretrained(models[model_type].id, models[model_type].modelConfig);
+
+    const processor = await AutoProcessor.from_pretrained(models[model_type].id, models[model_type].processorConfig);
+
+    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;
+    }
+}