From 799be4dd101abee1c67462ed2b5d97307f0bc035 Mon Sep 17 00:00:00 2001
From: ti <tiwang.cs@gmail.com>
Date: Tue, 10 Feb 2026 15:06:40 +0100
Subject: [PATCH 01/21] Refactor FMPose3D test script to use model_type instead
 of model_path

---
 scripts/FMPose3D_test.sh | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/scripts/FMPose3D_test.sh b/scripts/FMPose3D_test.sh
index 3d2a615..b83d65d 100755
--- a/scripts/FMPose3D_test.sh
+++ b/scripts/FMPose3D_test.sh
@@ -10,7 +10,7 @@ mode='exp'
 exp_temp=0.005
 folder_name=test_s${eval_multi_steps}_${mode}_h${num_hypothesis_list}_$(date +%Y%m%d_%H%M%S)
 
-model_path='./pre_trained_models/fmpose3d_h36m/model_GAMLP.py'
+model_type='fmpose3d'
 model_weights_path='./pre_trained_models/fmpose3d_h36m/FMpose3D_pretrained_weights.pth'
 
 #Test
@@ -20,7 +20,7 @@ python3 scripts/FMPose3D_main.py \
 --exp_temp ${exp_temp} \
 --folder_name ${folder_name} \
 --model_weights_path "${model_weights_path}" \
---model_path "${model_path}" \
+--model_type "${model_type}" \
 --eval_sample_steps ${eval_multi_steps} \
 --test_augmentation True \
 --batch_size ${batch_size} \

From 7acba1f81fab36492492bf2e49c4d7ba0a58ca7b Mon Sep 17 00:00:00 2001
From: ti <tiwang.cs@gmail.com>
Date: Tue, 10 Feb 2026 15:09:56 +0100
Subject: [PATCH 02/21] Refactor FMPose3D_main.py to load model using get_model
 from registry instead of model_path

---
 scripts/FMPose3D_main.py | 11 +++--------
 1 file changed, 3 insertions(+), 8 deletions(-)

diff --git a/scripts/FMPose3D_main.py b/scripts/FMPose3D_main.py
index e9adf3a..beb88da 100644
--- a/scripts/FMPose3D_main.py
+++ b/scripts/FMPose3D_main.py
@@ -41,8 +41,9 @@
     spec.loader.exec_module(module)
     CFM = getattr(module, "Model")
 else:
-    # Load model from installed fmpose package
-    from fmpose3d.models import Model as CFM
+    # Load model from registered model registry
+    from fmpose3d.models import get_model
+    CFM = get_model(args.model_type)
 
 
 def test_multi_hypothesis(
@@ -281,12 +282,6 @@ def print_error_action(action_error_sum, is_train):
             src=script_path,
             dst=os.path.join(args.checkpoint, args.create_time + "_" + script_name),
         )
-        if getattr(args, "model_path", ""):
-            model_src_path = os.path.abspath(args.model_path)
-            model_dst_name = f"{args.create_time}_" + os.path.basename(model_src_path)
-            shutil.copyfile(
-                src=model_src_path, dst=os.path.join(args.checkpoint, model_dst_name)
-            )
         sh_base = os.path.basename(args.sh_file)
         dst_name = f"{args.create_time}_" + sh_base
         shutil.copyfile(src=args.sh_file, dst=os.path.join(args.checkpoint, dst_name))

From f173935c9219a6bbf576b295483bc5746fab490d Mon Sep 17 00:00:00 2001
From: ti <tiwang.cs@gmail.com>
Date: Tue, 10 Feb 2026 15:10:03 +0100
Subject: [PATCH 03/21] Update FMPose3D_train.sh to use model_type for model
 selection instead of model_path

---
 scripts/FMPose3D_train.sh | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/scripts/FMPose3D_train.sh b/scripts/FMPose3D_train.sh
index 939658c..a677d9c 100755
--- a/scripts/FMPose3D_train.sh
+++ b/scripts/FMPose3D_train.sh
@@ -11,8 +11,7 @@ epochs=80
 num_saved_models=3
 frames=1
 channel_dim=512
-model_path="" # when the path is empty, the model will be loaded from the installed fmpose3d package
-# model_path='./models/model_GAMLP.py' # when the path is not empty, the model will be loaded from the local file path
+model_type='fmpose3d' # use registered model by default
 sh_file='scripts/FMPose3D_train.sh'
 folder_name=FMPose3D_layers${layers}_$(date +%Y%m%d_%H%M%S)
 
@@ -20,6 +19,7 @@ python3 scripts/FMPose3D_main.py \
   --train \
   --dataset h36m \
   --frames ${frames} \
+  --model_type "${model_type}" \
   ${model_path:+--model_path "$model_path"} \
   --gpu ${gpu_id} \
   --batch_size ${batch_size} \

From f878e7f89fcdb5cd56f314cc516cda9f98dad70b Mon Sep 17 00:00:00 2001
From: ti <tiwang.cs@gmail.com>
Date: Tue, 10 Feb 2026 17:38:06 +0100
Subject: [PATCH 04/21] Add model_type argument to opts for model registry
 selection

---
 fmpose3d/animals/common/arguments.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/fmpose3d/animals/common/arguments.py b/fmpose3d/animals/common/arguments.py
index 690c540..f465f17 100755
--- a/fmpose3d/animals/common/arguments.py
+++ b/fmpose3d/animals/common/arguments.py
@@ -68,6 +68,8 @@ def init(self):
         self.parser.add_argument("--model_dir", type=str, default="")
         # Optional: load model class from a specific file path
         self.parser.add_argument("--model_path", type=str, default="")
+        # Model registry name (e.g. "fmpose3d_animals"); used instead of --model_path
+        self.parser.add_argument("--model_type", type=str, default="fmpose3d_animals")
 
         self.parser.add_argument("--post_refine_reload", action="store_true")
         self.parser.add_argument("--checkpoint", type=str, default="")

From 321b28ad3b66fd500d09bfa9e3d4f3dfc03db095 Mon Sep 17 00:00:00 2001
From: ti <tiwang.cs@gmail.com>
Date: Tue, 10 Feb 2026 17:38:17 +0100
Subject: [PATCH 05/21] Remove unnecessary comment block in HRNet
 implementation file

---
 fmpose3d/lib/hrnet/hrnet.py | 2 --
 1 file changed, 2 deletions(-)

diff --git a/fmpose3d/lib/hrnet/hrnet.py b/fmpose3d/lib/hrnet/hrnet.py
index 1368b30..0d0b752 100644
--- a/fmpose3d/lib/hrnet/hrnet.py
+++ b/fmpose3d/lib/hrnet/hrnet.py
@@ -5,9 +5,7 @@
 "FMPose3D: monocular 3D Pose Estimation via Flow Matching"
 by Ti Wang, Xiaohang Yu, and Mackenzie Weygandt Mathis
 Licensed under Apache 2.0
-"""
 
-"""
 FMPose3D – clean HRNet 2D pose estimation API.
 
 Provides :class:`HRNetPose2d`, a self-contained wrapper around the

From ea1d3f70c0db4b6425bdff54bd81815a0a9044fe Mon Sep 17 00:00:00 2001
From: ti <tiwang.cs@gmail.com>
Date: Tue, 10 Feb 2026 17:38:26 +0100
Subject: [PATCH 06/21] Import animal models to ensure their registration in
 the model registry.

---
 fmpose3d/models/__init__.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/fmpose3d/models/__init__.py b/fmpose3d/models/__init__.py
index b9dc64a..dd13e64 100644
--- a/fmpose3d/models/__init__.py
+++ b/fmpose3d/models/__init__.py
@@ -15,6 +15,8 @@
 
 # Import model subpackages so their @register_model decorators execute.
 from .fmpose3d import Graph, Model
+# Import animal models so their @register_model decorators execute.
+from fmpose3d.animals import models as _animal_models  # noqa: F401
 
 __all__ = [
     "BaseModel",

From dd5be5d0769ec45dd585645025f308e090cca873 Mon Sep 17 00:00:00 2001
From: ti <tiwang.cs@gmail.com>
Date: Tue, 10 Feb 2026 17:38:40 +0100
Subject: [PATCH 07/21] Register Model class for animal3D in the model registry
 and update initialization to accept arguments.

---
 fmpose3d/animals/models/model_animal3d.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/fmpose3d/animals/models/model_animal3d.py b/fmpose3d/animals/models/model_animal3d.py
index 273b1dd..7ea8b95 100644
--- a/fmpose3d/animals/models/model_animal3d.py
+++ b/fmpose3d/animals/models/model_animal3d.py
@@ -16,6 +16,7 @@
 from timm.models.layers import DropPath
 
 from fmpose3d.animals.models.graph_frames import Graph
+from fmpose3d.models.base_model import BaseModel, register_model
 
 class TimeEmbedding(nn.Module):
     def __init__(self, dim: int, hidden_dim: int = 64):
@@ -207,9 +208,10 @@ def forward(self, x):
         x = self.fc5(x)
         return x
 
-class Model(nn.Module):
+@register_model("fmpose3d_animals")
+class Model(BaseModel):
     def __init__(self, args):
-        super().__init__()
+        super().__init__(args)
 
         self.graph = Graph('animal3d', 'spatial', pad=1)
         self.register_buffer('A', torch.tensor(self.graph.A, dtype=torch.float32))

From a7e25e9bfd0d951cc88fa5b6796e25f1783fecd2 Mon Sep 17 00:00:00 2001
From: ti <tiwang.cs@gmail.com>
Date: Tue, 10 Feb 2026 17:38:57 +0100
Subject: [PATCH 08/21] Refactor main_animal3d.py to load model from the
 registered model registry using get_model, replacing the previous model_path
 approach.

---
 animals/scripts/main_animal3d.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/animals/scripts/main_animal3d.py b/animals/scripts/main_animal3d.py
index 0436204..cb23f8a 100644
--- a/animals/scripts/main_animal3d.py
+++ b/animals/scripts/main_animal3d.py
@@ -38,8 +38,9 @@
     spec.loader.exec_module(module)
     CFM = getattr(module, "Model")
 else:
-    # Load model from installed fmpose package
-    from fmpose3d.animals.models import Model as CFM
+    # Load model from registered model registry
+    from fmpose3d.models import get_model
+    CFM = get_model(args.model_type)
      
 def train(opt, actions, train_loader, model, optimizer, epoch):
     return step('train', opt, actions, train_loader, model, optimizer, epoch)
@@ -98,7 +99,6 @@ def step(split, args, actions, dataLoader, model, optimizer=None, epoch=None, st
             gt_3D = gt_3D.clone()
             gt_3D[:, :, args.root_joint] = 0
             
-            
             # Conditional Flow Matching training
             # gt_3D, input_2D shape: (B,F,J,C)
             # vis_3D shape: (B,F,J,1) - visibility mask

From f1edf4489acc19b632b97cb92823296aab9c04e1 Mon Sep 17 00:00:00 2001
From: ti <tiwang.cs@gmail.com>
Date: Tue, 10 Feb 2026 17:39:16 +0100
Subject: [PATCH 09/21] Update test_animal3d.sh to modify eval_sample_steps,
 change model_type, and update saved_model_path for consistency with model
 registry usage.

---
 animals/scripts/test_animal3d.sh | 12 +++++++-----
 1 file changed, 7 insertions(+), 5 deletions(-)

diff --git a/animals/scripts/test_animal3d.sh b/animals/scripts/test_animal3d.sh
index 3c85885..ff3fcaa 100644
--- a/animals/scripts/test_animal3d.sh
+++ b/animals/scripts/test_animal3d.sh
@@ -2,7 +2,7 @@ layers=5
 batch_size=13
 lr=1e-3
 gpu_id=0
-eval_sample_steps=3
+eval_sample_steps=5
 num_saved_models=3
 frames=1
 large_decay_epoch=15
@@ -10,9 +10,10 @@ lr_decay_large=0.75
 n_joints=26
 out_joints=26
 epochs=300
-# model_path='models/model_animals.py'
-model_path='./pre_trained_models/animal3d_pretrained_weights/model_animal3d.py' # when the path is empty, the model will be loaded from the installed fmpose package
-saved_model_path='./pre_trained_models/animal3d_pretrained_weights/CFM_154_4403_best.pth'
+model_type='fmpose3d_animals'
+# model_path='' # set to a local file path to override the registry
+saved_model_path='./pre_trained_models/fmpose3d_animals/fmpose3d_animals_pretrained_weights.pth'
+
 # root path denotes the path to the original dataset
 root_path="./dataset/"
 train_dataset_paths=(
@@ -20,7 +21,7 @@ train_dataset_paths=(
   "./dataset/control_animal3dlatest/train.json"
 )
 test_dataset_paths=(
-  "./dataset/control_animal3dlatest/test.json"
+  "./dataset/animal3d/test.json"
 )
 
 folder_name="TestCtrlAni3D_L${layers}_lr${lr}_B${batch_size}_$(date +%Y%m%d_%H%M%S)"
@@ -33,6 +34,7 @@ python ./scripts/main_animal3d.py \
   --test 1 \
   --batch_size ${batch_size} \
   --lr ${lr} \
+  --model_type "${model_type}" \
   ${model_path:+--model_path "$model_path"} \
   --folder_name ${folder_name} \
   --layers ${layers} \

From 51b14ab07b5ad2a3c035f2ae46eb1e97aba49201 Mon Sep 17 00:00:00 2001
From: ti <tiwang.cs@gmail.com>
Date: Tue, 10 Feb 2026 17:42:36 +0100
Subject: [PATCH 10/21] Update train_animal3d.sh to modify eval_sample_steps,
 change model_type for model registry integration, and comment out model_path
 for clarity.

---
 animals/scripts/train_animal3d.sh | 10 ++++------
 1 file changed, 4 insertions(+), 6 deletions(-)

diff --git a/animals/scripts/train_animal3d.sh b/animals/scripts/train_animal3d.sh
index cdd6f8c..bcba1ee 100644
--- a/animals/scripts/train_animal3d.sh
+++ b/animals/scripts/train_animal3d.sh
@@ -2,16 +2,14 @@ layers=5
 batch_size=13
 lr=1e-3
 gpu_id=0
-eval_sample_steps=3
+eval_sample_steps=5
 num_saved_models=3
 frames=1
 large_decay_epoch=15
 lr_decay_large=0.75
-n_joints=26
-out_joints=26
 epochs=300
-# model_path='models/model_animals.py'
-model_path="" # when the path is empty, the model will be loaded from the installed fmpose package
+model_type='fmpose3d_animals'
+# model_path="" # set to a local file path to override the registry
 # root path denotes the path to the original dataset
 root_path="./dataset/"
 train_dataset_paths=(
@@ -32,7 +30,7 @@ python ./scripts/main_animal3d.py \
   --test 1 \
   --batch_size ${batch_size} \
   --lr ${lr} \
-  ${model_path:+--model_path "$model_path"} \
+  --model_type "${model_type}" \
   --folder_name ${folder_name} \
   --layers ${layers} \
   --gpu ${gpu_id} \

From 470248103cd81914f20ebfd2715bf086594f39af Mon Sep 17 00:00:00 2001
From: ti <tiwang.cs@gmail.com>
Date: Tue, 10 Feb 2026 17:53:46 +0100
Subject: [PATCH 11/21] Refactor vis_animals.py to load model using get_model
 from the registered model registry, enhancing integration with the model
 registry.

---
 animals/demo/vis_animals.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/animals/demo/vis_animals.py b/animals/demo/vis_animals.py
index c9fe438..459c2bc 100644
--- a/animals/demo/vis_animals.py
+++ b/animals/demo/vis_animals.py
@@ -42,8 +42,9 @@
     spec.loader.exec_module(module)
     CFM = getattr(module, "Model")
 else:
-    # Load model from installed fmpose package
-    from fmpose3d.models import Model as CFM
+    # Load model from registered model registry
+    from fmpose3d.models import get_model
+    CFM = get_model(args.model_type)
 
 from deeplabcut.pose_estimation_pytorch.apis import superanimal_analyze_images
 

From b02dedac5a6deab108f39cae24cb34ca9e826c7d Mon Sep 17 00:00:00 2001
From: ti <tiwang.cs@gmail.com>
Date: Tue, 10 Feb 2026 17:54:11 +0100
Subject: [PATCH 12/21] Update vis_animals.sh to set model_type for FMPose3D
 and adjust saved_model_path for consistency with the new model registry
 structure.

---
 animals/demo/vis_animals.sh | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/animals/demo/vis_animals.sh b/animals/demo/vis_animals.sh
index 3879391..e2944c2 100644
--- a/animals/demo/vis_animals.sh
+++ b/animals/demo/vis_animals.sh
@@ -7,8 +7,9 @@ sh_file='vis_animals.sh'
 # n_joints=26
 # out_joints=26
 
-model_path='../pre_trained_models/animal3d_pretrained_weights/model_animal3d.py'
-saved_model_path='../pre_trained_models/animal3d_pretrained_weights/CFM_154_4403_best.pth'
+model_type='fmpose3d_animals'
+# model_path=''  # set to a local file path to override the registry
+saved_model_path='../pre_trained_models/fmpose3d_animals/fmpose3d_animals_pretrained_weights.pth'
 
 # path='./images/image_00068.jpg'  # single image
 input_images_folder='./images/'  # folder containing multiple images
@@ -17,7 +18,8 @@ python3 vis_animals.py \
  --type 'image' \
  --path ${input_images_folder} \
  --saved_model_path "${saved_model_path}" \
- --model_path "${model_path}" \
+ ${model_path:+--model_path "$model_path"} \
+ --model_type "${model_type}" \
  --sample_steps ${sample_steps} \
  --batch_size ${batch_size} \
  --layers ${layers} \

From 16f340c6e8c8ff378123b00f4d1594b9ec246e0a Mon Sep 17 00:00:00 2001
From: ti <tiwang.cs@gmail.com>
Date: Tue, 10 Feb 2026 17:58:47 +0100
Subject: [PATCH 13/21] Update README.md with new download link for pre-trained
 model

---
 animals/README.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/animals/README.md b/animals/README.md
index af9da69..5121533 100644
--- a/animals/README.md
+++ b/animals/README.md
@@ -10,7 +10,7 @@ In this part, the FMPose3D model is trained on [Animal3D](https://xujiacong.gith
 This visualization script is designed for single-frame based model, allowing you to easily run 3D animal pose estimation on any single image.
 
 Before testing, make sure you have the pre-trained model ready.
-You may either use the model trained by your own or download ours from [here](https://drive.google.com/drive/folders/1fMKVaYziwFkAnFrtQZmoPOTfe7Hkl2at?usp=sharing) and place it in the `./pre_trained_models` directory.
+You may either use the model trained by your own or download ours from [here](https://drive.google.com/drive/folders/1kL4aOyWNq0o9zB0rSTRM8KYgkySVmUTk?usp=drive_link) and place it in the `./pre_trained_models` directory.
 
 Next, put your test images into folder `demo/images`. Then run the visualization script:
 ```bash

From 0512c2760de5ee274dbb296184e48dd94f30a946 Mon Sep 17 00:00:00 2001
From: ti <tiwang.cs@gmail.com>
Date: Wed, 11 Feb 2026 10:16:12 +0100
Subject: [PATCH 14/21] Refactor backup file handling in main_animal3d.py to
 enable file copying for checkpoints, including model and script files,
 improving data management during training.

---
 animals/scripts/main_animal3d.py | 27 ++++++++++++---------------
 1 file changed, 12 insertions(+), 15 deletions(-)

diff --git a/animals/scripts/main_animal3d.py b/animals/scripts/main_animal3d.py
index cb23f8a..c90bdea 100644
--- a/animals/scripts/main_animal3d.py
+++ b/animals/scripts/main_animal3d.py
@@ -217,21 +217,18 @@ def get_parameter_number(net):
             os.makedirs(args.checkpoint)
 
         # backup files
-        # import shutil
-        # file_path = os.path.abspath(__file__)
-        # file_name = os.path.basename(file_path)
-        # shutil.copyfile(src=file_path, dst=os.path.join(args.checkpoint, args.create_time + "_" + file_name))
-        # shutil.copyfile(src=os.path.abspath("common/arguments.py"), dst=os.path.join(args.checkpoint, args.create_time + "_arguments.py"))
-        # # backup the selected model file (from --model_path if provided)
-        # if getattr(args, 'model_path', ''):
-        #     model_src_path = os.path.abspath(args.model_path)
-        #     model_dst_name = f"{args.create_time}_" + os.path.basename(model_src_path)
-        #     shutil.copyfile(src=model_src_path, dst=os.path.join(args.checkpoint, model_dst_name))
-        # # shutil.copyfile(src="common/utils.py", dst = os.path.join(args.checkpoint, args.create_time + "_utils.py"))
-        # sh_base = os.path.basename(args.sh_file)
-        # dst_name = f"{args.create_time}_" + sh_base
-        # sh_src = os.path.abspath(args.sh_file)
-        # shutil.copyfile(src=sh_src, dst=os.path.join(args.checkpoint, dst_name))
+        import shutil
+        file_path = os.path.abspath(__file__)
+        file_name = os.path.basename(file_path)
+        shutil.copyfile(src=file_path, dst=os.path.join(args.checkpoint, args.create_time + "_" + file_name))
+        if getattr(args, 'model_path', ''):
+            model_src_path = os.path.abspath(args.model_path)
+            model_dst_name = f"{args.create_time}_" + os.path.basename(model_src_path)
+            shutil.copyfile(src=model_src_path, dst=os.path.join(args.checkpoint, model_dst_name))
+        sh_base = os.path.basename(args.sh_file)
+        dst_name = f"{args.create_time}_" + sh_base
+        sh_src = os.path.abspath(args.sh_file)
+        shutil.copyfile(src=sh_src, dst=os.path.join(args.checkpoint, dst_name))
 
         logging.basicConfig(format='%(asctime)s %(message)s', datefmt='%Y/%m/%d %H:%M:%S', \
             filename=os.path.join(args.checkpoint, 'train.log'), level=logging.INFO)

From a440a085b9044dff45d40d632e51f4d069c45e16 Mon Sep 17 00:00:00 2001
From: ti <tiwang.cs@gmail.com>
Date: Wed, 11 Feb 2026 10:16:23 +0100
Subject: [PATCH 15/21] Update test_animal3d.sh to change test dataset path and
 adjust script file reference for improved clarity and consistency.

---
 animals/scripts/test_animal3d.sh | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/animals/scripts/test_animal3d.sh b/animals/scripts/test_animal3d.sh
index ff3fcaa..207e332 100644
--- a/animals/scripts/test_animal3d.sh
+++ b/animals/scripts/test_animal3d.sh
@@ -21,11 +21,11 @@ train_dataset_paths=(
   "./dataset/control_animal3dlatest/train.json"
 )
 test_dataset_paths=(
-  "./dataset/animal3d/test.json"
+  "./dataset/control_animal3dlatest/test.json"
 )
 
 folder_name="TestCtrlAni3D_L${layers}_lr${lr}_B${batch_size}_$(date +%Y%m%d_%H%M%S)"
-sh_file='scripts/animals/test_animal3d.sh'
+sh_file='scripts/test_animal3d.sh'
 
 python ./scripts/main_animal3d.py \
   --root_path ${root_path} \

From 1f7bfa6891b12e83fb1215281dbdc493583069a7 Mon Sep 17 00:00:00 2001
From: Jaap de Ruyter <deruyter92@gmail.com>
Date: Wed, 11 Feb 2026 11:34:59 +0100
Subject: [PATCH 16/21] Add default FMPose3DConfig per model_type

---
 fmpose3d/common/config.py | 47 ++++++++++++++++++++++++++++++---------
 1 file changed, 37 insertions(+), 10 deletions(-)

diff --git a/fmpose3d/common/config.py b/fmpose3d/common/config.py
index b2980e1..5e2b503 100644
--- a/fmpose3d/common/config.py
+++ b/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

From 60557863b8c40924933de1d4d08571419b50bf67 Mon Sep 17 00:00:00 2001
From: Jaap de Ruyter <deruyter92@gmail.com>
Date: Wed, 11 Feb 2026 12:02:54 +0100
Subject: [PATCH 17/21] Update FMPose3DConfig and add SuperAnimalConfig

---
 fmpose3d/common/config.py | 44 +++++++++++++++++++++++++++++++++++++--
 1 file changed, 42 insertions(+), 2 deletions(-)

diff --git a/fmpose3d/common/config.py b/fmpose3d/common/config.py
index 5e2b503..39cbee5 100644
--- a/fmpose3d/common/config.py
+++ b/fmpose3d/common/config.py
@@ -31,12 +31,18 @@ class ModelConfig:
         "out_joints": 17,
         "dataset": "h36m",
         "sample_steps": 3,
+        "joints_left": [4, 5, 6, 11, 12, 13],
+        "joints_right": [1, 2, 3, 14, 15, 16],
+        "root_joint": 0,
     },
     "fmpose3d_animals": {
         "n_joints": 26,
         "out_joints": 26,
         "dataset": "animal3d",
         "sample_steps": 3,
+        "joints_left": [0, 3, 5, 8, 10, 12, 14, 16, 20, 22],
+        "joints_right": [1, 4, 6, 9, 11, 13, 15, 17, 21, 23],
+        "root_joint": 7,
     },
 }
 
@@ -53,6 +59,9 @@ class FMPose3DConfig(ModelConfig):
     token_dim: int = 256
     n_joints: int = INFER_FROM_MODEL_TYPE  # type: ignore[assignment]
     out_joints: int = INFER_FROM_MODEL_TYPE  # type: ignore[assignment]
+    joints_left: List[int] = INFER_FROM_MODEL_TYPE  # type: ignore[assignment]
+    joints_right: List[int] = INFER_FROM_MODEL_TYPE  # type: ignore[assignment]
+    root_joint: int = INFER_FROM_MODEL_TYPE  # type: ignore[assignment]
     in_channels: int = 2
     out_channels: int = 3
     frames: int = 1
@@ -207,6 +216,33 @@ class HRNetConfig(Pose2DConfig):
     hrnet_weights_path: str = ""
 
 
+@dataclass
+class SuperAnimalConfig(Pose2DConfig):
+    """DeepLabCut SuperAnimal 2D pose detector configuration.
+
+    Uses the DeepLabCut ``superanimal_analyze_images`` API to detect
+    animal keypoints in the quadruped80K format, then maps them to the
+    Animal3D 26-keypoint layout expected by the ``fmpose3d_animals``
+    3D lifter.
+
+    Attributes
+    ----------
+    superanimal_name : str
+        Name of the SuperAnimal model (default ``"superanimal_quadruped"``).
+    sa_model_name : str
+        Backbone architecture (default ``"hrnet_w32"``).
+    detector_name : str
+        Object detector used for animal bounding boxes.
+    max_individuals : int
+        Maximum number of individuals to detect per image (default 1).
+    """
+    pose2d_model: str = "superanimal"
+    superanimal_name: str = "superanimal_quadruped"
+    sa_model_name: str = "hrnet_w32"
+    detector_name: str = "fasterrcnn_resnet50_fpn_v2"
+    max_individuals: int = 1
+
+
 @dataclass
 class DemoConfig:
     """Demo / inference configuration."""
@@ -287,8 +323,12 @@ def _pick(dc_class, src: dict):
         for group_name, dc_class in _SUB_CONFIG_CLASSES.items():
             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
+            elif group_name == "pose2d_cfg":
+                p2d = raw.get("pose2d_model", "hrnet")
+                if p2d == "superanimal":
+                    dc_class = SuperAnimalConfig
+                elif p2d == "hrnet":
+                    dc_class = HRNetConfig
             kwargs[group_name] = _pick(dc_class, raw)
         return cls(**kwargs)
 

From f90ce602daac4c3ed2f7d63eff52fc1c0bb81d06 Mon Sep 17 00:00:00 2001
From: Jaap de Ruyter <deruyter92@gmail.com>
Date: Wed, 11 Feb 2026 11:35:25 +0100
Subject: [PATCH 18/21] expose model registry in main package

---
 fmpose3d/__init__.py | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/fmpose3d/__init__.py b/fmpose3d/__init__.py
index 563a140..7f23320 100644
--- a/fmpose3d/__init__.py
+++ b/fmpose3d/__init__.py
@@ -36,6 +36,9 @@
     Source,
 )
 
+# Model registry
+from .models import BaseModel, register_model, get_model, list_models
+
 # Import 2D pose detection utilities
 from .lib.hrnet.gen_kpts import gen_video_kpts
 from .lib.hrnet.hrnet import HRNetPose2d
@@ -59,6 +62,11 @@
     "average_aggregation",
     "aggregation_select_single_best_hypothesis_by_2D_error",
     "aggregation_RPEA_joint_level",
+    # Model registry
+    "BaseModel",
+    "register_model",
+    "get_model",
+    "list_models",
     # 2D pose detection
     "HRNetPose2d",
     "gen_video_kpts",

From 80b2420bf7473ccb27c5f907db58feb305cc6507 Mon Sep 17 00:00:00 2001
From: Jaap de Ruyter <deruyter92@gmail.com>
Date: Wed, 11 Feb 2026 12:04:05 +0100
Subject: [PATCH 19/21] update .gitignore: ignore predictions

---
 .gitignore | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/.gitignore b/.gitignore
index 90439ac..3f7f92b 100644
--- a/.gitignore
+++ b/.gitignore
@@ -49,4 +49,5 @@ htmlcov/
 # Excluded directories
 pre_trained_models/
 demo/predictions/
-demo/images/
\ No newline at end of file
+demo/images/
+**/predictions/
\ No newline at end of file

From 8f278ec5c8269b0d8d0bf8d1973caee4aebf904e Mon Sep 17 00:00:00 2001
From: ti <tiwang.cs@gmail.com>
Date: Wed, 11 Feb 2026 12:55:12 +0100
Subject: [PATCH 20/21] Update model_type to "fmpose3d_humans" across
 configuration and model files for consistency in the FMPose3D framework.

---
 fmpose3d/common/arguments.py            |  4 ++--
 fmpose3d/common/config.py               | 10 +++++-----
 fmpose3d/models/fmpose3d/model_GAMLP.py |  2 +-
 tests/test_config.py                    |  2 +-
 4 files changed, 9 insertions(+), 9 deletions(-)

diff --git a/fmpose3d/common/arguments.py b/fmpose3d/common/arguments.py
index b94db98..99d3d10 100755
--- a/fmpose3d/common/arguments.py
+++ b/fmpose3d/common/arguments.py
@@ -74,8 +74,8 @@ def init(self):
         self.parser.add_argument("--model_dir", type=str, default="")
         # Optional: load model class from a specific file path
         self.parser.add_argument("--model_path", type=str, default="")
-        # Model registry name (e.g. "fmpose3d"); used instead of --model_path
-        self.parser.add_argument("--model_type", type=str, default="fmpose3d")
+        # Model registry name (e.g. "fmpose3d_humans"); used instead of --model_path
+        self.parser.add_argument("--model_type", type=str, default="fmpose3d_humans")
         self.parser.add_argument("--model_weights_path", type=str, default="")
 
         self.parser.add_argument("--post_refine_reload", action="store_true")
diff --git a/fmpose3d/common/config.py b/fmpose3d/common/config.py
index 39cbee5..bddd739 100644
--- a/fmpose3d/common/config.py
+++ b/fmpose3d/common/config.py
@@ -19,14 +19,14 @@
 @dataclass
 class ModelConfig:
     """Model architecture configuration."""
-    model_type: str = "fmpose3d"
+    model_type: str = "fmpose3d_humans"
 
 
 # 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": {
+    "fmpose3d_humans": {
         "n_joints": 17,
         "out_joints": 17,
         "dataset": "h36m",
@@ -39,7 +39,7 @@ class ModelConfig:
         "n_joints": 26,
         "out_joints": 26,
         "dataset": "animal3d",
-        "sample_steps": 3,
+        "sample_steps": 5,
         "joints_left": [0, 3, 5, 8, 10, 12, 14, 16, 20, 22],
         "joints_right": [1, 4, 6, 9, 11, 13, 15, 17, 21, 23],
         "root_joint": 7,
@@ -51,7 +51,7 @@ class ModelConfig:
 
 @dataclass
 class FMPose3DConfig(ModelConfig):
-    model_type: str = "fmpose3d"
+    model_type: str = "fmpose3d_humans"
     model: str = ""
     layers: int = 5
     channel: int = 512
@@ -321,7 +321,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') in _FMPOSE3D_DEFAULTS:
+            if group_name == "model_cfg" and raw.get("model_type", 'fmpose3d_humans') in _FMPOSE3D_DEFAULTS:
                 dc_class = FMPose3DConfig
             elif group_name == "pose2d_cfg":
                 p2d = raw.get("pose2d_model", "hrnet")
diff --git a/fmpose3d/models/fmpose3d/model_GAMLP.py b/fmpose3d/models/fmpose3d/model_GAMLP.py
index 6657900..7176e26 100644
--- a/fmpose3d/models/fmpose3d/model_GAMLP.py
+++ b/fmpose3d/models/fmpose3d/model_GAMLP.py
@@ -212,7 +212,7 @@ def forward(self, x):
         x = self.fc2(x)
         return x
 
-@register_model("fmpose3d")
+@register_model("fmpose3d_humans")
 class Model(BaseModel):
     def __init__(self, args):
         super().__init__(args)
diff --git a/tests/test_config.py b/tests/test_config.py
index 0919b4e..e6cb9c4 100644
--- a/tests/test_config.py
+++ b/tests/test_config.py
@@ -250,7 +250,7 @@ def test_from_namespace_basic(self):
         ns = argparse.Namespace(
             # FMPose3DConfig
             model="test_model",
-            model_type="fmpose3d",
+            model_type="fmpose3d_humans",
             layers=5,
             channel=256,
             d_hid=512,

From 0138da635ac1d03106c497202116b9f6ef753de8 Mon Sep 17 00:00:00 2001
From: ti <tiwang.cs@gmail.com>
Date: Wed, 11 Feb 2026 12:55:31 +0100
Subject: [PATCH 21/21] Update model_type to "fmpose3d_humans" in demo and
 script files for consistency across the FMPose3D framework.

---
 demo/vis_in_the_wild.sh   | 2 +-
 scripts/FMPose3D_test.sh  | 2 +-
 scripts/FMPose3D_train.sh | 2 +-
 3 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/demo/vis_in_the_wild.sh b/demo/vis_in_the_wild.sh
index a6f3a4d..df88474 100755
--- a/demo/vis_in_the_wild.sh
+++ b/demo/vis_in_the_wild.sh
@@ -5,7 +5,7 @@ sample_steps=3
 batch_size=1
 sh_file='vis_in_the_wild.sh'
 
-model_type='fmpose3d'
+model_type='fmpose3d_humans'
 model_weights_path='../pre_trained_models/fmpose3d_h36m/FMpose3D_pretrained_weights.pth'
 
 target_path='./images/'  # folder containing multiple images
diff --git a/scripts/FMPose3D_test.sh b/scripts/FMPose3D_test.sh
index b83d65d..afc1f86 100755
--- a/scripts/FMPose3D_test.sh
+++ b/scripts/FMPose3D_test.sh
@@ -10,7 +10,7 @@ mode='exp'
 exp_temp=0.005
 folder_name=test_s${eval_multi_steps}_${mode}_h${num_hypothesis_list}_$(date +%Y%m%d_%H%M%S)
 
-model_type='fmpose3d'
+model_type='fmpose3d_humans'
 model_weights_path='./pre_trained_models/fmpose3d_h36m/FMpose3D_pretrained_weights.pth'
 
 #Test
diff --git a/scripts/FMPose3D_train.sh b/scripts/FMPose3D_train.sh
index a677d9c..e435482 100755
--- a/scripts/FMPose3D_train.sh
+++ b/scripts/FMPose3D_train.sh
@@ -11,7 +11,7 @@ epochs=80
 num_saved_models=3
 frames=1
 channel_dim=512
-model_type='fmpose3d' # use registered model by default
+model_type='fmpose3d_humans' # use registered model by default
 sh_file='scripts/FMPose3D_train.sh'
 folder_name=FMPose3D_layers${layers}_$(date +%Y%m%d_%H%M%S)