Add default FMPose3DConfig per model_type

Author: deruyter92

fmpose3d/common/config.py

@@ -9,8 +9,7 @@
 
 import math
 from dataclasses import dataclass, field, fields, asdict
-from typing import List
-
+from typing import Dict, List
 
 # ---------------------------------------------------------------------------
 # Dataclass configuration groups
@@ -23,21 +22,51 @@ class ModelConfig:
     model_type: str = "fmpose3d"
 
 
+# Per-model-type defaults for fields marked with INFER_FROM_MODEL_TYPE.
+# Also consumed by PipelineConfig.for_model_type to set cross-config
+# values (dataset, sample_steps, etc.).
+_FMPOSE3D_DEFAULTS: Dict[str, Dict] = {
+    "fmpose3d": {
+        "n_joints": 17,
+        "out_joints": 17,
+        "dataset": "h36m",
+        "sample_steps": 3,
+    },
+    "fmpose3d_animals": {
+        "n_joints": 26,
+        "out_joints": 26,
+        "dataset": "animal3d",
+        "sample_steps": 3,
+    },
+}
+
+# Sentinel object for defaults that are inferred from the model type.
+INFER_FROM_MODEL_TYPE = object()
+
 @dataclass
 class FMPose3DConfig(ModelConfig):
-    model: str = ""
     model_type: str = "fmpose3d"
-    layers: int = 3
+    model: str = ""
+    layers: int = 5
     channel: int = 512
     d_hid: int = 1024
     token_dim: int = 256
-    n_joints: int = 17
-    out_joints: int = 17
+    n_joints: int = INFER_FROM_MODEL_TYPE  # type: ignore[assignment]
+    out_joints: int = INFER_FROM_MODEL_TYPE  # type: ignore[assignment]
     in_channels: int = 2
     out_channels: int = 3
     frames: int = 1
-    """Optional: load model class from a specific file path."""
 
+    def __post_init__(self):
+        defaults = _FMPOSE3D_DEFAULTS.get(self.model_type)
+        if defaults is None:
+            supported = ", ".join(sorted(_FMPOSE3D_DEFAULTS))
+            raise ValueError(
+                f"Unknown model_type {self.model_type!r}; supported: {supported}"
+            )
+        for f in fields(self):
+            if getattr(self, f.name) is INFER_FROM_MODEL_TYPE:
+                setattr(self, f.name, defaults[f.name])
 
 @dataclass
 class DatasetConfig:
@@ -239,8 +268,6 @@ class PipelineConfig:
     demo_cfg: DemoConfig = field(default_factory=DemoConfig)
     runtime_cfg: RuntimeConfig = field(default_factory=RuntimeConfig)
 
-    # -- construction from argparse namespace ---------------------------------
-
     @classmethod
     def from_namespace(cls, ns) -> "PipelineConfig":
         """Build a :class:`PipelineConfig` from an ``argparse.Namespace``
@@ -258,7 +285,7 @@ def _pick(dc_class, src: dict):
 
         kwargs = {}
         for group_name, dc_class in _SUB_CONFIG_CLASSES.items():
-            if group_name == "model_cfg" and raw.get("model_type", "fmpose3d") == "fmpose3d":
+            if group_name == "model_cfg" and raw.get("model_type", 'fmpose3d') in _FMPOSE3D_DEFAULTS:
                 dc_class = FMPose3DConfig
             elif group_name == "pose2d_cfg" and raw.get("pose2d_model", "hrnet") == "hrnet":
                 dc_class = HRNetConfig