Merge initial mouse work. (#57)

* Working.

* Include ROOT_KEYPOINT PARAM

* Pre merge. (Working on viz improvements)

* Model generalization.

* Fit not working - killed.

* Still broken.

* Configs to try out. Rodent works, mouse doesn't

* Reverted root_kp_idx -- still broken for mouse, but works for rodent.

* Updated visuals includes new cam

* Notebooks.

* Viz changes (good).

* No parts.

* Make mouse work on workstation.

* Fix debug output.

* Post merge, not working.

* Constant  folding bug. (Doesn't fail, but veerry slow.)

* Remove root_optimization_keypoint.

* Working w/ hard coded. Clean up.

* More clean up.

* More clean up.

* More cleanup.

* CI changes.

* Black format.

* Add unit test.

* Build fixes.

* Fix vis.

* Revert

* Missed file.

* Misseed configs.

* Remove debug variable.

* Clean up plane estimation.

* Remove unused vars from configs.

* Address review comments.

* black reformat.

* Reinstate free joint check.

* Add free joint check to transform()

configs/mouse.yaml

@@ -0,0 +1,246 @@
+MJCF_PATH: "models/mouse_with_meshes.xml"
+
+# Frames per clip for transform.
+N_FRAMES_PER_CLIP: 360 # 1000 environments
+
+# Tolerance for the optimizations of the full model, limb, and root.
+# TODO: Re-implement optimizer loops to use these tolerances
+FTOL: 1.0e-04
+ROOT_FTOL: 1.0e-05
+LIMB_FTOL: 1.0e-06
+
+# Number of alternating pose and offset optimization rounds. 
+N_ITERS: 6
+
+KP_NAMES:
+  - Nose
+  - Ear_R
+  - Ear_L
+  - TTI
+  - TailTip
+  - Head
+  - Trunk
+  - Tail_0
+  - Tail_1
+  - Tail_2
+  - Shoulder_L
+  - Shoulder_R
+  - Knee_L
+  - Knee_R
+  - Neck
+  - Elbow_L
+  - Elbow_R
+  - Wrist_L
+  - Wrist_R
+  - Forepaw_L
+  - Forepaw_R
+  - Heel_L
+  - Heel_R
+  - Spine
+  - Eye_L
+  - Eye_R
+  - Tail_3
+  - Tail_4
+  - Hip_L
+  - Hip_R
+  - Lisfranc_R
+  - MTP_L
+  - Lisfranc_L
+  - MTP_R
+
+# Ambiguity, are keypoints joints or bones? (Nose is not a joint, but MTP is a joint.)
+KEYPOINT_MODEL_PAIRS:
+  Nose: mouse__Head
+  Ear_R: mouse__Head
+  Ear_L: mouse__Head
+  TTI: mouse__Tail1
+  Head: mouse__Head
+  Trunk: mouse__Thoracic1
+  Tail_0: mouse__Tail1
+  Tail_1: mouse__Tail2
+  Tail_2: mouse__Tail7
+  Tail_3: mouse__Tail13  
+  Tail_4: mouse__Tail16
+  TailTip: mouse__Tail20
+  Shoulder_L: mouse__LScapula
+  Shoulder_R: mouse__RScapula
+  Knee_L: mouse__LFemur
+  Knee_R: mouse__RFemur
+  Neck: mouse__Cervical8
+  Elbow_L: mouse__LHumerus
+  Elbow_R: mouse__RHumerus
+  Wrist_L: mouse__LMetaCarpus3
+  Wrist_R:  mouse__RMetaCarpus3
+  Forepaw_L: mouse__LMetaCarpus3
+  Forepaw_R: mouse__RMetaCarpus3
+  Heel_L: mouse__LTibia
+  Heel_R: mouse__RTibia
+  Spine: mouse__Thoracic13
+  Eye_L: mouse__Head
+  Eye_R: mouse__Head
+  Hip_L: mouse__LFemur_interim_0
+  Hip_R: mouse__RFemur_interim_0
+  Lisfranc_L: mouse__LBPhalange3
+  Lisfranc_R: mouse__RBPhalange3
+  MTP_L: mouse__LTarsus_interim_0
+  MTP_R: mouse__RTarsus_interim_0
+
+# The initial offsets for each keypoint in meters.
+KEYPOINT_INITIAL_OFFSETS:
+  Nose: 0.0 0.0 0.0
+  Ear_R: 0.0 0.0 0.0
+  Ear_L: 0.0 0.0 0.0
+  TTI: 0.0 0.0 0.0
+  TailTip: 0.0 0.0 0.0
+  Head: 0.0 0.0 0.0
+  Trunk: 0.0 0.0 0.0
+  Tail_0: 0.0 0.0 0.0
+  Tail_1: 0.0 0.0 0.0
+  Tail_2: 0.0 0.0 0.0
+  Shoulder_L: 0.0 0.0 0.0
+  Shoulder_R: 0.0 0.0 0.0
+  Knee_L: 0.0 0.0 0.0
+  Knee_R: 0.0 0.0 0.0
+  Neck: 0.0 0.0 0.0
+  Elbow_L: 0.0 0.0 0.0
+  Elbow_R: 0.0 0.0 0.0
+  Wrist_L: 0.0 0.0 0.0
+  Wrist_R: 0.0 0.0 0.0
+  Forepaw_L: 0.0 0.0 0.0
+  Forepaw_R: 0.0 0.0 0.0
+  Heel_L: 0.0 0.0 0.0
+  Heel_R: 0.0 0.0 0.0
+  Spine: 0.0 0.0 0.0
+  Eye_L: 0.0 0.0 0.0
+  Eye_R: 0.0 0.0 0.0
+  Tail_3: 0.0 0.0 0.0
+  Tail_4: 0.0 0.0 0.0
+  Hip_L: 0.0 0.0 0.0
+  Hip_R: 0.0 0.0 0.0
+  Lisfranc_R: 0.0 0.0 0.0
+  MTP_L: 0.0 0.0 0.0
+  Lisfranc_L: 0.0 0.0 0.0
+  MTP_R: 0.0 0.0 0.0
+
+TRUNK_OPTIMIZATION_KEYPOINTS:
+ - "Trunk"
+ - "HipL"
+ - "HipR"
+ - "ShoulderL"
+ - "ShoulderR"
+
+INDIVIDUAL_PART_OPTIMIZATION:
+#   "r_leg": [
+#     "mouse__Thoracic1",
+#     "mouse__Lumbar1",
+#     "mouse__Lumbar2",
+#     "mouse__Lumbar3",
+#     "mouse__Lumbar4",
+#     "mouse__Lumbar5",
+#     "mouse__RFemur",
+#     "mouse__RTibia",
+#     "mouse__RTarsus",
+#     "mouse__RBPhalange3",
+#     "mouse__RMetatarsus1",
+#     "mouse__RTarsus_interim_0"
+#   ]
+#   "l_leg": [
+#     "mouse__Thoracic1",
+#     "mouse__Lumbar1",
+#     "mouse__Lumbar2",
+#     "mouse__Lumbar3",
+#     "mouse__Lumbar4",
+#     "mouse__Lumbar5",
+#     "mouse__LFemur",
+#     "mouse__LTibia",
+#     "mouse__LTarsus",
+#     "mouse__RBPhalange3",
+#     "mouse__LMetatarsus1",
+#     "mouse__RTarsus_interim_0"
+#   ]
+#   "r_arm": [
+#     "mouse__RScapula",
+#     "mouse__RHumerus",
+#     "mouse__RRadius",
+#     "mouse__RCarpus",
+#     "mouse__RMetaCarpus1"
+#   ]
+#   "l_arm": [
+#     "mouse__LScapula",
+#     "mouse__LHumerus",
+#     "mouse__LRadius",
+#     "mouse__LCarpus",
+#     "mouse__LMetaCarpus1"
+#   ]
+#   "head": [
+#     "mouse__Cervical1",
+#     "mouse__Cervical2",
+#     "mouse__Cervical3",
+#     "mouse__Cervical4",
+#     "mouse__Cervical5",
+#     "mouse__Cervical6"
+#   ]
+
+# Color to use for each keypoint when visualizing the results
+KEYPOINT_COLOR_PAIRS:
+  Nose: 1 0 0 1  # Red
+  Ear_L: 0 1 0 1  # Green
+  Ear_R: 0 0.8 0 1  # Darker Green
+  TTI: 0 0 1 1  # Blue
+  TailTip: 0.58 0 0.82 1  # Purple
+  Head: 1 0.5 0 1  # Orange
+  Trunk: 0.5 0 0.5 1  # Plum
+  Tail_0: 1 1 0 1  # Yellow
+  Tail_1: 0.8 0.8 0 1  # Darker Yellow
+  Tail_2: 0.4 0.4 0 1  # Olive
+  Tail_3: 0.58 0 0.82 1  # Purple
+  Tail_4: 0.58 0 0.82 1  # Purple
+  Shoulder_L: 0.27 0.5 0.7 1  # Cerulean
+  Shoulder_R: 0.2 0.4 0.6 1  # Darker Cerulean
+  HaunchL: 0.85 0.75 0.85 1  # Thistle
+  HaunchR: 0.8 0.6 0.8 1  # Darker Thistle
+  Knee_L: 0.27 0.5 0.7 1  # Cerulean
+  Knee_R: 0.2 0.4 0.6 1  # Darker Cerulean
+  Neck: 0.5 0 0.5 1  # Plum
+  Elbow_L: 0 0.8 0 1  # Darker Green
+  Elbow_R: 0 1 0 1  # Green
+  Wrist_L: 0.4 0.4 0 1  # Olive
+  Wrist_R: 0.8 0.8 0 1  # Darker Yellow
+  Forepaw_L: 1 1 0 1  # Yellow
+  Forepaw_R: 1 0.5 0 1  # Orange
+  Heel_L: 0.85 0.75 0.85 1  # Thistle
+  Heel_R: 0.85 0.75 0.85 1  # Thistle
+  Spine: 0 0 1 1  # Blue
+  Eye_L: 1 0 0 1  # Red
+  Eye_R: 1 0 0 1  # Red
+  Hip_L: 0.85 0.75 0.85 1  # Thistle
+  Hip_R: 0.85 0.75 0.85 1  # Thistle
+  Lisfranc_R: 0 1 0 1  # Green
+  MTP_L: 1 0 0 1  # Red
+  Lisfranc_L: 0 0.8 0 1  # Darker Green
+  MTP_R: 0.2 0.4 0.6 1  # Darker Cerulean
+
+# What is the size of the animal you'd like to register, relative to the model?
+SCALE_FACTOR: 1.0
+
+# Multiplier to put the mocap data into the same scale as the data. Eg, if the
+# mocap data is known to be in millimeters and the model is in meters, this is
+# .001
+MOCAP_SCALE_FACTOR: .001
+
+SITES_TO_REGULARIZE:
+   - ShoulderL
+   - ShoulderR
+#   - SpineF
+#   - SpineL
+#   - SpineM
+#   - WristL
+#   - WristR
+
+RENDER_FPS: 30
+
+N_SAMPLE_FRAMES: 50
+
+# If you have reason to believe that the initial offsets are correct for particular keypoints, 
+# you can regularize those sites using _SITES_TO_REGULARIZE. 
+M_REG_COEF: 0.0

configs/rodent.yaml

@@ -194,5 +194,3 @@ N_SAMPLE_FRAMES: 100
 # If you have reason to believe that the initial offsets are correct for particular keypoints, 
 # you can regularize those sites using _SITES_TO_REGULARIZE. 
 M_REG_COEF: 1
-
-TIME_BINS: 0.02

configs/stac_mouse.yaml

@@ -0,0 +1,12 @@
+fit_path: "fit_mouse.p"
+transform_path: "transform_mouse.p"
+data_path: "tests/data/test_mouse_mocap_3600_frames.h5"
+
+n_fit_frames: 3600
+skip_fit: False
+skip_transform: True
+
+mujoco:
+  solver: "newton"
+  iterations: 1
+  ls_iterations: 4