Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merge initial mouse work. #57

Merged
merged 41 commits into from
Sep 14, 2024
Merged
Show file tree
Hide file tree
Changes from 39 commits
Commits
Show all changes
41 commits
Select commit Hold shift + click to select a range
e0a9cc1
Working.
jf514 Aug 28, 2024
37134fe
Merge branch 'main' into mouse-stac
jf514 Aug 28, 2024
73606d4
Include ROOT_KEYPOINT PARAM
jf514 Aug 29, 2024
bea03bd
Pre merge. (Working on viz improvements)
jf514 Aug 30, 2024
25aad4a
Model generalization.
jf514 Sep 3, 2024
ffa0ef8
Fit not working - killed.
jf514 Sep 4, 2024
7fe7765
Still broken.
jf514 Sep 4, 2024
1644c76
Configs to try out. Rodent works, mouse doesn't
jf514 Sep 4, 2024
bc62ada
Reverted root_kp_idx -- still broken for mouse, but works for rodent.
jf514 Sep 4, 2024
a024330
Updated visuals includes new cam
jf514 Sep 5, 2024
b0c1d7b
Notebooks.
jf514 Sep 5, 2024
6e2c351
Merge branch 'mouse-stac' of github.com:talmolab/stac-mjx into mouse-…
jf514 Sep 5, 2024
9a21244
Merge branch 'mouse-stac' of github.com:talmolab/stac-mjx into mouse-…
jf514 Sep 5, 2024
a0cda4b
Viz changes (good).
jf514 Sep 5, 2024
1fbeb44
Merge remote-tracking branch 'origin/mouse-stac-still-hard-coded' int…
jf514 Sep 5, 2024
3258259
No parts.
jf514 Sep 11, 2024
3285342
Make mouse work on workstation.
jf514 Sep 11, 2024
3abb5af
Merge branch 'main' into mouse-stac-still-hard-coded
jf514 Sep 11, 2024
b30a7e2
Fix debug output.
jf514 Sep 11, 2024
535dce9
Post merge, not working.
jf514 Sep 11, 2024
029d518
Constant folding bug. (Doesn't fail, but veerry slow.)
jf514 Sep 11, 2024
ef8477f
Remove root_optimization_keypoint.
jf514 Sep 12, 2024
633ce19
Working w/ hard coded. Clean up.
jf514 Sep 12, 2024
4dafd11
More clean up.
jf514 Sep 12, 2024
f26bfb6
More clean up.
jf514 Sep 12, 2024
15af551
More cleanup.
jf514 Sep 12, 2024
4d2849c
CI changes.
jf514 Sep 12, 2024
84ee264
Black format.
jf514 Sep 12, 2024
f7777d7
Add unit test.
jf514 Sep 12, 2024
31ff7ab
Build fixes.
jf514 Sep 12, 2024
917e8ca
Fix vis.
jf514 Sep 12, 2024
25a9c16
Revert
jf514 Sep 12, 2024
bb49b57
Missed file.
jf514 Sep 12, 2024
9cc0294
Misseed configs.
jf514 Sep 12, 2024
69d15ed
Remove debug variable.
jf514 Sep 12, 2024
4233964
Clean up plane estimation.
jf514 Sep 12, 2024
3f8ad20
Remove unused vars from configs.
jf514 Sep 13, 2024
92d0c34
Address review comments.
jf514 Sep 13, 2024
b653f36
black reformat.
jf514 Sep 13, 2024
5de7565
Reinstate free joint check.
jf514 Sep 13, 2024
b12e9ac
Add free joint check to transform()
jf514 Sep 13, 2024
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
246 changes: 246 additions & 0 deletions configs/mouse.yaml
Original file line number Diff line number Diff line change
@@ -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
jf514 marked this conversation as resolved.
Show resolved Hide resolved
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
jf514 marked this conversation as resolved.
Show resolved Hide resolved
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"
jf514 marked this conversation as resolved.
Show resolved Hide resolved
- "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
jf514 marked this conversation as resolved.
Show resolved Hide resolved
- 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
2 changes: 0 additions & 2 deletions configs/rodent.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -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
12 changes: 12 additions & 0 deletions configs/stac_mouse.yaml
Original file line number Diff line number Diff line change
@@ -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
jf514 marked this conversation as resolved.
Show resolved Hide resolved
Loading
Loading