Update config: extendable configs and changed name -> PipelineConfig

Author: deruyter92

fmpose3d/common/__init__.py

@@ -13,8 +13,11 @@
 
 from .arguments import opts
 from .config import (
-    FMPoseConfig,
+    PipelineConfig,
     ModelConfig,
+    FMPose3DConfig,
+    HRNetConfig,
+    Pose2DConfig,
     DatasetConfig,
     TrainingConfig,
     InferenceConfig,
@@ -40,7 +43,10 @@
 
 __all__ = [
     "opts",
-    "FMPoseConfig",
+    "PipelineConfig",
+    "FMPose3DConfig",
+    "HRNetConfig",
+    "Pose2DConfig",
     "ModelConfig",
     "DatasetConfig",
     "TrainingConfig",

fmpose3d/common/config.py

@@ -20,7 +20,11 @@
 @dataclass
 class ModelConfig:
     """Model architecture configuration."""
+    model_type: str = "fmpose3d"
+
 
+@dataclass
+class FMPose3DConfig(ModelConfig):
     model: str = ""
     model_type: str = "fmpose3d"
     layers: int = 3
@@ -141,6 +145,20 @@ class OutputConfig:
     sh_file: str = ""
 
 
+@dataclass
+class Pose2DConfig:
+    """2D pose estimator configuration."""
+    pose2d_model: str = "hrnet"
+
+
+@dataclass
+class HRNetConfig(Pose2DConfig):
+    """HRNet 2D pose detector configuration."""
+    pose2d_model: str = "hrnet"
+    det_dim: int = 416 # YOLO input dimension (HRNet-specific).
+    num_persons: int = 1 # Maximum number of persons to estimate per frame.
+
+
 @dataclass
 class DemoConfig:
     """Demo / inference configuration."""
@@ -174,53 +192,59 @@ class RuntimeConfig:
     "checkpoint_cfg": CheckpointConfig,
     "refinement_cfg": RefinementConfig,
     "output_cfg": OutputConfig,
+    "pose2d_cfg": Pose2DConfig,
     "demo_cfg": DemoConfig,
     "runtime_cfg": RuntimeConfig,
 }
 
 
 @dataclass
-class FMPoseConfig:
-    """Top-level configuration for FMPose3D.
+class PipelineConfig:
+    """Top-level configuration for FMPose3D pipeline.
 
     Groups related settings into sub-configs::
 
         config.model_cfg.layers
         config.training_cfg.lr
     """
 
-    model_cfg: ModelConfig = field(default_factory=ModelConfig)
+    model_cfg: ModelConfig = field(default_factory=FMPose3DConfig)
     dataset_cfg: DatasetConfig = field(default_factory=DatasetConfig)
     training_cfg: TrainingConfig = field(default_factory=TrainingConfig)
     inference_cfg: InferenceConfig = field(default_factory=InferenceConfig)
     aggregation_cfg: AggregationConfig = field(default_factory=AggregationConfig)
     checkpoint_cfg: CheckpointConfig = field(default_factory=CheckpointConfig)
     refinement_cfg: RefinementConfig = field(default_factory=RefinementConfig)
     output_cfg: OutputConfig = field(default_factory=OutputConfig)
+    pose2d_cfg: Pose2DConfig = field(default_factory=HRNetConfig)
     demo_cfg: DemoConfig = field(default_factory=DemoConfig)
     runtime_cfg: RuntimeConfig = field(default_factory=RuntimeConfig)
 
     # -- construction from argparse namespace ---------------------------------
 
     @classmethod
-    def from_namespace(cls, ns) -> "FMPoseConfig":
-        """Build a :class:`FMPoseConfig` from an ``argparse.Namespace``
+    def from_namespace(cls, ns) -> "PipelineConfig":
+        """Build a :class:`PipelineConfig` from an ``argparse.Namespace``
 
         Example::
 
             args = opts().parse()
-            cfg = FMPoseConfig.from_namespace(args)
+            cfg = PipelineConfig.from_namespace(args)
         """
         raw = vars(ns) if hasattr(ns, "__dict__") else dict(ns)
 
         def _pick(dc_class, src: dict):
             names = {f.name for f in fields(dc_class)}
             return dc_class(**{k: v for k, v in src.items() if k in names})
 
-        return cls(**{
-            group_name: _pick(dc_class, raw)
-            for group_name, dc_class in _SUB_CONFIG_CLASSES.items()
-        })
+        kwargs = {}
+        for group_name, dc_class in _SUB_CONFIG_CLASSES.items():
+            if group_name == "model_cfg" and raw.get("model_type", "fmpose3d") == "fmpose3d":
+                dc_class = FMPose3DConfig
+            elif group_name == "pose2d_cfg" and raw.get("pose2d_model", "hrnet") == "hrnet":
+                dc_class = HRNetConfig
+            kwargs[group_name] = _pick(dc_class, raw)
+        return cls(**kwargs)
 
     # -- utilities ------------------------------------------------------------
 

tests/test_config.py

@@ -13,8 +13,8 @@
 import pytest
 
 from fmpose3d.common.config import (
-    FMPoseConfig,
-    ModelConfig,
+    PipelineConfig,
+    FMPose3DConfig,
     DatasetConfig,
     TrainingConfig,
     InferenceConfig,
@@ -33,19 +33,18 @@
 # ---------------------------------------------------------------------------
 
 
-class TestModelConfig:
+class TestFMPose3DConfig:
     def test_defaults(self):
-        cfg = ModelConfig()
+        cfg = FMPose3DConfig()
         assert cfg.layers == 3
         assert cfg.channel == 512
         assert cfg.d_hid == 1024
         assert cfg.n_joints == 17
         assert cfg.out_joints == 17
         assert cfg.frames == 1
-        assert cfg.model_path == ""
 
     def test_custom_values(self):
-        cfg = ModelConfig(layers=5, channel=256, n_joints=26)
+        cfg = FMPose3DConfig(layers=5, channel=256, n_joints=26)
         assert cfg.layers == 5
         assert cfg.channel == 256
         assert cfg.n_joints == 26
@@ -167,15 +166,15 @@ def test_defaults(self):
 
 
 # ---------------------------------------------------------------------------
-# FMPoseConfig
+# PipelineConfig
 # ---------------------------------------------------------------------------
 
 
-class TestFMPoseConfig:
+class TestPipelineConfig:
     def test_default_construction(self):
         """All sub-configs are initialised with their defaults."""
-        cfg = FMPoseConfig()
-        assert isinstance(cfg.model_cfg, ModelConfig)
+        cfg = PipelineConfig()
+        assert isinstance(cfg.model_cfg, FMPose3DConfig)
         assert isinstance(cfg.dataset_cfg, DatasetConfig)
         assert isinstance(cfg.training_cfg, TrainingConfig)
         assert isinstance(cfg.inference_cfg, InferenceConfig)
@@ -188,8 +187,8 @@ def test_default_construction(self):
 
     def test_partial_construction(self):
         """Supplying only some sub-configs leaves the rest at defaults."""
-        cfg = FMPoseConfig(
-            model_cfg=ModelConfig(layers=5),
+        cfg = PipelineConfig(
+            model_cfg=FMPose3DConfig(layers=5),
             training_cfg=TrainingConfig(lr=2e-4),
         )
         assert cfg.model_cfg.layers == 5
@@ -200,21 +199,21 @@ def test_partial_construction(self):
 
     def test_sub_config_mutation(self):
         """Mutating a sub-config field is reflected on the config."""
-        cfg = FMPoseConfig()
+        cfg = PipelineConfig()
         cfg.training_cfg.lr = 0.01
         assert cfg.training_cfg.lr == pytest.approx(0.01)
 
     def test_sub_config_replacement(self):
         """Replacing an entire sub-config works."""
-        cfg = FMPoseConfig()
-        cfg.model_cfg = ModelConfig(layers=10, channel=1024)
+        cfg = PipelineConfig()
+        cfg.model_cfg = FMPose3DConfig(layers=10, channel=1024)
         assert cfg.model_cfg.layers == 10
         assert cfg.model_cfg.channel == 1024
 
     # -- to_dict --------------------------------------------------------------
 
     def test_to_dict_returns_flat_dict(self):
-        cfg = FMPoseConfig()
+        cfg = PipelineConfig()
         d = cfg.to_dict()
         assert isinstance(d, dict)
         # Spot-check keys from different groups
@@ -225,8 +224,8 @@ def test_to_dict_returns_flat_dict(self):
         assert "gpu" in d
 
     def test_to_dict_reflects_custom_values(self):
-        cfg = FMPoseConfig(
-            model_cfg=ModelConfig(layers=7),
+        cfg = PipelineConfig(
+            model_cfg=FMPose3DConfig(layers=7),
             aggregation_cfg=AggregationConfig(topk=5),
         )
         d = cfg.to_dict()
@@ -235,7 +234,7 @@ def test_to_dict_reflects_custom_values(self):
 
     def test_to_dict_no_duplicate_keys(self):
         """Every field name should be unique across all sub-configs."""
-        cfg = FMPoseConfig()
+        cfg = PipelineConfig()
         d = cfg.to_dict()
         all_field_names = []
         for dc_class in _SUB_CONFIG_CLASSES.values():
@@ -249,8 +248,9 @@ def test_to_dict_no_duplicate_keys(self):
 
     def test_from_namespace_basic(self):
         ns = argparse.Namespace(
-            # ModelConfig
+            # FMPose3DConfig
             model="test_model",
+            model_type="fmpose3d",
             layers=5,
             channel=256,
             d_hid=512,
@@ -260,7 +260,6 @@ def test_from_namespace_basic(self):
             in_channels=2,
             out_channels=3,
             frames=3,
-            model_path="/tmp/model.py",
             # DatasetConfig
             dataset="rat7m",
             keypoints="cpn",
@@ -339,7 +338,7 @@ def test_from_namespace_basic(self):
             single=True,
             reload_3d=False,
         )
-        cfg = FMPoseConfig.from_namespace(ns)
+        cfg = PipelineConfig.from_namespace(ns)
 
         # Verify a sample from each group
         assert cfg.model_cfg.layers == 5
@@ -363,15 +362,15 @@ def test_from_namespace_ignores_unknown_fields(self):
         ns = argparse.Namespace(
             layers=3, channel=512, unknown_field="should_be_ignored",
         )
-        cfg = FMPoseConfig.from_namespace(ns)
+        cfg = PipelineConfig.from_namespace(ns)
         assert cfg.model_cfg.layers == 3
         assert cfg.model_cfg.channel == 512
         assert not hasattr(cfg, "unknown_field")
 
     def test_from_namespace_partial_namespace(self):
         """A namespace missing some fields uses dataclass defaults for those."""
         ns = argparse.Namespace(layers=10, gpu="2")
-        cfg = FMPoseConfig.from_namespace(ns)
+        cfg = PipelineConfig.from_namespace(ns)
         assert cfg.model_cfg.layers == 10
         assert cfg.runtime_cfg.gpu == "2"
         # Unset fields keep defaults
@@ -385,7 +384,7 @@ def test_roundtrip_from_namespace_to_dict(self):
         ns = argparse.Namespace(
             layers=8, channel=1024, dataset="animal3d", lr=2e-4, topk=7, gpu="3",
         )
-        cfg = FMPoseConfig.from_namespace(ns)
+        cfg = PipelineConfig.from_namespace(ns)
         d = cfg.to_dict()
         assert d["layers"] == 8
         assert d["channel"] == 1024
@@ -396,10 +395,9 @@ def test_roundtrip_from_namespace_to_dict(self):
 
     def test_to_dict_after_mutation(self):
         """to_dict reflects in-place mutations on sub-configs."""
-        cfg = FMPoseConfig()
+        cfg = PipelineConfig()
         cfg.training_cfg.lr = 0.123
         cfg.model_cfg.layers = 99
         d = cfg.to_dict()
         assert d["lr"] == pytest.approx(0.123)
         assert d["layers"] == 99
-