Frank Schwenk
83d4468b69
Downscale to 2000px JPEG before rembg, random blend opacity, timeout retries, per-variant remix prefill, lightbox, and longer SFTP idle. Co-authored-by: Cursor <cursoragent@cursor.com>
116 lines
3.9 KiB
Python
116 lines
3.9 KiB
Python
#!/usr/bin/env python3
|
|
"""Compare background-removal backends on the live.f12.rocks inbox uploads.
|
|
|
|
Outputs land next to each other with explicit names so Fränky can pick a winner:
|
|
|
|
<stem>__rembg-u2net__no-alpha.png
|
|
<stem>__rembg-u2net__alpha.png
|
|
<stem>__rembg-default__no-alpha.png
|
|
<stem>__rembg-default__alpha.png
|
|
<stem>__rembg-birefnet-general__no-alpha.png
|
|
<stem>__rembg-birefnet-general__alpha.png
|
|
<stem>__withoutbg-open-weights.png
|
|
|
|
`rembg-default` = rembg.remove() with no session (library default model).
|
|
`rembg-u2net` = rembg.new_session("u2net") explicitly.
|
|
withoutbg embeds matting in its open-weights graph — one output only.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
import sys
|
|
import time
|
|
import traceback
|
|
from pathlib import Path
|
|
|
|
os.environ.setdefault("TQDM_DISABLE", "1")
|
|
|
|
INPUT_DIR = Path(os.environ.get("COMPARE_INPUT", "/data/inputs"))
|
|
OUTPUT_DIR = Path(os.environ.get("COMPARE_OUTPUT", "/data/outputs"))
|
|
EXTS = {".jpg", ".jpeg", ".png", ".webp"}
|
|
|
|
|
|
def log(msg: str) -> None:
|
|
print(msg, flush=True)
|
|
|
|
|
|
def run_rembg_named(data: bytes, label: str, model: str | None, alpha: bool, out: Path) -> None:
|
|
from rembg import new_session, remove
|
|
|
|
t0 = time.monotonic()
|
|
kwargs: dict = {"alpha_matting": alpha}
|
|
if model is None:
|
|
result = remove(data, **kwargs)
|
|
model_note = "default"
|
|
else:
|
|
session = new_session(model)
|
|
result = remove(data, session=session, **kwargs)
|
|
model_note = model
|
|
out.write_bytes(result)
|
|
log(f" OK {label} model={model_note} alpha={alpha} {time.monotonic()-t0:.1f}s -> {out.name} ({out.stat().st_size} bytes)")
|
|
|
|
|
|
def run_withoutbg(path: Path, out: Path) -> None:
|
|
from withoutbg import WithoutBG
|
|
|
|
t0 = time.monotonic()
|
|
model = WithoutBG.open_weights()
|
|
result = model.remove_background(str(path))
|
|
result.save(out)
|
|
log(f" OK withoutbg-open-weights {time.monotonic()-t0:.1f}s -> {out.name} ({out.stat().st_size} bytes)")
|
|
|
|
|
|
def process_one(path: Path) -> None:
|
|
stem = path.stem
|
|
log(f"=== {path.name} ===")
|
|
data = path.read_bytes()
|
|
jobs = [
|
|
("rembg-u2net", "u2net", False),
|
|
("rembg-u2net", "u2net", True),
|
|
("rembg-default", None, False),
|
|
("rembg-default", None, True),
|
|
("rembg-birefnet-general", "birefnet-general", False),
|
|
("rembg-birefnet-general", "birefnet-general", True),
|
|
]
|
|
for label, model, alpha in jobs:
|
|
tag = "alpha" if alpha else "no-alpha"
|
|
out = OUTPUT_DIR / f"{stem}__{label}__{tag}.png"
|
|
if out.exists() and out.stat().st_size > 0:
|
|
log(f" SKIP {out.name} (exists)")
|
|
continue
|
|
try:
|
|
run_rembg_named(data, label, model, alpha, out)
|
|
except Exception as exc: # noqa: BLE001
|
|
err = OUTPUT_DIR / f"{stem}__{label}__{tag}.ERROR.txt"
|
|
err.write_text(f"{type(exc).__name__}: {exc}\n\n{traceback.format_exc()}", encoding="utf-8")
|
|
log(f" FAIL {label} alpha={alpha}: {type(exc).__name__}: {exc}")
|
|
|
|
out_w = OUTPUT_DIR / f"{stem}__withoutbg-open-weights.png"
|
|
if out_w.exists() and out_w.stat().st_size > 0:
|
|
log(f" SKIP {out_w.name} (exists)")
|
|
else:
|
|
try:
|
|
run_withoutbg(path, out_w)
|
|
except Exception as exc: # noqa: BLE001
|
|
err = OUTPUT_DIR / f"{stem}__withoutbg-open-weights.ERROR.txt"
|
|
err.write_text(f"{type(exc).__name__}: {exc}\n\n{traceback.format_exc()}", encoding="utf-8")
|
|
log(f" FAIL withoutbg: {type(exc).__name__}: {exc}")
|
|
|
|
|
|
def main() -> int:
|
|
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
|
|
inputs = sorted(p for p in INPUT_DIR.iterdir() if p.is_file() and p.suffix.lower() in EXTS)
|
|
if not inputs:
|
|
log(f"No images in {INPUT_DIR}")
|
|
return 1
|
|
log(f"Comparing {len(inputs)} image(s); outputs -> {OUTPUT_DIR}")
|
|
for path in inputs:
|
|
process_one(path)
|
|
log("DONE")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|