talmolab · jf514 · Nov 1, 2024 · Oct 8, 2024 · Oct 25, 2024 · Oct 25, 2024
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
@@ -99,4 +99,9 @@ jobs:
         with:
           fail_ci_if_error: true
           verbose: false
-          token: ${{ secrets.CODECOV_TOKEN }}
+          token: ${{ secrets.CODECOV_TOKEN }}
+
+      # Test probably delete this
+      - name: Smoke Test
+        shell: bash -l {0}
+        run: python run_stac.py
diff --git a/Mat-to-Nwb-Synth-Data.ipynb b/Mat-to-Nwb-Synth-Data.ipynb
diff --git a/configs/config.yaml b/configs/config.yaml
@@ -1,6 +1,6 @@
 defaults:
   - stac: demo
-  - model: rodent
+  - model: synth
   - _self_
 
 ##FLY_MODEL

diff --git a/configs/model/synth.yaml b/configs/model/synth.yaml
@@ -0,0 +1,58 @@
+
+MJCF_PATH: 'models/synth_model.xml'
+
+# Frames per clip for transform.
+N_FRAMES_PER_CLIP: 1
+
+# Tolerance for the optimizations of the full model, limb, and root. 
+# TODO: Re-implement optimizer loops to use these tolerances
+FTOL: 5.0e-03
+ROOT_FTOL: 1.0e-05
+LIMB_FTOL: 1.0e-06
+
+# Number of alternating pose and offset optimization rounds. 
+N_ITERS: 1
+
+KP_NAMES: 
+  - 'part_0'
+
+# The model sites used to register the keypoints.
+KEYPOINT_MODEL_PAIRS:
+  part_0: base
+
+# The initial offsets for each keypoint in meters.
+KEYPOINT_INITIAL_OFFSETS:
+  part_0: 0 0 0.01
+
+TRUNK_OPTIMIZATION_KEYPOINTS:
+  -part_0
+
+INDIVIDUAL_PART_OPTIMIZATION: 
+  "T1R": ["base", "base"]
+
+# Color to use for each keypoint when visualizing the results
+KEYPOINT_COLOR_PAIRS:
+  base: 0 .5 1 1
+
+# What is the size of the animal you'd like to register, relative to the model?
+SCALE_FACTOR: 1
+
+# 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: 1
+
+# If you have reason to believe that the initial offsets are correct for particular keypoints, 
+# you can regularize those sites using this with M_REG_COEF. 
+SITES_TO_REGULARIZE: 
+  - part_0
+
+RENDER_FPS: 200
+
+N_SAMPLE_FRAMES: 1
+
+# 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
diff --git a/configs/stac/demo.yaml b/configs/stac/demo.yaml
@@ -1,10 +1,13 @@
-fit_path: "demo_fit.p"
+fit_offsets_path: "demo_fit.p"
 ik_only_path: "demo_ik_only.p"
 data_path: "tests/data/test_rodent_mocap_1000_frames.mat"
+fit_offsets_path: "synth_fit.p"
+ik_only_path: "synth_ik_only.p"
+data_path: "tests/data/save_data_AVG_2.nwb"
 
 n_fit_frames: 1
 skip_fit_offsets: False
-ik_only: True
+ik_only: False
 
 mujoco:
   solver: "newton"

diff --git a/configs/stac/stac.yaml b/configs/stac/stac.yaml
@@ -1,4 +1,4 @@
-fit_path: "fit.p"
+fit_offsets_path: "fit.p"
 ik_only_path: "ik_only.p"
 data_path: "tests/data/test_rodent_mocap_1000_frames.nwb"
 

diff --git a/configs/stac/stac_fly_tethered.yaml b/configs/stac/stac_fly_tethered.yaml
@@ -1,5 +1,5 @@
-fit_path: "Fruitfly_fit.p"
-transform_path: "transform_mocap_fly.h5"
+fit_offsets_path: "Fruitfly_fit.p"
+ik_only_path: "transform_mocap_fly.h5"
 # This data is unpublished, need to speak to Elliot Abe about how to obtain it.
 # A possible source is https://datadryad.org/stash/dataset/doi:10.5061/dryad.nzs7h44s4
 # See the discussion here for more info: 

diff --git a/configs/stac/stac_fly_treadmill.yaml b/configs/stac/stac_fly_treadmill.yaml
@@ -1,5 +1,5 @@
-fit_path: "fly_treadmill_fit.p"
-transform_path: "transform_treadmill.p"
+fit_offsets_path: "fly_treadmill_fit.p"
+ik_only_path: "transform_treadmill.p"
 # File is too large to commit
 # DL from: https://datadryad.org/stash/dataset/doi:10.5061/dryad.mpg4f4r73
 # Actual file: https://datadryad.org/stash/downloads/file_stream/3361804

diff --git a/configs/stac/stac_mouse.yaml b/configs/stac/stac_mouse.yaml
@@ -1,4 +1,4 @@
-fit_path: "fit_mouse.p"
+fit_offsets_path: "fit_mouse.p"
 ik_only_path: "ik_only_mouse.p"
 data_path: "tests/data/test_mouse_mocap_3600_frames.h5"
 

diff --git a/demos/create_synth_data.ipynb b/demos/create_synth_data.ipynb
diff --git a/models/synth_model.xml b/models/synth_model.xml
@@ -0,0 +1,23 @@
+<mujoco>
+  <option timestep=".001">
+    <!-- <flag energy="enable" contact="disable"/> -->
+  </option>
+
+  <default>
+    <joint type="hinge" axis="0 -1 0"/>
+    <geom type="capsule" size=".02"/>
+  </default>
+
+  <worldbody>
+    <light pos="0 -.4 1"/>
+    <camera name="fixed" pos="0 -1 0" xyaxes="1 0 0 0 0 1"/>
+    <body name="base" pos="0 0 .2">
+      <joint name="root"/>
+      <geom fromto="0 0 0 0 0 -.25" rgba="1 1 0 1"/>
+      <!-- <body name="3" pos="0 0 -.25">
+         <joint/>
+         <geom fromto="0 0 0 0 0 -.2" rgba="0 0 1 1"/>
+      </body> -->
+    </body>
+  </worldbody>
+</mujoco>
diff --git a/stac_mjx/io.py b/stac_mjx/io.py
@@ -67,6 +67,7 @@ def load_data(cfg: DictConfig, base_path: Union[Path, None] = None):
             f"Number of keypoint names ({len(kp_names)}) is not the same as the number of keypoints in data ({data.shape[1]})"
         )
 
+    print("kpnames ", kp_names)
     model_inds = [
         kp_names.index(src) for src, dst in cfg.model.KEYPOINT_MODEL_PAIRS.items()
     ]

diff --git a/stac_mjx/main.py b/stac_mjx/main.py
@@ -56,7 +56,7 @@ def run_stac(
     start_time = time.time()
 
     # Getting paths
-    fit_path = base_path / cfg.stac.fit_path
+    fit_offsets_path = base_path / cfg.stac.fit_offsets_path
     ik_only_path = base_path / cfg.stac.ik_only_path
 
     xml_path = base_path / cfg.model.MJCF_PATH
@@ -69,21 +69,21 @@ def run_stac(
         logging.info(f"Running fit. Mocap data shape: {fit_data.shape}")
         fit_data = stac.fit_offsets(fit_data)
 
-        logging.info(f"saving data to {fit_path}")
-        io.save(fit_data, fit_path)
+        logging.info(f"saving data to {fit_offsets_path}")
+        io.save(fit_data, fit_offsets_path)
 
     # Stop here if not doing ik only phase
     if cfg.stac.skip_ik_only == 1:
         logging.info("skipping ik_only()")
-        return fit_path, None
+        return fit_offsets_path, None
     # FLY_MODEL: The elif below must be commented out for fly_model.
     elif kp_data.shape[0] % cfg.model.N_FRAMES_PER_CLIP != 0:
         raise ValueError(
             f"N_FRAMES_PER_CLIP ({cfg.model.N_FRAMES_PER_CLIP}) must divide evenly with the total number of mocap frames({kp_data.shape[0]})"
         )
 
     logging.info("Running ik_only()")
-    with open(fit_path, "rb") as file:
+    with open(fit_offsets_path, "rb") as file:
         fit_data = pickle.load(file)
     offsets = fit_data["offsets"]
 
@@ -95,4 +95,4 @@ def run_stac(
     )
     io.save(ik_only_data, ik_only_path)
 
-    return fit_path, ik_only_path
+    return fit_offsets_path, ik_only_path
diff --git a/stac_mjx/stac.py b/stac_mjx/stac.py
@@ -390,7 +390,7 @@ def _package_data(self, mjx_model, q, x, walker_body_sites, kp_data, batched=Fal
             x = x.reshape(-1, x.shape[-1])
             q = q.reshape(-1, q.shape[-1])
         else:
-            offsets = self._offsets
+            offsets = self._offsets.reshape((-1, 3))
 
         kp_data = kp_data.reshape(-1, kp_data.shape[-1])
 

diff --git a/tests/configs/stac/test_stac.yaml b/tests/configs/stac/test_stac.yaml
@@ -1,4 +1,4 @@
-fit_path: "fit.p"
+fit_offsets_path: "fit.p"
 ik_only_path: "ik_only.p"
 data_path: "tests/data/test_rodent_mocap_1000_frames.nwb"
 

diff --git a/tests/data/save_data_AVG_2.nwb b/tests/data/save_data_AVG_2.nwb
diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py
@@ -18,5 +18,5 @@ def test_load_configs(config):
     assert isinstance(cfg, DictConfig)
 
     # Assert that the resulting configs contain the expected data
-    assert cfg.stac.fit_path == "fit.p"
+    assert cfg.stac.fit_offsets_path == "fit.p"
     assert cfg.model.N_FRAMES_PER_CLIP == 360