Merge branch 'jaap/add_config_and_registry' of github.com:deruyter92/FMPose3D into jaap/add_config_and_registry

Author: deruyter92

README.md

@@ -6,7 +6,7 @@
 
 This is the official implementation of the approach described in the preprint:
 
-[**FMPose3D: monocular 3D pose estimation via flow matching**](http://arxiv.org/abs/2602.05755)            
+[**FMPose3D: monocular 3D pose estimation via flow matching**](https://arxiv.org/abs/2602.05755)            
 Ti Wang, Xiaohang Yu, Mackenzie Weygandt Mathis
 
 <!-- <p align="center"><img src="./images/Frame 4.jpg" width="50%" alt="" /></p> -->

animals/demo/vis_animals.py

@@ -339,7 +339,7 @@ def get_pose3D(path, output_dir, type='image'):
     model_dict = model['CFM'].state_dict()
     model_path = args.saved_model_path
     print(f"Loading model from: {model_path}")
-    pre_dict = torch.load(model_path, weights_only=True)
+    pre_dict = torch.load(model_path, map_location=device, weights_only=True)
     for name, key in model_dict.items():
         model_dict[name] = pre_dict[name]
     model['CFM'].load_state_dict(model_dict)

animals/scripts/main_animal3d.py

@@ -274,7 +274,7 @@ def get_parameter_number(net):
         # Prefer explicit saved_model_path; otherwise fallback to previous_dir glob
         model_path = args.saved_model_path
         print(model_path)
-        pre_dict = torch.load(model_path, weights_only=True)
+        pre_dict = torch.load(model_path, weights_only=True, map_location=device)
         for name, key in model_dict.items():
             model_dict[name] = pre_dict[name]
         model['CFM'].load_state_dict(model_dict)

demo/vis_in_the_wild.py

@@ -280,7 +280,7 @@ def get_pose3D(path, output_dir, type='image'):
     model_dict = model['CFM'].state_dict()
     model_path = args.model_weights_path
     print(model_path)
-    pre_dict = torch.load(model_path, weights_only=True)
+    pre_dict = torch.load(model_path, map_location=device, weights_only=True)
     for name, key in model_dict.items():
         model_dict[name] = pre_dict[name]
     model['CFM'].load_state_dict(model_dict)

scripts/FMPose3D_main.py

@@ -344,7 +344,7 @@ def print_error_action(action_error_sum, is_train):
         model_dict = model["CFM"].state_dict()
         model_path = args.model_weights_path
         print(model_path)
-        pre_dict = torch.load(model_path, weights_only=True)
+        pre_dict = torch.load(model_path, map_location=device, weights_only=True)
         for name, key in model_dict.items():
             model_dict[name] = pre_dict[name]
         model["CFM"].load_state_dict(model_dict)