Feat: add extendable model registry

- New ABC base_model as template.
- Easy access to defined set of models.
- Modularly extendable with new implementations.

Author: deruyter92

fmpose3d/models/__init__.py

@@ -11,10 +11,16 @@
 FMPose3D models.
 """
 
-from .graph_frames import Graph
-from .model_GAMLP import Model
+from .base_model import BaseModel, register_model, get_model, list_models
+
+# Import model subpackages so their @register_model decorators execute.
+from .fmpose3d import Graph, Model
 
 __all__ = [
+    "BaseModel",
+    "register_model",
+    "get_model",
+    "list_models",
     "Graph",
     "Model",
 ]

fmpose3d/models/base_model.py

@@ -0,0 +1,114 @@
+"""
+FMPose3D: monocular 3D Pose Estimation via Flow Matching
+
+Official implementation of the paper:
+"FMPose3D: monocular 3D Pose Estimation via Flow Matching"
+by Ti Wang, Xiaohang Yu, and Mackenzie Weygandt Mathis
+Licensed under Apache 2.0
+"""
+
+from abc import ABC, abstractmethod
+import warnings
+
+import torch
+import torch.nn as nn
+
+
+# ---------------------------------------------------------------------------
+# Model registry
+# ---------------------------------------------------------------------------
+
+_MODEL_REGISTRY: dict[str, type["BaseModel"]] = {}
+
+
+def register_model(name: str):
+    """Class decorator that registers a model under *name*.
+
+    Usage::
+
+        @register_model("my_model")
+        class MyModel(BaseModel):
+            ...
+
+    The model can then be retrieved with :func:`get_model`.
+    """
+
+    def decorator(cls: type["BaseModel"]) -> type["BaseModel"]:
+        if name in _MODEL_REGISTRY:
+            warnings.warn(
+                f"Model '{name}' is already registered "
+                f"(existing: {_MODEL_REGISTRY[name].__qualname__}, "
+                f"new: {cls.__qualname__})"
+            )
+            # raise ValueError(
+            #     f"Model '{name}' is already registered "
+            #     f"(existing: {_MODEL_REGISTRY[name].__qualname__}, "
+            #     f"new: {cls.__qualname__})"
+            # )
+        _MODEL_REGISTRY[name] = cls
+        return cls
+
+    return decorator
+
+
+def get_model(name: str) -> type["BaseModel"]:
+    """Return the model class registered under *name*.
+
+    Raises :class:`KeyError` with a helpful message when the name is unknown.
+    """
+    if name not in _MODEL_REGISTRY:
+        available = ", ".join(sorted(_MODEL_REGISTRY)) or "(none)"
+        raise KeyError(
+            f"Unknown model '{name}'. Available models: {available}"
+        )
+    return _MODEL_REGISTRY[name]
+
+
+def list_models() -> list[str]:
+    """Return a sorted list of all registered model names."""
+    return sorted(_MODEL_REGISTRY)
+
+
+# ---------------------------------------------------------------------------
+# Base model
+# ---------------------------------------------------------------------------
+
+
+class BaseModel(ABC, nn.Module):
+    """Abstract base class for all FMPose3D lifting models.
+
+    Every model must accept a single *args* namespace / object in its
+    constructor and implement :meth:`forward` with the signature below.
+
+    Parameters expected on *args* (at minimum):
+        - ``channel``   – embedding dimension
+        - ``layers``    – number of transformer / GCN blocks
+        - ``d_hid``     – hidden MLP dimension
+        - ``token_dim`` – token dimension
+        - ``n_joints``  – number of body joints
+    """
+
+    @abstractmethod
+    def __init__(self, args) -> None:
+        super().__init__()
+
+    @abstractmethod
+    def forward(
+        self,
+        pose_2d: torch.Tensor,
+        y_t: torch.Tensor,
+        t: torch.Tensor,
+    ) -> torch.Tensor:
+        """Predict the velocity field for flow matching.
+
+        Args:
+            pose_2d: 2D keypoints, shape ``(B, F, J, 2)``.
+            y_t: Noisy 3D hypothesis at time *t*, shape ``(B, F, J, 3)``.
+            t: Diffusion / flow time, shape ``(B, F, 1, 1)`` with values
+               in ``[0, 1]``.
+
+        Returns:
+            Predicted velocity ``v``, shape ``(B, F, J, 3)``.
+        """
+        ...
+

fmpose3d/models/fmpose3d/__init__.py

@@ -0,0 +1,12 @@
+"""
+FMPose3D model subpackage.
+"""
+
+from .graph_frames import Graph
+from .model_GAMLP import Model
+
+__all__ = [
+    "Graph",
+    "Model",
+]
+

fmpose3d/models/graph_frames.py fmpose3d/models/fmpose3d/graph_frames.pyfmpose3d/models/graph_frames.py renamed to fmpose3d/models/fmpose3d/graph_frames.py

@@ -207,4 +207,5 @@ def normalize_undigraph(A):
 if __name__=="__main__":
     graph = Graph('hm36_gt', 'spatial', 1)
     print(graph.A.shape)
-    # print(graph)
+    # print(graph)
+

fmpose3d/models/model_GAMLP.py fmpose3d/models/fmpose3d/model_GAMLP.pyfmpose3d/models/model_GAMLP.py renamed to fmpose3d/models/fmpose3d/model_GAMLP.py

@@ -11,7 +11,8 @@
 import torch.nn as nn
 import math
 from einops import rearrange
-from fmpose3d.models.graph_frames import Graph
+from fmpose3d.models.fmpose3d.graph_frames import Graph
+from fmpose3d.models.base_model import BaseModel, register_model
 from functools import partial
 from einops import rearrange
 from timm.models.layers import DropPath
@@ -211,9 +212,10 @@ def forward(self, x):
         x = self.fc2(x)
         return x
 
-class Model(nn.Module):
+@register_model("fmpose3d")
+class Model(BaseModel):
     def __init__(self, args):
-        super().__init__()
+        super().__init__(args)
         ## GCN
         self.graph = Graph('hm36_gt', 'spatial', pad=1)
         # Register as buffer (not parameter) to follow module device automatically
@@ -264,4 +266,5 @@ class Args:
     y_t = torch.randn(1, 17, 17, 3, device=device)
     t = torch.randn(1, 1, 1, 1, device=device)
     v = model(x, y_t, t)
-    print(v.shape) 
+    print(v.shape) 
+