| """BEN2 background removal for the source image.""" |
|
|
| import os |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
| from huggingface_hub import snapshot_download |
| from PIL import Image |
| from safetensors.torch import load_file |
|
|
|
|
| REVISION = "e48a20765fb421d19dcdb0bf3cc61e802ca5ec8f" |
| MEAN = (0.485, 0.456, 0.406) |
| STD = (0.229, 0.224, 0.225) |
|
|
|
|
| def load_background(): |
| from ben2.model import BEN_Base |
|
|
| source = os.getenv("IDOL_BACKGROUND_MODEL", "PramaLLC/BEN2") |
| path = Path(source) |
| if not path.is_dir(): |
| path = Path( |
| snapshot_download( |
| source, |
| revision=os.getenv("IDOL_BACKGROUND_REVISION", REVISION), |
| allow_patterns=["model.safetensors"], |
| ) |
| ) |
| model = BEN_Base() |
| model.load_state_dict(load_file(str(path / "model.safetensors")), strict=True) |
| return model.eval() |
|
|
|
|
| def remove_background(model, image): |
| """Return the image as RGBA with BEN2's alpha matte, at its original size.""" |
| rgb = image.convert("RGB") |
| pixels = np.asarray(rgb.resize((1024, 1024), Image.Resampling.LANCZOS), dtype=np.float32) / 255 |
| |
| pixels = torch.from_numpy(pixels).permute(2, 0, 1).to("cuda", torch.float16) |
| mean = torch.tensor(MEAN, dtype=torch.float16, device="cuda").view(3, 1, 1) |
| std = torch.tensor(STD, dtype=torch.float16, device="cuda").view(3, 1, 1) |
| alpha = model(((pixels - mean) / std)[None]) |
| alpha = torch.nn.functional.interpolate(alpha, size=(rgb.height, rgb.width), mode="bilinear")[0, 0] |
| alpha = (alpha - alpha.min()) / (alpha.max() - alpha.min()) |
| matte = Image.fromarray((alpha * 255).cpu().numpy().astype(np.uint8)) |
| result = rgb.copy() |
| result.putalpha(matte) |
| return result |
|
|
