mirror of
https://github.com/snapotter-hq/SnapOtter.git
synced 2026-08-03 07:46:42 +02:00
feat: upgrade to BiRefNet SOTA background removal model
- Switch default model from U2-Net to BiRefNet (state-of-the-art) - Add 6 model options: BiRefNet, BiRefNet Lite, BiRefNet Portrait, BRIA RMBG, IS-Net, U2-Net - Add animated progress bar with stage indicators (loading model, analyzing, removing, refining edges) and elapsed timer - Add intuitive background color presets (Transparent, White, Black, Red, Green, Blue) as clickable buttons + custom color picker - Handle background color compositing in Python (PIL alpha composite) - Add checkerboard pattern to before/after slider for transparency - Pre-bake BiRefNet model (973MB) in Docker image for instant use
This commit is contained in:
@@ -1,17 +1,54 @@
|
|||||||
import { useState } from "react";
|
import { useState, useEffect, useRef } from "react";
|
||||||
import { useFileStore } from "@/stores/file-store";
|
import { useFileStore } from "@/stores/file-store";
|
||||||
import { useToolProcessor } from "@/hooks/use-tool-processor";
|
import { useToolProcessor } from "@/hooks/use-tool-processor";
|
||||||
import { Download, Loader2 } from "lucide-react";
|
import { Download, Loader2 } from "lucide-react";
|
||||||
|
|
||||||
type BgModel = "u2net" | "isnet";
|
type BgModel =
|
||||||
|
| "birefnet-general"
|
||||||
|
| "birefnet-general-lite"
|
||||||
|
| "birefnet-portrait"
|
||||||
|
| "bria-rmbg"
|
||||||
|
| "isnet-general-use"
|
||||||
|
| "u2net";
|
||||||
|
|
||||||
|
const MODELS: { value: BgModel; label: string; description: string }[] = [
|
||||||
|
{ value: "birefnet-general", label: "BiRefNet", description: "Best quality (recommended)" },
|
||||||
|
{ value: "birefnet-general-lite", label: "BiRefNet Lite", description: "Faster, slightly less accurate" },
|
||||||
|
{ value: "birefnet-portrait", label: "BiRefNet Portrait", description: "Optimized for people" },
|
||||||
|
{ value: "bria-rmbg", label: "BRIA RMBG", description: "Great for products" },
|
||||||
|
{ value: "isnet-general-use", label: "IS-Net", description: "Good general purpose" },
|
||||||
|
{ value: "u2net", label: "U2-Net", description: "Classic, fast" },
|
||||||
|
];
|
||||||
|
|
||||||
|
const BG_PRESETS = [
|
||||||
|
{ color: "", label: "Transparent", preview: "checkerboard" },
|
||||||
|
{ color: "#FFFFFF", label: "White", preview: "#FFFFFF" },
|
||||||
|
{ color: "#000000", label: "Black", preview: "#000000" },
|
||||||
|
{ color: "#FF0000", label: "Red", preview: "#FF0000" },
|
||||||
|
{ color: "#00FF00", label: "Green", preview: "#00FF00" },
|
||||||
|
{ color: "#0000FF", label: "Blue", preview: "#0000FF" },
|
||||||
|
];
|
||||||
|
|
||||||
export function RemoveBgSettings() {
|
export function RemoveBgSettings() {
|
||||||
const { files } = useFileStore();
|
const { files } = useFileStore();
|
||||||
const { processFiles, processing, error, downloadUrl, originalSize, processedSize } =
|
const { processFiles, processing, error, downloadUrl, originalSize, processedSize } =
|
||||||
useToolProcessor("remove-background");
|
useToolProcessor("remove-background");
|
||||||
|
|
||||||
const [model, setModel] = useState<BgModel>("u2net");
|
const [model, setModel] = useState<BgModel>("birefnet-general");
|
||||||
const [bgColor, setBgColor] = useState("");
|
const [bgColor, setBgColor] = useState("");
|
||||||
|
const [elapsed, setElapsed] = useState(0);
|
||||||
|
const timerRef = useRef<ReturnType<typeof setInterval> | null>(null);
|
||||||
|
|
||||||
|
// Progress timer
|
||||||
|
useEffect(() => {
|
||||||
|
if (processing) {
|
||||||
|
setElapsed(0);
|
||||||
|
timerRef.current = setInterval(() => setElapsed((e) => e + 1), 1000);
|
||||||
|
} else {
|
||||||
|
if (timerRef.current) clearInterval(timerRef.current);
|
||||||
|
}
|
||||||
|
return () => { if (timerRef.current) clearInterval(timerRef.current); };
|
||||||
|
}, [processing]);
|
||||||
|
|
||||||
const handleProcess = () => {
|
const handleProcess = () => {
|
||||||
const settings: Record<string, unknown> = { model };
|
const settings: Record<string, unknown> = { model };
|
||||||
@@ -21,6 +58,15 @@ export function RemoveBgSettings() {
|
|||||||
|
|
||||||
const hasFile = files.length > 0;
|
const hasFile = files.length > 0;
|
||||||
|
|
||||||
|
const progressStage =
|
||||||
|
elapsed < 3 ? "Loading AI model..." :
|
||||||
|
elapsed < 8 ? "Analyzing image..." :
|
||||||
|
elapsed < 15 ? "Removing background..." :
|
||||||
|
elapsed < 25 ? "Refining edges..." :
|
||||||
|
"Almost done...";
|
||||||
|
|
||||||
|
const progressPercent = Math.min(95, elapsed * 4);
|
||||||
|
|
||||||
return (
|
return (
|
||||||
<div className="space-y-4">
|
<div className="space-y-4">
|
||||||
{/* Model selector */}
|
{/* Model selector */}
|
||||||
@@ -31,51 +77,71 @@ export function RemoveBgSettings() {
|
|||||||
onChange={(e) => setModel(e.target.value as BgModel)}
|
onChange={(e) => setModel(e.target.value as BgModel)}
|
||||||
className="w-full mt-0.5 px-2 py-1.5 rounded border border-border bg-background text-sm text-foreground"
|
className="w-full mt-0.5 px-2 py-1.5 rounded border border-border bg-background text-sm text-foreground"
|
||||||
>
|
>
|
||||||
<option value="u2net">U2-Net (General purpose)</option>
|
{MODELS.map((m) => (
|
||||||
<option value="isnet">IS-Net (Higher accuracy)</option>
|
<option key={m.value} value={m.value}>
|
||||||
|
{m.label} — {m.description}
|
||||||
|
</option>
|
||||||
|
))}
|
||||||
</select>
|
</select>
|
||||||
</div>
|
</div>
|
||||||
|
|
||||||
{/* Background color */}
|
{/* Background color - intuitive preset buttons */}
|
||||||
<div>
|
<div>
|
||||||
<label className="text-xs text-muted-foreground">
|
<label className="text-sm font-medium text-muted-foreground">
|
||||||
Replacement Background (leave empty for transparent)
|
Output Background
|
||||||
</label>
|
</label>
|
||||||
<div className="flex gap-2 mt-0.5">
|
<div className="flex gap-1.5 mt-1.5 flex-wrap">
|
||||||
|
{BG_PRESETS.map((preset) => (
|
||||||
|
<button
|
||||||
|
key={preset.label}
|
||||||
|
onClick={() => setBgColor(preset.color)}
|
||||||
|
className={`flex items-center gap-1.5 px-2.5 py-1.5 rounded-lg border text-xs font-medium transition-colors ${
|
||||||
|
bgColor === preset.color
|
||||||
|
? "border-primary bg-primary/10 text-primary"
|
||||||
|
: "border-border text-muted-foreground hover:border-primary/50"
|
||||||
|
}`}
|
||||||
|
>
|
||||||
|
<span
|
||||||
|
className="w-4 h-4 rounded-sm border border-border shrink-0"
|
||||||
|
style={
|
||||||
|
preset.preview === "checkerboard"
|
||||||
|
? {
|
||||||
|
backgroundImage:
|
||||||
|
"linear-gradient(45deg, #ccc 25%, transparent 25%), linear-gradient(-45deg, #ccc 25%, transparent 25%), linear-gradient(45deg, transparent 75%, #ccc 75%), linear-gradient(-45deg, transparent 75%, #ccc 75%)",
|
||||||
|
backgroundSize: "8px 8px",
|
||||||
|
backgroundPosition: "0 0, 0 4px, 4px -4px, -4px 0px",
|
||||||
|
}
|
||||||
|
: { backgroundColor: preset.preview }
|
||||||
|
}
|
||||||
|
/>
|
||||||
|
{preset.label}
|
||||||
|
</button>
|
||||||
|
))}
|
||||||
|
</div>
|
||||||
|
|
||||||
|
{/* Custom color picker */}
|
||||||
|
<div className="flex items-center gap-2 mt-2">
|
||||||
<input
|
<input
|
||||||
type="color"
|
type="color"
|
||||||
value={bgColor || "#ffffff"}
|
value={bgColor || "#ffffff"}
|
||||||
onChange={(e) => setBgColor(e.target.value)}
|
onChange={(e) => setBgColor(e.target.value)}
|
||||||
className="w-10 h-8 rounded border border-border cursor-pointer"
|
className="w-8 h-8 rounded border border-border cursor-pointer"
|
||||||
/>
|
/>
|
||||||
<input
|
<input
|
||||||
type="text"
|
type="text"
|
||||||
value={bgColor}
|
value={bgColor}
|
||||||
onChange={(e) => setBgColor(e.target.value)}
|
onChange={(e) => setBgColor(e.target.value)}
|
||||||
placeholder="Transparent"
|
placeholder="Custom hex (#FF5500)"
|
||||||
className="flex-1 px-2 py-1.5 rounded border border-border bg-background text-sm text-foreground"
|
className="flex-1 px-2 py-1.5 rounded border border-border bg-background text-xs text-foreground"
|
||||||
/>
|
/>
|
||||||
{bgColor && (
|
|
||||||
<button
|
|
||||||
onClick={() => setBgColor("")}
|
|
||||||
className="text-xs text-muted-foreground hover:text-foreground"
|
|
||||||
>
|
|
||||||
Clear
|
|
||||||
</button>
|
|
||||||
)}
|
|
||||||
</div>
|
</div>
|
||||||
</div>
|
</div>
|
||||||
|
|
||||||
{/* Info */}
|
|
||||||
<p className="text-[10px] text-muted-foreground">
|
|
||||||
Requires Python with rembg installed. Works best with photos of people, products, and animals.
|
|
||||||
</p>
|
|
||||||
|
|
||||||
{/* Error */}
|
{/* Error */}
|
||||||
{error && <p className="text-xs text-red-500">{error}</p>}
|
{error && <p className="text-xs text-red-500">{error}</p>}
|
||||||
|
|
||||||
{/* Size info */}
|
{/* Size info */}
|
||||||
{originalSize != null && processedSize != null && (
|
{originalSize != null && processedSize != null && !processing && (
|
||||||
<div className="text-xs text-muted-foreground space-y-0.5">
|
<div className="text-xs text-muted-foreground space-y-0.5">
|
||||||
<p>Original: {(originalSize / 1024).toFixed(1)} KB</p>
|
<p>Original: {(originalSize / 1024).toFixed(1)} KB</p>
|
||||||
<p>Processed: {(processedSize / 1024).toFixed(1)} KB</p>
|
<p>Processed: {(processedSize / 1024).toFixed(1)} KB</p>
|
||||||
@@ -92,20 +158,27 @@ export function RemoveBgSettings() {
|
|||||||
{processing ? "Removing Background..." : "Remove Background"}
|
{processing ? "Removing Background..." : "Remove Background"}
|
||||||
</button>
|
</button>
|
||||||
|
|
||||||
{/* Progress indicator */}
|
{/* Animated progress bar with stages */}
|
||||||
{processing && (
|
{processing && (
|
||||||
<div className="space-y-2">
|
<div className="space-y-2 p-3 rounded-lg bg-muted/50 border border-border">
|
||||||
<div className="w-full bg-muted rounded-full h-2 overflow-hidden">
|
<div className="flex justify-between text-xs text-muted-foreground">
|
||||||
<div className="h-full bg-primary rounded-full animate-pulse" style={{ width: '100%' }} />
|
<span>{progressStage}</span>
|
||||||
|
<span>{elapsed}s</span>
|
||||||
</div>
|
</div>
|
||||||
<p className="text-xs text-muted-foreground text-center">
|
<div className="w-full bg-muted rounded-full h-2.5 overflow-hidden">
|
||||||
AI processing may take 10-30 seconds...
|
<div
|
||||||
|
className="h-full bg-primary rounded-full transition-all duration-1000 ease-out"
|
||||||
|
style={{ width: `${progressPercent}%` }}
|
||||||
|
/>
|
||||||
|
</div>
|
||||||
|
<p className="text-[10px] text-muted-foreground text-center">
|
||||||
|
First run may take longer as the model loads into memory
|
||||||
</p>
|
</p>
|
||||||
</div>
|
</div>
|
||||||
)}
|
)}
|
||||||
|
|
||||||
{/* Download */}
|
{/* Download */}
|
||||||
{downloadUrl && (
|
{downloadUrl && !processing && (
|
||||||
<a
|
<a
|
||||||
href={downloadUrl}
|
href={downloadUrl}
|
||||||
download
|
download
|
||||||
|
|||||||
+5
-3
@@ -71,10 +71,12 @@ RUN /opt/venv/bin/pip install --no-cache-dir --upgrade pip && \
|
|||||||
# This makes the Docker image fully self-contained — works offline
|
# This makes the Docker image fully self-contained — works offline
|
||||||
RUN /opt/venv/bin/python3 -c "\
|
RUN /opt/venv/bin/python3 -c "\
|
||||||
from rembg import new_session; \
|
from rembg import new_session; \
|
||||||
print('Downloading u2net model...'); \
|
print('Downloading BiRefNet model (SOTA)...'); \
|
||||||
|
new_session('birefnet-general'); \
|
||||||
|
print('Downloading u2net model (fallback)...'); \
|
||||||
new_session('u2net'); \
|
new_session('u2net'); \
|
||||||
print('u2net model ready') \
|
print('Background removal models ready') \
|
||||||
" 2>/dev/null || echo "WARNING: Could not pre-download u2net model"
|
" 2>/dev/null || echo "WARNING: Could not pre-download rembg models"
|
||||||
|
|
||||||
RUN /opt/venv/bin/python3 -c "\
|
RUN /opt/venv/bin/python3 -c "\
|
||||||
try: \
|
try: \
|
||||||
|
|||||||
@@ -1,4 +1,4 @@
|
|||||||
"""Background removal using rembg."""
|
"""Background removal using rembg with state-of-the-art BiRefNet models."""
|
||||||
import sys
|
import sys
|
||||||
import json
|
import json
|
||||||
|
|
||||||
@@ -8,25 +8,48 @@ def main():
|
|||||||
output_path = sys.argv[2]
|
output_path = sys.argv[2]
|
||||||
settings = json.loads(sys.argv[3]) if len(sys.argv) > 3 else {}
|
settings = json.loads(sys.argv[3]) if len(sys.argv) > 3 else {}
|
||||||
|
|
||||||
model = settings.get("model", "u2net")
|
model = settings.get("model", "birefnet-general")
|
||||||
|
bg_color = settings.get("backgroundColor", "")
|
||||||
|
|
||||||
try:
|
try:
|
||||||
from rembg import remove
|
from rembg import remove, new_session
|
||||||
|
from PIL import Image
|
||||||
|
import io
|
||||||
|
|
||||||
|
print(json.dumps({"progress": "loading_model"}), flush=True)
|
||||||
|
|
||||||
|
# Create a session with the selected model
|
||||||
|
session = new_session(model)
|
||||||
|
|
||||||
with open(input_path, "rb") as f:
|
with open(input_path, "rb") as f:
|
||||||
input_data = f.read()
|
input_data = f.read()
|
||||||
|
|
||||||
# Try with alpha matting first for better edges, but fall back
|
print(json.dumps({"progress": "processing"}), flush=True)
|
||||||
# without it if the image triggers the known rembg matting error
|
|
||||||
|
# Try with alpha matting first for better edges
|
||||||
try:
|
try:
|
||||||
output_data = remove(
|
output_data = remove(
|
||||||
input_data,
|
input_data,
|
||||||
|
session=session,
|
||||||
alpha_matting=True,
|
alpha_matting=True,
|
||||||
alpha_matting_foreground_threshold=240,
|
alpha_matting_foreground_threshold=240,
|
||||||
alpha_matting_background_threshold=10,
|
alpha_matting_background_threshold=10,
|
||||||
)
|
)
|
||||||
except Exception:
|
except Exception:
|
||||||
output_data = remove(input_data)
|
output_data = remove(input_data, session=session)
|
||||||
|
|
||||||
|
# If a background color is specified, composite onto it
|
||||||
|
if bg_color and bg_color.startswith("#"):
|
||||||
|
img = Image.open(io.BytesIO(output_data)).convert("RGBA")
|
||||||
|
hex_color = bg_color.lstrip("#")
|
||||||
|
r = int(hex_color[0:2], 16)
|
||||||
|
g = int(hex_color[2:4], 16)
|
||||||
|
b = int(hex_color[4:6], 16)
|
||||||
|
bg = Image.new("RGBA", img.size, (r, g, b, 255))
|
||||||
|
bg.paste(img, mask=img.split()[3])
|
||||||
|
buf = io.BytesIO()
|
||||||
|
bg.save(buf, format="PNG")
|
||||||
|
output_data = buf.getvalue()
|
||||||
|
|
||||||
with open(output_path, "wb") as f:
|
with open(output_path, "wb") as f:
|
||||||
f.write(output_data)
|
f.write(output_data)
|
||||||
|
|||||||
Reference in New Issue
Block a user