Add high-level inference API for FMPose3D (fmpose3d/fmpose3d.py)

Author: deruyter92

fmpose3d/__init__.py

@@ -18,13 +18,43 @@
     aggregation_RPEA_joint_level,
 )
 
+# Configuration dataclasses
+from .common.config import (
+    FMPose3DConfig,
+    HRNetConfig,
+    InferenceConfig,
+    ModelConfig,
+    PipelineConfig,
+)
+
+# High-level inference API
+from .fmpose3d import (
+    FMPose3DInference,
+    HRNetEstimator,
+    Pose2DResult,
+    Pose3DResult,
+    Source,
+)
+
 # Import 2D pose detection utilities
 from .lib.hrnet.gen_kpts import gen_video_kpts
 from .lib.hrnet.hrnet import HRNetPose2d
 from .lib.preprocess import h36m_coco_format, revise_kpts
 
 # Make commonly used classes/functions available at package level
 __all__ = [
+    # Inference API
+    "FMPose3DInference",
+    "HRNetEstimator",
+    "Pose2DResult",
+    "Pose3DResult",
+    "Source",
+    # Configuration
+    "FMPose3DConfig",
+    "HRNetConfig",
+    "InferenceConfig",
+    "ModelConfig",
+    "PipelineConfig",
     # Aggregation methods
     "average_aggregation",
     "aggregation_select_single_best_hypothesis_by_2D_error",

fmpose3d/common/utils.py

@@ -16,6 +16,44 @@
 import numpy as np
 import torch
 
+
+def euler_sample(
+    c_2d: torch.Tensor,
+    y: torch.Tensor,
+    steps: int,
+    model: torch.nn.Module,
+) -> torch.Tensor:
+    """Euler ODE sampler for Conditional Flow Matching at test time.
+
+    Integrates the learned velocity field from *t = 0* to *t = 1* using
+    ``steps`` uniform Euler steps.
+
+    Parameters
+    ----------
+    c_2d : Tensor
+        2-D conditioning input, shape ``(B, F, J, 2)``.
+    y : Tensor
+        Initial noise sample (same spatial dims as ``c_2d`` but with 3
+        output channels), shape ``(B, F, J, 3)``.
+    steps : int
+        Number of Euler integration steps.
+    model : nn.Module
+        The velocity-prediction network ``v(c_2d, y, t)``.
+
+    Returns
+    -------
+    Tensor
+        The denoised 3-D prediction, same shape as *y*.
+    """
+    dt = 1.0 / steps
+    for s in range(steps):
+        t_s = torch.full(
+            (c_2d.size(0), 1, 1, 1), s * dt, device=c_2d.device, dtype=c_2d.dtype
+        )
+        v_s = model(c_2d, y, t_s)
+        y = y + dt * v_s
+    return y
+
 def deterministic_random(min_value, max_value, data):
     digest = hashlib.sha256(data.encode()).digest()
     raw_value = int.from_bytes(digest[:4], byteorder="little", signed=False)