From 93dd37017c32787a94aff17526c5ce1ee7b99f36 Mon Sep 17 00:00:00 2001
From: Siddharth Kumar Sah <siddharth123sk@gmail.com>
Date: Sun, 12 Apr 2026 18:09:59 +0800
Subject: [PATCH] feat: add Ultra quality mode with BiRefNet-matting for people
 photos

Adds a new "Ultra" quality tier for People subject type that uses
BiRefNet-matting (ONNX, 928MB) for true alpha matting instead of
binary segmentation. Produces per-pixel transparency for hair wisps
and fine edges that standard models miss.

- Custom rembg session class loads BiRefNet-matting ONNX from GitHub releases
- Zero new Python dependencies (reuses existing onnxruntime)
- Model pre-downloaded in Docker build alongside existing models
- Ultra option only visible when subject is People
- Falls back to Best when switching to Products/General
---
 .../components/tools/remove-bg-settings.tsx   | 33 +++++++++++++----
 docker/download_models.py                     | 30 +++++++++++++++
 packages/ai/python/remove_bg.py               | 37 +++++++++++++++++++
 tests/e2e/remove-bg.spec.ts                   | 27 ++++++++++++++
 4 files changed, 119 insertions(+), 8 deletions(-)

diff --git a/apps/web/src/components/tools/remove-bg-settings.tsx b/apps/web/src/components/tools/remove-bg-settings.tsx
index a6a7a6e6..012654ac 100644
--- a/apps/web/src/components/tools/remove-bg-settings.tsx
+++ b/apps/web/src/components/tools/remove-bg-settings.tsx
@@ -13,18 +13,24 @@ import { useToolProcessor } from "@/hooks/use-tool-processor";
 import { useFileStore } from "@/stores/file-store";
 
 type SubjectType = "people" | "products" | "general";
-type Quality = "fast" | "balanced" | "best";
+type Quality = "fast" | "balanced" | "best" | "ultra";
 type BackgroundType = "transparent" | "color" | "gradient" | "image";
 
 type BgModel =
   | "birefnet-general"
   | "birefnet-general-lite"
+  | "birefnet-matting"
   | "birefnet-portrait"
   | "bria-rmbg"
   | "u2net";
 
-const MODEL_MAP: Record<SubjectType, Record<Quality, BgModel>> = {
-  people: { fast: "u2net", balanced: "birefnet-portrait", best: "birefnet-portrait" },
+const MODEL_MAP: Record<SubjectType, Partial<Record<Quality, BgModel>>> = {
+  people: {
+    fast: "u2net",
+    balanced: "birefnet-portrait",
+    best: "birefnet-portrait",
+    ultra: "birefnet-matting",
+  },
   products: { fast: "u2net", balanced: "bria-rmbg", best: "birefnet-general" },
   general: { fast: "u2net", balanced: "birefnet-general-lite", best: "birefnet-general" },
 };
@@ -35,10 +41,11 @@ const SUBJECT_OPTIONS: { value: SubjectType; label: string; icon: typeof User }[
   { value: "general", label: "General", icon: ImageIcon },
 ];
 
-const QUALITY_OPTIONS: { value: Quality; label: string }[] = [
+const ALL_QUALITY_OPTIONS: { value: Quality; label: string; peopleOnly?: boolean }[] = [
   { value: "fast", label: "Fast" },
   { value: "balanced", label: "Balanced" },
   { value: "best", label: "Best" },
+  { value: "ultra", label: "Ultra", peopleOnly: true },
 ];
 
 const COLOR_PRESETS = [
@@ -97,8 +104,18 @@ export function RemoveBgControls({ settings, onChange }: RemoveBgControlsProps)
   // Expandable sections
   const [effectsOpen, setEffectsOpen] = useState(false);
 
+  // Filter quality options based on subject (Ultra only for People)
+  const qualityOptions = ALL_QUALITY_OPTIONS.filter(
+    (opt) => !opt.peopleOnly || subject === "people",
+  );
+
+  // If switching away from People while on Ultra, fall back to Best
+  const effectiveQuality = quality === "ultra" && subject !== "people" ? "best" : quality;
+
   const model =
-    isPassport && subject === "people" ? "birefnet-portrait" : MODEL_MAP[subject][quality];
+    isPassport && subject === "people"
+      ? "birefnet-portrait"
+      : MODEL_MAP[subject][effectiveQuality] || "birefnet-general";
 
   const onChangeRef = useRef(onChange);
   useEffect(() => {
@@ -190,14 +207,14 @@ export function RemoveBgControls({ settings, onChange }: RemoveBgControlsProps)
 
       {/* Quality */}
       <SectionLabel>Quality</SectionLabel>
-      <div className="grid grid-cols-3 gap-1.5">
-        {QUALITY_OPTIONS.map((opt) => (
+      <div className={`grid gap-1.5 ${qualityOptions.length > 3 ? "grid-cols-4" : "grid-cols-3"}`}>
+        {qualityOptions.map((opt) => (
           <button
             key={opt.value}
             type="button"
             onClick={() => setQuality(opt.value)}
             className={`py-2 px-2 rounded-lg border text-xs font-medium transition-colors ${
-              quality === opt.value
+              effectiveQuality === opt.value
                 ? "border-primary bg-primary/10 text-primary"
                 : "border-border text-muted-foreground hover:border-primary/50"
             }`}
diff --git a/docker/download_models.py b/docker/download_models.py
index 19324d88..b3da74e9 100644
--- a/docker/download_models.py
+++ b/docker/download_models.py
@@ -27,6 +27,7 @@ REMBG_MODELS = [
     "birefnet-general-lite",
     "birefnet-portrait",
     "birefnet-general",
+    "birefnet-matting",
 ]
 
 # PaddleOCR language codes (not ISO). German/French/Spanish use "latin" model.
@@ -34,11 +35,40 @@ REMBG_MODELS = [
 PADDLEOCR_LANGUAGES = ["en", "ch", "japan", "korean", "latin"]
 
 
+def _register_birefnet_matting():
+    """Register BiRefNet-matting ONNX session for Ultra quality mode."""
+    import os
+    import pooch
+    from rembg.sessions import sessions_class
+    from rembg.sessions.birefnet_general import BiRefNetSessionGeneral
+
+    class BiRefNetMattingSession(BiRefNetSessionGeneral):
+        @classmethod
+        def download_models(cls, *args, **kwargs):
+            fname = f"{cls.name(*args, **kwargs)}.onnx"
+            pooch.retrieve(
+                "https://github.com/ZhengPeng7/BiRefNet/releases/download/v1/BiRefNet-matting-epoch_100.onnx",
+                None,  # Skip checksum for GitHub release assets
+                fname=fname,
+                path=cls.u2net_home(*args, **kwargs),
+                progressbar=True,
+            )
+            return os.path.join(cls.u2net_home(*args, **kwargs), fname)
+
+        @classmethod
+        def name(cls, *args, **kwargs):
+            return "birefnet-matting"
+
+    sessions_class.append(BiRefNetMattingSession)
+
+
 def download_rembg_models():
     """Download all rembg ONNX models."""
     print("=== Downloading rembg models ===")
     from rembg import new_session
 
+    _register_birefnet_matting()
+
     for model in REMBG_MODELS:
         print(f"  Downloading {model}...")
         new_session(model)
diff --git a/packages/ai/python/remove_bg.py b/packages/ai/python/remove_bg.py
index 0dd2b77c..26f6ab6b 100644
--- a/packages/ai/python/remove_bg.py
+++ b/packages/ai/python/remove_bg.py
@@ -9,6 +9,39 @@ def emit_progress(percent, stage):
     print(json.dumps({"progress": percent, "stage": stage}), file=sys.stderr, flush=True)
 
 
+_matting_registered = False
+
+def _register_matting_session(sessions_class):
+    """Register the BiRefNet-matting ONNX session for Ultra quality mode."""
+    global _matting_registered
+    if _matting_registered:
+        return
+    _matting_registered = True
+
+    import os
+    import pooch
+    from rembg.sessions.birefnet_general import BiRefNetSessionGeneral
+
+    class BiRefNetMattingSession(BiRefNetSessionGeneral):
+        @classmethod
+        def download_models(cls, *args, **kwargs):
+            fname = f"{cls.name(*args, **kwargs)}.onnx"
+            pooch.retrieve(
+                "https://github.com/ZhengPeng7/BiRefNet/releases/download/v1/BiRefNet-matting-epoch_100.onnx",
+                None,  # Skip checksum for GitHub release assets
+                fname=fname,
+                path=cls.u2net_home(*args, **kwargs),
+                progressbar=True,
+            )
+            return os.path.join(cls.u2net_home(*args, **kwargs), fname)
+
+        @classmethod
+        def name(cls, *args, **kwargs):
+            return "birefnet-matting"
+
+    sessions_class.append(BiRefNetMattingSession)
+
+
 def main():
     input_path = sys.argv[1]
     output_path = sys.argv[2]
@@ -23,8 +56,12 @@ def main():
 
     try:
         from rembg import remove, new_session
+        from rembg.sessions import sessions_class
         from gpu import onnx_providers
 
+        # Register BiRefNet-matting (Ultra quality) if not already present
+        _register_matting_session(sessions_class)
+
         emit_progress(10, "Loading model")
 
         session = new_session(model, providers=onnx_providers())
diff --git a/tests/e2e/remove-bg.spec.ts b/tests/e2e/remove-bg.spec.ts
index 4e4ae037..3666a395 100644
--- a/tests/e2e/remove-bg.spec.ts
+++ b/tests/e2e/remove-bg.spec.ts
@@ -108,6 +108,33 @@ test.describe("Remove Background tool", () => {
     await expect(page.locator("section[aria-label='Image area'] img").first()).toBeVisible();
   });
 
+  test("Ultra quality visible for People, hidden for Products", async ({ loggedInPage: page }) => {
+    await page.goto("/remove-background");
+
+    // People is default - Ultra should be visible
+    await expect(page.getByRole("button", { name: "Ultra" })).toBeVisible();
+
+    // Switch to Products - Ultra should disappear
+    await page.getByText("Products").click();
+    await expect(page.getByRole("button", { name: "Ultra" })).not.toBeVisible();
+
+    // Switch back to People - Ultra returns
+    await page.getByText("People").click();
+    await expect(page.getByRole("button", { name: "Ultra" })).toBeVisible();
+  });
+
+  test("Ultra quality processes JPG portrait", async ({ loggedInPage: page }) => {
+    await page.goto("/remove-background");
+    await uploadFile(page, fixturePath("test-portrait.jpg"));
+
+    // Select Ultra quality
+    await page.getByRole("button", { name: "Ultra" }).click();
+
+    await removeBgAndWait(page);
+    await expect(page.locator("section[aria-label='Image area'] img").first()).toBeVisible();
+    await expect(page.locator("text=Background removal failed")).not.toBeVisible();
+  });
+
   test("HEIC portrait - processes without error", async ({ loggedInPage: page }) => {
     await page.goto("/remove-background");
     await uploadFile(page, fixturePath("test-portrait.heic"));