Spaces:
Paused
Paused
File size: 4,280 Bytes
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 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 | """Run the FaceAnything model and package the raw predictions."""
from __future__ import annotations
from dataclasses import dataclass
import numpy as np
import torch
@dataclass
class FacePrediction:
"""Per-clip model outputs (all numpy, N = number of frames)."""
depth: np.ndarray # (N, H, W) float32
intrinsics: np.ndarray # (N, 3, 3) float32
extrinsics: np.ndarray # (N, 4, 4) float32 world-to-camera
images: np.ndarray # (N, H, W, 3) uint8 (model-processed)
canonical: np.ndarray | None # (N, H, W, 3) float32 canonical coords
conf: np.ndarray | None # (N, H, W) float32 depth confidence
valid: np.ndarray # (N, H, W) bool foreground/usable mask
def _to_4x4(ext: np.ndarray) -> np.ndarray:
if ext.shape[-2:] == (4, 4):
return ext.astype(np.float32)
out = np.tile(np.eye(4, dtype=np.float32), (ext.shape[0], 1, 1))
out[:, :3, :4] = ext
return out
def run_inference(model, frame_paths, mask_paths=None, process_res: int = 504,
use_ray_pose: bool = True, monocular: bool = True,
conf_percentile: float = 0.0, per_frame: bool = False) -> FacePrediction:
"""Run the model on a list of frame paths and assemble a ``FacePrediction``.
Args:
model: a loaded FaceAnything model.
frame_paths: list of image paths (one clip).
mask_paths: optional list (aligned with frames) of foreground mask paths;
masked-out pixels are dropped from all 3D products.
process_res: model processing resolution (square upper bound).
use_ray_pose: use ray-based pose instead of the camera-decoder pose.
monocular: if True, replace predicted extrinsics with identity so every
frame's cloud lives in its own camera frame (matches the released
evaluation pipeline). If False, keep predicted poses (multi-view
consistent world frame).
conf_percentile: drop pixels below this depth-confidence percentile
(0 disables).
"""
import cv2
def _infer(paths):
with torch.no_grad():
p = model.inference(paths, export_dir=None, use_ray_pose=use_ray_pose,
process_res=process_res)
defo = (np.asarray(p.deformation, np.float32)
if getattr(p, "deformation", None) is not None else None)
cf = (np.asarray(p.conf, np.float32)
if getattr(p, "conf", None) is not None else None)
return (np.asarray(p.depth, np.float32), np.asarray(p.processed_images),
np.asarray(p.intrinsics, np.float32),
np.asarray(p.extrinsics, np.float32), defo, cf)
if per_frame:
# one frame at a time: lower peak memory, so process_res can be larger
parts = [_infer([fp]) for fp in frame_paths]
cat = lambda j: np.concatenate([p[j] for p in parts], axis=0)
depth, images, intr, ext_raw = cat(0), cat(1), cat(2), cat(3)
canonical = cat(4) if parts[0][4] is not None else None
conf = cat(5) if parts[0][5] is not None else None
else:
depth, images, intr, ext_raw, canonical, conf = _infer(frame_paths)
ext = _to_4x4(ext_raw)
N, H, W = depth.shape
if monocular:
ext = np.tile(np.eye(4, dtype=np.float32), (N, 1, 1))
valid = np.isfinite(depth) & (depth > 0)
# Optional foreground masking (e.g. background removal).
if mask_paths is not None:
for i, mp in enumerate(mask_paths):
if mp is None:
continue
m = cv2.imread(mp)
if m is None:
continue
m = cv2.resize(m, (W, H), interpolation=cv2.INTER_AREA)
fg = m.mean(axis=2) >= 128
valid[i] &= fg
# Optional confidence thresholding.
if conf is not None and conf_percentile and conf_percentile > 0:
for i in range(N):
v = valid[i]
if v.any():
thr = np.percentile(conf[i][v], conf_percentile)
valid[i] &= conf[i] >= thr
return FacePrediction(depth=depth, intrinsics=intr, extrinsics=ext,
images=images, canonical=canonical, conf=conf,
valid=valid)
|