8a11325b2f
rembg with ONNX Runtime on CPU was chosen over BiRefNet: the latter requires deform_conv2d which crashes on ROCm bfloat16 and runs at 40s in float32. rembg delivers 0.23s per 1024x1024 image, no GPU deps, and isnet-anime is trained specifically on anime images — exactly the target domain. CLI and session API both support single-image and batch (plural-first) modes, reusing one model session across files.
100 lines
3.0 KiB
Python
100 lines
3.0 KiB
Python
"""Background removal for character sprites using rembg + isnet-anime.
|
||
|
||
Uses onnxruntime on CPU — the model is ~176 MB and runs in ~0.25s per
|
||
1024×1024 image on Strix Halo. No GPU dependency keeps things simple.
|
||
"""
|
||
import time
|
||
from pathlib import Path
|
||
from typing import Sequence
|
||
|
||
from PIL import Image
|
||
from rembg import new_session, remove
|
||
|
||
|
||
def _load_session(model: str = "isnet-anime"):
|
||
"""Load a rembg session for the given model. Cached in the caller."""
|
||
return new_session(model)
|
||
|
||
|
||
def remove_background(
|
||
input_path: str,
|
||
output_path: str,
|
||
model: str = "isnet-anime",
|
||
session=None,
|
||
) -> None:
|
||
"""Remove background from a single image. Output is RGBA PNG.
|
||
|
||
Args:
|
||
input_path: Path to the input RGB image.
|
||
output_path: Where to save the RGBA PNG.
|
||
model: rembg model name. Default "isnet-anime" is trained on anime
|
||
images. Other options: "u2net", "u2netp", "u2net_human_seg",
|
||
"isnet-general-use", "sam".
|
||
session: A rembg session (from ``new_session()``). If None, one
|
||
is created and immediately discarded. For batch use, create a
|
||
session once and pass it to avoid reloading the model.
|
||
|
||
Raises:
|
||
FileNotFoundError: If input_path doesn't exist.
|
||
"""
|
||
input_file = Path(input_path)
|
||
if not input_file.exists():
|
||
raise FileNotFoundError(f"Input image not found: {input_path}")
|
||
|
||
output_file = Path(output_path)
|
||
output_file.parent.mkdir(parents=True, exist_ok=True)
|
||
|
||
_session = session
|
||
close_session = session is None
|
||
if _session is None:
|
||
_session = _load_session(model)
|
||
|
||
img = Image.open(input_file).convert("RGB")
|
||
|
||
t0 = time.perf_counter()
|
||
result = remove(img, session=_session)
|
||
t_elapsed = round(time.perf_counter() - t0, 3)
|
||
|
||
result.save(output_file)
|
||
print(f"Saved {output_file} (bg removed, {t_elapsed}s)")
|
||
|
||
if close_session:
|
||
del _session
|
||
|
||
|
||
def remove_backgrounds(
|
||
input_paths: Sequence[str],
|
||
output_dir: str,
|
||
model: str = "isnet-anime",
|
||
) -> None:
|
||
"""Remove backgrounds from multiple images, reusing one model session.
|
||
|
||
Output files are named ``{stem}_nobg.png`` in ``output_dir``.
|
||
|
||
Args:
|
||
input_paths: List of input image paths.
|
||
output_dir: Directory for output RGBA PNGs.
|
||
model: rembg model name. Default "isnet-anime".
|
||
"""
|
||
out_dir = Path(output_dir)
|
||
out_dir.mkdir(parents=True, exist_ok=True)
|
||
|
||
session = _load_session(model)
|
||
try:
|
||
total = 0.0
|
||
for input_path in input_paths:
|
||
stem = Path(input_path).stem
|
||
output_path = out_dir / f"{stem}_nobg.png"
|
||
|
||
img = Image.open(input_path).convert("RGB")
|
||
t0 = time.perf_counter()
|
||
result = remove(img, session=session)
|
||
t_elapsed = time.perf_counter() - t0
|
||
total += t_elapsed
|
||
result.save(output_path)
|
||
print(f"Saved {output_path} ({t_elapsed:.3f}s)")
|
||
|
||
print(f"Done: {len(input_paths)} images, {total:.2f}s total")
|
||
finally:
|
||
del session
|