Add remove-background demo (#576)

* Add remove-background demo * Open links in new tab * Add mention of running locally
huggingface · Feb 7, 2024 · 66d59d7 · 66d59d7
1 parent 41f98b7
commit 66d59d7
Show file tree

Hide file tree

Showing 7 changed files with 1,741 additions and 0 deletions.
diff --git a/examples/remove-background-client/.gitignore b/examples/remove-background-client/.gitignore
@@ -0,0 +1,24 @@
+# Logs
+logs
+*.log
+npm-debug.log*
+yarn-debug.log*
+yarn-error.log*
+pnpm-debug.log*
+lerna-debug.log*
+
+node_modules
+dist
+dist-ssr
+*.local
+
+# Editor directories and files
+.vscode/*
+!.vscode/extensions.json
+.idea
+.DS_Store
+*.suo
+*.ntvs*
+*.njsproj
+*.sln
+*.sw?
diff --git a/examples/remove-background-client/index.html b/examples/remove-background-client/index.html
@@ -0,0 +1,33 @@
+<!DOCTYPE html>
+<html lang="en">
+
+<head>
+  <meta charset="UTF-8" />
+  <meta name="viewport" content="width=device-width, initial-scale=1.0" />
+  <title>Transformers.js - Background Removal</title>
+</head>
+
+<body>
+  <h1>Background Removal w/ <a href="http://github.com/xenova/transformers.js" target="_blank">🤗 Transformers.js</a>
+  </h1>
+  <h4>Runs locally in your browser, powered by the <a href="https://huggingface.co/briaai/RMBG-1.4" target="_blank">RMBG V1.4 model</a> from <a
+      href="https://bria.ai/" target="_blank">BRIA AI</a>
+  </h4>
+  <div id="container">
+    <label id="upload-button" for="upload">
+      <svg width="25" height="25" viewBox="0 0 25 25" fill="none" xmlns="http://www.w3.org/2000/svg">
+        <path fill="#000"
+          d="M3.5 24.3a3 3 0 0 1-1.9-.8c-.5-.5-.8-1.2-.8-1.9V2.9c0-.7.3-1.3.8-1.9.6-.5 1.2-.7 2-.7h18.6c.7 0 1.3.2 1.9.7.5.6.7 1.2.7 2v18.6c0 .7-.2 1.4-.7 1.9a3 3 0 0 1-2 .8H3.6Zm0-2.7h18.7V2.9H3.5v18.7Zm2.7-2.7h13.3c.3 0 .5 0 .6-.3v-.7l-3.7-5a.6.6 0 0 0-.6-.2c-.2 0-.4 0-.5.3l-3.5 4.6-2.4-3.3a.6.6 0 0 0-.6-.3c-.2 0-.4.1-.5.3l-2.7 3.6c-.1.2-.2.4 0 .7.1.2.3.3.6.3Z">
+        </path>
+      </svg>
+      Click to upload image
+      <label id="example">(or try example)</label>
+    </label>
+  </div>
+  <label id="status"></label>
+  <input id="upload" type="file" accept="image/*" />
+
+  <script type="module" src="/main.js"></script>
+</body>
+
+</html>
diff --git a/examples/remove-background-client/main.js b/examples/remove-background-client/main.js
@@ -0,0 +1,112 @@
+import './style.css';
+
+import { AutoModel, AutoProcessor, env, RawImage } from '@xenova/transformers';
+
+// Since we will download the model from the Hugging Face Hub, we can skip the local model check
+env.allowLocalModels = false;
+
+// Proxy the WASM backend to prevent the UI from freezing
+env.backends.onnx.wasm.proxy = true;
+
+// Constants
+const EXAMPLE_URL = 'https://images.pexels.com/photos/5965592/pexels-photo-5965592.jpeg?auto=compress&cs=tinysrgb&w=1024';
+
+// Reference the elements that we will need
+const status = document.getElementById('status');
+const fileUpload = document.getElementById('upload');
+const imageContainer = document.getElementById('container');
+const example = document.getElementById('example');
+
+// Load model and processor
+status.textContent = 'Loading model...';
+
+const model = await AutoModel.from_pretrained('briaai/RMBG-1.4', {
+    // Do not require config.json to be present in the repository
+    config: { model_type: 'custom' },
+});
+
+const processor = await AutoProcessor.from_pretrained('briaai/RMBG-1.4', {
+    // Do not require config.json to be present in the repository
+    config: {
+        do_normalize: true,
+        do_pad: false,
+        do_rescale: true,
+        do_resize: true,
+        image_mean: [0.5, 0.5, 0.5],
+        feature_extractor_type: "ImageFeatureExtractor",
+        image_std: [1, 1, 1],
+        resample: 2,
+        rescale_factor: 0.00392156862745098,
+        size: { width: 1024, height: 1024 },
+    }
+});
+
+status.textContent = 'Ready';
+
+example.addEventListener('click', (e) => {
+    e.preventDefault();
+    predict(EXAMPLE_URL);
+});
+
+fileUpload.addEventListener('change', function (e) {
+    const file = e.target.files[0];
+    if (!file) {
+        return;
+    }
+
+    const reader = new FileReader();
+
+    // Set up a callback when the file is loaded
+    reader.onload = e2 => predict(e2.target.result);
+
+    reader.readAsDataURL(file);
+});
+
+// Predict foreground of the given image
+async function predict(url) {
+    // Read image
+    const image = await RawImage.fromURL(url);
+
+    // Update UI
+    imageContainer.innerHTML = '';
+    imageContainer.style.backgroundImage = `url(${url})`;
+
+    // Set container width and height depending on the image aspect ratio
+    const ar = image.width / image.height;
+    const [cw, ch] = (ar > 720 / 480) ? [720, 720 / ar] : [480 * ar, 480];
+    imageContainer.style.width = `${cw}px`;
+    imageContainer.style.height = `${ch}px`;
+
+    status.textContent = 'Analysing...';
+
+    // Preprocess image
+    const { pixel_values } = await processor(image);
+
+    // Predict alpha matte
+    const { output } = await model({ input: pixel_values });
+
+    // Resize mask back to original size
+    const mask = await RawImage.fromTensor(output[0].mul(255).to('uint8')).resize(image.width, image.height);
+
+    // Create new canvas
+    const canvas = document.createElement('canvas');
+    canvas.width = image.width;
+    canvas.height = image.height;
+    const ctx = canvas.getContext('2d');
+
+    // Draw original image output to canvas
+    ctx.drawImage(image.toCanvas(), 0, 0);
+
+    // Update alpha channel
+    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);
+
+    // Update UI
+    imageContainer.append(canvas);
+    imageContainer.style.removeProperty('background-image');
+    imageContainer.style.background = `url("data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAABAAAAAQBAMAAADt3eJSAAAAAXNSR0IArs4c6QAAAARnQU1BAACxjwv8YQUAAAAGUExURb+/v////5nD/3QAAAAJcEhZcwAADsMAAA7DAcdvqGQAAAAUSURBVBjTYwABQSCglEENMxgYGAAynwRB8BEAgQAAAABJRU5ErkJggg==")`;
+    status.textContent = 'Done!';
+}