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}}]}))