1cc0b80d89
New operator: briaai/RMBG-2.0 image->transparent cutout, fully local on Apple Silicon (weights unlocked by owner HF license). Free local alternative to the cloud fal_bg_remove; the recommended pre-pass before SF3D / image->3D. Verified: clean 1024px cutout of the FLUX astrolabe (alpha 0-255, 86.6% removed, thin rings preserved). Same OpenMP guardrails as sf3d. scripts/install_rmbg.sh; venvs/rmbg. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
69 lines
2.2 KiB
Python
69 lines
2.2 KiB
Python
import argparse
|
|
import json
|
|
import os
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
# Same libomp collision guardrails as sf3d (kornia/timm/torch each may link one).
|
|
os.environ.setdefault("PYTORCH_ENABLE_MPS_FALLBACK", "1")
|
|
os.environ.setdefault("KMP_DUPLICATE_LIB_OK", "TRUE")
|
|
os.environ.setdefault("OMP_NUM_THREADS", "1")
|
|
os.environ.setdefault("MKL_NUM_THREADS", "1")
|
|
|
|
import torch # noqa: E402
|
|
from PIL import Image # noqa: E402
|
|
from torchvision import transforms # noqa: E402
|
|
from transformers import AutoModelForImageSegmentation # noqa: E402
|
|
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument("--input", action="append", default=[])
|
|
ap.add_argument("--outdir", required=True)
|
|
ap.add_argument("--params", default="{}")
|
|
a = ap.parse_args()
|
|
p = json.loads(a.params)
|
|
|
|
if not a.input:
|
|
print("ERROR: no input image")
|
|
sys.exit(1)
|
|
|
|
res = int(p.get("resolution", 1024))
|
|
bg = p.get("background", "transparent")
|
|
device = "mps" if torch.backends.mps.is_available() else "cpu"
|
|
print(f"loading briaai/RMBG-2.0 on {device} (first run downloads weights) ...", flush=True)
|
|
|
|
model = AutoModelForImageSegmentation.from_pretrained("briaai/RMBG-2.0", trust_remote_code=True)
|
|
try:
|
|
torch.set_float32_matmul_precision("high")
|
|
except Exception:
|
|
pass
|
|
model.eval().to(device)
|
|
|
|
transform = transforms.Compose([
|
|
transforms.Resize((res, res)),
|
|
transforms.ToTensor(),
|
|
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
|
|
])
|
|
|
|
src = Path(a.input[0])
|
|
image = Image.open(src).convert("RGB")
|
|
inp = transform(image).unsqueeze(0).to(device)
|
|
print("running segmentation ...", flush=True)
|
|
with torch.no_grad():
|
|
preds = model(inp)[-1].sigmoid().cpu()
|
|
mask = transforms.ToPILImage()(preds[0].squeeze()).resize(image.size)
|
|
|
|
outdir = Path(a.outdir)
|
|
name = f"{src.stem}_cutout.png"
|
|
if bg == "transparent":
|
|
out = image.copy()
|
|
out.putalpha(mask)
|
|
else:
|
|
fill = (255, 255, 255) if bg == "white" else (0, 0, 0)
|
|
canvas = Image.new("RGB", image.size, fill)
|
|
canvas.paste(image, (0, 0), mask)
|
|
out = canvas
|
|
out.save(outdir / name)
|
|
print(f"done: {name}", flush=True)
|
|
(outdir / "result.json").write_text(json.dumps(
|
|
{"outputs": [{"path": name, "meta": {"tool": "rmbg-2.0", "resolution": res, "background": bg}}]}))
|