Spaces:
Running on Zero
Running on Zero
| """Run the FaceAnything model and package the raw predictions.""" | |
| from __future__ import annotations | |
| from dataclasses import dataclass | |
| import numpy as np | |
| import torch | |
| 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) | |