Correctly add UniRig source files

UniRig/.gitattributes

@@ -0,0 +1,5 @@
+examples/*.glb filter=lfs diff=lfs merge=lfs -text
+examples/*.fbx filter=lfs diff=lfs merge=lfs -text
+examples/*.vrm filter=lfs diff=lfs merge=lfs -text
+examples/*.FBX filter=lfs diff=lfs merge=lfs -text
+examples/*.obj filter=lfs diff=lfs merge=lfs -text

UniRig/.gitignore

@@ -0,0 +1,59 @@
+# igonore all pychace
+**/__pycache__/
+*.py[cod]
+*$py.class
+
+# ignore tmp & output files
+_data/
+tmp/
+*.npz
+*.blend
+*.blend1
+*.blend2
+
+# ignore logs
+wandb/
+lightning_logs/
+*.log
+
+# ignore experiments
+experiments/
+results/
+dataset_clean/
+logs/
+datalist/
+dataset_inference/
+dataset_inference_clean/
+feature_viz/
+
+# Distribution / packaging
+dist/
+build/
+*.egg-info/
+*.egg
+*.whl
+
+# Virtual environments
+venv/
+env/
+.env/
+.venv/
+
+# IDE specific files
+.idea/
+.vscode/
+*.swp
+*.swo
+.DS_Store
+
+# Jupyter Notebook
+.ipynb_checkpoints
+*.ipynb
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.coverage
+.coverage.*
+coverage.xml
+*.cover

UniRig/LICENSE

@@ -0,0 +1,21 @@
+# MIT License
+
+# Copyright (c) 2025 VAST-AI-Research and contributors.
+
+# 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

UniRig/README.md

@@ -0,0 +1,164 @@
+# UniRig: One Model to Rig Them All
+
+<div align="center">
+
+[![Project Page](https://img.shields.io/badge/🏠-Project%20Page-blue.svg)](https://zjp-shadow.github.io/works/UniRig/)
+[![Paper](https://img.shields.io/badge/📑-Paper-green.svg)](https://arxiv.org/abs/2504.12451)
+[![Model](https://img.shields.io/badge/🤗-Model-yellow.svg)](https://huggingface.co/VAST-AI/UniRig)
+
+</div>
+
+![teaser](assets/doc/unirig_teaser.png)
+
+This repository contains the official implementation for the **SIGGRAPH'25 (TOG) UniRig** framework, a unified solution for automatic 3D model rigging, developed by Tsinghua University and [Tripo](https://www.tripo3d.ai).
+
+**Paper:** [One Model to Rig Them All: Diverse Skeleton Rigging with UniRig](https://arxiv.org/abs/2504.12451)
+
+## Overview
+
+Rigging 3D models – creating a skeleton and assigning skinning weights – is a crucial but often complex and time-consuming step in 3D animation. UniRig tackles this challenge by introducing a novel, unified framework leveraging large autoregressive models to automate the process for a diverse range of 3D assets.
+
+Combining UniRig with keyframe animation produces these following results:
+
+| ![devil](assets/doc/devil.gif) | ![dragon](assets/doc/dragon.gif) | ![rabbit](assets/doc/rabbit.gif) |
+|:-----------------------------:|:-------------------------------:|:-------------------------------:|
+
+The full UniRig system consists of two main stages:
+1.  **Skeleton Prediction:** An GPT-like transformer autoregressively predicts a topologically valid skeleton hierarchy using a novel **Skeleton Tree Tokenization** scheme.
+2.  **Skinning Weight & Attribute Prediction:** A **Bone-Point Cross Attention** mechanism predicts per-vertex skinning weights and relevant bone attributes (e.g., for physics simulation) based on the predicted skeleton and input mesh geometry.
+
+This repository provides the code implementation for the entire framework vision, with components being released progressively.
+
+## Key Features (Full UniRig Framework)
+
+*   **Unified Model:** Aims to handle diverse model categories (humans, animals, objects) with a single framework.
+*   **Automated Skeleton Generation:** Predicts topologically valid skeleton structures. **(✅ Available in current release)**
+*   **Automated Skinning Prediction:** Predicts per-vertex skinning weights. **(✅ Available in current release)**
+*   **Bone Attribute Prediction:** Predicts attributes like stiffness for physics-based secondary motion. **(⏳ Coming Soon)**
+*   **High Accuracy & Robustness:** Achieves state-of-the-art results on challenging datasets (as shown in the paper with Rig-XL/VRoid training).
+*   **Efficient Tokenization:** Uses Skeleton Tree Tokenization for compact representation and efficient processing.
+*   **Human-in-the-Loop Ready:** Designed to potentially support iterative refinement workflows.
+
+## 🚨 Current Release Status & Roadmap 🚨
+
+We are open-sourcing UniRig progressively. Please note the current status:
+
+**Available Now (Initial Release):**
+*   ✅ **Code:** Implementation for skeleton and skinning prediction.
+*   ✅ **Model:** Skeleton & Skinning Prediction checkpoint trained on [**Articulation-XL2.0**](https://huggingface.co/datasets/Seed3D/Articulation-XL2.0). Available on [Hugging Face](https://huggingface.co/VAST-AI/UniRig).
+
+**Planned Future Releases:**
+*   ⏳ Release of the **Rig-XL** and **VRoid** datasets used in the paper.
+*   ⏳ Full UniRig model checkpoints (Skeleton + Skinning) trained on Rig-XL/VRoid, replicating the paper's main results.
+
+We appreciate your patience as we prepare these components for release. Follow [VAST-AI-Research](https://github.com/orgs/VAST-AI-Research) announcements for updates!
+
+## Installation
+
+1.  **Prerequisites:**
+    *   Python 3.11
+    *   PyTorch (tested with version >=2.3.1)
+
+2.  **Clone the repository:**
+    ```bash
+    git clone https://github.com/VAST-AI-Research/UniRig
+    cd UniRig
+    ```
+
+3.  **Set up a virtual environment (recommended):**
+    ```bash
+    conda create -n UniRig python=3.11
+    conda activate UniRig
+    ```
+
+4.  **Install dependencies:**
+    ```bash
+    python -m pip install torch torchvision
+    python -m pip install -r requirements.txt
+    python -m pip install spconv-{you-cuda-version}
+    python -m pip install torch_scatter torch_cluster -f https://data.pyg.org/whl/torch-{your-torch-version}+{your-cuda-version}.html --no-cache-dir
+    python -m pip install numpy==1.26.4
+    ```
+
+5.  **Download Model Checkpoint:**
+    The currently available skeleton prediction model checkpoint is hosted on Hugging Face and will typically be downloaded automatically by the provided scripts/functions.
+
+6.  **(Optional, for importing/exporting .vrm) Install the blender addon:**
+    The blender addon is modifed from [VRM-Addon-for-Blender](https://github.com/saturday06/VRM-Addon-for-Blender).
+
+    Make sure you are in the root directory of the project, then:
+    ```bash
+    python -c "import bpy, os; bpy.ops.preferences.addon_install(filepath=os.path.abspath('blender/add-on-vrm-v2.20.77_modified.zip'))"
+    ```
+
+## Usage
+
+### Skeleton Prediction (Available Now)
+
+Generate a skeleton for your 3D model using our pre-trained model. The process automatically analyzes the geometry and predicts an appropriate skeletal structure.
+
+```bash
+# Process a single file
+bash launch/inference/generate_skeleton.sh --input examples/giraffe.glb --output results/giraffe_skeleton.fbx
+
+# Process multiple files in a directory
+bash launch/inference/generate_skeleton.sh --input_dir <your_input_directory> --output_dir <your_output_directory>
+
+# Try different skeleton variations by changing the random seed
+bash launch/inference/generate_skeleton.sh --input examples/giraffe.glb --output results/giraffe_skeleton.fbx --seed 42
+```
+
+Supported input formats: `.obj`, `.fbx`, `.glb`, and `.vrm`
+
+### Skinning Weight Prediction (Available Now)
+```bash
+# Skin a single file
+bash launch/inference/generate_skin.sh --input examples/skeleton/giraffe.fbx --output results/giraffe_skin.fbx
+
+# Process multiple files in a directory
+bash launch/inference/generate_skin.sh --input_dir <your_input_directory> --output_dir <your_output_directory>
+```
+
+Note that the command above uses an **edited-version** from skeleton phase. The results may degrade significantly if the skeleton is inaccurate — for example, if tail bones or wing bones are missing. Therefore, it is recommended to refine the skeleton before performing skinning in order to achieve better results.
+
+### Merge the Predicted Results
+
+Combine the predicted skeleton with your original 3D model to create a fully rigged asset:
+
+```bash
+# Merge skeleton from skeleton prediction
+bash launch/inference/merge.sh --source results/giraffe_skeleton.fbx --target examples/giraffe.glb --output results/giraffe_rigged.glb
+
+# Or merge skin from skin prediction
+bash launch/inference/merge.sh --source results/giraffe_skin.fbx --target examples/giraffe.glb --output results/giraffe_rigged.glb
+```
+
+## Models
+
+Available models are hosted on the: https://huggingface.co/VAST-AI/UniRig
+
+## System Requirements
+
+- CUDA-enabled GPU with at least 8GB VRAM
+
+## Citation
+
+```
+@article{zhang2025unirig,
+  title={One Model to Rig Them All: Diverse Skeleton Rigging with UniRig},
+  author={Zhang, Jia-Peng and Pu, Cheng-Feng and Guo, Meng-Hao and Cao, Yan-Pei and Hu, Shi-Min},
+  journal={arXiv preprint arXiv:2504.12451},
+  year={2025}
+}
+```
+
+## Acknowledgements
+
+We would like to thank the following open-source projects and research works:
+
+- [OPT](https://huggingface.co/facebook/opt-350m) for model architecture
+- [3DShape2VecSet](https://github.com/1zb/3DShape2VecSet) for 3D shape representation
+- [SAMPart3D](https://github.com/Pointcept/SAMPart3D) and [Michelangelo](https://github.com/NeuralCarver/Michelangelo/) for shape encoder implementation
+- [Articulation-XL2.0](https://huggingface.co/datasets/Seed3D/Articulation-XL2.0) for a curated dataset
+
+We are grateful to the broader research community for their open exploration and contributions to the field of 3D generation.

UniRig/blender/add-on-vrm-v2.20.77_modified.zip

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8fe1e8b7e31ec602d9b32db33c533b03d68df747837bfad3479e1057bc9937c5
+size 1331571

UniRig/configs/data/quick_inference.yaml

@@ -0,0 +1,16 @@
+input_dataset_dir: &input_dataset_dir ./dataset_inference
+output_dataset_dir: &output_dataset_dir ./dataset_inference_clean
+
+predict_dataset_config:
+  shuffle: False
+  batch_size: 1
+  num_workers: 1
+  pin_memory: False
+  persistent_workers: False
+  datapath_config:
+    input_dataset_dir: *output_dataset_dir
+    use_prob: False
+    data_path:
+      inference: [
+        [./dataset_inference_clean/inference_datalist.txt, 1.0],
+      ]

UniRig/configs/model/unirig_ar_350m_1024_81920_float32.yaml

@@ -0,0 +1,32 @@
+__target__: unirig_ar
+llm:
+  pretrained_model_name_or_path: facebook/opt-350m
+  n_positions: 3076
+  max_position_embeddings: 3076
+  hidden_size: 1024
+  word_embed_proj_dim: 1024
+  do_layer_norm_before: True
+  _attn_implementation: flash_attention_2
+
+mesh_encoder:
+  __target__: michelangelo_encoder
+  pretrained_path: ~
+  freeze_encoder: False
+  device: cpu
+  dtype: float32
+  num_latents: 512
+  embed_dim: 64
+  point_feats: 3
+  num_freqs: 8
+  include_pi: False
+  heads: 8
+  width: 512
+  num_encoder_layers: 16
+  use_ln_post: True
+  init_scale: 0.25
+  qkv_bias: False
+  use_checkpoint: False
+  flash: True
+  supervision_type: sdf
+  query_method: False
+  token_num: 1024

UniRig/configs/model/unirig_skin.yaml

@@ -0,0 +1,52 @@
+__target__: unirig_skin
+
+num_train_vertex: 512 # increase this for faster speed at the cost of memory
+num_heads: 16
+feat_dim: 768
+grid_size: 0.005
+mlp_dim: 512
+num_bone_attn: 8
+num_mesh_bone_attn: 16
+bone_embed_dim: 1024
+voxel_mask: 3.0
+
+mesh_encoder:
+  # vertex groups are handled in model
+  __target__: ptv3obj
+  pretrained_path: ~
+  freeze_encoder: False
+  in_channels: 9
+  cls_mode: False
+  shuffle_orders: True
+  drop_path: 0.0
+  upcast_attention: False
+  upcast_softmax: False
+  enc_depths: [3, 3, 3, 6, 16]
+  enc_channels: [32, 64, 128, 256, 384]
+  enc_num_head: [2, 4, 8, 16, 24]
+  enable_qknorm: True
+  layer_norm: False
+  res_linear: True
+
+global_encoder:
+  __target__: michelangelo_encoder
+  pretrained_path: ~
+  freeze_encoder: False
+  device: cpu
+  dtype: float32
+  num_latents: 512
+  embed_dim: 64
+  point_feats: 3
+  num_freqs: 8
+  include_pi: False
+  heads: 8
+  width: 512
+  num_encoder_layers: 16
+  use_ln_post: True
+  init_scale: 0.25
+  qkv_bias: False
+  use_checkpoint: False
+  flash: True
+  supervision_type: sdf
+  query_method: False
+  token_num: 1024

UniRig/configs/skeleton/mixamo.yaml

@@ -0,0 +1,59 @@
+parts_order: [body, hand]
+
+parts:
+  body: [
+    mixamorig:Hips,
+    mixamorig:Spine,
+    mixamorig:Spine1,
+    mixamorig:Spine2,
+    mixamorig:Neck,
+    mixamorig:Head,
+    mixamorig:LeftShoulder,
+    mixamorig:LeftArm,
+    mixamorig:LeftForeArm,
+    mixamorig:LeftHand,
+    mixamorig:RightShoulder,
+    mixamorig:RightArm,
+    mixamorig:RightForeArm,
+    mixamorig:RightHand,
+    mixamorig:LeftUpLeg,
+    mixamorig:LeftLeg,
+    mixamorig:LeftFoot,
+    mixamorig:LeftToeBase,
+    mixamorig:RightUpLeg,
+    mixamorig:RightLeg,
+    mixamorig:RightFoot,
+    mixamorig:RightToeBase,
+  ]
+  hand: [
+    mixamorig:LeftHandThumb1,
+    mixamorig:LeftHandThumb2,
+    mixamorig:LeftHandThumb3,
+    mixamorig:LeftHandIndex1,
+    mixamorig:LeftHandIndex2,
+    mixamorig:LeftHandIndex3,
+    mixamorig:LeftHandMiddle1,
+    mixamorig:LeftHandMiddle2,
+    mixamorig:LeftHandMiddle3,
+    mixamorig:LeftHandRing1,
+    mixamorig:LeftHandRing2,
+    mixamorig:LeftHandRing3,
+    mixamorig:LeftHandPinky1,
+    mixamorig:LeftHandPinky2,
+    mixamorig:LeftHandPinky3,
+    mixamorig:RightHandIndex1,
+    mixamorig:RightHandIndex2,
+    mixamorig:RightHandIndex3,
+    mixamorig:RightHandThumb1,
+    mixamorig:RightHandThumb2,
+    mixamorig:RightHandThumb3,
+    mixamorig:RightHandMiddle1,
+    mixamorig:RightHandMiddle2,
+    mixamorig:RightHandMiddle3,
+    mixamorig:RightHandRing1,
+    mixamorig:RightHandRing2,
+    mixamorig:RightHandRing3,
+    mixamorig:RightHandPinky1,
+    mixamorig:RightHandPinky2,
+    mixamorig:RightHandPinky3,
+  ]

UniRig/configs/skeleton/vroid.yaml

@@ -0,0 +1,59 @@
+parts_order: [body, hand]
+
+parts:
+  body: [
+    J_Bip_C_Hips,
+    J_Bip_C_Spine,
+    J_Bip_C_Chest,
+    J_Bip_C_UpperChest,
+    J_Bip_C_Neck,
+    J_Bip_C_Head,
+    J_Bip_L_Shoulder,
+    J_Bip_L_UpperArm,
+    J_Bip_L_LowerArm,
+    J_Bip_L_Hand,
+    J_Bip_R_Shoulder,
+    J_Bip_R_UpperArm,
+    J_Bip_R_LowerArm,
+    J_Bip_R_Hand,
+    J_Bip_L_UpperLeg,
+    J_Bip_L_LowerLeg,
+    J_Bip_L_Foot,
+    J_Bip_L_ToeBase,
+    J_Bip_R_UpperLeg,
+    J_Bip_R_LowerLeg,
+    J_Bip_R_Foot,
+    J_Bip_R_ToeBase,
+  ]
+  hand: [
+    J_Bip_L_Thumb1,
+    J_Bip_L_Thumb2,
+    J_Bip_L_Thumb3,
+    J_Bip_L_Index1,
+    J_Bip_L_Index2,
+    J_Bip_L_Index3,
+    J_Bip_L_Middle1,
+    J_Bip_L_Middle2,
+    J_Bip_L_Middle3,
+    J_Bip_L_Ring1,
+    J_Bip_L_Ring2,
+    J_Bip_L_Ring3,
+    J_Bip_L_Little1,
+    J_Bip_L_Little2,
+    J_Bip_L_Little3,
+    J_Bip_R_Index1,
+    J_Bip_R_Index2,
+    J_Bip_R_Index3,
+    J_Bip_R_Thumb1,
+    J_Bip_R_Thumb2,
+    J_Bip_R_Thumb3,
+    J_Bip_R_Middle1,
+    J_Bip_R_Middle2,
+    J_Bip_R_Middle3,
+    J_Bip_R_Ring1,
+    J_Bip_R_Ring2,
+    J_Bip_R_Ring3,
+    J_Bip_R_Little1,
+    J_Bip_R_Little2,
+    J_Bip_R_Little3,
+  ]

UniRig/configs/system/ar_inference_articulationxl.yaml

@@ -0,0 +1,14 @@
+__target__: ar
+val_interval: 1
+generate_kwargs:
+  max_new_tokens: 2048
+  num_return_sequences: 1
+  num_beams: 15
+  do_sample: True
+  top_k: 5
+  top_p: 0.95
+  repetition_penalty: 3.0
+  temperature: 1.5 # must be a float
+  no_cls: False
+  assign_cls: articulationxl
+  use_dir_cls: False

UniRig/configs/system/skin.yaml

@@ -0,0 +1,5 @@
+__target__: skin
+val_interval: 1
+val_start_from: 1
+output_path: tmp_skin
+record_res: True

UniRig/configs/task/quick_inference_skeleton_articulationxl_ar_256.yaml

@@ -0,0 +1,30 @@
+mode: predict
+debug: False
+experiment_name: quick_inference_skeleton_articulationxl_ar_256
+resume_from_checkpoint: experiments/skeleton/articulation-xl_quantization_256/model.ckpt
+
+components:
+  data: quick_inference
+  tokenizer: tokenizer_parts_articulationxl_256
+  transform: inference_ar_transform
+  model: unirig_ar_350m_1024_81920_float32
+  system: ar_inference_articulationxl
+  data_name: raw_data.npz
+
+writer:
+  __target__: ar
+  output_dir: ~ # export results into the same input folder
+  add_num: False
+  repeat: 1
+  export_npz: predict_skeleton
+  export_obj: skeleton
+  export_fbx: skeleton
+  # export_pc: pc
+
+trainer:
+  max_epochs: 1
+  num_nodes: 1
+  devices: 1
+  precision: bf16-mixed
+  accelerator: gpu
+  strategy: auto

UniRig/configs/task/quick_inference_unirig_skin.yaml

@@ -0,0 +1,28 @@
+mode: predict
+debug: False
+experiment_name: quick_inference_skin
+resume_from_checkpoint: experiments/skin/articulation-xl/model.ckpt
+
+components:
+  data: quick_inference
+  transform: inference_skin_transform
+  model: unirig_skin
+  system: skin
+  data_name: predict_skeleton.npz # capture data from ar phase
+
+writer:
+  __target__: skin
+  output_dir: results
+  add_num: False
+  repeat: 1
+  save_name: predict
+  export_npz: predict_skin # this must be specified if textured results are required
+  export_fbx: result_fbx
+
+trainer:
+  num_nodes: 1
+  devices: 1
+  precision: bf16-mixed
+  accelerator: gpu
+  strategy: auto
+  inference_mode: True

UniRig/configs/tokenizer/tokenizer_parts_articulationxl_256.yaml

@@ -0,0 +1,14 @@
+method: tokenizer_part
+num_discrete: 256
+continuous_range: [-1, 1]
+cls_token_id:
+  vroid: 0
+  mixamo: 1 # this is currently untrained, do not use it
+  articulationxl: 2
+parts_token_id:
+  body: 0
+  hand: 1
+order_config:
+  skeleton_path:
+    vroid: ./configs/skeleton/vroid.yaml
+    mixamo: ./configs/skeleton/mixamo.yaml

UniRig/configs/transform/inference_ar_transform.yaml

@@ -0,0 +1,30 @@
+sampler_config: &sampler_config
+  method: mix
+  num_samples: 65536
+  vertex_samples: 8192
+
+tail_config: &tail_config
+  copy_joint_to_tail: False # Be careful ! If tail is important, keep it False !!!
+  connect_tail_to_unique_son: True
+
+order_config: &order_config
+  skeleton_path:
+    vroid: ./configs/skeleton/vroid.yaml
+    mixamo: ./configs/skeleton/mixamo.yaml
+
+vertex_group_config: &vertex_group_config
+
+validate_transform_config: &validate_transform_config
+  augment_config:
+    augment_affine_config:
+      normalize_into: [-1.0, 1.0]
+      random_scale_p: 0.0
+      random_scale: [1.0, 1.0]
+      random_shift_p: 0.0
+      random_shift: [0.0, 0.0]
+  tail_config: *tail_config
+  order_config: *order_config
+  vertex_group_config: *vertex_group_config
+  sampler_config: *sampler_config
+
+predict_transform_config: *validate_transform_config

UniRig/configs/transform/inference_skin_transform.yaml

@@ -0,0 +1,32 @@
+sampler_config: &sampler_config
+  method: mix
+  num_samples: 32768
+  vertex_samples: 8192
+
+tail_config: &tail_config
+  copy_joint_to_tail: False # Be careful ! If tail is important, keep it False !!!
+  connect_tail_to_unique_son: True
+
+order_config: &order_config
+  skeleton_path:
+    vroid: ./configs/skeleton/vroid.yaml
+    mixamo: ./configs/skeleton/mixamo.yaml
+
+predict_transform_config: 
+  augment_config:
+    augment_affine_config:
+      normalize_into: [-1.0, 1.0]
+  tail_config: *tail_config
+  order_config: *order_config
+  vertex_group_config:
+    names: ['voxel_skin']
+    kwargs:
+      voxel_skin:
+        grid: 196 # increase this for better results
+        alpha: 0.5
+        link_dis: 0.00001
+        grid_query: 7
+        vertex_query: 1
+        grid_weight: 3.0
+        # mode: exp
+  sampler_config: *sampler_config

UniRig/examples/bird.glb

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bb59726f598ab4a4e4c431b9789317f5b2f4252a6fb57364d929f5e1ddd7b5bb
+size 8032388

UniRig/examples/giraffe.glb

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9a947cae00b169345802c08885c1e54313c6ce93c885bacf8e37e8a1f18f9e3b
+size 6310044

UniRig/examples/skeleton/bird.fbx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:885d432850506ab673d509ae2481067cc623452d5910090e1e15f323b9f83fa2
+size 401084

UniRig/examples/skeleton/giraffe.fbx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:73673a15c8103fbf9a6b39762768ae48e7c9404eab4850a60e4863ee400336fd
+size 759180

UniRig/examples/skeleton/tira.fbx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:57ed6969266d642da2d8a579059462c7cd4a36c9bd7a8415236dcbea36607fee
+size 1694668

UniRig/examples/skeleton/tripo_carrot.fbx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:50e08d70c7bdfa96eb842aed18564d8486a4622a474dbc0e0ef9304af1d4c6d3
+size 1879420

UniRig/examples/tira.glb

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8e5a282d0a99c61a8d2439496057b15b0c6ea02643c16da2dcbec85571157799
+size 32346060

UniRig/examples/tripo_carrot.glb

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3c00b7e6ff1e71a019a02128a5b37e8d2db259a402313ad9b4eaab4f183ae40b
+size 9511824

UniRig/launch/inference/extract.sh

@@ -0,0 +1,60 @@
+# extract mesh
+config="configs/data/quick_inference.yaml"
+require_suffix="obj,fbx,FBX,dae,glb,gltf,vrm"
+num_runs=1
+force_override="false"
+faces_target_count=50000
+
+while [[ "$#" -gt 0 ]]; do
+    case $1 in
+        --config) config="$2"; shift ;;
+        --require_suffix) require_suffix="$2"; shift ;;
+        --num_runs) num_runs="$2"; shift ;;
+        --force_override) force_override="$2"; shift ;;
+        --faces_target_count) faces_target_count="$2"; shift ;;
+        --time) time="$2"; shift ;;
+        --input) input="$2"; shift ;;
+        --input_dir) input_dir="$2"; shift ;;
+        --output_dir) output_dir="$2"; shift ;;
+        *) echo "Unknown parameter: $1"; exit 1 ;;
+    esac
+    shift
+done
+
+# ensure psutil is installed for memory management
+pip install psutil --quiet
+if [ $? -ne 0 ]; then
+    echo "Warning: Failed to install psutil. Memory management may not work properly."
+fi
+
+# set the time for all processes to use
+time=$(date "+%Y_%m_%d_%H_%M_%S")
+
+for (( i=0; i<num_runs; i++ ))
+do
+    cmd=" \
+    python -m src.data.extract \
+    --config=$config \
+    --require_suffix=$require_suffix \
+    --force_override=$force_override \
+    --num_runs=$num_runs \
+    --id=$i \
+    --time=$time \
+    --faces_target_count=$faces_target_count \
+    "
+    if [ -n "$input" ]; then
+        cmd="$cmd --input=$input"
+    fi
+    if [ -n "$input_dir" ]; then
+        cmd="$cmd --input_dir=$input_dir"
+    fi
+    if [ -n "$output_dir" ]; then
+        cmd="$cmd --output_dir=$output_dir"
+    fi
+    cmd="$cmd &"
+    eval $cmd
+done
+
+wait
+
+echo "done"

UniRig/launch/inference/generate_skeleton.sh

@@ -0,0 +1,81 @@
+# generate skeleton
+config="configs/data/quick_inference.yaml"
+require_suffix="obj,fbx,FBX,dae,glb,gltf,vrm"
+num_runs=1
+force_override="false"
+faces_target_count=50000
+skeleton_task="configs/task/quick_inference_skeleton_articulationxl_ar_256.yaml"
+add_root="false"
+seed=12345
+npz_dir="tmp"
+
+while [[ "$#" -gt 0 ]]; do
+    case $1 in
+        --config) config="$2"; shift ;;
+        --require_suffix) require_suffix="$2"; shift ;;
+        --num_runs) num_runs="$2"; shift ;;
+        --force_override) force_override="$2"; shift ;;
+        --faces_target_count) faces_target_count="$2"; shift ;;
+        --skeleton_task) skeleton_task="$2"; shift ;;
+        --add_root) add_root="$2"; shift ;;
+        --seed) seed="$2"; shift ;;
+        --input) input="$2"; shift ;;
+        --input_dir) input_dir="$2"; shift ;;
+        --output_dir) output_dir="$2"; shift ;;
+        --output) output="$2"; shift ;;
+        *) echo "Unknown parameter: $1"; exit 1 ;;
+    esac
+    shift
+done
+
+# 1. extract mesh
+cmd=" \
+    bash ./launch/inference/extract.sh \
+    --config $config \
+    --require_suffix $require_suffix \
+    --force_override $force_override \
+    --num_runs $num_runs \
+    --faces_target_count $faces_target_count \
+"
+if [ -n "$input" ]; then
+    cmd="$cmd --input $input"
+fi
+if [ -n "$input_dir" ]; then
+    cmd="$cmd --input_dir $input_dir"
+fi
+if [ -n "$npz_dir" ]; then
+    cmd="$cmd --output_dir $npz_dir"
+fi
+
+cmd="$cmd &"
+eval $cmd
+
+wait
+
+# 2. inference skeleton
+cmd="\
+    python run.py \
+    --task=$skeleton_task \
+    --seed=$seed \
+"
+if [ -n "$input" ]; then
+    cmd="$cmd --input=$input"
+fi
+if [ -n "$input_dir" ]; then
+    cmd="$cmd --input_dir=$input_dir"
+fi
+if [ -n "$output" ]; then
+    cmd="$cmd --output=$output"
+fi
+if [ -n "$output_dir" ]; then
+    cmd="$cmd --output_dir=$output_dir"
+fi
+if [ -n "$npz_dir" ]; then
+    cmd="$cmd --npz_dir=$npz_dir"
+fi
+
+eval $cmd
+
+wait
+
+echo "done"

UniRig/launch/inference/generate_skin.sh

@@ -0,0 +1,84 @@
+# generate skin
+config="configs/data/quick_inference.yaml"
+require_suffix="obj,fbx,FBX,dae,glb,gltf,vrm"
+num_runs=1
+force_override="true"
+faces_target_count=50000
+skin_task="configs/task/quick_inference_unirig_skin.yaml"
+seed=12345
+npz_dir="tmp"
+data_name="raw_data.npz"
+
+while [[ "$#" -gt 0 ]]; do
+    case $1 in
+        --config) config="$2"; shift ;;
+        --require_suffix) require_suffix="$2"; shift ;;
+        --num_runs) num_runs="$2"; shift ;;
+        --force_override) force_override="$2"; shift ;;
+        --faces_target_count) faces_target_count="$2"; shift ;;
+        --skin_task) skin_task="$2"; shift ;;
+        --seed) seed="$2"; shift ;;
+        --input) input="$2"; shift ;;
+        --input_dir) input_dir="$2"; shift ;;
+        --output_dir) output_dir="$2"; shift ;;
+        --output) output="$2"; shift ;;
+        --data_name) data_name="$2"; shift ;;
+        *) echo "Unknown parameter: $1"; exit 1 ;;
+    esac
+    shift
+done
+
+# 1. extract mesh
+cmd=" \
+    bash ./launch/inference/extract.sh \
+    --config $config \
+    --require_suffix $require_suffix \
+    --force_override $force_override \
+    --num_runs $num_runs \
+    --faces_target_count $faces_target_count \
+"
+if [ -n "$input" ]; then
+    cmd="$cmd --input $input"
+fi
+if [ -n "$input_dir" ]; then
+    cmd="$cmd --input_dir $input_dir"
+fi
+if [ -n "$npz_dir" ]; then
+    cmd="$cmd --output_dir $npz_dir"
+fi
+
+cmd="$cmd &"
+eval $cmd
+
+wait
+
+# 2. inference skin
+cmd="\
+    python run.py \
+    --task=$skin_task \
+    --seed=$seed \
+"
+if [ -n "$input" ]; then
+    cmd="$cmd --input=$input"
+fi
+if [ -n "$input_dir" ]; then
+    cmd="$cmd --input_dir=$input_dir"
+fi
+if [ -n "$output" ]; then
+    cmd="$cmd --output=$output"
+fi
+if [ -n "$output_dir" ]; then
+    cmd="$cmd --output_dir=$output_dir"
+fi
+if [ -n "$npz_dir" ]; then
+    cmd="$cmd --npz_dir=$npz_dir"
+fi
+if [ -n "$data_name" ]; then
+    cmd="$cmd --data_name=$data_name"
+fi
+
+eval $cmd
+
+wait
+
+echo "done"

UniRig/launch/inference/merge.sh

@@ -0,0 +1,33 @@
+# merge texture
+require_suffix="obj,fbx,FBX,dae,glb,gltf,vrm"
+source=""
+target=""
+output=""
+
+while [[ "$#" -gt 0 ]]; do
+    case $1 in
+        --require_suffix) require_suffix="$2"; shift ;;
+        --source) source="$2"; shift ;;
+        --target) target="$2"; shift ;;
+        --output) output="$2"; shift ;;
+        *) echo "Unknown parameter: $1"; exit 1 ;;
+    esac
+    shift
+done
+
+cmd=" \
+    python -m src.inference.merge \
+    --require_suffix=$require_suffix \
+    --num_runs=1 \
+    --id=0 \
+    --source=$source \
+    --target=$target \
+    --output=$output \
+"
+
+cmd="$cmd &"
+eval $cmd
+
+wait
+
+echo "done"

UniRig/requirements.txt

@@ -0,0 +1,15 @@
+transformers
+python-box
+einops
+omegaconf
+pytorch_lightning
+lightning
+addict
+timm
+fast-simplification
+bpy==4.2
+flash_attn
+trimesh
+open3d
+pyrender
+huggingface_hub

UniRig/run.py

@@ -0,0 +1,186 @@
+import argparse
+import yaml
+from box import Box
+import os
+import torch
+import lightning as L
+from lightning.pytorch.callbacks import ModelCheckpoint, Callback
+from typing import List
+from math import ceil
+import numpy as np
+from lightning.pytorch.strategies import FSDPStrategy, DDPStrategy
+from src.inference.download import download
+
+from src.data.asset import Asset
+from src.data.extract import get_files
+from src.data.dataset import UniRigDatasetModule, DatasetConfig, ModelInput
+from src.data.datapath import Datapath
+from src.data.transform import TransformConfig
+from src.tokenizer.spec import TokenizerConfig
+from src.tokenizer.parse import get_tokenizer
+from src.model.parse import get_model
+from src.system.parse import get_system, get_writer
+
+from tqdm import tqdm
+import time
+
+def load(task: str, path: str) -> Box:
+    if path.endswith('.yaml'):
+        path = path.removesuffix('.yaml')
+    path += '.yaml'
+    print(f"\033[92mload {task} config: {path}\033[0m")
+    return Box(yaml.safe_load(open(path, 'r')))
+
+def nullable_string(val):
+    if not val:
+        return None
+    return val
+
+if __name__ == "__main__":
+    torch.set_float32_matmul_precision('high')
+    
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--task", type=str, required=True)
+    parser.add_argument("--seed", type=int, required=False, default=123,
+                        help="random seed")
+    parser.add_argument("--input", type=nullable_string, required=False, default=None,
+                        help="a single input file or files splited by comma")
+    parser.add_argument("--input_dir", type=nullable_string, required=False, default=None,
+                        help="input directory")
+    parser.add_argument("--output", type=nullable_string, required=False, default=None,
+                        help="filename for a single output")
+    parser.add_argument("--output_dir", type=nullable_string, required=False, default=None,
+                        help="output directory")
+    parser.add_argument("--npz_dir", type=nullable_string, required=False, default='tmp',
+                        help="intermediate npz directory")
+    parser.add_argument("--cls", type=nullable_string, required=False, default=None,
+                        help="class name")
+    parser.add_argument("--data_name", type=nullable_string, required=False, default=None,
+                        help="npz filename from skeleton phase")
+    args = parser.parse_args()
+    
+    L.seed_everything(args.seed, workers=True)
+    
+    task = load('task', args.task)
+    mode = task.mode
+    assert mode in ['predict']
+    
+    if args.input is not None or args.input_dir is not None:
+        assert args.output_dir is not None or args.output is not None, 'output or output_dir must be specified'
+        assert args.npz_dir is not None, 'npz_dir must be specified'
+        files = get_files(
+            data_name=task.components.data_name,
+            inputs=args.input,
+            input_dataset_dir=args.input_dir,
+            output_dataset_dir=args.npz_dir,
+            force_override=True,
+            warning=False,
+        )
+        files = [f[1] for f in files]
+        if len(files) > 1 and args.output is not None:
+            print("\033[92mwarning: output is specified, but multiple files are detected. Output will be written.\033[0m")
+        datapath = Datapath(files=files, cls=args.cls)
+    else:
+        datapath = None
+    
+    data_config = load('data', os.path.join('configs/data', task.components.data))
+    transform_config = load('transform', os.path.join('configs/transform', task.components.transform))
+    
+    # get tokenizer
+    tokenizer_config = task.components.get('tokenizer', None)
+    if tokenizer_config is not None:
+        tokenizer_config = load('tokenizer', os.path.join('configs/tokenizer', task.components.tokenizer))
+        tokenizer_config = TokenizerConfig.parse(config=tokenizer_config)
+    
+    # get data name
+    data_name = task.components.get('data_name', 'raw_data.npz')
+    if args.data_name is not None:
+        data_name = args.data_name
+        
+    # get predict dataset
+    predict_dataset_config = data_config.get('predict_dataset_config', None)
+    if predict_dataset_config is not None:
+        predict_dataset_config = DatasetConfig.parse(config=predict_dataset_config).split_by_cls()
+    
+    # get predict transform
+    predict_transform_config = transform_config.get('predict_transform_config', None)
+    if predict_transform_config is not None:
+        predict_transform_config = TransformConfig.parse(config=predict_transform_config)
+    
+    # get model
+    model_config = task.components.get('model', None)
+    if model_config is not None:
+        model_config = load('model', os.path.join('configs/model', model_config))
+        if tokenizer_config is not None:
+            tokenizer = get_tokenizer(config=tokenizer_config)
+        else:
+            tokenizer = None
+        model = get_model(tokenizer=tokenizer, **model_config)
+    else:
+        model = None
+    
+    # set data
+    data = UniRigDatasetModule(
+        process_fn=None if model is None else model._process_fn,
+        predict_dataset_config=predict_dataset_config,
+        predict_transform_config=predict_transform_config,
+        tokenizer_config=tokenizer_config,
+        debug=False,
+        data_name=data_name,
+        datapath=datapath,
+        cls=args.cls,
+    )
+    
+    # add call backs
+    callbacks = []
+
+    ## get checkpoint callback
+    checkpoint_config = task.get('checkpoint', None)
+    if checkpoint_config is not None:
+        checkpoint_config['dirpath'] = os.path.join('experiments', task.experiment_name)
+        callbacks.append(ModelCheckpoint(**checkpoint_config))
+    
+    ## get writer callback
+    writer_config = task.get('writer', None)
+    if writer_config is not None:
+        assert predict_transform_config is not None, 'missing predict_transform_config in transform'
+        if args.output_dir is not None or args.output is not None:
+            if args.output is not None:
+                assert args.output.endswith('.fbx'), 'output must be .fbx'
+            writer_config['npz_dir'] = args.npz_dir
+            writer_config['output_dir'] = args.output_dir
+            writer_config['output_name'] = args.output
+            writer_config['user_mode'] = True
+        callbacks.append(get_writer(**writer_config, order_config=predict_transform_config.order_config))
+    
+    # get trainer
+    trainer_config = task.get('trainer', {})
+    
+    # get system
+    system_config = task.components.get('system', None)
+    if system_config is not None:
+        system_config = load('system', os.path.join('configs/system', system_config))
+        system = get_system(
+            **system_config,
+            model=model,
+            steps_per_epoch=1,
+        )
+    else:
+        system = None
+    
+    logger = None
+
+    # set ckpt path
+    resume_from_checkpoint = task.get('resume_from_checkpoint', None)
+    resume_from_checkpoint = download(resume_from_checkpoint)
+    trainer = L.Trainer(
+        callbacks=callbacks,
+        logger=logger,
+        **trainer_config,
+    )
+    
+    if mode == 'predict':
+        assert resume_from_checkpoint is not None, 'expect resume_from_checkpoint in task'
+        trainer.predict(system, datamodule=data, ckpt_path=resume_from_checkpoint, return_predictions=False)
+    else:
+        assert 0

UniRig/src/data/asset.py

@@ -0,0 +1,433 @@
+from collections import defaultdict
+from dataclasses import dataclass
+import numpy as np
+from numpy import ndarray
+
+from typing import Dict, Union, List, Tuple
+
+from .order import Order
+from .raw_data import RawData
+from .exporter import Exporter
+
+from ..tokenizer.spec import TokenizeInput
+from .utils import linear_blend_skinning
+
+import trimesh
+
+
+@dataclass
+class Asset(Exporter):
+    '''
+    Dataclass to handle data parsed from raw data.
+    '''
+    
+    # data class
+    cls: str
+    
+    # where is this asset from
+    path: str
+    
+    # data file name
+    data_name: str
+    
+    # vertices of the mesh, shape (N, 3), float32
+    vertices: ndarray
+    
+    # normals of vertices, shape (N, 3), float32
+    vertex_normals: ndarray
+    
+    # faces of mesh, shape (F, 3), face id starts from 0 to F-1, int64
+    faces: ndarray
+    
+    # face normal of mesh, shape (F, 3), float32
+    face_normals: ndarray
+    
+    # joints of bones, shape (J, 3), float32
+    joints: Union[ndarray, None]=None
+    
+    # tails of joints, shape (J, 3), float32
+    tails: Union[ndarray, None]=None
+    
+    # skinning of joints, shape (N, J), float32
+    skin: Union[ndarray, None]=None
+    
+    # whether the joint has skin, bool
+    no_skin: Union[ndarray, None]=None
+    
+    # vertex groups
+    vertex_groups: Union[Dict[str, ndarray], None]=None
+    
+    # parents of joints, None represents no parent(a root joint)
+    # make sure parent[k] < k
+    parents: Union[List[Union[int, None]], None]=None
+    
+    # names of joints
+    names: Union[List[str], None]=None
+    
+    # sampled vertices, shape (N, 3)
+    sampled_vertices: Union[ndarray, None]=None
+    
+    # sampled normals, shape (N, 3)
+    sampled_normals: Union[ndarray, None]=None
+    
+    # sampled vertex groups, every vertex group should be (N, J)
+    sampled_vertex_groups: Union[Dict[str, ndarray], None]=None
+    
+    # {id: part}, part==None -> a spring token
+    parts_bias: Union[Dict[int, Union[str, None]], None]=None
+    
+    # local coordinate, shape (J, 4, 4)
+    matrix_local: Union[ndarray, None]=None
+    
+    # pose matrix for skinning loss calculation, shape (J, 4, 4)
+    pose_matrix: Union[ndarray, None]=None
+    
+    meta: Union[Dict[str, ...], None]=None
+    
+    @property
+    def N(self):
+        '''
+        number of vertices
+        '''
+        return self.vertices.shape[0]
+    
+    @property
+    def F(self):
+        '''
+        number of faces
+        '''
+        return self.faces.shape[0]
+    
+    @property
+    def J(self):
+        '''
+        number of joints
+        '''
+        return self.joints.shape[0]
+    
+    def get_matrix(self, matrix_basis: ndarray, matrix_local: Union[ndarray, None]=None):
+        '''
+        get matrix
+        
+        matrix_basis: (J, 4, 4)
+        '''
+        if matrix_local is None:
+            assert self.joints is not None
+            matrix_local = self.matrix_local
+            if matrix_local is None:
+                matrix_local = np.zeros((self.J, 4, 4))
+                matrix_local[:, 0, 0] = 1.
+                matrix_local[:, 1, 1] = 1.
+                matrix_local[:, 2, 2] = 1.
+                matrix_local[:, 3, 3] = 1.
+                for i in range(self.J):
+                    matrix_local[i, :3, 3] = self.joints[i]
+        
+        matrix = np.zeros((self.J, 4, 4))
+        for i in range(self.J):
+            if i==0:
+                matrix[i] = matrix_local[i] @ matrix_basis[i]
+            else:
+                pid = self.parents[i]
+                matrix_parent = matrix[pid]
+                matrix_local_parent = matrix_local[pid]
+                
+                matrix[i] = (
+                    matrix_parent @
+                    (np.linalg.inv(matrix_local_parent) @ matrix_local[i]) @
+                    matrix_basis[i]
+                )
+        return matrix
+    
+    def apply_matrix_basis(self, matrix_basis: ndarray):
+        '''
+        apply a pose to armature
+        
+        matrix_basis: (J, 4, 4)
+        '''
+        matrix_local = self.matrix_local
+        if matrix_local is None:
+            matrix_local = np.zeros((self.J, 4, 4))
+            matrix_local[:, 0, 0] = 1.
+            matrix_local[:, 1, 1] = 1.
+            matrix_local[:, 2, 2] = 1.
+            matrix_local[:, 3, 3] = 1.
+            for i in range(self.J):
+                matrix_local[i, :3, 3] = self.joints[i].copy()
+        
+        matrix = self.get_matrix(matrix_basis=matrix_basis, matrix_local=matrix_local)
+        self.joints = matrix[:, :3, 3].copy()
+        vertices = linear_blend_skinning(self.vertices, matrix_local, matrix, self.skin, pad=1, value=1.)
+        # update matrix_local
+        self.matrix_local = matrix.copy()
+
+        # change tails
+        if self.tails is not None:
+            t_skin = np.eye(self.J)
+            self.tails = linear_blend_skinning(self.tails, matrix_local, matrix, t_skin, pad=1, value=1.)
+        # in accordance with trimesh's normals
+        mesh = trimesh.Trimesh(vertices=vertices, faces=self.faces, process=False)
+        self.vertices = vertices
+        self.vertex_normals = mesh.vertex_normals.copy()
+        self.face_normals = mesh.face_normals.copy()
+    
+    def set_order_by_names(self, new_names: List[str]):
+        assert len(new_names) == len(self.names)
+        name_to_id = {name: id for (id, name) in enumerate(self.names)}
+        new_name_to_id = {name: id for (id, name) in enumerate(new_names)}
+        perm = []
+        new_parents = []
+        for (new_id, name) in enumerate(new_names):
+            perm.append(name_to_id[name])
+            pid = self.parents[name_to_id[name]]
+            if new_id == 0:
+                assert pid is None, 'first bone is not root bone'
+            else:
+                pname = self.names[pid]
+                pid = new_name_to_id[pname]
+                assert pid < new_id, 'new order does not form a tree'
+            new_parents.append(pid)
+        
+        if self.joints is not None:
+            self.joints = self.joints[perm]
+        self.parents = new_parents
+        if self.tails is not None:
+            self.tails = self.tails[perm]
+        if self.skin is not None:
+            self.skin = self.skin[:, perm]
+        if self.no_skin is not None:
+            self.no_skin = self.no_skin[perm]
+        if self.matrix_local is not None:
+            self.matrix_local = self.matrix_local[perm]
+        self.names = new_names
+    
+    def set_order(self, order: Order):
+        if self.names is None or self.parents is None:
+            return
+        new_names, self.parts_bias = order.arrange_names(cls=self.cls, names=self.names, parents=self.parents)
+        self.set_order_by_names(new_names=new_names)
+    
+    def collapse(self, keep: List[str]):
+        dsu = [i for i in range(self.J)]
+        
+        def find(x: int) -> int:
+            if dsu[x] == x:
+                return x
+            y = find(dsu[x])
+            dsu[x] = y
+            return y
+        
+        def merge(x: int, y: int):
+            dsu[find(x)] = find(y)
+
+        if self.tails is not None:
+            new_tails = self.tails.copy()
+        else:
+            new_tails = None
+        if self.skin is not None:
+            new_skin = self.skin.copy()
+        else:
+            new_skin = None
+
+        if self.no_skin is not None:
+            new_no_skin = self.no_skin.copy()
+        else:
+            new_no_skin = None
+        
+        if self.matrix_local is not None:
+            matrix_local = self.matrix_local.copy()
+        else:
+            matrix_local = None
+        new_names = []
+        new_parents = []
+        perm = []
+        new_name_to_id = {}
+        tot = 0
+        for (i, name) in enumerate(self.names):
+            if name in keep:
+                new_names.append(name)
+                new_name_to_id[name] = tot
+                tot += 1
+                perm.append(i)
+                pid = self.parents[i]
+                if pid is None:
+                    new_parents.append(None)
+                else:
+                    pid = find(pid)
+                    new_parents.append(new_name_to_id[self.names[pid]])
+                continue
+            assert i != 0, 'cannot remove root'
+            id = find(i)
+            pid = find(self.parents[id])
+            # be careful !
+            # do not copy tail here because you dont know which child to inherit from
+            if new_skin is not None:
+                new_skin[:, pid] += new_skin[:, id]
+            if new_no_skin is not None:
+                new_no_skin[pid] &= new_no_skin[id]
+            merge(id, pid)
+        
+        if new_tails is not None:
+            new_tails = new_tails[perm]
+        if new_skin is not None:
+            new_skin = new_skin[:, perm]
+        if new_no_skin is not None:
+            new_no_skin = new_no_skin[perm]
+        if matrix_local is not None:
+            matrix_local = matrix_local[perm]
+        
+        if self.joints is not None:
+            self.joints = self.joints[perm]
+        self.parents        = new_parents
+        self.tails          = new_tails
+        self.skin           = new_skin
+        self.no_skin        = new_no_skin
+        self.names          = new_names
+        self.matrix_local   = matrix_local
+    
+    @staticmethod
+    def from_raw_data(
+        raw_data: RawData,
+        cls: str,
+        path: str,
+        data_name: str,
+    ) -> 'Asset':
+        '''
+        Return an asset initialized from raw data and do transform.
+        '''
+        return Asset(
+            cls=cls,
+            path=path,
+            data_name=data_name,
+            vertices=raw_data.vertices,
+            vertex_normals=raw_data.vertex_normals,
+            faces=raw_data.faces,
+            face_normals=raw_data.face_normals,
+            joints=raw_data.joints,
+            tails=raw_data.tails,
+            skin=raw_data.skin,
+            no_skin=raw_data.no_skin,
+            parents=raw_data.parents,
+            names=raw_data.names,
+            matrix_local=raw_data.matrix_local,
+            meta={},
+        )
+    
+    def get_tokenize_input(self) -> TokenizeInput:
+        children = defaultdict(list)
+        
+        for (id, p) in enumerate(self.parents):
+            if p is not None:
+                children[p].append(id)
+        bones = []
+        branch = []
+        is_leaf = []
+        last = None
+        for i in range(self.J):
+            is_leaf.append(len(children[i])==0)
+            if i == 0:
+                bones.append(np.concatenate([self.joints[i], self.joints[i]]))
+                branch.append(False)
+            else:
+                pid = self.parents[i]
+                bones.append(np.concatenate([self.joints[pid], self.joints[i]]))
+                branch.append(pid!=last)
+            last = i
+        bones = np.stack(bones)
+        branch = np.array(branch, dtype=bool)
+        is_leaf = np.array(is_leaf, dtype=bool)
+        return TokenizeInput(
+            bones=bones,
+            tails=self.tails,
+            branch=branch,
+            is_leaf=is_leaf,
+            no_skin=self.no_skin,
+            cls=self.cls,
+            parts_bias=self.parts_bias,
+        )
+    
+    def export_pc(self, path: str, with_normal: bool=True, normal_size=0.01):
+        '''
+        export point cloud
+        '''
+        vertices    = self.vertices
+        normals     = self.vertex_normals
+        if self.sampled_vertices is not None:
+            vertices    = self.sampled_vertices
+            normals     = self.sampled_normals
+        if with_normal == False:
+            normals = None
+        self._export_pc(vertices=vertices, path=path, vertex_normals=normals, normal_size=normal_size)
+    
+    def export_mesh(self, path: str):
+        '''
+        export mesh
+        '''
+        self._export_mesh(vertices=self.vertices, faces=self.faces, path=path)
+    
+    def export_skeleton(self, path: str):
+        '''
+        export spring
+        '''
+        self._export_skeleton(joints=self.joints, parents=self.parents, path=path)
+    
+    def export_skeleton_sequence(self, path: str):
+        '''
+        export spring
+        '''
+        self._export_skeleton_sequence(joints=self.joints, parents=self.parents, path=path)
+    
+    def export_fbx(
+        self,
+        path: str,
+        vertex_group_name: str,
+        extrude_size: float=0.03,
+        group_per_vertex: int=-1,
+        add_root: bool=False,
+        do_not_normalize: bool=False,
+        use_extrude_bone: bool=True,
+        use_connect_unique_child: bool=True,
+        extrude_from_parent: bool=True,
+        use_tail: bool=False,
+        use_origin: bool=False,
+    ):
+        '''
+        export the whole model with skining
+        '''
+        self._export_fbx(
+            path=path,
+            vertices=self.vertices if use_origin else self.sampled_vertices,
+            joints=self.joints,
+            skin=self.sampled_vertex_groups[vertex_group_name],
+            parents=self.parents,
+            names=self.names,
+            faces=self.faces if use_origin else None,
+            extrude_size=extrude_size,
+            group_per_vertex=group_per_vertex,
+            add_root=add_root,
+            do_not_normalize=do_not_normalize,
+            use_extrude_bone=use_extrude_bone,
+            use_connect_unique_child=use_connect_unique_child,
+            extrude_from_parent=extrude_from_parent,
+            tails=self.tails if use_tail else None,
+        )
+
+    def export_render(self, path: str, resolution: Tuple[int, int]=[256, 256], use_tail: bool=False):
+        if use_tail:
+            assert self.tails is not None
+            self._export_render(
+                path=path,
+                vertices=self.vertices,
+                faces=self.faces,
+                bones=np.concatenate([self.joints, self.tails], axis=-1),
+                resolution=resolution,
+            )
+        else:
+            pjoints = self.joints[self.parents[1:]]
+            self._export_render(
+                path=path,
+                vertices=self.vertices,
+                faces=self.faces,
+                bones=np.concatenate([pjoints, self.joints[1:]], axis=-1),
+                resolution=resolution,
+            )

UniRig/src/data/augment.py

@@ -0,0 +1,152 @@
+from dataclasses import dataclass
+from typing import Tuple, Union, List, Dict
+from numpy import ndarray
+import numpy as np
+from abc import ABC, abstractmethod
+from scipy.spatial.transform import Rotation as R
+
+from .spec import ConfigSpec
+from .asset import Asset
+from .utils import axis_angle_to_matrix
+
+@dataclass(frozen=True)
+class AugmentAffineConfig(ConfigSpec):
+    # final normalization cube
+    normalize_into: Tuple[float, float]
+
+    # randomly scale coordinates with probability p
+    random_scale_p: float
+    
+    # scale range (lower, upper)
+    random_scale: Tuple[float, float]
+    
+    # randomly shift coordinates with probability p
+    random_shift_p: float
+    
+    # shift range (lower, upper)
+    random_shift: Tuple[float, float]
+    
+    @classmethod
+    def parse(cls, config) -> Union['AugmentAffineConfig', None]:
+        if config is None:
+            return None
+        cls.check_keys(config)
+        return AugmentAffineConfig(
+            normalize_into=config.normalize_into,
+            random_scale_p=config.get('random_scale_p', 0.),
+            random_scale=config.get('random_scale', [1., 1.]),
+            random_shift_p=config.get('random_shift_p', 0.),
+            random_shift=config.get('random_shift', [0., 0.]),
+        )
+
+@dataclass(frozen=True)
+class AugmentConfig(ConfigSpec):
+    '''
+    Config to handle final easy augmentation of vertices, normals and bones before sampling.
+    '''    
+    augment_affine_config: Union[AugmentAffineConfig, None]
+    
+    @classmethod
+    def parse(cls, config) -> 'AugmentConfig':
+        cls.check_keys(config)
+        return AugmentConfig(
+            augment_affine_config=AugmentAffineConfig.parse(config.get('augment_affine_config', None)),
+        )
+
+class Augment(ABC):
+    '''
+    Abstract class for augmentation
+    '''
+    def __init__(self):
+        pass
+    
+    @abstractmethod
+    def transform(self, asset: Asset, **kwargs):
+        pass
+
+    @abstractmethod
+    def inverse(self, asset: Asset):
+        pass
+
+class AugmentAffine(Augment):
+    
+    def __init__(self, config: AugmentAffineConfig):
+        super().__init__()
+        self.config = config
+
+    def _apply(self, v: ndarray, trans: ndarray) -> ndarray:
+        return np.matmul(v, trans[:3, :3].transpose()) + trans[:3, 3]
+
+    def transform(self, asset: Asset, **kwargs):
+        bound_min = asset.vertices.min(axis=0)
+        bound_max = asset.vertices.max(axis=0)
+        if asset.joints is not None:
+            joints_bound_min = asset.joints.min(axis=0)
+            joints_bound_max = asset.joints.max(axis=0)            
+            bound_min = np.minimum(bound_min, joints_bound_min)
+            bound_max = np.maximum(bound_max, joints_bound_max)
+        
+        trans_vertex = np.eye(4, dtype=np.float32)
+        
+        trans_vertex = _trans_to_m(-(bound_max + bound_min)/2) @ trans_vertex
+        
+        # scale into the cube
+        normalize_into = self.config.normalize_into
+        scale = np.max((bound_max - bound_min) / (normalize_into[1] - normalize_into[0]))
+        trans_vertex = _scale_to_m(1. / scale) @ trans_vertex
+        
+        bias = (normalize_into[0] + normalize_into[1]) / 2
+        trans_vertex = _trans_to_m(np.array([bias, bias, bias], dtype=np.float32)) @ trans_vertex
+        
+        if np.random.rand() < self.config.random_scale_p:
+            scale = _scale_to_m(np.random.uniform(self.config.random_scale[0], self.config.random_scale[1]))
+            trans_vertex = scale @ trans_vertex
+            
+        if np.random.rand() < self.config.random_shift_p:
+            l, r = self.config.random_shift
+            shift = _trans_to_m(np.array([np.random.uniform(l, r), np.random.uniform(l, r), np.random.uniform(l, r)]), dtype=np.float32)
+            trans_vertex = shift @ trans_vertex
+        
+        asset.vertices = self._apply(asset.vertices, trans_vertex)
+        # do not affect scale in matrix
+        if asset.matrix_local is not None:
+            asset.matrix_local[:, :, 3:4] = trans_vertex @ asset.matrix_local[:, :, 3:4]
+        if asset.pose_matrix is not None:
+            asset.pose_matrix[:, :, 3:4] = trans_vertex @ asset.pose_matrix[:, :, 3:4]
+        # do not affect normal here
+        if asset.joints is not None:
+            asset.joints = self._apply(asset.joints, trans_vertex)
+        if asset.tails is not None:
+            asset.tails = self._apply(asset.tails, trans_vertex)
+        
+        self.trans_vertex = trans_vertex
+    
+    def inverse(self, asset: Asset):
+        m = np.linalg.inv(self.trans_vertex)
+        asset.vertices = self._apply(asset.vertices, m)
+        if asset.joints is not None:
+            asset.joints = self._apply(asset.joints, m)
+        if asset.tails is not None:
+            asset.tails = self._apply(asset.tails, m)
+
+def _trans_to_m(v: ndarray):
+    m = np.eye(4, dtype=np.float32)
+    m[0:3, 3] = v
+    return m
+
+def _scale_to_m(r: ndarray):
+    m = np.zeros((4, 4), dtype=np.float32)
+    m[0, 0] = r
+    m[1, 1] = r
+    m[2, 2] = r
+    m[3, 3] = 1.
+    return m
+
+def get_augments(config: AugmentConfig) -> Tuple[List[Augment], List[Augment]]:
+    first_augments  = [] # augments before sample
+    second_augments = [] # augments after sample
+    augment_affine_config           = config.augment_affine_config
+
+    if augment_affine_config is not None:
+        second_augments.append(AugmentAffine(config=augment_affine_config))
+    return first_augments, second_augments

UniRig/src/data/datapath.py

@@ -0,0 +1,149 @@
+from copy import deepcopy
+from collections import defaultdict
+from dataclasses import dataclass
+from typing import Dict, Union, Tuple, List
+import numpy as np
+from numpy import ndarray
+import os
+from random import shuffle
+from box import Box
+from torch.onnx.symbolic_opset11 import index_copy
+
+from .spec import ConfigSpec
+
+@dataclass
+class DatapathConfig(ConfigSpec):
+    '''
+    Config to handle input data paths.
+    '''
+    # root
+    input_dataset_dir: str
+    
+    # use proportion data sampling
+    use_prob: bool
+    
+    # cls: [(path_1, p_1), ...]
+    data_path: Dict[str, List[Tuple[str, float]]]
+    
+    # how many files to return when using data sampling
+    num_files: Union[int, None]
+    
+    @classmethod
+    def from_args(cls, **kwargs) -> 'DatapathConfig':
+        '''
+        Make a temporary datapath from user inputs.
+        '''
+        input = kwargs.get('input', None)
+        output = kwargs.get('output', None)
+        recursive = kwargs.get('recursive', False)
+        
+    
+    @classmethod
+    def parse(cls, config) -> 'DatapathConfig':
+        cls.check_keys(config)
+        return DatapathConfig(
+            input_dataset_dir=config.input_dataset_dir,
+            use_prob=config.get('use_prob', True),
+            data_path=config.data_path,
+            num_files=config.get('num_files', None),
+        )
+    
+    def split_by_cls(self) -> Dict[str, 'DatapathConfig']:
+        res: Dict[str, DatapathConfig] = {}
+        for cls in self.data_path:
+            res[cls] = deepcopy(self)
+            res[cls].data_path = {cls: self.data_path[cls]}
+        return res
+
+class Datapath():
+    def __init__(
+        self,
+        config: Union[DatapathConfig, None]=None,
+        files: Union[List[str], None]=None,
+        cls: Union[str, None]=None,
+    ):
+        if config is not None:
+            self.config     = config
+            self.file_list  = []
+            cls_probs_first = []
+            cls_first       = []
+            
+            self.files_by_class: Dict[str, List[Dict]] = defaultdict(list)
+            self.class_positions: Dict[str, List[int]] = defaultdict(list)
+            self.cls_probs_second: Dict[str, ndarray] = defaultdict(List)
+            
+            for cls in self.config.data_path:
+                prob = 0.
+                probs_second = []
+                for (path, p) in self.config.data_path[cls]:
+                    prob += p
+                    probs_second.append(p)
+                    with open(path, 'r') as f:
+                        file_items = []
+                        missing = 0
+                        for l in f.readlines():
+                            raw_data_path = os.path.join(self.config.input_dataset_dir, l.strip(), 'raw_data.npz')
+                            if not os.path.exists(raw_data_path):
+                                missing += 1
+                                continue
+                            file_items.append({
+                                'cls': cls,
+                                'path': os.path.join(self.config.input_dataset_dir, l.strip()),
+                                'prob': p
+                            })
+                        assert len(file_items) > 0, f"files in {path} are all missing! root: {self.config.input_dataset_dir}"
+                        if missing > 0:
+                            print(f"\033[31m{cls}: {missing} missing files\033[0m")
+                        self.files_by_class[cls].append(file_items)
+                        self.class_positions[cls].append(0)
+                        self.file_list.extend(file_items)
+                probs_second = np.array(probs_second)
+                self.cls_probs_second[cls] = probs_second / probs_second.sum()
+                cls_first.append(cls)
+                cls_probs_first.append(prob)
+            cls_probs_first = np.array(cls_probs_first)
+            self.cls_first: List[str] = cls_first
+            self.cls_probs_first: Dict[str, List[float]] = cls_probs_first / cls_probs_first.sum()
+        elif files is not None:
+            if cls is None:
+                cls = 'inference'
+            self.file_list  = [{'cls': cls, 'path': file} for file in files]
+            cls_probs_first = np.array([1.])
+            cls_first       = []
+            
+            self.files_by_class: Dict[str, List[Dict]] = {cls: self.file_list.copy()}
+            self.class_positions: Dict[str, List[int]] = {cls: [0]}
+            self.cls_probs_second: Dict[str, ndarray] = {cls: np.array([1.])}
+            self.config = Box({'use_prob': False})
+        else:
+            assert(0)
+    
+    def __len__(self):
+        if self.config.use_prob:
+            assert self.config.num_files is not None, 'num_files is not specified'
+            return self.config.num_files
+        return len(self.file_list)
+    
+    def __getitem__(self, index) -> Tuple[str, str]:
+        if self.config.use_prob:
+            # first sample a class
+            cls = np.random.choice(self.cls_first, p=self.cls_probs_first)
+            
+            # second sample in this class
+            idx = np.random.choice(len(self.files_by_class[cls]), p=self.cls_probs_second[cls])
+            
+            # get the current position
+            pos = self.class_positions[cls][idx]
+            files = self.files_by_class[cls][idx]
+            
+            # get the item andd update position
+            item = files[pos]
+            self.class_positions[cls][idx] = (pos + 1) % len(files)
+            if (pos + 1) % len(files) == 0:
+                shuffle(self.files_by_class[cls][idx])
+        else:
+            item = self.file_list[index]
+        return (item['cls'], item['path'])
+    
+    def get_data(self) -> List[Tuple[str, str]]:
+        return [self[i] for i in range(len(self))]

UniRig/src/data/dataset.py

@@ -0,0 +1,231 @@
+from copy import deepcopy
+from dataclasses import dataclass
+import lightning.pytorch as pl
+# from lightning.pytorch.utilities.types import EVAL_DATALOADERS, TRAIN_DATALOADERS
+import torch
+from torch import LongTensor
+from torch.utils import data
+from torch.utils.data import DataLoader, Dataset
+from typing import Dict, List, Tuple, Union, Callable
+import os
+import numpy as np
+
+from .raw_data import RawData
+from .asset import Asset
+from .transform import TransformConfig, transform_asset
+from .datapath import DatapathConfig, Datapath
+from .spec import ConfigSpec
+
+from ..tokenizer.spec import TokenizerSpec, TokenizerConfig
+from ..tokenizer.parse import get_tokenizer
+from ..model.spec import ModelInput
+
+@dataclass
+class DatasetConfig(ConfigSpec):
+    '''
+    Config to handle dataset format.
+    '''
+    # shuffle dataset
+    shuffle: bool
+
+    # batch size
+    batch_size: int
+
+    # number of workers
+    num_workers: int
+    
+    # datapath
+    datapath_config: DatapathConfig
+    
+    # use pin memory
+    pin_memory: bool = True
+    
+    # use persistent workers
+    persistent_workers: bool = True
+    
+    @classmethod
+    def parse(cls, config) -> 'DatapathConfig':
+        cls.check_keys(config)
+        return DatasetConfig(
+            shuffle=config.shuffle,
+            batch_size=config.batch_size,
+            num_workers=config.num_workers,
+            pin_memory=config.pin_memory,
+            persistent_workers=config.persistent_workers,
+            datapath_config=DatapathConfig.parse(config.datapath_config),
+        )
+    
+    def split_by_cls(self) -> Dict[str, 'DatasetConfig']:
+        res: Dict[str, DatasetConfig] = {}
+        datapath_config_dict = self.datapath_config.split_by_cls()
+        for cls in self.datapath_config.data_path:
+            res[cls] = deepcopy(self)
+            res[cls].datapath_config = datapath_config_dict[cls]
+        return res
+
+class UniRigDatasetModule(pl.LightningDataModule):  
+    def __init__(
+        self,
+        process_fn: Union[Callable[[List[ModelInput]], Dict]]=None,
+        predict_dataset_config: Union[Dict[str, DatasetConfig], None]=None,
+        predict_transform_config: Union[TransformConfig, None]=None,
+        tokenizer_config: Union[TokenizerConfig, None]=None,
+        debug: bool=False,
+        data_name: str='raw_data.npz',
+        datapath: Union[Datapath, None]=None,
+        cls: Union[str, None]=None,
+    ):
+        super().__init__()
+        self.process_fn                 = process_fn
+        self.predict_dataset_config     = predict_dataset_config
+        self.predict_transform_config   = predict_transform_config
+        self.tokenizer_config           = tokenizer_config
+        self.debug                      = debug
+        self.data_name                  = data_name
+        
+        if debug:
+            print("\033[31mWARNING: debug mode, dataloader will be extremely slow !!!\033[0m")
+        
+        if datapath is not None:
+            self.train_datapath = None
+            self.validate_datapath = None
+            self.predict_datapath = {
+                cls: deepcopy(datapath),
+            }
+            self.predict_dataset_config = {
+                cls: DatasetConfig(
+                    shuffle=False,
+                    batch_size=1,
+                    num_workers=0,
+                    datapath_config=deepcopy(datapath),
+                    pin_memory=False,
+                    persistent_workers=False,
+                )
+            }
+        else:
+            # build predict datapath
+            if self.predict_dataset_config is not None:
+                self.predict_datapath = {
+                    cls: Datapath(self.predict_dataset_config[cls].datapath_config)
+                    for cls in self.predict_dataset_config
+                }
+            else:
+                self.predict_datapath = None
+        
+        # get tokenizer
+        if tokenizer_config is None:
+            self.tokenizer = None
+        else:
+            self.tokenizer = get_tokenizer(config=tokenizer_config)
+
+    def prepare_data(self):
+        pass
+
+    def setup(self, stage=None):   
+        if self.predict_datapath is not None:
+            self._predict_ds = {}
+            for cls in self.predict_datapath:
+                self._predict_ds[cls] = UniRigDataset(
+                    process_fn=self.process_fn,
+                    data=self.predict_datapath[cls].get_data(),
+                    name=f"predict-{cls}",
+                    tokenizer=self.tokenizer,
+                    transform_config=self.predict_transform_config,
+                    debug=self.debug,
+                    data_name=self.data_name,
+                )
+    
+    def predict_dataloader(self):
+        if not hasattr(self, "_predict_ds"):
+            self.setup()
+        return self._create_dataloader(
+            dataset=self._predict_ds,
+            config=self.predict_dataset_config,
+            is_train=False,
+            drop_last=False,
+        )
+
+    def _create_dataloader(
+        self,
+        dataset: Union[Dataset, Dict[str, Dataset]],
+        config: DatasetConfig,
+        is_train: bool,
+        **kwargs,
+    ) -> Union[DataLoader, Dict[str, DataLoader]]:
+        def create_single_dataloader(dataset, config: Union[DatasetConfig, Dict[str, DatasetConfig]], **kwargs):
+            return DataLoader(
+                dataset,
+                batch_size=config.batch_size,
+                shuffle=config.shuffle,
+                num_workers=config.num_workers,
+                pin_memory=config.pin_memory,
+                persistent_workers=config.persistent_workers,
+                collate_fn=dataset.collate_fn,
+                **kwargs,
+            )
+        if isinstance(dataset, Dict):
+            return {k: create_single_dataloader(v, config[k], **kwargs) for k, v in dataset.items()}
+        else:
+            return create_single_dataloader(dataset, config, **kwargs)
+
+class UniRigDataset(Dataset):
+    def __init__(
+        self,
+        data: List[Tuple[str, str]], # (cls, part)
+        name: str,
+        process_fn: Union[Callable[[List[ModelInput]], Dict]]=None,
+        tokenizer: Union[TokenizerSpec, None]=None,
+        transform_config: Union[TransformConfig, None]=None,
+        debug: bool=False,
+        data_name: str='raw_data.npz',
+    ) -> None:
+        super().__init__()
+        
+        self.data               = data
+        self.name               = name
+        self.process_fn         = process_fn
+        self.tokenizer          = tokenizer
+        self.transform_config   = transform_config
+        self.debug              = debug
+        self.data_name          = data_name
+        
+        if not debug:
+            assert self.process_fn is not None, 'missing data processing function'
+
+    def __len__(self) -> int:
+        return len(self.data)
+
+    def __getitem__(self, idx) -> ModelInput:
+        cls, dir_path = self.data[idx]
+        raw_data = RawData.load(path=os.path.join(dir_path, self.data_name))
+        asset = Asset.from_raw_data(raw_data=raw_data, cls=cls, path=dir_path, data_name=self.data_name)
+        
+        first_augments, second_augments = transform_asset(
+            asset=asset,
+            transform_config=self.transform_config,
+        )
+        if self.tokenizer is not None and asset.parents is not None:
+            tokens = self.tokenizer.tokenize(input=asset.get_tokenize_input())
+        else:
+            tokens = None
+        return ModelInput(
+            tokens=tokens,
+            pad=None if self.tokenizer is None else self.tokenizer.pad,
+            vertices=asset.sampled_vertices.astype(np.float32),
+            normals=asset.sampled_normals.astype(np.float32),
+            joints=None if asset.joints is None else asset.joints.astype(np.float32),
+            tails=None if asset.tails is None else asset.tails.astype(np.float32),
+            asset=asset,
+            augments=None,
+        )
+
+    def _collate_fn_debug(self, batch):
+        return batch
+    
+    def _collate_fn(self, batch):
+        return data.dataloader.default_collate(self.process_fn(batch))
+
+    def collate_fn(self, batch):
+        if self.debug:
+            return self._collate_fn_debug(batch)
+        return self._collate_fn(batch)

UniRig/src/data/exporter.py

@@ -0,0 +1,486 @@
+import numpy as np
+from numpy import ndarray
+from typing import List, Union, Tuple
+from collections import defaultdict
+import os
+
+try:
+    import open3d as o3d
+    OPEN3D_EQUIPPED = True
+except:
+    print("do not have open3d")
+    OPEN3D_EQUIPPED = False
+
+class Exporter():
+    
+    def _safe_make_dir(self, path):
+        if os.path.dirname(path) == '':
+            return
+        os.makedirs(os.path.dirname(path), exist_ok=True)
+    
+    def _export_skeleton(self, joints: ndarray, parents: List[Union[int, None]], path: str):
+        format = path.split('.')[-1]
+        assert format in ['obj']
+        name = path.removesuffix('.obj')
+        path = name + ".obj"
+        self._safe_make_dir(path)
+        J = joints.shape[0]
+        with open(path, 'w') as file:
+            file.write("o spring_joint\n")
+            _joints = []
+            for id in range(J):
+                pid = parents[id]
+                if pid is None or pid == -1:
+                    continue
+                bx, by, bz = joints[id]
+                ex, ey, ez = joints[pid]
+                _joints.extend([
+                    f"v {bx} {bz} {-by}\n",
+                    f"v {ex} {ez} {-ey}\n",
+                    f"v {ex} {ez} {-ey + 0.00001}\n"
+                ])
+            file.writelines(_joints)
+            
+            _faces = [f"f {id*3+1} {id*3+2} {id*3+3}\n" for id in range(J)]
+            file.writelines(_faces)
+    
+    def _export_bones(self, bones: ndarray, path: str):
+        format = path.split('.')[-1]
+        assert format in ['obj']
+        name = path.removesuffix('.obj')
+        path = name + ".obj"
+        self._safe_make_dir(path)
+        J = bones.shape[0]
+        with open(path, 'w') as file:
+            file.write("o bones\n")
+            _joints = []
+            for bone in bones:
+                bx, by, bz = bone[:3]
+                ex, ey, ez = bone[3:]
+                _joints.extend([
+                    f"v {bx} {bz} {-by}\n",
+                    f"v {ex} {ez} {-ey}\n",
+                    f"v {ex} {ez} {-ey + 0.00001}\n"
+                ])
+            file.writelines(_joints)
+            
+            _faces = [f"f {id*3+1} {id*3+2} {id*3+3}\n" for id in range(J)]
+            file.writelines(_faces)
+    
+    def _export_skeleton_sequence(self, joints: ndarray, parents: List[Union[int, None]], path: str):
+        format = path.split('.')[-1]
+        assert format in ['obj']
+        name = path.removesuffix('.obj')
+        path = name + ".obj"
+        self._safe_make_dir(path)
+        J = joints.shape[0]
+        for i in range(J):
+            file = open(name + f"_{i}.obj", 'w')
+            file.write("o spring_joint\n")
+            _joints = []
+            for id in range(i + 1):
+                pid = parents[id]
+                if pid is None:
+                    continue
+                bx, by, bz = joints[id]
+                ex, ey, ez = joints[pid]
+                _joints.extend([
+                    f"v {bx} {bz} {-by}\n",
+                    f"v {ex} {ez} {-ey}\n",
+                    f"v {ex} {ez} {-ey + 0.00001}\n"
+                ])
+            file.writelines(_joints)
+            
+            _faces = [f"f {id*3+1} {id*3+2} {id*3+3}\n" for id in range(J)]
+            file.writelines(_faces)
+            file.close()
+    
+    def _export_mesh(self, vertices: ndarray, faces: ndarray, path: str):
+        format = path.split('.')[-1]
+        assert format in ['obj', 'ply']
+        if path.endswith('ply'):
+            if not OPEN3D_EQUIPPED:
+                raise RuntimeError("open3d is not available")
+            mesh = o3d.geometry.TriangleMesh()
+            mesh.vertices = o3d.utility.Vector3dVector(vertices)
+            mesh.triangles = o3d.utility.Vector3iVector(faces)
+            self._safe_make_dir(path)
+            o3d.io.write_triangle_mesh(path, mesh)
+            return
+        name = path.removesuffix('.obj')
+        path = name + ".obj"
+        self._safe_make_dir(path)
+        with open(path, 'w') as file:
+            file.write("o mesh\n")
+            _vertices = []
+            for co in vertices:
+                _vertices.append(f"v {co[0]} {co[2]} {-co[1]}\n")
+            file.writelines(_vertices)
+            _faces = []
+            for face in faces:
+                _faces.append(f"f {face[0]+1} {face[1]+1} {face[2]+1}\n")
+            file.writelines(_faces)
+            
+    def _export_pc(self, vertices: ndarray, path: str, vertex_normals: Union[ndarray, None]=None, normal_size: float=0.01):
+        if path.endswith('.ply'):
+            if vertex_normals is not None:
+                print("normal result will not be displayed in .ply format")
+            name = path.removesuffix('.ply')
+            path = name + ".ply"
+            pc = o3d.geometry.PointCloud()
+            pc.points = o3d.utility.Vector3dVector(vertices)
+            # segment fault when numpy >= 2.0 !! use torch environment
+            self._safe_make_dir(path)
+            o3d.io.write_point_cloud(path, pc)
+            return
+        name = path.removesuffix('.obj')
+        path = name + ".obj"
+        self._safe_make_dir(path)
+        with open(path, 'w') as file:
+            file.write("o pc\n")
+            _vertex = []
+            for co in vertices:
+                _vertex.append(f"v {co[0]} {co[2]} {-co[1]}\n")
+            file.writelines(_vertex)
+            if vertex_normals is not None:
+                new_path = path.replace('.obj', '_normal.obj')
+                nfile = open(new_path, 'w')
+                nfile.write("o normal\n")
+                _normal = []
+                for i in range(vertices.shape[0]):
+                    co = vertices[i]
+                    x = vertex_normals[i, 0]
+                    y = vertex_normals[i, 1]
+                    z = vertex_normals[i, 2]
+                    _normal.extend([
+                        f"v {co[0]} {co[2]} {-co[1]}\n",
+                        f"v {co[0]+0.0001} {co[2]} {-co[1]}\n",
+                        f"v {co[0]+x*normal_size} {co[2]+z*normal_size} {-(co[1]+y*normal_size)}\n",
+                        f"f {i*3+1} {i*3+2} {i*3+3}\n",
+                    ])
+                nfile.writelines(_normal)
+    
+    def _make_armature(
+        self,
+        vertices: Union[ndarray, None],
+        joints: ndarray,
+        skin: Union[ndarray, None],
+        parents: List[Union[int, None]],
+        names: List[str],
+        faces: Union[ndarray, None]=None,
+        extrude_size: float=0.03,
+        group_per_vertex: int=-1,
+        add_root: bool=False,
+        do_not_normalize: bool=False,
+        use_extrude_bone: bool=True,
+        use_connect_unique_child: bool=True,
+        extrude_from_parent: bool=True,
+        tails: Union[ndarray, None]=None,
+    ):
+        import bpy # type: ignore
+        from mathutils import Vector # type: ignore
+        
+        # make collection
+        collection = bpy.data.collections.new('new_collection')
+        bpy.context.scene.collection.children.link(collection)
+        
+        # make mesh
+        if vertices is not None:
+            mesh = bpy.data.meshes.new('mesh')
+            if faces is None:
+                faces = []
+            mesh.from_pydata(vertices, [], faces)
+            mesh.update()
+        
+            # make object from mesh
+            object = bpy.data.objects.new('character', mesh)
+        
+            # add object to scene collection
+            collection.objects.link(object)
+        
+        # deselect mesh
+        bpy.ops.object.armature_add(enter_editmode=True)
+        armature = bpy.data.armatures.get('Armature')
+        edit_bones = armature.edit_bones
+        
+        J = joints.shape[0]
+        if tails is None:
+            tails = joints.copy()
+            tails[:, 2] += extrude_size
+        connects = [False for _ in range(J)]
+        children = defaultdict(list)
+        for i in range(1, J):
+            children[parents[i]].append(i)
+        if tails is not None:
+            if use_extrude_bone:
+                for i in range(J):
+                    if len(children[i]) != 1 and extrude_from_parent and i != 0:
+                        pjoint = joints[parents[i]]
+                        joint = joints[i]
+                        d = joint - pjoint
+                        if np.linalg.norm(d) < 0.000001:
+                            d = np.array([0., 0., 1.]) # in case son.head == parent.head
+                        else:
+                            d = d / np.linalg.norm(d)
+                        tails[i] = joint + d * extrude_size
+            if use_connect_unique_child:
+                for i in range(J):
+                    if len(children[i]) == 1:
+                        child = children[i][0]
+                        tails[i] = joints[child]
+                    if parents[i] is not None and len(children[parents[i]]) == 1:
+                        connects[i] = True
+        
+        if add_root:
+            bone_root = edit_bones.get('Bone')
+            bone_root.name = 'Root'
+            bone_root.tail = Vector((joints[0, 0], joints[0, 1], joints[0, 2]))
+        else:
+            bone_root = edit_bones.get('Bone')
+            bone_root.name = names[0]
+            bone_root.head = Vector((joints[0, 0], joints[0, 1], joints[0, 2]))
+            bone_root.tail = Vector((joints[0, 0], joints[0, 1], joints[0, 2] + extrude_size))
+        
+        def extrude_bone(
+            edit_bones,
+            name: str,
+            parent_name: str,
+            head: Tuple[float, float, float],
+            tail: Tuple[float, float, float],
+            connect: bool
+        ):
+            bone = edit_bones.new(name)
+            bone.head = Vector((head[0], head[1], head[2]))
+            bone.tail = Vector((tail[0], tail[1], tail[2]))
+            bone.name = name
+            parent_bone = edit_bones.get(parent_name)
+            bone.parent = parent_bone
+            bone.use_connect = connect
+            assert not np.isnan(head).any(), f"nan found in head of bone {name}"
+            assert not np.isnan(tail).any(), f"nan found in tail of bone {name}"
+        
+        for i in range(J):
+            if add_root is False and i==0:
+                continue
+            edit_bones = armature.edit_bones
+            pname = 'Root' if parents[i] is None else names[parents[i]]
+            extrude_bone(edit_bones, names[i], pname, joints[i], tails[i], connects[i])
+        for i in range(J):
+            bone = edit_bones.get(names[i])
+            bone.head = Vector((joints[i, 0], joints[i, 1], joints[i, 2]))
+            bone.tail = Vector((tails[i, 0], tails[i, 1], tails[i, 2]))
+        
+        if vertices is None or skin is None:
+            return
+        # must set to object mode to enable parent_set
+        bpy.ops.object.mode_set(mode='OBJECT')
+        objects = bpy.data.objects
+        for o in bpy.context.selected_objects:
+            o.select_set(False)
+        ob = objects['character']
+        arm = bpy.data.objects['Armature']
+        ob.select_set(True)
+        arm.select_set(True)
+        bpy.ops.object.parent_set(type='ARMATURE_NAME')
+        vis = []
+        for x in ob.vertex_groups:
+            vis.append(x.name)
+        #sparsify
+        argsorted = np.argsort(-skin, axis=1)
+        vertex_group_reweight = skin[np.arange(skin.shape[0])[..., None], argsorted]
+        if group_per_vertex == -1:
+            group_per_vertex = vertex_group_reweight.shape[-1]
+        if not do_not_normalize:
+            vertex_group_reweight = vertex_group_reweight / vertex_group_reweight[..., :group_per_vertex].sum(axis=1)[...,None]
+
+        for v, w in enumerate(skin):
+            for ii in range(group_per_vertex):
+                i = argsorted[v, ii]
+                if i >= J:
+                    continue
+                n = names[i]
+                if n not in vis:
+                    continue
+                ob.vertex_groups[n].add([v], vertex_group_reweight[v, ii], 'REPLACE')
+
+    def _clean_bpy(self):
+        import bpy # type: ignore
+        for c in bpy.data.actions:
+            bpy.data.actions.remove(c)
+        for c in bpy.data.armatures:
+            bpy.data.armatures.remove(c)
+        for c in bpy.data.cameras:
+            bpy.data.cameras.remove(c)
+        for c in bpy.data.collections:
+            bpy.data.collections.remove(c)
+        for c in bpy.data.images:
+            bpy.data.images.remove(c)
+        for c in bpy.data.materials:
+            bpy.data.materials.remove(c)
+        for c in bpy.data.meshes:
+            bpy.data.meshes.remove(c)
+        for c in bpy.data.objects:
+            bpy.data.objects.remove(c)
+        for c in bpy.data.textures:
+            bpy.data.textures.remove(c)
+    
+    def _export_fbx(
+        self,
+        path: str,
+        vertices: Union[ndarray, None],
+        joints: ndarray,
+        skin: Union[ndarray, None],
+        parents: List[Union[int, None]],
+        names: List[str],
+        faces: Union[ndarray, None]=None,
+        extrude_size: float=0.03,
+        group_per_vertex: int=-1,
+        add_root: bool=False,
+        do_not_normalize: bool=False,
+        use_extrude_bone: bool=True,
+        use_connect_unique_child: bool=True,
+        extrude_from_parent: bool=True,
+        tails: Union[ndarray, None]=None,
+    ):
+        '''
+        Requires bpy installed
+        '''
+        import bpy # type: ignore
+        self._safe_make_dir(path)
+        self._clean_bpy()
+        self._make_armature(
+            vertices=vertices,
+            joints=joints,
+            skin=skin,
+            parents=parents,
+            names=names,
+            faces=faces,
+            extrude_size=extrude_size,
+            group_per_vertex=group_per_vertex,
+            add_root=add_root,
+            do_not_normalize=do_not_normalize,
+            use_extrude_bone=use_extrude_bone,
+            use_connect_unique_child=use_connect_unique_child,
+            extrude_from_parent=extrude_from_parent,
+            tails=tails,
+        )
+        
+        # always enable add_leaf_bones to keep leaf bones
+        bpy.ops.export_scene.fbx(filepath=path, check_existing=False, add_leaf_bones=False)
+    
+    def _export_render(
+        self,
+        path: str,
+        vertices: Union[ndarray, None],
+        faces: Union[ndarray, None],
+        bones: Union[ndarray, None],
+        resolution: Tuple[float, float]=[256, 256],
+    ):
+        import bpy # type: ignore
+        import bpy_extras # type: ignore
+        from mathutils import Vector # type: ignore
+        
+        self._safe_make_dir(path)
+        # normalize into [-1, 1]^3
+        # copied from augment
+        assert (vertices is not None) or (bones is not None)
+        bounds = []
+        if vertices is not None:
+            bounds.append(vertices)
+        if bones is not None:
+            bounds.append(bones[:, :3])
+            bounds.append(bones[:, 3:])
+        bounds = np.concatenate(bounds, axis=0)
+        bound_min = bounds.min(axis=0)
+        bound_max = bounds.max(axis=0)
+        
+        trans_vertex = np.eye(4)
+        
+        trans_vertex = _trans_to_m(-(bound_max + bound_min)/2) @ trans_vertex
+        
+        # scale into the cube [-1, 1]
+        scale = np.max((bound_max - bound_min) / 2)
+        trans_vertex = _scale_to_m(1. / scale) @ trans_vertex
+        
+        def _apply(v: ndarray, trans: ndarray) -> ndarray:
+            return np.matmul(v, trans[:3, :3].transpose()) + trans[:3, 3]
+        
+        if vertices is not None:
+            vertices = _apply(vertices, trans_vertex)
+        if bones is not None:
+            bones[:, :3] = _apply(bones[:, :3], trans_vertex)
+            bones[:, 3:] = _apply(bones[:, 3:], trans_vertex)
+        
+        # bpy api calls
+        self._clean_bpy()
+        bpy.context.scene.render.engine = 'BLENDER_WORKBENCH'
+        bpy.context.scene.render.film_transparent = True
+        bpy.context.scene.display.shading.background_type = 'VIEWPORT'
+        
+        collection = bpy.data.collections.new('new_collection')
+        bpy.context.scene.collection.children.link(collection)
+        
+        if vertices is not None:
+            mesh_data = bpy.data.meshes.new(name="MeshData")
+            mesh_obj = bpy.data.objects.new(name="MeshObject", object_data=mesh_data)
+            collection.objects.link(mesh_obj)
+
+            mesh_data.from_pydata((vertices).tolist(), [], faces.tolist())
+            mesh_data.update()
+
+        def look_at(camera, point):
+            direction = point - camera.location
+            rot_quat = direction.to_track_quat('-Z', 'Y')
+            camera.rotation_euler = rot_quat.to_euler()
+        
+        bpy.ops.object.camera_add(location=(4, -4, 2.5))
+        camera = bpy.context.object
+        camera.data.angle = np.radians(25.0)
+        look_at(camera, Vector((0, 0, -0.2)))
+        bpy.context.scene.camera = camera
+
+        bpy.context.scene.render.resolution_x = resolution[0]
+        bpy.context.scene.render.resolution_y = resolution[1]
+        bpy.context.scene.render.image_settings.file_format = 'PNG'
+        bpy.context.scene.render.filepath = path
+
+        bpy.ops.render.render(write_still=True)
+        # some AI generated code to draw bones over mesh
+        if bones is not None:
+            # TODO: do not save image after rendering
+            from PIL import Image, ImageDraw
+            img_pil = Image.open(path).convert("RGBA")
+            draw = ImageDraw.Draw(img_pil)
+            
+            from bpy_extras.image_utils import load_image  # type: ignore
+            bpy.context.scene.use_nodes = True
+            nodes = bpy.context.scene.node_tree.nodes
+            # nodes.clear()
+
+            img = load_image(path)
+            image_node = nodes.new(type='CompositorNodeImage')
+            image_node.image = img
+
+            for i, bone in enumerate(bones):
+                head, tail = bone[:3], bone[3:]
+                head_2d = bpy_extras.object_utils.world_to_camera_view(bpy.context.scene, camera, Vector(head))
+                tail_2d = bpy_extras.object_utils.world_to_camera_view(bpy.context.scene, camera, Vector(tail))
+
+                res_x, res_y = resolution
+                head_pix = (head_2d.x * res_x, (1 - head_2d.y) * res_y)
+                tail_pix = (tail_2d.x * res_x, (1 - tail_2d.y) * res_y)
+                draw.line([head_pix, tail_pix], fill=(255, 0, 0, 255), width=1)
+            img_pil.save(path)
+
+def _trans_to_m(v: ndarray):
+    m = np.eye(4)
+    m[0:3, 3] = v
+    return m
+
+def _scale_to_m(r: ndarray):
+    m = np.zeros((4, 4))
+    m[0, 0] = r
+    m[1, 1] = r
+    m[2, 2] = r
+    m[3, 3] = 1.
+    return m

UniRig/src/data/extract.py

@@ -0,0 +1,523 @@
+import bpy, os
+from collections import defaultdict
+from tqdm import tqdm
+import numpy as np
+from numpy import ndarray
+from typing import Dict, Tuple, List, Optional, Union
+import trimesh
+import fast_simplification
+from scipy.spatial import KDTree
+
+import argparse
+import yaml
+from box import Box
+import os
+
+from .log import new_entry, add_error, add_warning, new_log, end_log
+from .raw_data import RawData
+
+def load(filepath: str):
+    old_objs = set(bpy.context.scene.objects)
+    
+    if not os.path.exists(filepath):
+        raise ValueError(f'File {filepath} does not exist !')
+    
+    try:
+        if filepath.endswith(".vrm"):
+            # enable vrm addon and load vrm model
+            bpy.ops.preferences.addon_enable(module='vrm')
+            
+            bpy.ops.import_scene.vrm(
+                filepath=filepath,
+                use_addon_preferences=True,
+                extract_textures_into_folder=False,
+                make_new_texture_folder=False,
+                set_shading_type_to_material_on_import=False,
+                set_view_transform_to_standard_on_import=True,
+                set_armature_display_to_wire=True,
+                set_armature_display_to_show_in_front=True,
+                set_armature_bone_shape_to_default=True,
+                disable_bake=True, # customized option for better performance
+            )
+        elif filepath.endswith(".obj"):
+            bpy.ops.wm.obj_import(filepath=filepath)
+        elif filepath.endswith(".fbx") or filepath.endswith(".FBX"):
+            # end bone is removed using remove_dummy_bone
+            bpy.ops.import_scene.fbx(filepath=filepath, ignore_leaf_bones=False, use_image_search=False)
+        elif filepath.endswith(".glb") or filepath.endswith(".gltf"):
+            bpy.ops.import_scene.gltf(filepath=filepath, import_pack_images=False)
+        elif filepath.endswith(".dae"):
+            bpy.ops.wm.collada_import(filepath=filepath)
+        elif filepath.endswith(".blend"):
+            with bpy.data.libraries.load(filepath) as (data_from, data_to):
+                data_to.objects = data_from.objects
+            for obj in data_to.objects:
+                if obj is not None:
+                    bpy.context.collection.objects.link(obj)
+        else:
+            raise ValueError(f"not suported type {filepath}")
+    except:
+        raise ValueError(f"failed to load {filepath}")
+
+    armature = [x for x in set(bpy.context.scene.objects)-old_objs if x.type=="ARMATURE"]
+    if len(armature)==0:
+        return None
+    if len(armature)>1:
+        raise ValueError(f"multiple armatures found")
+    armature = armature[0]
+    
+    armature.select_set(True)
+    bpy.context.view_layer.objects.active = armature
+    bpy.ops.object.mode_set(mode='EDIT')
+    for bone in bpy.data.armatures[0].edit_bones:
+        bone.roll = 0. # change all roll to 0. to prevent weird behaviour
+
+    bpy.ops.object.mode_set(mode='OBJECT')
+    armature.select_set(False)
+    
+    bpy.ops.object.select_all(action='DESELECT')
+    return armature
+
+# remove all data in bpy
+def clean_bpy():
+    # First try to purge orphan data
+    try:
+        bpy.ops.outliner.orphans_purge(do_local_ids=True, do_linked_ids=True, do_recursive=True)
+    except Exception as e:
+        print(f"Warning: Could not purge orphans: {e}")
+        
+    # Then remove all data by type
+    data_types = [
+        bpy.data.actions,
+        bpy.data.armatures,
+        bpy.data.cameras,
+        bpy.data.collections,
+        bpy.data.curves,
+        bpy.data.images,
+        bpy.data.lights,
+        bpy.data.materials,
+        bpy.data.meshes,
+        bpy.data.objects,
+        bpy.data.textures,
+        bpy.data.worlds,
+        bpy.data.node_groups
+    ]
+    
+    for data_collection in data_types:
+        try:
+            for item in data_collection:
+                try:
+                    data_collection.remove(item)
+                except Exception as e:
+                    print(f"Warning: Could not remove {item.name} from {data_collection}: {e}")
+        except Exception as e:
+            print(f"Warning: Error processing {data_collection}: {e}")
+            
+    # Force garbage collection to free memory
+    import gc
+    gc.collect()
+
+def get_arranged_bones(armature):
+    matrix_world = armature.matrix_world
+    arranged_bones = []
+    root = armature.pose.bones[0]
+    while root.parent is not None:
+        root = root.parent
+    Q = [root]
+    rot = np.array(matrix_world)[:3, :3]
+    
+    # dfs and sort
+    while len(Q) != 0:
+        b = Q.pop(0)
+        arranged_bones.append(b)
+        children = []
+        for cb in b.children:
+            head = rot @ np.array(b.head)
+            children.append((cb, head[0], head[1], head[2]))
+        children = sorted(children, key=lambda x: (x[3], x[1], x[2]))
+        _c = [x[0] for x in children]
+        Q = _c + Q
+    return arranged_bones
+
+def process_mesh():
+    meshes = []
+    for v in bpy.data.objects:
+        if v.type == 'MESH':
+            meshes.append(v)
+    
+    _dict_mesh = {}
+    for obj in meshes:
+        m = np.array(obj.matrix_world)
+        matrix_world_rot = m[:3, :3]
+        matrix_world_bias = m[:3, 3]
+        rot = matrix_world_rot
+        total_vertices = len(obj.data.vertices)
+        vertex = np.zeros((4, total_vertices))
+        vertex_normal = np.zeros((total_vertices, 3))
+        obj_verts = obj.data.vertices
+        faces = []
+        normals = []
+        
+        for v in obj_verts:
+            vertex_normal[v.index] = rot @ np.array(v.normal) # be careful !
+            vv = rot @ v.co
+            vv = np.array(vv) + matrix_world_bias
+            vertex[0:3, v.index] = vv
+            vertex[3][v.index] = 1 # affine coordinate
+        
+        for polygon in obj.data.polygons:
+            edges = polygon.edge_keys
+            nodes = []
+            adj = {}
+            for edge in edges:
+                if adj.get(edge[0]) is None:
+                    adj[edge[0]] = []
+                adj[edge[0]].append(edge[1])
+                if adj.get(edge[1]) is None:
+                    adj[edge[1]] = []
+                adj[edge[1]].append(edge[0])
+                nodes.append(edge[0])
+                nodes.append(edge[1])
+            normal = polygon.normal
+            nodes = list(set(sorted(nodes)))
+            first = nodes[0]
+            loop = []
+            now = first
+            vis = {}
+            while True:
+                loop.append(now)
+                vis[now] = True
+                if vis.get(adj[now][0]) is None:
+                    now = adj[now][0]
+                elif vis.get(adj[now][1]) is None:
+                    now = adj[now][1]
+                else:
+                    break
+            for (second, third) in zip(loop[1:], loop[2:]):
+                faces.append((first + 1, second + 1, third + 1)) # the cursed +1
+                normals.append(rot @ normal) # and the cursed normal of BLENDER
+
+        correct_faces = []
+        for (i, face) in enumerate(faces):
+            normal = normals[i]
+            v0 = face[0] - 1
+            v1 = face[1] - 1
+            v2 = face[2] - 1
+            v = np.cross(
+                vertex[:3, v1] - vertex[:3, v0],
+                vertex[:3, v2] - vertex[:3, v0],
+            )
+            if (v*normal).sum() > 0:
+                correct_faces.append(face)
+            else:
+                correct_faces.append((face[0], face[2], face[1]))
+        if len(correct_faces) > 0:
+            _dict_mesh[obj.name] = {
+                'vertex': vertex,
+                'face': correct_faces,
+            }
+    
+    vertex = np.concatenate([_dict_mesh[name]['vertex'] for name in _dict_mesh], axis=1)[:3, :].transpose()
+    
+    total_faces = 0
+    now_bias = 0
+    for name in _dict_mesh:
+        total_faces += len(_dict_mesh[name]['face'])
+    faces = np.zeros((total_faces, 3), dtype=np.int64)
+    tot = 0
+    for name in _dict_mesh:
+        f = np.array(_dict_mesh[name]['face'], dtype=np.int64)
+        faces[tot:tot+f.shape[0]] = f + now_bias
+        now_bias += _dict_mesh[name]['vertex'].shape[1]
+        tot += f.shape[0]
+    
+    return vertex, faces
+
+def process_armature(
+    armature,
+    arranged_bones,
+) -> Tuple[np.ndarray, np.ndarray]:
+    matrix_world = armature.matrix_world
+    index = {}
+
+    for (id, pbone) in enumerate(arranged_bones):
+        index[pbone.name] = id
+    
+    root = armature.pose.bones[0]
+    while root.parent is not None:
+        root = root.parent
+    m = np.array(matrix_world.to_4x4())
+    scale_inv = np.linalg.inv(np.diag(matrix_world.to_scale()))
+    rot = m[:3, :3]
+    bias = m[:3, 3]
+    
+    s = []
+    bpy.ops.object.editmode_toggle()
+    edit_bones = armature.data.edit_bones
+    
+    J = len(arranged_bones)
+    joints = np.zeros((J, 3), dtype=np.float32)
+    tails = np.zeros((J, 3), dtype=np.float32)
+    parents = []
+    name_to_id = {}
+    names = []
+    matrix_local_stack = np.zeros((J, 4, 4), dtype=np.float32)
+    for (id, pbone) in enumerate(arranged_bones):
+        name = pbone.name
+        names.append(name)
+        matrix_local = np.array(pbone.bone.matrix_local)
+        use_inherit_rotation = pbone.bone.use_inherit_rotation
+        if use_inherit_rotation == False:
+            add_warning(f"use_inherit_rotation of bone {name} is False !")
+        head = rot @ matrix_local[0:3, 3] + bias
+        s.append(head)
+        edit_bone = edit_bones.get(name)
+        tail = rot @ np.array(edit_bone.tail) + bias
+        
+        name_to_id[name] = id
+        joints[id] = head
+        tails[id] = tail
+        parents.append(None if pbone.parent not in arranged_bones else name_to_id[pbone.parent.name])
+        # remove scale part
+        matrix_local[:, 3:4] = m @ matrix_local[:, 3:4]
+        matrix_local[:3, :3] = scale_inv @ matrix_local[:3, :3]
+        matrix_local_stack[id] = matrix_local
+    bpy.ops.object.editmode_toggle()
+    
+    return joints, tails, parents, names, matrix_local_stack
+
+def save_raw_data(
+    path: str,
+    vertices: ndarray,
+    faces: ndarray,
+    joints: Union[ndarray, None],
+    tails: Union[ndarray, None],
+    parents: Union[List[Union[int, None]], None],
+    names: Union[List[str], None],
+    matrix_local: Union[ndarray, None],
+    target_count: int,
+):
+    mesh = trimesh.Trimesh(vertices=vertices, faces=faces)
+    vertices = np.array(mesh.vertices, dtype=np.float32)
+    faces = np.array(mesh.faces, dtype=np.int64)
+    if faces.shape[0] > target_count:
+        vertices, faces = fast_simplification.simplify(vertices, faces, target_count=target_count)
+    mesh = trimesh.Trimesh(vertices=vertices, faces=faces)
+    
+    new_vertices = np.array(mesh.vertices, dtype=np.float32)
+    new_vertex_normals = np.array(mesh.vertex_normals, dtype=np.float32)
+    new_faces = np.array(mesh.faces, dtype=np.int64)
+    new_face_normals = np.array(mesh.face_normals, dtype=np.float32)
+    if joints is not None:
+        new_joints = np.array(joints, dtype=np.float32)
+    else:
+        new_joints = None
+    raw_data = RawData(
+        vertices=new_vertices,
+        vertex_normals=new_vertex_normals,
+        faces=new_faces,
+        face_normals=new_face_normals,
+        joints=new_joints,
+        tails=tails,
+        skin=None,
+        no_skin=None,
+        parents=parents,
+        names=names,
+        matrix_local=matrix_local,
+    )
+    raw_data.check()
+    raw_data.save(path=path)
+
+def extract_builtin(
+    output_folder: str,
+    target_count: int,
+    num_runs: int,
+    id: int,
+    time: str,
+    files: List[Union[str, str]],
+):
+    log_path = "./logs"
+    log_path = os.path.join(log_path, time)
+
+    num_files = len(files)
+    gap = num_files // num_runs
+    start = gap * id
+    end = gap * (id + 1)
+    if id+1==num_runs:
+        end = num_files
+    
+    files = sorted(files)
+    if end!=-1:
+        files = files[:end]
+    new_log(log_path, f"extract_builtin_{start}_{end}")
+    tot = 0
+    for file in tqdm(files[start:]):
+        input_file = file[0]
+        output_dir = file[1]
+        clean_bpy()
+        new_entry(input_file)
+        try:
+            print(f"Now processing {input_file}...")
+            
+            armature = load(input_file)
+            
+            print('save to:', output_dir)
+            os.makedirs(output_dir, exist_ok=True)
+            
+            vertices, faces = process_mesh()
+            if armature is not None:
+                arranged_bones = get_arranged_bones(armature)
+                joints, tails, parents, names, matrix_local = process_armature(armature, arranged_bones)
+                
+            else:
+                joints = None
+                tails = None
+                parents = None
+                names = None
+                matrix_local = None
+            
+            save_file = os.path.join(output_dir, 'raw_data.npz')
+            save_raw_data(
+                path=save_file,
+                vertices=vertices,
+                faces=faces-1,
+                joints=joints,
+                tails=tails,
+                parents=parents,
+                names=names,
+                matrix_local=matrix_local,
+                target_count=target_count,
+            )
+            
+            tot += 1
+
+        except ValueError as e:
+            add_error(str(e))
+            print(f"ValueError: {str(e)}")
+        except RuntimeError as e:
+            add_error(str(e))
+            print(f"RuntimeError: {str(e)}")
+        except TimeoutError as e:
+            add_error("time out")
+            print("TimeoutError: Processing timed out")
+        except Exception as e:
+            add_error(f"Unexpected error: {str(e)}")
+            print(f"Unexpected error: {str(e)}")
+    end_log()
+    print(f"{tot} models processed")
+
+def str2bool(v):
+    if isinstance(v, bool):
+        return v
+    if v.lower() in ('yes', 'true', 't', 'y', '1'):
+        return True
+    elif v.lower() in ('no', 'false', 'f', 'n', '0'):
+        return False
+    else:
+        raise argparse.ArgumentTypeError('Boolean value expected.')
+
+def nullable_string(val):
+    if not val:
+        return None
+    return val
+
+def get_files(
+    data_name: str,
+    input_dataset_dir: str,
+    output_dataset_dir: str,
+    inputs: Union[str, None]=None,
+    require_suffix: List[str]=['obj','fbx','FBX','dae','glb','gltf','vrm'],
+    force_override: bool=False,
+    warning: bool=True,
+) -> List[Tuple[str, str]]:
+    
+    files = [] # (input_file, output_dir)
+    if inputs is not None: # specified input file(s)
+        vis = {}
+        inputs = inputs.split(',')
+        for file in inputs:
+            file_name = file.removeprefix("./")
+            # remove suffix
+            file_name = '.'.join(file_name.split('.')[:-1])
+            output_dir = os.path.join(output_dataset_dir, file_name)
+            raw_data_npz = os.path.join(output_dir, data_name)
+            if not force_override and os.path.exists(raw_data_npz):
+                continue
+            if warning and output_dir in vis:
+                print(f"\033[33mWARNING: duplicate output directory: {output_dir}, you need to rename prefix of files to avoid ambiguity\033[0m")
+            vis[output_dir] = True
+            files.append((file, output_dir))
+    else:
+        vis = {}
+        for root, dirs, f in os.walk(input_dataset_dir):
+            for file in f:
+                if file.split('.')[-1] in require_suffix:
+                    file_name = file.removeprefix("./")
+                    # remove suffix
+                    file_name = '.'.join(file_name.split('.')[:-1])
+                    
+                    output_dir = os.path.join(output_dataset_dir, os.path.relpath(root, input_dataset_dir), file_name)
+                    raw_data_npz = os.path.join(output_dir, data_name)
+                    
+                    # Check if all required files exist
+                    if not force_override and os.path.exists(raw_data_npz):
+                        continue
+                    if warning and output_dir in vis:
+                        print(f"\033[33mWARNING: duplicate output directory: {output_dir}, you need to rename prefix of files to avoid ambiguity\033[0m")
+                    vis[output_dir] = True
+                    files.append((os.path.join(root, file), output_dir))
+
+    return files
+
+def parse():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--config', type=str, required=True)
+    parser.add_argument('--require_suffix', type=str, required=True)
+    parser.add_argument('--faces_target_count', type=int, required=True)
+    parser.add_argument('--num_runs', type=int, required=True)
+    parser.add_argument('--force_override', type=str2bool, required=True)
+    parser.add_argument('--id', type=int, required=True)
+    parser.add_argument('--time', type=str, required=True)
+
+    parser.add_argument('--input', type=nullable_string, required=False, default=None)
+    parser.add_argument('--input_dir', type=nullable_string, required=False, default=None)
+    parser.add_argument('--output_dir', type=nullable_string, required=False, default=None)
+    return parser.parse_args()
+
+if __name__ == "__main__":
+    args = parse()
+    
+    config = Box(yaml.safe_load(open(args.config, "r")))
+    
+    num_runs        = args.num_runs
+    id              = args.id
+    timestamp       = args.time
+    require_suffix  = args.require_suffix.split(',')
+    force_override  = args.force_override
+    target_count    = args.faces_target_count
+    
+    if args.input_dir:
+        config.input_dataset_dir = args.input_dir
+    if args.output_dir:
+        config.output_dataset_dir = args.output_dir
+    
+    assert config.input_dataset_dir is not None or args.input is None, 'you cannot specify both input and input_dir'
+
+    files = get_files(
+        data_name='raw_data.npz',
+        inputs=args.input,
+        input_dataset_dir=config.input_dataset_dir,
+        output_dataset_dir=config.output_dataset_dir,
+        require_suffix=require_suffix,
+        force_override=force_override,
+        warning=True,
+    )
+
+    extract_builtin(
+        output_folder=config.output_dataset_dir,
+        target_count=target_count,
+        num_runs=num_runs,
+        id=id,
+        time=timestamp,
+        files=files,
+    )

UniRig/src/data/log.py

@@ -0,0 +1,50 @@
+import os
+from typing import List
+
+login_time = ''
+log_filepath = ''
+
+class Entry:
+    def __init__(self, entry_name):
+        self.entry = entry_name
+        self.error = None
+        self.warning = []
+    
+    def have_error(self):
+        return self.error != None
+    
+    def have_warning(self):
+        return len(self.warning) != 0
+
+logs: List[Entry] = []
+
+def new_log(path, log_name):
+    global login_time, log_filepath
+    log_filepath = os.path.join(path, f"{log_name}.txt")
+    os.makedirs(path, exist_ok=True)
+    with open(log_filepath, 'a') as file:
+        file.write(f"Log: {log_name}\n")
+
+def end_log():
+    global log_filepath
+    with open(log_filepath, 'a') as file:
+        file.write(f"End of file\n")
+
+def new_entry(entry_name):
+    global log_filepath
+    print(f"\033[32mNow processing {entry_name}...\033[0m")
+    logs.append(Entry(entry_name))
+
+def add_error(error):
+    global log_filepath
+    print(f"\033[31mError found when processing {logs[-1].entry}: {error}\033[0m")
+    logs[-1].error = error
+    with open(log_filepath, 'a') as file:
+        file.write(f"Entry: {logs[-1].entry}, Error: {error}\n")
+
+def add_warning(warning):
+    global log_filepath
+    print(f"\033[33mWarning found when processing {logs[-1].entry}: {warning}\033[0m")
+    logs[-1].warning.append(warning)
+    with open(log_filepath, 'a') as file:
+        file.write(f"Entry: {logs[-1].entry}, Warning: {warning}\n")

UniRig/src/data/order.py

@@ -0,0 +1,112 @@
+from typing import Dict, List, Tuple, Union
+from collections import defaultdict
+from dataclasses import dataclass
+import yaml
+from box import Box
+
+from .spec import ConfigSpec
+
+@dataclass
+class OrderConfig(ConfigSpec):
+    '''
+    Config to handle bones re-ordering.
+    '''
+    
+    # {skeleton_name: path}
+    skeleton_path: Dict[str, str]
+    
+    # {cls: {part_name: [bone_name_1, bone_name_2, ...]}}
+    parts: Dict[str, Dict[str, List[str]]]
+    
+    # {cls: parts of bones to be arranged in [part_name_1, part_name_2, ...]}
+    parts_order: Dict[str, List[str]]
+    
+    @classmethod
+    def parse(cls, config):
+        cls.check_keys(config)
+        skeleton_path = config.skeleton_path
+        parts = {}
+        parts_order = {}
+        for (cls, path) in skeleton_path.items():
+            assert cls not in parts, 'cls conflicts'
+            d = Box(yaml.safe_load(open(path, 'r')))
+            parts[cls] = d.parts
+            parts_order[cls] = d.parts_order
+        return OrderConfig(
+            skeleton_path=skeleton_path,
+            parts=parts,
+            parts_order=parts_order,
+        )
+
+class Order():
+    
+    # {part_name: [bone_name_1, bone_name_2, ...]}
+    parts: Dict[str, Dict[str, List[str]]]
+    
+    # parts of bones to be arranged in [part_name_1, part_name_2, ...]
+    parts_order: Dict[str, List[str]]
+    
+    def __init__(self, config: OrderConfig):
+        self.parts          = config.parts
+        self.parts_order    = config.parts_order
+    
+    def part_exists(self, cls: str, part: str, names: List[str]) -> bool:
+        '''
+        Check if part exists.
+        '''
+        if part not in self.parts[cls]:
+            return False
+        for name in self.parts[cls][part]:
+            if name not in names:
+                return False
+        return True
+    
+    def make_names(self, cls: Union[str, None], parts: List[Union[str, None]], num_bones: int) -> List[str]:
+        '''
+        Get names for specified cls.
+        '''
+        names = []
+        for part in parts:
+            if part is None: # spring
+                continue
+            if cls in self.parts and part in self.parts[cls]:
+                names.extend(self.parts[cls][part])
+        assert len(names) <= num_bones, "number of bones in required skeleton is more than existing bones"
+        for i in range(len(names), num_bones):
+            names.append(f"bone_{i}")
+        return names
+    
+    def arrange_names(self, cls: str, names: List[str], parents: List[Union[int, None]]) -> Tuple[List[str], Dict[int, Union[str]]]:
+        '''
+        Arrange names according to required parts order.
+        '''
+        if cls not in self.parts_order:
+            return names, {0: None} # add a spring token
+        vis = defaultdict(bool)
+        name_to_id = {name: i for (i, name) in enumerate(names)}
+        new_names = []
+        parts_bias = {}
+        for part in self.parts_order[cls]:
+            if self.part_exists(cls=cls, part=part, names=names):
+                for name in self.parts[cls][part]:
+                    vis[name] = True
+                flag = False
+                for name in self.parts[cls][part]:
+                    pid = parents[name_to_id[name]]
+                    if pid is None:
+                        continue
+                    if not vis[names[pid]]:
+                        flag = True
+                        break
+                if flag: # incorrect parts order and should immediately add a spring token
+                    break
+                parts_bias[len(new_names)] = part
+                new_names.extend(self.parts[cls][part])
+        parts_bias[len(new_names)] = None # add a spring token
+        for name in names:
+            if name not in new_names:
+                new_names.append(name)
+        return new_names, parts_bias
+
+def get_order(config: OrderConfig) -> Order:
+    return Order(config=config)

UniRig/src/data/raw_data.py

@@ -0,0 +1,307 @@
+from dataclasses import dataclass
+import numpy as np
+from numpy import ndarray
+
+import os
+from typing import Union, List, Tuple
+
+from .exporter import Exporter
+
+from ..tokenizer.spec import DetokenzeOutput
+from .order import Order
+
+@dataclass(frozen=True)
+class RawData(Exporter):
+    '''
+    Dataclass to handle data from processed model files.
+    '''
+    
+    # vertices of the mesh, shape (N, 3), float32
+    vertices: Union[ndarray, None]
+    
+    # normals of vertices, shape (N, 3), float32
+    vertex_normals: Union[ndarray, None]
+    
+    # faces of mesh, shape (F, 3), face id starts from 0 to F-1, int64
+    faces: Union[ndarray, None]
+    
+    # face normal of mesh, shape (F, 3), float32
+    face_normals: Union[ndarray, None]
+    
+    # joints of bones, shape (J, 3), float32
+    joints: Union[ndarray, None]
+    
+    # tails of joints, shape (J, 3), float32
+    tails: Union[ndarray, None]
+    
+    # skinning of joints, shape (N, J), float32
+    skin: Union[ndarray, None]
+    
+    # whether the joint has skin, bool
+    no_skin: Union[ndarray, None]
+    
+    # parents of joints, None represents no parent(a root joint)
+    # make sure parent[k] < k
+    parents: Union[List[Union[int, None]], None]
+    
+    # names of joints
+    names: Union[List[str], None]
+    
+    # local coordinate
+    matrix_local: Union[ndarray, None]
+    
+    # path to data
+    path: Union[str, None]=None
+    
+    # data cls
+    cls: Union[str, None]=None
+    
+    @staticmethod
+    def load(path: str) -> 'RawData':
+        data = np.load(path, allow_pickle=True)
+        d = {name: data[name][()] for name in data}
+        d['path'] = path
+        return RawData(**d)
+    
+    def save(self, path: str):
+        os.makedirs(os.path.dirname(path), exist_ok=True)
+        np.savez(file=path, **self.__dict__)
+    
+    @property
+    def N(self):
+        '''
+        number of vertices
+        '''
+        return self.vertices.shape[0]
+    
+    @property
+    def F(self):
+        '''
+        number of faces
+        '''
+        return self.faces.shape[0]
+    
+    @property
+    def J(self):
+        '''
+        number of joints
+        '''
+        return self.joints.shape[0]
+    
+    def check(self):
+        if self.names is not None and self.joints is not None:
+            assert len(self.names) == self.J
+        if self.names is not None and self.parents is not None:
+            assert len(self.names) == len(self.parents)
+        if self.parents is not None:
+            for (i, pid) in enumerate(self.parents):
+                if i==0:
+                    assert pid is None
+                else:
+                    assert pid is not None
+                    assert pid < i
+    
+    def export_pc(self, path: str, with_normal: bool=True, normal_size=0.01):
+        '''
+        export point cloud
+        '''
+        if with_normal:
+            self._export_pc(vertices=self.vertices, path=path, vertex_normals=self.vertex_normals, normal_size=normal_size)
+        else:
+            self._export_pc(vertices=self.vertices, path=path, vertex_normals=None, normal_size=normal_size)
+    
+    def export_mesh(self, path: str):
+        '''
+        export mesh
+        '''
+        self._export_mesh(vertices=self.vertices, faces=self.faces, path=path)
+    
+    def export_skeleton(self, path: str):
+        '''
+        export spring
+        '''
+        self._export_skeleton(joints=self.joints, parents=self.parents, path=path)
+    
+    def export_skeleton_sequence(self, path: str):
+        '''
+        export spring
+        '''
+        self._export_skeleton_sequence(joints=self.joints, parents=self.parents, path=path)
+    
+    def export_fbx(
+        self,
+        path: str,
+        extrude_size: float=0.03,
+        group_per_vertex: int=-1,
+        add_root: bool=False,
+        do_not_normalize: bool=False,
+        use_extrude_bone: bool=True,
+        use_connect_unique_child: bool=True,
+        extrude_from_parent: bool=True,
+        use_tail: bool=False,
+        custom_vertex_group: Union[ndarray, None]=None,
+    ):
+        '''
+        export the whole model with skining
+        '''
+        self._export_fbx(
+            path=path,
+            vertices=self.vertices,
+            joints=self.joints,
+            skin=self.skin if custom_vertex_group is None else custom_vertex_group,
+            parents=self.parents,
+            names=self.names,
+            faces=self.faces,
+            extrude_size=extrude_size,
+            group_per_vertex=group_per_vertex,
+            add_root=add_root,
+            do_not_normalize=do_not_normalize,
+            use_extrude_bone=use_extrude_bone,
+            use_connect_unique_child=use_connect_unique_child,
+            extrude_from_parent=extrude_from_parent,
+            tails=self.tails if use_tail else None,
+        )
+    
+    def export_render(self, path: str, resolution: Tuple[int, int]=[256, 256]):
+        self._export_render(
+            path=path,
+            vertices=self.vertices,
+            faces=self.faces,
+            bones=np.concatenate([self.joints, self.tails], axis=-1),
+            resolution=resolution,
+        )
+
+@dataclass(frozen=True)
+class RawSkeleton(Exporter):
+    '''
+    Dataclass to handle skeleton from AR.
+    '''
+    # joints of bones, shape (J, 3), float32
+    joints: Union[ndarray, None]
+    
+    # tails of joints, shape (J, 3), float32
+    tails: Union[ndarray, None]
+    
+    # whether the joint has skin, bool
+    no_skin: Union[ndarray, None]
+    
+    # parents of joints, None represents no parent(a root joint)
+    # make sure parent[k] < k
+    parents: Union[List[Union[int, None]], None]
+    
+    # names of joints
+    names: Union[List[str], None]
+    
+    @staticmethod
+    def load(path: str) -> 'RawSkeleton':
+        data = np.load(path, allow_pickle=True)
+        return RawSkeleton(**{name: data[name][()] for name in data})
+    
+    def save(self, path: str):
+        os.makedirs(os.path.dirname(path), exist_ok=True)
+        np.savez(file=path, **self.__dict__)
+    
+    @staticmethod
+    def from_detokenize_output(res: DetokenzeOutput, order: Union[Order, None]) -> 'RawSkeleton':
+        J = len(res.bones)
+        names = order.make_names(cls=res.cls, parts=res.parts, num_bones=J)
+        joints = res.joints
+        p_joints = res.p_joints
+        parents = []
+        for (i, joint) in enumerate(joints):
+            if i == 0:
+                parents.append(None)
+                continue
+            p_joint = p_joints[i]
+            dis = 999999
+            pid = None
+            for j in reversed(range(i)):
+                n_dis = ((joints[j] - p_joint)**2).sum()
+                if n_dis < dis:
+                    pid = j
+                    dis = n_dis
+            parents.append(pid)
+        return RawSkeleton(
+            joints=joints,
+            tails=res.tails,
+            no_skin=res.no_skin,
+            parents=parents,
+            names=names,
+        )
+        
+    def export_skeleton(self, path: str):
+        '''
+        export spring
+        '''
+        self._export_skeleton(joints=self.joints, parents=self.parents, path=path)
+    
+    def export_skeleton_sequence(self, path: str):
+        '''
+        export spring
+        '''
+        self._export_skeleton_sequence(joints=self.joints, parents=self.parents, path=path)
+    
+    def export_fbx(
+        self,
+        path: str,
+        extrude_size: float=0.03,
+        group_per_vertex: int=-1,
+        add_root: bool=False,
+        do_not_normalize: bool=False,
+        use_extrude_bone: bool=True,
+        use_connect_unique_child: bool=True,
+        extrude_from_parent: bool=True,
+        use_tail: bool=False,
+    ):
+        '''
+        export the whole model with skining
+        '''
+        self._export_fbx(
+            path=path,
+            vertices=None,
+            joints=self.joints,
+            skin=None,
+            parents=self.parents,
+            names=self.names,
+            faces=None,
+            extrude_size=extrude_size,
+            group_per_vertex=group_per_vertex,
+            add_root=add_root,
+            do_not_normalize=do_not_normalize,
+            use_extrude_bone=use_extrude_bone,
+            use_connect_unique_child=use_connect_unique_child,
+            extrude_from_parent=extrude_from_parent,
+            tails=self.tails if use_tail else None,
+        )
+    
+    def export_render(self, path: str, resolution: Tuple[int, int]=[256, 256]):
+        self._export_render(
+            path=path,
+            vertices=None,
+            faces=None,
+            bones=np.concatenate([self.joints, self.tails], axis=-1),
+            resolution=resolution,
+        )
+
+@dataclass
+class RawSkin(Exporter):
+    '''
+    Dataclass to handle skeleton from AR.
+    '''
+    # skin, shape (J, N)
+    skin: ndarray
+    
+    # always sampled, shape (N, 3)
+    vertices: Union[ndarray, None]=None
+    
+    # for future use, shape (J, 3)
+    joints: Union[ndarray, None]=None
+    
+    @staticmethod
+    def load(path: str) -> 'RawSkin':
+        data = np.load(path, allow_pickle=True)
+        return RawSkin(**{name: data[name][()] for name in data})
+    
+    def save(self, path: str):
+        os.makedirs(os.path.dirname(path), exist_ok=True)
+        np.savez(file=path, **self.__dict__)

UniRig/src/data/sampler.py

@@ -0,0 +1,210 @@
+from typing import List
+from heapq import heappush, heappop, heapify
+from dataclasses import dataclass
+from abc import ABC, abstractmethod
+import numpy as np
+from numpy import ndarray
+
+from typing import Dict, Tuple
+
+from .asset import Asset
+from .spec import ConfigSpec
+
+@dataclass
+class SamplerConfig(ConfigSpec):
+    '''
+    Config to handle bones re-ordering.
+    '''
+    # which sampler to use
+    method: str
+    
+    # how many samples in total
+    num_samples: int
+    
+    # how many vertex samples
+    vertex_samples: int
+    
+    # kwargs
+    kwargs: Dict[str, Dict]
+    
+    @classmethod
+    def parse(cls, config) -> 'SamplerConfig':
+        cls.check_keys(config)
+        return SamplerConfig(
+            method=config.method,
+            num_samples=config.get('num_samples', 0),
+            vertex_samples=config.get('vertex_samples', 0),
+            kwargs=config.get('kwargs', {}),
+        )
+
+@dataclass
+class SamplerResult():
+    # sampled vertices
+    vertices: ndarray
+    
+    # sampled normals
+    normals: ndarray
+    
+    # sampled vertex groups
+    vertex_groups: Dict[str, ndarray]
+
+class Sampler(ABC):
+    '''
+    Abstract class for samplers.
+    '''
+    
+    def _sample_barycentric(
+        self,
+        vertex_group: ndarray,
+        faces: ndarray,
+        face_index: ndarray,
+        random_lengths: ndarray,
+    ):
+        v_origins = vertex_group[faces[face_index, 0]]
+        v_vectors = vertex_group[faces[face_index, 1:]]
+        v_vectors -= v_origins[:, np.newaxis, :]
+        
+        sample_vector = (v_vectors * random_lengths).sum(axis=1)
+        v_samples = sample_vector + v_origins
+        return v_samples
+    
+    @abstractmethod
+    def __init__(self, config: SamplerConfig):
+        pass
+    
+    @abstractmethod
+    def sample(
+        self,
+        asset: Asset,
+    ) -> SamplerResult:
+        '''
+        Return sampled vertices, sampled normals and vertex groups.
+        '''
+        pass
+
+class SamplerOrigin(Sampler):
+    def __init__(self, config: SamplerConfig):
+        super().__init__(config)
+        self.num_samples    = config.num_samples
+        self.vertex_samples = config.vertex_samples
+        
+    def sample(
+        self,
+        asset: Asset,
+    ) -> SamplerResult:
+        perm = np.random.permutation(asset.vertices.shape[0])
+        if asset.vertices.shape[0] < self.num_samples:
+            m = self.num_samples - asset.vertices.shape[0]
+            perm = np.concatenate([perm, np.random.randint(0, asset.vertices.shape[0], (m,))])
+        perm = perm[:self.num_samples]
+        n_v = asset.vertices[perm]
+        n_n = asset.vertex_normals[perm]
+        n_vg = {name: v[perm] for name, v in asset.vertex_groups.items()}
+        return SamplerResult(
+            vertices=n_v,
+            normals=n_n,
+            vertex_groups=n_vg,
+        )
+
+class SamplerMix(Sampler):
+    def __init__(self, config: SamplerConfig):
+        super().__init__(config)
+        self.num_samples    = config.num_samples
+        self.vertex_samples = config.vertex_samples
+        assert self.num_samples >= self.vertex_samples, 'num_samples should >= vertex_samples'
+    
+    @property
+    def mesh_preserve(self):
+        return self.num_samples==-1
+    
+    def sample(
+        self,
+        asset: Asset,
+    ) -> SamplerResult:
+        # 1. sample vertices
+        num_samples = self.num_samples
+        perm = np.random.permutation(asset.vertices.shape[0])
+        vertex_samples = min(self.vertex_samples, asset.vertices.shape[0])
+        num_samples -= vertex_samples
+        perm = perm[:vertex_samples]
+        n_vertex = asset.vertices[perm]
+        n_normal = asset.vertex_normals[perm]
+        n_v = {name: v[perm] for name, v in asset.vertex_groups.items()}
+        
+        # 2. sample surface
+        perm = np.random.permutation(num_samples)
+        vertex_samples, face_index, random_lengths = sample_surface(
+            num_samples=num_samples,
+            vertices=asset.vertices,
+            faces=asset.faces,
+            return_weight=True,
+        )
+        vertex_samples = np.concatenate([n_vertex, vertex_samples], axis=0)
+        normal_samples = np.concatenate([n_normal, asset.face_normals[face_index]], axis=0)
+        vertex_group_samples = {}
+        for n, v in asset.vertex_groups.items():
+            g = self._sample_barycentric(
+                vertex_group=v,
+                faces=asset.faces,
+                face_index=face_index,
+                random_lengths=random_lengths,
+            )
+            vertex_group_samples[n] = np.concatenate([n_v[n], g], axis=0)
+        return SamplerResult(
+            vertices=vertex_samples,
+            normals=normal_samples,
+            vertex_groups=vertex_group_samples,
+        )
+
+def sample_surface(
+    num_samples: int,
+    vertices: ndarray,
+    faces: ndarray,
+    return_weight: bool=False,
+):
+    '''
+    Randomly pick samples according to face area.
+    
+    See sample_surface: https://github.com/mikedh/trimesh/blob/main/trimesh/sample.py
+    '''
+    # get face area
+    offset_0 = vertices[faces[:, 1]] - vertices[faces[:, 0]]
+    offset_1 = vertices[faces[:, 2]] - vertices[faces[:, 0]]
+    face_weight = np.cross(offset_0, offset_1, axis=-1)
+    face_weight = (face_weight * face_weight).sum(axis=1)
+    
+    weight_cum = np.cumsum(face_weight, axis=0)
+    face_pick = np.random.rand(num_samples) * weight_cum[-1]
+    face_index = np.searchsorted(weight_cum, face_pick)
+    
+    # pull triangles into the form of an origin + 2 vectors
+    tri_origins = vertices[faces[:, 0]]
+    tri_vectors = vertices[faces[:, 1:]]
+    tri_vectors -= np.tile(tri_origins, (1, 2)).reshape((-1, 2, 3))
+
+    # pull the vectors for the faces we are going to sample from
+    tri_origins = tri_origins[face_index]
+    tri_vectors = tri_vectors[face_index]
+    
+    # randomly generate two 0-1 scalar components to multiply edge vectors b
+    random_lengths = np.random.rand(len(tri_vectors), 2, 1)
+    
+    random_test = random_lengths.sum(axis=1).reshape(-1) > 1.0
+    random_lengths[random_test] -= 1.0
+    random_lengths = np.abs(random_lengths)
+    
+    sample_vector = (tri_vectors * random_lengths).sum(axis=1)
+    vertex_samples = sample_vector + tri_origins
+    if not return_weight:
+        return vertex_samples
+    return vertex_samples, face_index, random_lengths
+
+def get_sampler(config: SamplerConfig) -> Sampler:
+    method = config.method
+    if method=='origin':
+        sampler = SamplerOrigin(config)
+    elif method=='mix':
+        sampler = SamplerMix(config)
+    else:
+        raise ValueError(f"sampler method {method} not supported")
+    return sampler

UniRig/src/data/spec.py

@@ -0,0 +1,15 @@
+from abc import ABC, abstractmethod
+from dataclasses import fields
+
+class ConfigSpec(ABC):
+    @classmethod
+    def check_keys(cls, config):
+        expect = [field.name for field in fields(cls)]
+        for key in config.keys():
+            if key not in expect:
+                raise ValueError(f"expect names {expect} in {cls.__name__}, found {key}")
+    
+    @classmethod
+    @abstractmethod
+    def parse(cls, config):
+        pass

UniRig/src/data/tail.py

@@ -0,0 +1,50 @@
+from collections import defaultdict
+from dataclasses import dataclass
+import numpy as np
+from numpy import ndarray
+
+from typing import Tuple
+
+from .asset import Asset
+from .spec import ConfigSpec
+
+@dataclass
+class TailConfig(ConfigSpec):
+    '''
+    Config to handle tails.
+    '''
+    
+    # copy joints to tails
+    copy_joint_to_tail: bool
+    
+    # if the joint has only one son, then connect tail to son's joint
+    connect_tail_to_unique_son: bool
+    
+    @classmethod
+    def parse(cls, config) -> 'TailConfig':
+        cls.check_keys(config)
+        return TailConfig(
+            copy_joint_to_tail=config.copy_joint_to_tail,
+            connect_tail_to_unique_son=config.connect_tail_to_unique_son,
+        )
+
+class Tail():
+    
+    def __init__(self, config: TailConfig):
+        self.config = config
+    
+    def process_tail(self, asset: Asset):
+        if self.config.copy_joint_to_tail:
+            assert asset.tails is None, 'copying joints to existing tails is not permitted, please change copy_joint_to_tail to False in transform config'
+            asset.tails = asset.joints.copy()
+        if self.config.connect_tail_to_unique_son and asset.tails is not None:
+            children = defaultdict(list)
+            for (id, p) in enumerate(asset.parents):
+                if p is not None:
+                    children[p].append(id)
+            for i in range(asset.J):
+                if len(children[i]) == 1:
+                    asset.tails[i] = asset.joints[children[i][0]]
+
+def get_tail(config: TailConfig) -> Tail:
+    return Tail(config=config)

UniRig/src/data/transform.py

@@ -0,0 +1,107 @@
+from collections import defaultdict
+from dataclasses import dataclass
+from typing import Union, List, Tuple
+from copy import deepcopy
+
+from .asset import Asset
+from .augment import AugmentConfig, Augment, get_augments
+from .order import OrderConfig, Order, get_order
+from .sampler import SamplerConfig, get_sampler
+from .vertex_group import VertexGroupConfig, get_vertex_groups
+from .tail import TailConfig, get_tail
+from .spec import ConfigSpec
+
+@dataclass
+class TransformConfig(ConfigSpec):
+    
+    tail_config: Union[TailConfig, None]=None,
+    
+    order_config: Union[OrderConfig, None]=None,
+    
+    vertex_group_config: Union[VertexGroupConfig, None]=None,
+    
+    augment_config: Union[AugmentConfig, None]=None,
+    
+    sampler_config: Union[SamplerConfig, None]=None,
+    
+    @classmethod
+    def parse(cls, config) -> 'TransformConfig':
+        cls.check_keys(config)
+        tail_config = config.get('tail_config', None)
+        order_config = config.get('order_config', None)
+        vertex_group_config = config.get('vertex_group_config', None)
+        augment_config = config.get('augment_config', None)
+        sampler_config = config.get('sampler_config', None)
+        
+        if tail_config is not None:
+            tail_config = TailConfig.parse(config=tail_config)
+        if order_config is not None:
+            order_config = OrderConfig.parse(config=order_config)
+        if vertex_group_config is not None:
+            vertex_group_config = VertexGroupConfig.parse(config=vertex_group_config)
+        if augment_config is not None:
+            augment_config = AugmentConfig.parse(config=augment_config)
+        if sampler_config is not None:
+            sampler_config = SamplerConfig.parse(config=sampler_config)
+        
+        return TransformConfig(
+            tail_config=tail_config,
+            order_config=order_config,
+            vertex_group_config=vertex_group_config,
+            augment_config=augment_config,
+            sampler_config=sampler_config,
+        )
+
+def transform_asset(
+    asset: Asset,
+    transform_config: TransformConfig,
+) -> Tuple[List[Augment], List[Augment]]:
+    assert isinstance(transform_config, TransformConfig), f"found {type(transform_config)}"
+    # 1. try processing tails
+    # TODO: use a better method
+    if transform_config.tail_config is not None:
+        tail = get_tail(config=transform_config.tail_config)
+        tail.process_tail(asset=asset)
+    
+    # 2. arrange bones
+    if transform_config.order_config is not None:
+        order = get_order(config=transform_config.order_config)
+        asset.set_order(order=order)
+    
+    # 3. collapse must perform first
+    if transform_config.augment_config:
+        first_augments, second_augments = get_augments(config=transform_config.augment_config)
+    else:
+        first_augments = []
+        second_augments = []
+        
+    kwargs = {}
+    for augment in first_augments:
+        augment.transform(asset=asset, **kwargs)
+    
+    # 4. get vertex groups
+    if transform_config.vertex_group_config is not None:
+        vertex_groups = get_vertex_groups(config=transform_config.vertex_group_config)
+        d = {}
+        for v in vertex_groups:
+            d.update(v.get_vertex_group(asset=asset))
+        asset.vertex_groups = d
+    else:
+        asset.vertex_groups = {}
+    
+    # 5. regular augments
+    for augment in second_augments:
+        augment.transform(asset=asset, **kwargs)
+    
+    # 6. sample
+    if transform_config.sampler_config is not None:
+        sampler = get_sampler(config=transform_config.sampler_config)
+        res = sampler.sample(asset=asset)
+        asset.sampled_vertices      = res.vertices
+        asset.sampled_normals       = res.normals
+        asset.sampled_vertex_groups = res.vertex_groups
+    else:
+        asset.sampled_vertices      = asset.vertices.copy()
+        asset.sampled_normals       = asset.vertex_normals.copy()
+        asset.sampled_vertex_groups = deepcopy(asset.vertex_groups)
+    return first_augments, second_augments

UniRig/src/data/utils.py

@@ -0,0 +1,258 @@
+import torch
+import numpy as np
+from numpy import ndarray
+from torch import Tensor, FloatTensor
+from typing import Tuple, Union
+
+from scipy.spatial.transform import Rotation as R
+from scipy.sparse import csc_matrix
+import numpy as np
+
+def quaternion_to_matrix(x, use_4x4=True) -> FloatTensor:
+    """
+    Ref: https://pytorch3d.readthedocs.io/en/latest/_modules/pytorch3d/transforms/rotation_conversions.html#quaternion_to_matrix
+    """
+    if not isinstance(x, Tensor):
+        quaternions = torch.tensor(x, dtype=torch.float32)
+    else:
+        quaternions = x
+    r, i, j, k = torch.unbind(quaternions, -1)
+    two_s = 2.0 / (quaternions * quaternions).sum(-1)
+    device = quaternions.device
+    
+    if use_4x4:
+        o = torch.stack(
+            (
+                1 - two_s * (j * j + k * k),
+                two_s * (i * j - k * r),
+                two_s * (i * k + j * r),
+                torch.zeros(quaternions.shape[:-1], device=device, dtype=torch.float32),
+                two_s * (i * j + k * r),
+                1 - two_s * (i * i + k * k),
+                two_s * (j * k - i * r),
+                torch.zeros(quaternions.shape[:-1], device=device, dtype=torch.float32),
+                two_s * (i * k - j * r),
+                two_s * (j * k + i * r),
+                1 - two_s * (i * i + j * j),
+                torch.zeros(quaternions.shape[:-1], device=device, dtype=torch.float32),
+                torch.zeros(quaternions.shape[:-1], device=device, dtype=torch.float32),
+                torch.zeros(quaternions.shape[:-1], device=device, dtype=torch.float32),
+                torch.zeros(quaternions.shape[:-1], device=device, dtype=torch.float32),
+                torch.ones(quaternions.shape[:-1], device=device, dtype=torch.float32),
+            ),
+            -1,
+        )
+        return o.reshape(quaternions.shape[:-1] + (4, 4))
+    else:
+        o = torch.stack(
+            (
+                1 - two_s * (j * j + k * k),
+                two_s * (i * j - k * r),
+                two_s * (i * k + j * r),
+                two_s * (i * j + k * r),
+                1 - two_s * (i * i + k * k),
+                two_s * (j * k - i * r),
+                two_s * (i * k - j * r),
+                two_s * (j * k + i * r),
+                1 - two_s * (i * i + j * j),
+            ),
+            -1,
+        )
+        return o.reshape(quaternions.shape[:-1] + (3, 3))
+
+def axis_angle_to_quaternion(axis_angle: FloatTensor) -> FloatTensor:
+    """
+    Ref: https://pytorch3d.readthedocs.io/en/latest/_modules/pytorch3d/transforms/rotation_conversions.html#axis_angle_to_quaternion
+    """
+    angles = torch.norm(axis_angle, p=2, dim=-1, keepdim=True)
+    half_angles = angles * 0.5
+    eps = 1e-6
+    small_angles = angles.abs() < eps
+    sin_half_angles_over_angles = torch.empty_like(angles)
+    sin_half_angles_over_angles[~small_angles] = (
+        torch.sin(half_angles[~small_angles]) / angles[~small_angles]
+    )
+    # for x small, sin(x/2) is about x/2 - (x/2)^3/6
+    # so sin(x/2)/x is about 1/2 - (x*x)/48
+    sin_half_angles_over_angles[small_angles] = (
+        0.5 - (angles[small_angles] * angles[small_angles]) / 48
+    )
+    quaternions = torch.cat(
+        [torch.cos(half_angles), axis_angle * sin_half_angles_over_angles], dim=-1
+    )
+    return quaternions
+
+def axis_angle_to_matrix(axis_angle: Union[FloatTensor, ndarray]) -> Union[FloatTensor, ndarray]:
+    """
+    Ref: https://pytorch3d.readthedocs.io/en/latest/_modules/pytorch3d/transforms/rotation_conversions.html#axis_angle_to_matrix
+    """
+    if isinstance(axis_angle, FloatTensor):
+        return quaternion_to_matrix(axis_angle_to_quaternion(axis_angle))
+    else:
+        res = np.pad(R.from_rotvec(axis_angle).as_matrix(), ((0, 0), (0, 1), (0, 1)), 'constant', constant_values=((0, 0), (0, 0), (0, 0)))
+        assert res.ndim == 3
+        res[:, -1, -1] = 1
+        return res
+
+def _sqrt_positive_part(x: torch.Tensor) -> torch.Tensor:
+    """
+    Returns torch.sqrt(torch.max(0, x))
+    but with a zero subgradient where x is 0.
+    """
+    ret = torch.zeros_like(x)
+    positive_mask = x > 0
+    if torch.is_grad_enabled():
+        ret[positive_mask] = torch.sqrt(x[positive_mask])
+    else:
+        ret = torch.where(positive_mask, torch.sqrt(x), ret)
+    return ret
+
+def standardize_quaternion(quaternions: torch.Tensor) -> torch.Tensor:
+    """
+    Convert a unit quaternion to a standard form: one in which the real
+    part is non negative.
+
+    Args:
+        quaternions: Quaternions with real part first,
+            as tensor of shape (..., 4).
+
+    Returns:
+        Standardized quaternions as tensor of shape (..., 4).
+    """
+    return torch.where(quaternions[..., 0:1] < 0, -quaternions, quaternions)
+
+def matrix_to_quaternion(matrix: torch.Tensor) -> torch.Tensor:
+    """
+    Convert rotations given as rotation matrices to quaternions.
+
+    Args:
+        matrix: Rotation matrices as tensor of shape (..., 3, 3).
+
+    Returns:
+        quaternions with real part first, as tensor of shape (..., 4).
+    """
+    if matrix.size(-1) != 3 or matrix.size(-2) != 3:
+        raise ValueError(f"Invalid rotation matrix shape {matrix.shape}.")
+
+    batch_dim = matrix.shape[:-2]
+    m00, m01, m02, m10, m11, m12, m20, m21, m22 = torch.unbind(
+        matrix.reshape(batch_dim + (9,)), dim=-1
+    )
+
+    q_abs = _sqrt_positive_part(
+        torch.stack(
+            [
+                1.0 + m00 + m11 + m22,
+                1.0 + m00 - m11 - m22,
+                1.0 - m00 + m11 - m22,
+                1.0 - m00 - m11 + m22,
+            ],
+            dim=-1,
+        )
+    )
+
+    # we produce the desired quaternion multiplied by each of r, i, j, k
+    quat_by_rijk = torch.stack(
+        [
+            # pyre-fixme[58]: `**` is not supported for operand types `Tensor` and
+            #  `int`.
+            torch.stack([q_abs[..., 0] ** 2, m21 - m12, m02 - m20, m10 - m01], dim=-1),
+            # pyre-fixme[58]: `**` is not supported for operand types `Tensor` and
+            #  `int`.
+            torch.stack([m21 - m12, q_abs[..., 1] ** 2, m10 + m01, m02 + m20], dim=-1),
+            # pyre-fixme[58]: `**` is not supported for operand types `Tensor` and
+            #  `int`.
+            torch.stack([m02 - m20, m10 + m01, q_abs[..., 2] ** 2, m12 + m21], dim=-1),
+            # pyre-fixme[58]: `**` is not supported for operand types `Tensor` and
+            #  `int`.
+            torch.stack([m10 - m01, m20 + m02, m21 + m12, q_abs[..., 3] ** 2], dim=-1),
+        ],
+        dim=-2,
+    )
+
+    # We floor here at 0.1 but the exact level is not important; if q_abs is small,
+    # the candidate won't be picked.
+    flr = torch.tensor(0.1).to(dtype=q_abs.dtype, device=q_abs.device)
+    quat_candidates = quat_by_rijk / (2.0 * q_abs[..., None].max(flr))
+
+    # if not for numerical problems, quat_candidates[i] should be same (up to a sign),
+    # forall i; we pick the best-conditioned one (with the largest denominator)
+    out = quat_candidates[
+        torch.nn.functional.one_hot(q_abs.argmax(dim=-1), num_classes=4) > 0.5, :
+    ].reshape(batch_dim + (4,))
+    return standardize_quaternion(out)
+
+def linear_blend_skinning(
+    vertex: Union[FloatTensor, ndarray],
+    matrix_local: Union[FloatTensor, ndarray],
+    matrix: Union[FloatTensor, ndarray],
+    skin: Union[FloatTensor, ndarray],
+    pad: int=0,
+    value: float=0.,
+) -> Union[FloatTensor, ndarray]:
+    '''
+    Args:
+        vertex: (B, N, 4-pad) or (N, 4-pad)
+        matrix_local: (B, J, 4, 4) or (J, 4, 4)
+        matrix: (B, J, 4, 4) or (J, 4, 4)
+        skin: (B, N, J) or (N, J), value of pseudo bones should be 0
+    Returns:
+        (B, N, 3) or (N, 3)
+    '''
+    assert vertex.shape[-1] + pad == 4
+    if isinstance(vertex, Tensor):
+        dims = vertex.dim()
+    elif isinstance(vertex, ndarray):
+        dims = vertex.ndim
+    else:
+        raise NotImplementedError()
+    if dims == 3:  # Case: (B, N, 3+pad)
+        assert isinstance(vertex, Tensor)
+        J = matrix_local.shape[1]
+        # (B, J, 3+pad, N)
+        offset = (
+            matrix_local.inverse() @
+            torch.nn.functional.pad(vertex, (0, pad, 0, 0, 0, 0), value=value).unsqueeze(1).transpose(2, 3).repeat(1, J, 1, 1)
+        )
+        # (B, J, 4, N)
+        per_bone_matrix = matrix @ offset
+        # (B, J, 4, N)
+        weighted_per_bone_matrix = skin.transpose(1, 2).unsqueeze(2) * per_bone_matrix
+        # (B, 3, N)
+        g = weighted_per_bone_matrix.sum(dim=1)
+        # (B, 3, N)
+        final = g[:, 0:3, :] / (skin.transpose(1, 2).sum(dim=1) + 1e-8).unsqueeze(1)
+        return final.permute(0, 2, 1)
+    
+    elif dims == 2:  # Case: (N, 3+pad)
+        if isinstance(vertex, Tensor):
+            J = matrix_local.shape[0]
+            offset = (
+                matrix_local.inverse() @
+                torch.nn.functional.pad(vertex, (0, pad, 0, 0), value=value).unsqueeze(0).transpose(1, 2).repeat(J, 1, 1)
+            )
+            per_bone_matrix = matrix @ offset
+            weighted_per_bone_matrix = skin.transpose(0, 1).unsqueeze(1) * per_bone_matrix
+            g = weighted_per_bone_matrix.sum(dim=0)
+            final = g[0:3, :] / (skin.transpose(0, 1).sum(dim=0) + 1e-8).unsqueeze(0)
+            return final.permute(1, 0)  # Output shape (N, 3)
+        else:
+            J = matrix_local.shape[0]
+            N = vertex.shape[0]
+            # (4, N)
+            padded = np.pad(vertex, ((0, 0), (0, pad)), 'constant', constant_values=(0, value)).T
+            # (J, 4, 4)
+            trans = matrix @ np.linalg.inv(matrix_local)
+            weighted_per_bone_matrix = []
+            # (J, N)
+            mask = (skin > 0).T
+            for i in range(J):
+                offset = np.zeros((4, N), dtype=np.float32)
+                offset[:, mask[i]] = (trans[i] @ padded[:, mask[i]]) * skin.T[i, mask[i]]
+                weighted_per_bone_matrix.append(offset)
+            weighted_per_bone_matrix = np.stack(weighted_per_bone_matrix)
+            g = np.sum(weighted_per_bone_matrix, axis=0)
+            final = g[:3, :] / (np.sum(skin, axis=1) + 1e-8)
+            return final.T
+    else:
+        assert 0, f'unsupported shape: {vertex.shape}'