Spaces:
Paused
Paused
File size: 2,535 Bytes
4db294e d35e4df 4db294e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 | """Foreground/background segmentation via Robust Video Matting (RVM).
This isolates the subject (head + shoulders) so the reconstruction, canonical
map and point tracks are not polluted by the background. It reuses the same RVM
model the released pipeline was evaluated with
(``PeterL1n/RobustVideoMatting``, loaded through ``torch.hub``). The model is
recurrent: it carries temporal state across frames, so pass frames in order.
"""
from __future__ import annotations
import os
import numpy as np
import torch
from PIL import Image
def generate_masks(image_paths, output_dir, model_name: str = "resnet50",
downsample_ratio: float = 0.25, warmup_frames: int = 25,
device: str = "cuda", verbose: bool = True):
"""Run RVM on ``image_paths`` (in order) and save an alpha mask per frame.
Masks are written to ``output_dir`` with the same basename as each input
image (single-channel PNG, 255 = foreground).
Returns the list of written mask paths (aligned with ``image_paths``).
"""
os.makedirs(output_dir, exist_ok=True)
dev = torch.device(device if torch.cuda.is_available() else "cpu")
if verbose:
print(f"[faceanything] loading Robust Video Matting ({model_name}) ...", flush=True)
# trust_repo=True avoids the interactive "trust this repo?" prompt, which
# would otherwise raise EOFError in a non-interactive Space.
model = torch.hub.load("PeterL1n/RobustVideoMatting", model_name, trust_repo=True)
model = model.to(dev).eval()
rec = [None] * 4 # recurrent state
def _load(path):
img = Image.open(path).convert("RGB")
t = torch.from_numpy(np.array(img)).permute(2, 0, 1)[None].float().to(dev) / 255.0
return t
# warm up the recurrent state on the first frame
with torch.no_grad():
rgb0 = _load(image_paths[0])
for _ in range(warmup_frames):
_, _, *rec = model(rgb0, *rec, downsample_ratio)
mask_paths = []
with torch.no_grad():
for path in image_paths:
rgb = _load(path)
_, pha, *rec = model(rgb, *rec, downsample_ratio)
alpha = (pha[0, 0].cpu().numpy() * 255.0).astype(np.uint8)
out = os.path.join(output_dir, os.path.basename(path))
out = os.path.splitext(out)[0] + ".png"
Image.fromarray(alpha).save(out)
mask_paths.append(out)
if verbose:
print(f"[faceanything] wrote {len(mask_paths)} masks -> {output_dir}", flush=True)
return mask_paths
|