feat: CPU demo

.gitignore

@@ -0,0 +1,17 @@
+# Caches.
+.DS_Store
+__pycache__
+*.egg-info
+
+# VEnv.
+.hsmr_env
+
+# Hydra dev.
+outputs
+
+# Development tools.
+.vscode
+pyrightconfig.json
+
+# Logs.
+gpu_monitor.log

Dockerfile

@@ -0,0 +1,37 @@
+# read the doc: https://huggingface.co/docs/hub/spaces-sdks-docker
+# you will also find guides on how best to write your Dockerfile
+
+FROM python:3.10
+
+WORKDIR /code
+
+COPY ./requirements_part1.txt /code/requirements_part1.txt
+COPY ./requirements_part2.txt /code/requirements_part2.txt
+
+RUN apt-get update -y && \
+    apt-get upgrade -y && \
+    apt-get install -y libglfw3-dev && \
+    apt-get install -y libgles2-mesa-dev && \
+    apt-get install -y aria2 && \
+    pip install --no-cache-dir --upgrade -r /code/requirements_part1.txt && \
+    pip install --no-cache-dir --upgrade -r /code/requirements_part2.txt
+
+# Set up a new user named "user" with user ID 1000
+RUN useradd -m -u 1000 user
+
+# Switch to the "user" user
+USER user
+
+
+# Set home to the user's home directory
+ENV HOME=/home/user \
+    PATH=/home/user/.local/bin:$PATH
+
+# Set the working directory to the user's home directory
+WORKDIR $HOME/app
+
+# Copy the current directory contents into the container at $HOME/app setting the owner to the user
+COPY --chown=user . $HOME/app
+
+
+CMD ["bash", "tools/start.sh"]

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Yan XIA
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -1,11 +1,10 @@
 ---
 title: HSMR
 emoji: 💀
-colorFrom: purple
-colorTo: indigo
-sdk: gradio
-sdk_version: 5.20.0
-app_file: tools/service.py
+colorFrom: blue
+colorTo: pink
+sdk: docker
+app_port: 7860
 pinned: true
 ---
 

configs/README.md

@@ -0,0 +1,31 @@
+# Instructions for the Configuration System
+
+The configuration system I used here is based on [Hydra](https://hydra.cc/). However, I made some small 'hack' to achieve some better features, such as `_hub_`. It might be a little bit hard to understand at first, but a comprehensive guidance is provided. Check `README.md` in each directory for more details.
+
+## Philosophy
+
+- Easy to modify and maintain.
+- Help you to code with clear structure.
+- Consistency.
+- Easy to trace and identify specific item.
+
+## Some Ideas
+
+- Less ListConfig, or ListConfig only for real list data.
+    - Dumped list will be unfolded, each element occupies one line, and it's annoying when presenting.
+    - List things are not friendly to command line arguments supports.
+- For defaults list, `_self_` must be explicitly specified.
+    - Items before `_self_` means 'based on those items'.
+    - Items after `_self_` means 'import those items'.
+
+### COMPOSITION OVER INHERITANCE
+
+Do not use "overrides" AS MUCH AS POSSIBLE, except when the changes are really tiny. Since it's hard to identify which term is actually used without running the code. Instead, the `default.yaml` serves like a template, you are supposed to copy it and modify it to create a new configuration.
+
+### REFERENCE OVER COPY
+
+If you want to use one things for many times (across each experiments or across each components in one experiment), you'd better use `${...}` to reference the Hydra object. So that you only need to modify one place while ensuring the consistency. Think about where to put the source of the object, `_hub_` is recommended but may not be the best choice every time.
+
+### PREPARE EVERYTHING YOU NEED LOCALLY
+
+Sometimes you might want to use some configurations outside the class configuration (I mean in the coding process). In that case, I recommend you to reference these things again in local configuration package. It might be too redundant, but it will bring cleaner code.

configs/_hub_/README.md

@@ -0,0 +1,7 @@
+# _hub
+
+Configs here shouldn't be used as what Hydra calls a "config group". Configs here serve as a hub for other configs to reference. Most things here usually won't be used in a certain experiment.
+
+For example, the details of each datasets can be defined here, and than I only need to reference the `Human3.6M` dataset through `${_hub_.datasets.h36m}` in other configs.
+
+Each `.yaml` file here will be loaded separately in `base.yaml`. Check the `base.yaml` to understand how it works.

configs/_hub_/datasets.yaml

@@ -0,0 +1,71 @@
+train:
+  hsmr:  # Fitted according to the SMPL's vertices.
+    # Standard 4 datasets.
+    mpi_inf_3dhp:
+      name: 'HSMR-MPI-INF-3DHP-train-pruned'
+      urls: ${_pm_.inputs}/hsmr_training_data/mpi_inf_3dhp-tars/{000000..000012}.tar
+      epoch_size: 12_000
+    h36m:
+      name: 'HSMR-H36M-train'
+      urls: ${_pm_.inputs}/hsmr_training_data/h36m-tars/{000000..000312}.tar
+      epoch_size: 314_000
+    mpii:
+      name: 'HSMR-MPII-train'
+      urls: ${_pm_.inputs}/hsmr_training_data/mpii-tars/{000000..000009}.tar
+      epoch_size: 10_000
+    coco14:
+      name: 'HSMR-COCO14-train'
+      urls: ${_pm_.inputs}/hsmr_training_data/coco14-tars/{000000..000017}.tar
+      epoch_size: 18_000
+    # The rest full datasets for HMR2.0.
+    coco14vit:
+      name: 'HSMR-COCO14-vit-train'
+      urls: ${_pm_.inputs}/hsmr_training_data/coco14vit-tars/{000000..000044}.tar
+      epoch_size: 45_000
+    aic:
+      name: 'HSMR-AIC-train'
+      urls: ${_pm_.inputs}/hsmr_training_data/aic-tars/{000000..000209}.tar
+      epoch_size: 210_000
+    ava:
+      name: 'HSMR-AVA-train'
+      urls: ${_pm_.inputs}/hsmr_training_data/ava-tars/{000000..000184}.tar
+      epoch_size: 185_000
+    insta:
+      name: 'HSMR-INSTA-train'
+      urls: ${_pm_.inputs}/hsmr_training_data/insta-tars/{000000..003657}.tar
+      epoch_size: 3_658_000
+
+mocap:
+  bioamass_v1:
+    dataset_file: ${_pm_.inputs}/datasets/amass_skel/data_v1.npz
+    pve_threshold: 0.05
+  cmu_mocap:
+    dataset_file: ${_pm_.inputs}/hmr2_training_data/cmu_mocap.npz
+
+eval:
+  h36m_val_p2:
+    dataset_file: ${_pm_.inputs}/hmr2_evaluation_data/h36m_val_p2.npz
+    img_root: ${_pm_.inputs}/datasets/h36m/images
+    kp_list: [25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 43]
+    use_hips: true
+
+  3dpw_test:
+    dataset_file: ${_pm_.inputs}/hmr2_evaluation_data/3dpw_test.npz
+    img_root: ${_pm_.inputs}/datasets/3dpw
+    kp_list: [25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 43]
+    use_hips: false
+
+  posetrack_val:
+    dataset_file: ${_pm_.inputs}/hmr2_evaluation_data/posetrack_2018_val.npz
+    img_root: ${_pm_.inputs}/datasets/posetrack/posetrack2018/posetrack_data/
+    kp_list: [0] # dummy
+
+  lsp_extended:
+    dataset_file: ${_pm_.inputs}/hmr2_evaluation_data/hr-lspet_train.npz
+    img_root: ${_pm_.inputs}/datasets/hr-lspet
+    kp_list: [0] # dummy
+
+  coco_val:
+    dataset_file: ${_pm_.inputs}/hmr2_evaluation_data/coco_val.npz
+    img_root: ${_pm_.inputs}/datasets/coco
+    kp_list: [0] # dummy

configs/_hub_/models.yaml

@@ -0,0 +1,57 @@
+body_models:
+
+  skel_mix_hsmr:
+    _target_: lib.body_models.skel_wrapper.SKELWrapper
+    model_path: '${_pm_.inputs}/body_models/skel'
+    gender: male  # Use male since we don't have neutral model.
+    joint_regressor_extra: '${_pm_.inputs}/body_models/SMPL_to_J19.pkl'
+    joint_regressor_custom: '${_pm_.inputs}/body_models/J_regressor_SKIL_mix_MALE.pkl'
+
+  skel_hsmr:
+    _target_: lib.body_models.skel_wrapper.SKELWrapper
+    model_path: '${_pm_.inputs}/body_models/skel'
+    gender: male  # Use male since we don't have neutral model.
+    joint_regressor_extra: '${_pm_.inputs}/body_models/SMPL_to_J19.pkl'
+    # joint_regressor_custom: '${_pm_.inputs}/body_models/J_regressor_SMPL_MALE.pkl'
+
+  smpl_hsmr:
+    _target_: lib.body_models.smpl_wrapper.SMPLWrapper
+    model_path: '${_pm_.inputs}/body_models/smpl'
+    gender: male  # align with skel_hsmr
+    num_body_joints: 23
+    joint_regressor_extra: '${_pm_.inputs}/body_models/SMPL_to_J19.pkl'
+
+  smpl_hsmr_neutral:
+    _target_: lib.body_models.smpl_wrapper.SMPLWrapper
+    model_path: '${_pm_.inputs}/body_models/smpl'
+    gender: neutral  # align with skel_hsmr
+    num_body_joints: 23
+    joint_regressor_extra: '${_pm_.inputs}/body_models/SMPL_to_J19.pkl'
+
+backbones:
+
+  vit_b:
+    _target_: lib.modeling.networks.backbones.ViT
+    img_size: [256, 192]
+    patch_size: 16
+    embed_dim: 768
+    depth: 12
+    num_heads: 12
+    ratio: 1
+    use_checkpoint: False
+    mlp_ratio: 4
+    qkv_bias: True
+    drop_path_rate: 0.3
+
+  vit_h:
+    _target_: lib.modeling.networks.backbones.ViT
+    img_size: [256, 192]
+    patch_size: 16
+    embed_dim: 1280
+    depth: 32
+    num_heads: 16
+    ratio: 1
+    use_checkpoint: False
+    mlp_ratio: 4
+    qkv_bias: True
+    drop_path_rate: 0.55

configs/base.yaml

@@ -0,0 +1,33 @@
+defaults:
+  # Load each sub-hub.
+  - _hub_/datasets@_hub_.datasets
+  - _hub_/models@_hub_.models
+
+  # Set up the template.
+  - _self_
+
+  # Register some defaults.
+  - exp: default
+  - policy: default
+
+
+# exp_name: !!null  # Name of experiment will determine the output folder.
+exp_name: '${exp_topic}-${exp_tag}'
+exp_topic: !!null  # Theme of the experiment.
+exp_tag: 'debug'  # Tag of the experiment.
+
+
+output_dir: ${_pm_.root}/data_outputs/exp/${exp_name}  # Output directory for the experiment.
+# output_dir: ${_pm_.root}/data_outputs/exp/${now:%Y-%m-%d}/${exp_name}  # Output directory for the experiment.
+
+hydra:
+  run:
+    dir: ${output_dir}
+
+
+# Information from `PathManager`, which will be automatically set by entrypoint wrapper.
+# The related implementation is in `lib/utils/cfg_utils.py`:`entrypoint_with_args`.
+_pm_:
+  root: !!null
+  inputs: !!null
+  outputs: !!null

configs/callback/ckpt/all-p1k.yaml

@@ -0,0 +1,5 @@
+_target_: pytorch_lightning.callbacks.ModelCheckpoint
+dirpath: '${output_dir}/checkpoints'
+save_last: True
+every_n_train_steps: 1000
+save_top_k: -1  # save all ckpts

configs/callback/ckpt/top1-p1k.yaml

@@ -0,0 +1,5 @@
+_target_: pytorch_lightning.callbacks.ModelCheckpoint
+dirpath: '${output_dir}/checkpoints'
+save_last: True
+every_n_train_steps: 1000
+save_top_k: 1  # save only the last checkpoint

configs/callback/skelify-spin/i10kb1.yaml

@@ -0,0 +1,19 @@
+defaults:
+  - _self_
+  # Import the SKELify-SPIN callback.
+  - /pipeline/skelify-refiner@skelify
+
+_target_: lib.modeling.callbacks.SKELifySPIN
+
+# This configuration are for frozen backbone exp with batch_size = 24000
+
+cfg:
+  interval: 10
+  batch_size: 24000  # when greater than `interval * dataloader's batch_size`, it's equivalent to that
+  max_batches_per_round: 1  # only the latest k * batch_size items are SPINed to save time
+  # better_pgt_fn: '${_pm_.inputs}/datasets/skel_training_data/spin/better_pseudo_gt.npz'
+  better_pgt_fn: '${output_dir}/better_pseudo_gt.npz'
+  skip_warm_up_steps: 10000
+  # skip_warm_up_steps: 0
+  update_better_pgt: True
+  valid_betas_threshold: 2

configs/callback/skelify-spin/i230kb1.yaml

@@ -0,0 +1,16 @@
+defaults:
+  - _self_
+  # Import the SKELify-SPIN callback.
+  - /pipeline/skelify-refiner@skelify
+
+_target_: lib.modeling.callbacks.SKELifySPIN
+
+cfg:
+  interval: 230
+  batch_size: 18000  # when greater than `interval * dataloader's batch_size`, it's equivalent to that
+  max_batches_per_round: 1  # only the latest k * batch_size items are SPINed to save time
+  # better_pgt_fn: '${_pm_.inputs}/datasets/skel_training_data/spin/better_pseudo_gt.npz'
+  better_pgt_fn: '${output_dir}/better_pseudo_gt.npz'
+  skip_warm_up_steps: 5000
+  update_better_pgt: True
+  valid_betas_threshold: 2

configs/callback/skelify-spin/i80.yaml

@@ -0,0 +1,15 @@
+defaults:
+  - _self_
+  # Import the SKELify-SPIN callback.
+  - /pipeline/skelify-refiner@skelify
+
+_target_: lib.modeling.callbacks.SKELifySPIN
+
+cfg:
+  interval: 80
+  batch_size: 24000  # when greater than `interval * dataloader's batch_size`, it's equivalent to that
+  # better_pgt_fn: '${_pm_.inputs}/datasets/skel_training_data/spin/better_pseudo_gt.npz'
+  better_pgt_fn: '${output_dir}/better_pseudo_gt.npz'
+  skip_warm_up_steps: 5000
+  update_better_pgt: True
+  valid_betas_threshold: 2

configs/callback/skelify-spin/read_only.yaml

@@ -0,0 +1,18 @@
+defaults:
+  - _self_
+  # Import the SKELify-SPIN callback.
+  - /pipeline/skelify-refiner@skelify
+
+_target_: lib.modeling.callbacks.SKELifySPIN
+
+# This configuration are for frozen backbone exp with batch_size = 24000
+
+cfg:
+  interval: 0
+  batch_size: 0  # when greater than `interval * dataloader's batch_size`, it's equivalent to that
+  max_batches_per_round: 0  # only the latest k * batch_size items are SPINed to save time
+  # better_pgt_fn: '${_pm_.inputs}/datasets/skel_training_data/spin/better_pseudo_gt.npz'
+  better_pgt_fn: '${output_dir}/better_pseudo_gt.npz'
+  skip_warm_up_steps: 0
+  update_better_pgt: True
+  valid_betas_threshold: 2

configs/data/skel-hmr2_fashion.yaml

@@ -0,0 +1,71 @@
+_target_: lib.data.modules.hmr2_fashion.skel_wds.DataModule
+
+name: HMR2_fashion_WDS
+
+cfg:
+
+  train:
+    shared_ds_opt:  # TODO: modify this
+      SUPPRESS_KP_CONF_THRESH: 0.3
+      FILTER_NUM_KP: 4
+      FILTER_NUM_KP_THRESH: 0.0
+      FILTER_REPROJ_THRESH: 31000
+      SUPPRESS_BETAS_THRESH: 3.0
+      SUPPRESS_BAD_POSES: False
+      POSES_BETAS_SIMULTANEOUS: True
+      FILTER_NO_POSES: False
+      BETAS_REG: True
+
+    datasets:
+      - name: 'H36M'
+        item: ${_hub_.datasets.train.hsmr.h36m}
+        weight: 0.3
+      - name: 'MPII'
+        item: ${_hub_.datasets.train.hsmr.mpii}
+        weight: 0.1
+      - name: 'COCO14'
+        item: ${_hub_.datasets.train.hsmr.coco14}
+        weight: 0.4
+      - name: 'MPI-INF-3DHP'
+        item: ${_hub_.datasets.train.hsmr.mpi_inf_3dhp}
+        weight: 0.2
+
+    dataloader:
+      drop_last: True
+      batch_size: 300
+      num_workers: 6
+      prefetch_factor: 2
+
+  eval:
+    datasets:
+      - name: 'LSP-EXTENDED'
+        item: ${_hub_.datasets.eval.lsp_extended}
+      - name: 'H36M-VAL-P2'
+        item: ${_hub_.datasets.eval.h36m_val_p2}
+      - name: '3DPW-TEST'
+        item: ${_hub_.datasets.eval.3dpw_test}
+      - name: 'POSETRACK-VAL'
+        item: ${_hub_.datasets.eval.posetrack_val}
+      - name: 'COCO-VAL'
+        item: ${_hub_.datasets.eval.coco_val}
+
+
+    dataloader:
+      shuffle: False
+      batch_size: 300
+      num_workers: 6
+
+  policy: ${policy}
+
+  # TODO: modify this
+  augm:
+    SCALE_FACTOR: 0.3
+    ROT_FACTOR: 30
+    TRANS_FACTOR: 0.02
+    COLOR_SCALE: 0.2
+    ROT_AUG_RATE: 0.6
+    TRANS_AUG_RATE: 0.5
+    DO_FLIP: True
+    FLIP_AUG_RATE: 0.5
+    EXTREME_CROP_AUG_RATE: 0.10
+    EXTREME_CROP_AUG_LEVEL: 1

configs/data/skel-hsr_v1_4ds.yaml

@@ -0,0 +1,84 @@
+_target_: lib.data.modules.hsmr_v1.data_module.DataModule
+
+name: SKEL_HSMR_V1
+
+cfg:
+
+  train:
+    cfg:
+      # Loader settings.
+      suppress_pgt_params_pve_max_thresh: 0.06  # Mark PVE-MAX larger than this value as invalid.
+      suppress_kp_conf_thresh: 0.3  # Mark key-point confidence smaller than this value as invalid.
+      suppress_betas_thresh: 3.0  # Mark betas having components larger than this value as invalid.
+      poses_betas_simultaneous: True  # Sync poses and betas for the small person.
+      filter_insufficient_kp_cnt: 4
+      suppress_insufficient_kp_thresh: 0.0
+      filter_reproj_err_thresh: 31000
+      regularize_invalid_betas: True
+      # Others.
+      image_augmentation: ${...image_augmentation}
+      policy: ${...policy}
+
+    datasets:
+      - name: 'H36M'
+        item: ${_hub_.datasets.train.hsmr.h36m}
+        weight: 0.3
+      - name: 'MPII'
+        item: ${_hub_.datasets.train.hsmr.mpii}
+        weight: 0.1
+      - name: 'COCO14'
+        item: ${_hub_.datasets.train.hsmr.coco14}
+        weight: 0.4
+      - name: 'MPI-INF-3DHP'
+        item: ${_hub_.datasets.train.hsmr.mpi_inf_3dhp}
+        weight: 0.2
+
+    dataloader:
+      drop_last: True
+      batch_size: 300
+      num_workers: 6
+      prefetch_factor: 2
+
+  # mocap:
+  #   cfg: ${_hub_.datasets.mocap.bioamass_v1}
+  #   dataloader:
+  #     batch_size: 600  # num_train:2 * batch_size:300 (from HMR2.0's cfg)
+  #     drop_last: True
+  #     shuffle: True
+  #     num_workers: 1
+
+  eval:
+    cfg:
+      image_augmentation: ${...image_augmentation}
+      policy: ${...policy}
+
+    datasets:
+      - name: 'LSP-EXTENDED'
+        item: ${_hub_.datasets.eval.lsp_extended}
+      - name: 'H36M-VAL-P2'
+        item: ${_hub_.datasets.eval.h36m_val_p2}
+      - name: '3DPW-TEST'
+        item: ${_hub_.datasets.eval.3dpw_test}
+      - name: 'POSETRACK-VAL'
+        item: ${_hub_.datasets.eval.posetrack_val}
+      - name: 'COCO-VAL'
+        item: ${_hub_.datasets.eval.coco_val}
+
+    dataloader:
+      shuffle: False
+      batch_size: 300
+      num_workers: 6
+
+  # Augmentation settings.
+  image_augmentation:
+    trans_factor: 0.02
+    bbox_scale_factor: 0.3
+    rot_aug_rate: 0.6
+    rot_factor: 30
+    do_flip: True
+    flip_aug_rate: 0.5
+    extreme_crop_aug_rate: 0.10
+    half_color_scale: 0.2
+
+  # Others.
+  policy: ${policy}

configs/data/skel-hsr_v1_full.yaml

@@ -0,0 +1,96 @@
+_target_: lib.data.modules.hsmr_v1.data_module.DataModule
+
+name: SKEL_HSMR_V1
+
+cfg:
+
+  train:
+    cfg:
+      # Loader settings.
+      suppress_pgt_params_pve_max_thresh: 0.06  # Mark PVE-MAX larger than this value as invalid.
+      suppress_kp_conf_thresh: 0.3  # Mark key-point confidence smaller than this value as invalid.
+      suppress_betas_thresh: 3.0  # Mark betas having components larger than this value as invalid.
+      poses_betas_simultaneous: True  # Sync poses and betas for the small person.
+      filter_insufficient_kp_cnt: 4
+      suppress_insufficient_kp_thresh: 0.0
+      filter_reproj_err_thresh: 31000
+      regularize_invalid_betas: True
+      # Others.
+      image_augmentation: ${...image_augmentation}
+      policy: ${...policy}
+
+    datasets:
+      - name: 'H36M'
+        item: ${_hub_.datasets.train.hsmr.h36m}
+        weight: 0.1
+      - name: 'MPII'
+        item: ${_hub_.datasets.train.hsmr.mpii}
+        weight: 0.1
+      - name: 'COCO14'
+        item: ${_hub_.datasets.train.hsmr.coco14}
+        weight: 0.1
+      - name: 'COCO14-ViTPose'
+        item: ${_hub_.datasets.train.hsmr.coco14vit}
+        weight: 0.1
+      - name: 'MPI-INF-3DHP'
+        item: ${_hub_.datasets.train.hsmr.mpi_inf_3dhp}
+        weight: 0.02
+      - name: 'AVA'
+        item: ${_hub_.datasets.train.hsmr.ava}
+        weight: 0.19
+      - name: 'AIC'
+        item: ${_hub_.datasets.train.hsmr.aic}
+        weight: 0.19
+      - name: 'INSTA'
+        item: ${_hub_.datasets.train.hsmr.insta}
+        weight: 0.2
+
+    dataloader:
+      drop_last: True
+      batch_size: 300
+      num_workers: 6
+      prefetch_factor: 2
+
+  mocap:
+    cfg: ${_hub_.datasets.mocap.bioamass_v1}
+    dataloader:
+      batch_size: 600  # num_train:2 * batch_size:300 (from HMR2.0's cfg)
+      drop_last: True
+      shuffle: True
+      num_workers: 1
+
+  eval:
+    cfg:
+      image_augmentation: ${...image_augmentation}
+      policy: ${...policy}
+
+    datasets:
+      - name: 'LSP-EXTENDED'
+        item: ${_hub_.datasets.eval.lsp_extended}
+      - name: 'H36M-VAL-P2'
+        item: ${_hub_.datasets.eval.h36m_val_p2}
+      - name: '3DPW-TEST'
+        item: ${_hub_.datasets.eval.3dpw_test}
+      - name: 'POSETRACK-VAL'
+        item: ${_hub_.datasets.eval.posetrack_val}
+      - name: 'COCO-VAL'
+        item: ${_hub_.datasets.eval.coco_val}
+
+    dataloader:
+      shuffle: False
+      batch_size: 300
+      num_workers: 6
+
+  # Augmentation settings.
+  image_augmentation:
+    trans_factor: 0.02
+    bbox_scale_factor: 0.3
+    rot_aug_rate: 0.6
+    rot_factor: 30
+    do_flip: True
+    flip_aug_rate: 0.5
+    extreme_crop_aug_rate: 0.10
+    half_color_scale: 0.2
+
+  # Others.
+  policy: ${policy}

configs/exp/default.yaml

@@ -0,0 +1,16 @@
+# @package _global_
+defaults:
+  # Use absolute path, cause the root in this file is `configs/exp/` rather than `configs/`.
+  # - /data: ...
+
+  # Configurations in this file should have the highest priority.
+  - _self_
+
+# ======= Overwrite Section =======
+# (Do not use as much as possible!)
+
+
+
+# ====== Main Section ======
+
+exp_topic: !!null  # Name of experiment will determine the output folder.

configs/exp/hsr/train.yaml

@@ -0,0 +1,47 @@
+# @package _global_
+defaults:
+  - /pipeline: hsmr
+  - /data: skel-hsmr_v1_full
+  # Configurations in this file
+  - _self_
+  # Import callbacks.
+  - /callback/ckpt/[email protected]
+  # - /callback/skelify-spin/[email protected]
+  - /callback/skelify-spin/[email protected]
+
+# ======= Overwrite Section =======
+# (Do not use as much as possible!)
+
+pipeline:
+  cfg:
+    backbone: ${_hub_.models.backbones.vit_h}
+    backbone_ckpt: '${_pm_.inputs}/backbone/vitpose_backbone.pth'
+    freeze_backbone: False
+
+data:
+  cfg:
+    train:
+      dataloader:
+        batch_size: 78
+
+
+# ====== Main Section ======
+
+exp_topic: 'HSMR-train-vit_h'
+
+enable_time_monitor: False
+
+seed: NULL
+ckpt_path: NULL
+
+logger:
+  interval: 1000
+  interval_skelify: 10
+  samples_per_record: 5
+
+task: 'fit'
+pl_trainer:
+  devices: 8
+  max_epochs: 100
+  deterministic: false
+  precision: 16

configs/exp/skelify-full.yaml

@@ -0,0 +1,21 @@
+# @package _global_
+defaults:
+  - /pipeline: skelify-full
+  - /data: skel-hsmr_v1_full
+  # Configurations in this file should have the highest priority.
+  - _self_
+
+# ======= Overwrite Section =======
+# (Do not use as much as possible!)
+
+
+
+# ====== Main Section ======
+
+exp_topic: 'SKELify-Full'
+
+enable_time_monitor: True
+
+logger:
+  interval: 1
+  samples_per_record: 8

configs/exp/skelify-refiner.yaml

@@ -0,0 +1,21 @@
+# @package _global_
+defaults:
+  - /pipeline: skelify-refiner
+  - /data: skel-hsmr_v1_full
+  # Configurations in this file should have the highest priority.
+  - _self_
+
+# ======= Overwrite Section =======
+# (Do not use as much as possible!)
+
+
+
+# ====== Main Section ======
+
+exp_topic: 'SKELify-Refiner'
+
+enable_time_monitor: True
+
+logger:
+  interval: 2
+  samples_per_record: 8

configs/pipeline/hmr2.yaml

@@ -0,0 +1,43 @@
+_target_: lib.modeling.pipelines.HMR2Pipeline
+
+name: HMR2
+
+cfg:
+  # Body Models
+  SMPL: ${_hub_.models.body_models.smpl_hsmr_neutral}
+
+  # Backbone and its checkpoint.
+  backbone: ${_hub_.models.backbones.vit_h}
+  backbone_ckpt: ???
+
+  # Head to get the parameters.
+  head:
+    _target_: lib.modeling.networks.heads.SMPLTransformerDecoderHead
+    cfg:
+      transformer_decoder:
+        depth: 6
+        heads: 8
+        mlp_dim: 1024
+        dim_head: 64
+        dropout: 0.0
+        emb_dropout: 0.0
+        norm: 'layer'
+        context_dim: ${....backbone.embed_dim}
+
+  optimizer:
+    _target_: torch.optim.AdamW
+    lr: 1e-5
+    weight_decay: 1e-4
+
+  # This may be redesigned, e.g., we can add a loss object to maintain the calculation of the loss, or a callback.
+  loss_weights:
+    kp3d: 0.05
+    kp2d: 0.01
+    poses_orient: 0.002
+    poses_body: 0.001
+    betas: 0.0005
+    adversarial: 0.0005
+    # adversarial: 0.0
+
+  policy: ${policy}
+  logger: ${logger}

configs/pipeline/hsr.yaml

@@ -0,0 +1,49 @@
+_target_: lib.modeling.pipelines.HSMRPipeline
+
+name: HSMR
+
+cfg:
+  pd_poses_repr: 'rotation_6d'  # poses representation for prediction, choices: 'euler_angle' | 'rotation_6d'
+  sp_poses_repr: 'rotation_matrix'  # poses representation for supervision, choices: 'euler_angle' | 'rotation_matrix'
+
+  # Body Models
+  # SKEL: ${_hub_.models.body_models.skel_hsmr}
+  SKEL: ${_hub_.models.body_models.skel_mix_hsmr}
+
+  # Backbone and its checkpoint.
+  backbone: ???
+  backbone_ckpt: ???
+
+  # Head to get the parameters.
+  head:
+    _target_: lib.modeling.networks.heads.SKELTransformerDecoderHead
+    cfg:
+      pd_poses_repr: ${...pd_poses_repr}
+      transformer_decoder:
+        depth: 6
+        heads: 8
+        mlp_dim: 1024
+        dim_head: 64
+        dropout: 0.0
+        emb_dropout: 0.0
+        norm: 'layer'
+        context_dim: ${....backbone.embed_dim}
+
+  optimizer:
+    _target_: torch.optim.AdamW
+    lr: 1e-5
+    weight_decay: 1e-4
+
+  # This may be redesigned, e.g., we can add a loss object to maintain the calculation of the loss, or a callback.
+  loss_weights:
+    kp3d: 0.05
+    kp2d: 0.01
+    # prior: 0.0005
+    prior: 0.0
+    poses_orient: 0.002
+    poses_body: 0.001
+    betas: 0.0005
+    # adversarial: 0.0005
+
+  policy: ${policy}
+  logger: ${logger}

configs/pipeline/skelify-full.yaml

@@ -0,0 +1,98 @@
+_target_: lib.modeling.optim.SKELify
+
+name: SKELify
+
+cfg:
+  skel_model: ${_hub_.models.body_models.skel_mix_hsmr}
+
+  _f_normalize_kp2d: True
+  _f_normalize_kp2d_to_mean: False
+  _w_angle_prior_scale: 1.7
+
+  phases:
+    # ================================
+    # ⛩️ Part 1: Camera initialization.
+    # --------------------------------
+    STAGE-camera-init:
+      max_loop: 30
+      params_keys: ['cam_t', 'poses_orient']
+      parts: ['torso']
+      optimizer: ${...optimizer}
+      losses:
+        f_normalize_kp2d: ${...._f_normalize_kp2d}
+        f_normalize_kp2d_to_mean: ${...._f_normalize_kp2d_to_mean}
+        w_depth: 100.0
+        w_reprojection: 1.78
+    # ================================
+
+    # ================================
+    # ⛩️ Part 2: Overall optimization.
+    # --------------------------------
+    STAGE-overall-1:
+      max_loop: 30
+      params_keys: ['cam_t', 'poses_orient', 'poses_body', 'betas']
+      parts: ['all']
+      optimizer: ${...optimizer}
+      losses:
+        f_normalize_kp2d: ${...._f_normalize_kp2d}
+        f_normalize_kp2d_to_mean: ${...._f_normalize_kp2d_to_mean}
+        w_reprojection: 1.0
+        w_shape_prior: 100.0
+        w_angle_prior: 404.0
+        w_angle_prior_scale: ${...._w_angle_prior_scale}  # TODO: Finalize it.
+    # --------------------------------
+    STAGE-overall-2:
+      max_loop: 30
+      params_keys: ['cam_t', 'poses_orient', 'poses_body', 'betas']
+      optimizer: ${...optimizer}
+      parts: ['all']
+      losses:
+        f_normalize_kp2d: ${...._f_normalize_kp2d}
+        f_normalize_kp2d_to_mean: ${...._f_normalize_kp2d_to_mean}
+        w_reprojection: 1.0
+        w_shape_prior: 50.0
+        w_angle_prior: 404.0
+        w_angle_prior_scale: ${...._w_angle_prior_scale}  # TODO: Finalize it.
+    # --------------------------------
+    STAGE-overall-3:
+      max_loop: 30
+      params_keys: ['cam_t', 'poses_orient', 'poses_body', 'betas']
+      parts: ['all']
+      optimizer: ${...optimizer}
+      losses:
+        f_normalize_kp2d: ${...._f_normalize_kp2d}
+        f_normalize_kp2d_to_mean: ${...._f_normalize_kp2d_to_mean}
+        w_reprojection: 1.0
+        w_shape_prior: 10.0
+        w_angle_prior: 57.4
+        w_angle_prior_scale: ${...._w_angle_prior_scale}  # TODO: Finalize it.
+    # --------------------------------
+    STAGE-overall-4:
+      max_loop: 30
+      params_keys: ['cam_t', 'poses_orient', 'poses_body', 'betas']
+      parts: ['all']
+      optimizer: ${...optimizer}
+      losses:
+        f_normalize_kp2d: ${...._f_normalize_kp2d}
+        f_normalize_kp2d_to_mean: ${...._f_normalize_kp2d_to_mean}
+        w_reprojection: 1.0
+        w_shape_prior: 5.0
+        w_angle_prior: 4.78
+        w_angle_prior_scale: ${...._w_angle_prior_scale}  # TODO: Finalize it.
+    # ================================
+
+  optimizer:
+    _target_: torch.optim.LBFGS
+    lr: 1
+    line_search_fn: 'strong_wolfe'
+    tolerance_grad: ${..early_quit_thresholds.abs}
+    tolerance_change: ${..early_quit_thresholds.rel}
+
+  early_quit_thresholds:
+    abs: 1e-9
+    rel: 1e-9
+
+  img_patch_size: ${policy.img_patch_size}
+  focal_length: ${policy.focal_length}
+
+  logger: ${logger}

configs/pipeline/skelify-refiner.yaml

@@ -0,0 +1,37 @@
+_target_: lib.modeling.optim.SKELify
+
+name: SKELify-Refiner
+
+cfg:
+  skel_model: ${_hub_.models.body_models.skel_mix_hsmr}
+
+  phases:
+    STAGE-refine:
+      max_loop: 10
+      params_keys: ['cam_t', 'poses_orient', 'poses_body', 'betas']
+      optimizer: ${...optimizer}
+      losses:
+        f_normalize_kp2d: True
+        f_normalize_kp2d_to_mean: False
+        w_reprojection: 1.0
+        w_shape_prior: 5.0
+        w_angle_prior: 4.78
+        w_angle_prior_scale: 0.17
+      parts: ['all']
+    # ================================
+
+  optimizer:
+    _target_: torch.optim.LBFGS
+    lr: 1
+    line_search_fn: 'strong_wolfe'
+    tolerance_grad: ${..early_quit_thresholds.abs}
+    tolerance_change: ${..early_quit_thresholds.rel}
+
+  early_quit_thresholds:
+    abs: 1e-7
+    rel: 1e-9
+
+  img_patch_size: ${policy.img_patch_size}
+  focal_length: ${policy.focal_length}
+
+  logger: ${logger}

configs/policy/README.md

@@ -0,0 +1,3 @@
+# policy
+
+Configs here define frequently used values that are shared across multiple sub-modules. They usually cannot belongs to any specific sub-module.

configs/policy/default.yaml

@@ -0,0 +1,5 @@
+img_patch_size: 256
+focal_length: 5000
+img_mean: [0.485, 0.456, 0.406]
+img_std: [0.229, 0.224, 0.225]
+bbox_shape: null

data_inputs/.gitignore

@@ -0,0 +1,2 @@
+body_models
+released_models

data_inputs/description.md

@@ -0,0 +1,21 @@
+<center><b><font size='10'>
+  H<font color='CAB6A5'>S</font><font color='98BACB'>M</font>R
+</font></b></center>
+<center><b><font size='6'>
+  Reconstructing <font color='98BACB'>Humans</font> with a <br> <font color='CAB6A5'>Biomechanically</font> Accurate Skeleton
+</font></b></center>
+<center><b>
+<a href=#>Project Page</a>
+|
+<a href=#>GitHub Repo</a>
+|
+<a href=#>Paper</a>
+</b></center>
+
+> Reconstructing Humans with a Biomechanically Accurate Skeleton <br>
+> [Yan Xia](https://scholar.isshikih.top),
+> [Xiaowei Zhou](https://xzhou.me),
+> [Etienne Vouga](https://www.cs.utexas.edu/~evouga/),
+> [Qixing Huang](https://www.cs.utexas.edu/~huangqx/),
+> [Georgios Pavlakos](https://geopavlakos.github.io/) <br>
+> *CVPR 2025*

lib/__init__.py

@@ -0,0 +1 @@
+from .version import __version__

lib/body_models/abstract_skeletons.py

@@ -0,0 +1,123 @@
+# "Skeletons" here means virtual "bones" between joints. It is used to draw the skeleton on the image.
+
+class Skeleton():
+    bones = []
+    bone_colors = []
+    chains = []
+    parent = []
+
+
+class Skeleton_SMPL24(Skeleton):
+    # The joints definition are copied from
+    # [ROMP](https://github.com/Arthur151/ROMP/blob/4eebd3647f57d291d26423e51f0d514ff7197cb3/romp/lib/constants.py#L58).
+    chains = [
+            [0, 1, 4, 7, 10],         # left leg
+            [0, 2, 5, 8, 11],         # right leg
+            [0, 3, 6, 9, 12, 15],     # spine & head
+            [12, 13, 16, 18, 20, 22], # left arm
+            [12, 14, 17, 19, 21, 23], # right arm
+        ]
+    bones = [
+            [ 0,  1], [ 1,  4], [ 4,  7], [ 7, 10],           # left leg
+            [ 0,  2], [ 2,  5], [ 5,  8], [ 8, 11],           # right leg
+            [ 0,  3], [ 3,  6], [ 6,  9], [ 9, 12], [12, 15], # spine & head
+            [12, 13], [13, 16], [16, 18], [18, 20], [20, 22], # left arm
+            [12, 14], [14, 17], [17, 19], [19, 21], [21, 23], # right arm
+        ]
+    bone_colors = [
+            [127,   0,   0], [148,  21,  21], [169,  41,  41], [191,  63,  63],                  # red
+            [  0, 127,   0], [ 21, 148,  21], [ 41, 169,  41], [ 63, 191,  63],                  # green
+            [  0,   0, 127], [ 15,  15, 143], [ 31,  31, 159], [ 47,  47, 175], [ 63,  63, 191], # blue
+            [  0, 127, 127], [ 15, 143, 143], [ 31, 159, 159], [ 47, 175, 175], [ 63, 191, 191], # cyan
+            [127,   0, 127], [143,  15, 143], [159,  31, 159], [175,  47, 175], [191,  63, 191], # magenta
+        ]
+    parent = [-1, 0, 0, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 12, 12, 12, 13, 14, 16, 17, 18, 19, 20, 21]
+
+
+class Skeleton_SMPL22(Skeleton):
+    chains = [
+            [0, 1, 4, 7, 10],     # left leg
+            [0, 2, 5, 8, 11],     # right leg
+            [0, 3, 6, 9, 12, 15], # spine & head
+            [12, 13, 16, 18, 20], # left arm
+            [12, 14, 17, 19, 21], # right arm
+        ]
+    bones = [
+            [ 0,  1], [ 1,  4], [ 4,  7], [ 7, 10],           # left leg
+            [ 0,  2], [ 2,  5], [ 5,  8], [ 8, 11],           # right leg
+            [ 0,  3], [ 3,  6], [ 6,  9], [ 9, 12], [12, 15], # spine & head
+            [12, 13], [13, 16], [16, 18], [18, 20],           # left arm
+            [12, 14], [14, 17], [17, 19], [19, 21],           # right arm
+        ]
+    bone_colors = [
+            [127,   0,   0], [148,  21,  21], [169,  41,  41], [191,  63,  63],                  # red
+            [  0, 127,   0], [ 21, 148,  21], [ 41, 169,  41], [ 63, 191,  63],                  # green
+            [  0,   0, 127], [ 15,  15, 143], [ 31,  31, 159], [ 47,  47, 175], [ 63,  63, 191], # blue
+            [  0, 127, 127], [ 15, 143, 143], [ 31, 159, 159], [ 47, 175, 175],                  # cyan
+            [127,   0, 127], [143,  15, 143], [159,  31, 159], [175,  47, 175],                  # magenta
+        ]
+    parent = [-1, 0, 0, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 12, 12, 12, 13, 14, 16, 17, 18, 19]
+
+
+class Skeleton_SKEL24(Skeleton):
+    chains = [
+            [ 0,  6,  7,  8,  9, 10], # left leg
+            [ 0,  1,  2,  3,  4,  5], # right leg
+            [ 0, 11, 12, 13],         # spine & head
+            [12, 19, 20, 21, 22, 23], # left arm
+            [12, 14, 15, 16, 17, 18], # right arm
+        ]
+    bones = [
+            [ 0,  6], [ 6,  7], [ 7,  8], [ 8,  9], [ 9, 10], # left leg
+            [ 0,  1], [ 1,  2], [ 2,  3], [ 3,  4], [ 4,  5], # right leg
+            [ 0, 11], [11, 12], [12, 13],                     # spine & head
+            [12, 19], [19, 20], [20, 21], [21, 22], [22, 23], # left arm
+            [12, 14], [14, 15], [15, 16], [16, 17], [17, 18], # right arm
+        ]
+    bone_colors = [
+            [127,   0,   0], [148,  21,  21], [169,  41,  41], [191,  63,  63], [191,  63,  63], # red
+            [  0, 127,   0], [ 21, 148,  21], [ 41, 169,  41], [ 63, 191,  63], [ 63, 191,  63], # green
+            [  0,   0, 127], [ 31,  31, 159], [ 63,  63, 191],                                   # blue
+            [  0, 127, 127], [ 15, 143, 143], [ 31, 159, 159], [ 47, 175, 175], [ 63, 191, 191], # cyan
+            [127,   0, 127], [143,  15, 143], [159,  31, 159], [175,  47, 175], [191,  63, 191], # magenta
+        ]
+    parent = [-1, 0, 1, 2, 3, 4, 0, 6, 7, 8, 9, 0, 11, 12, 12, 19, 20, 21, 22, 12, 14, 15, 16, 17]
+
+
+class Skeleton_OpenPose25(Skeleton):
+    ''' https://www.researchgate.net/figure/Twenty-five-keypoints-of-the-OpenPose-software-model_fig1_374116819 '''
+    chain = [
+            [ 8, 12, 13, 14, 19, 20],  # left leg
+            [14, 21],                  # left heel
+            [ 8,  9, 10, 11, 22, 23],  # right leg
+            [11, 24],                  # right heel
+            [ 8,  1,  0],              # spine & head
+            [ 0, 16, 18],              # left face
+            [ 0, 15, 17],              # right face
+            [ 1,  5,  6,  7],          # left arm
+            [ 1,  2,  3,  4],          # right arm
+        ]
+    bones = [
+            [ 8, 12], [12, 13], [13, 14],  # left leg
+            [14, 19], [19, 20], [14, 21],  # left foot
+            [ 8,  9], [ 9, 10], [10, 11],  # right leg
+            [11, 22], [22, 23], [11, 24],  # right foot
+            [ 8,  1], [ 1,  0],            # spine & head
+            [ 0, 16], [16, 18],            # left face
+            [ 0, 15], [15, 17],            # right face
+            [ 1,  5], [ 5,  6], [ 6,  7],  # left arm
+            [ 1,  2], [ 2,  3], [ 3,  4],  # right arm
+        ]
+    bone_colors = [
+            [ 95,   0, 255], [ 79,   0, 255], [ 83,   0, 255],  # dark blue
+            [ 31,   0, 255], [ 15,   0, 255], [  0,   0, 255],  # dark blue
+            [127, 205, 255], [127, 205, 255], [ 95, 205, 255],  # light blue
+            [ 63, 205, 255], [ 31, 205, 255], [  0, 205, 255],  # light blue
+            [255,   0,   0], [255,   0,   0],                   # red
+            [191,  63,  63], [191,  63, 191],                   # magenta
+            [255,   0, 127], [255,   0, 255],                   # purple
+            [127, 255,   0], [ 63, 255,   0], [  0, 255,   0],  # green
+            [255, 127,   0], [255, 191,   0], [255, 255,   0],  # yellow
+
+        ]
+    parent = [1, 8, 1, 2, 3, 1, 5, 6, -1, 8, 9, 10, 8, 12, 13, 0, 0, 15, 16, 14, 19, 14, 11, 22, 11]

lib/body_models/common.py

@@ -0,0 +1,69 @@
+from lib.kits.basic import *
+
+from smplx import SMPL
+
+from lib.platform import PM
+from lib.body_models.skel_wrapper import SKELWrapper as SKEL
+from lib.body_models.smpl_wrapper import SMPLWrapper
+
+def make_SMPL(gender='neutral', device='cuda:0'):
+    return SMPL(
+        gender = gender,
+        model_path = PM.inputs / 'body_models' / 'smpl',
+    ).to(device)
+
+
+def make_SMPL_hmr2(gender='neutral', device='cuda:0'):
+    ''' SKEL doesn't have neutral model, so align with SKEL, using male. '''
+    return SMPLWrapper(
+        gender = gender,
+        model_path = PM.inputs / 'body_models' / 'smpl',
+        num_body_joints = 23,
+        joint_regressor_extra = PM.inputs / 'body_models/SMPL_to_J19.pkl',
+    ).to(device)
+
+
+
+def make_SKEL(gender='male', device='cuda:0'):
+    ''' We don't have neutral model for SKEL, so use male for now. '''
+    return make_SKEL_mix_joints(gender, device)
+
+
+def make_SKEL_smpl_joints(gender='male', device='cuda:0'):
+    ''' We don't have neutral model for SKEL, so use male for now. '''
+    return SKEL(
+        gender = gender,
+        model_path = PM.inputs / 'body_models' / 'skel',
+        joint_regressor_extra = PM.inputs / 'body_models' / 'SMPL_to_J19.pkl',
+        joint_regressor_custom = PM.inputs / 'body_models' / 'J_regressor_SMPL_MALE.pkl',
+    ).to(device)
+
+
+def make_SKEL_skel_joints(gender='male', device='cuda:0'):
+    ''' We don't have neutral model for SKEL, so use male for now. '''
+    return SKEL(
+        gender = gender,
+        model_path = PM.inputs / 'body_models' / 'skel',
+        joint_regressor_extra = PM.inputs / 'body_models' / 'SMPL_to_J19.pkl',
+    ).to(device)
+
+
+def make_SKEL_mix_joints(gender='male', device='cuda:0'):
+    ''' We don't have neutral model for SKEL, so use male for now. '''
+    return SKEL(
+        gender = gender,
+        model_path = PM.inputs / 'body_models' / 'skel',
+        joint_regressor_extra = PM.inputs / 'body_models' / 'SMPL_to_J19.pkl',
+        joint_regressor_custom = PM.inputs / 'body_models' / 'J_regressor_SKEL_mix_MALE.pkl',
+    ).to(device)
+
+
+def make_SMPLX_moyo(v_template_path:Union[str, Path], batch_size:int=1, device='cuda:0'):
+    from lib.body_models.moyo_smplx_wrapper import MoYoSMPLX
+
+    return MoYoSMPLX(
+        model_path      = PM.inputs / 'body_models' / 'smplx',
+        v_template_path = v_template_path,
+        batch_size      = batch_size,
+        device          = device,
+    )

lib/body_models/moyo_smplx_wrapper.py

@@ -0,0 +1,77 @@
+from lib.kits.basic import *
+
+import smplx
+from psbody.mesh import Mesh
+
+class MoYoSMPLX(smplx.SMPLX):
+
+    def __init__(
+        self,
+        model_path      : Union[str, Path],
+        v_template_path : Union[str, Path],
+        batch_size = 1,
+        n_betas    = 10,
+        device     = 'cpu'
+    ):
+
+        if isinstance(v_template_path, Path):
+            v_template_path = str(v_template_path)
+
+        # Load the `v_template`.
+        v_template_mesh = Mesh(filename=v_template_path)
+        v_template = to_tensor(v_template_mesh.v, device=device)
+
+        self.n_betas = n_betas
+
+        # Create the `body_model_params`.
+        body_model_params = {
+                'model_path'             : model_path,
+                'gender'                 : 'neutral',
+                'v_template'             : v_template.float(),
+                'batch_size'             : batch_size,
+                'create_global_orient'   : True,
+                'create_body_pose'       : True,
+                'create_betas'           : True,
+                'num_betas'              : self.n_betas,  # They actually don't use num_betas.
+                'create_left_hand_pose'  : True,
+                'create_right_hand_pose' : True,
+                'create_expression'      : True,
+                'create_jaw_pose'        : True,
+                'create_leye_pose'       : True,
+                'create_reye_pose'       : True,
+                'create_transl'          : True,
+                'use_pca'                : False,
+                'flat_hand_mean'         : True,
+                'dtype'                  : torch.float32,
+            }
+
+        super().__init__(**body_model_params)
+        self = self.to(device)
+
+    def forward(self, **kwargs):
+        ''' Only all parameters are passed, the batch_size will be flexible adjusted. '''
+        assert 'global_orient' in kwargs, '`global_orient` is required for the forward pass.'
+        assert 'body_pose' in kwargs, '`body_pose` is required for the forward pass.'
+        B = kwargs['global_orient'].shape[0]
+        body_pose = kwargs['body_pose']
+
+        if 'left_hand_pose' not in kwargs:
+            kwargs['left_hand_pose'] = body_pose.new_zeros((B, 45))
+            get_logger().warning('`left_hand_pose` is not provided, but it\'s expected, set to zeros.')
+        if 'right_hand_pose' not in kwargs:
+            kwargs['right_hand_pose'] = body_pose.new_zeros((B, 45))
+            get_logger().warning('`left_hand_pose` is not provided, but it\'s expected, set to zeros.')
+        if 'transl' not in kwargs:
+            kwargs['transl'] = body_pose.new_zeros((B, 3))
+        if 'betas' not in kwargs:
+            kwargs['betas'] = body_pose.new_zeros((B, self.n_betas))
+        if 'expression' not in kwargs:
+            kwargs['expression'] = body_pose.new_zeros((B, 10))
+        if 'jaw_pose' not in kwargs:
+            kwargs['jaw_pose'] = body_pose.new_zeros((B, 3))
+        if 'leye_pose' not in kwargs:
+            kwargs['leye_pose'] = body_pose.new_zeros((B, 3))
+        if 'reye_pose' not in kwargs:
+            kwargs['reye_pose'] = body_pose.new_zeros((B, 3))
+
+        return super().forward(**kwargs)

lib/body_models/skel/alignment/align_config.py

@@ -0,0 +1,72 @@
+""" Optimization config file. In 'optim_steps' we define each optimization steps.
+Each step inherit and overwrite the parameters of the previous step."""
+
+config = {
+    'keepalive_meshviewer': False,
+    'optim_steps':
+        [
+            {
+            'description' : 'Adjust the root orientation and translation',
+            'use_basic_loss': True,
+            'lr': 1,
+            'max_iter': 20,
+            'num_steps': 10,
+            'line_search_fn': 'strong_wolfe', #'strong_wolfe',
+            'tolerance_change': 1e-7,# 1e-4, #0.01
+            'mode' : 'root_only',
+
+            'l_verts_loose': 300,
+            'l_time_loss': 0,#5e2,
+
+            'l_joint': 0.0,
+            'l_verts': 0,
+            'l_scapula_loss': 0.0,
+            'l_spine_loss': 0.0,
+            'l_pose_loss': 0.0,
+
+
+            },
+            # Adjust the upper limbs
+            {
+            'description' : 'Adjust the upper limbs pose',
+            'lr': 0.1,
+            'max_iter': 20,
+            'num_steps': 10,
+            'tolerance_change': 1e-7,
+            'mode' : 'fixed_upper_limbs', #'fixed_root',
+
+            'l_verts_loose': 600,
+            'l_joint': 1e3,
+            'l_time_loss': 0,# 5e2,
+            'l_pose_loss': 1e-4,
+            },
+            # Adjust the whole body
+            {
+            'description' : 'Adjust the whole body pose with fixed root',
+            'lr': 0.1,
+            'max_iter': 20,
+            'num_steps': 10,
+            'tolerance_change': 1e-7,
+            'mode' : 'fixed_root', #'fixed_root',
+
+            'l_verts_loose': 600,
+            'l_joint': 1e3,
+            'l_time_loss': 0,
+            'l_pose_loss': 1e-4,
+            },
+            #
+            {
+            'description' : 'Free optimization',
+            'lr': 0.1,
+            'max_iter': 20,
+            'num_steps': 10,
+            'tolerance_change': 1e-7,
+            'mode' : 'free', #'fixed_root',
+
+            'l_verts_loose': 600,
+            'l_joint': 1e3,
+            'l_time_loss':0,
+            'l_pose_loss': 1e-4,
+            },
+    ]
+}

lib/body_models/skel/alignment/align_config_joint.py

@@ -0,0 +1,72 @@
+""" Optimization config file. In 'optim_steps' we define each optimization steps.
+Each step inherit and overwrite the parameters of the previous step."""
+
+config = {
+    'keepalive_meshviewer': False,
+    'optim_steps':
+        [
+            {
+            'description' : 'Adjust the root orientation and translation',
+            'use_basic_loss': True,
+            'lr': 1,
+            'max_iter': 20,
+            'num_steps': 10,
+            'line_search_fn': 'strong_wolfe', #'strong_wolfe',
+            'tolerance_change': 1e-7,# 1e-4, #0.01
+            'mode' : 'root_only',
+
+            'l_verts_loose': 300,
+            'l_time_loss': 0,#5e2,
+
+            'l_joint': 0.0,
+            'l_verts': 0,
+            'l_scapula_loss': 0.0,
+            'l_spine_loss': 0.0,
+            'l_pose_loss': 0.0,
+
+
+            },
+            # Adjust the upper limbs
+            {
+            'description' : 'Adjust the upper limbs pose',
+            'lr': 0.1,
+            'max_iter': 20,
+            'num_steps': 10,
+            'tolerance_change': 1e-7,
+            'mode' : 'fixed_upper_limbs', #'fixed_root',
+
+            'l_verts_loose': 600,
+            'l_joint': 2e3,
+            'l_time_loss': 0,# 5e2,
+            'l_pose_loss': 1e-3,
+            },
+            # Adjust the whole body
+            {
+            'description' : 'Adjust the whole body pose with fixed root',
+            'lr': 0.1,
+            'max_iter': 20,
+            'num_steps': 10,
+            'tolerance_change': 1e-7,
+            'mode' : 'fixed_root', #'fixed_root',
+
+            'l_verts_loose': 600,
+            'l_joint': 1e3,
+            'l_time_loss': 0,
+            'l_pose_loss': 1e-4,
+            },
+            #
+            {
+            'description' : 'Free optimization',
+            'lr': 0.1,
+            'max_iter': 20,
+            'num_steps': 10,
+            'tolerance_change': 1e-7,
+            'mode' : 'free', #'fixed_root',
+
+            'l_verts_loose': 300,
+            'l_joint': 2e3,
+            'l_time_loss':0,
+            'l_pose_loss': 1e-4,
+            },
+    ]
+}

lib/body_models/skel/alignment/aligner.py

@@ -0,0 +1,469 @@
+
+"""
+Copyright©2023 Max-Planck-Gesellschaft zur Förderung
+der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+for Intelligent Systems. All rights reserved.
+
+Author: Soyong Shin, Marilyn Keller
+See https://skel.is.tue.mpg.de/license.html for licensing and contact information.
+"""
+
+import traceback
+import math
+import os
+import pickle
+import torch
+import smplx
+import omegaconf
+import torch.nn.functional as F
+from psbody.mesh import Mesh, MeshViewer, MeshViewers
+from tqdm import trange
+from pathlib import Path
+
+import lib.body_models.skel.config as cg
+from lib.body_models.skel.skel_model import SKEL
+from .losses import compute_anchor_pose, compute_anchor_trans, compute_pose_loss, compute_scapula_loss, compute_spine_loss, compute_time_loss, pretty_loss_print
+from .utils import location_to_spheres, to_numpy, to_params, to_torch
+from .align_config import config
+from .align_config_joint import config as config_joint
+
+class SkelFitter(object):
+
+    def __init__(self, gender, device, num_betas=10, export_meshes=False, joint_optim=False) -> None:
+
+        self.smpl = smplx.create(cg.smpl_folder, model_type='smpl', gender=gender, num_betas=num_betas, batch_size=1, export_meshes=False).to(device)
+        self.skel = SKEL(gender).to(device)
+        self.gender = gender
+        self.device = device
+        self.num_betas = num_betas
+        # Instanciate masks used for the vertex to vertex fitting
+        fitting_mask_file = Path(__file__).parent / 'riggid_parts_mask.pkl'
+        fitting_indices = pickle.load(open(fitting_mask_file, 'rb'))
+        fitting_mask = torch.zeros(6890, dtype=torch.bool, device=self.device)
+        fitting_mask[fitting_indices] = 1
+        self.fitting_mask = fitting_mask.reshape(1, -1, 1).to(self.device) # 1xVx1 to be applied to verts that are BxVx3
+
+        smpl_torso_joints = [0,3]
+        verts_mask = (self.smpl.lbs_weights[:,smpl_torso_joints]>0.5).sum(dim=-1)>0
+        self.torso_verts_mask = verts_mask.unsqueeze(0).unsqueeze(-1) # Because verts are of shape BxVx3
+
+        self.export_meshes = export_meshes
+
+
+        # make the cfg being an object using omegaconf
+        if joint_optim:
+            self.cfg = omegaconf.OmegaConf.create(config_joint)
+        else:
+            self.cfg = omegaconf.OmegaConf.create(config)
+
+        # Instanciate the mesh viewer to visualize the fitting
+        if('DISABLE_VIEWER' in os.environ):
+            self.mv = None
+            print("\n DISABLE_VIEWER flag is set, running in headless mode")
+        else:
+            self.mv = MeshViewers((1,2),  keepalive=self.cfg.keepalive_meshviewer)
+
+
+    def run_fit(self,
+            trans_in,
+            betas_in,
+            poses_in,
+            batch_size=20,
+            skel_data_init=None,
+            force_recompute=False,
+            debug=False,
+            watch_frame=0,
+            freevert_mesh=None,
+            opt_sequence=False,
+            fix_poses=False,
+            variant_exp=''):
+        """Align SKEL to a SMPL sequence."""
+
+        self.nb_frames = poses_in.shape[0]
+        self.watch_frame = watch_frame
+        self.is_skel_data_init = skel_data_init is not None
+        self.force_recompute = force_recompute
+
+        print('Fitting {} frames'.format(self.nb_frames))
+        print('Watching frame: {}'.format(watch_frame))
+
+        # Initialize SKEL torch params
+        body_params = self._init_params(betas_in, poses_in, trans_in, skel_data_init, variant_exp)
+
+        # We cut the whole sequence in batches for parallel optimization
+        if batch_size > self.nb_frames:
+            batch_size = self.nb_frames
+            print('Batch size is larger than the number of frames. Setting batch size to {}'.format(batch_size))
+
+        n_batch = math.ceil(self.nb_frames/batch_size)
+        pbar = trange(n_batch, desc='Running batch optimization')
+
+        # Initialize the res dict to store the per frame result skel parameters
+        out_keys = ['poses', 'betas', 'trans']
+        if self.export_meshes:
+            out_keys += ['skel_v', 'skin_v', 'smpl_v']
+        res_dict = {key: [] for key in out_keys}
+
+        res_dict['gender'] = self.gender
+        if self.export_meshes:
+            res_dict['skel_f'] = self.skel.skel_f.cpu().numpy().copy()
+            res_dict['skin_f'] = self.skel.skin_f.cpu().numpy().copy()
+            res_dict['smpl_f'] = self.smpl.faces
+
+        # Iterate over the batches to fit the whole sequence
+        for i in pbar:
+
+            if debug:
+                # Only run the first batch to test, ignore the rest
+                if i > 1:
+                    continue
+
+            # Get batch start and end indices
+            i_start = i * batch_size
+            i_end = min((i+1) * batch_size, self.nb_frames)
+
+            # Fit the batch
+            betas, poses, trans, verts = self._fit_batch(body_params, i, i_start, i_end, enable_time=opt_sequence, fix_poses=fix_poses)
+            # if torch.isnan(betas).any() \
+            # or torch.isnan(poses).any() \
+            # or torch.isnan(trans).any():
+            #     print(f'Nan values detected.')
+            #     raise ValueError('Nan values detected in the output.')
+
+            # Store ethe results
+            res_dict['poses'].append(poses)
+            res_dict['betas'].append(betas)
+            res_dict['trans'].append(trans)
+            if self.export_meshes:
+                # Store the meshes vertices
+                skel_output = self.skel.forward(poses=poses, betas=betas, trans=trans, poses_type='skel', skelmesh=True)
+                res_dict['skel_v'].append(skel_output.skel_verts)
+                res_dict['skin_v'].append(skel_output.skin_verts)
+                res_dict['smpl_v'].append(verts)
+
+            if opt_sequence:
+                # Initialize the next frames with current frame
+                body_params['poses_skel'][i_end:] = poses[-1:].detach()
+                body_params['trans_skel'][i_end:] = trans[-1].detach()
+                body_params['betas_skel'][i_end:] = betas[-1:].detach()
+
+        # Concatenate the batches and convert to numpy
+        for key, val in res_dict.items():
+            if isinstance(val, list):
+                res_dict[key] = torch.cat(val, dim=0).detach().cpu().numpy()
+
+        return res_dict
+
+    def _init_params(self, betas_smpl, poses_smpl, trans_smpl, skel_data_init=None, variant_exp=''):
+        """ Return initial SKEL parameters from SMPL data dictionary and an optional SKEL data dictionary."""
+
+        if skel_data_init is None or self.force_recompute:
+
+            poses_skel = torch.zeros((self.nb_frames, self.skel.num_q_params), device=self.device)
+            if variant_exp == '' or variant_exp == '_official_old':
+                poses_skel[:, :3] = poses_smpl[:, :3] # Global orient are similar between SMPL and SKEL, so init with SMPL angles
+            elif variant_exp == '_official_fix':
+                # https://github.com/MarilynKeller/SKEL/commit/d1f6ff62235c142ba010158e00e21fd4fe25807f#diff-09188717a56a42e9589e9bd289f9ddb4fb53160e03c81a7ced70b3a84c1d9d0bR157
+                pass
+            elif variant_exp == '_my_fix':
+                gt_orient_aa = poses_smpl[:, :3]
+                # IMPORTANT: The alignment comes from `exp/inspect_skel/archive/orientation.py`.
+                from lib.utils.geometry.rotation import axis_angle_to_matrix, matrix_to_euler_angles
+                gt_orient_mat = axis_angle_to_matrix(gt_orient_aa)
+                gt_orient_ea = matrix_to_euler_angles(gt_orient_mat, 'YXZ')
+                flip = torch.tensor([-1, 1, 1], device=self.device)
+                poses_skel[:, :3] = gt_orient_ea[:, [2, 1, 0]] * flip
+            else:
+                raise ValueError(f'Unknown variant_exp {variant_exp}')
+
+            betas_skel = torch.zeros((self.nb_frames, 10), device=self.device)
+            betas_skel[:] = betas_smpl[..., :10]
+
+            trans_skel = trans_smpl # Translation is similar between SMPL and SKEL, so init with SMPL translation
+
+        else:
+            # Load from previous alignment
+            betas_skel = to_torch(skel_data_init['betas'], self.device)
+            poses_skel = to_torch(skel_data_init['poses'], self.device)
+            trans_skel = to_torch(skel_data_init['trans'], self.device)
+
+        # Make a dictionary out of the necessary body parameters
+        body_params = {
+            'betas_skel': betas_skel,
+            'poses_skel': poses_skel,
+            'trans_skel': trans_skel,
+            'betas_smpl': betas_smpl,
+            'poses_smpl': poses_smpl,
+            'trans_smpl': trans_smpl
+        }
+
+        return body_params
+
+
+    def _fit_batch(self, body_params, i, i_start, i_end, enable_time=False, fix_poses=False):
+        """ Create parameters for the batch and run the optimization."""
+
+        # Sample a batch ver
+        body_params = { key: val[i_start:i_end] for key, val in body_params.items()}
+
+        # SMPL params
+        betas_smpl = body_params['betas_smpl']
+        poses_smpl = body_params['poses_smpl']
+        trans_smpl = body_params['trans_smpl']
+
+        # SKEL params
+        betas = to_params(body_params['betas_skel'], device=self.device)
+        poses = to_params(body_params['poses_skel'], device=self.device)
+        trans = to_params(body_params['trans_skel'], device=self.device)
+
+        if 'verts' in body_params:
+            verts = body_params['verts']
+        else:
+            # Run a SMPL forward pass to get the SMPL body vertices
+            smpl_output = self.smpl(betas=betas_smpl, body_pose=poses_smpl[:,3:], transl=trans_smpl, global_orient=poses_smpl[:,:3])
+            verts = smpl_output.vertices
+
+        # Optimize
+        config = self.cfg.optim_steps
+        current_cfg = config[0]
+
+        # from lib.kits.debug import set_trace
+        # set_trace()
+
+        try:
+            if fix_poses:
+                # for ci, cfg in enumerate(config[1:]):
+                for ci, cfg in enumerate([config[-1]]): # To debug, only run the last step
+                    current_cfg.update(cfg)
+                    print(f'Step {ci+1}: {current_cfg.description}')
+                    self._optim([trans,betas], poses, betas, trans, verts, current_cfg, enable_time)
+            else:
+                if not enable_time or not self.is_skel_data_init:
+                    # Optimize the global rotation and translation for the initial fitting
+                    print(f'Step 0: {current_cfg.description}')
+                    self._optim([trans,poses], poses, betas, trans, verts, current_cfg, enable_time)
+
+                for ci, cfg in enumerate(config[1:]):
+                # for ci, cfg in enumerate([config[-1]]): # To debug, only run the last step
+                    current_cfg.update(cfg)
+                    print(f'Step {ci+1}: {current_cfg.description}')
+                    self._optim([poses], poses, betas, trans, verts, current_cfg, enable_time)
+
+
+                # # Refine by optimizing the whole body
+                # cfg.update(self.cfg_optim[])
+                # cfg.update({'mode' : 'free', 'tolerance_change': 0.0001, 'l_joint': 0.2e4})
+                # self._optim([trans, poses], poses, betas, trans, verts, cfg)
+        except Exception as e:
+            print(e)
+            traceback.print_exc()
+            # from lib.kits.debug import set_trace
+            # set_trace()
+
+        return betas, poses, trans, verts
+
+    def _optim(self,
+            params,
+            poses,
+            betas,
+            trans,
+            verts,
+            cfg,
+            enable_time=False,
+            ):
+
+            # regress anatomical joints from SMPL's vertices
+            anat_joints = torch.einsum('bik,ji->bjk', [verts, self.skel.J_regressor_osim])
+            dJ=torch.zeros((poses.shape[0], 24, 3), device=betas.device)
+
+            # Create the optimizer
+            optimizer = torch.optim.LBFGS(params,
+                                          lr=cfg.lr,
+                                          max_iter=cfg.max_iter,
+                                          line_search_fn=cfg.line_search_fn,
+                                          tolerance_change=cfg.tolerance_change)
+
+            poses_init = poses.detach().clone()
+            trans_init = trans.detach().clone()
+
+            def closure():
+                optimizer.zero_grad()
+
+                # fi = self.watch_frame #frame of the batch to display
+                # output = self.skel.forward(poses=poses[fi:fi+1],
+                #                             betas=betas[fi:fi+1],
+                #                             trans=trans[fi:fi+1],
+                #                             poses_type='skel',
+                #                             dJ=dJ[fi:fi+1],
+                #                             skelmesh=True)
+
+                # self._fstep_plot(output, cfg, verts[fi:fi+1], anat_joints[fi:fi+1], )
+
+                loss_dict = self._fitting_loss(poses,
+                                        poses_init,
+                                        betas,
+                                        trans,
+                                        trans_init,
+                                        dJ,
+                                        anat_joints,
+                                        verts,
+                                        cfg,
+                                        enable_time)
+
+                # print(pretty_loss_print(loss_dict))
+
+                loss = sum(loss_dict.values())
+                loss.backward()
+
+                return loss
+
+            for step_i in range(cfg.num_steps):
+                loss = optimizer.step(closure).item()
+
+    def _get_masks(self, cfg):
+        pose_mask = torch.ones((self.skel.num_q_params)).to(self.device).unsqueeze(0)
+        verts_mask = torch.ones_like(self.fitting_mask)
+        joint_mask = torch.ones((self.skel.num_joints, 3)).to(self.device).unsqueeze(0).bool()
+
+        # Mask vertices
+        if cfg.mode=='root_only':
+            # Only optimize the global rotation of the body, i.e. the first 3 angles of the pose
+            pose_mask[:] = 0 # Only optimize for the global rotation
+            pose_mask[:,:3] = 1
+            # Only fit the thorax vertices to recover the proper body orientation and translation
+            verts_mask = self.torso_verts_mask
+
+        elif cfg.mode=='fixed_upper_limbs':
+            upper_limbs_joints = [0,1,2,3,6,9,12,15,17]
+            verts_mask = (self.smpl.lbs_weights[:,upper_limbs_joints]>0.5).sum(dim=-1)>0
+            verts_mask = verts_mask.unsqueeze(0).unsqueeze(-1)
+
+            joint_mask[:, [3,4,5,8,9,10,18,23], :] = 0 # Do not try to match the joints of the upper limbs
+
+            pose_mask[:] = 1
+            pose_mask[:,:3] = 0    # Block the global rotation
+            pose_mask[:,19] = 0  # block the lumbar twist
+            # pose_mask[:, 36:39] = 0
+            # pose_mask[:, 43:46] = 0
+            # pose_mask[:, 62:65] = 0
+            # pose_mask[:, 62:65] = 0
+
+        elif cfg.mode=='fixed_root':
+            pose_mask[:] = 1
+            pose_mask[:,:3] = 0  # Block the global rotation
+            # pose_mask[:,19] = 0  # block the lumbar twist
+
+            # The orientation of the upper limbs is often wrong in SMPL so ignore these vertices for the finale step
+            upper_limbs_joints = [1,2,16,17]
+            verts_mask = (self.smpl.lbs_weights[:,upper_limbs_joints]>0.5).sum(dim=-1)>0
+            verts_mask = torch.logical_not(verts_mask)
+            verts_mask = verts_mask.unsqueeze(0).unsqueeze(-1)
+
+        elif cfg.mode=='free':
+            verts_mask = torch.ones_like(self.fitting_mask )
+
+            joint_mask[:]=0
+            joint_mask[:, [19,14], :] = 1 # Only fir the scapula join to avoid collapsing shoulders
+
+        else:
+            raise ValueError(f'Unknown mode {cfg.mode}')
+
+        return pose_mask, verts_mask, joint_mask
+
+    def _fitting_loss(self,
+                    poses,
+                    poses_init,
+                    betas,
+                    trans,
+                    trans_init,
+                    dJ,
+                    anat_joints,
+                    verts,
+                    cfg,
+                    enable_time=False):
+
+        loss_dict = {}
+
+
+        pose_mask, verts_mask, joint_mask = self._get_masks(cfg)
+        poses = poses * pose_mask + poses_init * (1-pose_mask)
+
+        # Mask joints to not optimize before computing the losses
+
+        output = self.skel.forward(poses=poses, betas=betas, trans=trans, poses_type='skel', dJ=dJ, skelmesh=False)
+
+        # Fit the SMPL vertices
+        # We know the skinning of the forearm and the neck are not perfect,
+        # so we create a mask of the SMPL vertices that are important to fit, like the hands and the head
+        loss_dict['verts_loss_loose'] = cfg.l_verts_loose * (verts_mask  * (output.skin_verts - verts)**2).sum() / (((verts_mask).sum()*self.nb_frames))
+
+        # Fit the regressed joints, this avoids collapsing shoulders
+        # loss_dict['joint_loss'] = cfg.l_joint * F.mse_loss(output.joints, anat_joints)
+        loss_dict['joint_loss'] = cfg.l_joint * (joint_mask * (output.joints - anat_joints)**2).mean()
+
+        # Time consistancy
+        if poses.shape[0] > 1 and enable_time:
+            # This avoids unstable hips orientationZ
+            loss_dict['time_loss'] = cfg.l_time_loss * F.mse_loss(poses[1:], poses[:-1])
+
+        loss_dict['pose_loss'] = cfg.l_pose_loss * compute_pose_loss(poses, poses_init)
+
+        if cfg.use_basic_loss is False:
+            # These losses can be used to regularize the optimization but are not always necessary
+            loss_dict['anch_rot'] = cfg.l_anch_pose * compute_anchor_pose(poses, poses_init)
+            loss_dict['anch_trans'] = cfg.l_anch_trans * compute_anchor_trans(trans, trans_init)
+
+            loss_dict['verts_loss'] = cfg.l_verts * (verts_mask * self.fitting_mask * (output.skin_verts - verts)**2).sum() / (self.fitting_mask*verts_mask).sum()
+
+            # Regularize the pose
+            loss_dict['scapula_loss'] = cfg.l_scapula_loss * compute_scapula_loss(poses)
+            loss_dict['spine_loss'] = cfg.l_spine_loss * compute_spine_loss(poses)
+
+            # Adjust the losses of all the pose regularizations sub losses with the pose_reg_factor value
+            for key in ['scapula_loss', 'spine_loss', 'pose_loss']:
+                loss_dict[key] = cfg.pose_reg_factor * loss_dict[key]
+
+        return loss_dict
+
+    def _fstep_plot(self, output, cfg, verts, anat_joints):
+        "Function to plot each step"
+
+        if('DISABLE_VIEWER' in os.environ):
+            return
+
+        pose_mask, verts_mask, joint_mask = self._get_masks(cfg)
+
+        skin_err_value = ((output.skin_verts[0] - verts[0])**2).sum(dim=-1).sqrt()
+        skin_err_value = skin_err_value / 0.05
+        skin_err_value = to_numpy(skin_err_value)
+
+        skin_mesh = Mesh(v=to_numpy(output.skin_verts[0]), f=[], vc='white')
+        skel_mesh = Mesh(v=to_numpy(output.skel_verts[0]), f=self.skel.skel_f.cpu().numpy(), vc='white')
+
+        # Display vertex distance on SMPL
+        smpl_verts = to_numpy(verts[0])
+        smpl_mesh = Mesh(v=smpl_verts, f=self.smpl.faces)
+        smpl_mesh.set_vertex_colors_from_weights(skin_err_value, scale_to_range_1=False)
+
+        smpl_mesh_masked = Mesh(v=smpl_verts[to_numpy(verts_mask[0,:,0])], f=[], vc='green')
+        smpl_mesh_pc = Mesh(v=smpl_verts, f=[], vc='green')
+
+        skin_mesh_err = Mesh(v=to_numpy(output.skin_verts[0]), f=self.skel.skin_f.cpu().numpy(), vc='white')
+        skin_mesh_err.set_vertex_colors_from_weights(skin_err_value, scale_to_range_1=False)
+        # List the meshes to display
+        meshes_left = [skin_mesh_err, smpl_mesh_pc]
+        meshes_right = [smpl_mesh_masked, skin_mesh, skel_mesh]
+
+        if cfg.l_joint > 0:
+            # Plot the joints
+            meshes_right += location_to_spheres(to_numpy(output.joints[joint_mask[:,:,0]]), color=(1,0,0), radius=0.02)
+            meshes_right += location_to_spheres(to_numpy(anat_joints[joint_mask[:,:,0]]), color=(0,1,0), radius=0.02) \
+
+
+        self.mv[0][0].set_dynamic_meshes(meshes_left)
+        self.mv[0][1].set_dynamic_meshes(meshes_right)
+
+        # print(poses[frame_to_watch, :3])
+        # print(trans[frame_to_watch])
+        # print(betas[frame_to_watch, :3])
+        # mv.get_keypress()

lib/body_models/skel/alignment/losses.py

@@ -0,0 +1,47 @@
+import torch
+
+def compute_scapula_loss(poses):
+    
+    scapula_indices = [26, 27, 28, 36, 37, 38]
+    
+    scapula_poses = poses[:, scapula_indices]
+    scapula_loss = torch.linalg.norm(scapula_poses, ord=2)
+    return scapula_loss
+
+def compute_spine_loss(poses):
+    
+    spine_indices = range(17, 25)
+    
+    spine_poses = poses[:, spine_indices]
+    spine_loss = torch.linalg.norm(spine_poses, ord=2)
+    return spine_loss
+
+def compute_pose_loss(poses, pose_init):
+    
+    pose_loss = torch.linalg.norm(poses[:, 3:], ord=2) # The global rotation should not be constrained
+    return pose_loss
+
+def compute_anchor_pose(poses, pose_init):
+    
+    pose_loss = torch.nn.functional.mse_loss(poses[:, :3], pose_init[:, :3])
+    return pose_loss 
+
+def compute_anchor_trans(trans, trans_init):
+
+    trans_loss = torch.nn.functional.mse_loss(trans, trans_init)
+    return trans_loss 
+
+def compute_time_loss(poses):
+    
+    pose_delta = poses[1:] - poses[:-1]
+    time_loss = torch.linalg.norm(pose_delta, ord=2)
+    return time_loss
+
+def pretty_loss_print(loss_dict):
+    # Pretty print the loss on the form loss val | loss1 val1 | loss2 val2 
+    # Start with the total loss
+    loss = sum(loss_dict.values())
+    pretty_loss = f'{loss:.4f}'
+    for key, val in loss_dict.items():
+        pretty_loss += f' | {key} {val:.4f}'
+    return pretty_loss   

lib/body_models/skel/alignment/utils.py

@@ -0,0 +1,119 @@
+
+import os
+import pickle
+import torch
+import numpy as np
+from psbody.mesh.sphere import Sphere
+
+# to_params = lambda x: torch.from_numpy(x).float().to(self.device).requires_grad_(True)
+# to_torch = lambda x: torch.from_numpy(x).float().to(self.device)
+
+def to_params(x, device):
+    return x.to(device).requires_grad_(True)
+
+def to_torch(x, device):
+    return torch.from_numpy(x).float().to(device)
+
+def to_numpy(x):
+    return x.detach().cpu().numpy()
+
+def load_smpl_seq(smpl_seq_path, gender=None, straighten_hands=False):
+
+    if not os.path.exists(smpl_seq_path):
+        raise Exception('Path does not exist: {}'.format(smpl_seq_path))
+    
+    if smpl_seq_path.endswith('.pkl'):
+        data_dict = pickle.load(open(smpl_seq_path, 'rb'))
+    
+    elif smpl_seq_path.endswith('.npz'):
+        data_dict = np.load(smpl_seq_path, allow_pickle=True)
+        
+        if data_dict.files == ['pred_smpl_parms', 'verts', 'pred_cam_t']:
+            data_dict = data_dict['pred_smpl_parms'].item()# ['global_orient', 'body_pose', 'body_pose_axis_angle', 'global_orient_axis_angle', 'betas']
+        else:
+            data_dict = {key: data_dict[key] for key in data_dict.keys()} # convert to python dict
+    else:
+        raise Exception('Unknown file format: {}. Supported formats are .pkl and .npz'.format(smpl_seq_path))
+        
+    # Instanciate a dictionary with the keys expected by the fitter
+    data_fixed = {}  
+    
+    # Get gender     
+    if 'gender' not in data_dict:
+        assert gender is not None, f"The provided SMPL data dictionary does not contain gender, you need to pass it in command line"
+        data_fixed['gender'] = gender
+    elif not isinstance(data_dict['gender'], str):
+            # In some npz, the gender type happens to be: array('male', dtype='<U4'). So we convert it to string
+            data_fixed['gender'] = str(data_dict['gender'])
+    else:
+        data_fixed['gender'] = gender
+            
+    # convert tensors to numpy arrays 
+    for key, val in data_dict.items():
+        if isinstance(val, torch.Tensor):
+            data_dict[key] = val.detach().cpu().numpy()
+
+    # Get the SMPL pose
+    if 'poses' in data_dict: 
+        poses = data_dict['poses']
+    elif 'body_pose_axis_angle' in data_dict and 'global_orient_axis_angle' in data_dict:
+        # assert 'global_orient' in data_dict and 'body_pose' in data_dict, f"Could not find poses in {smpl_seq_path}. Available keys: {data_dict.keys()})"
+        poses = np.concatenate([data_dict['global_orient_axis_angle'], data_dict['body_pose_axis_angle']], axis=1)
+        poses = poses.reshape(-1, 72)
+    elif 'body_pose' in data_dict and 'global_orient' in data_dict:
+        poses = np.concatenate([data_dict['global_orient_axis_angle'], data_dict['body_pose_axis_angle']], axis=-1)
+    else: 
+        raise Exception(f"Could not find poses in {smpl_seq_path}. Available keys: {data_dict.keys()})")
+        
+    if poses.shape[1] == 156:
+        # Those are SMPL+H poses, we remove the hand poses to keep only the body poses
+        smpl_poses = np.zeros((poses.shape[0], 72))
+        smpl_poses[:, :72-2*3] = poses[:, :72-2*3] # We leave params for SMPL joints 22 and 23 to zero as these DOF are not present in SMPLH
+        poses = smpl_poses
+    
+    # Set SMPL joints 22 and 23 to zero as SKEL has rigid hands 
+    if straighten_hands:      
+        poses[:, 72-2*3:] = 0
+        
+    data_fixed['poses'] = poses
+        
+    # Translation
+    if 'trans' not in data_dict:
+        data_fixed['trans'] = np.zeros((poses.shape[0], 3))
+    else:
+        data_fixed['trans'] = data_dict['trans']
+        
+    # Get betas 
+    betas = data_dict['betas'][..., :10] # Keep only the 10 first betas
+    if len(betas.shape) == 1 and len(poses.shape) == 2:
+        betas = betas[None, :] # Add a batch dimension
+    data_fixed['betas'] = betas
+     
+    for key in ['trans', 'poses', 'betas', 'gender']:
+        assert key in data_fixed.keys(), f'Could not find {key} in {smpl_seq_path}. Available keys: {data_fixed.keys()})'
+        
+    out_dict = {}
+    out_dict['trans'] = data_fixed['trans']
+    out_dict['poses'] = data_fixed['poses']
+    out_dict['betas'] = data_fixed['betas']
+    out_dict['gender'] = data_fixed['gender']
+    
+    return out_dict
+        
+        
+def location_to_spheres(loc, color=(1,0,0), radius=0.02):
+    """Given an array of 3D points, return a list of spheres located at those positions.
+
+    Args:
+        loc (numpy.array): Nx3 array giving 3D positions
+        color (tuple, optional): One RGB float color vector to color the spheres. Defaults to (1,0,0).
+        radius (float, optional): Radius of the spheres in meters. Defaults to 0.02.
+
+    Returns:
+        list: List of spheres Mesh
+    """
+
+    cL = [Sphere(np.asarray([loc[i, 0], loc[i, 1], loc[i, 2]]), radius).to_mesh() for i in range(loc.shape[0])]
+    for spL in cL:
+        spL.set_vertex_colors(np.array(color)) 
+    return cL

lib/body_models/skel/config.py

@@ -0,0 +1,2 @@
+skel_folder = '/u/yanxia/workspace/data/body_models/skel_models_v1.1'
+smpl_folder = '/data/yanxia/code/SKEL/models/'

lib/body_models/skel/fit_osim/Fitting_SKEL_to_osim.md

@@ -0,0 +1,20 @@
+# Fitting SKEL to OpenSim
+
+Here we explain how to fit the SKEL model to an OpenSim model sequence to get a surface mesh. We do not provide fitting code yet but if you are interested in this direction, this folder should give you a good starting point.
+
+As an example, we will use the OpenSim model from https://simtk.org/projects/full_body.
+Download the model and sample data from (this page)[https://simtk.org/projects/full_body]. You should get the zip folder: `FullBodyModelwSampleSim.-latest.zip` and unzip it. This folder contains a .osim OpenSim model and example .mot motion sequences.
+
+You can visualize the sequence with: 
+
+```bash
+python examples/load_osim.py --osim /path/to/FullBodyModelwSampleSim.-latest/ModelWithSampleSimulations-4.0/SimulationDataAndSetupFiles-4.0/Rajagopal2015.osim --mot /path/to/FullBodyModelwSampleSim.-latest/ModelWithSampleSimulations-4.0/SimulationDataAndSetupFiles-4.0/IK/results_run/ik_output_run.mot
+```
+
+To fit SKEL to this sequence, we need to map each joint of this OpenSim model to a joint of the SKEL model. For the current model, the mapping is in `SKEL/skel/fit_osim/mappings/full_body.yaml`. If you use another model, you will need to create your own mapping yaml for that model.
+
+The following script will then show you how to get the needed joints trajectories to fit SKEL to this OpenSim sequence:
+
+```
+python SKEL/skel/fit_osim/osim_fitter.py
+```

lib/body_models/skel/fit_osim/mappings/full_body.yaml

@@ -0,0 +1,29 @@
+# Mapping to SKEL
+
+joints_mapping: # For each joint of the OpenSim model, specify the corresponding joint in SKEL
+  pelvis : pelvis
+  femur_r : femur_r
+  tibia_r : tibia_r
+  talus_r : talus_r
+  calcn_r : calcn_r
+  toes_r : toes_r
+  femur_l : femur_l
+  tibia_l : tibia_l
+  talus_l : talus_l
+  calcn_l : calcn_l
+  toes_l : toes_l
+  torso : thorax # This changes
+  humerus_r : humerus_r
+  ulna_r : ulna_r
+  radius_r : radius_r
+  hand_r : hand_r
+  humerus_l : humerus_l
+  ulna_l : ulna_l
+  radius_l : radius_l
+  hand_l : hand_l
+
+markers_mapping: #Given specific markers of the OpenSim model, specify the corresponding vertex index on SMPL/SKEL skin
+  RFAradius : 5573 #RWRA
+  RFAulna : 5568 #RWRB
+  LFAradius : 2112 #LWRA
+  LFAulna : 2108 #LWRB