Initial snapshot: hand motion pipeline (Dyn-HaMR + dex-retargeting + SPIDER)
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,3 @@
|
||||
"""
|
||||
Utility functions module.
|
||||
"""
|
||||
@@ -0,0 +1,68 @@
|
||||
"""
|
||||
Checkpoint loading utilities.
|
||||
"""
|
||||
|
||||
import gc
|
||||
from pathlib import Path
|
||||
from typing import Dict, Optional, Tuple
|
||||
|
||||
import torch
|
||||
|
||||
from model.flow_matching.denoiser import HandPoseDenoiser
|
||||
|
||||
|
||||
def _clean_state_dict(state_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
|
||||
return {
|
||||
(k[len("module.") :] if k.startswith("module.") else k): v
|
||||
for k, v in state_dict.items()
|
||||
}
|
||||
|
||||
|
||||
def _load_state_dict(
|
||||
ckpt_path: Path, device: torch.device
|
||||
) -> Tuple[Dict[str, torch.Tensor], Optional[int]]:
|
||||
# Load to CPU first to avoid a large checkpoint filling GPU memory
|
||||
raw = None
|
||||
if ckpt_path.is_dir():
|
||||
state_file = ckpt_path / "mp_rank_00_model_states.pt"
|
||||
if not state_file.exists():
|
||||
raise FileNotFoundError(f"DeepSpeed checkpoint not found: {state_file}")
|
||||
raw = torch.load(state_file, map_location="cpu", weights_only=False)
|
||||
else:
|
||||
raw = torch.load(str(ckpt_path), map_location="cpu", weights_only=False)
|
||||
|
||||
state_dict = raw.get("module", raw) if isinstance(raw, dict) else raw
|
||||
epoch = raw.get("epoch") if isinstance(raw, dict) else None
|
||||
del raw
|
||||
gc.collect()
|
||||
|
||||
# Move state_dict to the target device
|
||||
state_dict = {k: v.to(device) for k, v in state_dict.items()}
|
||||
return _clean_state_dict(state_dict), epoch
|
||||
|
||||
|
||||
def load_denoiser_from_ckpt(
|
||||
config,
|
||||
ckpt_path: str,
|
||||
device: torch.device,
|
||||
) -> HandPoseDenoiser:
|
||||
"""Load HandPoseDenoiser, supporting either a DeepSpeed directory or a .pt file."""
|
||||
denoiser = HandPoseDenoiser(config).to(device)
|
||||
ckpt = Path(ckpt_path)
|
||||
state_dict, epoch = _load_state_dict(ckpt, device)
|
||||
result = denoiser.load_state_dict(state_dict, strict=False)
|
||||
# The frozen HaMeR backbone is loaded separately from HAMER_CKPT by HaMeRBackbone
|
||||
# and is intentionally absent from the released denoiser checkpoint; filter its keys
|
||||
# out of the warning so the log stays focused on genuinely unexpected mismatches.
|
||||
missing = [k for k in result.missing_keys if not k.startswith("hamer_backbone.")]
|
||||
unexpected = [k for k in result.unexpected_keys if not k.startswith("hamer_backbone.")]
|
||||
if missing:
|
||||
print(f"[Denoiser] Missing keys: {missing}")
|
||||
if unexpected:
|
||||
print(f"[Denoiser] Unexpected keys: {unexpected}")
|
||||
epoch_str = epoch if epoch is not None else "?"
|
||||
print(f"[Denoiser] Loaded from {ckpt} (epoch={epoch_str})")
|
||||
del state_dict
|
||||
gc.collect()
|
||||
torch.cuda.empty_cache()
|
||||
return denoiser
|
||||
@@ -0,0 +1,59 @@
|
||||
"""Dual-hand confidence veto for target-box selection.
|
||||
|
||||
Background: the YOLO detector (WiLoR) occasionally assigns both left and right labels to the
|
||||
same hand (highly overlapping boxes). In that case we arbitrate by confidence: if the
|
||||
opposite-side box has higher confidence, the target box is unreliable and is vetoed (treated
|
||||
as no detection). Non-overlapping left/right boxes (genuinely different hands in the frame) do
|
||||
not affect the target box, so true two-hand frames are not wrongly suppressed.
|
||||
|
||||
Used by the box-selection logic in preprocessing/hamer_crop_and_skeleton.py and
|
||||
utils/online_hamer.py. Once a frame is vetoed it has no detection (conf=0); at inference time
|
||||
the cmask mechanism falls back to the no-information state automatically.
|
||||
"""
|
||||
import numpy as np
|
||||
|
||||
|
||||
def iou(a, b):
|
||||
"""a, b: xyxy of shape (4,). Returns the IoU."""
|
||||
ix1, iy1 = max(a[0], b[0]), max(a[1], b[1])
|
||||
ix2, iy2 = min(a[2], b[2]), min(a[3], b[3])
|
||||
iw, ih = max(0.0, ix2 - ix1), max(0.0, iy2 - iy1)
|
||||
inter = iw * ih
|
||||
ua = (a[2] - a[0]) * (a[3] - a[1]) + (b[2] - b[0]) * (b[3] - b[1]) - inter
|
||||
return inter / ua if ua > 0 else 0.0
|
||||
|
||||
|
||||
def select_target_box_dual_conf(clses, confs, bboxes, target_cls, iou_thr=0.7):
|
||||
"""Select the highest-confidence box of the target class; veto it if an opposite-side box wins.
|
||||
|
||||
Mechanism: among boxes of ``target_cls`` pick the highest-confidence "main" box, then find the
|
||||
opposite-side box (cls = 1 - target_cls) with the largest IoU to it. If that IoU exceeds
|
||||
``iou_thr`` and the opposite-side confidence is higher than the main box's, the main box is
|
||||
vetoed (the location actually belongs to the opposite-side hand) and None is returned.
|
||||
Otherwise the main box index is returned.
|
||||
|
||||
Args:
|
||||
clses: (N,) YOLO box classes (WiLoR: 0=left, 1=right).
|
||||
confs: (N,) YOLO box confidences.
|
||||
bboxes: (N, 4) xyxy.
|
||||
target_cls: target hand class (1=right, 0=left).
|
||||
iou_thr: IoU threshold above which two boxes count as overlapping (same hand). Default 0.7.
|
||||
Returns:
|
||||
The main box index (int), or None (no target box / vetoed).
|
||||
"""
|
||||
target_mask = clses == target_cls
|
||||
if not target_mask.any():
|
||||
return None
|
||||
sel = int(np.argmax(np.where(target_mask, confs, -1.0)))
|
||||
other_mask = clses == (1 - target_cls)
|
||||
if other_mask.any():
|
||||
main_box = bboxes[sel]
|
||||
best_iou, best_other_conf = 0.0, 0.0
|
||||
for oi in np.where(other_mask)[0]:
|
||||
v = iou(main_box, bboxes[oi])
|
||||
if v > best_iou:
|
||||
best_iou = v
|
||||
best_other_conf = float(confs[oi])
|
||||
if best_iou > iou_thr and best_other_conf > float(confs[sel]):
|
||||
return None # vetoed by the opposite side
|
||||
return sel
|
||||
@@ -0,0 +1,570 @@
|
||||
"""
|
||||
Hand motion evaluation utilities.
|
||||
|
||||
Provides:
|
||||
- A-MPJPE, RA-MPJPE, PA-MPJPE, W-MPJPE, WA-MPJPE (joints metrics)
|
||||
- MANO-PSNR (compression fidelity, dB)
|
||||
- Temporal Incremental RMSE (TI-RMSE) (motion smoothness, mm)
|
||||
"""
|
||||
|
||||
from typing import Dict, Optional, Tuple
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Alignment helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _compute_procrustes_transform(
|
||||
pred_c: np.ndarray,
|
||||
gt_c: np.ndarray
|
||||
) -> tuple[np.ndarray, float]:
|
||||
"""
|
||||
Compute the Procrustes transform (rotation + scale).
|
||||
|
||||
Args:
|
||||
pred_c: (N, 3) centered predicted points
|
||||
gt_c: (N, 3) centered GT points
|
||||
|
||||
Returns:
|
||||
R: (3, 3) rotation matrix
|
||||
scale: float scale factor
|
||||
"""
|
||||
H = pred_c.T @ gt_c
|
||||
U, _, Vt = np.linalg.svd(H)
|
||||
d = np.linalg.det(Vt.T @ U.T)
|
||||
R = Vt.T @ np.diag([1, 1, d]) @ U.T
|
||||
scale = np.trace((pred_c @ R.T).T @ gt_c) / (np.linalg.norm(pred_c) ** 2 + 1e-8)
|
||||
return R, scale
|
||||
|
||||
|
||||
def procrustes_align(pred: torch.Tensor, gt: torch.Tensor) -> torch.Tensor:
|
||||
"""Per-frame Procrustes alignment (rotation + scale + translation).
|
||||
Input: (N, 21, 3). Used for PA-MPJPE."""
|
||||
pred_np = pred.cpu().numpy()
|
||||
gt_np = gt.cpu().numpy()
|
||||
aligned = np.zeros_like(pred_np)
|
||||
for i in range(pred_np.shape[0]):
|
||||
p, g = pred_np[i], gt_np[i]
|
||||
p_mean = p.mean(axis=0)
|
||||
g_mean = g.mean(axis=0)
|
||||
R, scale = _compute_procrustes_transform(p - p_mean, g - g_mean)
|
||||
aligned[i] = scale * ((p - p_mean) @ R.T) + g_mean
|
||||
return torch.from_numpy(aligned).to(pred.device)
|
||||
|
||||
|
||||
def first_n_frames_procrustes(
|
||||
joints_pred: np.ndarray,
|
||||
joints_gt: np.ndarray,
|
||||
valid_frames: np.ndarray,
|
||||
n_anchor: int = 2,
|
||||
) -> np.ndarray:
|
||||
"""
|
||||
W-MPJPE alignment: Procrustes (scale + rotation + translation) estimated
|
||||
from the first n_anchor valid frames, then applied globally to the entire
|
||||
sequence.
|
||||
|
||||
The anchor frames are stacked into a single point cloud (n*21 points) and
|
||||
one optimal (s, R, t) is solved. The same transform is applied to every
|
||||
frame — inter-frame relative motion is fully preserved.
|
||||
|
||||
Args:
|
||||
joints_pred : (T, 21, 3) metres
|
||||
joints_gt : (T, 21, 3) metres
|
||||
valid_frames : 1-D array of valid frame indices (sorted)
|
||||
n_anchor : number of leading valid frames to use for alignment
|
||||
|
||||
Returns:
|
||||
aligned_pred : (T, 21, 3) metres
|
||||
"""
|
||||
anchor_frames = valid_frames[:n_anchor]
|
||||
|
||||
pred_anchor = joints_pred[anchor_frames].reshape(-1, 3)
|
||||
gt_anchor = joints_gt[anchor_frames].reshape(-1, 3)
|
||||
|
||||
pred_mean = pred_anchor.mean(axis=0)
|
||||
gt_mean = gt_anchor.mean(axis=0)
|
||||
pred_c = pred_anchor - pred_mean
|
||||
gt_c = gt_anchor - gt_mean
|
||||
|
||||
R, scale = _compute_procrustes_transform(pred_c, gt_c)
|
||||
|
||||
T_seq, J, _ = joints_pred.shape
|
||||
pred_flat = joints_pred.reshape(-1, 3)
|
||||
aligned_flat = scale * (R @ (pred_flat - pred_mean).T).T + gt_mean
|
||||
return aligned_flat.reshape(T_seq, J, 3)
|
||||
|
||||
|
||||
# Keep old name as alias for backward compatibility
|
||||
def first_two_frames_procrustes(joints_pred, joints_gt, valid_frames):
|
||||
return first_n_frames_procrustes(joints_pred, joints_gt, valid_frames, n_anchor=2)
|
||||
|
||||
|
||||
def whole_seq_procrustes(
|
||||
joints_pred: np.ndarray,
|
||||
joints_gt: np.ndarray,
|
||||
valid_mask: np.ndarray,
|
||||
) -> np.ndarray:
|
||||
"""
|
||||
Global Procrustes alignment (rotation + scale + translation) estimated from
|
||||
all valid frames treated as a single point cloud.
|
||||
|
||||
Inter-frame relative motion is 100% preserved because every frame is
|
||||
transformed by the same (R, scale, translation).
|
||||
|
||||
Args:
|
||||
joints_pred : (T, 21, 3) metres
|
||||
joints_gt : (T, 21, 3) metres
|
||||
valid_mask : (T,) bool
|
||||
|
||||
Returns:
|
||||
aligned_pred : (T, 21, 3) metres
|
||||
"""
|
||||
pred_valid = joints_pred[valid_mask].reshape(-1, 3) # (V*21, 3)
|
||||
gt_valid = joints_gt[valid_mask].reshape(-1, 3)
|
||||
|
||||
pred_mean = pred_valid.mean(axis=0)
|
||||
gt_mean = gt_valid.mean(axis=0)
|
||||
pred_c = pred_valid - pred_mean
|
||||
gt_c = gt_valid - gt_mean
|
||||
|
||||
R, scale = _compute_procrustes_transform(pred_c, gt_c)
|
||||
|
||||
T_seq, J, _ = joints_pred.shape
|
||||
pred_flat = joints_pred.reshape(-1, 3)
|
||||
aligned_flat = scale * (R @ (pred_flat - pred_mean).T).T + gt_mean
|
||||
return aligned_flat.reshape(T_seq, J, 3)
|
||||
|
||||
|
||||
def batched_procrustes_align(pred: torch.Tensor, gt: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Batched per-frame Procrustes alignment (rotation + scale + translation).
|
||||
GPU-accelerated replacement for the per-frame numpy loop in procrustes_align.
|
||||
Input: (N, K, 3). Returns: (N, K, 3).
|
||||
"""
|
||||
pred_mean = pred.mean(dim=1, keepdim=True)
|
||||
gt_mean = gt.mean(dim=1, keepdim=True)
|
||||
pred_c = pred - pred_mean
|
||||
gt_c = gt - gt_mean
|
||||
H = torch.bmm(pred_c.transpose(1, 2), gt_c) # (N, 3, 3)
|
||||
U, S, Vt = torch.linalg.svd(H)
|
||||
d = torch.det(torch.bmm(Vt.transpose(1, 2), U.transpose(1, 2)))
|
||||
D = torch.ones_like(S)
|
||||
D[:, -1] = d.sign()
|
||||
R = torch.bmm(Vt.transpose(1, 2), torch.bmm(torch.diag_embed(D), U.transpose(1, 2)))
|
||||
pred_R = torch.bmm(pred_c, R.transpose(1, 2)) # (N, K, 3)
|
||||
scale = (pred_R * gt_c).sum(dim=(1, 2)) / ((pred_c ** 2).sum(dim=(1, 2)) + 1e-8)
|
||||
return scale.unsqueeze(-1).unsqueeze(-1) * pred_R + gt_mean
|
||||
|
||||
|
||||
def _single_procrustes_torch(
|
||||
pred: torch.Tensor, gt: torch.Tensor
|
||||
) -> tuple:
|
||||
"""Procrustes (R + scale + translation) for a single point cloud pair."""
|
||||
pred_mean = pred.mean(dim=0)
|
||||
gt_mean = gt.mean(dim=0)
|
||||
pred_c = pred - pred_mean
|
||||
gt_c = gt - gt_mean
|
||||
H = pred_c.T @ gt_c
|
||||
U, S, Vt = torch.linalg.svd(H)
|
||||
d = torch.det(Vt.T @ U.T)
|
||||
D = torch.diag(torch.tensor([1.0, 1.0, d.sign().item()], device=pred.device, dtype=pred.dtype))
|
||||
R = Vt.T @ D @ U.T
|
||||
scale = ((pred_c @ R.T) * gt_c).sum() / ((pred_c ** 2).sum() + 1e-8)
|
||||
return R, scale, pred_mean, gt_mean
|
||||
|
||||
|
||||
def first_n_frames_procrustes_torch(
|
||||
joints_pred: torch.Tensor,
|
||||
joints_gt: torch.Tensor,
|
||||
n_anchor: int = 2,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
W-MPJPE alignment, torch version. Single sequence, valid frames only.
|
||||
Input: (V, 21, 3). Returns: (V, 21, 3).
|
||||
"""
|
||||
anchor_pred = joints_pred[:n_anchor].reshape(-1, 3)
|
||||
anchor_gt = joints_gt[:n_anchor].reshape(-1, 3)
|
||||
R, scale, pred_mean, gt_mean = _single_procrustes_torch(anchor_pred, anchor_gt)
|
||||
flat = joints_pred.reshape(-1, 3)
|
||||
aligned = scale * (flat - pred_mean) @ R.T + gt_mean
|
||||
return aligned.reshape_as(joints_pred)
|
||||
|
||||
|
||||
def whole_seq_procrustes_torch(
|
||||
joints_pred: torch.Tensor,
|
||||
joints_gt: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
WA-MPJPE alignment, torch version. Single sequence, valid frames only.
|
||||
Input: (V, 21, 3). Returns: (V, 21, 3).
|
||||
"""
|
||||
R, scale, pred_mean, gt_mean = _single_procrustes_torch(
|
||||
joints_pred.reshape(-1, 3), joints_gt.reshape(-1, 3)
|
||||
)
|
||||
flat = joints_pred.reshape(-1, 3)
|
||||
aligned = scale * (flat - pred_mean) @ R.T + gt_mean
|
||||
return aligned.reshape_as(joints_pred)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Joints metrics
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def compute_mpjpe_metrics(
|
||||
joints_gt_m: np.ndarray,
|
||||
joints_pred_m: np.ndarray,
|
||||
valid_mask: Optional[np.ndarray] = None,
|
||||
) -> Dict[str, float]:
|
||||
"""
|
||||
Compute five MPJPE variants.
|
||||
|
||||
Alignment definitions:
|
||||
A-MPJPE : absolute, no alignment at all — raw joint coordinate difference
|
||||
RA-MPJPE : per-frame root-relative (subtract wrist joint)
|
||||
PA-MPJPE : per-frame Procrustes (rotation + scale + translation)
|
||||
W-MPJPE : Procrustes (scale + rotation + translation) estimated from the
|
||||
first 2 valid frames (stacked as one 42-point cloud), then the
|
||||
same (s, R, t) applied to the entire sequence
|
||||
WA-MPJPE : whole-sequence global Procrustes (rotation + scale + translation)
|
||||
→ all valid frames treated as one point cloud, one alignment for all
|
||||
→ inter-frame relative motion 100% preserved
|
||||
|
||||
Args:
|
||||
joints_gt_m : (T, 21, 3) GT joints in metres — FULL sequence
|
||||
joints_pred_m : (T, 21, 3) predicted joints in metres — FULL sequence
|
||||
valid_mask : (T,) bool. If None, all frames used.
|
||||
|
||||
Returns dict with scalar means (mm) and per-valid-frame arrays.
|
||||
"""
|
||||
T = len(joints_gt_m)
|
||||
vm = valid_mask if valid_mask is not None else np.ones(T, dtype=bool)
|
||||
|
||||
gt_v = joints_gt_m[vm] # (V, 21, 3)
|
||||
pred_v = joints_pred_m[vm]
|
||||
|
||||
# ── A-MPJPE: absolute, no alignment ──────────────────────────────────
|
||||
a_err = np.linalg.norm((pred_v - gt_v) * 1000.0, axis=-1)
|
||||
a_mpjpe_pf = a_err.mean(axis=-1)
|
||||
|
||||
# ── RA-MPJPE: per-frame root-relative ────────────────────────────────
|
||||
gt_rel = gt_v - gt_v[:, 0:1, :]
|
||||
pred_rel = pred_v - pred_v[:, 0:1, :]
|
||||
ra_err = np.linalg.norm((pred_rel - gt_rel) * 1000.0, axis=-1)
|
||||
ra_mpjpe_pf = ra_err.mean(axis=-1)
|
||||
|
||||
# ── PA-MPJPE: per-frame Procrustes (rotation + scale + translation) ──
|
||||
pred_pa = procrustes_align(
|
||||
torch.from_numpy(pred_v), torch.from_numpy(gt_v)
|
||||
).numpy()
|
||||
pa_err = np.linalg.norm((pred_pa - gt_v) * 1000.0, axis=-1)
|
||||
pa_mpjpe_pf = pa_err.mean(axis=-1)
|
||||
|
||||
# ── W-MPJPE: first-2-frames Procrustes (s + R + t), applied globally ─
|
||||
valid_frames = np.where(vm)[0]
|
||||
pred_w = first_two_frames_procrustes(joints_pred_m, joints_gt_m, valid_frames)
|
||||
w_err = np.linalg.norm((pred_w[vm] - gt_v) * 1000.0, axis=-1)
|
||||
w_mpjpe_pf = w_err.mean(axis=-1)
|
||||
|
||||
# ── WA-MPJPE: whole-sequence global Procrustes (R + scale + t) ───────
|
||||
pred_wa = whole_seq_procrustes(joints_pred_m, joints_gt_m, vm)
|
||||
wa_err = np.linalg.norm((pred_wa[vm] - gt_v) * 1000.0, axis=-1)
|
||||
wa_mpjpe_pf = wa_err.mean(axis=-1)
|
||||
|
||||
return {
|
||||
"a_mpjpe": float(a_mpjpe_pf.mean()),
|
||||
"ra_mpjpe": float(ra_mpjpe_pf.mean()),
|
||||
"pa_mpjpe": float(pa_mpjpe_pf.mean()),
|
||||
"w_mpjpe": float(w_mpjpe_pf.mean()),
|
||||
"wa_mpjpe": float(wa_mpjpe_pf.mean()),
|
||||
"a_mpjpe_per_frame": a_mpjpe_pf,
|
||||
"ra_mpjpe_per_frame": ra_mpjpe_pf,
|
||||
"pa_mpjpe_per_frame": pa_mpjpe_pf,
|
||||
"w_mpjpe_per_frame": w_mpjpe_pf,
|
||||
"wa_mpjpe_per_frame": wa_mpjpe_pf,
|
||||
}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# MANO-PSNR and Temporal Incremental RMSE
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def compute_mano_psnr(
|
||||
joints_gt_m: np.ndarray,
|
||||
joints_pred_m: np.ndarray,
|
||||
valid_mask: Optional[np.ndarray] = None,
|
||||
max_val_mm: float = 1000.0,
|
||||
) -> float:
|
||||
"""
|
||||
MANO-PSNR over valid frames (no alignment).
|
||||
|
||||
MSE is computed over all valid frames × 21 joints × 3 coords in mm.
|
||||
PSNR = 10 * log10(MAX² / MSE), MAX = 1000 mm (joint coords in [-500, 500] mm).
|
||||
|
||||
Args:
|
||||
joints_gt_m: (T, 21, 3) GT joints in metres
|
||||
joints_pred_m: (T, 21, 3) predicted joints in metres
|
||||
valid_mask: (T,) bool. If None all frames are used.
|
||||
max_val_mm: signal range upper bound (mm)
|
||||
|
||||
Returns:
|
||||
PSNR in dB. Higher is better. Returns +inf if MSE == 0, nan if no valid frames.
|
||||
"""
|
||||
if valid_mask is not None:
|
||||
if valid_mask.sum() == 0:
|
||||
return float("nan")
|
||||
joints_gt_m = joints_gt_m[valid_mask]
|
||||
joints_pred_m = joints_pred_m[valid_mask]
|
||||
|
||||
gt_mm = joints_gt_m * 1000.0
|
||||
pred_mm = joints_pred_m * 1000.0
|
||||
|
||||
mse = float(np.mean((pred_mm - gt_mm) ** 2))
|
||||
if mse == 0.0:
|
||||
return float("inf")
|
||||
return float(10.0 * np.log10(max_val_mm ** 2 / mse))
|
||||
|
||||
|
||||
def compute_ti_rmse(
|
||||
joints_gt_m: np.ndarray,
|
||||
joints_pred_m: np.ndarray,
|
||||
valid_mask: Optional[np.ndarray] = None,
|
||||
) -> float:
|
||||
"""
|
||||
Temporal Incremental RMSE (TI-RMSE) — motion smoothness metric (no alignment).
|
||||
|
||||
Computes RMSE between GT and predicted frame-to-frame motion increments.
|
||||
Only consecutive pairs of valid frames are included.
|
||||
|
||||
Args:
|
||||
joints_gt_m: (T, 21, 3) GT joints in metres
|
||||
joints_pred_m: (T, 21, 3) predicted joints in metres
|
||||
valid_mask: (T,) bool. If None all frames are used.
|
||||
|
||||
Returns:
|
||||
TI-RMSE in mm. Lower is better. Returns nan if no valid pairs.
|
||||
"""
|
||||
T = len(joints_gt_m)
|
||||
gt_mm = joints_gt_m * 1000.0
|
||||
pred_mm = joints_pred_m * 1000.0
|
||||
|
||||
gt_deltas, pred_deltas = [], []
|
||||
for t in range(T - 1):
|
||||
if valid_mask is None or (valid_mask[t] and valid_mask[t + 1]):
|
||||
gt_deltas.append(gt_mm[t + 1] - gt_mm[t])
|
||||
pred_deltas.append(pred_mm[t + 1] - pred_mm[t])
|
||||
|
||||
if len(gt_deltas) == 0:
|
||||
return float("nan")
|
||||
|
||||
gt_deltas = np.stack(gt_deltas)
|
||||
pred_deltas = np.stack(pred_deltas)
|
||||
return float(np.sqrt(np.mean((pred_deltas - gt_deltas) ** 2)))
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Combined evaluation
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def evaluate_sequence(
|
||||
x_pred_norm: np.ndarray, # [T, 61] normalized predictions
|
||||
x_gt_norm: np.ndarray, # [T, 61] normalized GT
|
||||
x_pred_raw: np.ndarray, # [T, 61] unnormalized predictions
|
||||
x_gt_raw: np.ndarray, # [T, 61] unnormalized GT
|
||||
joints_pred: np.ndarray, # [T, 21, 3] meters
|
||||
joints_gt: np.ndarray, # [T, 21, 3] meters
|
||||
valid_mask: np.ndarray, # [T] bool
|
||||
) -> Dict[str, float]:
|
||||
"""Full evaluation of one sequence."""
|
||||
metrics = {}
|
||||
|
||||
# Joints metrics
|
||||
jm = compute_mpjpe_metrics(joints_gt, joints_pred, valid_mask)
|
||||
metrics.update({k: v for k, v in jm.items() if not k.endswith("_per_frame")})
|
||||
|
||||
# MANO-PSNR and TI-RMSE
|
||||
metrics["mano_psnr_db"] = compute_mano_psnr(joints_gt, joints_pred, valid_mask)
|
||||
metrics["ti_rmse_mm"] = compute_ti_rmse(joints_gt, joints_pred, valid_mask)
|
||||
|
||||
return metrics
|
||||
|
||||
|
||||
def aggregate_metrics(list_of_dicts: list) -> Dict[str, float]:
|
||||
"""Average scalar metrics across a list of per-sequence metric dicts."""
|
||||
if not list_of_dicts:
|
||||
return {}
|
||||
keys = [k for k, v in list_of_dicts[0].items() if np.isscalar(v)]
|
||||
result = {}
|
||||
for k in keys:
|
||||
vals = [d[k] for d in list_of_dicts if k in d]
|
||||
result[k] = float(np.mean(vals))
|
||||
return result
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Accel / RTE / AUC
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def compute_accel_error(
|
||||
joints_gt_m: np.ndarray,
|
||||
joints_pred_m: np.ndarray,
|
||||
valid_mask: Optional[np.ndarray] = None,
|
||||
fps: float = 30.0,
|
||||
) -> Tuple[np.ndarray, float]:
|
||||
"""
|
||||
Acceleration error (Accel Error).
|
||||
|
||||
Definition: ||(a_pred - a_gt)||_2, averaged over all joints, in m/s².
|
||||
Acceleration = second-order difference * fps².
|
||||
Only frames where all three (t-1, t, t+1) are valid are considered.
|
||||
|
||||
Returns:
|
||||
per_frame: (T-2,) per-frame error (valid triplets only), m/s²
|
||||
mean: scalar mean
|
||||
"""
|
||||
T = len(joints_gt_m)
|
||||
vm = valid_mask if valid_mask is not None else np.ones(T, dtype=bool)
|
||||
|
||||
errors = []
|
||||
for t in range(1, T - 1):
|
||||
if vm[t - 1] and vm[t] and vm[t + 1]:
|
||||
accel_gt = joints_gt_m[t - 1] - 2 * joints_gt_m[t] + joints_gt_m[t + 1]
|
||||
accel_pred = joints_pred_m[t - 1] - 2 * joints_pred_m[t] + joints_pred_m[t + 1]
|
||||
diff = np.linalg.norm(accel_pred - accel_gt, axis=-1).mean()
|
||||
errors.append(diff * fps ** 2)
|
||||
|
||||
if len(errors) == 0:
|
||||
return np.array([]), float("nan")
|
||||
arr = np.array(errors)
|
||||
return arr, float(arr.mean())
|
||||
|
||||
|
||||
def compute_rte(
|
||||
joints_gt_m: np.ndarray,
|
||||
joints_pred_m: np.ndarray,
|
||||
valid_mask: Optional[np.ndarray] = None,
|
||||
) -> Tuple[np.ndarray, float]:
|
||||
"""
|
||||
Root Translation Error (RTE), following the HaWoR definition.
|
||||
|
||||
First aligns the predicted root translation with a fixed-scale global
|
||||
rotation + translation (Umeyama, fixed_scale=True), then computes the
|
||||
per-frame root position error divided by the total GT trajectory length,
|
||||
reported as a percentage (x100).
|
||||
|
||||
joints_gt_m / joints_pred_m: (T, 21, 3) metres, root = joint[0]
|
||||
|
||||
Returns:
|
||||
per_frame: (V,) RTE per valid frame (%)
|
||||
mean: scalar mean
|
||||
"""
|
||||
T = len(joints_gt_m)
|
||||
vm = valid_mask if valid_mask is not None else np.ones(T, dtype=bool)
|
||||
|
||||
gt_root = joints_gt_m[vm, 0, :] # (V, 3)
|
||||
pred_root = joints_pred_m[vm, 0, :] # (V, 3)
|
||||
|
||||
if len(gt_root) < 2:
|
||||
return np.array([]), float("nan")
|
||||
|
||||
gt_t = torch.from_numpy(gt_root).float()
|
||||
pred_t = torch.from_numpy(pred_root).float()
|
||||
|
||||
# Umeyama fixed-scale alignment
|
||||
gt_mean = gt_t.mean(0)
|
||||
pred_mean = pred_t.mean(0)
|
||||
gt_c = gt_t - gt_mean
|
||||
pred_c = pred_t - pred_mean
|
||||
|
||||
H = pred_c.T @ gt_c
|
||||
U, _, Vt = torch.linalg.svd(H)
|
||||
d = torch.det(Vt.T @ U.T)
|
||||
D = torch.diag(torch.tensor([1.0, 1.0, d.sign().item()]))
|
||||
R = Vt.T @ D @ U.T
|
||||
t = gt_mean - (R @ pred_mean.unsqueeze(-1)).squeeze(-1)
|
||||
pred_aligned = (R @ pred_t.T).T + t
|
||||
|
||||
# Total length of the GT trajectory
|
||||
disp = float((gt_t[1:] - gt_t[:-1]).norm(dim=-1).sum())
|
||||
if disp < 1e-8:
|
||||
return np.array([]), float("nan")
|
||||
|
||||
per_frame = (gt_t - pred_aligned).norm(dim=-1).numpy() / disp * 100.0
|
||||
return per_frame, float(per_frame.mean())
|
||||
|
||||
|
||||
def compute_pck_auc(
|
||||
joints_gt_m: np.ndarray,
|
||||
joints_pred_m: np.ndarray,
|
||||
valid_mask: Optional[np.ndarray] = None,
|
||||
max_thresh_mm: float = 50.0,
|
||||
n_steps: int = 50,
|
||||
) -> float:
|
||||
"""
|
||||
PCK-AUC (based on PA-MPJPE): area under the PCK curve over [0, max_thresh_mm],
|
||||
normalized to [0, 1].
|
||||
|
||||
PCK@t = fraction of frames whose per-frame PA-MPJPE (Procrustes-aligned)
|
||||
mean is < t.
|
||||
AUC = trapz(PCK) / max_thresh_mm, range [0, 1], higher is better.
|
||||
"""
|
||||
T = len(joints_gt_m)
|
||||
vm = valid_mask if valid_mask is not None else np.ones(T, dtype=bool)
|
||||
|
||||
gt_v = joints_gt_m[vm]
|
||||
pred_v = joints_pred_m[vm]
|
||||
|
||||
if len(gt_v) == 0:
|
||||
return float("nan")
|
||||
|
||||
# per-frame Procrustes alignment (rotation + scale + translation)
|
||||
pred_pa = batched_procrustes_align(
|
||||
torch.from_numpy(pred_v), torch.from_numpy(gt_v)
|
||||
).numpy()
|
||||
|
||||
# per-frame mean joint error (mm)
|
||||
per_frame_err = np.linalg.norm((pred_pa - gt_v) * 1000.0, axis=-1).mean(axis=-1) # (V,)
|
||||
|
||||
thresholds = np.linspace(0, max_thresh_mm, n_steps + 1)
|
||||
pck = np.array([(per_frame_err < thr).mean() for thr in thresholds])
|
||||
auc = float(np.trapz(pck, thresholds) / max_thresh_mm)
|
||||
return auc
|
||||
|
||||
|
||||
def compute_ra_pck_auc(
|
||||
joints_gt_m: np.ndarray,
|
||||
joints_pred_m: np.ndarray,
|
||||
valid_mask: Optional[np.ndarray] = None,
|
||||
max_thresh_mm: float = 50.0,
|
||||
n_steps: int = 50,
|
||||
) -> float:
|
||||
"""
|
||||
PCK-AUC (based on RA-MPJPE): area under the PCK curve over [0, max_thresh_mm],
|
||||
normalized to [0, 1].
|
||||
|
||||
PCK@t = fraction of frames whose per-frame RA-MPJPE (root-relative)
|
||||
mean is < t.
|
||||
AUC = trapz(PCK) / max_thresh_mm, range [0, 1], higher is better.
|
||||
"""
|
||||
T = len(joints_gt_m)
|
||||
vm = valid_mask if valid_mask is not None else np.ones(T, dtype=bool)
|
||||
|
||||
gt_v = joints_gt_m[vm]
|
||||
pred_v = joints_pred_m[vm]
|
||||
|
||||
if len(gt_v) == 0:
|
||||
return float("nan")
|
||||
|
||||
# per-frame root-relative (wrist joint = index 0)
|
||||
gt_rel = gt_v - gt_v[:, 0:1, :]
|
||||
pred_rel = pred_v - pred_v[:, 0:1, :]
|
||||
|
||||
# per-frame mean joint error (mm)
|
||||
per_frame_err = np.linalg.norm((pred_rel - gt_rel) * 1000.0, axis=-1).mean(axis=-1) # (V,)
|
||||
|
||||
thresholds = np.linspace(0, max_thresh_mm, n_steps + 1)
|
||||
pck = np.array([(per_frame_err < thr).mean() for thr in thresholds])
|
||||
auc = float(np.trapz(pck, thresholds) / max_thresh_mm)
|
||||
return auc
|
||||
@@ -0,0 +1,495 @@
|
||||
"""
|
||||
Inference and MPJPE evaluation utilities.
|
||||
|
||||
Provides:
|
||||
- Overlapping-window inference (stride stitching + per-field overlap averaging)
|
||||
- Sequence-level MPJPE evaluation
|
||||
"""
|
||||
|
||||
from typing import Dict, List, Optional, Tuple
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from tqdm import tqdm
|
||||
|
||||
from utils.eval_utils import compute_mpjpe_metrics
|
||||
from utils.mano_utils import MANOForwardKinematics, compute_joints_global
|
||||
|
||||
|
||||
def split_sequence_into_windows(
|
||||
seq_batch: dict,
|
||||
window_size: int,
|
||||
stride: int,
|
||||
device: torch.device,
|
||||
) -> Tuple[dict, list[int]]:
|
||||
"""
|
||||
Split a variable-length sequence batch into overlapping windows,
|
||||
returning a window batch and a list of sequence memberships.
|
||||
|
||||
Args:
|
||||
seq_batch: collated sequence batch; temporal dim shape (B, L, ...)
|
||||
window_size: window size
|
||||
stride: window stride = window_size - overlap_size
|
||||
device: target device
|
||||
|
||||
Returns:
|
||||
win_batch: dict of (N_wins_total, ...) tensors
|
||||
seq_ids: list of length N_wins_total
|
||||
"""
|
||||
import math
|
||||
|
||||
TIME_KEYS = {
|
||||
"mano_params", "mano_trans", "padding_mask",
|
||||
"hamer_landmarks", "crop_intrinsics", "hamer_confidence",
|
||||
"images", "joint_3d", "image_tokens",
|
||||
}
|
||||
B = seq_batch["mano_params"].shape[0]
|
||||
|
||||
all_windows: list[dict] = []
|
||||
seq_ids: list[int] = []
|
||||
|
||||
def pad_sequence_length(nf: int, window_size: int, stride: int) -> int:
|
||||
if nf <= window_size:
|
||||
return window_size
|
||||
k = math.ceil((nf - window_size) / stride)
|
||||
return window_size + stride * k
|
||||
|
||||
for b in range(B):
|
||||
pm = seq_batch.get("padding_mask")
|
||||
if pm is not None:
|
||||
nf = int((~pm[b]).sum().item())
|
||||
nf = max(nf, 1)
|
||||
else:
|
||||
nf = seq_batch["mano_params"].shape[1]
|
||||
|
||||
valid_len = pad_sequence_length(nf, window_size, stride)
|
||||
num_wins = (valid_len - window_size) // stride + 1
|
||||
|
||||
for w in range(num_wins):
|
||||
start = w * stride
|
||||
end = start + window_size
|
||||
win = {}
|
||||
for key, val in seq_batch.items():
|
||||
if not isinstance(val, torch.Tensor):
|
||||
if isinstance(val, list):
|
||||
win[key] = val[b]
|
||||
else:
|
||||
win[key] = val
|
||||
continue
|
||||
if key in TIME_KEYS and val.ndim >= 2:
|
||||
seq_len = val.shape[1] # batched: dim 1 is time
|
||||
if val.ndim >= 3:
|
||||
if end <= seq_len:
|
||||
win[key] = val[b, start:end]
|
||||
else:
|
||||
chunk = val[b, start:min(end, seq_len)]
|
||||
pad_len = end - min(end, seq_len)
|
||||
pad_shape = (pad_len,) + chunk.shape[1:]
|
||||
pad_val = True if key == "padding_mask" else 0.0
|
||||
pad_t = torch.full(pad_shape, pad_val, dtype=chunk.dtype, device=chunk.device)
|
||||
win[key] = torch.cat([chunk, pad_t], dim=0)
|
||||
else:
|
||||
# ndim == 2: (B, L)
|
||||
if end <= seq_len:
|
||||
win[key] = val[b, start:end]
|
||||
else:
|
||||
chunk = val[b, start:min(end, seq_len)]
|
||||
pad_len = end - min(end, seq_len)
|
||||
pad_val = True if key == "padding_mask" else 0.0
|
||||
pad_t = torch.full((pad_len,), pad_val, dtype=chunk.dtype, device=chunk.device)
|
||||
win[key] = torch.cat([chunk, pad_t], dim=0)
|
||||
else:
|
||||
win[key] = val[b] if val.ndim >= 1 else val
|
||||
|
||||
# Skip windows that are all padding
|
||||
pm_win = win.get("padding_mask")
|
||||
if pm_win is not None and pm_win.all():
|
||||
continue
|
||||
|
||||
all_windows.append(win)
|
||||
seq_ids.append(b)
|
||||
|
||||
def collate_fn(batch):
|
||||
out = {}
|
||||
for key in batch[0]:
|
||||
vals = [b[key] for b in batch]
|
||||
if isinstance(vals[0], torch.Tensor):
|
||||
out[key] = torch.stack(vals)
|
||||
else:
|
||||
out[key] = vals
|
||||
return out
|
||||
|
||||
if not all_windows:
|
||||
return collate_fn([{
|
||||
"mano_params": torch.zeros(1, window_size, 48),
|
||||
"mano_trans": torch.zeros(1, window_size, 3),
|
||||
"mano_betas": torch.zeros(1, 10),
|
||||
"padding_mask": torch.ones(1, window_size, dtype=torch.bool),
|
||||
"images": torch.zeros(1, window_size, 3, 256, 256, dtype=torch.uint8),
|
||||
"hamer_landmarks": torch.zeros(1, window_size, 21, 2),
|
||||
"crop_intrinsics": torch.zeros(1, window_size, 4),
|
||||
"hamer_confidence": torch.zeros(1, window_size),
|
||||
"side": "right",
|
||||
"source": "dexycb",
|
||||
}]), [0]
|
||||
|
||||
win_batch = collate_fn(all_windows)
|
||||
return win_batch, seq_ids
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def run_fm_inference_with_overlap(
|
||||
model_engine,
|
||||
denoiser,
|
||||
batch: Dict[str, torch.Tensor],
|
||||
window_size: int,
|
||||
overlap_size: int,
|
||||
ode_steps: int,
|
||||
device: torch.device,
|
||||
show_progress: bool = False,
|
||||
guidance=None,
|
||||
overlap_method: str = "vblend",
|
||||
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
Flow Matching inference with overlapping windows.
|
||||
|
||||
overlap_method:
|
||||
"avg" : run each window's ODE independently, then uniformly average
|
||||
in x₁ space at the end
|
||||
"vblend" : blend velocity by center distance during ODE integration,
|
||||
keeping a single trajectory
|
||||
|
||||
Noise strategy (MagicAnimate style):
|
||||
- Initial noise is fixed per frame index: the pose/trans noise of frame f
|
||||
is globally unique, and all windows covering frame f share the same
|
||||
initial noise component.
|
||||
- beta is a single token for the whole sequence; all windows share the
|
||||
same beta noise.
|
||||
|
||||
Args:
|
||||
model_engine: DeepSpeed engine or plain model
|
||||
denoiser: HandPoseDenoiser
|
||||
batch: window batch (N_wins, T, ...)
|
||||
window_size: window size
|
||||
overlap_size: number of overlapping frames
|
||||
ode_steps: number of ODE integration steps
|
||||
device: torch device
|
||||
show_progress: whether to show a progress bar
|
||||
guidance: 2D reprojection guidance
|
||||
overlap_method: "vblend" or "avg"
|
||||
|
||||
Returns:
|
||||
pose_seq: (nf, 48) predicted pose (global_aa + joints_aa)
|
||||
trans_seq: (nf, 3) predicted trans
|
||||
betas_pred: (nf, 10) predicted beta (same value for the whole sequence)
|
||||
"""
|
||||
stride = window_size - overlap_size
|
||||
N_wins = batch["mano_params"].shape[0]
|
||||
T = window_size
|
||||
|
||||
# Sequence length
|
||||
pm = batch.get("padding_mask")
|
||||
if pm is not None:
|
||||
last_win_valid = int((~pm[-1]).sum().item())
|
||||
nf = (N_wins - 1) * stride + last_win_valid
|
||||
else:
|
||||
nf = (N_wins - 1) * stride + T
|
||||
|
||||
ms = denoiser.mano_slice
|
||||
flat_dim = ms.flat_dim
|
||||
beta_dim = ms.beta_dim # 10
|
||||
pose_dim = ms.pose_dim # 48
|
||||
trans_dim = ms.trans_dim # 3
|
||||
|
||||
dtype = next(model_engine.parameters()).dtype
|
||||
|
||||
# ── Build global initial noise per frame index ────────────────────────
|
||||
# beta: a single noise vector shared across the whole sequence
|
||||
# pose/trans: independent per frame; window w takes [w*stride, w*stride+T)
|
||||
n_noise_frames = (N_wins - 1) * stride + T # covers all window positions
|
||||
global_noise_beta = torch.randn(beta_dim, device=device, dtype=dtype)
|
||||
global_noise_pose = torch.randn(n_noise_frames, pose_dim, device=device, dtype=dtype)
|
||||
global_noise_trans = torch.randn(n_noise_frames, trans_dim, device=device, dtype=dtype)
|
||||
|
||||
# ── Select overlap inference path ─────────────────────────────────────
|
||||
if overlap_size > 0 and overlap_method == "vblend":
|
||||
# vblend: per-frame state; blend velocity by center distance during ODE
|
||||
x_pose = global_noise_pose[:nf].clone()
|
||||
x_trans = global_noise_trans[:nf].clone()
|
||||
x_beta = global_noise_beta.clone()
|
||||
|
||||
center = (T - 1) / 2.0
|
||||
w_table = torch.zeros(N_wins, T, device=device, dtype=dtype)
|
||||
for w in range(N_wins):
|
||||
for t in range(T):
|
||||
w_table[w, t] = max(0.01, 1.0 - abs(t - center) / center)
|
||||
|
||||
dt = 1.0 / ode_steps
|
||||
iterator = tqdm(range(ode_steps), desc="ODE", leave=False) if show_progress else range(ode_steps)
|
||||
|
||||
for i in iterator:
|
||||
t_step = torch.full((N_wins,), i * dt, device=device, dtype=dtype)
|
||||
|
||||
# Assemble window batch from per-frame state
|
||||
x_win = torch.empty(N_wins, flat_dim, device=device, dtype=dtype)
|
||||
for w in range(N_wins):
|
||||
s = w * stride
|
||||
e = min(s + T, nf)
|
||||
n = e - s
|
||||
x_win[w, ms.beta] = x_beta
|
||||
x_win[w, ms.poses] = torch.cat([
|
||||
x_pose[s:e].reshape(-1),
|
||||
global_noise_pose[s + n:s + T].reshape(-1),
|
||||
])[:T * pose_dim]
|
||||
x_win[w, ms.trans] = torch.cat([
|
||||
x_trans[s:e].reshape(-1),
|
||||
global_noise_trans[s + n:s + T].reshape(-1),
|
||||
])[:T * trans_dim]
|
||||
|
||||
v_win = model_engine(x_win, t_step, batch)
|
||||
if guidance is not None and guidance.active:
|
||||
with torch.enable_grad():
|
||||
v_win = guidance.compute_velocity_correction(v_win, x_win, t_step, batch)
|
||||
|
||||
v_pose_w = v_win[:, ms.poses].reshape(N_wins, T, pose_dim)
|
||||
v_trans_w = v_win[:, ms.trans].reshape(N_wins, T, trans_dim)
|
||||
v_beta_w = v_win[:, ms.beta]
|
||||
|
||||
# Weighted blending of per-frame velocity
|
||||
vp_acc = torch.zeros_like(x_pose)
|
||||
vt_acc = torch.zeros_like(x_trans)
|
||||
wt_acc = torch.zeros(nf, device=device, dtype=dtype)
|
||||
for w in range(N_wins):
|
||||
s = w * stride
|
||||
e = min(s + T, nf)
|
||||
for t in range(e - s):
|
||||
gf = s + t
|
||||
wgt = w_table[w, t]
|
||||
vp_acc[gf] += v_pose_w[w, t] * wgt
|
||||
vt_acc[gf] += v_trans_w[w, t] * wgt
|
||||
wt_acc[gf] += wgt
|
||||
|
||||
x_pose = x_pose + (vp_acc / wt_acc.unsqueeze(-1).clamp(min=0.01)) * dt
|
||||
x_trans = x_trans + (vt_acc / wt_acc.unsqueeze(-1).clamp(min=0.01)) * dt
|
||||
x_beta = x_beta + v_beta_w.mean(dim=0) * dt
|
||||
|
||||
acc_pose = x_pose
|
||||
acc_trans = x_trans
|
||||
acc_beta = x_beta
|
||||
|
||||
else:
|
||||
# avg: run each window's ODE independently, then uniformly average in x₁ space
|
||||
x = torch.empty(N_wins, flat_dim, device=device, dtype=dtype)
|
||||
for w in range(N_wins):
|
||||
s = w * stride
|
||||
x[w, ms.beta] = global_noise_beta
|
||||
x[w, ms.poses] = global_noise_pose[s:s + T].reshape(-1)
|
||||
x[w, ms.trans] = global_noise_trans[s:s + T].reshape(-1)
|
||||
|
||||
dt = 1.0 / ode_steps
|
||||
iterator = tqdm(range(ode_steps), desc="ODE", leave=False) if show_progress else range(ode_steps)
|
||||
|
||||
for i in iterator:
|
||||
t_step = torch.full((N_wins,), i * dt, device=device, dtype=dtype)
|
||||
v = model_engine(x, t_step, batch)
|
||||
if guidance is not None and guidance.active:
|
||||
with torch.enable_grad():
|
||||
v = guidance.compute_velocity_correction(v, x, t_step, batch)
|
||||
x = x + v * dt
|
||||
|
||||
# Accumulate per frame in normalized x₁ space; average overlapping frames
|
||||
acc_pose = torch.zeros(nf, pose_dim, device=device, dtype=dtype)
|
||||
acc_trans = torch.zeros(nf, trans_dim, device=device, dtype=dtype)
|
||||
acc_beta = torch.zeros(beta_dim, device=device, dtype=dtype)
|
||||
cnt = torch.zeros(nf, device=device, dtype=dtype)
|
||||
|
||||
for w in range(N_wins):
|
||||
s = w * stride
|
||||
e = min(s + T, nf)
|
||||
n = e - s
|
||||
acc_pose[s:e] += x[w, ms.poses].reshape(T, pose_dim)[:n]
|
||||
acc_trans[s:e] += x[w, ms.trans].reshape(T, trans_dim)[:n]
|
||||
acc_beta += x[w, ms.beta]
|
||||
cnt[s:e] += 1
|
||||
|
||||
acc_pose /= cnt.unsqueeze(-1).clamp(min=1)
|
||||
acc_trans /= cnt.unsqueeze(-1).clamp(min=1)
|
||||
acc_beta /= N_wins
|
||||
|
||||
# ── Denormalize ──────────────────────────────────────────────────────
|
||||
pose_seq = acc_pose * ms.pose_std_ + ms.pose_mean_ # (nf, 48)
|
||||
trans_seq = acc_trans * ms.trans_std_ + ms.trans_mean_ # (nf, 3)
|
||||
beta_orig = acc_beta * ms.beta_std_ + ms.beta_mean_ # (10,)
|
||||
betas_pred = beta_orig.unsqueeze(0).expand(nf, -1) # (nf, 10)
|
||||
|
||||
return pose_seq, trans_seq, betas_pred
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def evaluate_sequence_mpjpe(
|
||||
pose_pred: torch.Tensor,
|
||||
trans_pred: torch.Tensor,
|
||||
betas_pred: torch.Tensor,
|
||||
pose_gt: torch.Tensor,
|
||||
trans_gt: torch.Tensor,
|
||||
betas_gt: torch.Tensor,
|
||||
sides: List[str],
|
||||
mano_fk: MANOForwardKinematics,
|
||||
valid_mask: Optional[np.ndarray] = None,
|
||||
) -> Dict[str, float]:
|
||||
"""
|
||||
Sequence-level MPJPE evaluation.
|
||||
|
||||
Args:
|
||||
pose_pred: (T, 48) predicted pose
|
||||
trans_pred: (T, 3) predicted trans
|
||||
betas_pred: (T, 10) predicted beta
|
||||
pose_gt: (T, 48) GT pose
|
||||
trans_gt: (T, 3) GT trans
|
||||
betas_gt: (T, 10) GT beta
|
||||
sides: list of length T
|
||||
mano_fk: MANO FK module
|
||||
valid_mask: (T,) bool numpy array
|
||||
"""
|
||||
joints_gt_seq = compute_joints_global(
|
||||
pose_gt, betas_gt, trans_gt, sides, mano_fk
|
||||
)
|
||||
joints_pred_seq = compute_joints_global(
|
||||
pose_pred, betas_pred, trans_pred, sides, mano_fk
|
||||
)
|
||||
metrics = compute_mpjpe_metrics(
|
||||
joints_gt_seq, joints_pred_seq, valid_mask=valid_mask
|
||||
)
|
||||
return {
|
||||
"a_mpjpe_mm": metrics["a_mpjpe"],
|
||||
"ra_mpjpe_mm": metrics["ra_mpjpe"],
|
||||
"pa_mpjpe_mm": metrics["pa_mpjpe"],
|
||||
"w_mpjpe_mm": metrics["w_mpjpe"],
|
||||
"wa_mpjpe_mm": metrics["wa_mpjpe"],
|
||||
"a_mpjpe_per_frame": metrics["a_mpjpe_per_frame"],
|
||||
"ra_mpjpe_per_frame": metrics["ra_mpjpe_per_frame"],
|
||||
"pa_mpjpe_per_frame": metrics["pa_mpjpe_per_frame"],
|
||||
"w_mpjpe_per_frame": metrics["w_mpjpe_per_frame"],
|
||||
"wa_mpjpe_per_frame": metrics["wa_mpjpe_per_frame"],
|
||||
}
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def evaluate_fm_on_loader(
|
||||
model_engine,
|
||||
denoiser,
|
||||
mano_fk: MANOForwardKinematics,
|
||||
loader,
|
||||
device: torch.device,
|
||||
ode_steps: int,
|
||||
max_batches: Optional[int],
|
||||
is_master: bool,
|
||||
window_size: int,
|
||||
overlap_size: int,
|
||||
guidance=None,
|
||||
overlap_method: str = "vblend",
|
||||
) -> Dict[str, float]:
|
||||
"""
|
||||
Evaluate the Flow Matching model's MPJPE on a DataLoader.
|
||||
|
||||
Args:
|
||||
model_engine: DeepSpeed engine or plain model
|
||||
denoiser: HandPoseDenoiser
|
||||
mano_fk: MANO FK module
|
||||
loader: DataLoader (full_sequence mode)
|
||||
device: torch device
|
||||
ode_steps: number of ODE integration steps
|
||||
max_batches: maximum number of batches (None = all)
|
||||
is_master: whether this is the master process
|
||||
window_size: window size
|
||||
overlap_size: number of overlapping frames
|
||||
"""
|
||||
model_engine.eval()
|
||||
n_frames = 0
|
||||
sum_ra = 0.0
|
||||
sum_pa = 0.0
|
||||
sum_w = 0.0
|
||||
sum_wa = 0.0
|
||||
stride = window_size - overlap_size
|
||||
|
||||
for batch_idx, batch in enumerate(tqdm(loader, desc="MPJPE eval", leave=False, disable=not is_master)):
|
||||
if max_batches is not None and batch_idx >= max_batches:
|
||||
break
|
||||
|
||||
batch = {k: v.to(device) if isinstance(v, torch.Tensor) else v
|
||||
for k, v in batch.items()}
|
||||
B = batch["mano_params"].shape[0]
|
||||
|
||||
for b in range(B):
|
||||
pm_b = batch.get("padding_mask")
|
||||
if pm_b is not None:
|
||||
nf = int((~pm_b[b]).sum().item())
|
||||
nf = max(nf, 1)
|
||||
else:
|
||||
nf = batch["mano_params"].shape[1]
|
||||
|
||||
single = {k: (v[b:b+1] if isinstance(v, torch.Tensor) else v[b] if isinstance(v, list) else v)
|
||||
for k, v in batch.items()}
|
||||
win_batch, _ = split_sequence_into_windows(single, window_size, stride, device)
|
||||
win_batch = {k: v.to(device) if isinstance(v, torch.Tensor) else v
|
||||
for k, v in win_batch.items()}
|
||||
|
||||
pose_seq, trans_seq, betas_pred_b = run_fm_inference_with_overlap(
|
||||
model_engine, denoiser, win_batch,
|
||||
window_size, overlap_size, ode_steps, device,
|
||||
show_progress=False,
|
||||
guidance=guidance,
|
||||
overlap_method=overlap_method,
|
||||
)
|
||||
|
||||
betas_gt_b = batch["mano_betas"][b:b+1].expand(nf, -1)
|
||||
sides_b = [batch.get("side", ["right"] * nf)]
|
||||
if isinstance(sides_b[0], list):
|
||||
sides_b = sides_b[0]
|
||||
else:
|
||||
sides_b = [sides_b[0]] * nf
|
||||
gt_pose = batch["mano_params"][b, :nf]
|
||||
gt_trans = batch["mano_trans"][b, :nf]
|
||||
gt_pm = batch.get("padding_mask")
|
||||
gt_valid = (~gt_pm[b, :nf]) if gt_pm is not None else torch.ones(nf, dtype=torch.bool, device=device)
|
||||
|
||||
metrics = evaluate_sequence_mpjpe(
|
||||
pose_seq, trans_seq, betas_pred_b,
|
||||
gt_pose, gt_trans, betas_gt_b,
|
||||
sides_b, mano_fk, valid_mask=gt_valid.cpu().numpy(),
|
||||
)
|
||||
|
||||
n_valid = gt_valid.sum().item()
|
||||
sum_ra += metrics["ra_mpjpe_per_frame"].sum()
|
||||
sum_pa += metrics["pa_mpjpe_per_frame"].sum()
|
||||
sum_w += metrics["w_mpjpe_per_frame"].sum()
|
||||
sum_wa += metrics["wa_mpjpe_per_frame"].sum()
|
||||
n_frames += n_valid
|
||||
|
||||
# Distributed reduction
|
||||
try:
|
||||
from deepspeed import comm as dist
|
||||
if dist.is_initialized():
|
||||
metrics_tensor = torch.tensor(
|
||||
[sum_ra, sum_pa, sum_w, sum_wa, float(n_frames)],
|
||||
dtype=torch.float32, device=device,
|
||||
)
|
||||
dist.all_reduce(metrics_tensor)
|
||||
sum_ra, sum_pa, sum_w, sum_wa, n_frames = metrics_tensor.tolist()
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
if n_frames == 0:
|
||||
return {
|
||||
"ra_mpjpe_mm": float("nan"),
|
||||
"pa_mpjpe_mm": float("nan"),
|
||||
"w_mpjpe_mm": float("nan"),
|
||||
"wa_mpjpe_mm": float("nan"),
|
||||
}
|
||||
|
||||
return {
|
||||
"ra_mpjpe_mm": sum_ra / n_frames,
|
||||
"pa_mpjpe_mm": sum_pa / n_frames,
|
||||
"w_mpjpe_mm": sum_w / n_frames,
|
||||
"wa_mpjpe_mm": sum_wa / n_frames,
|
||||
}
|
||||
@@ -0,0 +1,182 @@
|
||||
"""
|
||||
MANO FK and joint computation utilities.
|
||||
|
||||
pose format: 48D axis-angle = [global_aa(3) | joints_aa(45)]
|
||||
"""
|
||||
|
||||
from typing import List
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from manopth.manolayer import ManoLayer
|
||||
|
||||
|
||||
class MANOForwardKinematics:
|
||||
"""Batched FK, supports mixed left/right hands."""
|
||||
|
||||
def __init__(self, mano_root: str, device: torch.device):
|
||||
self.device = device
|
||||
self._layers = {}
|
||||
for side in ("right", "left"):
|
||||
layer = ManoLayer(
|
||||
mano_root=mano_root,
|
||||
side=side,
|
||||
use_pca=False,
|
||||
flat_hand_mean=True,
|
||||
).to(device).eval()
|
||||
for p in layer.parameters():
|
||||
p.requires_grad_(False)
|
||||
self._layers[side] = layer
|
||||
|
||||
def joints_torch(
|
||||
self,
|
||||
pose: torch.Tensor,
|
||||
betas: torch.Tensor,
|
||||
trans: torch.Tensor,
|
||||
sides: List[str],
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Args:
|
||||
pose: (N, 48) — global_aa(3) + joints_aa(45)
|
||||
betas: (N, 10)
|
||||
trans: (N, 3) — translation in meters
|
||||
sides: list of "right"/"left", length N
|
||||
Returns:
|
||||
joints: (N, 21, 3) in mm
|
||||
"""
|
||||
n_samples = pose.shape[0]
|
||||
joints_out = torch.zeros(
|
||||
n_samples, 21, 3, device=pose.device, dtype=pose.dtype,
|
||||
)
|
||||
for side in ("right", "left"):
|
||||
idx = [i for i, s in enumerate(sides) if s == side]
|
||||
if not idx:
|
||||
continue
|
||||
idx_t = torch.tensor(idx, device=pose.device, dtype=torch.long)
|
||||
mano_pose = pose[idx_t] # (K, 48) fed directly into ManoLayer
|
||||
_, joints = self._layers[side](
|
||||
mano_pose,
|
||||
th_betas=betas[idx_t],
|
||||
th_trans=trans[idx_t],
|
||||
)
|
||||
joints_out[idx_t] = joints.to(dtype=pose.dtype)
|
||||
return joints_out
|
||||
|
||||
@torch.no_grad()
|
||||
def joints(
|
||||
self,
|
||||
pose: torch.Tensor,
|
||||
betas: torch.Tensor,
|
||||
trans: torch.Tensor,
|
||||
sides: List[str],
|
||||
) -> torch.Tensor:
|
||||
return self.joints_torch(pose, betas, trans, sides)
|
||||
|
||||
def verts_torch(
|
||||
self,
|
||||
pose: torch.Tensor,
|
||||
betas: torch.Tensor,
|
||||
trans: torch.Tensor,
|
||||
sides: List[str],
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Args:
|
||||
pose: (N, 48) — global_aa(3) + joints_aa(45)
|
||||
betas: (N, 10)
|
||||
trans: (N, 3) — translation in meters
|
||||
sides: list of "right"/"left", length N
|
||||
Returns:
|
||||
verts: (N, 778, 3) in mm
|
||||
"""
|
||||
n_samples = pose.shape[0]
|
||||
verts_out = torch.zeros(
|
||||
n_samples, 778, 3, device=pose.device, dtype=pose.dtype,
|
||||
)
|
||||
for side in ("right", "left"):
|
||||
idx = [i for i, s in enumerate(sides) if s == side]
|
||||
if not idx:
|
||||
continue
|
||||
idx_t = torch.tensor(idx, device=pose.device, dtype=torch.long)
|
||||
mano_pose = pose[idx_t]
|
||||
verts, _ = self._layers[side](
|
||||
mano_pose,
|
||||
th_betas=betas[idx_t],
|
||||
th_trans=trans[idx_t],
|
||||
)
|
||||
verts_out[idx_t] = verts.to(dtype=pose.dtype)
|
||||
return verts_out
|
||||
|
||||
@torch.no_grad()
|
||||
def verts(
|
||||
self,
|
||||
pose: torch.Tensor,
|
||||
betas: torch.Tensor,
|
||||
trans: torch.Tensor,
|
||||
sides: List[str],
|
||||
) -> torch.Tensor:
|
||||
return self.verts_torch(pose, betas, trans, sides)
|
||||
|
||||
def get_faces(self, side: str = "right") -> np.ndarray:
|
||||
"""Return MANO face indices as (1538, 3) numpy int32 array."""
|
||||
f = self._layers[side].th_faces
|
||||
return f.cpu().numpy().reshape(-1, 3).astype(np.int32)
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def compute_joints_global(
|
||||
pose: torch.Tensor,
|
||||
betas: torch.Tensor,
|
||||
trans: torch.Tensor,
|
||||
sides: List[str],
|
||||
mano_fk: MANOForwardKinematics,
|
||||
) -> np.ndarray:
|
||||
"""
|
||||
manopth FK → global 3D joints (metres, camera coordinate system).
|
||||
|
||||
Args:
|
||||
pose: (N, 48) — global_aa(3) + joints_aa(45)
|
||||
betas: (N, 10)
|
||||
trans: (N, 3) in metres
|
||||
sides: length N
|
||||
Returns:
|
||||
joints_m: (N, 21, 3) float32, in metres
|
||||
"""
|
||||
joints_mm = mano_fk.joints(pose, betas, trans, sides)
|
||||
return (joints_mm / 1000.0).cpu().numpy().astype(np.float32)
|
||||
|
||||
|
||||
def compute_joints_global_torch(
|
||||
pose: torch.Tensor,
|
||||
betas: torch.Tensor,
|
||||
trans: torch.Tensor,
|
||||
sides: List[str],
|
||||
mano_fk: MANOForwardKinematics,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
manopth FK → global 3D joints (metres), supports backpropagation.
|
||||
|
||||
Args:
|
||||
pose: (N, 48)
|
||||
betas: (N, 10)
|
||||
trans: (N, 3) in metres
|
||||
sides: length N
|
||||
Returns:
|
||||
joints_m: (N, 21, 3), in metres
|
||||
"""
|
||||
joints_mm = mano_fk.joints_torch(pose, betas, trans, sides)
|
||||
return joints_mm / 1000.0
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def compute_joints_global_sequence(
|
||||
pose: torch.Tensor,
|
||||
betas: torch.Tensor,
|
||||
trans: torch.Tensor,
|
||||
side: str,
|
||||
mano_fk: MANOForwardKinematics,
|
||||
) -> np.ndarray:
|
||||
"""Single-side sequence version: input (T, *), output (T, 21, 3)."""
|
||||
T = pose.shape[0]
|
||||
betas_seq = betas.unsqueeze(0).expand(T, -1)
|
||||
sides = [side] * T
|
||||
return compute_joints_global(pose, betas_seq, trans, sides, mano_fk)
|
||||
@@ -0,0 +1,438 @@
|
||||
"""
|
||||
Online HaMeR pipeline — YOLO detection + full HaMeR inference.
|
||||
|
||||
Starting from raw images, a single forward pass simultaneously outputs:
|
||||
- backbone features (ViT-H 192 tokens × 1280D)
|
||||
- 2D landmarks (crop coordinate system [0,1])
|
||||
- crop intrinsics
|
||||
- detection confidence
|
||||
|
||||
Reuses the crop utility functions and coordinate conversion logic from preprocessing/hamer_crop_and_skeleton.py.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
# Project root directory
|
||||
_PROJECT_ROOT = Path(__file__).resolve().parent.parent
|
||||
sys.path.insert(0, str(_PROJECT_ROOT))
|
||||
|
||||
from utils.dual_conf import select_target_box_dual_conf
|
||||
|
||||
# HaMeR is a submodule located at third_party/hamer; checkpoint path is set via environment variable (see README)
|
||||
HAMER_ROOT = _PROJECT_ROOT / "third_party" / "hamer"
|
||||
HAMER_CKPT = Path(os.environ.get(
|
||||
"HAMER_CKPT",
|
||||
str(HAMER_ROOT / "_DATA" / "hamer_ckpts" / "checkpoints" / "hamer.ckpt"),
|
||||
))
|
||||
# _DATA: prefer environment variable; otherwise inferred from HAMER_CKPT (ckpt is at <_DATA>/hamer_ckpts/checkpoints/)
|
||||
HAMER_DATA = Path(os.environ.get(
|
||||
"HAMER_DATA",
|
||||
str(HAMER_CKPT.parents[2]) if os.environ.get("HAMER_CKPT") else str(HAMER_ROOT / "_DATA"),
|
||||
))
|
||||
DETECTOR_PT = Path(os.environ["DETECTOR_CKPT"]) if "DETECTOR_CKPT" in os.environ else None
|
||||
|
||||
# WiLoR hand detector class mapping: cls 0=left, 1=right
|
||||
_SIDE_TO_CLS = {"left": 0, "right": 1}
|
||||
|
||||
|
||||
# ── crop utility functions (imported from preprocessing) ─────────────────────────────────
|
||||
|
||||
def _import_crop_utils():
|
||||
"""Lazy import to avoid import errors in non-GPU environments."""
|
||||
sys.path.insert(0, str(_PROJECT_ROOT))
|
||||
from preprocessing.hamer_crop_and_skeleton import (
|
||||
expand_to_aspect_ratio,
|
||||
gen_trans_from_patch_cv,
|
||||
compute_crop_intrinsics,
|
||||
_cam_crop_to_full,
|
||||
_project_full_img,
|
||||
)
|
||||
return expand_to_aspect_ratio, gen_trans_from_patch_cv, compute_crop_intrinsics, _cam_crop_to_full, _project_full_img
|
||||
|
||||
|
||||
class OnlineHaMeRPipeline:
|
||||
"""
|
||||
Online HaMeR pipeline: raw image -> YOLO detection -> crop -> full HaMeR inference.
|
||||
|
||||
Captures backbone features via a forward hook; a single forward pass jointly produces
|
||||
backbone features + 2D landmarks + crop intrinsics + confidence.
|
||||
"""
|
||||
|
||||
# ImageNet normalization (x255 variant, consistent with the HAMER model)
|
||||
_IMAGENET_MEAN = torch.tensor([123.675, 116.28, 103.53])
|
||||
_IMAGENET_STD = torch.tensor([58.395, 57.12, 57.375])
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
hamer_ckpt: str | None = None,
|
||||
detector_ckpt: str | None = None,
|
||||
device: str = "cuda",
|
||||
det_conf: float = 0.3,
|
||||
hamer_batch_size: int = 32,
|
||||
image_size: int = 256,
|
||||
):
|
||||
self.device = device
|
||||
self.det_conf = det_conf
|
||||
self.hamer_batch_size = hamer_batch_size
|
||||
self.image_size = image_size
|
||||
|
||||
# Lazy-import utility functions
|
||||
(
|
||||
self._expand_to_aspect_ratio,
|
||||
self._gen_trans_from_patch_cv,
|
||||
self._compute_crop_intrinsics,
|
||||
self._cam_crop_to_full,
|
||||
self._project_full_img,
|
||||
) = _import_crop_utils()
|
||||
|
||||
# Load YOLO + HAMER
|
||||
_hamer = hamer_ckpt or str(HAMER_CKPT)
|
||||
_detector = detector_ckpt or (str(DETECTOR_PT) if DETECTOR_PT else None)
|
||||
if _detector is None:
|
||||
raise RuntimeError(
|
||||
"YOLO detector weights not specified: please export DETECTOR_CKPT=/path/to/detector.pt, "
|
||||
"or pass the detector_ckpt argument (weights come from the WiLoR repo)"
|
||||
)
|
||||
self.detector, self.model, self.model_cfg = self._setup_models(
|
||||
_hamer, _detector, device,
|
||||
)
|
||||
|
||||
# backbone hook storage
|
||||
self._backbone_features: Optional[torch.Tensor] = None
|
||||
self._hook_handle = self.model.backbone.register_forward_hook(
|
||||
self._backbone_hook
|
||||
)
|
||||
|
||||
# normalization buffer
|
||||
self.register_buffer = lambda name, val: setattr(self, name, val.to(device))
|
||||
self._mean = self._IMAGENET_MEAN.view(1, 3, 1, 1).to(device)
|
||||
self._std = self._IMAGENET_STD.view(1, 3, 1, 1).to(device)
|
||||
|
||||
def _backbone_hook(self, module, input, output):
|
||||
"""Forward hook: captures backbone output (B, C, Hp, Wp) -> (B, 192, 1280)."""
|
||||
feat_map = output # (B, 1280, Hp, Wp)
|
||||
B, C, Hp, Wp = feat_map.shape
|
||||
self._backbone_features = feat_map.flatten(2).transpose(1, 2) # (B, Hp*Wp, 1280)
|
||||
|
||||
# ── model initialization ────────────────────────────────────────────────
|
||||
|
||||
@staticmethod
|
||||
def _setup_models(hamer_ckpt: str, detector_ckpt: str, device: str):
|
||||
"""Load the YOLO detector + the full HAMER model."""
|
||||
import gc
|
||||
|
||||
sys.path.insert(0, str(HAMER_ROOT))
|
||||
os.environ["HAMER_DATA"] = str(HAMER_DATA)
|
||||
|
||||
# Temporarily patch torch.load for compatibility with legacy ckpts
|
||||
_orig_load = torch.load
|
||||
def _patched_load(*a, **k):
|
||||
k.setdefault("weights_only", False)
|
||||
return _orig_load(*a, **k)
|
||||
torch.load = _patched_load
|
||||
|
||||
from ultralytics import YOLO
|
||||
detector = YOLO(str(detector_ckpt))
|
||||
detector.to(device)
|
||||
|
||||
torch.load = _orig_load # restore
|
||||
|
||||
import hamer.configs as _hcfg
|
||||
_hcfg.CACHE_DIR_HAMER = str(HAMER_DATA)
|
||||
|
||||
from hamer.configs import get_config
|
||||
from hamer.models.hamer import HAMER
|
||||
|
||||
model_cfg_path = str(Path(hamer_ckpt).parent.parent / "model_config.yaml")
|
||||
model_cfg = get_config(model_cfg_path, update_cachedir=True)
|
||||
|
||||
if (model_cfg.MODEL.BACKBONE.TYPE == "vit") and ("BBOX_SHAPE" not in model_cfg.MODEL):
|
||||
model_cfg.defrost()
|
||||
model_cfg.MODEL.BBOX_SHAPE = [192, 256]
|
||||
model_cfg.freeze()
|
||||
if "PRETRAINED_WEIGHTS" in model_cfg.MODEL.BACKBONE:
|
||||
model_cfg.defrost()
|
||||
model_cfg.MODEL.BACKBONE.pop("PRETRAINED_WEIGHTS")
|
||||
model_cfg.freeze()
|
||||
|
||||
# Manually load to CPU to avoid PL load_from_checkpoint keeping the entire ckpt on GPU
|
||||
model = HAMER(cfg=model_cfg, init_renderer=False)
|
||||
ckpt_dict = torch.load(str(hamer_ckpt), map_location="cpu", weights_only=False)
|
||||
state_dict = ckpt_dict.get("state_dict", ckpt_dict)
|
||||
# Clean up the ckpt dict, keeping only the state_dict
|
||||
del ckpt_dict
|
||||
gc.collect()
|
||||
model.load_state_dict(state_dict, strict=False)
|
||||
del state_dict
|
||||
gc.collect()
|
||||
model = model.to(device).eval()
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
# Print actual GPU memory usage
|
||||
if torch.cuda.is_available():
|
||||
free, total = torch.cuda.mem_get_info(0)
|
||||
print(f"[OnlineHaMeR] YOLO + HAMER loaded on {device}, "
|
||||
f"GPU memory: {(total-free)/1024**2:.0f} / {total/1024**2:.0f} MB")
|
||||
|
||||
return detector, model, model_cfg
|
||||
|
||||
# ── main interface ─────────────────────────────────────────────────────
|
||||
|
||||
@torch.no_grad()
|
||||
def process_sequence(
|
||||
self,
|
||||
raw_images_bgr: list[np.ndarray],
|
||||
raw_intrinsics_4d: list[np.ndarray],
|
||||
target_side: str = "right",
|
||||
) -> dict:
|
||||
"""
|
||||
Process the raw images of a full sequence.
|
||||
|
||||
Args:
|
||||
raw_images_bgr: list of (H, W, 3) BGR uint8 images
|
||||
raw_intrinsics_4d: list of (4,) [fx, fy, cx, cy] raw intrinsics
|
||||
target_side: "right"/"left", lock onto a single hand -- only select the YOLO box of this hand side;
|
||||
detections of the other hand are always ignored (prevents frames with both hands co-occurring from sliding onto the wrong hand). WiLoR: 0=left, 1=right.
|
||||
|
||||
Returns:
|
||||
dict with:
|
||||
crop_images: (T, 3, 256, 256) uint8
|
||||
backbone_features: (T, 192, 1280) float32 — ViT backbone output
|
||||
hamer_landmarks: (T, 21, 2) float32 — crop coordinate system [0,1]
|
||||
crop_intrinsics: (T, 4) float32
|
||||
hamer_confidence: (T,) float32
|
||||
detection_valid: (T,) bool
|
||||
det_time_s: float — YOLO detection time
|
||||
hamer_time_s: float — HaMeR inference time
|
||||
"""
|
||||
import time
|
||||
|
||||
T = len(raw_images_bgr)
|
||||
dev = self.device
|
||||
img_size = self.image_size
|
||||
|
||||
# ── 1. YOLO detection ────────────────────────────────────────────
|
||||
t0 = time.perf_counter()
|
||||
detection_batch_size = 32
|
||||
yolo_out = []
|
||||
for s in range(0, T, detection_batch_size):
|
||||
e = min(s + detection_batch_size, T)
|
||||
yolo_out.extend(
|
||||
self.detector(raw_images_bgr[s:e], conf=self.det_conf, verbose=False)
|
||||
)
|
||||
|
||||
# Parse detection results
|
||||
valid_fi = [] # frame indices with successful detections
|
||||
valid_bboxes = [] # (1, 4) bbox
|
||||
valid_confs = [] # float
|
||||
valid_hw = [] # (H, W)
|
||||
valid_sides = [] # hand side "left"/"right" for each valid frame
|
||||
detection_valid = np.zeros(T, dtype=bool)
|
||||
|
||||
for fi, yres in enumerate(yolo_out):
|
||||
if len(yres.boxes) == 0:
|
||||
continue
|
||||
clses = yres.boxes.cls.cpu().numpy()
|
||||
confs = yres.boxes.conf.cpu().numpy()
|
||||
bboxes = yres.boxes.xyxy.cpu().numpy()
|
||||
|
||||
# Lock onto target_side: among boxes of this hand side only, pick the highest conf (WiLoR: 0=left, 1=right)
|
||||
# Detections of the other hand are always ignored, preventing frames with both hands co-occurring from sliding onto the wrong hand
|
||||
target_cls = _SIDE_TO_CLS[target_side]
|
||||
# Pick highest-conf box of target hand side; but if it heavily overlaps (IoU>thr)
|
||||
# a box of the other side with higher conf, the target box is defeated (actually the
|
||||
# other hand) -> skip this frame (cmask will fallback to no-info at inference).
|
||||
sel = select_target_box_dual_conf(clses, confs, bboxes, target_cls=target_cls)
|
||||
if sel is None:
|
||||
continue # target not detected, or defeated by overlapping other-hand box
|
||||
valid_fi.append(fi)
|
||||
valid_bboxes.append(bboxes[sel:sel + 1].astype(np.float32))
|
||||
valid_confs.append(float(confs[sel]))
|
||||
valid_hw.append(raw_images_bgr[fi].shape[:2])
|
||||
valid_sides.append(target_side) # after locking, always the target hand side
|
||||
detection_valid[fi] = True
|
||||
|
||||
t_det = time.perf_counter()
|
||||
|
||||
# ── 2. generate crops + affine transform ─────────────────────────────
|
||||
BBOX_SHAPE = self.model_cfg.MODEL.get("BBOX_SHAPE", None)
|
||||
rescale_factor = 2.0
|
||||
crop_bgr_list = []
|
||||
crop_intr_list = []
|
||||
box_center_list = []
|
||||
bbox_size_list = []
|
||||
|
||||
for vi, fi in enumerate(valid_fi):
|
||||
img = raw_images_bgr[fi]
|
||||
bbox = valid_bboxes[vi]
|
||||
|
||||
box = bbox[0].astype(np.float32)
|
||||
center = (box[2:4] + box[0:2]) / 2.0
|
||||
scale = rescale_factor * (box[2:4] - box[0:2]) / 200.0
|
||||
bbox_size = float(
|
||||
self._expand_to_aspect_ratio(scale * 200, target_aspect_ratio=BBOX_SHAPE).max()
|
||||
)
|
||||
center_x, center_y = center
|
||||
|
||||
trans = self._gen_trans_from_patch_cv(
|
||||
center_x, center_y, bbox_size, bbox_size,
|
||||
img_size, img_size, 1.0, 0,
|
||||
)
|
||||
crop_bgr = cv2.warpAffine(
|
||||
img, trans, (img_size, img_size),
|
||||
flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_CONSTANT,
|
||||
)
|
||||
crop_bgr_list.append(crop_bgr)
|
||||
box_center_list.append(center.astype(np.float32))
|
||||
bbox_size_list.append(float(bbox_size))
|
||||
|
||||
# crop intrinsics
|
||||
intr = raw_intrinsics_4d[fi] if fi < len(raw_intrinsics_4d) else np.zeros(4, np.float32)
|
||||
crop_intr_list.append(
|
||||
self._compute_crop_intrinsics(trans, intr)
|
||||
)
|
||||
|
||||
# ── 3. HaMeR batched inference ───────────────────────────────────────
|
||||
backbone_features_all = np.zeros((T, 192, 1280), dtype=np.float32)
|
||||
landmarks_all = np.zeros((T, 21, 2), dtype=np.float32)
|
||||
crop_intrinsics_all = np.zeros((T, 4), dtype=np.float32)
|
||||
confidence_all = np.zeros(T, dtype=np.float32)
|
||||
crop_images_all = np.zeros((T, 3, 256, 256), dtype=np.uint8)
|
||||
|
||||
for vi, fi in enumerate(valid_fi):
|
||||
crop_intrinsics_all[fi] = crop_intr_list[vi]
|
||||
confidence_all[fi] = valid_confs[vi]
|
||||
# crop image: BGR → RGB → (3, H, W)
|
||||
crop_images_all[fi] = crop_bgr_list[vi][:, :, ::-1].transpose(2, 0, 1)
|
||||
|
||||
if len(valid_fi) > 0:
|
||||
box_center_t = torch.tensor(
|
||||
np.stack(box_center_list), dtype=torch.float32, device=dev,
|
||||
)
|
||||
bbox_size_t = torch.tensor(
|
||||
bbox_size_list, dtype=torch.float32, device=dev,
|
||||
)
|
||||
img_size_t = torch.tensor(
|
||||
[[float(w), float(h)] for h, w in valid_hw],
|
||||
dtype=torch.float32, device=dev,
|
||||
)
|
||||
|
||||
scaled_fl = float(
|
||||
self.model_cfg.EXTRA.FOCAL_LENGTH
|
||||
/ self.model_cfg.MODEL.IMAGE_SIZE
|
||||
* img_size_t[0].max().item()
|
||||
)
|
||||
|
||||
# Batched inference
|
||||
for s in range(0, len(crop_bgr_list), self.hamer_batch_size):
|
||||
e = min(s + self.hamer_batch_size, len(crop_bgr_list))
|
||||
crops_np = np.stack(crop_bgr_list[s:e])[:, :, :, ::-1].astype(np.float32)
|
||||
crops_t = torch.from_numpy(crops_np.copy()).permute(0, 3, 1, 2).to(dev)
|
||||
crops_t = (crops_t - self._mean) / self._std
|
||||
out = self.model({"img": crops_t})
|
||||
|
||||
# backbone features (captured by the hook, output for the current sub-batch)
|
||||
bf = self._backbone_features.cpu().numpy() # (batch, 192, 1280)
|
||||
for bi in range(e - s):
|
||||
backbone_features_all[valid_fi[s + bi]] = bf[bi]
|
||||
|
||||
# 2D landmarks: 3D joints -> full image -> normalize -> transform to crop
|
||||
pred_cam = out["pred_cam"]
|
||||
joints_3d = out["pred_keypoints_3d"].detach().cpu().numpy()
|
||||
|
||||
cam_t_full = self._cam_crop_to_full(
|
||||
pred_cam, box_center_t[s:e], bbox_size_t[s:e],
|
||||
img_size_t[s:e], scaled_fl,
|
||||
).detach().cpu().numpy()
|
||||
|
||||
for bi in range(e - s):
|
||||
fi = valid_fi[s + bi]
|
||||
H, W = valid_hw[s + bi]
|
||||
|
||||
# Project onto the original image
|
||||
kp2d = self._project_full_img(
|
||||
joints_3d[bi], cam_t_full[bi], scaled_fl, W, H,
|
||||
)
|
||||
# Normalize to [0,1] (original image coordinates)
|
||||
kp_n_orig = np.clip(
|
||||
kp2d / np.array([W, H], dtype=np.float32), 0.0, 1.0,
|
||||
).astype(np.float32)
|
||||
|
||||
# Transform to crop coordinate system [0,1]
|
||||
K_crop = crop_intrinsics_all[fi]
|
||||
K_orig = raw_intrinsics_4d[fi] if fi < len(raw_intrinsics_4d) else np.zeros(4)
|
||||
kp_n_crop = self._transform_landmarks_to_crop(
|
||||
kp_n_orig, K_orig, K_crop, W, H,
|
||||
)
|
||||
landmarks_all[fi] = kp_n_crop
|
||||
|
||||
# Save box_center/box_size/img_size/sides as full-length-T arrays (invalid frames are 0/default)
|
||||
box_center_all = np.zeros((T, 2), dtype=np.float32)
|
||||
bbox_size_all = np.zeros(T, dtype=np.float32)
|
||||
img_size_all = np.zeros((T, 2), dtype=np.float32)
|
||||
bbox_xyxy_all = np.zeros((T, 4), dtype=np.float32)
|
||||
conf_all = np.zeros(T, dtype=np.float32)
|
||||
sides_all = ["right"] * T
|
||||
for vi, fi in enumerate(valid_fi):
|
||||
box_center_all[fi] = box_center_list[vi]
|
||||
bbox_size_all[fi] = bbox_size_list[vi]
|
||||
img_size_all[fi] = [float(valid_hw[vi][1]), float(valid_hw[vi][0])] # (W, H)
|
||||
bbox_xyxy_all[fi] = valid_bboxes[vi][0]
|
||||
conf_all[fi] = valid_confs[vi]
|
||||
sides_all[fi] = valid_sides[vi]
|
||||
|
||||
t_hamer = time.perf_counter()
|
||||
|
||||
return {
|
||||
"crop_images": torch.from_numpy(crop_images_all),
|
||||
"sides": sides_all,
|
||||
"bbox_xyxy": torch.from_numpy(bbox_xyxy_all),
|
||||
"bbox_conf": torch.from_numpy(conf_all),
|
||||
"backbone_features": torch.from_numpy(backbone_features_all),
|
||||
"hamer_landmarks": torch.from_numpy(landmarks_all),
|
||||
"crop_intrinsics": torch.from_numpy(crop_intrinsics_all),
|
||||
"hamer_confidence": torch.from_numpy(confidence_all),
|
||||
"detection_valid": torch.from_numpy(detection_valid),
|
||||
"box_center": torch.from_numpy(box_center_all),
|
||||
"box_size": torch.from_numpy(bbox_size_all),
|
||||
"img_size": torch.from_numpy(img_size_all),
|
||||
"det_time_s": t_det - t0,
|
||||
"hamer_time_s": t_hamer - t_det,
|
||||
}
|
||||
|
||||
# ── coordinate conversion (consistent with _transform_landmarks_to_crop in dataset.py) ──
|
||||
|
||||
@staticmethod
|
||||
def _transform_landmarks_to_crop(
|
||||
landmarks_orig_n: np.ndarray,
|
||||
K_orig: np.ndarray,
|
||||
K_crop: np.ndarray,
|
||||
W: int,
|
||||
H: int,
|
||||
) -> np.ndarray:
|
||||
"""Original-image normalized [0,1] -> crop normalized [0,1]."""
|
||||
fx, fy, cx, cy = K_orig
|
||||
fxc, fyc, cxc, cyc = K_crop
|
||||
|
||||
sx = fxc / fx
|
||||
tx = cxc - fxc * cx / fx
|
||||
sy = fyc / fy
|
||||
ty = cyc - fyc * cy / fy
|
||||
|
||||
pts_x = landmarks_orig_n[:, 0] * W
|
||||
pts_y = landmarks_orig_n[:, 1] * H
|
||||
|
||||
crop_x = sx * pts_x + tx
|
||||
crop_y = sy * pts_y + ty
|
||||
|
||||
result = np.stack([crop_x / 256.0, crop_y / 256.0], axis=-1)
|
||||
return result.astype(np.float32)
|
||||
@@ -0,0 +1,179 @@
|
||||
"""ViPE SLAM worker — invoked via subprocess + conda run -n vipe.
|
||||
|
||||
Provides:
|
||||
- run_vipe_slam(): caller side (runs in the handflow env); launches a separate vipe env via subprocess
|
||||
- main(): worker side (runs in the vipe env); actually loads ViPE, runs SLAM, and writes JSON output
|
||||
|
||||
demo.py calls run_vipe_slam only in moving-camera mode (i.e. without --fix_camera): with
|
||||
explicit/default intrinsics it uses the gt_intr variant (c2w only), and with --intrinsics auto
|
||||
it uses the default variant (intrinsics + c2w).
|
||||
The worker is triggered by `conda run -n vipe python utils/vipe_worker.py --worker ...`,
|
||||
so vipe imports are deliberately deferred into main() (top-level imports would fail in the handflow env).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import subprocess
|
||||
import time
|
||||
import traceback
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
import numpy as np
|
||||
|
||||
# ViPE submodule root directory (can be overridden by the VIPE_ROOT environment variable)
|
||||
_PROJECT_ROOT = Path(__file__).resolve().parent.parent
|
||||
VIPE_ROOT = Path(os.environ.get("VIPE_ROOT", str(_PROJECT_ROOT / "third_party" / "vipe")))
|
||||
|
||||
|
||||
def run_vipe_slam(
|
||||
images_dir: str,
|
||||
output_dir: str,
|
||||
variant: str = "default",
|
||||
gt_intrinsics_4d: Optional[np.ndarray] = None,
|
||||
vipe_env: str = "vipe",
|
||||
seq_name: Optional[str] = None,
|
||||
timeout: int = 3600,
|
||||
) -> dict:
|
||||
"""Caller side: run ViPE SLAM via `conda run -n <vipe_env>`.
|
||||
|
||||
Args:
|
||||
images_dir: directory of frame images (already extracted into images)
|
||||
output_dir: output directory (slam_result.json is written here)
|
||||
variant: "default" (ViPE estimates intrinsics + c2w)
|
||||
| "gt_intr" (given intrinsics, estimate c2w only)
|
||||
gt_intrinsics_4d: [fx, fy, cx, cy] used when variant=gt_intr
|
||||
vipe_env: name of the conda environment containing ViPE (default "vipe")
|
||||
Returns:
|
||||
{"poses": (T,4,4) c2w, "intrinsics": (T,4) or (4,), "slam_time_s": float}
|
||||
"""
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
output_json = str(Path(output_dir) / "slam_result.json")
|
||||
seq_name = seq_name or Path(images_dir).name
|
||||
|
||||
vipe_python = os.environ.get("VIPE_PYTHON")
|
||||
launcher = [vipe_python] if vipe_python else [
|
||||
"conda", "run", "-n", vipe_env, "--no-capture-output", "python",
|
||||
]
|
||||
cmd = launcher + [
|
||||
str(Path(__file__).resolve()), "--worker",
|
||||
"--img_dir", images_dir,
|
||||
"--seq_name", seq_name,
|
||||
"--variant", variant,
|
||||
"--output", output_json,
|
||||
]
|
||||
if variant == "gt_intr" and gt_intrinsics_4d is not None:
|
||||
intr = np.asarray(gt_intrinsics_4d).ravel()[:4]
|
||||
cmd += ["--gt_intr_str", ",".join(str(float(x)) for x in intr)]
|
||||
|
||||
result = subprocess.run(
|
||||
cmd, capture_output=True, text=True, timeout=timeout, cwd=str(_PROJECT_ROOT)
|
||||
)
|
||||
if result.returncode != 0:
|
||||
raise RuntimeError(
|
||||
f"ViPE SLAM failed (returncode={result.returncode}):\n"
|
||||
f"{result.stderr.strip()[-800:]}"
|
||||
)
|
||||
if not os.path.exists(output_json):
|
||||
raise RuntimeError(f"ViPE worker produced no output: {result.stderr[-400:]}")
|
||||
|
||||
with open(output_json) as f:
|
||||
data = json.load(f)
|
||||
if not data.get("ok"):
|
||||
raise RuntimeError(f"ViPE worker error: {data.get('error')}")
|
||||
|
||||
return {
|
||||
"poses": np.array(data["poses"], dtype=np.float32), # (T, 4, 4) c2w
|
||||
"intrinsics": np.array(data["intrinsics"], dtype=np.float32), # (T, 4) or (4,)
|
||||
"slam_time_s": float(data.get("slam_time_s", 0.0)),
|
||||
}
|
||||
|
||||
|
||||
# ── Worker (executed inside the vipe environment) ──────────────────────────────────────────
|
||||
|
||||
def _worker_main() -> None:
|
||||
"""Actually loads ViPE, runs SLAM, and writes results to JSON. Only runnable in the vipe environment."""
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--worker", action="store_true")
|
||||
parser.add_argument("--img_dir", required=True)
|
||||
parser.add_argument("--seq_name", default="seq")
|
||||
parser.add_argument("--variant", default="default", choices=["default", "gt_intr"])
|
||||
parser.add_argument("--gt_intr_str", default="")
|
||||
parser.add_argument("--output", required=True)
|
||||
args = parser.parse_args()
|
||||
|
||||
from omegaconf import OmegaConf
|
||||
import torch
|
||||
from vipe.streams.frame_dir_stream import FrameDirStream
|
||||
from vipe.slam.system import SLAMSystem
|
||||
from vipe.streams.base import ProcessedVideoStream, StreamProcessor, FrameAttribute
|
||||
from vipe.utils.cameras import CameraType
|
||||
|
||||
slam_cfg = OmegaConf.load(str(VIPE_ROOT / "configs" / "slam" / "default.yaml"))
|
||||
slam_cfg.optimize_intrinsics = (args.variant != "gt_intr")
|
||||
|
||||
try:
|
||||
stream = FrameDirStream(path=Path(args.img_dir),
|
||||
name=args.seq_name.replace("/", "_"))
|
||||
|
||||
if args.variant == "gt_intr":
|
||||
fx, fy, cx, cy = [float(x) for x in args.gt_intr_str.split(",")]
|
||||
|
||||
class GTIntrinsicsProcessor(StreamProcessor):
|
||||
def update_attributes(self, prev):
|
||||
return prev | {FrameAttribute.INTRINSICS}
|
||||
|
||||
def __call__(self, frame_idx, frame):
|
||||
frame.intrinsics = torch.as_tensor([fx, fy, cx, cy]).float()
|
||||
frame.camera_type = CameraType.PINHOLE
|
||||
return frame
|
||||
|
||||
stream = ProcessedVideoStream(stream, [GTIntrinsicsProcessor()])
|
||||
else:
|
||||
from vipe.pipeline.processors import GeoCalibIntrinsicsProcessor
|
||||
stream = ProcessedVideoStream(
|
||||
stream, [GeoCalibIntrinsicsProcessor(stream, camera_type=CameraType.PINHOLE)],
|
||||
)
|
||||
|
||||
device = torch.device("cuda")
|
||||
slam_system = SLAMSystem(device=device, config=slam_cfg)
|
||||
t0 = time.perf_counter()
|
||||
slam_output = slam_system.run([stream], camera_type=CameraType.PINHOLE)
|
||||
slam_time = time.perf_counter() - t0
|
||||
|
||||
traj_mat = slam_output.trajectory.matrix()
|
||||
if hasattr(traj_mat, "detach"):
|
||||
traj_mat = traj_mat.detach().cpu().numpy()
|
||||
intr_tensor = slam_output.intrinsics
|
||||
if hasattr(intr_tensor, "detach"):
|
||||
intr_tensor = intr_tensor.detach().cpu().numpy()
|
||||
|
||||
n_frames = len(traj_mat)
|
||||
if len(intr_tensor) == 1 and n_frames > 1:
|
||||
intr_tensor = np.tile(intr_tensor, (n_frames, 1))
|
||||
N = min(len(traj_mat), len(intr_tensor))
|
||||
|
||||
result = {
|
||||
"poses": traj_mat[:N].tolist(),
|
||||
"intrinsics": intr_tensor[:N].tolist(),
|
||||
"slam_time_s": slam_time,
|
||||
"ok": True,
|
||||
}
|
||||
except Exception as e:
|
||||
result = {
|
||||
"ok": False,
|
||||
"error": str(e),
|
||||
"traceback": traceback.format_exc()[-500:],
|
||||
"slam_time_s": 0.0,
|
||||
}
|
||||
|
||||
os.makedirs(os.path.dirname(args.output), exist_ok=True)
|
||||
with open(args.output, "w") as f:
|
||||
json.dump(result, f)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
_worker_main()
|
||||
Reference in New Issue
Block a user