Initial snapshot: hand motion pipeline (Dyn-HaMR + dex-retargeting + SPIDER)
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Vendored
+21
@@ -0,0 +1,21 @@
|
||||
MIT License
|
||||
|
||||
Copyright (c) 2024 Zhengdi Yu
|
||||
|
||||
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.
|
||||
Vendored
+241
@@ -0,0 +1,241 @@
|
||||
<p align="center">
|
||||
|
||||
<h1 align="center">Dyn-HaMR: Recovering 4D Interacting Hand Motion from a Dynamic Camera</h1>
|
||||
<p align="center">
|
||||
<a href="https://github.com/ZhengdiYu"><strong>Zhengdi Yu</strong></a>
|
||||
·
|
||||
<a href="https://scholar.google.com/citations?user=QKOH5iYAAAAJ&hl=en&oi=ao"><strong>Stefanos Zafeiriou</strong></a>
|
||||
·
|
||||
<a href="https://tolgabirdal.github.io/"><strong>Tolga Birdal</strong></a>
|
||||
</p>
|
||||
<p align="center">
|
||||
<strong>Imperial College London</strong>
|
||||
</p>
|
||||
</p>
|
||||
<p align="center">
|
||||
<strong>CVPR 2025 (Highlight)</strong>
|
||||
</p>
|
||||
|
||||
<div align="center" style="display: flex; justify-content: center; gap: 20px;">
|
||||
<img src="./assets/wild.gif" alt="Wild GIF" width="45%" style="border-radius: 10px;">
|
||||
<img src="./assets/global.gif" alt="Global GIF" width="45%" style="border-radius: 10px;">
|
||||
</div>
|
||||
|
||||
<p align="center">
|
||||
<a href='https://arxiv.org/abs/2412.12861'>
|
||||
<img src='https://img.shields.io/badge/Arxiv-2412.12861-A42C25?style=flat&logo=arXiv&logoColor=A42C25'>
|
||||
</a>
|
||||
<a href='https://arxiv.org/pdf/2412.12861'>
|
||||
<img src='https://img.shields.io/badge/Paper-PDF-green?style=flat&logo=arXiv&logoColor=green'>
|
||||
</a>
|
||||
<a href='https://dyn-hamr.github.io/'>
|
||||
<img src='https://img.shields.io/badge/Project-Page-blue?style=flat&logo=Google%20chrome&logoColor=blue'></a>
|
||||
<a href='https://youtu.be/n25NGIWiA7M'>
|
||||
<img src='https://img.shields.io/badge/YouTube-Video-EA3323?style=flat&logo=youtube&logoColor=EA3323'></a>
|
||||
<img src="https://visitor-badge.laobi.icu/badge?page_id=ZhengdiYu.Dyn-HaMR&left_color=gray&right_color=orange">
|
||||
</a>
|
||||
</p>
|
||||
</p>
|
||||
|
||||
|
||||
## Introduction
|
||||
<img src="./assets/teaser.png">
|
||||
|
||||
We propose **Dyn-HaMR** to reconstruct 4D global hand motion from monocular videos recorded by dynamic cameras in the wild, as a remedy for the motion entanglement in the wild.
|
||||
|
||||
<!-- TABLE OF CONTENTS -->
|
||||
<details open="open" style='padding: 10px; border-radius:5px 30px 30px 5px; border-style: solid; border-width: 1px;'>
|
||||
<summary>Table of Contents</summary>
|
||||
<ol>
|
||||
<li>
|
||||
<a href="#installation">Installation</a>
|
||||
</li>
|
||||
<li>
|
||||
<a href="#get-started">Get Started</a>
|
||||
</li>
|
||||
<li>
|
||||
<a href="#citation">Citation</a>
|
||||
</li>
|
||||
</ol>
|
||||
</details>
|
||||
<br/>
|
||||
|
||||
## News :triangular_flag_on_post:
|
||||
- [2025/11/20] 🚀 **Major Update**:
|
||||
- **Integrated [VIPE](https://github.com/nv-tlabs/vipe)** for camera estimation, significantly improving reconstruction quality over DROID-SLAM
|
||||
- **Enhanced Hand Tracker** with robust hallucination prevention and handedness correction for better hand tracking and **significantly** Improved temporal consistency. Please `pip install ultralytics==8.1.34` since YOLO is using in this version (Thanks to [WiloR](https://github.com/rolpotamias/WiLoR)). Please download the checkpoint from [here](https://huggingface.co/spaces/rolpotamias/WiLoR/resolve/main/pretrained_models/detector.pt) and put it under `third-party/hamer/pretrained_models`.
|
||||
|
||||
See comparison below:
|
||||
|
||||
<table>
|
||||
<tr>
|
||||
<th>Before: Jitter from DROID-SLAM</th>
|
||||
<th>New: VIPE + enhanced HaMeR (Recommended)</th>
|
||||
</tr>
|
||||
<tr>
|
||||
<td align="center">
|
||||
<img src="./assets/droid_result.gif" width="100%">
|
||||
</td>
|
||||
<td align="center">
|
||||
<img src="./assets/vipe_result.gif" width="100%">
|
||||
</td>
|
||||
</tr>
|
||||
</table>
|
||||
|
||||
<table>
|
||||
<tr>
|
||||
<th>Before: Handedness shifting with original HaMeR</th>
|
||||
<th>New: Enhanced hand tracker (Recommended)</th>
|
||||
</tr>
|
||||
<tr>
|
||||
<td align="center">
|
||||
<img src="./assets/handedness1.gif" width="100%">
|
||||
</td>
|
||||
<td align="center">
|
||||
<img src="./assets/handedness2.gif" width="100%">
|
||||
</td>
|
||||
</tr>
|
||||
</table>
|
||||
|
||||
<table>
|
||||
<tr>
|
||||
<th>Before: Handedness shifting with original HaMeR</th>
|
||||
<th>New: Enhanced hand tracker (Recommended)</th>
|
||||
</tr>
|
||||
<tr>
|
||||
<td align="center">
|
||||
<img src="./assets/handedness3.gif" width="100%">
|
||||
</td>
|
||||
<td align="center">
|
||||
<img src="./assets/handedness4.gif" width="100%">
|
||||
</td>
|
||||
</tr>
|
||||
</table>
|
||||
|
||||
- [2025/06/04] Code released.
|
||||
- [2024/12/18] [Paper](https://arxiv.org/abs/2412.12861) is now available on arXiv. ⭐
|
||||
|
||||
## Installation
|
||||
|
||||
### Environment setup
|
||||
1. Clone the repository with submodules with the following command:
|
||||
```bash
|
||||
git clone --recursive https://github.com/ZhengdiYu/Dyn-HaMR.git
|
||||
cd Dyn-HaMR
|
||||
```
|
||||
You can also run the following command to fetch the submodules:
|
||||
```bash
|
||||
git submodule update --init --recursive .
|
||||
```
|
||||
|
||||
2. To set up the virtual environment for Dyn-HaMR, we provide the integrated commands in `\scripts`. You can create the environment from
|
||||
```bash
|
||||
source install_pip.sh
|
||||
```
|
||||
|
||||
Or, alternatively, create the environment from conda:
|
||||
```bash
|
||||
source install_conda.sh
|
||||
```
|
||||
|
||||
### Model checkpoints download
|
||||
Please run the following command to fetch the data dependencies. This will create a folder in `_DATA`:
|
||||
```bash
|
||||
source prepare.sh
|
||||
```
|
||||
After processing, the folder layout should be:
|
||||
```
|
||||
|-- _DATA
|
||||
| |-- data/
|
||||
| | |-- mano/
|
||||
| | | |-- MANO_RIGHT.pkl
|
||||
| | |-- mano_mean_params.npz
|
||||
| |-- BMC/
|
||||
| |-- hamer_ckpts/
|
||||
| |-- vitpose_ckpts/
|
||||
| |-- <SLAM model .pkl>
|
||||
```
|
||||
|
||||
### Prerequisites
|
||||
We use [MANO](https://mano.is.tue.mpg.de) model for hand mesh representation. Please visit the [MANO website](https://mano.is.tue.mpg.de) for registration and the model downloading. Please download `MANO_RIGHT.pkl` and put under the `_DATA/data/mano` folder.
|
||||
|
||||
## Get Started🚀
|
||||
|
||||
### Preparation
|
||||
Please follow the instructions [here](https://github.com/MengHao666/Hand-BMC-pytorch) to calculate the below `.npz` files in order `dyn-hamr/optim/BMC/`:
|
||||
```
|
||||
|-- BMC
|
||||
| |-- bone_len_max.npy
|
||||
| |-- bone_len_min.npy
|
||||
| |-- CONVEX_HULLS.npy
|
||||
| |-- curvatures_max.npy
|
||||
| |-- curvatures_min.npy
|
||||
| |-- joint_angles.npy
|
||||
| |-- PHI_max.npy
|
||||
| |-- PHI_min.npy
|
||||
```
|
||||
|
||||
> [!NOTE]
|
||||
> If accurate camera parameters are available, please follow the format of `Dyn-HaMR/test/dynhamr/cameras/demo/shot-0/cameras.npz` to prepare the camera parameters for loading. Similarly, you can use Dyn-HaMR to refine and recover the hand mesh in the world coordinate system initializing from your own 2D & 3D motion data.
|
||||
|
||||
### Customize configurations
|
||||
| Config | Operation |
|
||||
|--------|-----------------|
|
||||
| GPU | Edit in [`<CONFIG_GPU>`](https://github.com/ZhengdiYu/Dyn-HaMR/blob/main/dyn-hamr/confs/config.yaml#L56) |
|
||||
| Video info | Edit in [`<VIDEO_SEQ>`](https://github.com/ZhengdiYu/Dyn-HaMR/blob/main/dyn-hamr/confs/data/video.yaml#L5) |
|
||||
| Interval | Edit in [`<VIDEO_START_END>`](https://github.com/ZhengdiYu/Dyn-HaMR/blob/main/dyn-hamr/confs/data/video.yaml#L16-L17) |
|
||||
| Optimization configurations | Edit in [`<OPT_WEIGHTS>`](https://github.com/ZhengdiYu/Dyn-HaMR/blob/main/dyn-hamr/confs/optim.yaml#L29-L49) |
|
||||
| General configurations | Edit in [`<GENERAL_CONFIG>`](https://github.com/ZhengdiYu/Dyn-HaMR/blob/main/dyn-hamr/confs/config.yaml) |
|
||||
|
||||
### Fitting on RGB-(D) videos 🎮
|
||||
To run the optimization pipeline for fitting on arbitrary RGB-(D) videos, please first edit the path information here in `dyn-hamr/confs/data/video.yaml`, where `root` is the root folder to all of your datasets. `video_dir` is the corresponding folder that contains the videos. The key `seq` represents the video name you wanted to process. For example, you can run the following command to recover the global motion for `test/videos/demo1.mp4`:
|
||||
|
||||
|
||||
#### 🌟 Using VIPE for Camera Estimation (Recommended)
|
||||
For significantly better camera estimation quality, use VIPE instead of DROID-SLAM:
|
||||
|
||||
```bash
|
||||
python run_opt.py data=video_vipe run_opt=True data.seq=demo1 is_static=False
|
||||
```
|
||||
|
||||
#### 🌟 Using original DROID-SLAM for Camera Estimation
|
||||
```bash
|
||||
python run_opt.py data=video run_opt=True data.seq=demo1 is_static=<True or False>
|
||||
```
|
||||
|
||||
VIPE will automatically run if results are not found. Make sure you have:
|
||||
1. Installed VIPE in `third-party/vipe/` with conda environment named `vipe`
|
||||
2. Set `src_path` in `dyn-hamr/confs/data/video_vipe.yaml` to your video file
|
||||
|
||||
By default, the camera parameters will be predicted during the process and assumes a moving camera (`is_static=False`). If your video is recorded with a static camera, you can add `is_static=True` for more stable optimization. The result will be saved to `outputs/logs/video-custom/<DATE>/<VIDEO_NAME>-<tracklet>-shot-<shot_id>-<start_frame_id>-<end_frame_id>`. After optimization, you can specify the output log dir and visualize the results by running the following command:
|
||||
```
|
||||
python run_vis.py --log_root <LOG_ROOT>
|
||||
```
|
||||
This will visualize all log subdirectories and save the rendered videos and images, as well as saved 3D meshes in the world space in `<LOG_ROOT>`. Please visit `run_vis.py` for further details. Alternatively, you can also use the following command to run and visualize the results in one-stage:
|
||||
```
|
||||
python -u run_opt.py data=video_vipe run_opt=True run_vis=True is_static=<True of False>
|
||||
```
|
||||
As a multi-stage pipeline, you can customize the optimization process. Add `is_static=True` for static camera videos. Adding `run_prior=True` can activate the motion prior in stage III. Please note that in the current version, each motion chunk size needs to be set to 128 to be compatible with the original setting of HMP only when the prior module is activated.
|
||||
|
||||
### Blender Addon
|
||||
Coming soon.
|
||||
|
||||
## Acknowledgements
|
||||
The PyTorch implementation of MANO is based on [manopth](https://github.com/hassony2/manopth). Part of the fitting and optimization code of this repository is borrowed from [SLAHMR](https://github.com/vye16/slahmr). For data preprocessing and observation, [ViTPose](https://github.com/ViTAE-Transformer/ViTPose) and [HaMeR](https://github.com/geopavlakos/hamer/) is used for 2D keypoints detection and MANO parameter initilization. For camera motion estimation, we support [VIPE](https://github.com/facebookresearch/vipe) (recommended), [DPVO](https://github.com/princeton-vl/DPVO), and [DROID-SLAM](https://github.com/princeton-vl/DROID-SLAM). For biomechanical constraints and motion prior, we use the code from [here](https://github.com/MengHao666/Hand-BMC-pytorch) and [HMP](https://hmp.is.tue.mpg.de/). We thank all the authors for their impressive work!
|
||||
|
||||
## License
|
||||
Please see [License](https://github.com/ZhengdiYu/Dyn-HaMR/blob/main/LICENSE) for details of Dyn-HaMR. This code and model are available only for non-commercial research purposes as defined in the LICENSE (i.e., MIT LICENSE). Note that, for MANO you must agree with the LICENSE of it. You can check the LICENSE of MANO from https://mano.is.tue.mpg.de/license.html.
|
||||
|
||||
## Citation
|
||||
```bibtex
|
||||
@inproceedings{yu2025dynhamr,
|
||||
title={Dyn-HaMR: Recovering 4D Interacting Hand Motion from a Dynamic Camera},
|
||||
author={Yu, Zhengdi and Zafeiriou, Stefanos and Birdal, Tolga},
|
||||
booktitle = {Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
|
||||
month = {June},
|
||||
year = {2025},
|
||||
}
|
||||
```
|
||||
|
||||
## Contact
|
||||
For any technical questions, please contact z.yu23@imperial.ac.uk or ZhengdiYu@hotmail.com.
|
||||
+11
@@ -0,0 +1,11 @@
|
||||
dataset_dir: ./data/amass/generative
|
||||
data:
|
||||
fps: 30
|
||||
clip_length: 128
|
||||
gender: neutral
|
||||
up: z
|
||||
|
||||
canonical: False
|
||||
unified_orientation: True
|
||||
single_motion: False
|
||||
normalize: True
|
||||
+36
@@ -0,0 +1,36 @@
|
||||
import os
|
||||
|
||||
import numpy as np
|
||||
import yaml
|
||||
|
||||
|
||||
class Struct:
|
||||
def __init__(self, **entries):
|
||||
self.__dict__.update(entries)
|
||||
|
||||
|
||||
class Arguments:
|
||||
def __init__(self, base_dir, config_path, filename='default.yaml'):
|
||||
with open(os.path.join(config_path, 'mano.yaml'), 'r') as f:
|
||||
smpl = yaml.safe_load(f)
|
||||
|
||||
self.smpl = Struct(**smpl)
|
||||
self.smpl.offsets['right'] = np.array(self.smpl.offsets['right'])
|
||||
self.smpl.offsets['left'] = np.array(self.smpl.offsets['left'])
|
||||
self.smpl.parents = np.array(self.smpl.parents).astype(np.int32)
|
||||
self.smpl.joint_num = 16 # len(self.smpl.joints_to_use)
|
||||
self.smpl.joints_to_use = np.array(self.smpl.joints_to_use)
|
||||
self.smpl.joints_to_use = np.arange(0, 63).reshape((-1, 3))[self.smpl.joints_to_use].reshape(-1)
|
||||
self.smpl.smpl_body_model = os.path.join(base_dir, '_DATA/data/mano')
|
||||
|
||||
self.filename = os.path.splitext(filename)[0]
|
||||
with open(os.path.join(config_path, filename), 'r') as f:
|
||||
config = yaml.safe_load(f)
|
||||
|
||||
for key, value in config.items():
|
||||
if isinstance(value, dict):
|
||||
setattr(self, key, Struct(**value))
|
||||
else:
|
||||
setattr(self, key, value)
|
||||
|
||||
self.json = config
|
||||
+639
@@ -0,0 +1,639 @@
|
||||
import os
|
||||
import glob
|
||||
import json
|
||||
import torch
|
||||
import joblib
|
||||
import random
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
from scipy import linalg
|
||||
from arguments import Arguments
|
||||
from nemf.losses import GeodesicLoss
|
||||
from nemf.generative import Architecture
|
||||
from nemf.fk import ForwardKinematicsLayer
|
||||
from scripts.train_gmm import MaxMixturePrior
|
||||
from human_body_prior.tools.omni_tools import copy2cpu as c2c
|
||||
from utils import FOOT_IDX, estimate_angular_velocity, estimate_linear_velocity
|
||||
from rotations import axis_angle_to_matrix, matrix_to_rotation_6d, rotation_6d_to_matrix
|
||||
|
||||
|
||||
def padding_prior_data(data):
|
||||
N, T, J = data['pos'].shape[:3]
|
||||
# pad inputs to match the training sequence length
|
||||
zero_padding = torch.zeros(N, args.data.clip_length - T, J, 3).to(prior.device)
|
||||
identity_padding = torch.zeros(N, args.data.clip_length - T, 3, 3).to(prior.device)
|
||||
identity_padding[:, :, 0, 0] = 1
|
||||
identity_padding[:, :, 1, 1] = 1
|
||||
identity_padding[:, :, 2, 2] = 1
|
||||
rotation_6d_padding = torch.zeros(N, args.data.clip_length - T, 6).to(prior.device)
|
||||
rotation_6d_padding[..., 0] = 1
|
||||
rotation_6d_padding[..., 4] = 1
|
||||
|
||||
pos = torch.cat((data['pos'], zero_padding), dim=1)
|
||||
velocity = torch.cat((data['velocity'], zero_padding), dim=1)
|
||||
angular = torch.cat((data['angular'], zero_padding), dim=1)
|
||||
root_orient = torch.cat((data['root_orient'], rotation_6d_padding), dim=1)
|
||||
root_vel = torch.cat((data['root_vel'], zero_padding[:, :, 0]), dim=1)
|
||||
rotmat_padding = rotation_6d_padding.unsqueeze(2).expand(-1, -1, J, -1)
|
||||
global_xform = torch.cat((data['global_xform'], rotmat_padding), dim=1)
|
||||
|
||||
data['pos'] = pos
|
||||
data['velocity'] = velocity
|
||||
data['global_xform'] = global_xform
|
||||
data['angular'] = angular
|
||||
data['root_orient'] = root_orient
|
||||
data['root_vel'] = root_vel
|
||||
|
||||
return data
|
||||
|
||||
|
||||
def load_data(files, padding=False, max_samples=400):
|
||||
poses = []
|
||||
trans = []
|
||||
visibility = []
|
||||
assert len(files) != 0, 'files not found'
|
||||
|
||||
max_samples = min(max_samples, len(files))
|
||||
|
||||
# import ipdb; ipdb.set_trace()
|
||||
|
||||
for f in files[:max_samples]:
|
||||
bdata = dict(np.load(f, allow_pickle=True))
|
||||
if 'r_wrist_orient' in bdata.keys() and 'pose_right_hand' in bdata.keys():
|
||||
|
||||
coef = len(bdata['r_wrist_orient']) // args.data.clip_length
|
||||
|
||||
for k_ in range(coef):
|
||||
root_orient = bdata['r_wrist_orient'][k_*args.data.clip_length:(k_+1)*args.data.clip_length]
|
||||
pose_body = bdata['pose_right_hand'][k_*args.data.clip_length:(k_+1)*args.data.clip_length]
|
||||
trans.append(bdata['r_wrist_location'][k_*args.data.clip_length:(k_+1)*args.data.clip_length])
|
||||
|
||||
poses.append(np.concatenate((root_orient, pose_body), axis=-1))
|
||||
|
||||
elif 'pose_right_hand' in bdata.keys():
|
||||
pass
|
||||
# poses.append(bdata['pose_right_hand'][:, :72])
|
||||
else:
|
||||
raise RuntimeError(f'missing pose parameters in the file: {f}')
|
||||
|
||||
|
||||
if 'visibility' in bdata.keys():
|
||||
visibility.append(bdata['visibility'])
|
||||
|
||||
# if 'poses' in bdata.keys():
|
||||
# poses.append(bdata['poses'][:, :72])
|
||||
# elif 'root_orient' in bdata.keys() and 'pose_body' in bdata.keys():
|
||||
# root_orient = bdata['root_orient']
|
||||
# pose_body = bdata['pose_body']
|
||||
# poses.append(np.concatenate((root_orient, pose_body), axis=-1))
|
||||
# else:
|
||||
# raise RuntimeError(f'missing pose parameters in the file: {f}')
|
||||
# trans.append(bdata['trans'])
|
||||
# if 'visibility' in bdata.keys():
|
||||
# visibility.append(bdata['visibility'])
|
||||
|
||||
trans = torch.from_numpy(np.asarray(trans, np.float32)).cuda() # global translation (N, T, 3)
|
||||
N, T = trans.shape[:2]
|
||||
poses = torch.from_numpy(np.asarray(poses, np.float32)).cuda()
|
||||
poses = poses.view(N, T, -1, 3) # axis-angle (N, T, J, 3)
|
||||
|
||||
|
||||
# poses = poses.view(N, T, 24, 3) # axis-angle (N, T, J, 3)
|
||||
|
||||
root_orient = poses[:, :, 0].clone()
|
||||
root_rotation = axis_angle_to_matrix(root_orient) # (N, T, 3, 3)
|
||||
poses[:, :, 0] = 0
|
||||
|
||||
rotmat = axis_angle_to_matrix(poses) # (N, T, J, 3, 3)
|
||||
angular = estimate_angular_velocity(rotmat.clone(), dt=1.0 / args.data.fps) # angular velocity of all the joints (N, T, J, 3)
|
||||
pos, global_xform = fk(rotmat.view(-1, 16, 3, 3))
|
||||
pos = pos.contiguous().view(N, T, 16, 3) # local joint positions (N, T, J, 3)
|
||||
global_xform = global_xform.view(N, T, 16, 4, 4)
|
||||
global_xform = global_xform[:, :, :, :3, :3] # global transformation matrix for each joint (N, T, J, 3, 3)
|
||||
velocity = estimate_linear_velocity(pos, dt=1.0 / args.data.fps) # linear velocity of all the joints (N, T, J, 3)
|
||||
|
||||
root_vel = estimate_linear_velocity(trans, dt=1.0 / args.data.fps) # linear velocity of the root joint (N, T, 3)
|
||||
|
||||
global_pos = torch.matmul(root_rotation.unsqueeze(2), pos.unsqueeze(-1)).squeeze(-1) # (N, T, J, 3)
|
||||
global_pos = global_pos + trans.unsqueeze(2)
|
||||
|
||||
data = {
|
||||
'pos': pos,
|
||||
'velocity': velocity,
|
||||
'global_xform': matrix_to_rotation_6d(global_xform),
|
||||
'angular': angular,
|
||||
'root_orient': matrix_to_rotation_6d(root_rotation),
|
||||
'root_vel': root_vel,
|
||||
'global_pos': global_pos,
|
||||
'rotmat': rotmat,
|
||||
'trans': trans
|
||||
}
|
||||
|
||||
if len(visibility) != 0:
|
||||
visibility = torch.from_numpy(np.asarray(visibility, bool)).cuda()
|
||||
padding = torch.zeros(visibility.shape[:2] + torch.Size([2]), dtype=torch.bool).cuda()
|
||||
visibility = torch.cat((visibility, padding), dim=-1) # (T, 24)
|
||||
print(f'visibility: {visibility.shape}')
|
||||
data['visibility'] = visibility
|
||||
|
||||
if T < args.data.clip_length:
|
||||
padding = True
|
||||
|
||||
if padding:
|
||||
if T == args.data.clip_length:
|
||||
data['pos'] = data['pos'].view(2 * N, T // 2, 24, 3)
|
||||
data['velocity'] = data['velocity'].view(2 * N, T // 2, 24, 3)
|
||||
data['global_xform'] = data['global_xform'].view(2 * N, T // 2, 24, 6)
|
||||
data['angular'] = data['angular'].view(2 * N, T // 2, 24, 3)
|
||||
data['root_orient'] = data['root_orient'].view(2 * N, T // 2, 6)
|
||||
data['root_vel'] = data['root_vel'].view(2 * N, T // 2, 3)
|
||||
data['global_pos'] = data['global_pos'].view(2 * N, T // 2, 24, 3)
|
||||
data['rotmat'] = data['rotmat'].view(2 * N, T // 2, 24, 3, 3)
|
||||
data['trans'] = data['trans'].view(2 * N, T // 2, 3)
|
||||
|
||||
padding_data = padding_prior_data(data)
|
||||
for key in padding_data.keys():
|
||||
data[key] = padding_data[key]
|
||||
|
||||
return data
|
||||
|
||||
|
||||
def calculate_feature_statistics(feature):
|
||||
feature = feature.cpu().numpy()
|
||||
mu = np.mean(feature, axis=0)
|
||||
sigma = np.cov(feature, rowvar=False)
|
||||
|
||||
return mu, sigma
|
||||
|
||||
|
||||
def calculate_frechet_distance(mu1, sigma1, mu2, sigma2, eps=1e-6):
|
||||
"""Numpy implementation of the Frechet Distance.
|
||||
The Frechet distance between two multivariate Gaussians X_1 ~ N(mu_1, C_1)
|
||||
and X_2 ~ N(mu_2, C_2) is
|
||||
d^2 = ||mu_1 - mu_2||^2 + Tr(C_1 + C_2 - 2*sqrt(C_1*C_2)).
|
||||
Stable version by Dougal J. Sutherland.
|
||||
Params:
|
||||
-- mu1 : Numpy array containing the activations of a layer of the
|
||||
inception net (like returned by the function 'get_predictions')
|
||||
for generated samples.
|
||||
-- mu2 : The sample mean over activations, precalculated on an
|
||||
representative data set.
|
||||
-- sigma1: The covariance matrix over activations for generated samples.
|
||||
-- sigma2: The covariance matrix over activations, precalculated on an
|
||||
representative data set.
|
||||
Returns:
|
||||
-- : The Frechet Distance.
|
||||
"""
|
||||
|
||||
mu1 = np.atleast_1d(mu1)
|
||||
mu2 = np.atleast_1d(mu2)
|
||||
|
||||
sigma1 = np.atleast_2d(sigma1)
|
||||
sigma2 = np.atleast_2d(sigma2)
|
||||
|
||||
assert mu1.shape == mu2.shape, \
|
||||
'Training and test mean vectors have different lengths'
|
||||
assert sigma1.shape == sigma2.shape, \
|
||||
'Training and test covariances have different dimensions'
|
||||
|
||||
diff = mu1 - mu2
|
||||
|
||||
# Product might be almost singular
|
||||
covmean, _ = linalg.sqrtm(sigma1.dot(sigma2), disp=False)
|
||||
if not np.isfinite(covmean).all():
|
||||
msg = ('fid calculation produces singular product; '
|
||||
'adding %s to diagonal of cov estimates') % eps
|
||||
print(msg)
|
||||
offset = np.eye(sigma1.shape[0]) * eps
|
||||
covmean = linalg.sqrtm((sigma1 + offset).dot(sigma2 + offset))
|
||||
|
||||
# Numerical error might give slight imaginary component
|
||||
if np.iscomplexobj(covmean):
|
||||
if not np.allclose(np.diagonal(covmean).imag, 0, atol=1e-3):
|
||||
m = np.max(np.abs(covmean.imag))
|
||||
raise ValueError('Imaginary component {}'.format(m))
|
||||
covmean = covmean.real
|
||||
|
||||
tr_covmean = np.trace(covmean)
|
||||
|
||||
return (diff.dot(diff) + np.trace(sigma1) + np.trace(sigma2) - 2 * tr_covmean)
|
||||
|
||||
|
||||
def evaluate_fid(data, gt_data, num_motions=None, mask=None):
|
||||
if num_motions is None:
|
||||
index = list(range(data['pos'].shape[0]))
|
||||
index_gt = list(range(gt_data['pos'].shape[0]))
|
||||
else:
|
||||
# randomly pick [num_motions] motions from the data
|
||||
assert num_motions <= data['pos'].shape[0]
|
||||
assert num_motions <= gt_data['pos'].shape[0]
|
||||
|
||||
index = list(range(data['pos'].shape[0]))
|
||||
index = random.sample(index, num_motions)
|
||||
index_gt = list(range(gt_data['pos'].shape[0]))
|
||||
index_gt = random.sample(index_gt, num_motions)
|
||||
|
||||
data_selected = {}
|
||||
gt_data_selected = {}
|
||||
for key in data.keys():
|
||||
data_selected[key] = data[key][index].clone()
|
||||
gt_data_selected[key] = gt_data[key][index_gt].clone()
|
||||
|
||||
if mask is not None:
|
||||
# replace with ground truth
|
||||
for key in data.keys():
|
||||
data_selected[key][:, mask] = gt_data[key][:, mask].clone()
|
||||
|
||||
with torch.no_grad():
|
||||
prior.set_input(data_selected)
|
||||
f_l, _, _ = prior.encode_local()
|
||||
f_g, _, _ = prior.encode_global()
|
||||
feature = torch.cat((f_l, f_g), dim=-1)
|
||||
|
||||
prior.set_input(gt_data_selected)
|
||||
f_l_gt, _, _ = prior.encode_local()
|
||||
f_g_gt, _, _ = prior.encode_global()
|
||||
feature_gt = torch.cat((f_l_gt, f_g_gt), dim=-1)
|
||||
|
||||
statistics = calculate_feature_statistics(feature.view(feature.shape[0], -1))
|
||||
statistics_gt = calculate_feature_statistics(feature_gt.view(feature_gt.shape[0], -1))
|
||||
|
||||
fid = calculate_frechet_distance(statistics[0], statistics[1], statistics_gt[0], statistics_gt[1])
|
||||
print(f'FID: {fid:.4f}')
|
||||
|
||||
return fid
|
||||
|
||||
|
||||
def evaluate_diversity(data):
|
||||
index = list(range(data['pos'].shape[0]))
|
||||
random.shuffle(index)
|
||||
index1 = index[:data['pos'].shape[0] // 2]
|
||||
index2 = index[data['pos'].shape[0] // 2:]
|
||||
|
||||
with torch.no_grad():
|
||||
prior.set_input(data)
|
||||
f_l, _, _ = prior.encode_local()
|
||||
# f_g, _, _ = prior.encode_global()
|
||||
# feature = torch.cat((f_l, f_g), dim=-1)
|
||||
|
||||
feature = f_l.clone()
|
||||
|
||||
|
||||
diversity = 0
|
||||
for first_idx, second_idx in zip(index1, index2):
|
||||
diversity += torch.dist(feature[first_idx], feature[second_idx])
|
||||
diversity /= len(index1)
|
||||
|
||||
print(f'Diversity: {diversity.item():.4f}')
|
||||
|
||||
return diversity.item()
|
||||
|
||||
|
||||
def evaluate_recon_error(data, gt_data):
|
||||
criterion_geo = GeodesicLoss()
|
||||
|
||||
root_orient = rotation_6d_to_matrix(data['root_orient'])
|
||||
root_orient_gt = rotation_6d_to_matrix(gt_data['root_orient'])
|
||||
orientation_error = criterion_geo(root_orient.view(-1, 3, 3), root_orient_gt.view(-1, 3, 3))
|
||||
|
||||
rotation_error = criterion_geo(data['rotmat'][:, :, 1:].reshape(-1, 3, 3), gt_data['rotmat'][:, :, 1:].reshape(-1, 3, 3))
|
||||
position_error = torch.linalg.norm((data['pos'] - gt_data['pos']), dim=-1).mean()
|
||||
translation_error = torch.linalg.norm((data['trans'] - gt_data['trans']), dim=-1).mean()
|
||||
final_displacement_error = torch.linalg.norm(data['global_pos'][:, -1] - gt_data['global_pos'][:, -1], dim=-1).mean()
|
||||
|
||||
return {
|
||||
'rotation': c2c(rotation_error) * 180.0 / np.pi,
|
||||
'position': c2c(position_error) * 100.0,
|
||||
'orientation': c2c(orientation_error) * 180.0 / np.pi,
|
||||
'translation': c2c(translation_error) * 100.0,
|
||||
'final_displacement': c2c(final_displacement_error) * 100.0
|
||||
}
|
||||
|
||||
|
||||
def synthetic_trajectory(N_motions):
|
||||
|
||||
z_l = torch.normal(mean=0, std=1, size=(N_motions, args.local_prior.z_dim)).to(prior.device)
|
||||
z_g = torch.normal(mean=0, std=1, size=(N_motions, args.global_prior.z_dim)).to(prior.device)
|
||||
# z_g = torch.zeros((N_motions, args.global_prior.z_dim)).to(prior.device)
|
||||
|
||||
|
||||
with torch.no_grad():
|
||||
output = prior.decode(z_l, z_g, length=args.data.clip_length, step=1)
|
||||
|
||||
rotmat = output["rotmat"] # (N, T, J, 3, 3)
|
||||
angular = estimate_angular_velocity(rotmat.clone(), dt=1.0 / args.data.fps) # angular velocity of all the joints (N, T, J, 3)
|
||||
pos, global_xform = fk(rotmat.view(-1, 16, 3, 3))
|
||||
pos = pos.contiguous().view(N_motions, args.data.clip_length, 16, 3) # local joint positions (N, T, J, 3)
|
||||
global_xform = global_xform.view(N_motions, args.data.clip_length, 16, 4, 4)
|
||||
|
||||
global_xform = global_xform[:, :, :, :3, :3] # global transformation matrix for each joint (N, T, J, 3, 3)
|
||||
velocity = estimate_linear_velocity(pos, dt=1.0 / args.data.fps) # linear velocity of all the joints (N, T, J, 3)
|
||||
|
||||
root_rotation = rotmat[:, :, 0] # (N, T, 3, 3)
|
||||
global_pos = torch.matmul(root_rotation.unsqueeze(2), pos.unsqueeze(-1)).squeeze(-1) # (N, T, J, 3)
|
||||
|
||||
output_ = {'pos': pos,
|
||||
'velocity': velocity,
|
||||
'global_xform': matrix_to_rotation_6d(global_xform),
|
||||
'angular': angular,
|
||||
'root_orient': matrix_to_rotation_6d(root_rotation),
|
||||
'rotmat': rotmat}
|
||||
|
||||
return output_
|
||||
|
||||
|
||||
def ablation_study(path, subtask):
|
||||
|
||||
num_motions = 100
|
||||
print(os.path.join(path, 'TCDHands', '*.npz'))
|
||||
|
||||
test_data = load_data(sorted(glob.glob(os.path.join('./outputs/generative_copy/TCDHands', '*.npz')))) # test data TCDHands
|
||||
val_data = load_data(sorted(glob.glob(os.path.join('./outputs/generative_copy/GRAB', '*.npz')))) # test data TCDHands
|
||||
train_data = load_data(sorted(glob.glob(os.path.join('./outputs/generative_copy/SAMP', '*.npz')))) # test data TCDHands
|
||||
synthetic_data = synthetic_trajectory(N_motions=num_motions)
|
||||
|
||||
pca_aa = joblib.load("data/pca.pkl")
|
||||
pca_mean = torch.tensor(pca_aa.mean_).to("cuda")
|
||||
pca_cov = torch.tensor(pca_aa.get_covariance()).to("cuda")
|
||||
pca_sv = torch.tensor(pca_aa.singular_values_).to("cuda")
|
||||
pca_pc = torch.tensor(pca_aa.components_).to("cuda")
|
||||
|
||||
random_samples_low_dim = np.random.randn(num_motions, pca_aa.n_components_)
|
||||
pca_data = pca_aa.inverse_transform(random_samples_low_dim)
|
||||
|
||||
gmm_aa = MaxMixturePrior()
|
||||
gmm_data = gmm_aa.gmm.sample(num_motions)
|
||||
|
||||
identity = torch.eye(3)
|
||||
identity = identity.view(1, 1, 3, 3).repeat(num_motions, args.data.clip_length, 1, 1)
|
||||
|
||||
# convert to tensor
|
||||
gmm_data = torch.tensor(gmm_data[0].reshape(num_motions, args.data.clip_length, -1, 3))
|
||||
gmm_data = torch.cat((identity.unsqueeze(2), axis_angle_to_matrix(gmm_data)), dim=2)
|
||||
gmm_rot6d = matrix_to_rotation_6d(gmm_data).reshape(num_motions * args.data.clip_length, -1, 6)
|
||||
|
||||
# convert to tensor
|
||||
pca_data = torch.tensor(pca_data.reshape(num_motions, args.data.clip_length, -1, 3))
|
||||
pca_data = torch.cat((identity.unsqueeze(2), axis_angle_to_matrix(pca_data)), dim=2)
|
||||
pca_rot6d = matrix_to_rotation_6d(pca_data).reshape(num_motions * args.data.clip_length, -1, 6)
|
||||
|
||||
gmm_pos, gmm_xform = fk(gmm_rot6d.to(fk.device))
|
||||
pca_pos, pca_xform = fk(pca_rot6d.to(fk.device))
|
||||
|
||||
gmm_xform = matrix_to_rotation_6d(gmm_xform[:, :, :3, :3])
|
||||
pca_xform = matrix_to_rotation_6d(pca_xform[:, :, :3, :3])
|
||||
|
||||
gmm_velocity = estimate_linear_velocity(gmm_pos.unsqueeze(0), dt=1.0 / args.data.fps).squeeze(0)
|
||||
pca_velocity = estimate_linear_velocity(pca_pos.unsqueeze(0), dt=1.0 / args.data.fps).squeeze(0)
|
||||
|
||||
gmm_angular = estimate_angular_velocity(gmm_data.unsqueeze(0), dt=1.0 / args.data.fps).squeeze(0)
|
||||
pca_angular = estimate_angular_velocity(pca_data.unsqueeze(0), dt=1.0 / args.data.fps).squeeze(0)
|
||||
|
||||
gmm_dict = {'pos': gmm_pos.reshape(num_motions, args.data.clip_length, -1, 3),
|
||||
'global_xform': gmm_xform.reshape(num_motions, args.data.clip_length, -1, 6),
|
||||
'velocity': gmm_velocity.reshape(num_motions, args.data.clip_length, -1, 3),
|
||||
'angular': gmm_angular}
|
||||
|
||||
pca_dict = {'pos': pca_pos.reshape(num_motions, args.data.clip_length, -1, 3),
|
||||
'global_xform': pca_xform.reshape(num_motions, args.data.clip_length, -1, 6),
|
||||
'velocity': pca_velocity.reshape(num_motions, args.data.clip_length, -1, 3),
|
||||
'angular': pca_angular}
|
||||
|
||||
metrics = dict()
|
||||
|
||||
def f(data, num_motions, key):
|
||||
div = evaluate_diversity(data)
|
||||
if key in metrics.keys():
|
||||
pass
|
||||
else:
|
||||
metrics[key] = {'DIV': [div]}
|
||||
|
||||
print('Motion Synthesis:')
|
||||
div_test = f(test_data, num_motions, 'TEST_DATA (TDCHANDS)')
|
||||
div_val = f(val_data, num_motions, 'VAL_DATA (SAMP)')
|
||||
div_train = f(train_data, num_motions, 'TRAIN_DATA (GRAB)')
|
||||
div_synt = f(synthetic_data, num_motions, 'SYNTHETIC')
|
||||
div_gmm = f(gmm_dict, num_motions, 'GMM')
|
||||
div_pca = f(pca_dict, num_motions, 'PCA')
|
||||
|
||||
df = pd.DataFrame.from_dict(metrics)
|
||||
df.to_csv(os.path.join(path, f'diversity.csv'), index=False)
|
||||
|
||||
|
||||
|
||||
def comparison(path, subtask):
|
||||
gt_data = load_data(sorted(glob.glob(os.path.join(path, 'gt', '*.npz'))))
|
||||
|
||||
file_path = os.path.join(path, subtask)
|
||||
|
||||
metrics = dict()
|
||||
if subtask == 'recon':
|
||||
humor_data = load_data(sorted(glob.glob(os.path.join(file_path, 'humor', '*.npz'))))
|
||||
hm_vae_data = load_data(sorted(glob.glob(os.path.join(file_path, 'hm-vae', '*.npz'))))
|
||||
nemf_data = load_data(sorted(glob.glob(os.path.join(file_path, 'nemf', '*.npz'))))
|
||||
|
||||
metrics['HUMOR'] = evaluate_recon_error(humor_data, gt_data)
|
||||
metrics['HMVAE'] = evaluate_recon_error(hm_vae_data, gt_data)
|
||||
metrics['NEMF'] = evaluate_recon_error(nemf_data, gt_data)
|
||||
print('Motion Reconstruction:')
|
||||
print(json.dumps(metrics, sort_keys=True, indent=4))
|
||||
|
||||
df = pd.DataFrame.from_dict(metrics)
|
||||
df.to_csv(os.path.join(file_path, f'{subtask}.csv'))
|
||||
elif subtask == 'synth':
|
||||
humor_data = load_data(sorted(glob.glob(os.path.join(file_path, 'humor', '.*.npz'))), padding=True)
|
||||
hm_vae_data = load_data(sorted(glob.glob(os.path.join(file_path, 'hm-vae', '*.npz'))), padding=True)
|
||||
nemf_data = load_data(sorted(glob.glob(os.path.join(file_path, 'nemf', '*.npz'))), padding=True)
|
||||
|
||||
num_motions = 100
|
||||
|
||||
def f(data, num_motions, key):
|
||||
fid = evaluate_fid(data, gt_data, num_motions)
|
||||
div = evaluate_diversity(data)
|
||||
index = list(range(data['global_pos'].shape[0]))
|
||||
index = random.sample(index, num_motions)
|
||||
s = foot_skating(data['global_pos'][index])
|
||||
if key in metrics.keys():
|
||||
metrics[key]['FID'].append(fid)
|
||||
metrics[key]['DIV'].append(div)
|
||||
metrics[key]['FOOT'].append(s)
|
||||
else:
|
||||
metrics[key] = {'FID': [fid], 'DIV': [div], 'FOOT': [s]}
|
||||
|
||||
print('Motion Synthesis:')
|
||||
for replication in range(20):
|
||||
print(f'==================== Replication {replication} ====================')
|
||||
f(humor_data, num_motions, 'HUMOR')
|
||||
f(hm_vae_data, num_motions, 'HMVAE')
|
||||
f(nemf_data, num_motions, 'NEMF')
|
||||
|
||||
for key in metrics.keys():
|
||||
df = pd.DataFrame.from_dict(metrics[key])
|
||||
df.to_csv(os.path.join(file_path, f'{subtask}_{key}.csv'), index=False)
|
||||
else:
|
||||
raise RuntimeError(f'unsupported evaluation task for comparison: {subtask}')
|
||||
|
||||
|
||||
def inbetween_benchmarks(path, keyframes=False):
|
||||
trans_lengths = [5, 10, 15, 20] if keyframes else [10, 20, 30]
|
||||
|
||||
gt_data = load_data(sorted(glob.glob(os.path.join(path, 'gt', '*.npz'))))
|
||||
s = foot_skating(gt_data['global_pos'])
|
||||
print(f'average foot skating of the real motions: {s:.4f}')
|
||||
|
||||
def f(data, key, mask):
|
||||
fid = evaluate_fid(data, gt_data, mask=mask)
|
||||
s = foot_skating(data['global_pos'], mask=mask)
|
||||
if key in statistics.keys():
|
||||
statistics[key]['FID'].append(fid)
|
||||
statistics[key]['FOOT'].append(s)
|
||||
else:
|
||||
statistics[key] = {'FID': [fid], 'FOOT': [s]}
|
||||
|
||||
file_path = os.path.join(path, 'keyframe') if keyframes else os.path.join(path, 'clip')
|
||||
statistics = {}
|
||||
for t in trans_lengths:
|
||||
print(f'==================== Length {t} ====================')
|
||||
if keyframes:
|
||||
mask = None # don't mask data for keyframe inbetweening
|
||||
else:
|
||||
end_frame = (args.data.clip_length - t) // 2
|
||||
mask = list(range(end_frame)) + list(range(end_frame + t, args.data.clip_length))
|
||||
print(f'mask: {mask}')
|
||||
|
||||
slerp_data = load_data(sorted(glob.glob(os.path.join(file_path, 'slerp', f'{t}_frames', '*.npz'))))
|
||||
f(slerp_data, 'SLERP', mask)
|
||||
|
||||
if not keyframes:
|
||||
inertia_data = load_data(sorted(glob.glob(os.path.join(file_path, 'inertialization', f'{t}_frames', '*.npz'))))
|
||||
f(inertia_data, 'INERTIALIZATION', mask)
|
||||
|
||||
rmi_data = load_data(sorted(glob.glob(os.path.join(file_path, 'rmi', f'{t}_frames', '*.npz'))))
|
||||
f(rmi_data, 'RMI', mask)
|
||||
|
||||
hm_vae_data = load_data(sorted(glob.glob(os.path.join(file_path, 'hm-vae', f'{t}_frames', '*.npz'))))
|
||||
f(hm_vae_data, 'HMVAE', mask)
|
||||
|
||||
nemf_data = load_data(sorted(glob.glob(os.path.join(file_path, 'nemf', f'{t}_frames', '*.npz'))))
|
||||
f(nemf_data, 'NEMF', mask)
|
||||
|
||||
for key in statistics.keys():
|
||||
df = pd.DataFrame.from_dict(statistics[key])
|
||||
df.to_csv(os.path.join(file_path, f'inbetween_{key}.csv'))
|
||||
|
||||
|
||||
def compute_joint_accel(joint_seq, fps):
|
||||
''' Magnitude of joint accelerations for joint_seq : B x T x J x 3 '''
|
||||
joint_accel = joint_seq[:, :-2] - (2 * joint_seq[:, 1:-1]) + joint_seq[:, 2:]
|
||||
DATA_h = 1.0 / fps
|
||||
joint_accel = joint_accel / ((DATA_h**2))
|
||||
joint_accel_mag = torch.norm(joint_accel, dim=-1)
|
||||
|
||||
return joint_accel, joint_accel_mag
|
||||
|
||||
|
||||
def evaluate_super_sampling(path):
|
||||
fps_30_data = load_data(sorted(glob.glob(os.path.join(path, '30fps/*.npz'))))
|
||||
fps_60_data = load_data(sorted(glob.glob(os.path.join(path, '60fps/*.npz'))))
|
||||
fps_120_data = load_data(sorted(glob.glob(os.path.join(path, '120fps/*.npz'))))
|
||||
fps_240_data = load_data(sorted(glob.glob(os.path.join(path, '240fps/*.npz'))))
|
||||
|
||||
smoothness = dict()
|
||||
joint_vel = estimate_linear_velocity(fps_30_data['global_pos'], dt=1.0 / 30.0)
|
||||
joint_vel_mag = torch.norm(joint_vel, dim=-1).mean()
|
||||
smoothness['30FPS'] = [c2c(joint_vel_mag * 100.0)]
|
||||
joint_vel = estimate_linear_velocity(fps_60_data['global_pos'], dt=1.0 / 60.0)
|
||||
joint_vel_mag = torch.norm(joint_vel, dim=-1).mean()
|
||||
smoothness['60FPS'] = [c2c(joint_vel_mag * 100.0)]
|
||||
joint_vel = estimate_linear_velocity(fps_120_data['global_pos'], dt=1.0 / 120.0)
|
||||
joint_vel_mag = torch.norm(joint_vel, dim=-1).mean()
|
||||
smoothness['120FPS'] = [c2c(joint_vel_mag * 100.0)]
|
||||
joint_vel = estimate_linear_velocity(fps_240_data['global_pos'], dt=1.0 / 240.0)
|
||||
joint_vel_mag = torch.norm(joint_vel, dim=-1).mean()
|
||||
smoothness['240FPS'] = [c2c(joint_vel_mag * 100.0)]
|
||||
|
||||
df = pd.DataFrame.from_dict(smoothness)
|
||||
df.to_csv(os.path.join(path, f'smoothness.csv'))
|
||||
|
||||
|
||||
def evaluate_smoothness_test(path):
|
||||
smoothness = dict()
|
||||
|
||||
for fps in [30, 60]:
|
||||
smoothness[f'{fps}FPS'] = dict()
|
||||
for L in range(1, 23, 2):
|
||||
data = load_data([os.path.join(path, f'L{L}', f'test_{fps}fps.npz')])
|
||||
# _, joint_accel_mag = compute_joint_accel(data['global_pos'], fps=fps)
|
||||
# smoothness[f'{fps}FPS'][L] = c2c(joint_accel_mag.mean() * 100.0)
|
||||
joint_vel = estimate_linear_velocity(data['global_pos'], dt=1.0 / fps)
|
||||
joint_vel_mag = torch.norm(joint_vel, dim=-1).mean()
|
||||
smoothness[f'{fps}FPS'][L] = c2c(joint_vel_mag.mean() * 100.0)
|
||||
|
||||
df = pd.DataFrame.from_dict(smoothness)
|
||||
df.to_csv(os.path.join(path, f'smoothness.csv'))
|
||||
|
||||
import matplotlib.pyplot as plt
|
||||
plt.rcParams['text.usetex'] = True
|
||||
|
||||
smoothness_30fps = np.array(list(smoothness['30FPS'].values()))
|
||||
smoothness_60fps = np.array(list(smoothness['60FPS'].values()))
|
||||
|
||||
plt.plot(np.arange(1, 23, 2), smoothness_30fps, marker='o', label='30fps')
|
||||
plt.plot(np.arange(1, 23, 2), smoothness_60fps, marker='o', label='60fps')
|
||||
plt.legend()
|
||||
plt.grid(linestyle='--')
|
||||
plt.xticks(np.arange(1, 23, 2))
|
||||
plt.xlabel(r'$L$')
|
||||
# plt.yscale('log')
|
||||
plt.ylabel('Mean Per-Joint Velocity')
|
||||
plt.savefig(os.path.join(path, 'smoothness.pdf'))
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
from argparse import ArgumentParser
|
||||
|
||||
parser = ArgumentParser()
|
||||
parser.add_argument('--config', type=str, default='evaluation.yaml', help='name of the configutation file')
|
||||
parser.add_argument('--task', type=str, help='name of application tasks', required=True)
|
||||
parser.add_argument('--load_path', type=str, help='path to load results for evaluation', required=True)
|
||||
|
||||
opt = parser.parse_args()
|
||||
args = Arguments('./configs', filename=opt.config)
|
||||
|
||||
torch.set_default_dtype(torch.float32)
|
||||
|
||||
torch.manual_seed(0)
|
||||
np.random.seed(0)
|
||||
random.seed(0)
|
||||
|
||||
ngpu = 1
|
||||
if args.multi_gpu is True:
|
||||
ngpu = torch.cuda.device_count()
|
||||
if ngpu == 1:
|
||||
args.multi_gpu = False
|
||||
print(f'Number of GPUs: {ngpu}')
|
||||
|
||||
prior = Architecture(args, ngpu)
|
||||
prior.load(optimal=True)
|
||||
prior.eval()
|
||||
fk = ForwardKinematicsLayer(args)
|
||||
|
||||
if opt.task == 'ablation_study':
|
||||
print('Evaluating ablation study ...')
|
||||
ablation_study(opt.load_path, subtask='synth')
|
||||
elif opt.task == 'comparison':
|
||||
print('Evaluating comparison with other methods ...')
|
||||
comparison(opt.load_path, subtask='recon')
|
||||
comparison(opt.load_path, subtask='synth')
|
||||
elif opt.task == 'inbetween_benchmark':
|
||||
print('Evaluating motion in-betweening ...')
|
||||
inbetween_benchmarks(opt.load_path, keyframes=False)
|
||||
inbetween_benchmarks(opt.load_path, keyframes=True)
|
||||
elif opt.task == 'super_sampling':
|
||||
print('Evaluating super sampling ...')
|
||||
evaluate_super_sampling(opt.load_path)
|
||||
elif opt.task == 'smoothness_test':
|
||||
print('Evaluating smoothness testing ...')
|
||||
evaluate_smoothness_test(opt.load_path)
|
||||
else:
|
||||
raise NotImplementedError(f'unknown evaluation: {opt.task}')
|
||||
+1731
File diff suppressed because it is too large
Load Diff
+1119
File diff suppressed because it is too large
Load Diff
+150
@@ -0,0 +1,150 @@
|
||||
batch_size: 1
|
||||
bvh_viz: false
|
||||
|
||||
data:
|
||||
clip_length: 128
|
||||
fps: 30
|
||||
gender: neutral
|
||||
normalize:
|
||||
- pos
|
||||
- velocity
|
||||
- global_xform
|
||||
- angular
|
||||
- height
|
||||
- root_orient
|
||||
- root_vel
|
||||
root_transform: true
|
||||
up: z
|
||||
|
||||
root: None
|
||||
dataset_dir: _DATA/hmp_model
|
||||
dtw_loss: true
|
||||
geodesic_loss: true
|
||||
|
||||
global_prior:
|
||||
activation: tanh
|
||||
in_channels: 6
|
||||
kernel_size: 4
|
||||
temporal_scale: 8
|
||||
z_dim: 256
|
||||
|
||||
initialization: true
|
||||
is_train: false
|
||||
keypoint_source: mmpose
|
||||
|
||||
l1_loss: true
|
||||
lambda_angle: 1
|
||||
lambda_dtw: 0.5
|
||||
lambda_kl: 0.0001
|
||||
|
||||
local_prior:
|
||||
activation: tanh
|
||||
channel_base: 15
|
||||
extra_conv: 0
|
||||
kernel_size: 4
|
||||
num_layers: 4
|
||||
padding_mode: reflect
|
||||
skeleton_dist: 2
|
||||
skeleton_pool: mean
|
||||
temporal_scale: 8
|
||||
upsampling: linear
|
||||
use_residual_blocks: true
|
||||
z_dim: 1024
|
||||
|
||||
log: false
|
||||
multi_gpu: false
|
||||
|
||||
nemf:
|
||||
bandwidth: 7
|
||||
global_output: 6
|
||||
global_z: 256
|
||||
hidden_neuron: 1024
|
||||
local_output: 144
|
||||
local_z: 1024
|
||||
norm_layer: true
|
||||
siren: false
|
||||
skip_connection: true
|
||||
nsubject: 1
|
||||
num_workers: 0
|
||||
opt_betas: false
|
||||
output_trans: false
|
||||
overlap_len: 16
|
||||
pretrained_gmp: false
|
||||
save_dir: _DATA/hmp_model
|
||||
|
||||
scheduler:
|
||||
gamma: 0.7
|
||||
step_size: 200
|
||||
|
||||
stg1:
|
||||
betas_prior: 10
|
||||
keypoint_source: mmpose
|
||||
lambda_batch_cs: 0
|
||||
lambda_init_z_prior: 0
|
||||
lambda_j3d_smooth: 0
|
||||
lambda_motion_prior: 0
|
||||
lambda_pose_prior: 0
|
||||
lambda_orient: 0
|
||||
lambda_orient_smooth: 1
|
||||
lambda_pos: 0
|
||||
lambda_reproj: 0.05
|
||||
lambda_bio: 1
|
||||
lambda_consistency: 0
|
||||
lambda_rot: 0
|
||||
lambda_rot_smooth: 0
|
||||
lambda_trans: 10
|
||||
lambda_trans_smooth: 10
|
||||
lr: 0.05
|
||||
niters: 400
|
||||
opt_params:
|
||||
- betas
|
||||
- trans
|
||||
- root_orient
|
||||
|
||||
stg2:
|
||||
betas_prior: 10
|
||||
keypoint_source: mmpose
|
||||
lambda_batch_cs: 0
|
||||
lambda_init_z_prior: 0
|
||||
lambda_j3d_smooth: 0
|
||||
lambda_motion_prior: 200
|
||||
lambda_pose_prior: 0
|
||||
lambda_orient: 0
|
||||
lambda_orient_smooth: 1
|
||||
lambda_pos: 0
|
||||
lambda_reproj: 0.05
|
||||
lambda_bio: 1
|
||||
lambda_consistency: 0
|
||||
lambda_rot: 0
|
||||
lambda_rot_smooth: 0
|
||||
lambda_trans: 10
|
||||
lambda_trans_smooth: 10
|
||||
lr: 0.05
|
||||
niters: 400
|
||||
opt_params:
|
||||
- betas
|
||||
- trans
|
||||
- root_orient
|
||||
- z_l
|
||||
|
||||
stg3:
|
||||
betas_prior: 0
|
||||
lambda_batch_cs: 0
|
||||
lambda_init_z_prior: 0
|
||||
lambda_j3d_smooth: 0
|
||||
lambda_motion_prior: 0
|
||||
lambda_orient: 0
|
||||
lambda_orient_smooth: 0
|
||||
lambda_pos: 0
|
||||
lambda_reproj: 0
|
||||
lambda_bio: 0
|
||||
lambda_rot: 0
|
||||
lambda_rot_smooth: 0
|
||||
lambda_trans: 0
|
||||
lambda_trans_smooth: 0
|
||||
lr: 0.0
|
||||
niters: 0
|
||||
opt_params:
|
||||
- root_orient
|
||||
- z_l
|
||||
verbose: false
|
||||
@@ -0,0 +1,675 @@
|
||||
import operator
|
||||
|
||||
import numpy as np
|
||||
import numpy.core.umath as ut # (for numpy 1.25)
|
||||
# import numpy.core.umath_tests as ut
|
||||
|
||||
import AnimationStructure
|
||||
from Quaternions import Quaternions
|
||||
|
||||
|
||||
class Animation:
|
||||
"""
|
||||
Animation is a numpy-like wrapper for animation data
|
||||
|
||||
Animation data consists of several arrays consisting
|
||||
of F frames and J joints.
|
||||
|
||||
The animation is specified by
|
||||
|
||||
rotations : (F, J) Quaternions | Joint Rotations
|
||||
positions : (F, J, 3) ndarray | Joint Positions
|
||||
|
||||
The base pose is specified by
|
||||
|
||||
orients : (J) Quaternions | Joint Orientations
|
||||
offsets : (J, 3) ndarray | Joint Offsets
|
||||
|
||||
And the skeletal structure is specified by
|
||||
|
||||
parents : (J) ndarray | Joint Parents
|
||||
"""
|
||||
|
||||
def __init__(self, rotations, positions, orients, offsets, parents):
|
||||
|
||||
self.rotations = rotations
|
||||
self.positions = positions
|
||||
self.orients = orients
|
||||
self.offsets = offsets
|
||||
self.parents = parents
|
||||
|
||||
def __op__(self, op, other):
|
||||
return Animation(
|
||||
op(self.rotations, other.rotations),
|
||||
op(self.positions, other.positions),
|
||||
op(self.orients, other.orients),
|
||||
op(self.offsets, other.offsets),
|
||||
op(self.parents, other.parents))
|
||||
|
||||
def __iop__(self, op, other):
|
||||
self.rotations = op(self.roations, other.rotations)
|
||||
self.positions = op(self.roations, other.positions)
|
||||
self.orients = op(self.orients, other.orients)
|
||||
self.offsets = op(self.offsets, other.offsets)
|
||||
self.parents = op(self.parents, other.parents)
|
||||
return self
|
||||
|
||||
def __sop__(self, op):
|
||||
return Animation(
|
||||
op(self.rotations),
|
||||
op(self.positions),
|
||||
op(self.orients),
|
||||
op(self.offsets),
|
||||
op(self.parents))
|
||||
|
||||
def __add__(self, other): return self.__op__(operator.add, other)
|
||||
def __sub__(self, other): return self.__op__(operator.sub, other)
|
||||
def __mul__(self, other): return self.__op__(operator.mul, other)
|
||||
def __div__(self, other): return self.__op__(operator.div, other)
|
||||
|
||||
def __abs__(self): return self.__sop__(operator.abs)
|
||||
def __neg__(self): return self.__sop__(operator.neg)
|
||||
|
||||
def __iadd__(self, other): return self.__iop__(operator.iadd, other)
|
||||
def __isub__(self, other): return self.__iop__(operator.isub, other)
|
||||
def __imul__(self, other): return self.__iop__(operator.imul, other)
|
||||
def __idiv__(self, other): return self.__iop__(operator.idiv, other)
|
||||
|
||||
def __len__(self): return len(self.rotations)
|
||||
|
||||
def __getitem__(self, k):
|
||||
if isinstance(k, tuple):
|
||||
return Animation(
|
||||
self.rotations[k],
|
||||
self.positions[k],
|
||||
self.orients[k[1:]],
|
||||
self.offsets[k[1:]],
|
||||
self.parents[k[1:]])
|
||||
else:
|
||||
return Animation(
|
||||
self.rotations[k],
|
||||
self.positions[k],
|
||||
self.orients,
|
||||
self.offsets,
|
||||
self.parents)
|
||||
|
||||
def __setitem__(self, k, v):
|
||||
if isinstance(k, tuple):
|
||||
self.rotations.__setitem__(k, v.rotations)
|
||||
self.positions.__setitem__(k, v.positions)
|
||||
self.orients.__setitem__(k[1:], v.orients)
|
||||
self.offsets.__setitem__(k[1:], v.offsets)
|
||||
self.parents.__setitem__(k[1:], v.parents)
|
||||
else:
|
||||
self.rotations.__setitem__(k, v.rotations)
|
||||
self.positions.__setitem__(k, v.positions)
|
||||
self.orients.__setitem__(k, v.orients)
|
||||
self.offsets.__setitem__(k, v.offsets)
|
||||
self.parents.__setitem__(k, v.parents)
|
||||
|
||||
@property
|
||||
def shape(self): return (self.rotations.shape[0], self.rotations.shape[1])
|
||||
|
||||
def copy(self): return Animation(
|
||||
self.rotations.copy(), self.positions.copy(),
|
||||
self.orients.copy(), self.offsets.copy(),
|
||||
self.parents.copy())
|
||||
|
||||
def repeat(self, *args, **kw):
|
||||
return Animation(
|
||||
self.rotations.repeat(*args, **kw),
|
||||
self.positions.repeat(*args, **kw),
|
||||
self.orients, self.offsets, self.parents)
|
||||
|
||||
def ravel(self):
|
||||
return np.hstack([
|
||||
self.rotations.log().ravel(),
|
||||
self.positions.ravel(),
|
||||
self.orients.log().ravel(),
|
||||
self.offsets.ravel()])
|
||||
|
||||
@classmethod
|
||||
def unravel(clas, anim, shape, parents):
|
||||
nf, nj = shape
|
||||
rotations = anim[nf*nj*0:nf*nj*3]
|
||||
positions = anim[nf*nj*3:nf*nj*6]
|
||||
orients = anim[nf*nj*6+nj*0:nf*nj*6+nj*3]
|
||||
offsets = anim[nf*nj*6+nj*3:nf*nj*6+nj*6]
|
||||
return cls(
|
||||
Quaternions.exp(rotations), positions,
|
||||
Quaternions.exp(orients), offsets,
|
||||
parents.copy())
|
||||
|
||||
|
||||
""" Maya Interaction """
|
||||
|
||||
|
||||
def load_to_maya(anim, names=None, radius=0.5):
|
||||
"""
|
||||
Load Animation Object into Maya as Joint Skeleton
|
||||
loads each frame as a new keyfame in maya.
|
||||
|
||||
If the animation is too slow or too fast perhaps
|
||||
the framerate needs adjusting before being loaded
|
||||
such that it matches the maya scene framerate.
|
||||
|
||||
|
||||
Parameters
|
||||
----------
|
||||
|
||||
anim : Animation
|
||||
Animation to load into Scene
|
||||
|
||||
names : [str]
|
||||
Optional list of Joint names for Skeleton
|
||||
|
||||
Returns
|
||||
-------
|
||||
|
||||
List of Maya Joint Nodes loaded into scene
|
||||
"""
|
||||
|
||||
import pymel.core as pm
|
||||
|
||||
joints = []
|
||||
frames = range(1, len(anim)+1)
|
||||
|
||||
if names is None:
|
||||
names = ["joint_" + str(i) for i in range(len(anim.parents))]
|
||||
|
||||
for i, offset, orient, parent, name in zip(range(len(anim.offsets)), anim.offsets, anim.orients, anim.parents, names):
|
||||
|
||||
if parent < 0:
|
||||
pm.select(d=True)
|
||||
else:
|
||||
pm.select(joints[parent])
|
||||
|
||||
joint = pm.joint(n=name, p=offset, relative=True, radius=radius)
|
||||
joint.setOrientation([orient[1], orient[2], orient[3], orient[0]])
|
||||
|
||||
curvex = pm.nodetypes.AnimCurveTA(n=name + "_rotateX")
|
||||
curvey = pm.nodetypes.AnimCurveTA(n=name + "_rotateY")
|
||||
curvez = pm.nodetypes.AnimCurveTA(n=name + "_rotateZ")
|
||||
|
||||
jrotations = (-Quaternions(orient[np.newaxis]) * anim.rotations[:, i]).euler()
|
||||
curvex.addKeys(frames, jrotations[:, 0])
|
||||
curvey.addKeys(frames, jrotations[:, 1])
|
||||
curvez.addKeys(frames, jrotations[:, 2])
|
||||
|
||||
pm.connectAttr(curvex.output, joint.rotateX)
|
||||
pm.connectAttr(curvey.output, joint.rotateY)
|
||||
pm.connectAttr(curvez.output, joint.rotateZ)
|
||||
|
||||
offsetx = pm.nodetypes.AnimCurveTU(n=name + "_translateX")
|
||||
offsety = pm.nodetypes.AnimCurveTU(n=name + "_translateY")
|
||||
offsetz = pm.nodetypes.AnimCurveTU(n=name + "_translateZ")
|
||||
|
||||
offsetx.addKeys(frames, anim.positions[:, i, 0])
|
||||
offsety.addKeys(frames, anim.positions[:, i, 1])
|
||||
offsetz.addKeys(frames, anim.positions[:, i, 2])
|
||||
|
||||
pm.connectAttr(offsetx.output, joint.translateX)
|
||||
pm.connectAttr(offsety.output, joint.translateY)
|
||||
pm.connectAttr(offsetz.output, joint.translateZ)
|
||||
|
||||
joints.append(joint)
|
||||
|
||||
return joints
|
||||
|
||||
|
||||
def load_from_maya(root, start, end):
|
||||
"""
|
||||
Load Animation Object from Maya Joint Skeleton
|
||||
|
||||
Parameters
|
||||
----------
|
||||
|
||||
root : PyNode
|
||||
Root Joint of Maya Skeleton
|
||||
|
||||
start, end : int, int
|
||||
Start and End frame index of Maya Animation
|
||||
|
||||
Returns
|
||||
-------
|
||||
|
||||
animation : Animation
|
||||
Loaded animation from maya
|
||||
|
||||
names : [str]
|
||||
Joint names from maya
|
||||
"""
|
||||
|
||||
import pymel.core as pm
|
||||
|
||||
original_time = pm.currentTime(q=True)
|
||||
pm.currentTime(start)
|
||||
|
||||
""" Build Structure """
|
||||
|
||||
names, parents = AnimationStructure.load_from_maya(root)
|
||||
descendants = AnimationStructure.descendants_list(parents)
|
||||
orients = Quaternions.id(len(names))
|
||||
offsets = np.array([pm.xform(j, q=True, translation=True) for j in names])
|
||||
|
||||
for j, name in enumerate(names):
|
||||
scale = pm.xform(pm.PyNode(name), q=True, scale=True, relative=True)
|
||||
if len(descendants[j]) == 0:
|
||||
continue
|
||||
offsets[descendants[j]] *= scale
|
||||
|
||||
""" Load Animation """
|
||||
|
||||
eulers = np.zeros((end-start, len(names), 3))
|
||||
positions = np.zeros((end-start, len(names), 3))
|
||||
rotations = Quaternions.id((end-start, len(names)))
|
||||
|
||||
for i in range(end-start):
|
||||
|
||||
pm.currentTime(start+i+1, u=True)
|
||||
|
||||
scales = {}
|
||||
|
||||
for j, name, parent in zip(range(len(names)), names, parents):
|
||||
|
||||
node = pm.PyNode(name)
|
||||
|
||||
if i == 0 and pm.hasAttr(node, 'jointOrient'):
|
||||
ort = node.getOrientation()
|
||||
orients[j] = Quaternions(np.array([ort[3], ort[0], ort[1], ort[2]]))
|
||||
|
||||
if pm.hasAttr(node, 'rotate'):
|
||||
eulers[i, j] = np.radians(pm.xform(node, q=True, rotation=True))
|
||||
if pm.hasAttr(node, 'translate'):
|
||||
positions[i, j] = pm.xform(node, q=True, translation=True)
|
||||
if pm.hasAttr(node, 'scale'):
|
||||
scales[j] = pm.xform(node, q=True, scale=True, relative=True)
|
||||
|
||||
for j in scales:
|
||||
if len(descendants[j]) == 0:
|
||||
continue
|
||||
positions[i, descendants[j]] *= scales[j]
|
||||
|
||||
positions[i, 0] = pm.xform(root, q=True, translation=True, worldSpace=True)
|
||||
|
||||
rotations = orients[np.newaxis] * Quaternions.from_euler(eulers, order='xyz', world=True)
|
||||
|
||||
""" Done """
|
||||
|
||||
pm.currentTime(original_time)
|
||||
|
||||
return Animation(rotations, positions, orients, offsets, parents), names
|
||||
|
||||
# local transformation matrices
|
||||
|
||||
|
||||
def transforms_local(anim):
|
||||
"""
|
||||
Computes Animation Local Transforms
|
||||
|
||||
As well as a number of other uses this can
|
||||
be used to compute global joint transforms,
|
||||
which in turn can be used to compete global
|
||||
joint positions
|
||||
|
||||
Parameters
|
||||
----------
|
||||
|
||||
anim : Animation
|
||||
Input animation
|
||||
|
||||
Returns
|
||||
-------
|
||||
|
||||
transforms : (F, J, 4, 4) ndarray
|
||||
|
||||
For each frame F, joint local
|
||||
transforms for each joint J
|
||||
"""
|
||||
|
||||
transforms = anim.rotations.transforms()
|
||||
transforms = np.concatenate([transforms, np.zeros(transforms.shape[:2] + (3, 1))], axis=-1)
|
||||
transforms = np.concatenate([transforms, np.zeros(transforms.shape[:2] + (1, 4))], axis=-2)
|
||||
# the last column is filled with the joint positions!
|
||||
transforms[:, :, 0:3, 3] = anim.positions
|
||||
transforms[:, :, 3:4, 3] = 1.0
|
||||
return transforms
|
||||
|
||||
|
||||
def transforms_multiply(t0s, t1s):
|
||||
"""
|
||||
Transforms Multiply
|
||||
|
||||
Multiplies two arrays of animation transforms
|
||||
|
||||
Parameters
|
||||
----------
|
||||
|
||||
t0s, t1s : (F, J, 4, 4) ndarray
|
||||
Two arrays of transforms
|
||||
for each frame F and each
|
||||
joint J
|
||||
|
||||
Returns
|
||||
-------
|
||||
|
||||
transforms : (F, J, 4, 4) ndarray
|
||||
Array of transforms for each
|
||||
frame F and joint J multiplied
|
||||
together
|
||||
"""
|
||||
|
||||
return ut.matrix_multiply(t0s, t1s)
|
||||
|
||||
|
||||
def transforms_inv(ts):
|
||||
fts = ts.reshape(-1, 4, 4)
|
||||
fts = np.array(list(map(lambda x: np.linalg.inv(x), fts)))
|
||||
return fts.reshape(ts.shape)
|
||||
|
||||
|
||||
def transforms_blank(anim):
|
||||
"""
|
||||
Blank Transforms
|
||||
|
||||
Parameters
|
||||
----------
|
||||
|
||||
anim : Animation
|
||||
Input animation
|
||||
|
||||
Returns
|
||||
-------
|
||||
|
||||
transforms : (F, J, 4, 4) ndarray
|
||||
Array of identity transforms for
|
||||
each frame F and joint J
|
||||
"""
|
||||
|
||||
ts = np.zeros(anim.shape + (4, 4))
|
||||
ts[:, :, 0, 0] = 1.0
|
||||
ts[:, :, 1, 1] = 1.0
|
||||
ts[:, :, 2, 2] = 1.0
|
||||
ts[:, :, 3, 3] = 1.0
|
||||
return ts
|
||||
|
||||
# global transformation matrices
|
||||
|
||||
|
||||
def transforms_global(anim):
|
||||
"""
|
||||
Global Animation Transforms
|
||||
|
||||
This relies on joint ordering
|
||||
being incremental. That means a joint
|
||||
J1 must not be a ancestor of J0 if
|
||||
J0 appears before J1 in the joint
|
||||
ordering.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
|
||||
anim : Animation
|
||||
Input animation
|
||||
|
||||
Returns
|
||||
------
|
||||
|
||||
transforms : (F, J, 4, 4) ndarray
|
||||
Array of global transforms for
|
||||
each frame F and joint J
|
||||
"""
|
||||
|
||||
joints = np.arange(anim.shape[1])
|
||||
parents = np.arange(anim.shape[1])
|
||||
locals = transforms_local(anim)
|
||||
globals = transforms_blank(anim)
|
||||
|
||||
globals[:, 0] = locals[:, 0]
|
||||
|
||||
for i in range(1, anim.shape[1]):
|
||||
globals[:, i] = transforms_multiply(globals[:, anim.parents[i]], locals[:, i])
|
||||
|
||||
return globals
|
||||
|
||||
# !!! useful!
|
||||
|
||||
|
||||
def positions_global(anim):
|
||||
"""
|
||||
Global Joint Positions
|
||||
|
||||
Given an animation compute the global joint
|
||||
positions at at every frame
|
||||
|
||||
Parameters
|
||||
----------
|
||||
|
||||
anim : Animation
|
||||
Input animation
|
||||
|
||||
Returns
|
||||
-------
|
||||
|
||||
positions : (F, J, 3) ndarray
|
||||
Positions for every frame F
|
||||
and joint position J
|
||||
"""
|
||||
|
||||
# get the last column -- corresponding to the coordinates
|
||||
positions = transforms_global(anim)[:, :, :, 3]
|
||||
return positions[:, :, :3] / positions[:, :, 3, np.newaxis]
|
||||
|
||||
|
||||
""" Rotations """
|
||||
|
||||
|
||||
def rotations_global(anim):
|
||||
"""
|
||||
Global Animation Rotations
|
||||
|
||||
This relies on joint ordering
|
||||
being incremental. That means a joint
|
||||
J1 must not be a ancestor of J0 if
|
||||
J0 appears before J1 in the joint
|
||||
ordering.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
|
||||
anim : Animation
|
||||
Input animation
|
||||
|
||||
Returns
|
||||
-------
|
||||
|
||||
points : (F, J) Quaternions
|
||||
global rotations for every frame F
|
||||
and joint J
|
||||
"""
|
||||
|
||||
joints = np.arange(anim.shape[1])
|
||||
parents = np.arange(anim.shape[1])
|
||||
locals = anim.rotations
|
||||
globals = Quaternions.id(anim.shape)
|
||||
|
||||
globals[:, 0] = locals[:, 0]
|
||||
|
||||
for i in range(1, anim.shape[1]):
|
||||
globals[:, i] = globals[:, anim.parents[i]] * locals[:, i]
|
||||
|
||||
return globals
|
||||
|
||||
|
||||
def rotations_parents_global(anim):
|
||||
rotations = rotations_global(anim)
|
||||
rotations = rotations[:, anim.parents]
|
||||
rotations[:, 0] = Quaternions.id(len(anim))
|
||||
return rotations
|
||||
|
||||
|
||||
def rotations_load_to_maya(rotations, positions, names=None):
|
||||
"""
|
||||
Load Rotations into Maya
|
||||
|
||||
Loads a Quaternions array into the scene
|
||||
via the representation of axis
|
||||
|
||||
Parameters
|
||||
----------
|
||||
|
||||
rotations : (F, J) Quaternions
|
||||
array of rotations to load
|
||||
into the scene where
|
||||
F = number of frames
|
||||
J = number of joints
|
||||
|
||||
positions : (F, J, 3) ndarray
|
||||
array of positions to load
|
||||
rotation axis at where:
|
||||
F = number of frames
|
||||
J = number of joints
|
||||
|
||||
names : [str]
|
||||
List of joint names
|
||||
|
||||
Returns
|
||||
-------
|
||||
|
||||
maxies : Group
|
||||
Grouped Maya Node of all Axis nodes
|
||||
"""
|
||||
|
||||
import pymel.core as pm
|
||||
|
||||
if names is None:
|
||||
names = ["joint_" + str(i) for i in range(rotations.shape[1])]
|
||||
|
||||
maxis = []
|
||||
frames = range(1, len(positions)+1)
|
||||
for i, name in enumerate(names):
|
||||
|
||||
name = name + "_axis"
|
||||
axis = pm.group(
|
||||
pm.curve(p=[(0, 0, 0), (1, 0, 0)], d=1, n=name+'_axis_x'),
|
||||
pm.curve(p=[(0, 0, 0), (0, 1, 0)], d=1, n=name+'_axis_y'),
|
||||
pm.curve(p=[(0, 0, 0), (0, 0, 1)], d=1, n=name+'_axis_z'),
|
||||
n=name)
|
||||
|
||||
axis.rotatePivot.set((0, 0, 0))
|
||||
axis.scalePivot.set((0, 0, 0))
|
||||
axis.childAtIndex(0).overrideEnabled.set(1)
|
||||
axis.childAtIndex(0).overrideColor.set(13)
|
||||
axis.childAtIndex(1).overrideEnabled.set(1)
|
||||
axis.childAtIndex(1).overrideColor.set(14)
|
||||
axis.childAtIndex(2).overrideEnabled.set(1)
|
||||
axis.childAtIndex(2).overrideColor.set(15)
|
||||
|
||||
curvex = pm.nodetypes.AnimCurveTA(n=name + "_rotateX")
|
||||
curvey = pm.nodetypes.AnimCurveTA(n=name + "_rotateY")
|
||||
curvez = pm.nodetypes.AnimCurveTA(n=name + "_rotateZ")
|
||||
|
||||
arotations = rotations[:, i].euler()
|
||||
curvex.addKeys(frames, arotations[:, 0])
|
||||
curvey.addKeys(frames, arotations[:, 1])
|
||||
curvez.addKeys(frames, arotations[:, 2])
|
||||
|
||||
pm.connectAttr(curvex.output, axis.rotateX)
|
||||
pm.connectAttr(curvey.output, axis.rotateY)
|
||||
pm.connectAttr(curvez.output, axis.rotateZ)
|
||||
|
||||
offsetx = pm.nodetypes.AnimCurveTU(n=name + "_translateX")
|
||||
offsety = pm.nodetypes.AnimCurveTU(n=name + "_translateY")
|
||||
offsetz = pm.nodetypes.AnimCurveTU(n=name + "_translateZ")
|
||||
|
||||
offsetx.addKeys(frames, positions[:, i, 0])
|
||||
offsety.addKeys(frames, positions[:, i, 1])
|
||||
offsetz.addKeys(frames, positions[:, i, 2])
|
||||
|
||||
pm.connectAttr(offsetx.output, axis.translateX)
|
||||
pm.connectAttr(offsety.output, axis.translateY)
|
||||
pm.connectAttr(offsetz.output, axis.translateZ)
|
||||
|
||||
maxis.append(axis)
|
||||
|
||||
return pm.group(*maxis, n='RotationAnimation')
|
||||
|
||||
|
||||
""" Offsets & Orients """
|
||||
|
||||
|
||||
def orients_global(anim):
|
||||
|
||||
joints = np.arange(anim.shape[1])
|
||||
parents = np.arange(anim.shape[1])
|
||||
locals = anim.orients
|
||||
globals = Quaternions.id(anim.shape[1])
|
||||
|
||||
globals[:, 0] = locals[:, 0]
|
||||
|
||||
for i in range(1, anim.shape[1]):
|
||||
globals[:, i] = globals[:, anim.parents[i]] * locals[:, i]
|
||||
|
||||
return globals
|
||||
|
||||
|
||||
def offsets_transforms_local(anim):
|
||||
|
||||
transforms = anim.orients[np.newaxis].transforms()
|
||||
transforms = np.concatenate([transforms, np.zeros(transforms.shape[:2] + (3, 1))], axis=-1)
|
||||
transforms = np.concatenate([transforms, np.zeros(transforms.shape[:2] + (1, 4))], axis=-2)
|
||||
transforms[:, :, 0:3, 3] = anim.offsets[np.newaxis]
|
||||
transforms[:, :, 3:4, 3] = 1.0
|
||||
return transforms
|
||||
|
||||
|
||||
def offsets_transforms_global(anim):
|
||||
|
||||
joints = np.arange(anim.shape[1])
|
||||
parents = np.arange(anim.shape[1])
|
||||
locals = offsets_transforms_local(anim)
|
||||
globals = transforms_blank(anim)
|
||||
|
||||
globals[:, 0] = locals[:, 0]
|
||||
|
||||
for i in range(1, anim.shape[1]):
|
||||
globals[:, i] = transforms_multiply(globals[:, anim.parents[i]], locals[:, i])
|
||||
|
||||
return globals
|
||||
|
||||
|
||||
def offsets_global(anim):
|
||||
offsets = offsets_transforms_global(anim)[:, :, :, 3]
|
||||
return offsets[0, :, :3] / offsets[0, :, 3, np.newaxis]
|
||||
|
||||
|
||||
""" Lengths """
|
||||
|
||||
|
||||
def offset_lengths(anim):
|
||||
return np.sum(anim.offsets[1:]**2.0, axis=1)**0.5
|
||||
|
||||
|
||||
def position_lengths(anim):
|
||||
return np.sum(anim.positions[:, 1:]**2.0, axis=2)**0.5
|
||||
|
||||
|
||||
""" Skinning """
|
||||
|
||||
|
||||
def skin(anim, rest, weights, mesh, maxjoints=4):
|
||||
|
||||
full_transforms = transforms_multiply(
|
||||
transforms_global(anim),
|
||||
transforms_inv(transforms_global(rest[0:1])))
|
||||
|
||||
weightids = np.argsort(-weights, axis=1)[:, :maxjoints]
|
||||
weightvls = np.array(list(map(lambda w, i: w[i], weights, weightids)))
|
||||
weightvls = weightvls / weightvls.sum(axis=1)[..., np.newaxis]
|
||||
|
||||
verts = np.hstack([mesh, np.ones((len(mesh), 1))])
|
||||
verts = verts[np.newaxis, :, np.newaxis, :, np.newaxis]
|
||||
verts = transforms_multiply(full_transforms[:, weightids], verts)
|
||||
verts = (verts[:, :, :, :3] / verts[:, :, :, 3:4])[:, :, :, :, 0]
|
||||
|
||||
return np.sum(weightvls[np.newaxis, :, :, np.newaxis] * verts, axis=2)
|
||||
@@ -0,0 +1,451 @@
|
||||
import numpy as np
|
||||
import scipy.sparse as sparse
|
||||
import Animation as Animation
|
||||
|
||||
|
||||
""" Maya Functions """
|
||||
|
||||
def load_from_maya(root):
|
||||
"""
|
||||
Load joint parents and names from maya
|
||||
|
||||
Parameters
|
||||
----------
|
||||
|
||||
root : PyNode
|
||||
Root Maya Node
|
||||
|
||||
Returns
|
||||
-------
|
||||
|
||||
(names, parents) : ([str], (J) ndarray)
|
||||
List of joint names and array
|
||||
of indices representing the parent
|
||||
joint for each joint J.
|
||||
|
||||
Joint index -1 is used to represent
|
||||
that there is no parent joint
|
||||
"""
|
||||
|
||||
import pymel.core as pm
|
||||
|
||||
names = []
|
||||
parents = []
|
||||
|
||||
def unload_joint(j, parents, par):
|
||||
|
||||
id = len(names)
|
||||
names.append(j)
|
||||
parents.append(par)
|
||||
|
||||
children = [c for c in j.getChildren() if
|
||||
isinstance(c, pm.nt.Transform) and
|
||||
not isinstance(c, pm.nt.Constraint) and
|
||||
not any(pm.listRelatives(c, s=True)) and
|
||||
(any(pm.listRelatives(c, ad=True, ap=False, type='joint')) or isinstance(c, pm.nt.Joint))]
|
||||
|
||||
map(lambda c: unload_joint(c, parents, id), children)
|
||||
|
||||
unload_joint(root, parents, -1)
|
||||
|
||||
return (names, parents)
|
||||
|
||||
|
||||
""" Family Functions """
|
||||
|
||||
def joints(parents):
|
||||
"""
|
||||
Parameters
|
||||
----------
|
||||
|
||||
parents : (J) ndarray
|
||||
parents array
|
||||
|
||||
Returns
|
||||
-------
|
||||
|
||||
joints : (J) ndarray
|
||||
Array of joint indices
|
||||
"""
|
||||
return np.arange(len(parents), dtype=int)
|
||||
|
||||
def joints_list(parents):
|
||||
"""
|
||||
Parameters
|
||||
----------
|
||||
|
||||
parents : (J) ndarray
|
||||
parents array
|
||||
|
||||
Returns
|
||||
-------
|
||||
|
||||
joints : [ndarray]
|
||||
List of arrays of joint idices for
|
||||
each joint
|
||||
"""
|
||||
return list(joints(parents)[:,np.newaxis])
|
||||
|
||||
def parents_list(parents):
|
||||
"""
|
||||
Parameters
|
||||
----------
|
||||
|
||||
parents : (J) ndarray
|
||||
parents array
|
||||
|
||||
Returns
|
||||
-------
|
||||
|
||||
parents : [ndarray]
|
||||
List of arrays of joint idices for
|
||||
the parents of each joint
|
||||
"""
|
||||
return list(parents[:,np.newaxis])
|
||||
|
||||
|
||||
def children_list(parents):
|
||||
"""
|
||||
Parameters
|
||||
----------
|
||||
|
||||
parents : (J) ndarray
|
||||
parents array
|
||||
|
||||
Returns
|
||||
-------
|
||||
|
||||
children : [ndarray]
|
||||
List of arrays of joint indices for
|
||||
the children of each joint
|
||||
"""
|
||||
|
||||
def joint_children(i):
|
||||
return [j for j, p in enumerate(parents) if p == i]
|
||||
|
||||
return list(map(lambda j: np.array(joint_children(j)), joints(parents)))
|
||||
|
||||
|
||||
def descendants_list(parents):
|
||||
"""
|
||||
Parameters
|
||||
----------
|
||||
|
||||
parents : (J) ndarray
|
||||
parents array
|
||||
|
||||
Returns
|
||||
-------
|
||||
|
||||
descendants : [ndarray]
|
||||
List of arrays of joint idices for
|
||||
the descendants of each joint
|
||||
"""
|
||||
|
||||
children = children_list(parents)
|
||||
|
||||
def joint_descendants(i):
|
||||
return sum([joint_descendants(j) for j in children[i]], list(children[i]))
|
||||
|
||||
return list(map(lambda j: np.array(joint_descendants(j)), joints(parents)))
|
||||
|
||||
|
||||
def ancestors_list(parents):
|
||||
"""
|
||||
Parameters
|
||||
----------
|
||||
|
||||
parents : (J) ndarray
|
||||
parents array
|
||||
|
||||
Returns
|
||||
-------
|
||||
|
||||
ancestors : [ndarray]
|
||||
List of arrays of joint idices for
|
||||
the ancestors of each joint
|
||||
"""
|
||||
|
||||
decendants = descendants_list(parents)
|
||||
|
||||
def joint_ancestors(i):
|
||||
return [j for j in joints(parents) if i in decendants[j]]
|
||||
|
||||
return list(map(lambda j: np.array(joint_ancestors(j)), joints(parents)))
|
||||
|
||||
|
||||
""" Mask Functions """
|
||||
|
||||
def mask(parents, filter):
|
||||
"""
|
||||
Constructs a Mask for a give filter
|
||||
|
||||
A mask is a (J, J) ndarray truth table for a given
|
||||
condition over J joints. For example there
|
||||
may be a mask specifying if a joint N is a
|
||||
child of another joint M.
|
||||
|
||||
This could be constructed into a mask using
|
||||
`m = mask(parents, children_list)` and the condition
|
||||
of childhood tested using `m[N, M]`.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
|
||||
parents : (J) ndarray
|
||||
parents array
|
||||
|
||||
filter : (J) ndarray -> [ndarray]
|
||||
function that outputs a list of arrays
|
||||
of joint indices for some condition
|
||||
|
||||
Returns
|
||||
-------
|
||||
|
||||
mask : (N, N) ndarray
|
||||
boolean truth table of given condition
|
||||
"""
|
||||
m = np.zeros((len(parents), len(parents))).astype(bool)
|
||||
jnts = joints(parents)
|
||||
fltr = filter(parents)
|
||||
for i,f in enumerate(fltr): m[i,:] = np.any(jnts[:,np.newaxis] == f[np.newaxis,:], axis=1)
|
||||
return m
|
||||
|
||||
def joints_mask(parents): return np.eye(len(parents)).astype(bool)
|
||||
def children_mask(parents): return mask(parents, children_list)
|
||||
def parents_mask(parents): return mask(parents, parents_list)
|
||||
def descendants_mask(parents): return mask(parents, descendants_list)
|
||||
def ancestors_mask(parents): return mask(parents, ancestors_list)
|
||||
|
||||
""" Search Functions """
|
||||
|
||||
def joint_chain_ascend(parents, start, end):
|
||||
chain = []
|
||||
while start != end:
|
||||
chain.append(start)
|
||||
start = parents[start]
|
||||
chain.append(end)
|
||||
return np.array(chain, dtype=int)
|
||||
|
||||
|
||||
""" Constraints """
|
||||
|
||||
def constraints(anim, **kwargs):
|
||||
"""
|
||||
Constraint list for Animation
|
||||
|
||||
This constraint list can be used in the
|
||||
VerletParticle solver to constrain
|
||||
a animation global joint positions.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
|
||||
anim : Animation
|
||||
Input animation
|
||||
|
||||
masses : (F, J) ndarray
|
||||
Optional list of masses
|
||||
for joints J across frames F
|
||||
defaults to weighting by
|
||||
vertical height
|
||||
|
||||
Returns
|
||||
-------
|
||||
|
||||
constraints : [(int, int, (F, J) ndarray, (F, J) ndarray, (F, J) ndarray)]
|
||||
A list of constraints in the format:
|
||||
(Joint1, Joint2, Masses1, Masses2, Lengths)
|
||||
|
||||
"""
|
||||
|
||||
masses = kwargs.pop('masses', None)
|
||||
|
||||
children = children_list(anim.parents)
|
||||
constraints = []
|
||||
|
||||
points_offsets = Animation.offsets_global(anim)
|
||||
points = Animation.positions_global(anim)
|
||||
|
||||
if masses is None:
|
||||
masses = 1.0 / (0.1 + np.absolute(points_offsets[:,1]))
|
||||
masses = masses[np.newaxis].repeat(len(anim), axis=0)
|
||||
|
||||
for j in xrange(anim.shape[1]):
|
||||
|
||||
""" Add constraints between all joints and their children """
|
||||
for c0 in children[j]:
|
||||
|
||||
dists = np.sum((points[:, c0] - points[:, j])**2.0, axis=1)**0.5
|
||||
constraints.append((c0, j, masses[:,c0], masses[:,j], dists))
|
||||
|
||||
""" Add constraints between all children of joint """
|
||||
for c1 in children[j]:
|
||||
if c0 == c1: continue
|
||||
|
||||
dists = np.sum((points[:, c0] - points[:, c1])**2.0, axis=1)**0.5
|
||||
constraints.append((c0, c1, masses[:,c0], masses[:,c1], dists))
|
||||
|
||||
return constraints
|
||||
|
||||
""" Graph Functions """
|
||||
|
||||
def graph(anim):
|
||||
"""
|
||||
Generates a weighted adjacency matrix
|
||||
using local joint distances along
|
||||
the skeletal structure.
|
||||
|
||||
Joints which are not connected
|
||||
are assigned the weight `0`.
|
||||
|
||||
Joints which actually have zero distance
|
||||
between them, but are still connected, are
|
||||
perturbed by some minimal amount.
|
||||
|
||||
The output of this routine can be used
|
||||
with the `scipy.sparse.csgraph`
|
||||
routines for graph analysis.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
|
||||
anim : Animation
|
||||
input animation
|
||||
|
||||
Returns
|
||||
-------
|
||||
|
||||
graph : (N, N) ndarray
|
||||
weight adjacency matrix using
|
||||
local distances along the
|
||||
skeletal structure from joint
|
||||
N to joint M. If joints are not
|
||||
directly connected are assigned
|
||||
the weight `0`.
|
||||
"""
|
||||
|
||||
graph = np.zeros(anim.shape[1], anim.shape[1])
|
||||
lengths = np.sum(anim.offsets**2.0, axis=1)**0.5 + 0.001
|
||||
|
||||
for i,p in enumerate(anim.parents):
|
||||
if p == -1: continue
|
||||
graph[i,p] = lengths[p]
|
||||
graph[p,i] = lengths[p]
|
||||
|
||||
return graph
|
||||
|
||||
|
||||
def distances(anim):
|
||||
"""
|
||||
Generates a distance matrix for
|
||||
pairwise joint distances along
|
||||
the skeletal structure
|
||||
|
||||
Parameters
|
||||
----------
|
||||
|
||||
anim : Animation
|
||||
input animation
|
||||
|
||||
Returns
|
||||
-------
|
||||
|
||||
distances : (N, N) ndarray
|
||||
array of pairwise distances
|
||||
along skeletal structure
|
||||
from some joint N to some
|
||||
joint M
|
||||
"""
|
||||
|
||||
distances = np.zeros((anim.shape[1], anim.shape[1]))
|
||||
generated = distances.copy().astype(bool)
|
||||
|
||||
joint_lengths = np.sum(anim.offsets**2.0, axis=1)**0.5
|
||||
joint_children = children_list(anim)
|
||||
joint_parents = parents_list(anim)
|
||||
|
||||
def find_distance(distances, generated, prev, i, j):
|
||||
|
||||
""" If root, identity, or already generated, return """
|
||||
if j == -1: return (0.0, True)
|
||||
if j == i: return (0.0, True)
|
||||
if generated[i,j]: return (distances[i,j], True)
|
||||
|
||||
""" Find best distances along parents and children """
|
||||
par_dists = [(joint_lengths[j], find_distance(distances, generated, j, i, p)) for p in joint_parents[j] if p != prev]
|
||||
out_dists = [(joint_lengths[c], find_distance(distances, generated, j, i, c)) for c in joint_children[j] if c != prev]
|
||||
|
||||
""" Check valid distance and not dead end """
|
||||
par_dists = [a + d for (a, (d, f)) in par_dists if f]
|
||||
out_dists = [a + d for (a, (d, f)) in out_dists if f]
|
||||
|
||||
""" All dead ends """
|
||||
if (out_dists + par_dists) == []: return (0.0, False)
|
||||
|
||||
""" Get minimum path """
|
||||
dist = min(out_dists + par_dists)
|
||||
distances[i,j] = dist; distances[j,i] = dist
|
||||
generated[i,j] = True; generated[j,i] = True
|
||||
|
||||
for i in xrange(anim.shape[1]):
|
||||
for j in xrange(anim.shape[1]):
|
||||
find_distance(distances, generated, -1, i, j)
|
||||
|
||||
return distances
|
||||
|
||||
def edges(parents):
|
||||
"""
|
||||
Animation structure edges
|
||||
|
||||
Parameters
|
||||
----------
|
||||
|
||||
parents : (J) ndarray
|
||||
parents array
|
||||
|
||||
Returns
|
||||
-------
|
||||
|
||||
edges : (M, 2) ndarray
|
||||
array of pairs where each
|
||||
pair contains two indices of a joints
|
||||
which corrisponds to an edge in the
|
||||
joint structure going from parent to child.
|
||||
"""
|
||||
|
||||
return np.array(list(zip(parents, joints(parents)))[1:])
|
||||
|
||||
|
||||
def incidence(parents):
|
||||
"""
|
||||
Incidence Matrix
|
||||
|
||||
Parameters
|
||||
----------
|
||||
|
||||
parents : (J) ndarray
|
||||
parents array
|
||||
|
||||
Returns
|
||||
-------
|
||||
|
||||
incidence : (N, M) ndarray
|
||||
|
||||
Matrix of N joint positions by
|
||||
M edges which each entry is either
|
||||
1 or -1 and multiplication by the
|
||||
joint positions returns the an
|
||||
array of vectors along each edge
|
||||
of the structure
|
||||
"""
|
||||
|
||||
es = edges(parents)
|
||||
|
||||
inc = np.zeros((len(parents)-1, len(parents))).astype(np.int)
|
||||
for i, e in enumerate(es):
|
||||
inc[i,e[0]] = 1
|
||||
inc[i,e[1]] = -1
|
||||
|
||||
return inc.T
|
||||
|
||||
+471
@@ -0,0 +1,471 @@
|
||||
from Quaternions import Quaternions
|
||||
from Animation import Animation
|
||||
import re
|
||||
import numpy as np
|
||||
import sys
|
||||
|
||||
|
||||
channelmap = {
|
||||
'Xrotation': 'x',
|
||||
'Yrotation': 'y',
|
||||
'Zrotation': 'z'
|
||||
}
|
||||
|
||||
channelmap_inv = {
|
||||
'x': 'Xrotation',
|
||||
'y': 'Yrotation',
|
||||
'z': 'Zrotation',
|
||||
}
|
||||
|
||||
ordermap = {
|
||||
'x': 0,
|
||||
'y': 1,
|
||||
'z': 2,
|
||||
}
|
||||
|
||||
|
||||
def load(filename, start=None, end=None, order=None, world=False):
|
||||
"""
|
||||
Reads a BVH file and constructs an animation
|
||||
|
||||
Parameters
|
||||
----------
|
||||
filename: str
|
||||
File to be opened
|
||||
|
||||
start : int
|
||||
Optional Starting Frame
|
||||
|
||||
end : int
|
||||
Optional Ending Frame
|
||||
|
||||
order : str
|
||||
Optional Specifier for joint order.
|
||||
Given as string E.G 'xyz', 'zxy'
|
||||
|
||||
world : bool
|
||||
If set to true euler angles are applied
|
||||
together in world space rather than local
|
||||
space
|
||||
|
||||
Returns
|
||||
-------
|
||||
|
||||
(animation, joint_names, frametime)
|
||||
Tuple of loaded animation and joint names
|
||||
"""
|
||||
|
||||
f = open(filename, "r")
|
||||
|
||||
i = 0
|
||||
active = -1
|
||||
end_site = False
|
||||
|
||||
names = []
|
||||
orients = Quaternions.id(0)
|
||||
offsets = np.array([]).reshape((0, 3))
|
||||
parents = np.array([], dtype=int)
|
||||
|
||||
for line in f:
|
||||
|
||||
if "HIERARCHY" in line:
|
||||
continue
|
||||
if "MOTION" in line:
|
||||
continue
|
||||
|
||||
rmatch = re.match(r"ROOT (\w+)", line)
|
||||
if rmatch:
|
||||
names.append(rmatch.group(1))
|
||||
offsets = np.append(offsets, np.array([[0, 0, 0]]), axis=0)
|
||||
orients.qs = np.append(orients.qs, np.array([[1, 0, 0, 0]]), axis=0)
|
||||
parents = np.append(parents, active)
|
||||
active = (len(parents)-1)
|
||||
continue
|
||||
|
||||
if "{" in line:
|
||||
continue
|
||||
|
||||
if "}" in line:
|
||||
if end_site:
|
||||
end_site = False
|
||||
else:
|
||||
active = parents[active]
|
||||
continue
|
||||
|
||||
offmatch = re.match(r"\s*OFFSET\s+([\-\d\.e]+)\s+([\-\d\.e]+)\s+([\-\d\.e]+)", line)
|
||||
if offmatch:
|
||||
if not end_site:
|
||||
offsets[active] = np.array([list(map(float, offmatch.groups()))])
|
||||
continue
|
||||
|
||||
chanmatch = re.match(r"\s*CHANNELS\s+(\d+)", line)
|
||||
if chanmatch:
|
||||
channels = int(chanmatch.group(1))
|
||||
if order is None:
|
||||
channelis = 0 if channels == 3 else 3
|
||||
channelie = 3 if channels == 3 else 6
|
||||
parts = line.split()[2+channelis:2+channelie]
|
||||
if any([p not in channelmap for p in parts]):
|
||||
continue
|
||||
order = "".join([channelmap[p] for p in parts])
|
||||
continue
|
||||
|
||||
jmatch = re.match("\s*JOINT\s+(\w+)", line)
|
||||
if jmatch:
|
||||
names.append(jmatch.group(1))
|
||||
offsets = np.append(offsets, np.array([[0, 0, 0]]), axis=0)
|
||||
orients.qs = np.append(orients.qs, np.array([[1, 0, 0, 0]]), axis=0)
|
||||
parents = np.append(parents, active)
|
||||
active = (len(parents)-1)
|
||||
continue
|
||||
|
||||
if "End Site" in line:
|
||||
end_site = True
|
||||
continue
|
||||
|
||||
fmatch = re.match("\s*Frames:\s+(\d+)", line)
|
||||
if fmatch:
|
||||
if start and end:
|
||||
fnum = (end - start)-1
|
||||
else:
|
||||
fnum = int(fmatch.group(1))
|
||||
jnum = len(parents)
|
||||
# result: [fnum, J, 3]
|
||||
positions = offsets[np.newaxis].repeat(fnum, axis=0)
|
||||
# result: [fnum, len(orients), 3]
|
||||
rotations = np.zeros((fnum, len(orients), 3))
|
||||
continue
|
||||
|
||||
fmatch = re.match("\s*Frame Time:\s+([\d\.]+)", line)
|
||||
if fmatch:
|
||||
frametime = float(fmatch.group(1))
|
||||
continue
|
||||
|
||||
if (start and end) and (i < start or i >= end-1):
|
||||
i += 1
|
||||
continue
|
||||
|
||||
dmatch = line.strip().split()
|
||||
if dmatch:
|
||||
data_block = np.array(list(map(float, dmatch)))
|
||||
N = len(parents)
|
||||
fi = i - start if start else i
|
||||
if channels == 3:
|
||||
# This should be root positions[0:1] & all rotations
|
||||
positions[fi, 0:1] = data_block[0:3]
|
||||
rotations[fi, :] = data_block[3:].reshape(N, 3)
|
||||
elif channels == 6:
|
||||
data_block = data_block.reshape(N, 6)
|
||||
# fill in all positions
|
||||
positions[fi, :] = data_block[:, 0:3]
|
||||
rotations[fi, :] = data_block[:, 3:6]
|
||||
elif channels == 9:
|
||||
positions[fi, 0] = data_block[0:3]
|
||||
data_block = data_block[3:].reshape(N-1, 9)
|
||||
rotations[fi, 1:] = data_block[:, 3:6]
|
||||
positions[fi, 1:] += data_block[:, 0:3] * data_block[:, 6:9]
|
||||
else:
|
||||
raise Exception("Too many channels! %i" % channels)
|
||||
|
||||
i += 1
|
||||
|
||||
f.close()
|
||||
|
||||
rotations = Quaternions.from_euler(np.radians(rotations), order=order, world=world)
|
||||
|
||||
return (Animation(rotations, positions, orients, offsets, parents), names, frametime)
|
||||
|
||||
|
||||
def load_bfa(filename, start=None, end=None, order=None, world=False):
|
||||
"""
|
||||
Reads a BVH file and constructs an animation
|
||||
|
||||
!!! Read from bfa, will replace the end sites of arms by two joints (w/ unit rotation)
|
||||
|
||||
Parameters
|
||||
----------
|
||||
filename: str
|
||||
File to be opened
|
||||
|
||||
start : int
|
||||
Optional Starting Frame
|
||||
|
||||
end : int
|
||||
Optional Ending Frame
|
||||
|
||||
order : str
|
||||
Optional Specifier for joint order.
|
||||
Given as string E.G 'xyz', 'zxy'
|
||||
|
||||
world : bool
|
||||
If set to true euler angles are applied
|
||||
together in world space rather than local
|
||||
space
|
||||
|
||||
Returns
|
||||
-------
|
||||
|
||||
(animation, joint_names, frametime)
|
||||
Tuple of loaded animation and joint names
|
||||
"""
|
||||
|
||||
f = open(filename, "r")
|
||||
|
||||
i = 0
|
||||
active = -1
|
||||
end_site = False
|
||||
|
||||
hand_idx = [9, 14]
|
||||
|
||||
names = []
|
||||
orients = Quaternions.id(0)
|
||||
offsets = np.array([]).reshape((0, 3))
|
||||
parents = np.array([], dtype=int)
|
||||
|
||||
for line in f:
|
||||
|
||||
if "HIERARCHY" in line:
|
||||
continue
|
||||
if "MOTION" in line:
|
||||
continue
|
||||
|
||||
rmatch = re.match(r"ROOT (\w+)", line)
|
||||
if rmatch:
|
||||
names.append(rmatch.group(1))
|
||||
offsets = np.append(offsets, np.array([[0, 0, 0]]), axis=0)
|
||||
orients.qs = np.append(orients.qs, np.array([[1, 0, 0, 0]]), axis=0)
|
||||
parents = np.append(parents, active)
|
||||
active = (len(parents)-1)
|
||||
continue
|
||||
|
||||
if "{" in line:
|
||||
continue
|
||||
|
||||
if "}" in line:
|
||||
if end_site:
|
||||
end_site = False
|
||||
else:
|
||||
active = parents[active]
|
||||
continue
|
||||
|
||||
offmatch = re.match(r"\s*OFFSET\s+([\-\d\.e]+)\s+([\-\d\.e]+)\s+([\-\d\.e]+)", line)
|
||||
if offmatch:
|
||||
if not end_site:
|
||||
offsets[active] = np.array([list(map(float, offmatch.groups()))])
|
||||
"""
|
||||
else:
|
||||
print("active = ", active)
|
||||
if active in hand_idx:
|
||||
offsets[active] = np.array([list(map(float, offmatch.groups()))])
|
||||
"""
|
||||
continue
|
||||
|
||||
chanmatch = re.match(r"\s*CHANNELS\s+(\d+)", line)
|
||||
if chanmatch:
|
||||
channels = int(chanmatch.group(1))
|
||||
if order is None:
|
||||
channelis = 0 if channels == 3 else 3
|
||||
channelie = 3 if channels == 3 else 6
|
||||
parts = line.split()[2+channelis:2+channelie]
|
||||
if any([p not in channelmap for p in parts]):
|
||||
continue
|
||||
order = "".join([channelmap[p] for p in parts])
|
||||
continue
|
||||
|
||||
jmatch = re.match("\s*JOINT\s+(\w+)", line)
|
||||
if jmatch:
|
||||
names.append(jmatch.group(1))
|
||||
offsets = np.append(offsets, np.array([[0, 0, 0]]), axis=0)
|
||||
orients.qs = np.append(orients.qs, np.array([[1, 0, 0, 0]]), axis=0)
|
||||
parents = np.append(parents, active)
|
||||
active = (len(parents)-1)
|
||||
continue
|
||||
|
||||
if "End Site" in line:
|
||||
if active + 1 in hand_idx:
|
||||
print("parent:", names[-1])
|
||||
name = "LeftHandIndex" if active + 1 == hand_idx[0] else "RightHandIndex"
|
||||
names.append(name)
|
||||
offsets = np.append(offsets, np.array([[0, 0, 0]]), axis=0)
|
||||
orients.qs = np.append(orients.qs, np.array([[1, 0, 0, 0]]), axis=0)
|
||||
parents = np.append(parents, active)
|
||||
active = (len(parents)-1)
|
||||
else:
|
||||
end_site = True
|
||||
continue
|
||||
|
||||
fmatch = re.match("\s*Frames:\s+(\d+)", line)
|
||||
if fmatch:
|
||||
if start and end:
|
||||
fnum = (end - start)-1
|
||||
else:
|
||||
fnum = int(fmatch.group(1))
|
||||
jnum = len(parents)
|
||||
# result: [fnum, J, 3]
|
||||
positions = offsets[np.newaxis].repeat(fnum, axis=0)
|
||||
# result: [fnum, len(orients), 3]
|
||||
rotations = np.zeros((fnum, len(orients), 3))
|
||||
continue
|
||||
|
||||
fmatch = re.match("\s*Frame Time:\s+([\d\.]+)", line)
|
||||
if fmatch:
|
||||
frametime = float(fmatch.group(1))
|
||||
continue
|
||||
|
||||
if (start and end) and (i < start or i >= end-1):
|
||||
i += 1
|
||||
continue
|
||||
|
||||
dmatch = line.strip().split()
|
||||
if dmatch:
|
||||
data_block = np.array(list(map(float, dmatch)))
|
||||
N = len(parents)
|
||||
fi = i - start if start else i
|
||||
if channels == 3:
|
||||
# This should be root positions[0:1] & all rotations
|
||||
positions[fi, 0:1] = data_block[0:3]
|
||||
tmp = data_block[3:].reshape(N - 2, 3)
|
||||
tmp = np.concatenate([tmp[:hand_idx[0]],
|
||||
np.array([[0, 0, 0]]),
|
||||
tmp[hand_idx[0]: hand_idx[1] - 1],
|
||||
np.array([[0, 0, 0]]),
|
||||
tmp[hand_idx[1] - 1:]], axis=0)
|
||||
rotations[fi, :] = tmp.reshape(N, 3)
|
||||
elif channels == 6:
|
||||
data_block = data_block.reshape(N, 6)
|
||||
# fill in all positions
|
||||
positions[fi, :] = data_block[:, 0:3]
|
||||
rotations[fi, :] = data_block[:, 3:6]
|
||||
elif channels == 9:
|
||||
positions[fi, 0] = data_block[0:3]
|
||||
data_block = data_block[3:].reshape(N-1, 9)
|
||||
rotations[fi, 1:] = data_block[:, 3:6]
|
||||
positions[fi, 1:] += data_block[:, 0:3] * data_block[:, 6:9]
|
||||
else:
|
||||
raise Exception("Too many channels! %i" % channels)
|
||||
|
||||
i += 1
|
||||
|
||||
f.close()
|
||||
|
||||
rotations = Quaternions.from_euler(np.radians(rotations), order=order, world=world)
|
||||
|
||||
return (Animation(rotations, positions, orients, offsets, parents), names, frametime)
|
||||
|
||||
|
||||
def save(filename, anim, names=None, frametime=1.0/24.0, order='zyx', positions=False, orients=True):
|
||||
"""
|
||||
Saves an Animation to file as BVH
|
||||
|
||||
Parameters
|
||||
----------
|
||||
filename: str
|
||||
File to be saved to
|
||||
|
||||
anim : Animation
|
||||
Animation to save
|
||||
|
||||
names : [str]
|
||||
List of joint names
|
||||
|
||||
order : str
|
||||
Optional Specifier for joint order.
|
||||
Given as string E.G 'xyz', 'zxy'
|
||||
|
||||
frametime : float
|
||||
Optional Animation Frame time
|
||||
|
||||
positions : bool
|
||||
Optional specfier to save bone
|
||||
positions for each frame
|
||||
|
||||
orients : bool
|
||||
Multiply joint orients to the rotations
|
||||
before saving.
|
||||
|
||||
"""
|
||||
|
||||
if names is None:
|
||||
names = ["joint_" + str(i) for i in range(len(anim.parents))]
|
||||
|
||||
with open(filename, 'w') as f:
|
||||
|
||||
t = ""
|
||||
f.write("%sHIERARCHY\n" % t)
|
||||
f.write("%sROOT %s\n" % (t, names[0]))
|
||||
f.write("%s{\n" % t)
|
||||
t += '\t'
|
||||
|
||||
f.write("%sOFFSET %f %f %f\n" % (t, anim.offsets[0, 0], anim.offsets[0, 1], anim.offsets[0, 2]))
|
||||
f.write("%sCHANNELS 6 Xposition Yposition Zposition %s %s %s \n" %
|
||||
(t, channelmap_inv[order[0]], channelmap_inv[order[1]], channelmap_inv[order[2]]))
|
||||
|
||||
for i in range(anim.shape[1]):
|
||||
if anim.parents[i] == 0:
|
||||
t = save_joint(f, anim, names, t, i, order=order, positions=positions)
|
||||
|
||||
t = t[:-1]
|
||||
f.write("%s}\n" % t)
|
||||
|
||||
f.write("MOTION\n")
|
||||
f.write("Frames: %i\n" % anim.shape[0])
|
||||
f.write("Frame Time: %f\n" % frametime)
|
||||
|
||||
# if orients:
|
||||
# rots = np.degrees((-anim.orients[np.newaxis] * anim.rotations).euler(order=order[::-1]))
|
||||
# else:
|
||||
# rots = np.degrees(anim.rotations.euler(order=order[::-1]))
|
||||
rots = np.degrees(anim.rotations.euler(order=order[::-1]))
|
||||
poss = anim.positions
|
||||
|
||||
for i in range(anim.shape[0]):
|
||||
for j in range(anim.shape[1]):
|
||||
|
||||
if positions or j == 0:
|
||||
|
||||
f.write("%f %f %f %f %f %f " % (
|
||||
poss[i, j, 0], poss[i, j, 1], poss[i, j, 2],
|
||||
rots[i, j, ordermap[order[0]]], rots[i, j, ordermap[order[1]]], rots[i, j, ordermap[order[2]]]))
|
||||
|
||||
else:
|
||||
|
||||
f.write("%f %f %f " % (
|
||||
rots[i, j, ordermap[order[0]]], rots[i, j, ordermap[order[1]]], rots[i, j, ordermap[order[2]]]))
|
||||
|
||||
f.write("\n")
|
||||
|
||||
|
||||
def save_joint(f, anim, names, t, i, order='zyx', positions=False):
|
||||
|
||||
f.write("%sJOINT %s\n" % (t, names[i]))
|
||||
f.write("%s{\n" % t)
|
||||
t += '\t'
|
||||
|
||||
f.write("%sOFFSET %f %f %f\n" % (t, anim.offsets[i, 0], anim.offsets[i, 1], anim.offsets[i, 2]))
|
||||
|
||||
if positions:
|
||||
f.write("%sCHANNELS 6 Xposition Yposition Zposition %s %s %s \n" % (t,
|
||||
channelmap_inv[order[0]], channelmap_inv[order[1]], channelmap_inv[order[2]]))
|
||||
else:
|
||||
f.write("%sCHANNELS 3 %s %s %s\n" % (t,
|
||||
channelmap_inv[order[0]], channelmap_inv[order[1]], channelmap_inv[order[2]]))
|
||||
|
||||
end_site = True
|
||||
|
||||
for j in range(anim.shape[1]):
|
||||
if anim.parents[j] == i:
|
||||
t = save_joint(f, anim, names, t, j, order=order, positions=positions)
|
||||
end_site = False
|
||||
|
||||
if end_site:
|
||||
f.write("%sEnd Site\n" % t)
|
||||
f.write("%s{\n" % t)
|
||||
t += '\t'
|
||||
f.write("%sOFFSET %f %f %f\n" % (t, 0.0, 0.0, 0.0))
|
||||
t = t[:-1]
|
||||
f.write("%s}\n" % t)
|
||||
|
||||
t = t[:-1]
|
||||
f.write("%s}\n" % t)
|
||||
|
||||
return t
|
||||
@@ -0,0 +1,496 @@
|
||||
import numpy as np
|
||||
import scipy.linalg as linalg
|
||||
|
||||
import Animation
|
||||
import AnimationStructure
|
||||
|
||||
from Quaternions import Quaternions
|
||||
|
||||
|
||||
class BasicInverseKinematics:
|
||||
"""
|
||||
Basic Inverse Kinematics Solver
|
||||
|
||||
This is an extremely simple full body IK
|
||||
solver.
|
||||
|
||||
It works given the following conditions:
|
||||
|
||||
* All joint targets must be specified
|
||||
* All joint targets must be in reach
|
||||
* All joint targets must not differ
|
||||
extremely from the starting pose
|
||||
* No bone length constraints can be violated
|
||||
* The root translation and rotation are
|
||||
set to good initial values
|
||||
|
||||
It works under the observation that if the
|
||||
_directions_ the joints are pointing toward
|
||||
match the _directions_ of the vectors between
|
||||
the target joints then the pose should match
|
||||
that of the target pose.
|
||||
|
||||
Therefore it iterates over joints rotating
|
||||
each joint such that the vectors between it
|
||||
and it's children match that of the target
|
||||
positions.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
|
||||
animation : Animation
|
||||
animation input
|
||||
|
||||
positions : (F, J, 3) ndarray
|
||||
target positions for each frame F
|
||||
and each joint J
|
||||
|
||||
iterations : int
|
||||
Optional number of iterations.
|
||||
If the above conditions are met
|
||||
1 iteration should be enough,
|
||||
therefore the default is 1
|
||||
|
||||
silent : bool
|
||||
Optional if to suppress output
|
||||
defaults to False
|
||||
"""
|
||||
|
||||
def __init__(self, animation, positions, iterations=1, silent=True):
|
||||
|
||||
self.animation = animation
|
||||
self.positions = positions
|
||||
self.iterations = iterations
|
||||
self.silent = silent
|
||||
|
||||
def __call__(self):
|
||||
|
||||
children = AnimationStructure.children_list(self.animation.parents)
|
||||
|
||||
for i in range(self.iterations):
|
||||
|
||||
for j in AnimationStructure.joints(self.animation.parents):
|
||||
|
||||
c = np.array(children[j])
|
||||
if len(c) == 0:
|
||||
continue
|
||||
|
||||
anim_transforms = Animation.transforms_global(self.animation)
|
||||
anim_positions = anim_transforms[:, :, :3, 3]
|
||||
anim_rotations = Quaternions.from_transforms(anim_transforms)
|
||||
|
||||
jdirs = anim_positions[:, c] - anim_positions[:, np.newaxis, j]
|
||||
ddirs = self.positions[:, c] - anim_positions[:, np.newaxis, j]
|
||||
|
||||
jsums = np.sqrt(np.sum(jdirs**2.0, axis=-1)) + 1e-10
|
||||
dsums = np.sqrt(np.sum(ddirs**2.0, axis=-1)) + 1e-10
|
||||
|
||||
jdirs = jdirs / jsums[:, :, np.newaxis]
|
||||
ddirs = ddirs / dsums[:, :, np.newaxis]
|
||||
|
||||
angles = np.arccos(np.sum(jdirs * ddirs, axis=2).clip(-1, 1))
|
||||
axises = np.cross(jdirs, ddirs)
|
||||
axises = -anim_rotations[:, j, np.newaxis] * axises
|
||||
|
||||
rotations = Quaternions.from_angle_axis(angles, axises)
|
||||
|
||||
if rotations.shape[1] == 1:
|
||||
averages = rotations[:, 0]
|
||||
else:
|
||||
averages = Quaternions.exp(rotations.log().mean(axis=-2))
|
||||
|
||||
self.animation.rotations[:, j] = self.animation.rotations[:, j] * averages
|
||||
|
||||
if not self.silent:
|
||||
anim_positions = Animation.positions_global(self.animation)
|
||||
error = np.mean(np.sum((anim_positions - self.positions)**2.0, axis=-1)**0.5, axis=-1)
|
||||
print('[BasicInverseKinematics] Iteration %i Error: %f' % (i+1, error))
|
||||
|
||||
return self.animation
|
||||
|
||||
|
||||
class JacobianInverseKinematics:
|
||||
"""
|
||||
Jacobian Based Full Body IK Solver
|
||||
|
||||
This is a full body IK solver which
|
||||
uses the dampened least squares inverse
|
||||
jacobian method.
|
||||
|
||||
It should remain fairly stable and effective
|
||||
even for joint positions which are out of
|
||||
reach and it can also take any number of targets
|
||||
to treat as end effectors.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
|
||||
animation : Animation
|
||||
animation to solve inverse problem on
|
||||
|
||||
targets : {int : (F, 3) ndarray}
|
||||
Dictionary of target positions for each
|
||||
frame F, mapping joint index to
|
||||
a target position
|
||||
|
||||
references : (F, 3)
|
||||
Optional list of J joint position
|
||||
references for which the result
|
||||
should bias toward
|
||||
|
||||
iterations : int
|
||||
Optional number of iterations to
|
||||
compute. More iterations results in
|
||||
better accuracy but takes longer to
|
||||
compute. Default is 10.
|
||||
|
||||
recalculate : bool
|
||||
Optional if to recalcuate jacobian
|
||||
each iteration. Gives better accuracy
|
||||
but slower to compute. Defaults to True
|
||||
|
||||
damping : float
|
||||
Optional damping constant. Higher
|
||||
damping increases stability but
|
||||
requires more iterations to converge.
|
||||
Defaults to 5.0
|
||||
|
||||
secondary : float
|
||||
Force, or bias toward secondary target.
|
||||
Defaults to 0.25
|
||||
|
||||
silent : bool
|
||||
Optional if to suppress output
|
||||
defaults to False
|
||||
"""
|
||||
|
||||
def __init__(self, animation, targets,
|
||||
references=None, iterations=10,
|
||||
recalculate=True, damping=2.0,
|
||||
secondary=0.25, translate=False,
|
||||
silent=False, weights=None,
|
||||
weights_translate=None):
|
||||
|
||||
self.animation = animation
|
||||
self.targets = targets
|
||||
self.references = references
|
||||
|
||||
self.iterations = iterations
|
||||
self.recalculate = recalculate
|
||||
self.damping = damping
|
||||
self.secondary = secondary
|
||||
self.translate = translate
|
||||
self.silent = silent
|
||||
self.weights = weights
|
||||
self.weights_translate = weights_translate
|
||||
|
||||
def cross(self, a, b):
|
||||
o = np.empty(b.shape)
|
||||
o[..., 0] = a[..., 1]*b[..., 2] - a[..., 2]*b[..., 1]
|
||||
o[..., 1] = a[..., 2]*b[..., 0] - a[..., 0]*b[..., 2]
|
||||
o[..., 2] = a[..., 0]*b[..., 1] - a[..., 1]*b[..., 0]
|
||||
return o
|
||||
|
||||
def jacobian(self, x, fp, fr, ts, dsc, tdsc):
|
||||
""" Find parent rotations """
|
||||
prs = fr[:, self.animation.parents]
|
||||
prs[:, 0] = Quaternions.id((1))
|
||||
|
||||
""" Find global positions of target joints """
|
||||
tps = fp[:, np.array(list(ts.keys()))]
|
||||
|
||||
""" Get partial rotations """
|
||||
qys = Quaternions.from_angle_axis(x[:, 1:prs.shape[1]*3:3], np.array([[[0, 1, 0]]]))
|
||||
qzs = Quaternions.from_angle_axis(x[:, 2:prs.shape[1]*3:3], np.array([[[0, 0, 1]]]))
|
||||
|
||||
""" Find axis of rotations """
|
||||
es = np.empty((len(x), fr.shape[1]*3, 3))
|
||||
es[:, 0::3] = ((prs * qzs) * qys) * np.array([[[1, 0, 0]]])
|
||||
es[:, 1::3] = ((prs * qzs) * np.array([[[0, 1, 0]]]))
|
||||
es[:, 2::3] = ((prs * np.array([[[0, 0, 1]]])))
|
||||
|
||||
""" Construct Jacobian """
|
||||
j = fp.repeat(3, axis=1)
|
||||
j = dsc[np.newaxis, :, :, np.newaxis] * (tps[:, np.newaxis, :] - j[:, :, np.newaxis])
|
||||
j = self.cross(es[:, :, np.newaxis, :], j)
|
||||
j = np.swapaxes(j.reshape((len(x), fr.shape[1]*3, len(ts)*3)), 1, 2)
|
||||
|
||||
if self.translate:
|
||||
|
||||
es = np.empty((len(x), fr.shape[1]*3, 3))
|
||||
es[:, 0::3] = prs * np.array([[[1, 0, 0]]])
|
||||
es[:, 1::3] = prs * np.array([[[0, 1, 0]]])
|
||||
es[:, 2::3] = prs * np.array([[[0, 0, 1]]])
|
||||
|
||||
jt = tdsc[np.newaxis, :, :, np.newaxis] * es[:, :, np.newaxis, :].repeat(tps.shape[1], axis=2)
|
||||
jt = np.swapaxes(jt.reshape((len(x), fr.shape[1]*3, len(ts)*3)), 1, 2)
|
||||
|
||||
j = np.concatenate([j, jt], axis=-1)
|
||||
|
||||
return j
|
||||
|
||||
# @profile(immediate=True)
|
||||
def __call__(self, descendants=None, gamma=1.0):
|
||||
|
||||
self.descendants = descendants
|
||||
|
||||
""" Calculate Masses """
|
||||
if self.weights is None:
|
||||
self.weights = np.ones(self.animation.shape[1])
|
||||
|
||||
if self.weights_translate is None:
|
||||
self.weights_translate = np.ones(self.animation.shape[1])
|
||||
|
||||
""" Calculate Descendants """
|
||||
if self.descendants is None:
|
||||
self.descendants = AnimationStructure.descendants_mask(self.animation.parents)
|
||||
|
||||
self.tdescendants = np.eye(self.animation.shape[1]) + self.descendants
|
||||
|
||||
self.first_descendants = self.descendants[:, np.array(list(self.targets.keys()))].repeat(3, axis=0).astype(int)
|
||||
self.first_tdescendants = self.tdescendants[:, np.array(list(self.targets.keys()))].repeat(3, axis=0).astype(int)
|
||||
|
||||
""" Calculate End Effectors """
|
||||
self.endeff = np.array(list(self.targets.values()))
|
||||
self.endeff = np.swapaxes(self.endeff, 0, 1)
|
||||
|
||||
if not self.references is None:
|
||||
self.second_descendants = self.descendants.repeat(3, axis=0).astype(int)
|
||||
self.second_tdescendants = self.tdescendants.repeat(3, axis=0).astype(int)
|
||||
self.second_targets = dict([(i, self.references[:, i]) for i in range(self.references.shape[1])])
|
||||
|
||||
nf = len(self.animation)
|
||||
nj = self.animation.shape[1]
|
||||
|
||||
if not self.silent:
|
||||
gp = Animation.positions_global(self.animation)
|
||||
gp = gp[:, np.array(list(self.targets.keys()))]
|
||||
error = np.mean(np.sqrt(np.sum((self.endeff - gp)**2.0, axis=2)))
|
||||
print('[JacobianInverseKinematics] Start | Error: %f' % error)
|
||||
|
||||
for i in range(self.iterations):
|
||||
|
||||
""" Get Global Rotations & Positions """
|
||||
gt = Animation.transforms_global(self.animation)
|
||||
gp = gt[:, :, :, 3]
|
||||
gp = gp[:, :, :3] / gp[:, :, 3, np.newaxis]
|
||||
gr = Quaternions.from_transforms(gt)
|
||||
|
||||
x = self.animation.rotations.euler().reshape(nf, -1)
|
||||
w = self.weights.repeat(3)
|
||||
|
||||
if self.translate:
|
||||
x = np.hstack([x, self.animation.positions.reshape(nf, -1)])
|
||||
w = np.hstack([w, self.weights_translate.repeat(3)])
|
||||
|
||||
""" Generate Jacobian """
|
||||
if self.recalculate or i == 0:
|
||||
j = self.jacobian(x, gp, gr, self.targets, self.first_descendants, self.first_tdescendants)
|
||||
|
||||
""" Update Variables """
|
||||
l = self.damping * (1.0 / (w + 0.001))
|
||||
d = (l*l) * np.eye(x.shape[1])
|
||||
e = gamma * (self.endeff.reshape(nf, -1) - gp[:, np.array(list(self.targets.keys()))].reshape(nf, -1))
|
||||
|
||||
j = np.nan_to_num(j)
|
||||
e = np.nan_to_num(e)
|
||||
x += np.array(list(map(lambda jf, ef:
|
||||
linalg.lu_solve(linalg.lu_factor(jf.T.dot(jf) + d), jf.T.dot(ef)), j, e)))
|
||||
|
||||
""" Generate Secondary Jacobian """
|
||||
if self.references is not None:
|
||||
|
||||
ns = np.array(list(map(lambda jf:
|
||||
np.eye(x.shape[1]) - linalg.solve(jf.T.dot(jf) + d, jf.T.dot(jf)), j)))
|
||||
|
||||
if self.recalculate or i == 0:
|
||||
j2 = self.jacobian(x, gp, gr, self.second_targets, self.second_descendants, self.second_tdescendants)
|
||||
|
||||
e2 = self.secondary * (self.references.reshape(nf, -1) - gp.reshape(nf, -1))
|
||||
|
||||
x += np.array(list(map(lambda nsf, j2f, e2f:
|
||||
nsf.dot(linalg.lu_solve(linalg.lu_factor(j2f.T.dot(j2f) + d), j2f.T.dot(e2f))), ns, j2, e2)))
|
||||
|
||||
""" Set Back Rotations / Translations """
|
||||
self.animation.rotations = Quaternions.from_euler(
|
||||
x[:, :nj*3].reshape((nf, nj, 3)), order='xyz', world=True)
|
||||
|
||||
if self.translate:
|
||||
self.animation.positions = x[:, nj*3:].reshape((nf, nj, 3))
|
||||
|
||||
""" Generate Error """
|
||||
|
||||
if not self.silent:
|
||||
gp = Animation.positions_global(self.animation)
|
||||
gp = gp[:, np.array(list(self.targets.keys()))]
|
||||
error = np.mean(np.sum((self.endeff - gp)**2.0, axis=2)**0.5)
|
||||
print('[JacobianInverseKinematics] Iteration %i | Error: %f' % (i+1, error))
|
||||
|
||||
|
||||
class BasicJacobianIK:
|
||||
"""
|
||||
Same interface as BasicInverseKinematics
|
||||
but uses the Jacobian IK Solver Instead
|
||||
"""
|
||||
|
||||
def __init__(self, animation, positions, iterations=10, silent=True, **kw):
|
||||
|
||||
targets = dict([(i, positions[:, i]) for i in range(positions.shape[1])])
|
||||
self.ik = JacobianInverseKinematics(animation, targets, iterations=iterations, silent=silent, **kw)
|
||||
|
||||
def __call__(self, **kw):
|
||||
return self.ik(**kw)
|
||||
|
||||
|
||||
class ICP:
|
||||
|
||||
def __init__(self,
|
||||
anim, rest, weights, mesh, goal,
|
||||
find_closest=True, damping=10,
|
||||
iterations=10, silent=True,
|
||||
translate=True, recalculate=True,
|
||||
weights_translate=None):
|
||||
|
||||
self.animation = anim
|
||||
self.rest = rest
|
||||
self.vweights = weights
|
||||
self.mesh = mesh
|
||||
self.goal = goal
|
||||
self.find_closest = find_closest
|
||||
self.iterations = iterations
|
||||
self.silent = silent
|
||||
self.translate = translate
|
||||
self.damping = damping
|
||||
self.weights = None
|
||||
self.weights_translate = weights_translate
|
||||
self.recalculate = recalculate
|
||||
|
||||
def cross(self, a, b):
|
||||
o = np.empty(b.shape)
|
||||
o[..., 0] = a[..., 1]*b[..., 2] - a[..., 2]*b[..., 1]
|
||||
o[..., 1] = a[..., 2]*b[..., 0] - a[..., 0]*b[..., 2]
|
||||
o[..., 2] = a[..., 0]*b[..., 1] - a[..., 1]*b[..., 0]
|
||||
return o
|
||||
|
||||
def jacobian(self, x, fp, fr, goal, weights, des_r, des_t):
|
||||
""" Find parent rotations """
|
||||
prs = fr[:, self.animation.parents]
|
||||
prs[:, 0] = Quaternions.id((1))
|
||||
|
||||
""" Get partial rotations """
|
||||
qys = Quaternions.from_angle_axis(x[:, 1:prs.shape[1]*3:3], np.array([[[0, 1, 0]]]))
|
||||
qzs = Quaternions.from_angle_axis(x[:, 2:prs.shape[1]*3:3], np.array([[[0, 0, 1]]]))
|
||||
|
||||
""" Find axis of rotations """
|
||||
es = np.empty((len(x), fr.shape[1]*3, 3))
|
||||
es[:, 0::3] = ((prs * qzs) * qys) * np.array([[[1, 0, 0]]])
|
||||
es[:, 1::3] = ((prs * qzs) * np.array([[[0, 1, 0]]]))
|
||||
es[:, 2::3] = ((prs * np.array([[[0, 0, 1]]])))
|
||||
|
||||
""" Construct Jacobian """
|
||||
j = fp.repeat(3, axis=1)
|
||||
j = des_r[np.newaxis, :, :, :, np.newaxis] * (goal[:, np.newaxis, :, np.newaxis] - j[:, :, np.newaxis, np.newaxis])
|
||||
j = np.sum(j * weights[np.newaxis, np.newaxis, :, :, np.newaxis], 3)
|
||||
j = self.cross(es[:, :, np.newaxis, :], j)
|
||||
j = np.swapaxes(j.reshape((len(x), fr.shape[1]*3, goal.shape[1]*3)), 1, 2)
|
||||
|
||||
if self.translate:
|
||||
|
||||
es = np.empty((len(x), fr.shape[1]*3, 3))
|
||||
es[:, 0::3] = prs * np.array([[[1, 0, 0]]])
|
||||
es[:, 1::3] = prs * np.array([[[0, 1, 0]]])
|
||||
es[:, 2::3] = prs * np.array([[[0, 0, 1]]])
|
||||
|
||||
jt = des_t[np.newaxis, :, :, :, np.newaxis] * es[:, :, np.newaxis, np.newaxis, :].repeat(goal.shape[1], axis=2)
|
||||
jt = np.sum(jt * weights[np.newaxis, np.newaxis, :, :, np.newaxis], 3)
|
||||
jt = np.swapaxes(jt.reshape((len(x), fr.shape[1]*3, goal.shape[1]*3)), 1, 2)
|
||||
|
||||
j = np.concatenate([j, jt], axis=-1)
|
||||
|
||||
return j
|
||||
|
||||
# @profile(immediate=True)
|
||||
def __call__(self, descendants=None, maxjoints=4, gamma=1.0, transpose=False):
|
||||
""" Calculate Masses """
|
||||
if self.weights is None:
|
||||
self.weights = np.ones(self.animation.shape[1])
|
||||
|
||||
if self.weights_translate is None:
|
||||
self.weights_translate = np.ones(self.animation.shape[1])
|
||||
|
||||
nf = len(self.animation)
|
||||
nj = self.animation.shape[1]
|
||||
nv = self.goal.shape[1]
|
||||
|
||||
weightids = np.argsort(-self.vweights, axis=1)[:, :maxjoints]
|
||||
weightvls = np.array(list(map(lambda w, i: w[i], self.vweights, weightids)))
|
||||
weightvls = weightvls / weightvls.sum(axis=1)[..., np.newaxis]
|
||||
|
||||
if descendants is None:
|
||||
self.descendants = AnimationStructure.descendants_mask(self.animation.parents)
|
||||
else:
|
||||
self.descendants = descendants
|
||||
|
||||
des_r = np.eye(nj) + self.descendants
|
||||
des_r = des_r[:, weightids].repeat(3, axis=0)
|
||||
|
||||
des_t = np.eye(nj) + self.descendants
|
||||
des_t = des_t[:, weightids].repeat(3, axis=0)
|
||||
|
||||
if not self.silent:
|
||||
curr = Animation.skin(self.animation, self.rest, self.vweights, self.mesh, maxjoints=maxjoints)
|
||||
error = np.mean(np.sqrt(np.sum((curr - self.goal)**2.0, axis=-1)))
|
||||
print('[ICP] Start | Error: %f' % error)
|
||||
|
||||
for i in range(self.iterations):
|
||||
|
||||
""" Get Global Rotations & Positions """
|
||||
gt = Animation.transforms_global(self.animation)
|
||||
gp = gt[:, :, :, 3]
|
||||
gp = gp[:, :, :3] / gp[:, :, 3, np.newaxis]
|
||||
gr = Quaternions.from_transforms(gt)
|
||||
|
||||
x = self.animation.rotations.euler().reshape(nf, -1)
|
||||
w = self.weights.repeat(3)
|
||||
|
||||
if self.translate:
|
||||
x = np.hstack([x, self.animation.positions.reshape(nf, -1)])
|
||||
w = np.hstack([w, self.weights_translate.repeat(3)])
|
||||
|
||||
""" Get Current State """
|
||||
curr = Animation.skin(self.animation, self.rest, self.vweights, self.mesh, maxjoints=maxjoints)
|
||||
|
||||
""" Find Cloest Points """
|
||||
if self.find_closest:
|
||||
mapping = np.argmin(
|
||||
(curr[:, :, np.newaxis] -
|
||||
self.goal[:, np.newaxis, :])**2.0, axis=2)
|
||||
e = gamma * (np.array(list(map(lambda g, m: g[m], self.goal, mapping))) - curr).reshape(nf, -1)
|
||||
else:
|
||||
e = gamma * (self.goal - curr).reshape(nf, -1)
|
||||
|
||||
""" Generate Jacobian """
|
||||
if self.recalculate or i == 0:
|
||||
j = self.jacobian(x, gp, gr, self.goal, weightvls, des_r, des_t)
|
||||
|
||||
""" Update Variables """
|
||||
l = self.damping * (1.0 / (w + 1e-10))
|
||||
d = (l*l) * np.eye(x.shape[1])
|
||||
|
||||
if transpose:
|
||||
x += np.array(list(map(lambda jf, ef: jf.T.dot(ef), j, e)))
|
||||
else:
|
||||
x += np.array(list(map(lambda jf, ef:
|
||||
linalg.lu_solve(linalg.lu_factor(jf.T.dot(jf) + d), jf.T.dot(ef)), j, e)))
|
||||
|
||||
""" Set Back Rotations / Translations """
|
||||
self.animation.rotations = Quaternions.from_euler(
|
||||
x[:, :nj*3].reshape((nf, nj, 3)), order='xyz', world=True)
|
||||
|
||||
if self.translate:
|
||||
self.animation.positions = x[:, nj*3:].reshape((nf, nj, 3))
|
||||
|
||||
if not self.silent:
|
||||
curr = Animation.skin(self.animation, self.rest, self.vweights, self.mesh)
|
||||
error = np.mean(np.sqrt(np.sum((curr - self.goal)**2.0, axis=-1)))
|
||||
print('[ICP] Iteration %i | Error: %f' % (i+1, error))
|
||||
@@ -0,0 +1,91 @@
|
||||
import numpy as np
|
||||
|
||||
from Quaternions import Quaternions
|
||||
|
||||
|
||||
class Pivots:
|
||||
"""
|
||||
Pivots is an ndarray of angular rotations
|
||||
|
||||
This wrapper provides some functions for
|
||||
working with pivots.
|
||||
|
||||
These are particularly useful as a number
|
||||
of atomic operations (such as adding or
|
||||
subtracting) cannot be achieved using
|
||||
the standard arithmatic and need to be
|
||||
defined differently to work correctly
|
||||
"""
|
||||
|
||||
def __init__(self, ps): self.ps = np.array(ps)
|
||||
def __str__(self): return "Pivots(" + str(self.ps) + ")"
|
||||
def __repr__(self): return "Pivots(" + repr(self.ps) + ")"
|
||||
|
||||
def __add__(self, other): return Pivots(np.arctan2(np.sin(self.ps + other.ps), np.cos(self.ps + other.ps)))
|
||||
def __sub__(self, other): return Pivots(np.arctan2(np.sin(self.ps - other.ps), np.cos(self.ps - other.ps)))
|
||||
def __mul__(self, other): return Pivots(self.ps * other.ps)
|
||||
def __div__(self, other): return Pivots(self.ps / other.ps)
|
||||
def __mod__(self, other): return Pivots(self.ps % other.ps)
|
||||
def __pow__(self, other): return Pivots(self.ps ** other.ps)
|
||||
|
||||
def __lt__(self, other): return self.ps < other.ps
|
||||
def __le__(self, other): return self.ps <= other.ps
|
||||
def __eq__(self, other): return self.ps == other.ps
|
||||
def __ne__(self, other): return self.ps != other.ps
|
||||
def __ge__(self, other): return self.ps >= other.ps
|
||||
def __gt__(self, other): return self.ps > other.ps
|
||||
|
||||
def __abs__(self): return Pivots(abs(self.ps))
|
||||
def __neg__(self): return Pivots(-self.ps)
|
||||
|
||||
def __iter__(self): return iter(self.ps)
|
||||
def __len__(self): return len(self.ps)
|
||||
|
||||
def __getitem__(self, k): return Pivots(self.ps[k])
|
||||
def __setitem__(self, k, v): self.ps[k] = v.ps
|
||||
|
||||
def _ellipsis(self): return tuple(map(lambda x: slice(None), self.shape))
|
||||
|
||||
def quaternions(self, plane='xz'):
|
||||
fa = self._ellipsis()
|
||||
axises = np.ones(self.ps.shape + (3,))
|
||||
axises[fa + ("xyz".index(plane[0]),)] = 0.0
|
||||
axises[fa + ("xyz".index(plane[1]),)] = 0.0
|
||||
return Quaternions.from_angle_axis(self.ps, axises)
|
||||
|
||||
def directions(self, plane='xz'):
|
||||
dirs = np.zeros((len(self.ps), 3))
|
||||
dirs[..., "xyz".index(plane[0])] = np.sin(self.ps)
|
||||
dirs[..., "xyz".index(plane[1])] = np.cos(self.ps)
|
||||
return dirs
|
||||
|
||||
def normalized(self):
|
||||
xs = np.copy(self.ps)
|
||||
while np.any(xs > np.pi):
|
||||
xs[xs > np.pi] = xs[xs > np.pi] - 2 * np.pi
|
||||
while np.any(xs < -np.pi):
|
||||
xs[xs < -np.pi] = xs[xs < -np.pi] + 2 * np.pi
|
||||
return Pivots(xs)
|
||||
|
||||
def interpolate(self, ws):
|
||||
dir = np.average(self.directions, weights=ws, axis=0)
|
||||
return np.arctan2(dir[2], dir[0])
|
||||
|
||||
def copy(self):
|
||||
return Pivots(np.copy(self.ps))
|
||||
|
||||
@property
|
||||
def shape(self):
|
||||
return self.ps.shape
|
||||
|
||||
@classmethod
|
||||
def from_quaternions(cls, qs, forward='z', plane='xz'):
|
||||
ds = np.zeros(qs.shape + (3,))
|
||||
ds[..., 'xyz'.index(forward)] = 1.0
|
||||
return Pivots.from_directions(qs * ds, plane=plane)
|
||||
|
||||
@classmethod
|
||||
def from_directions(cls, ds, plane='xz'):
|
||||
ys = ds[..., 'xyz'.index(plane[0])]
|
||||
xs = ds[..., 'xyz'.index(plane[1])]
|
||||
return Pivots(np.arctan2(ys, xs))
|
||||
@@ -0,0 +1,481 @@
|
||||
import numpy as np
|
||||
|
||||
|
||||
class Quaternions:
|
||||
"""
|
||||
Quaternions is a wrapper around a numpy ndarray
|
||||
that allows it to act as if it were an narray of
|
||||
a quaternion data type.
|
||||
|
||||
Therefore addition, subtraction, multiplication,
|
||||
division, negation, absolute, are all defined
|
||||
in terms of quaternion operations such as quaternion
|
||||
multiplication.
|
||||
|
||||
This allows for much neater code and many routines
|
||||
which conceptually do the same thing to be written
|
||||
in the same way for point data and for rotation data.
|
||||
|
||||
The Quaternions class has been desgined such that it
|
||||
should support broadcasting and slicing in all of the
|
||||
usual ways.
|
||||
"""
|
||||
|
||||
def __init__(self, qs):
|
||||
if isinstance(qs, np.ndarray):
|
||||
|
||||
if len(qs.shape) == 1:
|
||||
qs = np.array([qs])
|
||||
self.qs = qs
|
||||
return
|
||||
|
||||
if isinstance(qs, Quaternions):
|
||||
self.qs = qs.qs
|
||||
return
|
||||
|
||||
raise TypeError('Quaternions must be constructed from iterable, numpy array, or Quaternions, not %s' % type(qs))
|
||||
|
||||
def __str__(self): return "Quaternions(" + str(self.qs) + ")"
|
||||
def __repr__(self): return "Quaternions(" + repr(self.qs) + ")"
|
||||
""" Helper Methods for Broadcasting and Data extraction """
|
||||
|
||||
@classmethod
|
||||
def _broadcast(cls, sqs, oqs, scalar=False):
|
||||
if isinstance(oqs, float):
|
||||
return sqs, oqs * np.ones(sqs.shape[:-1])
|
||||
|
||||
ss = np.array(sqs.shape) if not scalar else np.array(sqs.shape[:-1])
|
||||
os = np.array(oqs.shape)
|
||||
|
||||
if len(ss) != len(os):
|
||||
raise TypeError('Quaternions cannot broadcast together shapes %s and %s' % (sqs.shape, oqs.shape))
|
||||
|
||||
if np.all(ss == os):
|
||||
return sqs, oqs
|
||||
|
||||
if not np.all((ss == os) | (os == np.ones(len(os))) | (ss == np.ones(len(ss)))):
|
||||
raise TypeError('Quaternions cannot broadcast together shapes %s and %s' % (sqs.shape, oqs.shape))
|
||||
|
||||
sqsn, oqsn = sqs.copy(), oqs.copy()
|
||||
|
||||
for a in np.where(ss == 1)[0]:
|
||||
sqsn = sqsn.repeat(os[a], axis=a)
|
||||
for a in np.where(os == 1)[0]:
|
||||
oqsn = oqsn.repeat(ss[a], axis=a)
|
||||
|
||||
return sqsn, oqsn
|
||||
|
||||
""" Adding Quaterions is just Defined as Multiplication """
|
||||
|
||||
def __add__(self, other): return self * other
|
||||
def __sub__(self, other): return self / other
|
||||
""" Quaterion Multiplication """
|
||||
|
||||
def __mul__(self, other):
|
||||
"""
|
||||
Quaternion multiplication has three main methods.
|
||||
|
||||
When multiplying a Quaternions array by Quaternions
|
||||
normal quaternion multiplication is performed.
|
||||
|
||||
When multiplying a Quaternions array by a vector
|
||||
array of the same shape, where the last axis is 3,
|
||||
it is assumed to be a Quaternion by 3D-Vector
|
||||
multiplication and the 3D-Vectors are rotated
|
||||
in space by the Quaternions.
|
||||
|
||||
When multipplying a Quaternions array by a scalar
|
||||
or vector of different shape it is assumed to be
|
||||
a Quaternions by Scalars multiplication and the
|
||||
Quaternions are scaled using Slerp and the identity
|
||||
quaternions.
|
||||
"""
|
||||
|
||||
""" If Quaternions type do Quaternions * Quaternions """
|
||||
if isinstance(other, Quaternions):
|
||||
sqs, oqs = Quaternions._broadcast(self.qs, other.qs)
|
||||
|
||||
q0 = sqs[..., 0]
|
||||
q1 = sqs[..., 1]
|
||||
q2 = sqs[..., 2]
|
||||
q3 = sqs[..., 3]
|
||||
r0 = oqs[..., 0]
|
||||
r1 = oqs[..., 1]
|
||||
r2 = oqs[..., 2]
|
||||
r3 = oqs[..., 3]
|
||||
|
||||
qs = np.empty(sqs.shape)
|
||||
qs[..., 0] = r0 * q0 - r1 * q1 - r2 * q2 - r3 * q3
|
||||
qs[..., 1] = r0 * q1 + r1 * q0 - r2 * q3 + r3 * q2
|
||||
qs[..., 2] = r0 * q2 + r1 * q3 + r2 * q0 - r3 * q1
|
||||
qs[..., 3] = r0 * q3 - r1 * q2 + r2 * q1 + r3 * q0
|
||||
|
||||
return Quaternions(qs)
|
||||
|
||||
""" If array type do Quaternions * Vectors """
|
||||
if isinstance(other, np.ndarray) and other.shape[-1] == 3:
|
||||
vs = Quaternions(np.concatenate([np.zeros(other.shape[:-1] + (1,)), other], axis=-1))
|
||||
|
||||
return (self * (vs * -self)).imaginaries
|
||||
|
||||
""" If float do Quaternions * Scalars """
|
||||
if isinstance(other, np.ndarray) or isinstance(other, float):
|
||||
return Quaternions.slerp(Quaternions.id_like(self), self, other)
|
||||
|
||||
raise TypeError('Cannot multiply/add Quaternions with type %s' % str(type(other)))
|
||||
|
||||
def __div__(self, other):
|
||||
"""
|
||||
When a Quaternion type is supplied, division is defined
|
||||
as multiplication by the inverse of that Quaternion.
|
||||
|
||||
When a scalar or vector is supplied it is defined
|
||||
as multiplicaion of one over the supplied value.
|
||||
Essentially a scaling.
|
||||
"""
|
||||
|
||||
if isinstance(other, Quaternions):
|
||||
return self * (-other)
|
||||
if isinstance(other, np.ndarray):
|
||||
return self * (1.0 / other)
|
||||
if isinstance(other, float):
|
||||
return self * (1.0 / other)
|
||||
raise TypeError('Cannot divide/subtract Quaternions with type %s' + str(type(other)))
|
||||
|
||||
def __eq__(self, other): return self.qs == other.qs
|
||||
def __ne__(self, other): return self.qs != other.qs
|
||||
|
||||
def __neg__(self):
|
||||
""" Invert Quaternions """
|
||||
return Quaternions(self.qs * np.array([[1, -1, -1, -1]]))
|
||||
|
||||
def __abs__(self):
|
||||
""" Unify Quaternions To Single Pole """
|
||||
qabs = self.normalized().copy()
|
||||
top = np.sum((qabs.qs) * np.array([1, 0, 0, 0]), axis=-1)
|
||||
bot = np.sum((-qabs.qs) * np.array([1, 0, 0, 0]), axis=-1)
|
||||
qabs.qs[top < bot] = -qabs.qs[top < bot]
|
||||
return qabs
|
||||
|
||||
def __iter__(self): return iter(self.qs)
|
||||
def __len__(self): return len(self.qs)
|
||||
|
||||
def __getitem__(self, k): return Quaternions(self.qs[k])
|
||||
def __setitem__(self, k, v): self.qs[k] = v.qs
|
||||
|
||||
@property
|
||||
def lengths(self):
|
||||
return np.sum(self.qs**2.0, axis=-1)**0.5
|
||||
|
||||
@property
|
||||
def reals(self):
|
||||
return self.qs[..., 0]
|
||||
|
||||
@property
|
||||
def imaginaries(self):
|
||||
return self.qs[..., 1:4]
|
||||
|
||||
@property
|
||||
def shape(self): return self.qs.shape[:-1]
|
||||
|
||||
def repeat(self, n, **kwargs):
|
||||
return Quaternions(self.qs.repeat(n, **kwargs))
|
||||
|
||||
def normalized(self):
|
||||
return Quaternions(self.qs / self.lengths[..., np.newaxis])
|
||||
|
||||
def log(self):
|
||||
norm = abs(self.normalized())
|
||||
imgs = norm.imaginaries
|
||||
lens = np.sqrt(np.sum(imgs**2, axis=-1))
|
||||
lens = np.arctan2(lens, norm.reals) / (lens + 1e-10)
|
||||
return imgs * lens[..., np.newaxis]
|
||||
|
||||
def constrained(self, axis):
|
||||
|
||||
rl = self.reals
|
||||
im = np.sum(axis * self.imaginaries, axis=-1)
|
||||
|
||||
t1 = -2 * np.arctan2(rl, im) + np.pi
|
||||
t2 = -2 * np.arctan2(rl, im) - np.pi
|
||||
|
||||
top = Quaternions.exp(axis[np.newaxis] * (t1[:, np.newaxis] / 2.0))
|
||||
bot = Quaternions.exp(axis[np.newaxis] * (t2[:, np.newaxis] / 2.0))
|
||||
img = self.dot(top) > self.dot(bot)
|
||||
|
||||
ret = top.copy()
|
||||
ret[img] = top[img]
|
||||
ret[~img] = bot[~img]
|
||||
return ret
|
||||
|
||||
def constrained_x(self): return self.constrained(np.array([1, 0, 0]))
|
||||
def constrained_y(self): return self.constrained(np.array([0, 1, 0]))
|
||||
def constrained_z(self): return self.constrained(np.array([0, 0, 1]))
|
||||
|
||||
def dot(self, q): return np.sum(self.qs * q.qs, axis=-1)
|
||||
|
||||
def copy(self): return Quaternions(np.copy(self.qs))
|
||||
|
||||
def reshape(self, s):
|
||||
self.qs.reshape(s)
|
||||
return self
|
||||
|
||||
def interpolate(self, ws):
|
||||
return Quaternions.exp(np.average(abs(self).log, axis=0, weights=ws))
|
||||
|
||||
def euler(self, order='xyz'):
|
||||
|
||||
q = self.normalized().qs
|
||||
q0 = q[..., 0]
|
||||
q1 = q[..., 1]
|
||||
q2 = q[..., 2]
|
||||
q3 = q[..., 3]
|
||||
es = np.zeros(self.shape + (3,))
|
||||
|
||||
if order == 'xyz':
|
||||
es[..., 0] = np.arctan2(2 * (q0 * q1 + q2 * q3), 1 - 2 * (q1 * q1 + q2 * q2))
|
||||
es[..., 1] = np.arcsin((2 * (q0 * q2 - q3 * q1)).clip(-1, 1))
|
||||
es[..., 2] = np.arctan2(2 * (q0 * q3 + q1 * q2), 1 - 2 * (q2 * q2 + q3 * q3))
|
||||
elif order == 'yzx':
|
||||
es[..., 0] = np.arctan2(2 * (q1 * q0 - q2 * q3), -q1 * q1 + q2 * q2 - q3 * q3 + q0 * q0)
|
||||
es[..., 1] = np.arctan2(2 * (q2 * q0 - q1 * q3), q1 * q1 - q2 * q2 - q3 * q3 + q0 * q0)
|
||||
es[..., 2] = np.arcsin((2 * (q1 * q2 + q3 * q0)).clip(-1, 1))
|
||||
else:
|
||||
raise NotImplementedError('Cannot convert from ordering %s' % order)
|
||||
|
||||
"""
|
||||
|
||||
# These conversion don't appear to work correctly for Maya.
|
||||
# http://bediyap.com/programming/convert-quaternion-to-euler-rotations/
|
||||
|
||||
if order == 'xyz':
|
||||
es[fa + (0,)] = np.arctan2(2 * (q0 * q3 - q1 * q2), q0 * q0 + q1 * q1 - q2 * q2 - q3 * q3)
|
||||
es[fa + (1,)] = np.arcsin((2 * (q1 * q3 + q0 * q2)).clip(-1,1))
|
||||
es[fa + (2,)] = np.arctan2(2 * (q0 * q1 - q2 * q3), q0 * q0 - q1 * q1 - q2 * q2 + q3 * q3)
|
||||
elif order == 'yzx':
|
||||
es[fa + (0,)] = np.arctan2(2 * (q0 * q1 - q2 * q3), q0 * q0 - q1 * q1 + q2 * q2 - q3 * q3)
|
||||
es[fa + (1,)] = np.arcsin((2 * (q1 * q2 + q0 * q3)).clip(-1,1))
|
||||
es[fa + (2,)] = np.arctan2(2 * (q0 * q2 - q1 * q3), q0 * q0 + q1 * q1 - q2 * q2 - q3 * q3)
|
||||
elif order == 'zxy':
|
||||
es[fa + (0,)] = np.arctan2(2 * (q0 * q2 - q1 * q3), q0 * q0 - q1 * q1 - q2 * q2 + q3 * q3)
|
||||
es[fa + (1,)] = np.arcsin((2 * (q0 * q1 + q2 * q3)).clip(-1,1))
|
||||
es[fa + (2,)] = np.arctan2(2 * (q0 * q3 - q1 * q2), q0 * q0 - q1 * q1 + q2 * q2 - q3 * q3)
|
||||
elif order == 'xzy':
|
||||
es[fa + (0,)] = np.arctan2(2 * (q0 * q2 + q1 * q3), q0 * q0 + q1 * q1 - q2 * q2 - q3 * q3)
|
||||
es[fa + (1,)] = np.arcsin((2 * (q0 * q3 - q1 * q2)).clip(-1,1))
|
||||
es[fa + (2,)] = np.arctan2(2 * (q0 * q1 + q2 * q3), q0 * q0 - q1 * q1 + q2 * q2 - q3 * q3)
|
||||
elif order == 'yxz':
|
||||
es[fa + (0,)] = np.arctan2(2 * (q1 * q2 + q0 * q3), q0 * q0 - q1 * q1 + q2 * q2 - q3 * q3)
|
||||
es[fa + (1,)] = np.arcsin((2 * (q0 * q1 - q2 * q3)).clip(-1,1))
|
||||
es[fa + (2,)] = np.arctan2(2 * (q1 * q3 + q0 * q2), q0 * q0 - q1 * q1 - q2 * q2 + q3 * q3)
|
||||
elif order == 'zyx':
|
||||
es[fa + (0,)] = np.arctan2(2 * (q0 * q1 + q2 * q3), q0 * q0 - q1 * q1 - q2 * q2 + q3 * q3)
|
||||
es[fa + (1,)] = np.arcsin((2 * (q0 * q2 - q1 * q3)).clip(-1,1))
|
||||
es[fa + (2,)] = np.arctan2(2 * (q0 * q3 + q1 * q2), q0 * q0 + q1 * q1 - q2 * q2 - q3 * q3)
|
||||
else:
|
||||
raise KeyError('Unknown ordering %s' % order)
|
||||
|
||||
"""
|
||||
|
||||
# https://github.com/ehsan/ogre/blob/master/OgreMain/src/OgreMatrix3.cpp
|
||||
# Use this class and convert from matrix
|
||||
|
||||
return es
|
||||
|
||||
def average(self):
|
||||
|
||||
if len(self.shape) == 1:
|
||||
|
||||
import numpy.core.umath_tests as ut
|
||||
system = ut.matrix_multiply(self.qs[:, :, np.newaxis], self.qs[:, np.newaxis, :]).sum(axis=0)
|
||||
w, v = np.linalg.eigh(system)
|
||||
qiT_dot_qref = (self.qs[:, :, np.newaxis] * v[np.newaxis, :, :]).sum(axis=1)
|
||||
return Quaternions(v[:, np.argmin((1.-qiT_dot_qref**2).sum(axis=0))])
|
||||
|
||||
else:
|
||||
|
||||
raise NotImplementedError('Cannot average multi-dimensionsal Quaternions')
|
||||
|
||||
def angle_axis(self):
|
||||
|
||||
norm = self.normalized()
|
||||
s = np.sqrt(1 - (norm.reals**2.0))
|
||||
s[s == 0] = 0.001
|
||||
|
||||
angles = 2.0 * np.arccos(norm.reals)
|
||||
axis = norm.imaginaries / s[..., np.newaxis]
|
||||
|
||||
return angles, axis
|
||||
|
||||
def transforms(self):
|
||||
|
||||
qw = self.qs[..., 0]
|
||||
qx = self.qs[..., 1]
|
||||
qy = self.qs[..., 2]
|
||||
qz = self.qs[..., 3]
|
||||
|
||||
x2 = qx + qx
|
||||
y2 = qy + qy
|
||||
z2 = qz + qz
|
||||
xx = qx * x2
|
||||
yy = qy * y2
|
||||
wx = qw * x2
|
||||
xy = qx * y2
|
||||
yz = qy * z2
|
||||
wy = qw * y2
|
||||
xz = qx * z2
|
||||
zz = qz * z2
|
||||
wz = qw * z2
|
||||
|
||||
m = np.empty(self.shape + (3, 3))
|
||||
m[..., 0, 0] = 1.0 - (yy + zz)
|
||||
m[..., 0, 1] = xy - wz
|
||||
m[..., 0, 2] = xz + wy
|
||||
m[..., 1, 0] = xy + wz
|
||||
m[..., 1, 1] = 1.0 - (xx + zz)
|
||||
m[..., 1, 2] = yz - wx
|
||||
m[..., 2, 0] = xz - wy
|
||||
m[..., 2, 1] = yz + wx
|
||||
m[..., 2, 2] = 1.0 - (xx + yy)
|
||||
|
||||
return m
|
||||
|
||||
def ravel(self):
|
||||
return self.qs.ravel()
|
||||
|
||||
@classmethod
|
||||
def id(cls, n):
|
||||
|
||||
if isinstance(n, tuple):
|
||||
qs = np.zeros(n + (4,))
|
||||
qs[..., 0] = 1.0
|
||||
return Quaternions(qs)
|
||||
|
||||
if isinstance(n, int):
|
||||
qs = np.zeros((n, 4))
|
||||
qs[:, 0] = 1.0
|
||||
return Quaternions(qs)
|
||||
|
||||
raise TypeError('Cannot Construct Quaternion from %s type' % str(type(n)))
|
||||
|
||||
@classmethod
|
||||
def id_like(cls, a):
|
||||
qs = np.zeros(a.shape + (4,))
|
||||
qs[..., 0] = 1.0
|
||||
return Quaternions(qs)
|
||||
|
||||
@classmethod
|
||||
def exp(cls, ws):
|
||||
|
||||
ts = np.sum(ws**2.0, axis=-1)**0.5
|
||||
ts[ts == 0] = 0.001
|
||||
ls = np.sin(ts) / ts
|
||||
|
||||
qs = np.empty(ws.shape[:-1] + (4,))
|
||||
qs[..., 0] = np.cos(ts)
|
||||
qs[..., 1] = ws[..., 0] * ls
|
||||
qs[..., 2] = ws[..., 1] * ls
|
||||
qs[..., 3] = ws[..., 2] * ls
|
||||
|
||||
return Quaternions(qs).normalized()
|
||||
|
||||
@classmethod
|
||||
def slerp(cls, q0s, q1s, a):
|
||||
|
||||
fst, snd = cls._broadcast(q0s.qs, q1s.qs)
|
||||
fst, a = cls._broadcast(fst, a, scalar=True)
|
||||
snd, a = cls._broadcast(snd, a, scalar=True)
|
||||
|
||||
len = np.sum(fst * snd, axis=-1)
|
||||
|
||||
neg = len < 0.0
|
||||
len[neg] = -len[neg]
|
||||
snd[neg] = -snd[neg]
|
||||
|
||||
amount0 = np.zeros(a.shape)
|
||||
amount1 = np.zeros(a.shape)
|
||||
|
||||
linear = (1.0 - len) < 0.01
|
||||
omegas = np.arccos(len[~linear])
|
||||
sinoms = np.sin(omegas)
|
||||
|
||||
amount0[linear] = 1.0 - a[linear]
|
||||
amount1[linear] = a[linear]
|
||||
amount0[~linear] = np.sin((1.0 - a[~linear]) * omegas) / sinoms
|
||||
amount1[~linear] = np.sin(a[~linear] * omegas) / sinoms
|
||||
|
||||
return Quaternions(
|
||||
amount0[..., np.newaxis] * fst +
|
||||
amount1[..., np.newaxis] * snd)
|
||||
|
||||
@classmethod
|
||||
def between(cls, v0s, v1s):
|
||||
a = np.cross(v0s, v1s)
|
||||
w = np.sqrt((v0s**2).sum(axis=-1) * (v1s**2).sum(axis=-1)) + (v0s * v1s).sum(axis=-1)
|
||||
return Quaternions(np.concatenate([w[..., np.newaxis], a], axis=-1)).normalized()
|
||||
|
||||
@classmethod
|
||||
def from_angle_axis(cls, angles, axis):
|
||||
axis = axis / (np.sqrt(np.sum(axis**2, axis=-1)) + 1e-10)[..., np.newaxis]
|
||||
sines = np.sin(angles / 2.0)[..., np.newaxis]
|
||||
cosines = np.cos(angles / 2.0)[..., np.newaxis]
|
||||
return Quaternions(np.concatenate([cosines, axis * sines], axis=-1))
|
||||
|
||||
@classmethod
|
||||
def from_euler(cls, es, order='xyz', world=False):
|
||||
|
||||
axis = {
|
||||
'x': np.array([1, 0, 0]),
|
||||
'y': np.array([0, 1, 0]),
|
||||
'z': np.array([0, 0, 1]),
|
||||
}
|
||||
|
||||
q0s = Quaternions.from_angle_axis(es[..., 0], axis[order[0]])
|
||||
q1s = Quaternions.from_angle_axis(es[..., 1], axis[order[1]])
|
||||
q2s = Quaternions.from_angle_axis(es[..., 2], axis[order[2]])
|
||||
|
||||
return (q2s * (q1s * q0s)) if world else (q0s * (q1s * q2s))
|
||||
|
||||
@classmethod
|
||||
def from_transforms(cls, ts):
|
||||
|
||||
d0, d1, d2 = ts[..., 0, 0], ts[..., 1, 1], ts[..., 2, 2]
|
||||
|
||||
q0 = (d0 + d1 + d2 + 1.0) / 4.0
|
||||
q1 = (d0 - d1 - d2 + 1.0) / 4.0
|
||||
q2 = (-d0 + d1 - d2 + 1.0) / 4.0
|
||||
q3 = (-d0 - d1 + d2 + 1.0) / 4.0
|
||||
|
||||
q0 = np.sqrt(q0.clip(0, None))
|
||||
q1 = np.sqrt(q1.clip(0, None))
|
||||
q2 = np.sqrt(q2.clip(0, None))
|
||||
q3 = np.sqrt(q3.clip(0, None))
|
||||
|
||||
c0 = (q0 >= q1) & (q0 >= q2) & (q0 >= q3)
|
||||
c1 = (q1 >= q0) & (q1 >= q2) & (q1 >= q3)
|
||||
c2 = (q2 >= q0) & (q2 >= q1) & (q2 >= q3)
|
||||
c3 = (q3 >= q0) & (q3 >= q1) & (q3 >= q2)
|
||||
|
||||
q1[c0] *= np.sign(ts[c0, 2, 1] - ts[c0, 1, 2])
|
||||
q2[c0] *= np.sign(ts[c0, 0, 2] - ts[c0, 2, 0])
|
||||
q3[c0] *= np.sign(ts[c0, 1, 0] - ts[c0, 0, 1])
|
||||
|
||||
q0[c1] *= np.sign(ts[c1, 2, 1] - ts[c1, 1, 2])
|
||||
q2[c1] *= np.sign(ts[c1, 1, 0] + ts[c1, 0, 1])
|
||||
q3[c1] *= np.sign(ts[c1, 0, 2] + ts[c1, 2, 0])
|
||||
|
||||
q0[c2] *= np.sign(ts[c2, 0, 2] - ts[c2, 2, 0])
|
||||
q1[c2] *= np.sign(ts[c2, 1, 0] + ts[c2, 0, 1])
|
||||
q3[c2] *= np.sign(ts[c2, 2, 1] + ts[c2, 1, 2])
|
||||
|
||||
q0[c3] *= np.sign(ts[c3, 1, 0] - ts[c3, 0, 1])
|
||||
q1[c3] *= np.sign(ts[c3, 2, 0] + ts[c3, 0, 2])
|
||||
q2[c3] *= np.sign(ts[c3, 2, 1] + ts[c3, 1, 2])
|
||||
|
||||
qs = np.empty(ts.shape[:-2] + (4,))
|
||||
qs[..., 0] = q0
|
||||
qs[..., 1] = q1
|
||||
qs[..., 2] = q2
|
||||
qs[..., 3] = q3
|
||||
|
||||
return cls(qs)
|
||||
@@ -0,0 +1,4 @@
|
||||
import sys
|
||||
import os
|
||||
BASEPATH = os.path.dirname(__file__)
|
||||
sys.path.insert(0, BASEPATH)
|
||||
+60
@@ -0,0 +1,60 @@
|
||||
smpl_body_model: None # _DATA/data/mano
|
||||
|
||||
offsets:
|
||||
right: [[ 0.0957, 0.0064, 0.0062],
|
||||
[ 0.0076, 0.0012, 0.0269],
|
||||
[-0.0251, 0.0052, 0.0291],
|
||||
[-0.0473, 0.0039, 0.0290],
|
||||
[ 0.0010, 0.0049, 0.0028],
|
||||
[-0.0302, 0.0068, -0.0028],
|
||||
[-0.0531, 0.0055, -0.0067],
|
||||
[ 0.0269, -0.0036, -0.0370],
|
||||
[ 0.0099, -0.0035, -0.0495],
|
||||
[-0.0060, -0.0042, -0.0599],
|
||||
[ 0.0139, 0.0024, -0.0205],
|
||||
[-0.0144, 0.0045, -0.0256],
|
||||
[-0.0379, 0.0028, -0.0332],
|
||||
[ 0.0716, -0.0091, 0.0320],
|
||||
[ 0.0519, -0.0082, 0.0557],
|
||||
[ 0.0297, -0.0137, 0.0702]]
|
||||
|
||||
|
||||
left: [
|
||||
[-0.000876, 0.909315, 0.027821],
|
||||
[ 0.071361, -0.089584, -0.008046],
|
||||
[ 0.030669, -0.364209, -0.006689],
|
||||
[-0.011554, -0.383348, -0.043502],
|
||||
[ 0.023338, -0.054645, 0.114370],
|
||||
[-0.069012, -0.088960, -0.004796],
|
||||
[-0.036152, -0.370650, -0.009185],
|
||||
[ 0.014029, -0.383638, -0.041892],
|
||||
[-0.022043, -0.046410, 0.117900],
|
||||
[-0.002508, 0.103257, -0.022185],
|
||||
[ 0.003581, 0.127658, -0.001713],
|
||||
[ 0.002027, 0.049072, 0.027867],
|
||||
[-0.001963, 0.208243, -0.049765],
|
||||
[ 0.003517, 0.062322, 0.050205],
|
||||
[ 0.075298, 0.117780, -0.036875],
|
||||
[ 0.085317, 0.031739, -0.007293],
|
||||
[ 0.251247, -0.011967, -0.027518],
|
||||
[ 0.238019, 0.009007, 0.000044],
|
||||
[ 0.079668, -0.009683, -0.013206],
|
||||
[-0.077033, 0.115606, -0.041811],
|
||||
[-0.089203, 0.032785, -0.009802],
|
||||
[-0.245990, -0.013152, -0.020162],
|
||||
[-0.245177, 0.008622, -0.003532],
|
||||
[-0.080400, -0.007248, -0.010419]
|
||||
]
|
||||
|
||||
parents: [-1, 0, 1, 2, 0, 4, 5, 0, 7, 8, 0, 10, 11, 0, 13, 14]
|
||||
|
||||
joint_names: ['wrist', 'index1', 'index2', 'index3', 'middle1', 'middle2', 'middle3', 'pinky1', 'pinky2', 'pinky3',
|
||||
'ring1', 'ring2', 'ring3', 'thumb1', 'thumb2', 'thumb3']
|
||||
|
||||
|
||||
joints_to_use: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]
|
||||
|
||||
leaf_joints: [16, 17, 18, 19, 20]
|
||||
|
||||
wrist_index: [0]
|
||||
|
||||
@@ -0,0 +1,165 @@
|
||||
import math
|
||||
import os
|
||||
from abc import ABC, abstractmethod
|
||||
|
||||
import torch
|
||||
import torch.nn.init
|
||||
import torch.optim
|
||||
|
||||
|
||||
def weights_init(init_type='default'):
|
||||
def init_fun(m):
|
||||
classname = m.__class__.__name__
|
||||
if (classname.find('Conv') == 0 or classname.find('Linear') == 0) and hasattr(m, 'weight'):
|
||||
if init_type == 'gaussian':
|
||||
torch.nn.init.normal_(m.weight.data)
|
||||
elif init_type == 'xavier':
|
||||
torch.nn.init.xavier_normal_(m.weight.data, gain=math.sqrt(2))
|
||||
elif init_type == 'kaiming':
|
||||
torch.nn.init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')
|
||||
elif init_type == 'orthogonal':
|
||||
torch.nn.init.orthogonal_(m.weight.data, gain=math.sqrt(2))
|
||||
elif init_type == 'default':
|
||||
torch.nn.init.normal_(m.weight.data)
|
||||
weight_norm = m.weight.pow(2).sum(2, keepdim=True).sum(1, keepdim=True).add(1e-8).sqrt()
|
||||
m.weight.data.div_(weight_norm)
|
||||
else:
|
||||
raise NotImplementedError(f"Unsupported initialization: {init_type}")
|
||||
if hasattr(m, 'bias') and m.bias is not None:
|
||||
torch.nn.init.constant_(m.bias.data, 0.0)
|
||||
return init_fun
|
||||
|
||||
|
||||
class BaseModel(ABC):
|
||||
"""This class is an abstract base class (ABC) for models.
|
||||
To create a subclass, you need to implement the following five functions:
|
||||
-- <__init__>: initialize the class; first call BaseModel.__init__(self, opt).
|
||||
-- <set_input>: unpack data from dataset and apply preprocessing.
|
||||
-- <forward>: produce intermediate results.
|
||||
-- <optimize_parameters>: calculate losses, gradients, and update network weights.
|
||||
"""
|
||||
|
||||
def __init__(self, args):
|
||||
self.args = args
|
||||
self.is_train = args.is_train
|
||||
self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
||||
self.model_save_dir = os.path.join(args.save_dir, 'checkpoints') # save all the checkpoints to save_dir
|
||||
|
||||
if self.is_train and args.log:
|
||||
import time
|
||||
from tensorboardX import SummaryWriter
|
||||
from .loss_record import LossRecorder
|
||||
|
||||
log_dir = os.path.join(args.save_dir, 'logs')
|
||||
if args.epoch_begin != 0:
|
||||
all = [d for d in os.listdir(log_dir)]
|
||||
if len(all) == 0:
|
||||
raise RuntimeError(f'Empty logging path {log_dir}')
|
||||
timestamp = sorted(all)[-1]
|
||||
else:
|
||||
timestamp = time.strftime("%Y-%m-%d-%H:%M:%S", time.localtime(time.time()))
|
||||
self.log_path = os.path.join(log_dir, timestamp)
|
||||
os.makedirs(self.log_path, exist_ok=True)
|
||||
self.writer = SummaryWriter(self.log_path)
|
||||
self.loss_recorder = LossRecorder(self.writer)
|
||||
|
||||
self.epoch_cnt = 0
|
||||
self.schedulers = []
|
||||
self.optimizers = []
|
||||
self.models = []
|
||||
|
||||
@abstractmethod
|
||||
def set_input(self, input):
|
||||
"""Unpack input data from the dataloader and perform necessary pre-processing steps.
|
||||
Parameters:
|
||||
input (dict): includes the data itself and its metadata information.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def compute_test_result(self):
|
||||
"""
|
||||
After forward, do something like output bvh, get error value
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def forward(self):
|
||||
"""Run forward pass; called by both functions <optimize_parameters> and <test>."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def optimize_parameters(self):
|
||||
"""Calculate losses, gradients, and update network weights; called in every training iteration"""
|
||||
pass
|
||||
|
||||
def get_scheduler(self, optimizer):
|
||||
if self.args.scheduler.name == 'linear':
|
||||
def lambda_rule(epoch):
|
||||
lr_l = 1.0 - max(0, epoch - self.args.n_epochs_origin) / float(self.args.n_epochs_decay + 1)
|
||||
return lr_l
|
||||
return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda_rule)
|
||||
if self.args.scheduler.name == 'StepLR':
|
||||
print('StepLR scheduler set')
|
||||
return torch.optim.lr_scheduler.StepLR(optimizer, self.args.scheduler.step_size, self.args.scheduler.gamma, verbose=self.args.verbose)
|
||||
if self.args.scheduler.name == 'Plateau':
|
||||
print('ReduceLROnPlateau shceduler set')
|
||||
return torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=self.args.scheduler.factor, threshold=self.args.scheduler.threshold, patience=5, verbose=True)
|
||||
if self.args.scheduler.name == 'MultiStep':
|
||||
print('MultiStep shceduler set')
|
||||
return torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=self.args.scheduler.milestones, gamma=self.args.scheduler.gamma, verbose=self.args.verbose)
|
||||
else:
|
||||
print('No scheduler set')
|
||||
return None
|
||||
|
||||
def setup(self):
|
||||
"""Load and print networks; create schedulers
|
||||
Parameters:
|
||||
opt (Option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions
|
||||
"""
|
||||
if self.is_train:
|
||||
self.schedulers = [self.get_scheduler(optimizer) for optimizer in self.optimizers]
|
||||
|
||||
def epoch(self):
|
||||
if self.args.verbose:
|
||||
self.loss_recorder.epoch()
|
||||
for scheduler in self.schedulers:
|
||||
if scheduler is not None:
|
||||
if self.args.scheduler.name == 'Plateau':
|
||||
loss = self.loss_recorder.losses['total_loss'].current()
|
||||
scheduler.step(loss[1])
|
||||
else:
|
||||
scheduler.step()
|
||||
self.epoch_cnt += 1
|
||||
|
||||
def test(self):
|
||||
"""Forward function used in test time.
|
||||
This function wraps <forward> function in no_grad() so we don't save intermediate steps for backprop
|
||||
It also calls <compute_visuals> to produce additional visualization results
|
||||
"""
|
||||
with torch.no_grad():
|
||||
self.forward()
|
||||
self.compute_test_result()
|
||||
|
||||
def train(self):
|
||||
for model in self.models:
|
||||
model.train()
|
||||
for param in model.parameters():
|
||||
param.requires_grad = True
|
||||
|
||||
def eval(self):
|
||||
for model in self.models:
|
||||
model.eval()
|
||||
for param in model.parameters():
|
||||
param.requires_grad = False
|
||||
|
||||
def print(self):
|
||||
for model in self.models:
|
||||
print(model)
|
||||
|
||||
def count_params(self):
|
||||
params = 0
|
||||
for model in self.models:
|
||||
params += sum(p.numel() for p in model.parameters() if p.requires_grad)
|
||||
|
||||
return params
|
||||
+291
@@ -0,0 +1,291 @@
|
||||
import os
|
||||
|
||||
import holden.BVH as BVH
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from holden.Animation import Animation
|
||||
from holden.Quaternions import Quaternions
|
||||
from human_body_prior.tools.omni_tools import copy2cpu as c2c
|
||||
from rotations import matrix_to_quaternion, quaternion_to_axis_angle, rotation_6d_to_matrix
|
||||
from utils import align_joints, compute_trajectory
|
||||
|
||||
from .base_model import BaseModel
|
||||
from .fk import ForwardKinematicsLayer
|
||||
from .losses import GeodesicLoss
|
||||
from .neural_motion import NeuralMotionField
|
||||
|
||||
|
||||
class Architecture(BaseModel):
|
||||
def __init__(self, args, ngpu):
|
||||
super(Architecture, self).__init__(args)
|
||||
|
||||
self.args = args
|
||||
|
||||
if args.amass_data:
|
||||
self.fk = ForwardKinematicsLayer(args)
|
||||
else:
|
||||
self.fk = ForwardKinematicsLayer(parents=args.animal.parents, positions=args.animal.offsets)
|
||||
|
||||
self.field = NeuralMotionField(args.nemf).to(self.device)
|
||||
if args.multi_gpu is True:
|
||||
self.field = nn.DataParallel(self.field, device_ids=range(ngpu))
|
||||
|
||||
self.models = [self.field]
|
||||
|
||||
ordermap = {
|
||||
'x': 0,
|
||||
'y': 1,
|
||||
'z': 2,
|
||||
}
|
||||
self.up_axis = ordermap[self.args.data.up]
|
||||
self.v_axis = [x for x in ordermap.values() if x != self.up_axis]
|
||||
|
||||
print(self.up_axis)
|
||||
print(self.v_axis)
|
||||
|
||||
self.input_data = dict()
|
||||
self.recon_data = dict()
|
||||
|
||||
if self.is_train:
|
||||
self.optimizer = torch.optim.Adam(self.field.parameters(), args.learning_rate)
|
||||
self.optimizers = [self.optimizer]
|
||||
self.criterion_rec = nn.L1Loss() if args.l1_loss else nn.MSELoss()
|
||||
|
||||
self.criterion_geo = GeodesicLoss().to(self.device)
|
||||
self.fps = args.data.fps
|
||||
if args.bvh_viz:
|
||||
self.viz_dir = os.path.join(args.save_dir, 'results', 'bvh')
|
||||
else:
|
||||
self.viz_dir = os.path.join(args.save_dir, 'results', 'smpl')
|
||||
|
||||
def set_input(self, input):
|
||||
self.input_data = {k: v.float().to(self.device) for k, v in input.items() if k in ['pos', 'global_xform', 'root_orient', 'root_vel', 'trans']}
|
||||
self.input_data['rotmat'] = rotation_6d_to_matrix(self.input_data['global_xform'])
|
||||
|
||||
def forward(self, step=1):
|
||||
b_size, t_length, n_joints = self.input_data['pos'].shape[:3]
|
||||
|
||||
t = torch.arange(start=0, end=t_length, step=step).unsqueeze(0) # (1, T)
|
||||
t = t / t_length * 2 - 1 # map t to [-1, 1]
|
||||
t = t.expand(b_size, -1).unsqueeze(-1).to(self.device) # (B, T, 1)
|
||||
f, _ = self.field(t) # (B, T, D), rot6d + root orient + root vel + root height
|
||||
# import ipdb; ipdb.set_trace()
|
||||
self.recon_data = dict()
|
||||
rot6d_recon = f[:, :, :n_joints * 6] # (B, T, J x 6)
|
||||
rot6d_recon = rot6d_recon.view(b_size, -1, n_joints, 6) # (B, T, J, 6)
|
||||
rotmat_recon = rotation_6d_to_matrix(rot6d_recon) # (B, T, J, 3, 3)
|
||||
if self.args.data.root_transform:
|
||||
identity = torch.zeros(b_size, t.shape[1], 3, 3).to(self.device)
|
||||
identity[:, :, 0, 0] = 1
|
||||
identity[:, :, 1, 1] = 1
|
||||
identity[:, :, 2, 2] = 1
|
||||
rotmat_recon[:, :, 0] = identity
|
||||
# rotmat_recon[:, :, -2] = rotmat_recon[:, :, -4]
|
||||
# rotmat_recon[:, :, -1] = rotmat_recon[:, :, -3]
|
||||
local_rotmat = self.fk.global_to_local(rotmat_recon.view(-1, n_joints, 3, 3)) # (B x T, J, 3. 3)
|
||||
pos_recon, _ = self.fk(local_rotmat) # (B x T, J, 3)
|
||||
pos_recon = pos_recon.contiguous().view(b_size, -1, n_joints, 3) # (B, T, J, 3)
|
||||
self.recon_data['rotmat'] = rotmat_recon
|
||||
self.recon_data['pos'] = pos_recon
|
||||
|
||||
if self.args.data.root_transform:
|
||||
root_orient = f[:, :, n_joints * 6:n_joints * 6 + 6] # (B, T, 6)
|
||||
self.recon_data['root_orient'] = root_orient
|
||||
else:
|
||||
self.recon_data['root_orient'] = rot6d_recon[:, :, 0] # (B, T, 6)
|
||||
|
||||
root_vel = f[:, :, -4:-1] # (B, T, 3)
|
||||
root_height = f[:, :, -1] # (B, T)
|
||||
|
||||
self.recon_data['root_vel'] = root_vel
|
||||
self.recon_data['root_height'] = root_height
|
||||
|
||||
dt = 1.0 / self.args.data.fps * step
|
||||
origin = torch.zeros(b_size, 3).to(self.device)
|
||||
trans = compute_trajectory(root_vel, root_height, origin, dt, up_axis=self.args.data.up)
|
||||
self.recon_data['trans'] = trans
|
||||
|
||||
def backward(self, validation=False):
|
||||
loss = 0
|
||||
|
||||
root_orient = rotation_6d_to_matrix(self.recon_data['root_orient'])
|
||||
root_orient_gt = rotation_6d_to_matrix(self.input_data['root_orient'])
|
||||
if self.args.geodesic_loss:
|
||||
orient_recon_loss = self.criterion_geo(root_orient.view(-1, 3, 3), root_orient_gt.view(-1, 3, 3))
|
||||
else:
|
||||
orient_recon_loss = self.criterion_rec(root_orient, root_orient_gt)
|
||||
self.loss_recorder.add_scalar('orient_recon_loss', orient_recon_loss, validation=validation)
|
||||
loss += self.args.lambda_orient * orient_recon_loss
|
||||
|
||||
if self.args.geodesic_loss:
|
||||
rotmat_recon_loss = self.criterion_geo(self.recon_data['rotmat'].view(-1, 3, 3), self.input_data['rotmat'].view(-1, 3, 3))
|
||||
else:
|
||||
rotmat_recon_loss = self.criterion_rec(self.recon_data['rotmat'], self.input_data['rotmat'])
|
||||
self.loss_recorder.add_scalar('rotmat_recon_loss', rotmat_recon_loss, validation=validation)
|
||||
loss += self.args.lambda_rotmat * rotmat_recon_loss
|
||||
|
||||
pos_recon_loss = self.criterion_rec(self.recon_data['pos'], self.input_data['pos'])
|
||||
self.loss_recorder.add_scalar('pos_recon_loss', pos_recon_loss, validation=validation)
|
||||
loss += self.args.lambda_pos * pos_recon_loss
|
||||
|
||||
v_loss = self.criterion_rec(self.recon_data['root_vel'], self.input_data['root_vel'])
|
||||
up_loss = self.criterion_rec(self.recon_data['root_height'], self.input_data['trans'][..., 'xyz'.index(self.args.data.up)])
|
||||
self.loss_recorder.add_scalar('v_loss', v_loss, validation=validation)
|
||||
self.loss_recorder.add_scalar('up_loss', up_loss, validation=validation)
|
||||
loss += self.args.lambda_v * v_loss + self.args.lambda_up * up_loss
|
||||
|
||||
origin = self.input_data['trans'][:, 0]
|
||||
trans = self.recon_data['trans']
|
||||
trans[..., self.v_axis] = trans[..., self.v_axis] + origin[..., self.v_axis].unsqueeze(1)
|
||||
trans_loss = self.criterion_rec(trans, self.input_data['trans'])
|
||||
self.loss_recorder.add_scalar('trans_loss', trans_loss, validation=validation)
|
||||
loss += self.args.lambda_trans * trans_loss
|
||||
|
||||
self.loss_recorder.add_scalar('total_loss', loss, validation=validation)
|
||||
|
||||
if not validation:
|
||||
loss.backward()
|
||||
|
||||
def optimize_parameters(self):
|
||||
self.optimizer.zero_grad()
|
||||
self.forward()
|
||||
self.backward(validation=False)
|
||||
self.optimizer.step()
|
||||
|
||||
def report_errors(self):
|
||||
root_orient = rotation_6d_to_matrix(self.recon_data['root_orient']) # (B, T, 3, 3)
|
||||
root_orient_gt = rotation_6d_to_matrix(self.input_data['root_orient']) # (B, T, 3, 3)
|
||||
orientation_error = self.criterion_geo(root_orient.view(-1, 3, 3), root_orient_gt.view(-1, 3, 3))
|
||||
|
||||
local_rotmat = self.fk.global_to_local(self.recon_data['rotmat'].view(-1, self.args.smpl.joint_num, 3, 3)) # (B x T, J, 3, 3)
|
||||
local_rotmat_gt = self.fk.global_to_local(self.input_data['rotmat'].view(-1, self.args.smpl.joint_num, 3, 3)) # (B x T, J, 3, 3)
|
||||
rotation_error = self.criterion_geo(local_rotmat[:, 1:].reshape(-1, 3, 3), local_rotmat_gt[:, 1:].reshape(-1, 3, 3))
|
||||
|
||||
pos = self.recon_data['pos'] # (B, T, J, 3)
|
||||
pos_gt = self.input_data['pos'] # (B, T, J, 3)
|
||||
position_error = torch.linalg.norm((pos - pos_gt), dim=-1).mean()
|
||||
|
||||
origin = self.input_data['trans'][:, 0]
|
||||
trans = self.recon_data['trans']
|
||||
trans[..., self.v_axis] = trans[..., self.v_axis] + origin[..., self.v_axis].unsqueeze(1)
|
||||
trans_gt = self.input_data['trans'] # (B, T, 3)
|
||||
translation_error = torch.linalg.norm((trans - trans_gt), dim=-1).mean()
|
||||
|
||||
return {
|
||||
'rotation': c2c(rotation_error),
|
||||
'position': c2c(position_error),
|
||||
'orientation': c2c(orientation_error),
|
||||
'translation': c2c(translation_error)
|
||||
}
|
||||
|
||||
def verbose(self):
|
||||
res = {}
|
||||
for loss in self.loss_recorder.losses.values():
|
||||
res[loss.name] = {'train': loss.current()[0]}
|
||||
|
||||
return res
|
||||
|
||||
def save(self, optimal=False):
|
||||
if optimal:
|
||||
path = os.path.join(self.args.save_dir, 'results', 'model')
|
||||
else:
|
||||
path = os.path.join(self.model_save_dir, f'{self.epoch_cnt:04d}')
|
||||
|
||||
os.makedirs(path, exist_ok=True)
|
||||
nemf = self.field.module if isinstance(self.field, nn.DataParallel) else self.field
|
||||
torch.save(nemf.state_dict(), os.path.join(path, 'nemf.pth'))
|
||||
torch.save(self.optimizer.state_dict(), os.path.join(path, 'optimizer.pth'))
|
||||
if self.args.scheduler.name:
|
||||
torch.save(self.schedulers[0].state_dict(), os.path.join(path, 'scheduler.pth'))
|
||||
self.loss_recorder.save(path)
|
||||
|
||||
print(f'Save at {path} succeeded')
|
||||
|
||||
def load(self, epoch=None, optimal=False):
|
||||
if optimal:
|
||||
path = os.path.join(self.args.save_dir, 'results', 'model')
|
||||
else:
|
||||
if epoch is None:
|
||||
all = [int(q) for q in os.listdir(self.model_save_dir)]
|
||||
if len(all) == 0:
|
||||
raise RuntimeError(f'Empty loading path {self.model_save_dir}')
|
||||
epoch = sorted(all)[-1]
|
||||
path = os.path.join(self.model_save_dir, f'{epoch:04d}')
|
||||
|
||||
print(f'Loading from {path}')
|
||||
nemf = self.field.module if isinstance(self.field, nn.DataParallel) else self.field
|
||||
nemf.load_state_dict(torch.load(os.path.join(path, 'nemf.pth'), map_location=self.device))
|
||||
if self.is_train:
|
||||
self.optimizer.load_state_dict(torch.load(os.path.join(path, 'optimizer.pth')))
|
||||
self.loss_recorder.load(path)
|
||||
if self.args.scheduler.name:
|
||||
self.schedulers[0].load_state_dict(torch.load(os.path.join(path, 'scheduler.pth')))
|
||||
self.epoch_cnt = epoch if not optimal else 0
|
||||
|
||||
print('Load succeeded')
|
||||
|
||||
def super_sampling(self, step=1):
|
||||
with torch.no_grad():
|
||||
self.forward(step)
|
||||
self.fps = self.args.data.fps / step
|
||||
self.compute_test_result()
|
||||
|
||||
def compute_test_result(self):
|
||||
os.makedirs(self.viz_dir, exist_ok=True)
|
||||
|
||||
rotmat = self.recon_data['rotmat'][0] # (T, J, 3, 3)
|
||||
rotmat_gt = self.input_data['rotmat'][0] # (T, J, 3, 3)
|
||||
local_rotmat = self.fk.global_to_local(rotmat) # (T, J, 3, 3)
|
||||
local_rotmat_gt = self.fk.global_to_local(rotmat_gt) # (T, J, 3, 3)
|
||||
|
||||
if self.args.data.root_transform:
|
||||
root_orient = rotation_6d_to_matrix(self.recon_data['root_orient'][0]) # (T, 3, 3)
|
||||
root_orient_gt = rotation_6d_to_matrix(self.input_data['root_orient'][0]) # (T, 3, 3)
|
||||
local_rotmat[:, 0] = root_orient
|
||||
local_rotmat_gt[:, 0] = root_orient_gt
|
||||
rotation = matrix_to_quaternion(local_rotmat) # (T, J, 4)
|
||||
rotation_gt = matrix_to_quaternion(local_rotmat_gt) # (T, J, 4)
|
||||
|
||||
origin = self.input_data['trans'][0, 0].clone() # [3]
|
||||
position = self.recon_data['trans'][0].clone() # (T, 3)
|
||||
position[:, self.v_axis] = position[:, self.v_axis] + origin[self.v_axis].unsqueeze(0)
|
||||
position_gt = self.input_data['trans'][0] # (T, 3)
|
||||
|
||||
if self.args.bvh_viz:
|
||||
if self.args.amass_data:
|
||||
rotation = align_joints(rotation)
|
||||
rotation_gt = align_joints(rotation_gt)
|
||||
|
||||
position = position.unsqueeze(1) # (T, 1, 3)
|
||||
position_gt = position_gt.unsqueeze(1) # (T, 1, 3)
|
||||
|
||||
# export data to BVH files
|
||||
if self.args.amass_data:
|
||||
offsets = self.args.smpl.offsets[self.args.data.gender]
|
||||
parents = self.args.smpl.parents
|
||||
names = self.args.smpl.joint_names
|
||||
else:
|
||||
offsets = self.args.animal.offsets
|
||||
parents = self.args.animal.parents
|
||||
names = self.args.animal.joint_names
|
||||
|
||||
anim = Animation(Quaternions(c2c(rotation)), c2c(position), None, offsets=offsets, parents=parents)
|
||||
BVH.save(os.path.join(self.viz_dir, f'test_{int(self.fps)}fps.bvh'), anim, names=names, frametime=1 / self.fps)
|
||||
|
||||
anim.rotations = Quaternions(c2c(rotation_gt))
|
||||
anim.positions = c2c(position_gt)
|
||||
BVH.save(os.path.join(self.viz_dir, 'test_gt.bvh'), anim, names=names, frametime=1 / self.args.data.fps)
|
||||
else:
|
||||
poses = c2c(quaternion_to_axis_angle(rotation)) # (T, J, 3)
|
||||
poses_gt = c2c(quaternion_to_axis_angle(rotation_gt)) # (T, J, 3)
|
||||
|
||||
poses = poses.reshape((poses.shape[0], -1)) # (T, J x 3)
|
||||
poses = np.pad(poses, [(0, 0), (0, 93)], mode='constant')
|
||||
poses_gt = poses_gt.reshape((poses_gt.shape[0], -1)) # (T, J x 3)
|
||||
poses_gt = np.pad(poses_gt, [(0, 0), (0, 93)], mode='constant')
|
||||
|
||||
np.savez(os.path.join(self.viz_dir, f'test_{int(self.fps)}fps.npz'),
|
||||
poses=poses, trans=c2c(position), betas=np.zeros(10), gender=self.args.data.gender, mocap_framerate=self.fps)
|
||||
np.savez(os.path.join(self.viz_dir, 'test_gt.npz'),
|
||||
poses=poses_gt, trans=c2c(position_gt), betas=np.zeros(10), gender=self.args.data.gender, mocap_framerate=self.args.data.fps)
|
||||
+152
@@ -0,0 +1,152 @@
|
||||
"""Based on Daniel Holden code from:
|
||||
A Deep Learning Framework for Character Motion Synthesis and Editing
|
||||
(http://www.ipab.inf.ed.ac.uk/cgvu/motionsynthesis.pdf)
|
||||
"""
|
||||
|
||||
import os
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from rotations import euler_angles_to_matrix, quaternion_to_matrix, rotation_6d_to_matrix
|
||||
|
||||
|
||||
class ForwardKinematicsLayer(nn.Module):
|
||||
""" Forward Kinematics Layer Class """
|
||||
|
||||
def __init__(self, args=None, parents=None, positions=None, device=None):
|
||||
super().__init__()
|
||||
self.b_idxs = None
|
||||
if device is None:
|
||||
self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
||||
else:
|
||||
self.device = device
|
||||
if parents is None and positions is None:
|
||||
# Load SMPL skeleton (their joint order is different from the one we use for bvh export)
|
||||
# smpl_fname = os.path.join(args.smpl.smpl_body_model, args.data.gender, 'model.npz')
|
||||
smpl_fname = os.path.join(args.smpl.smpl_body_model, "MANO_RIGHT.pkl")
|
||||
if not os.path.exists(smpl_fname):
|
||||
raise FileNotFoundError(f'SMPL file {os.path.abspath(smpl_fname)} does not exist. Please check the path., pwd: {os.getcwd()}')
|
||||
smpl_data = np.load(smpl_fname, encoding='latin1', allow_pickle=True)
|
||||
self.parents = torch.from_numpy(smpl_data['kintree_table'][0].astype(np.int32)).to(self.device)
|
||||
self.parents = self.parents.long()
|
||||
self.positions = torch.from_numpy(smpl_data['J'].astype(np.float32)).to(self.device)
|
||||
self.positions[1:] -= self.positions[self.parents[1:]]
|
||||
else:
|
||||
self.parents = torch.from_numpy(parents).to(self.device)
|
||||
self.parents = self.parents.long()
|
||||
self.positions = torch.from_numpy(positions).to(self.device)
|
||||
self.positions = self.positions.float()
|
||||
self.positions[0] = 0
|
||||
|
||||
def rotate(self, t0s, t1s):
|
||||
return torch.matmul(t0s, t1s)
|
||||
|
||||
def identity_rotation(self, rotations):
|
||||
diagonal = torch.diag(torch.tensor([1.0, 1.0, 1.0, 1.0])).to(self.device)
|
||||
diagonal = torch.reshape(
|
||||
diagonal, torch.Size([1] * len(rotations.shape[:2]) + [4, 4]))
|
||||
ts = diagonal.repeat(rotations.shape[:2] + torch.Size([1, 1]))
|
||||
return ts
|
||||
|
||||
def make_fast_rotation_matrices(self, positions, rotations):
|
||||
if len(rotations.shape) == 4 and rotations.shape[-2:] == torch.Size([3, 3]):
|
||||
rot_matrices = rotations
|
||||
elif rotations.shape[-1] == 3:
|
||||
rot_matrices = euler_angles_to_matrix(rotations, convention='XYZ')
|
||||
elif rotations.shape[-1] == 4:
|
||||
rot_matrices = quaternion_to_matrix(rotations)
|
||||
elif rotations.shape[-1] == 6:
|
||||
rot_matrices = rotation_6d_to_matrix(rotations)
|
||||
else:
|
||||
raise NotImplementedError(f'Unimplemented rotation representation in FK layer, shape of {rotations.shape}')
|
||||
|
||||
rot_matrices = torch.cat([rot_matrices, positions[..., None]], dim=-1)
|
||||
zeros = torch.zeros(rot_matrices.shape[:-2] + torch.Size([1, 3])).to(self.device)
|
||||
ones = torch.ones(rot_matrices.shape[:-2] + torch.Size([1, 1])).to(self.device)
|
||||
zerosones = torch.cat([zeros, ones], dim=-1)
|
||||
rot_matrices = torch.cat([rot_matrices, zerosones], dim=-2)
|
||||
return rot_matrices
|
||||
|
||||
def rotate_global(self, parents, positions, rotations):
|
||||
locals = self.make_fast_rotation_matrices(positions, rotations)
|
||||
globals = self.identity_rotation(rotations)
|
||||
|
||||
globals = torch.cat([locals[:, 0:1], globals[:, 1:]], dim=1)
|
||||
b_size = positions.shape[0]
|
||||
if self.b_idxs is None:
|
||||
self.b_idxs = torch.LongTensor(np.arange(b_size)).to(self.device)
|
||||
elif self.b_idxs.shape[-1] != b_size:
|
||||
self.b_idxs = torch.LongTensor(np.arange(b_size)).to(self.device)
|
||||
|
||||
for i in range(1, positions.shape[1]):
|
||||
globals[:, i] = self.rotate(
|
||||
globals[self.b_idxs, parents[i]], locals[:, i])
|
||||
|
||||
return globals
|
||||
|
||||
def get_tpose_joints(self, offsets, parents):
|
||||
num_joints = len(parents)
|
||||
joints = [offsets[:, 0]]
|
||||
for j in range(1, len(parents)):
|
||||
joints.append(joints[parents[j]] + offsets[:, j])
|
||||
|
||||
return torch.stack(joints, dim=1)
|
||||
|
||||
def canonical_to_local(self, canonical_xform, global_orient=None):
|
||||
"""
|
||||
Args:
|
||||
canonical_xform: (B, J, 3, 3)
|
||||
global_orient: (B, 3, 3)
|
||||
|
||||
Returns:
|
||||
local_xform: (B, J, 3, 3)
|
||||
"""
|
||||
local_xform = torch.zeros_like(canonical_xform)
|
||||
|
||||
if global_orient is None:
|
||||
global_xform = canonical_xform
|
||||
else:
|
||||
global_xform = torch.matmul(global_orient.unsqueeze(1), canonical_xform)
|
||||
for i in range(global_xform.shape[1]):
|
||||
if i == 0:
|
||||
local_xform[:, i] = global_xform[:, i]
|
||||
else:
|
||||
local_xform[:, i] = torch.bmm(torch.linalg.inv(global_xform[:, self.parents[i]]), global_xform[:, i])
|
||||
|
||||
return local_xform
|
||||
|
||||
def global_to_local(self, global_xform):
|
||||
"""
|
||||
Args:
|
||||
global_xform: (B, J, 3, 3)
|
||||
|
||||
Returns:
|
||||
local_xform: (B, J, 3, 3)
|
||||
"""
|
||||
local_xform = torch.zeros_like(global_xform)
|
||||
|
||||
for i in range(global_xform.shape[1]):
|
||||
if i == 0:
|
||||
local_xform[:, i] = global_xform[:, i]
|
||||
else:
|
||||
local_xform[:, i] = torch.bmm(torch.linalg.inv(global_xform[:, self.parents[i]]), global_xform[:, i])
|
||||
|
||||
return local_xform
|
||||
|
||||
def forward(self, rotations, positions=None):
|
||||
"""
|
||||
Args:
|
||||
rotations (B, J, D)
|
||||
|
||||
Returns:
|
||||
The global position of each joint after FK (B, J, 3)
|
||||
"""
|
||||
# Get the full transform with rotations for skinning
|
||||
b_size = rotations.shape[0]
|
||||
if positions is None:
|
||||
positions = self.positions.repeat(b_size, 1, 1)
|
||||
transforms = self.rotate_global(self.parents, positions, rotations)
|
||||
coordinates = transforms[:, :, :3, 3] / transforms[:, :, 3:, 3]
|
||||
|
||||
return coordinates, transforms
|
||||
@@ -0,0 +1,525 @@
|
||||
import os
|
||||
import yaml
|
||||
|
||||
import holden.BVH as BVH
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from holden.Animation import Animation
|
||||
from holden.Quaternions import Quaternions
|
||||
from human_body_prior.tools.omni_tools import copy2cpu as c2c
|
||||
from rotations import matrix_to_quaternion, matrix_to_rotation_6d, quaternion_to_axis_angle, rotation_6d_to_matrix
|
||||
from utils import align_joints, compute_trajectory, estimate_angular_velocity, estimate_linear_velocity, normalize
|
||||
|
||||
from .base_model import BaseModel
|
||||
from .fk import ForwardKinematicsLayer
|
||||
from .global_motion import GlobalMotionPredictor
|
||||
from .losses import GeodesicLoss
|
||||
from .neural_motion import NeuralMotionField
|
||||
from .prior import GlobalEncoder, LocalEncoder
|
||||
from .skeleton import build_edge_topology
|
||||
|
||||
|
||||
|
||||
class Struct:
|
||||
def __init__(self, **entries):
|
||||
self.__dict__.update(entries)
|
||||
|
||||
class Arguments:
|
||||
def __init__(self, config_path, filename='default.yaml'):
|
||||
with open(os.path.join(config_path, 'mano.yaml'), 'r') as f:
|
||||
smpl = yaml.safe_load(f)
|
||||
|
||||
self.smpl = Struct(**smpl)
|
||||
self.smpl.offsets['right'] = np.array(self.smpl.offsets['right'])
|
||||
self.smpl.offsets['left'] = np.array(self.smpl.offsets['left'])
|
||||
self.smpl.parents = np.array(self.smpl.parents).astype(np.int32)
|
||||
self.smpl.joint_num = 16 # len(self.smpl.joints_to_use)
|
||||
self.smpl.joints_to_use = np.array(self.smpl.joints_to_use)
|
||||
self.smpl.joints_to_use = np.arange(0, 63).reshape((-1, 3))[self.smpl.joints_to_use].reshape(-1)
|
||||
|
||||
self.filename = os.path.splitext(filename)[0]
|
||||
with open(os.path.join(config_path, filename), 'r') as f:
|
||||
config = yaml.safe_load(f)
|
||||
|
||||
for key, value in config.items():
|
||||
if isinstance(value, dict):
|
||||
setattr(self, key, Struct(**value))
|
||||
else:
|
||||
setattr(self, key, value)
|
||||
|
||||
self.json = config
|
||||
|
||||
|
||||
class Architecture(BaseModel):
|
||||
def __init__(self, args, ngpu, batch_num=0):
|
||||
super(Architecture, self).__init__(args)
|
||||
|
||||
self.batch_num = batch_num
|
||||
self.args = args
|
||||
#from pprint import pprint
|
||||
#pprint(vars(self.args))
|
||||
|
||||
self.fk = ForwardKinematicsLayer(args)
|
||||
|
||||
smpl_fname = os.path.join(args.smpl.smpl_body_model, "MANO_RIGHT.pkl")
|
||||
smpl_data = np.load(smpl_fname, allow_pickle=True, encoding='latin1')
|
||||
parents = smpl_data['kintree_table'][0].astype(np.int32)
|
||||
edges = build_edge_topology(parents)
|
||||
|
||||
self.local_encoder = LocalEncoder(args.local_prior, edges).to(self.device)
|
||||
self.global_encoder = GlobalEncoder(args.global_prior).to(self.device)
|
||||
self.field = NeuralMotionField(args.nemf).to(self.device)
|
||||
|
||||
if args.multi_gpu is True:
|
||||
self.local_encoder = nn.DataParallel(self.local_encoder, device_ids=range(ngpu))
|
||||
self.global_encoder = nn.DataParallel(self.global_encoder, device_ids=range(ngpu))
|
||||
self.field = nn.DataParallel(self.field, device_ids=range(ngpu))
|
||||
|
||||
self.models = [self.local_encoder, self.field]
|
||||
if args.data.root_transform:
|
||||
self.models.append(self.global_encoder)
|
||||
|
||||
ordermap = {
|
||||
'x': 0,
|
||||
'y': 1,
|
||||
'z': 2,
|
||||
}
|
||||
self.up_axis = ordermap[self.args.data.up]
|
||||
self.v_axis = [x for x in ordermap.values() if x != self.up_axis]
|
||||
|
||||
print(self.up_axis)
|
||||
print(self.v_axis)
|
||||
|
||||
self.input_data = dict()
|
||||
self.recon_data = dict()
|
||||
print(os.getcwd(), '...', os.path.join(args.dataset_dir, f'mean-{args.data.gender}-{args.data.clip_length}-{args.data.fps}fps.pt'))
|
||||
self.mean = torch.load(os.path.join(args.dataset_dir, f'mean-{args.data.gender}-{args.data.clip_length}-{args.data.fps}fps.pt'), map_location=self.device)
|
||||
print(f'mean: {list(self.mean.keys())}')
|
||||
self.std = torch.load(os.path.join(args.dataset_dir, f'std-{args.data.gender}-{args.data.clip_length}-{args.data.fps}fps.pt'), map_location=self.device)
|
||||
print(f'std: {list(self.std.keys())}')
|
||||
|
||||
#for k, v in self.mean.items():
|
||||
# print(k, v.shape)
|
||||
|
||||
if args.pretrained_gmp:
|
||||
assert args.output_trans is False, 'already output global translation in this model'
|
||||
gmp_args = Arguments('./configs', filename=args.pretrained_gmp)
|
||||
self.gmp = GlobalMotionPredictor(gmp_args, ngpu)
|
||||
self.gmp.load(optimal=True)
|
||||
self.gmp.eval()
|
||||
|
||||
self.criterion_geo = GeodesicLoss().to(self.device)
|
||||
if self.is_train:
|
||||
parameters = list(self.local_encoder.parameters()) + list(self.field.parameters())
|
||||
if args.data.root_transform:
|
||||
parameters += list(self.global_encoder.parameters())
|
||||
|
||||
if args.adam_optimizer:
|
||||
self.optimizer = torch.optim.Adam(parameters, args.learning_rate, weight_decay=args.weight_decay)
|
||||
else:
|
||||
self.optimizer = torch.optim.Adamax(parameters, args.learning_rate, weight_decay=args.weight_decay)
|
||||
self.optimizers = [self.optimizer]
|
||||
|
||||
self.criterion_rec = nn.L1Loss() if args.l1_loss else nn.MSELoss()
|
||||
self.criterion_bce = nn.BCEWithLogitsLoss()
|
||||
|
||||
self.iteration = 0
|
||||
else:
|
||||
self.fps = args.data.fps
|
||||
self.test_index = 0
|
||||
if args.bvh_viz:
|
||||
self.viz_dir = os.path.join(args.save_dir, 'results', 'bvh')
|
||||
else:
|
||||
self.viz_dir = os.path.join(args.save_dir, 'results', 'smpl')
|
||||
|
||||
def set_input(self, input):
|
||||
self.input_data = {k: v.float().to(self.device) for k, v in input.items() if k in [
|
||||
'pos', 'velocity', 'global_xform', 'angular', 'root_orient', 'root_vel', 'trans', 'contacts'
|
||||
]}
|
||||
self.input_data['rotmat'] = rotation_6d_to_matrix(self.input_data['global_xform'])
|
||||
|
||||
def reparameterize(self, mu, logvar):
|
||||
eps = torch.randn_like(mu)
|
||||
std = torch.exp(0.5 * logvar)
|
||||
return mu + eps * std
|
||||
|
||||
def encode_local(self):
|
||||
b_size, t_length = self.input_data['pos'].shape[:2]
|
||||
|
||||
if 'pos' in self.args.data.normalize:
|
||||
pos = normalize(self.input_data['pos'], mean=self.mean['pos'], std=self.std['pos'])
|
||||
else:
|
||||
pos = self.input_data['pos']
|
||||
if 'velocity' in self.args.data.normalize:
|
||||
velocity = normalize(self.input_data['velocity'], mean=self.mean['velocity'], std=self.std['velocity'])
|
||||
else:
|
||||
velocity = self.input_data['velocity']
|
||||
if 'global_xform' in self.args.data.normalize:
|
||||
global_xform = normalize(self.input_data['global_xform'], mean=self.mean['global_xform'], std=self.std['global_xform'])
|
||||
else:
|
||||
global_xform = self.input_data['global_xform']
|
||||
if 'angular' in self.args.data.normalize:
|
||||
angular = normalize(self.input_data['angular'], mean=self.mean['angular'], std=self.std['angular'])
|
||||
else:
|
||||
angular = self.input_data['angular']
|
||||
|
||||
x = torch.cat((pos, velocity, global_xform, angular), dim=-1)
|
||||
x = x.view(b_size, t_length, -1) # (B, T, J x D)
|
||||
x = x.permute(0, 2, 1) # (B, J x D, T)
|
||||
|
||||
mu, logvar = self.local_encoder(x) # (B, D)
|
||||
if self.args.lambda_kl != 0:
|
||||
z = self.reparameterize(mu, logvar)
|
||||
else:
|
||||
z = mu
|
||||
|
||||
return z, mu, logvar
|
||||
|
||||
def encode_global(self):
|
||||
if 'root_orient' in self.args.data.normalize:
|
||||
root_orient = normalize(self.input_data['root_orient'], mean=self.mean['root_orient'], std=self.std['root_orient'])
|
||||
else:
|
||||
root_orient = self.input_data['root_orient']
|
||||
if self.args.global_prior.in_channels == 9:
|
||||
if 'root_vel' in self.args.data.normalize:
|
||||
root_vel = normalize(self.input_data['root_vel'], mean=self.mean['root_vel'], std=self.std['root_vel'])
|
||||
else:
|
||||
root_vel = self.input_data['root_vel']
|
||||
x = torch.cat((root_orient, root_vel), dim=-1) # (B, T, D)
|
||||
elif self.args.global_prior.in_channels == 6:
|
||||
x = root_orient
|
||||
else:
|
||||
raise RuntimeError('unsupported input dimensions for global prior')
|
||||
x = x.permute(0, 2, 1) # (B, D, T)
|
||||
|
||||
mu, logvar = self.global_encoder(x) # (B, D)
|
||||
if self.args.lambda_kl != 0:
|
||||
z = self.reparameterize(mu, logvar)
|
||||
else:
|
||||
z = mu
|
||||
|
||||
return z, mu, logvar
|
||||
|
||||
def forward(self, step=1):
|
||||
z_l, mu_l, logvar_l = self.encode_local() # (B, D)
|
||||
if self.args.data.root_transform:
|
||||
z_g, mu_g, logvar_g = self.encode_global() # (B, D)
|
||||
else:
|
||||
z_g = None
|
||||
|
||||
self.recon_data.clear()
|
||||
self.recon_data = self.decode(z_l, z_g, length=self.args.data.clip_length, step=step)
|
||||
|
||||
self.recon_data['mu_l'] = mu_l
|
||||
self.recon_data['logvar_l'] = logvar_l
|
||||
if self.args.data.root_transform:
|
||||
self.recon_data['mu_g'] = mu_g
|
||||
self.recon_data['logvar_g'] = logvar_g
|
||||
|
||||
def decode(self, z_l, z_g, length, step=1):
|
||||
b_size = z_l.shape[0]
|
||||
n_joints = self.args.smpl.joint_num
|
||||
n_joints = 16
|
||||
|
||||
# generate f(t; z)
|
||||
t = torch.arange(start=0, end=length, step=step).unsqueeze(0) # (1, T)
|
||||
t = t / self.args.data.clip_length * 2 - 1 # map t to [-1, 1]
|
||||
t = t.expand(b_size, -1).unsqueeze(-1).to(self.device) # (B, T, 1)
|
||||
local_motion, global_motion = self.field(t, z_l, z_g) # (B, T, N)
|
||||
|
||||
rot6d_recon = local_motion[:, :, :n_joints * 6].contiguous() # (B, T, J x 6)
|
||||
rot6d_recon = rot6d_recon.view(b_size, -1, n_joints, 6) # (B, T, J, 6)
|
||||
rotmat_recon = rotation_6d_to_matrix(rot6d_recon) # (B, T, J, 3, 3)
|
||||
|
||||
if self.args.data.root_transform:
|
||||
identity = torch.zeros(b_size, t.shape[1], 3, 3).to(self.device)
|
||||
identity[:, :, 0, 0] = 1
|
||||
identity[:, :, 1, 1] = 1
|
||||
identity[:, :, 2, 2] = 1
|
||||
rotmat_recon[:, :, 0] = identity
|
||||
# rotmat_recon[:, :, -2] = rotmat_recon[:, :, -4]
|
||||
# rotmat_recon[:, :, -1] = rotmat_recon[:, :, -3]
|
||||
local_rotmat = self.fk.global_to_local(rotmat_recon.view(-1, n_joints, 3, 3)) # (B x T, J, 3. 3)
|
||||
pos_recon, _ = self.fk(local_rotmat) # (B x T, J, 3)
|
||||
pos_recon = pos_recon.contiguous().view(b_size, -1, n_joints, 3) # (B, T, J, 3)
|
||||
|
||||
output = dict()
|
||||
output['rotmat'] = rotmat_recon
|
||||
output['pos'] = pos_recon
|
||||
|
||||
if self.args.data.root_transform:
|
||||
root_orient = global_motion[:, :, :6] # (B, T, 6)
|
||||
output['root_orient'] = root_orient
|
||||
else:
|
||||
output['root_orient'] = rot6d_recon[:, :, 0] # (B, T, 6)
|
||||
|
||||
dt = 1.0 / self.args.data.fps * step
|
||||
if self.args.output_trans:
|
||||
root_vel = global_motion[:, :, 6:9] # (B, T, 3)
|
||||
root_height = global_motion[:, :, 9] # (B, T)
|
||||
|
||||
output['root_vel'] = root_vel
|
||||
output['root_height'] = root_height
|
||||
|
||||
contacts = global_motion[:, :, 10:] # (B, T, 8)
|
||||
output['contacts'] = contacts
|
||||
|
||||
origin = torch.zeros(b_size, 3).to(self.device)
|
||||
trans = compute_trajectory(root_vel, root_height, origin, dt, up_axis=self.args.data.up)
|
||||
output['trans'] = trans
|
||||
elif self.args.pretrained_gmp:
|
||||
local_rotmat = local_rotmat.view(b_size, -1, n_joints, 3, 3) # (B, T, J, 3, 3)
|
||||
if self.args.data.root_transform:
|
||||
root_orient_rotmat = rotation_6d_to_matrix(output['root_orient']) # (B, T, 3, 3)
|
||||
local_rotmat[:, :, 0] = root_orient_rotmat
|
||||
pos = torch.matmul(root_orient_rotmat.unsqueeze(2), pos_recon.unsqueeze(-1)).squeeze(-1)
|
||||
else:
|
||||
pos = pos_recon
|
||||
|
||||
gmp_data = dict()
|
||||
gmp_data['rot6d'] = matrix_to_rotation_6d(local_rotmat) # (B, T, J, 6)
|
||||
gmp_data['angular'] = estimate_angular_velocity(local_rotmat.clone(), dt=dt) # (B, T, J, 3)
|
||||
|
||||
gmp_data['pos'] = pos
|
||||
gmp_data['velocity'] = estimate_linear_velocity(pos, dt=dt)
|
||||
|
||||
pred_data = self.gmp.predict(gmp_data, dt=dt)
|
||||
|
||||
output |= pred_data # merge two dicts
|
||||
|
||||
return output
|
||||
|
||||
def kl_scheduler(self):
|
||||
"""
|
||||
Cyclical Annealing Schedule
|
||||
"""
|
||||
if (self.epoch_cnt % self.args.annealing_cycles == 0) and (self.iteration % self.batch_num == 0):
|
||||
print('KL annealing restart')
|
||||
return 0.01
|
||||
|
||||
return min(1, 0.01 + self.iteration / (self.args.annealing_warmup * self.batch_num))
|
||||
|
||||
def backward(self, validation=False):
|
||||
loss = 0
|
||||
|
||||
root_orient = rotation_6d_to_matrix(self.recon_data['root_orient'])
|
||||
root_orient_gt = rotation_6d_to_matrix(self.input_data['root_orient'])
|
||||
|
||||
if self.args.geodesic_loss:
|
||||
orient_recon_loss = self.criterion_geo(root_orient.view(-1, 3, 3), root_orient_gt.view(-1, 3, 3))
|
||||
else:
|
||||
orient_recon_loss = self.criterion_rec(root_orient, root_orient_gt)
|
||||
self.loss_recorder.add_scalar('orient_recon_loss', orient_recon_loss, validation=validation)
|
||||
loss += self.args.lambda_orient * orient_recon_loss
|
||||
|
||||
if self.args.geodesic_loss:
|
||||
rotmat_recon_loss = self.criterion_geo(self.recon_data['rotmat'].view(-1, 3, 3), self.input_data['rotmat'].view(-1, 3, 3))
|
||||
else:
|
||||
rotmat_recon_loss = self.criterion_rec(self.recon_data['rotmat'], self.input_data['rotmat'])
|
||||
self.loss_recorder.add_scalar('rotmat_recon_loss', rotmat_recon_loss, validation=validation)
|
||||
loss += self.args.lambda_rotmat * rotmat_recon_loss
|
||||
|
||||
pos_recon_loss = self.criterion_rec(self.recon_data['pos'], self.input_data['pos'])
|
||||
self.loss_recorder.add_scalar('pos_recon_loss', pos_recon_loss, validation=validation)
|
||||
loss += self.args.lambda_pos * pos_recon_loss
|
||||
|
||||
if self.args.output_trans:
|
||||
v_loss = self.criterion_rec(self.recon_data['root_vel'], self.input_data['root_vel'])
|
||||
up_loss = self.criterion_rec(self.recon_data['root_height'], self.input_data['trans'][..., 'xyz'.index(self.args.data.up)])
|
||||
self.loss_recorder.add_scalar('v_loss', v_loss, validation=validation)
|
||||
self.loss_recorder.add_scalar('up_loss', up_loss, validation=validation)
|
||||
loss += self.args.lambda_v * v_loss + self.args.lambda_up * up_loss
|
||||
|
||||
origin = self.input_data['trans'][:, 0]
|
||||
trans = self.recon_data['trans']
|
||||
trans[..., self.v_axis] = trans[..., self.v_axis] + origin[..., self.v_axis].unsqueeze(1)
|
||||
trans_loss = self.criterion_rec(trans, self.input_data['trans'])
|
||||
self.loss_recorder.add_scalar('trans_loss', trans_loss, validation=validation)
|
||||
loss += self.args.lambda_trans * trans_loss
|
||||
|
||||
contacts_loss = self.criterion_bce(self.recon_data['contacts'], self.input_data['contacts'])
|
||||
self.loss_recorder.add_scalar('contacts_loss', contacts_loss, validation=validation)
|
||||
loss += self.args.lambda_contacts * contacts_loss
|
||||
|
||||
if self.args.lambda_kl != 0:
|
||||
mu_l = self.recon_data['mu_l']
|
||||
logvar_l = self.recon_data['logvar_l']
|
||||
local_kl_loss = torch.mean(-0.5 * torch.sum(1 + logvar_l - mu_l.pow(2) - logvar_l.exp(), dim=-1))
|
||||
self.loss_recorder.add_scalar('local_kl_loss', local_kl_loss, validation=validation)
|
||||
|
||||
mu_g = self.recon_data['mu_g']
|
||||
logvar_g = self.recon_data['logvar_g']
|
||||
global_kl_loss = torch.mean(-0.5 * torch.sum(1 + logvar_g - mu_g.pow(2) - logvar_g.exp(), dim=-1))
|
||||
self.loss_recorder.add_scalar('global_kl_loss', global_kl_loss, validation=validation)
|
||||
if not validation:
|
||||
loss += self.args.lambda_kl * self.kl_scheduler() * (local_kl_loss + global_kl_loss)
|
||||
|
||||
self.loss_recorder.add_scalar('total_loss', loss, validation=validation)
|
||||
|
||||
if not validation:
|
||||
loss.backward()
|
||||
|
||||
def optimize_parameters(self):
|
||||
self.optimizer.zero_grad()
|
||||
self.forward()
|
||||
self.backward(validation=False)
|
||||
self.optimizer.step()
|
||||
|
||||
self.iteration += 1
|
||||
|
||||
def report_errors(self):
|
||||
root_orient = rotation_6d_to_matrix(self.recon_data['root_orient']) # (B, T, 3, 3)
|
||||
root_orient_gt = rotation_6d_to_matrix(self.input_data['root_orient']) # (B, T, 3, 3)
|
||||
orientation_error = self.criterion_geo(root_orient.view(-1, 3, 3), root_orient_gt.view(-1, 3, 3))
|
||||
|
||||
local_rotmat = self.fk.global_to_local(self.recon_data['rotmat'].view(-1, self.args.smpl.joint_num, 3, 3)) # (B x T, J, 3, 3)
|
||||
local_rotmat_gt = self.fk.global_to_local(self.input_data['rotmat'].view(-1, self.args.smpl.joint_num, 3, 3)) # (B x T, J, 3, 3)
|
||||
rotation_error = self.criterion_geo(local_rotmat[:, 1:].reshape(-1, 3, 3), local_rotmat_gt[:, 1:].reshape(-1, 3, 3))
|
||||
|
||||
pos = self.recon_data['pos'] # (B, T, J, 3)
|
||||
pos_gt = self.input_data['pos'] # (B, T, J, 3)
|
||||
position_error = torch.linalg.norm((pos - pos_gt), dim=-1).mean()
|
||||
|
||||
if 'trans' in self.recon_data.keys():
|
||||
origin = self.input_data['trans'][:, 0]
|
||||
trans = self.recon_data['trans']
|
||||
trans[..., self.v_axis] = trans[..., self.v_axis] + origin[..., self.v_axis].unsqueeze(1)
|
||||
trans_gt = self.input_data['trans'] # (B, T, 3)
|
||||
translation_error = torch.linalg.norm((trans - trans_gt), dim=-1).mean()
|
||||
else:
|
||||
translation_error = torch.zeros(1).mean()
|
||||
|
||||
return {
|
||||
'rotation': c2c(rotation_error),
|
||||
'position': c2c(position_error),
|
||||
'orientation': c2c(orientation_error),
|
||||
'translation': c2c(translation_error)
|
||||
}
|
||||
|
||||
def verbose(self):
|
||||
res = {}
|
||||
for loss in self.loss_recorder.losses.values():
|
||||
res[loss.name] = {'train': loss.current()[0], 'val': loss.current()[1]}
|
||||
|
||||
return res
|
||||
|
||||
def validate(self):
|
||||
with torch.no_grad():
|
||||
self.forward()
|
||||
self.backward(validation=True)
|
||||
|
||||
def save(self, optimal=False):
|
||||
if optimal:
|
||||
path = os.path.join(self.args.save_dir, 'results', 'model')
|
||||
else:
|
||||
path = os.path.join(self.model_save_dir, f'{self.epoch_cnt:04d}')
|
||||
|
||||
os.makedirs(path, exist_ok=True)
|
||||
local_encoder = self.local_encoder.module if isinstance(self.local_encoder, nn.DataParallel) else self.local_encoder
|
||||
nemf = self.field.module if isinstance(self.field, nn.DataParallel) else self.field
|
||||
torch.save(local_encoder.state_dict(), os.path.join(path, 'local_encoder.pth'))
|
||||
torch.save(nemf.state_dict(), os.path.join(path, 'nemf.pth'))
|
||||
if self.args.data.root_transform:
|
||||
global_encoder = self.global_encoder.module if isinstance(self.global_encoder, nn.DataParallel) else self.global_encoder
|
||||
torch.save(global_encoder.state_dict(), os.path.join(path, 'global_encoder.pth'))
|
||||
torch.save(self.optimizer.state_dict(), os.path.join(path, 'optimizer.pth'))
|
||||
if self.args.scheduler.name:
|
||||
torch.save(self.schedulers[0].state_dict(), os.path.join(path, 'scheduler.pth'))
|
||||
self.loss_recorder.save(path)
|
||||
|
||||
print(f'Save at {path} succeeded')
|
||||
|
||||
def load(self, epoch=None, optimal=False):
|
||||
if optimal:
|
||||
path = os.path.join(self.args.save_dir, 'results', 'model')
|
||||
else:
|
||||
if epoch is None:
|
||||
all = [int(q) for q in os.listdir(self.model_save_dir)]
|
||||
if len(all) == 0:
|
||||
raise RuntimeError(f'Empty loading path {self.model_save_dir}')
|
||||
epoch = sorted(all)[-1]
|
||||
path = os.path.join(self.model_save_dir, f'{epoch:04d}')
|
||||
|
||||
print(f'Loading from {path}')
|
||||
local_encoder = self.local_encoder.module if isinstance(self.local_encoder, nn.DataParallel) else self.local_encoder
|
||||
nemf = self.field.module if isinstance(self.field, nn.DataParallel) else self.field
|
||||
local_encoder.load_state_dict(torch.load(os.path.join(path, 'local_encoder.pth'), map_location=self.device))
|
||||
nemf.load_state_dict(torch.load(os.path.join(path, 'nemf.pth'), map_location=self.device))
|
||||
if self.args.data.root_transform:
|
||||
global_encoder = self.global_encoder.module if isinstance(self.global_encoder, nn.DataParallel) else self.global_encoder
|
||||
global_encoder.load_state_dict(torch.load(os.path.join(path, 'global_encoder.pth'), map_location=self.device))
|
||||
if self.is_train:
|
||||
self.optimizer.load_state_dict(torch.load(os.path.join(path, 'optimizer.pth')))
|
||||
self.loss_recorder.load(path)
|
||||
self.iteration = len(self.loss_recorder.losses['total_loss'].loss_step)
|
||||
print(f'trained iterations: {self.iteration}')
|
||||
if self.args.scheduler.name:
|
||||
self.schedulers[0].load_state_dict(torch.load(os.path.join(path, 'scheduler.pth')))
|
||||
self.epoch_cnt = epoch if not optimal else 0
|
||||
|
||||
print('Load succeeded')
|
||||
|
||||
def super_sampling(self, step=1):
|
||||
with torch.no_grad():
|
||||
self.forward(step)
|
||||
self.fps = self.args.data.fps / step
|
||||
self.compute_test_result()
|
||||
|
||||
def compute_test_result(self):
|
||||
os.makedirs(self.viz_dir, exist_ok=True)
|
||||
|
||||
b_size = self.input_data['pos'].shape[0]
|
||||
n_joints = self.input_data['pos'].shape[2]
|
||||
|
||||
rotmat = self.recon_data['rotmat'] # (B, T, J, 3, 3)
|
||||
local_rotmat = self.fk.global_to_local(rotmat.view(-1, n_joints, 3, 3)) # (B x T, J, 3, 3)
|
||||
local_rotmat = local_rotmat.view(b_size, -1, n_joints, 3, 3) # (B, T, J, 3, 3)
|
||||
rotmat_gt = self.input_data['rotmat'] # (B, T, J, 3, 3)
|
||||
local_rotmat_gt = self.fk.global_to_local(rotmat_gt.view(-1, n_joints, 3, 3)) # (B x T, J, 3, 3)
|
||||
local_rotmat_gt = local_rotmat_gt.view(b_size, -1, n_joints, 3, 3) # (B, T, J, 3, 3)
|
||||
|
||||
if self.args.data.root_transform:
|
||||
root_orient = rotation_6d_to_matrix(self.recon_data['root_orient']) # (B, T, 3, 3)
|
||||
root_orient_gt = rotation_6d_to_matrix(self.input_data['root_orient']) # (B, T, 3, 3)
|
||||
local_rotmat[:, :, 0] = root_orient
|
||||
local_rotmat_gt[:, :, 0] = root_orient_gt
|
||||
rotations = matrix_to_quaternion(local_rotmat) # (B, T, J, 4)
|
||||
rotations_gt = matrix_to_quaternion(local_rotmat_gt) # (B, T, J, 4)
|
||||
|
||||
if 'trans' in self.recon_data.keys():
|
||||
origin = self.input_data['trans'][:, 0]
|
||||
positions = self.recon_data['trans']
|
||||
positions[..., self.v_axis] = positions[..., self.v_axis] + origin[..., self.v_axis].unsqueeze(1)
|
||||
positions_gt = self.input_data['trans'] # (B, T, 3)
|
||||
else:
|
||||
positions = torch.zeros_like(self.input_data['trans']) # (B, T, 3)
|
||||
positions_gt = torch.zeros_like(self.input_data['trans']) # (B, T, 3)
|
||||
|
||||
for i in range(b_size):
|
||||
if self.args.bvh_viz:
|
||||
rotation = align_joints(rotations[i])
|
||||
position = positions[i].unsqueeze(1) # (T, 1, 3)
|
||||
anim = Animation(Quaternions(c2c(rotation)), c2c(position), None, offsets=self.args.smpl.offsets[self.args.data.gender], parents=self.args.smpl.parents)
|
||||
BVH.save(os.path.join(self.viz_dir, f'test_{self.test_index + i:03d}_{int(self.fps)}fps.bvh'), anim, names=self.args.smpl.joint_names, frametime=1 / self.fps)
|
||||
|
||||
rotation_gt = align_joints(rotations_gt[i])
|
||||
position_gt = positions_gt[i].unsqueeze(1) # (T, 1, 3)
|
||||
anim.rotations = Quaternions(c2c(rotation_gt))
|
||||
anim.positions = c2c(position_gt)
|
||||
BVH.save(os.path.join(self.viz_dir, f'test_{self.test_index + i:03d}_gt.bvh'), anim, names=self.args.smpl.joint_names, frametime=1 / self.args.data.fps)
|
||||
else:
|
||||
poses = c2c(quaternion_to_axis_angle(rotations[i])) # (T, J, 3)
|
||||
poses_gt = c2c(quaternion_to_axis_angle(rotations_gt[i])) # (T, J, 3)
|
||||
|
||||
poses = poses.reshape((poses.shape[0], -1)) # (T, J x 3)
|
||||
poses = np.pad(poses, [(0, 0), (0, 93)], mode='constant')
|
||||
poses_gt = poses_gt.reshape((poses_gt.shape[0], -1)) # (T, J x 3)
|
||||
poses_gt = np.pad(poses_gt, [(0, 0), (0, 93)], mode='constant')
|
||||
|
||||
np.savez(os.path.join(self.viz_dir, f'test_{self.test_index + i:03d}_{int(self.fps)}fps.npz'),
|
||||
poses=poses, trans=c2c(positions[i]), betas=np.zeros(10), gender=self.args.data.gender, mocap_framerate=self.fps)
|
||||
np.savez(os.path.join(self.viz_dir, f'test_{self.test_index + i:03d}_gt.npz'),
|
||||
poses=poses_gt, trans=c2c(positions_gt[i]), betas=np.zeros(10), gender=self.args.data.gender, mocap_framerate=self.args.data.fps)
|
||||
|
||||
self.test_index += b_size
|
||||
@@ -0,0 +1,309 @@
|
||||
import os
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from human_body_prior.tools.omni_tools import copy2cpu as c2c
|
||||
from rotations import matrix_to_axis_angle, rotation_6d_to_matrix
|
||||
from utils import compute_trajectory, export_ply_trajectory, normalize
|
||||
|
||||
from .base_model import BaseModel
|
||||
from .fk import ForwardKinematicsLayer
|
||||
from .residual_blocks import ResidualBlock, SkeletonResidual, residual_ratio
|
||||
from .skeleton import SkeletonConv, SkeletonPool, build_edge_topology, find_neighbor
|
||||
|
||||
|
||||
class Predictor(nn.Module):
|
||||
def __init__(self, args, topology):
|
||||
super(Predictor, self).__init__()
|
||||
self.topologies = [topology]
|
||||
self.channel_base = [args.channel_base]
|
||||
self.channel_list = []
|
||||
self.edge_num = [len(topology)]
|
||||
self.pooling_list = []
|
||||
self.layers = nn.ModuleList()
|
||||
self.args = args
|
||||
# self.convs = []
|
||||
|
||||
kernel_size = args.kernel_size
|
||||
padding = (kernel_size - 1) // 2
|
||||
bias = True
|
||||
|
||||
for _ in range(args.num_layers):
|
||||
self.channel_base.append(self.channel_base[-1] * 2)
|
||||
|
||||
for i in range(args.num_layers):
|
||||
seq = []
|
||||
neighbour_list = find_neighbor(self.topologies[i], args.skeleton_dist)
|
||||
in_channels = self.channel_base[i] * self.edge_num[i]
|
||||
out_channels = self.channel_base[i + 1] * self.edge_num[i]
|
||||
if i == 0:
|
||||
self.channel_list.append(in_channels)
|
||||
self.channel_list.append(out_channels)
|
||||
last_pool = True if i == args.num_layers - 1 else False
|
||||
|
||||
# (T, J, D) => (T, J', D)
|
||||
pool = SkeletonPool(edges=self.topologies[i], pooling_mode=args.skeleton_pool,
|
||||
channels_per_edge=out_channels // len(neighbour_list), last_pool=last_pool)
|
||||
|
||||
if args.use_residual_blocks:
|
||||
# (T, J, D) => (T, J', 2D)
|
||||
seq.append(SkeletonResidual(self.topologies[i], neighbour_list, joint_num=self.edge_num[i], in_channels=in_channels, out_channels=out_channels,
|
||||
kernel_size=kernel_size, stride=1, padding=padding, padding_mode=args.padding_mode, bias=bias,
|
||||
extra_conv=args.extra_conv, pooling_mode=args.skeleton_pool, activation=args.activation, last_pool=last_pool))
|
||||
else:
|
||||
for _ in range(args.extra_conv):
|
||||
# (T, J, D) => (T, J, D)
|
||||
seq.append(SkeletonConv(neighbour_list, in_channels=in_channels, out_channels=in_channels,
|
||||
joint_num=self.edge_num[i], kernel_size=kernel_size, stride=1,
|
||||
padding=padding, padding_mode=args.padding_mode, bias=bias))
|
||||
seq.append(nn.PReLU())
|
||||
# (T, J, D) => (T, J, 2D)
|
||||
seq.append(SkeletonConv(neighbour_list, in_channels=in_channels, out_channels=out_channels,
|
||||
joint_num=self.edge_num[i], kernel_size=kernel_size, stride=1,
|
||||
padding=padding, padding_mode=args.padding_mode, bias=bias, add_offset=False,
|
||||
in_offset_channel=3 * self.channel_base[i] // self.channel_base[0]))
|
||||
# self.convs.append(seq[-1])
|
||||
|
||||
seq.append(pool)
|
||||
seq.append(nn.PReLU())
|
||||
self.layers.append(nn.Sequential(*seq))
|
||||
|
||||
self.topologies.append(pool.new_edges)
|
||||
self.pooling_list.append(pool.pooling_list)
|
||||
self.edge_num.append(len(self.topologies[-1]))
|
||||
|
||||
in_channels = self.channel_base[-1] * len(self.pooling_list[-1])
|
||||
out_channels = args.out_channels # root orient (6) + root vel (3) + root height (1) + contacts (8)
|
||||
|
||||
self.global_motion = nn.Sequential(
|
||||
ResidualBlock(in_channels=in_channels, out_channels=512, kernel_size=kernel_size, stride=1, padding=padding, residual_ratio=residual_ratio(1), activation=args.activation),
|
||||
ResidualBlock(in_channels=512, out_channels=256, kernel_size=kernel_size, stride=1, padding=padding, residual_ratio=residual_ratio(2), activation=args.activation),
|
||||
ResidualBlock(in_channels=256, out_channels=128, kernel_size=kernel_size, stride=1, padding=padding, residual_ratio=residual_ratio(3), activation=args.activation),
|
||||
ResidualBlock(in_channels=128, out_channels=out_channels, kernel_size=kernel_size, stride=1, padding=padding,
|
||||
residual_ratio=residual_ratio(4), activation=args.activation, last_layer=True)
|
||||
)
|
||||
|
||||
def forward(self, input):
|
||||
feature = input
|
||||
for layer in self.layers:
|
||||
feature = layer(feature)
|
||||
|
||||
return self.global_motion(feature)
|
||||
|
||||
|
||||
class GlobalMotionPredictor(BaseModel):
|
||||
def __init__(self, args, ngpu):
|
||||
super(GlobalMotionPredictor, self).__init__(args)
|
||||
|
||||
self.args = args
|
||||
# smpl_fname = os.path.join(args.smpl.smpl_body_model, args.data.gender, 'model.npz')
|
||||
# smpl_data = np.load(smpl_fname, encoding='latin1')
|
||||
smpl_fname = os.path.join(args.smpl.smpl_body_model, "MANO_RIGHT.pkl")
|
||||
smpl_data = np.load(smpl_fname, encoding='latin1', allow_pickle=True)
|
||||
parents = smpl_data['kintree_table'][0].astype(np.int32)
|
||||
edges = build_edge_topology(parents)
|
||||
|
||||
self.fk = ForwardKinematicsLayer(args)
|
||||
|
||||
self.predictor = Predictor(args.global_motion, edges).to(self.device)
|
||||
if args.multi_gpu is True:
|
||||
self.predictor = nn.DataParallel(self.predictor, device_ids=range(ngpu))
|
||||
self.models = [self.predictor]
|
||||
|
||||
ordermap = {
|
||||
'x': 0,
|
||||
'y': 1,
|
||||
'z': 2,
|
||||
}
|
||||
self.up_axis = ordermap[self.args.data.up]
|
||||
self.v_axis = [x for x in ordermap.values() if x != self.up_axis]
|
||||
|
||||
print(self.up_axis)
|
||||
print(self.v_axis)
|
||||
|
||||
self.input_data = {}
|
||||
self.pred_data = {}
|
||||
|
||||
self.mean = torch.load(os.path.join(args.dataset_dir, f'mean-{args.data.gender}-{args.data.clip_length}-{args.data.fps}fps.pt'), map_location=self.device)
|
||||
print(f'mean: {list(self.mean.keys())}')
|
||||
self.std = torch.load(os.path.join(args.dataset_dir, f'std-{args.data.gender}-{args.data.clip_length}-{args.data.fps}fps.pt'), map_location=self.device)
|
||||
print(f'std: {list(self.std.keys())}')
|
||||
|
||||
if self.is_train:
|
||||
self.optimizer = torch.optim.Adam(self.predictor.parameters(), args.learning_rate, weight_decay=args.weight_decay)
|
||||
self.optimizers = [self.optimizer]
|
||||
self.criterion_pred = nn.L1Loss() if args.l1_loss else nn.MSELoss()
|
||||
self.criterion_bce = nn.BCEWithLogitsLoss()
|
||||
else:
|
||||
self.test_index = 0
|
||||
self.smpl_dir = os.path.join(args.save_dir, 'results', 'smpl')
|
||||
|
||||
def set_input(self, input):
|
||||
self.input_data = {k: v.to(self.device) for k, v in input.items() if k in ['pos', 'velocity', 'rot6d', 'angular', 'global_xform',
|
||||
'root_vel', 'trans', 'contacts']}
|
||||
|
||||
def forward(self):
|
||||
self.pred_data.clear()
|
||||
self.input_data['origin'] = self.input_data['trans'][:, 0]
|
||||
self.pred_data = self.predict(data=self.input_data, dt=1.0 / self.args.data.fps)
|
||||
|
||||
def predict(self, data, dt):
|
||||
b_size, t_length = data['pos'].shape[:2]
|
||||
|
||||
if 'pos' in self.args.data.normalize:
|
||||
pos = normalize(data['pos'], mean=self.mean['pos'], std=self.std['pos'])
|
||||
else:
|
||||
pos = data['pos']
|
||||
if 'velocity' in self.args.data.normalize:
|
||||
velocity = normalize(data['velocity'], mean=self.mean['velocity'], std=self.std['velocity'])
|
||||
else:
|
||||
velocity = data['velocity']
|
||||
if 'rot6d' in self.args.data.normalize:
|
||||
rot6d = normalize(data['rot6d'], mean=self.mean['rot6d'], std=self.std['rot6d'])
|
||||
else:
|
||||
rot6d = data['rot6d']
|
||||
if 'angular' in self.args.data.normalize:
|
||||
angular = normalize(data['angular'], mean=self.mean['angular'], std=self.std['angular'])
|
||||
else:
|
||||
angular = data['angular']
|
||||
|
||||
x = torch.cat((pos, velocity, rot6d, angular), dim=-1)
|
||||
x = x.view(b_size, t_length, -1) # (B, T, J x D)
|
||||
x = x.permute(0, 2, 1) # (B, J x D, T)
|
||||
|
||||
global_motion = self.predictor(x)
|
||||
global_motion = global_motion.permute(0, 2, 1) # (B, T, D)
|
||||
|
||||
output = dict()
|
||||
output['root_vel'] = global_motion[:, :, :3]
|
||||
output['root_height'] = global_motion[:, :, 3]
|
||||
output['contacts'] = global_motion[:, :, 4:]
|
||||
if 'origin' in data.keys():
|
||||
origin = data['origin']
|
||||
else:
|
||||
origin = torch.zeros(b_size, 3).to(self.device)
|
||||
output['trans'] = compute_trajectory(output['root_vel'], output['root_height'], origin, dt, up_axis=self.args.data.up)
|
||||
|
||||
return output
|
||||
|
||||
def backward(self, validation=False):
|
||||
loss = 0
|
||||
|
||||
v_loss = self.criterion_pred(self.pred_data['root_vel'], self.input_data['root_vel'])
|
||||
up_loss = self.criterion_pred(self.pred_data['root_height'], self.input_data['trans'][..., self.up_axis])
|
||||
self.loss_recorder.add_scalar('v_loss', v_loss, validation=validation)
|
||||
self.loss_recorder.add_scalar('up_loss', up_loss, validation=validation)
|
||||
loss += self.args.lambda_v * v_loss + self.args.lambda_up * up_loss
|
||||
|
||||
trans_loss = self.criterion_pred(self.pred_data['trans'], self.input_data['trans'])
|
||||
self.loss_recorder.add_scalar('trans_loss', trans_loss, validation=validation)
|
||||
loss += self.args.lambda_trans * trans_loss
|
||||
|
||||
contacts_loss = self.criterion_bce(self.pred_data['contacts'], self.input_data['contacts'])
|
||||
self.loss_recorder.add_scalar('contacts_loss', contacts_loss, validation=validation)
|
||||
loss += self.args.lambda_contacts * contacts_loss
|
||||
|
||||
self.loss_recorder.add_scalar('total_loss', loss, validation=validation)
|
||||
|
||||
if not validation:
|
||||
loss.backward()
|
||||
|
||||
def optimize_parameters(self):
|
||||
self.optimizer.zero_grad()
|
||||
self.forward()
|
||||
self.backward(validation=False)
|
||||
self.optimizer.step()
|
||||
|
||||
def report_errors(self):
|
||||
trans = self.pred_data['trans'] # (B, T, 3)
|
||||
trans_gt = self.input_data['trans'] # (B, T, 3)
|
||||
translation_error = torch.linalg.norm((trans - trans_gt), dim=-1).mean()
|
||||
|
||||
return {
|
||||
'translation': c2c(translation_error)
|
||||
}
|
||||
|
||||
def verbose(self):
|
||||
res = {}
|
||||
for loss in self.loss_recorder.losses.values():
|
||||
res[loss.name] = {'train': loss.current()[0], 'val': loss.current()[1]}
|
||||
|
||||
return res
|
||||
|
||||
def validate(self):
|
||||
with torch.no_grad():
|
||||
self.forward()
|
||||
self.backward(validation=True)
|
||||
|
||||
def save(self, optimal=False):
|
||||
if optimal:
|
||||
path = os.path.join(self.args.save_dir, 'results', 'model')
|
||||
else:
|
||||
path = os.path.join(self.model_save_dir, f'{self.epoch_cnt:04d}')
|
||||
|
||||
os.makedirs(path, exist_ok=True)
|
||||
predictor = self.predictor.module if isinstance(self.predictor, nn.DataParallel) else self.predictor
|
||||
torch.save(predictor.state_dict(), os.path.join(path, 'predictor.pth'))
|
||||
torch.save(self.optimizer.state_dict(), os.path.join(path, 'optimizer.pth'))
|
||||
if self.args.scheduler.name:
|
||||
torch.save(self.schedulers[0].state_dict(), os.path.join(path, 'scheduler.pth'))
|
||||
self.loss_recorder.save(path)
|
||||
|
||||
print(f'Save at {path} succeeded')
|
||||
|
||||
def load(self, epoch=None, optimal=False):
|
||||
if optimal:
|
||||
path = os.path.join(self.args.save_dir, 'results', 'model')
|
||||
else:
|
||||
if epoch is None:
|
||||
all = [int(q) for q in os.listdir(self.model_save_dir)]
|
||||
if len(all) == 0:
|
||||
raise RuntimeError(f'Empty loading path {self.model_save_dir}')
|
||||
epoch = sorted(all)[-1]
|
||||
path = os.path.join(self.model_save_dir, f'{epoch:04d}')
|
||||
|
||||
print(f'Loading from {path}')
|
||||
predictor = self.predictor.module if isinstance(self.predictor, nn.DataParallel) else self.predictor
|
||||
predictor.load_state_dict(torch.load(os.path.join(path, 'predictor.pth'), map_location=self.device))
|
||||
if self.is_train:
|
||||
self.optimizer.load_state_dict(torch.load(os.path.join(path, 'optimizer.pth')))
|
||||
self.loss_recorder.load(path)
|
||||
if self.args.scheduler.name:
|
||||
self.schedulers[0].load_state_dict(torch.load(os.path.join(path, 'scheduler.pth')))
|
||||
self.epoch_cnt = epoch if not optimal else 0
|
||||
|
||||
print('Load succeeded')
|
||||
|
||||
def compute_test_result(self):
|
||||
os.makedirs(self.smpl_dir, exist_ok=True)
|
||||
|
||||
b_size = self.input_data['trans'].shape[0]
|
||||
n_joints = self.input_data['pos'].shape[2]
|
||||
|
||||
trans = self.pred_data['trans']
|
||||
trans_gt = self.input_data['trans']
|
||||
|
||||
rotmat_gt = rotation_6d_to_matrix(self.input_data['global_xform']) # (B, T, J, 3, 3)
|
||||
local_rotmat_gt = self.fk.global_to_local(rotmat_gt.view(-1, n_joints, 3, 3)) # (B x T, J, 3, 3)
|
||||
local_rotmat_gt = local_rotmat_gt.view(b_size, -1, n_joints, 3, 3) # (B, T, J, 3, 3)
|
||||
local_rotmat = local_rotmat_gt.clone()
|
||||
|
||||
for i in range(b_size):
|
||||
export_ply_trajectory(points=trans[i], color=(255, 0, 0), ply_fname=os.path.join(self.smpl_dir, f'test_{self.test_index + i:03d}.ply'))
|
||||
export_ply_trajectory(points=trans_gt[i], color=(0, 255, 0), ply_fname=os.path.join(self.smpl_dir, f'test_{self.test_index + i:03d}_gt.ply'))
|
||||
|
||||
poses = c2c(matrix_to_axis_angle(local_rotmat[i])) # (T, J, 3)
|
||||
poses_gt = c2c(matrix_to_axis_angle(local_rotmat_gt[i])) # (T, J, 3)
|
||||
|
||||
poses = poses.reshape((poses.shape[0], -1)) # (T, J x 3)
|
||||
poses = np.pad(poses, [(0, 0), (0, 93)], mode='constant')
|
||||
poses_gt = poses_gt.reshape((poses_gt.shape[0], -1)) # (T, J x 3)
|
||||
poses_gt = np.pad(poses_gt, [(0, 0), (0, 93)], mode='constant')
|
||||
|
||||
np.savez(os.path.join(self.smpl_dir, f'test_{self.test_index + i:03d}.npz'),
|
||||
poses=poses, trans=c2c(trans[i]), betas=np.zeros(10), gender=self.args.data.gender, mocap_framerate=self.args.data.fps)
|
||||
np.savez(os.path.join(self.smpl_dir, f'test_{self.test_index + i:03d}_gt.npz'),
|
||||
poses=poses_gt, trans=c2c(trans_gt[i]), betas=np.zeros(10), gender=self.args.data.gender, mocap_framerate=self.args.data.fps)
|
||||
|
||||
self.test_index += b_size
|
||||
@@ -0,0 +1,82 @@
|
||||
import os
|
||||
import pickle
|
||||
|
||||
import torch
|
||||
from tensorboardX import SummaryWriter
|
||||
|
||||
|
||||
class SingleLoss:
|
||||
def __init__(self, name: str, writer: SummaryWriter):
|
||||
self.name = name
|
||||
self.loss_step = []
|
||||
self.loss_epoch = []
|
||||
self.loss_epoch_tmp = []
|
||||
self.loss_epoch_val = []
|
||||
self.writer = writer
|
||||
|
||||
def add_scalar(self, val, step=None, validation=False):
|
||||
if validation:
|
||||
self.loss_epoch_val.append(val)
|
||||
else:
|
||||
if step is None:
|
||||
step = len(self.loss_step)
|
||||
self.loss_step.append(val)
|
||||
self.loss_epoch_tmp.append(val)
|
||||
self.writer.add_scalar('Loss/step_' + self.name, val, step)
|
||||
|
||||
def current(self):
|
||||
return self.loss_epoch[-1]
|
||||
|
||||
def epoch(self, step=None):
|
||||
if step is None:
|
||||
step = len(self.loss_epoch)
|
||||
if self.loss_epoch_tmp:
|
||||
loss_avg = sum(self.loss_epoch_tmp) / len(self.loss_epoch_tmp)
|
||||
else:
|
||||
loss_avg = 0
|
||||
self.loss_epoch_tmp = []
|
||||
if self.loss_epoch_val:
|
||||
loss_avg_val = sum(self.loss_epoch_val) / len(self.loss_epoch_val)
|
||||
else:
|
||||
loss_avg_val = 0
|
||||
self.loss_epoch_val = []
|
||||
self.loss_epoch.append([loss_avg, loss_avg_val])
|
||||
self.writer.add_scalars('Loss/epoch_' + self.name, {'train': loss_avg, 'val': loss_avg_val}, step)
|
||||
|
||||
def save(self):
|
||||
return {
|
||||
'loss_step': self.loss_step,
|
||||
'loss_epoch': self.loss_epoch
|
||||
}
|
||||
|
||||
|
||||
class LossRecorder:
|
||||
def __init__(self, writer: SummaryWriter):
|
||||
self.losses = {}
|
||||
self.writer = writer
|
||||
|
||||
def add_scalar(self, name, val, step=None, validation=False):
|
||||
if isinstance(val, torch.Tensor):
|
||||
val = val.item()
|
||||
if name not in self.losses:
|
||||
self.losses[name] = SingleLoss(name, self.writer)
|
||||
self.losses[name].add_scalar(val, step, validation)
|
||||
|
||||
def epoch(self, step=None):
|
||||
for loss in self.losses.values():
|
||||
loss.epoch(step)
|
||||
|
||||
def save(self, path):
|
||||
data = {}
|
||||
for key, value in self.losses.items():
|
||||
data[key] = value.save()
|
||||
with open(os.path.join(path, 'loss.pkl'), 'wb') as f:
|
||||
pickle.dump(data, f)
|
||||
|
||||
def load(self, path):
|
||||
with open(os.path.join(path, 'loss.pkl'), 'rb') as f:
|
||||
data = pickle.load(f)
|
||||
for key, value in data.items():
|
||||
self.losses[key] = SingleLoss(key, self.writer)
|
||||
self.losses[key].loss_step = value['loss_step']
|
||||
self.losses[key].loss_epoch = value['loss_epoch']
|
||||
@@ -0,0 +1,88 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
class GeodesicLoss(nn.Module):
|
||||
def __init__(self):
|
||||
super(GeodesicLoss, self).__init__()
|
||||
|
||||
def compute_geodesic_distance(self, m1, m2):
|
||||
""" Compute the geodesic distance between two rotation matrices.
|
||||
|
||||
Args:
|
||||
m1, m2: Two rotation matrices with the shape (batch x 3 x 3).
|
||||
|
||||
Returns:
|
||||
The minimal angular difference between two rotation matrices in radian form [0, pi].
|
||||
"""
|
||||
batch = m1.shape[0]
|
||||
m = torch.bmm(m1, m2.transpose(1, 2)) # batch*3*3
|
||||
|
||||
cos = (m[:, 0, 0] + m[:, 1, 1] + m[:, 2, 2] - 1) / 2
|
||||
# cos = torch.min(cos, torch.autograd.Variable(torch.ones(batch).cuda()))
|
||||
# cos = torch.max(cos, torch.autograd.Variable(torch.ones(batch).cuda()) * -1)
|
||||
eps = 1e-7
|
||||
cos = torch.clamp(cos, -1 + eps, 1 - eps)
|
||||
|
||||
theta = torch.acos(cos)
|
||||
|
||||
return theta
|
||||
|
||||
def rot_smooth_loss(self, rotmat):
|
||||
|
||||
diff = self.compute_geodesic_distance(rotmat[:, 1:], rotmat[:, :-1])
|
||||
vel = diff * 30
|
||||
loss = vel ** 2
|
||||
loss = 0.5 * torch.mean(loss)
|
||||
return loss
|
||||
|
||||
def __call__(self, m1, m2, reduction='mean'):
|
||||
loss = self.compute_geodesic_distance(m1, m2)
|
||||
|
||||
if reduction == 'mean':
|
||||
return loss.mean()
|
||||
elif reduction == 'none':
|
||||
return loss
|
||||
else:
|
||||
raise RuntimeError(f'unsupported reduction: {reduction}')
|
||||
|
||||
|
||||
def pos_smooth_loss(pos):
|
||||
# minimize delta steps
|
||||
diff = (pos[:, 1:] - pos[:, :-1])
|
||||
vel = diff * 30
|
||||
loss = vel ** 2
|
||||
loss = 0.5 * torch.mean(loss)
|
||||
return loss
|
||||
|
||||
def compute_geodesic_distance(m1, m2):
|
||||
""" Compute the geodesic distance between two rotation matrices.
|
||||
|
||||
Args:
|
||||
m1, m2: Two rotation matrices with the shape (batch x 3 x 3).
|
||||
|
||||
Returns:
|
||||
The minimal angular difference between two rotation matrices in radian form [0, pi].
|
||||
"""
|
||||
batch = m1.shape[:-2]
|
||||
m = torch.matmul(m1, m2.transpose(-2, -1)) # batch*3*3
|
||||
|
||||
cos = (m[..., 0, 0] + m[..., 1, 1] + m[..., 2, 2] - 1) / 2
|
||||
cos = torch.min(cos, torch.ones(*batch).cuda())
|
||||
cos = torch.max(cos, torch.ones(*batch).cuda() * -1)
|
||||
|
||||
eps = 1e-7
|
||||
cos = torch.clamp(cos, -1 + eps, 1 - eps)
|
||||
|
||||
theta = torch.acos(cos)
|
||||
return theta
|
||||
|
||||
|
||||
def rot_smooth_loss(rotmat):
|
||||
|
||||
diff = compute_geodesic_distance(rotmat[:, 1:], rotmat[:, :-1])
|
||||
vel = diff * 30
|
||||
loss = vel ** 2
|
||||
loss = 0.5 * torch.mean(loss)
|
||||
return loss
|
||||
|
||||
@@ -0,0 +1,161 @@
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
|
||||
class PositionalEncoding(nn.Module):
|
||||
def __init__(self, n_functions):
|
||||
super(PositionalEncoding, self).__init__()
|
||||
|
||||
self.register_buffer('frequencies', 2.0 ** torch.arange(n_functions))
|
||||
|
||||
def forward(self, x):
|
||||
"""
|
||||
Args:
|
||||
x: tensor of shape [..., dim]
|
||||
|
||||
Returns:
|
||||
embedding: a temporal embedding of `x` of shape [..., n_functions * dim * 2]
|
||||
"""
|
||||
freq = (x[..., None] * self.frequencies).view(*x.shape[:-1], -1)
|
||||
|
||||
embedding = torch.zeros(*freq.shape[:-1], freq.shape[-1] * 2).to(device)
|
||||
embedding[..., 0::2] = freq.sin()
|
||||
embedding[..., 1::2] = freq.cos()
|
||||
|
||||
return embedding
|
||||
|
||||
|
||||
class Sine(nn.Module):
|
||||
def __init__(self, w0=1.):
|
||||
super(Sine, self).__init__()
|
||||
|
||||
self.w0 = w0
|
||||
|
||||
def forward(self, x):
|
||||
return torch.sin(self.w0 * x)
|
||||
|
||||
|
||||
class SirenBlock(nn.Module):
|
||||
def __init__(self, in_features, out_features, w0=30, c=6, is_first=False, use_bias=True, activation=None):
|
||||
super(SirenBlock, self).__init__()
|
||||
|
||||
self.in_features = in_features
|
||||
self.is_first = is_first
|
||||
|
||||
weight = torch.zeros(out_features, in_features)
|
||||
bias = torch.zeros(out_features) if use_bias else None
|
||||
self.init(weight, bias, c=c, w0=w0)
|
||||
|
||||
self.weight = nn.Parameter(weight)
|
||||
self.bias = nn.Parameter(bias) if use_bias else None
|
||||
self.activation = Sine(w0) if activation is None else activation
|
||||
|
||||
def init(self, weight, bias, c, w0):
|
||||
n = self.in_features
|
||||
|
||||
w_std = (1 / n) if self.is_first else (np.sqrt(c / n) / w0)
|
||||
weight.uniform_(-w_std, w_std)
|
||||
|
||||
if bias is not None:
|
||||
bias.uniform_(-w_std, w_std)
|
||||
|
||||
def forward(self, x):
|
||||
out = F.linear(x, self.weight, self.bias)
|
||||
|
||||
return self.activation(out)
|
||||
|
||||
|
||||
class FCBlock(nn.Module):
|
||||
def __init__(self, in_features, out_features, norm_layer=False, activation=None):
|
||||
super(FCBlock, self).__init__()
|
||||
|
||||
self.fc = nn.Linear(in_features, out_features)
|
||||
self.residual = (in_features == out_features) # when the input and output have the same dimensions, build a residual block
|
||||
self.norm_layer = nn.LayerNorm(out_features) if norm_layer else None
|
||||
self.activation = nn.ReLU(inplace=True) if activation is None else activation
|
||||
|
||||
def forward(self, x):
|
||||
"""
|
||||
Args:
|
||||
x: (B, T, D), features are in the last dimension.
|
||||
"""
|
||||
out = self.fc(x)
|
||||
|
||||
if self.norm_layer is not None:
|
||||
out = self.norm_layer(out)
|
||||
|
||||
if self.residual:
|
||||
return self.activation(out) + x
|
||||
|
||||
return self.activation(out)
|
||||
|
||||
|
||||
class NeuralMotionField(nn.Module):
|
||||
def __init__(self, args):
|
||||
super(NeuralMotionField, self).__init__()
|
||||
|
||||
self.args = args
|
||||
|
||||
if args.bandwidth != 0:
|
||||
self.positional_encoding = PositionalEncoding(args.bandwidth)
|
||||
embedding_dim = args.bandwidth * 2
|
||||
else:
|
||||
embedding_dim = 1
|
||||
|
||||
hidden_neuron = args.hidden_neuron
|
||||
in_features = 1 + args.local_z if args.siren else embedding_dim + args.local_z
|
||||
local_in_features = hidden_neuron + in_features if args.skip_connection else hidden_neuron
|
||||
global_in_features = local_in_features + args.global_z if args.skip_connection else hidden_neuron
|
||||
|
||||
layers = [
|
||||
SirenBlock(in_features, hidden_neuron, is_first=True) if args.siren else FCBlock(in_features, hidden_neuron, norm_layer=args.norm_layer),
|
||||
SirenBlock(local_in_features, hidden_neuron) if args.siren else FCBlock(local_in_features, hidden_neuron, norm_layer=args.norm_layer),
|
||||
SirenBlock(local_in_features, hidden_neuron) if args.siren else FCBlock(local_in_features, hidden_neuron, norm_layer=args.norm_layer),
|
||||
SirenBlock(local_in_features, hidden_neuron) if args.siren else FCBlock(local_in_features, hidden_neuron, norm_layer=args.norm_layer),
|
||||
SirenBlock(local_in_features, hidden_neuron) if args.siren else FCBlock(local_in_features, hidden_neuron, norm_layer=args.norm_layer),
|
||||
SirenBlock(local_in_features, hidden_neuron) if args.siren else FCBlock(local_in_features, hidden_neuron, norm_layer=args.norm_layer),
|
||||
SirenBlock(local_in_features, hidden_neuron) if args.siren else FCBlock(local_in_features, hidden_neuron, norm_layer=args.norm_layer),
|
||||
SirenBlock(local_in_features, hidden_neuron) if args.siren else FCBlock(local_in_features, hidden_neuron, norm_layer=args.norm_layer),
|
||||
SirenBlock(global_in_features, hidden_neuron) if args.siren else FCBlock(global_in_features, hidden_neuron, norm_layer=args.norm_layer),
|
||||
SirenBlock(global_in_features, hidden_neuron) if args.siren else FCBlock(global_in_features, hidden_neuron, norm_layer=args.norm_layer),
|
||||
SirenBlock(global_in_features, hidden_neuron) if args.siren else FCBlock(global_in_features, hidden_neuron, norm_layer=args.norm_layer)
|
||||
]
|
||||
self.mlp = nn.ModuleList(layers)
|
||||
self.skip_layers = [] if not args.skip_connection else list(range(1, len(self.mlp)))
|
||||
self.local_layers = list(range(8))
|
||||
self.local_linear = nn.Sequential(nn.Linear(hidden_neuron, args.local_output))
|
||||
self.global_layers = list(range(8, len(self.mlp)))
|
||||
self.global_linear = nn.Sequential(nn.Linear(hidden_neuron, args.global_output))
|
||||
|
||||
def forward(self, t, z_l=None, z_g=None):
|
||||
if self.args.bandwidth != 0 and self.args.siren is False:
|
||||
t = self.positional_encoding(t)
|
||||
|
||||
if z_l is not None:
|
||||
z_l = z_l.unsqueeze(1).expand(-1, t.shape[1], -1)
|
||||
x = torch.cat([t, z_l], dim=-1)
|
||||
else:
|
||||
x = t
|
||||
|
||||
skip_in = x
|
||||
for i in self.local_layers:
|
||||
if i in self.skip_layers:
|
||||
x = torch.cat((x, skip_in), dim=-1)
|
||||
x = self.mlp[i](x)
|
||||
local_output = self.local_linear(x)
|
||||
|
||||
if z_g is not None:
|
||||
z_g = z_g.unsqueeze(1).expand(-1, t.shape[1], -1)
|
||||
for i in self.global_layers:
|
||||
if i in self.skip_layers:
|
||||
x = torch.cat((x, skip_in, z_g), dim=-1)
|
||||
x = self.mlp[i](x)
|
||||
global_output = self.global_linear(x)
|
||||
else:
|
||||
global_output = None
|
||||
|
||||
return local_output, global_output
|
||||
+109
@@ -0,0 +1,109 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from .residual_blocks import ResidualBlock, SkeletonResidual, residual_ratio
|
||||
from .skeleton import SkeletonConv, SkeletonPool, find_neighbor
|
||||
|
||||
|
||||
class LocalEncoder(nn.Module):
|
||||
def __init__(self, args, topology):
|
||||
super(LocalEncoder, self).__init__()
|
||||
self.topologies = [topology]
|
||||
self.channel_base = [args.channel_base]
|
||||
|
||||
self.channel_list = []
|
||||
self.edge_num = [len(topology)]
|
||||
self.pooling_list = []
|
||||
self.layers = nn.ModuleList()
|
||||
self.args = args
|
||||
# self.convs = []
|
||||
|
||||
kernel_size = args.kernel_size
|
||||
kernel_even = False if kernel_size % 2 else True
|
||||
padding = (kernel_size - 1) // 2
|
||||
bias = True
|
||||
|
||||
for _ in range(args.num_layers):
|
||||
self.channel_base.append(self.channel_base[-1] * 2)
|
||||
|
||||
for i in range(args.num_layers):
|
||||
seq = []
|
||||
neighbour_list = find_neighbor(self.topologies[i], args.skeleton_dist)
|
||||
in_channels = self.channel_base[i] * self.edge_num[i]
|
||||
out_channels = self.channel_base[i + 1] * self.edge_num[i]
|
||||
if i == 0:
|
||||
self.channel_list.append(in_channels)
|
||||
self.channel_list.append(out_channels)
|
||||
last_pool = True if i == args.num_layers - 1 else False
|
||||
|
||||
# (T, J, D) => (T, J', D)
|
||||
pool = SkeletonPool(edges=self.topologies[i], pooling_mode=args.skeleton_pool,
|
||||
channels_per_edge=out_channels // len(neighbour_list), last_pool=last_pool)
|
||||
|
||||
if args.use_residual_blocks:
|
||||
# (T, J, D) => (T/2, J', 2D)
|
||||
seq.append(SkeletonResidual(self.topologies[i], neighbour_list, joint_num=self.edge_num[i], in_channels=in_channels, out_channels=out_channels,
|
||||
kernel_size=kernel_size, stride=2, padding=padding, padding_mode=args.padding_mode, bias=bias,
|
||||
extra_conv=args.extra_conv, pooling_mode=args.skeleton_pool, activation=args.activation, last_pool=last_pool))
|
||||
else:
|
||||
for _ in range(args.extra_conv):
|
||||
# (T, J, D) => (T, J, D)
|
||||
seq.append(SkeletonConv(neighbour_list, in_channels=in_channels, out_channels=in_channels,
|
||||
joint_num=self.edge_num[i], kernel_size=kernel_size - 1 if kernel_even else kernel_size,
|
||||
stride=1,
|
||||
padding=padding, padding_mode=args.padding_mode, bias=bias))
|
||||
seq.append(nn.PReLU() if args.activation == 'relu' else nn.Tanh())
|
||||
# (T, J, D) => (T/2, J, 2D)
|
||||
seq.append(SkeletonConv(neighbour_list, in_channels=in_channels, out_channels=out_channels,
|
||||
joint_num=self.edge_num[i], kernel_size=kernel_size, stride=2,
|
||||
padding=padding, padding_mode=args.padding_mode, bias=bias, add_offset=False,
|
||||
in_offset_channel=3 * self.channel_base[i] // self.channel_base[0]))
|
||||
# self.convs.append(seq[-1])
|
||||
|
||||
seq.append(pool)
|
||||
seq.append(nn.PReLU() if args.activation == 'relu' else nn.Tanh())
|
||||
self.layers.append(nn.Sequential(*seq))
|
||||
|
||||
self.topologies.append(pool.new_edges)
|
||||
self.pooling_list.append(pool.pooling_list)
|
||||
self.edge_num.append(len(self.topologies[-1]))
|
||||
|
||||
in_features = self.channel_base[-1] * len(self.pooling_list[-1])
|
||||
in_features *= args.temporal_scale
|
||||
self.mu = nn.Linear(in_features, args.z_dim)
|
||||
self.logvar = nn.Linear(in_features, args.z_dim)
|
||||
|
||||
def forward(self, input):
|
||||
output = input
|
||||
for layer in self.layers:
|
||||
output = layer(output)
|
||||
output = output.view(output.shape[0], -1)
|
||||
|
||||
return self.mu(output), self.logvar(output)
|
||||
|
||||
|
||||
class GlobalEncoder(nn.Module):
|
||||
def __init__(self, args):
|
||||
super(GlobalEncoder, self).__init__()
|
||||
|
||||
in_channels = args.in_channels # root orientation + root velocity
|
||||
out_channels = 512
|
||||
kernel_size = args.kernel_size
|
||||
padding = (kernel_size - 1) // 2
|
||||
|
||||
self.encoder = nn.Sequential(
|
||||
ResidualBlock(in_channels=in_channels, out_channels=128, kernel_size=kernel_size, stride=2, padding=padding, residual_ratio=residual_ratio(1), activation=args.activation),
|
||||
ResidualBlock(in_channels=128, out_channels=256, kernel_size=kernel_size, stride=2, padding=padding, residual_ratio=residual_ratio(2), activation=args.activation),
|
||||
ResidualBlock(in_channels=256, out_channels=512, kernel_size=kernel_size, stride=2, padding=padding, residual_ratio=residual_ratio(3), activation=args.activation),
|
||||
ResidualBlock(in_channels=512, out_channels=out_channels, kernel_size=kernel_size, stride=2, padding=padding, residual_ratio=residual_ratio(4), activation=args.activation)
|
||||
)
|
||||
|
||||
in_features = out_channels * args.temporal_scale
|
||||
self.mu = nn.Linear(in_features, args.z_dim)
|
||||
self.logvar = nn.Linear(in_features, args.z_dim)
|
||||
|
||||
def forward(self, input):
|
||||
feature = self.encoder(input)
|
||||
feature = feature.view(feature.shape[0], -1)
|
||||
|
||||
return self.mu(feature), self.logvar(feature)
|
||||
@@ -0,0 +1,197 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from .skeleton import SkeletonConv, SkeletonPool, SkeletonUnpool
|
||||
|
||||
|
||||
def residual_ratio(k):
|
||||
return 1 / (k + 1)
|
||||
|
||||
|
||||
class Affine(nn.Module):
|
||||
def __init__(self, num_parameters, scale=True, bias=True, scale_init=1.0):
|
||||
super(Affine, self).__init__()
|
||||
if scale:
|
||||
self.scale = nn.Parameter(torch.ones(num_parameters) * scale_init)
|
||||
else:
|
||||
self.register_parameter('scale', None)
|
||||
|
||||
if bias:
|
||||
self.bias = nn.Parameter(torch.zeros(num_parameters))
|
||||
else:
|
||||
self.register_parameter('bias', None)
|
||||
|
||||
def forward(self, input):
|
||||
output = input
|
||||
if self.scale is not None:
|
||||
scale = self.scale.unsqueeze(0)
|
||||
while scale.dim() < input.dim():
|
||||
scale = scale.unsqueeze(2)
|
||||
output = output.mul(scale)
|
||||
|
||||
if self.bias is not None:
|
||||
bias = self.bias.unsqueeze(0)
|
||||
while bias.dim() < input.dim():
|
||||
bias = bias.unsqueeze(2)
|
||||
output += bias
|
||||
|
||||
return output
|
||||
|
||||
|
||||
class BatchStatistics(nn.Module):
|
||||
def __init__(self, affine=-1):
|
||||
super(BatchStatistics, self).__init__()
|
||||
self.affine = nn.Sequential() if affine == -1 else Affine(affine)
|
||||
self.loss = 0
|
||||
|
||||
def clear_loss(self):
|
||||
self.loss = 0
|
||||
|
||||
def compute_loss(self, input):
|
||||
input_flat = input.view(input.size(1), input.numel() // input.size(1))
|
||||
mu = input_flat.mean(1)
|
||||
logvar = (input_flat.pow(2).mean(1) - mu.pow(2)).sqrt().log()
|
||||
|
||||
self.loss = mu.pow(2).mean() + logvar.pow(2).mean()
|
||||
|
||||
def forward(self, input):
|
||||
self.compute_loss(input)
|
||||
return self.affine(input)
|
||||
|
||||
|
||||
class ResidualBlock(nn.Module):
|
||||
def __init__(self, in_channels, out_channels, kernel_size, stride, padding, residual_ratio, activation, batch_statistics=False, last_layer=False):
|
||||
super(ResidualBlock, self).__init__()
|
||||
|
||||
self.residual_ratio = residual_ratio
|
||||
self.shortcut_ratio = 1 - residual_ratio
|
||||
|
||||
residual = []
|
||||
residual.append(nn.Conv1d(in_channels, out_channels, kernel_size, stride, padding))
|
||||
if batch_statistics:
|
||||
residual.append(BatchStatistics(out_channels))
|
||||
if not last_layer:
|
||||
residual.append(nn.PReLU() if activation == 'relu' else nn.Tanh())
|
||||
self.residual = nn.Sequential(*residual)
|
||||
|
||||
self.shortcut = nn.Sequential(
|
||||
nn.AvgPool1d(kernel_size=2) if stride == 2 else nn.Sequential(),
|
||||
nn.Conv1d(in_channels, out_channels, kernel_size=1, stride=1, padding=0),
|
||||
BatchStatistics(out_channels) if (in_channels != out_channels and batch_statistics is True) else nn.Sequential()
|
||||
)
|
||||
|
||||
def forward(self, input):
|
||||
return self.residual(input).mul(self.residual_ratio) + self.shortcut(input).mul(self.shortcut_ratio)
|
||||
|
||||
|
||||
class ResidualBlockTranspose(nn.Module):
|
||||
def __init__(self, in_channels, out_channels, kernel_size, stride, padding, residual_ratio, activation):
|
||||
super(ResidualBlockTranspose, self).__init__()
|
||||
|
||||
self.residual_ratio = residual_ratio
|
||||
self.shortcut_ratio = 1 - residual_ratio
|
||||
|
||||
self.residual = nn.Sequential(
|
||||
nn.ConvTranspose1d(in_channels, out_channels, kernel_size, stride, padding),
|
||||
nn.PReLU() if activation == 'relu' else nn.Tanh()
|
||||
)
|
||||
|
||||
self.shortcut = nn.Sequential(
|
||||
nn.Upsample(scale_factor=2, mode='linear', align_corners=False) if stride == 2 else nn.Sequential(),
|
||||
nn.Conv1d(in_channels, out_channels, kernel_size=1, stride=1, padding=0)
|
||||
)
|
||||
|
||||
def forward(self, input):
|
||||
return self.residual(input).mul(self.residual_ratio) + self.shortcut(input).mul(self.shortcut_ratio)
|
||||
|
||||
|
||||
class SkeletonResidual(nn.Module):
|
||||
def __init__(self, topology, neighbour_list, joint_num, in_channels, out_channels, kernel_size, stride, padding, padding_mode, bias, extra_conv, pooling_mode, activation, last_pool):
|
||||
super(SkeletonResidual, self).__init__()
|
||||
|
||||
kernel_even = False if kernel_size % 2 else True
|
||||
|
||||
seq = []
|
||||
for _ in range(extra_conv):
|
||||
# (T, J, D) => (T, J, D)
|
||||
seq.append(SkeletonConv(neighbour_list, in_channels=in_channels, out_channels=in_channels,
|
||||
joint_num=joint_num, kernel_size=kernel_size - 1 if kernel_even else kernel_size,
|
||||
stride=1,
|
||||
padding=padding, padding_mode=padding_mode, bias=bias))
|
||||
seq.append(nn.PReLU() if activation == 'relu' else nn.Tanh())
|
||||
# (T, J, D) => (T/2, J, 2D)
|
||||
seq.append(SkeletonConv(neighbour_list, in_channels=in_channels, out_channels=out_channels,
|
||||
joint_num=joint_num, kernel_size=kernel_size, stride=stride,
|
||||
padding=padding, padding_mode=padding_mode, bias=bias, add_offset=False))
|
||||
seq.append(nn.GroupNorm(4, out_channels)) # FIXME: REMEMBER TO CHANGE BACK !!!
|
||||
self.residual = nn.Sequential(*seq)
|
||||
|
||||
# (T, J, D) => (T/2, J, 2D)
|
||||
self.shortcut = SkeletonConv(neighbour_list, in_channels=in_channels, out_channels=out_channels,
|
||||
joint_num=joint_num, kernel_size=1, stride=stride, padding=0,
|
||||
bias=True, add_offset=False)
|
||||
|
||||
seq = []
|
||||
# (T/2, J, 2D) => (T/2, J', 2D)
|
||||
pool = SkeletonPool(edges=topology, pooling_mode=pooling_mode,
|
||||
channels_per_edge=out_channels // len(neighbour_list), last_pool=last_pool)
|
||||
if len(pool.pooling_list) != pool.edge_num:
|
||||
seq.append(pool)
|
||||
seq.append(nn.PReLU() if activation == 'relu' else nn.Tanh())
|
||||
self.common = nn.Sequential(*seq)
|
||||
|
||||
def forward(self, input):
|
||||
output = self.residual(input) + self.shortcut(input)
|
||||
|
||||
return self.common(output)
|
||||
|
||||
|
||||
class SkeletonResidualTranspose(nn.Module):
|
||||
def __init__(self, neighbour_list, joint_num, in_channels, out_channels, kernel_size, padding, padding_mode, bias, extra_conv, pooling_list, upsampling, activation, last_layer):
|
||||
super(SkeletonResidualTranspose, self).__init__()
|
||||
|
||||
kernel_even = False if kernel_size % 2 else True
|
||||
|
||||
seq = []
|
||||
# (T, J, D) => (2T, J, D)
|
||||
if upsampling is not None:
|
||||
seq.append(nn.Upsample(scale_factor=2, mode=upsampling, align_corners=False))
|
||||
# (2T, J, D) => (2T, J', D)
|
||||
unpool = SkeletonUnpool(pooling_list, in_channels // len(neighbour_list))
|
||||
if unpool.input_edge_num != unpool.output_edge_num:
|
||||
seq.append(unpool)
|
||||
self.common = nn.Sequential(*seq)
|
||||
|
||||
seq = []
|
||||
for _ in range(extra_conv):
|
||||
# (2T, J', D) => (2T, J', D)
|
||||
seq.append(SkeletonConv(neighbour_list, in_channels=in_channels, out_channels=in_channels,
|
||||
joint_num=joint_num, kernel_size=kernel_size - 1 if kernel_even else kernel_size,
|
||||
stride=1,
|
||||
padding=padding, padding_mode=padding_mode, bias=bias))
|
||||
seq.append(nn.PReLU() if activation == 'relu' else nn.Tanh())
|
||||
# (2T, J', D) => (2T, J', D/2)
|
||||
seq.append(SkeletonConv(neighbour_list, in_channels=in_channels, out_channels=out_channels,
|
||||
joint_num=joint_num, kernel_size=kernel_size - 1 if kernel_even else kernel_size,
|
||||
stride=1,
|
||||
padding=padding, padding_mode=padding_mode, bias=bias, add_offset=False))
|
||||
self.residual = nn.Sequential(*seq)
|
||||
|
||||
# (2T, J', D) => (2T, J', D/2)
|
||||
self.shortcut = SkeletonConv(neighbour_list, in_channels=in_channels, out_channels=out_channels,
|
||||
joint_num=joint_num, kernel_size=1, stride=1, padding=0,
|
||||
bias=True, add_offset=False)
|
||||
|
||||
if activation == 'relu':
|
||||
self.activation = nn.PReLU() if not last_layer else None
|
||||
else:
|
||||
self.activation = nn.Tanh() if not last_layer else None
|
||||
|
||||
def forward(self, input):
|
||||
output = self.common(input)
|
||||
output = self.residual(output) + self.shortcut(output)
|
||||
|
||||
if self.activation is not None:
|
||||
return self.activation(output)
|
||||
else:
|
||||
return output
|
||||
+443
@@ -0,0 +1,443 @@
|
||||
import math
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
class SkeletonConv(nn.Module):
|
||||
def __init__(self, neighbour_list, in_channels, out_channels, kernel_size, joint_num, stride=1, padding=0,
|
||||
bias=True, padding_mode='zeros', add_offset=False, in_offset_channel=0):
|
||||
self.in_channels_per_joint = in_channels // joint_num
|
||||
self.out_channels_per_joint = out_channels // joint_num
|
||||
if in_channels % joint_num != 0 or out_channels % joint_num != 0:
|
||||
raise Exception('BAD')
|
||||
super(SkeletonConv, self).__init__()
|
||||
|
||||
if padding_mode == 'zeros':
|
||||
padding_mode = 'constant'
|
||||
if padding_mode == 'reflection':
|
||||
padding_mode = 'reflect'
|
||||
|
||||
self.expanded_neighbour_list = []
|
||||
self.expanded_neighbour_list_offset = []
|
||||
self.neighbour_list = neighbour_list
|
||||
self.add_offset = add_offset
|
||||
self.joint_num = joint_num
|
||||
|
||||
self.stride = stride
|
||||
self.dilation = 1
|
||||
self.groups = 1
|
||||
self.padding = padding
|
||||
self.padding_mode = padding_mode
|
||||
self._padding_repeated_twice = (padding, padding)
|
||||
|
||||
for neighbour in neighbour_list:
|
||||
expanded = []
|
||||
for k in neighbour:
|
||||
for i in range(self.in_channels_per_joint):
|
||||
expanded.append(k * self.in_channels_per_joint + i)
|
||||
self.expanded_neighbour_list.append(expanded)
|
||||
|
||||
if self.add_offset:
|
||||
self.offset_enc = SkeletonLinear(neighbour_list, in_offset_channel * len(neighbour_list), out_channels)
|
||||
|
||||
for neighbour in neighbour_list:
|
||||
expanded = []
|
||||
for k in neighbour:
|
||||
for i in range(add_offset):
|
||||
expanded.append(k * in_offset_channel + i)
|
||||
self.expanded_neighbour_list_offset.append(expanded)
|
||||
|
||||
self.weight = torch.zeros(out_channels, in_channels, kernel_size)
|
||||
if bias:
|
||||
self.bias = torch.zeros(out_channels)
|
||||
else:
|
||||
self.register_parameter('bias', None)
|
||||
|
||||
self.mask = torch.zeros_like(self.weight)
|
||||
for i, neighbour in enumerate(self.expanded_neighbour_list):
|
||||
self.mask[self.out_channels_per_joint * i: self.out_channels_per_joint * (i + 1), neighbour, ...] = 1
|
||||
self.mask = nn.Parameter(self.mask, requires_grad=False)
|
||||
|
||||
self.description = 'SkeletonConv(in_channels_per_armature={}, out_channels_per_armature={}, kernel_size={}, ' \
|
||||
'joint_num={}, stride={}, padding={}, bias={})'.format(
|
||||
in_channels // joint_num, out_channels // joint_num, kernel_size, joint_num, stride, padding, bias
|
||||
)
|
||||
|
||||
self.reset_parameters()
|
||||
|
||||
def reset_parameters(self):
|
||||
for i, neighbour in enumerate(self.expanded_neighbour_list):
|
||||
""" Use temporary variable to avoid assign to copy of slice, which might lead to unexpected result """
|
||||
tmp = torch.zeros_like(self.weight[self.out_channels_per_joint * i: self.out_channels_per_joint * (i + 1),
|
||||
neighbour, ...])
|
||||
nn.init.kaiming_uniform_(tmp, a=math.sqrt(5))
|
||||
self.weight[self.out_channels_per_joint * i: self.out_channels_per_joint * (i + 1),
|
||||
neighbour, ...] = tmp
|
||||
if self.bias is not None:
|
||||
fan_in, _ = nn.init._calculate_fan_in_and_fan_out(
|
||||
self.weight[self.out_channels_per_joint * i: self.out_channels_per_joint * (i + 1), neighbour, ...])
|
||||
bound = 1 / math.sqrt(fan_in)
|
||||
tmp = torch.zeros_like(
|
||||
self.bias[self.out_channels_per_joint * i: self.out_channels_per_joint * (i + 1)])
|
||||
nn.init.uniform_(tmp, -bound, bound)
|
||||
self.bias[self.out_channels_per_joint * i: self.out_channels_per_joint * (i + 1)] = tmp
|
||||
|
||||
self.weight = nn.Parameter(self.weight)
|
||||
if self.bias is not None:
|
||||
self.bias = nn.Parameter(self.bias)
|
||||
|
||||
def set_offset(self, offset):
|
||||
if not self.add_offset:
|
||||
raise Exception('Wrong Combination of Parameters')
|
||||
self.offset = offset.reshape(offset.shape[0], -1)
|
||||
|
||||
def forward(self, input):
|
||||
# print('SkeletonConv')
|
||||
weight_masked = self.weight * self.mask
|
||||
# print(f'input: {input.size()}')
|
||||
res = F.conv1d(F.pad(input, self._padding_repeated_twice, mode=self.padding_mode),
|
||||
weight_masked, self.bias, self.stride,
|
||||
0, self.dilation, self.groups)
|
||||
|
||||
if self.add_offset:
|
||||
offset_res = self.offset_enc(self.offset)
|
||||
offset_res = offset_res.reshape(offset_res.shape + (1, ))
|
||||
res += offset_res / 100
|
||||
# print(f'res: {res.size()}')
|
||||
return res
|
||||
|
||||
|
||||
class SkeletonLinear(nn.Module):
|
||||
def __init__(self, neighbour_list, in_channels, out_channels, extra_dim1=False):
|
||||
super(SkeletonLinear, self).__init__()
|
||||
self.neighbour_list = neighbour_list
|
||||
self.in_channels = in_channels
|
||||
self.out_channels = out_channels
|
||||
self.in_channels_per_joint = in_channels // len(neighbour_list)
|
||||
self.out_channels_per_joint = out_channels // len(neighbour_list)
|
||||
self.extra_dim1 = extra_dim1
|
||||
self.expanded_neighbour_list = []
|
||||
|
||||
for neighbour in neighbour_list:
|
||||
expanded = []
|
||||
for k in neighbour:
|
||||
for i in range(self.in_channels_per_joint):
|
||||
expanded.append(k * self.in_channels_per_joint + i)
|
||||
self.expanded_neighbour_list.append(expanded)
|
||||
|
||||
self.weight = torch.zeros(out_channels, in_channels)
|
||||
self.mask = torch.zeros(out_channels, in_channels)
|
||||
self.bias = nn.Parameter(torch.Tensor(out_channels))
|
||||
|
||||
self.reset_parameters()
|
||||
|
||||
def reset_parameters(self):
|
||||
for i, neighbour in enumerate(self.expanded_neighbour_list):
|
||||
tmp = torch.zeros_like(
|
||||
self.weight[i*self.out_channels_per_joint: (i + 1)*self.out_channels_per_joint, neighbour]
|
||||
)
|
||||
self.mask[i*self.out_channels_per_joint: (i + 1)*self.out_channels_per_joint, neighbour] = 1
|
||||
nn.init.kaiming_uniform_(tmp, a=math.sqrt(5))
|
||||
self.weight[i*self.out_channels_per_joint: (i + 1)*self.out_channels_per_joint, neighbour] = tmp
|
||||
|
||||
fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weight)
|
||||
bound = 1 / math.sqrt(fan_in)
|
||||
nn.init.uniform_(self.bias, -bound, bound)
|
||||
|
||||
self.weight = nn.Parameter(self.weight)
|
||||
self.mask = nn.Parameter(self.mask, requires_grad=False)
|
||||
|
||||
def forward(self, input):
|
||||
input = input.reshape(input.shape[0], -1)
|
||||
weight_masked = self.weight * self.mask
|
||||
res = F.linear(input, weight_masked, self.bias)
|
||||
if self.extra_dim1:
|
||||
res = res.reshape(res.shape + (1,))
|
||||
return res
|
||||
|
||||
|
||||
class SkeletonPool(nn.Module):
|
||||
def __init__(self, edges, pooling_mode, channels_per_edge, last_pool=False):
|
||||
super(SkeletonPool, self).__init__()
|
||||
|
||||
if pooling_mode != 'mean':
|
||||
raise Exception('Unimplemented pooling mode in matrix_implementation')
|
||||
|
||||
self.channels_per_edge = channels_per_edge
|
||||
self.pooling_mode = pooling_mode
|
||||
self.edge_num = len(edges)
|
||||
# self.edge_num = len(edges) + 1
|
||||
self.seq_list = []
|
||||
self.pooling_list = []
|
||||
self.new_edges = []
|
||||
degree = [0] * 100 # each element represents the degree of the corresponding joint
|
||||
|
||||
for edge in edges:
|
||||
degree[edge[0]] += 1
|
||||
degree[edge[1]] += 1
|
||||
|
||||
# seq_list contains multiple sub-lists where each sub-list is an edge chain from the joint whose degree > 2 to the end effectors or joints whose degree > 2.
|
||||
def find_seq(j, seq):
|
||||
nonlocal self, degree, edges
|
||||
|
||||
if degree[j] > 2 and j != 0:
|
||||
self.seq_list.append(seq)
|
||||
seq = []
|
||||
|
||||
if degree[j] == 1:
|
||||
self.seq_list.append(seq)
|
||||
return
|
||||
|
||||
for idx, edge in enumerate(edges):
|
||||
if edge[0] == j:
|
||||
find_seq(edge[1], seq + [idx])
|
||||
|
||||
find_seq(0, [])
|
||||
# print(f'self.seq_list: {self.seq_list}')
|
||||
|
||||
for seq in self.seq_list:
|
||||
if last_pool:
|
||||
self.pooling_list.append(seq)
|
||||
continue
|
||||
if len(seq) % 2 == 1:
|
||||
self.pooling_list.append([seq[0]])
|
||||
self.new_edges.append(edges[seq[0]])
|
||||
seq = seq[1:]
|
||||
for i in range(0, len(seq), 2):
|
||||
self.pooling_list.append([seq[i], seq[i + 1]])
|
||||
self.new_edges.append([edges[seq[i]][0], edges[seq[i + 1]][1]])
|
||||
# print(f'self.pooling_list: {self.pooling_list}')
|
||||
# print(f'self.new_egdes: {self.new_edges}')
|
||||
|
||||
# add global position
|
||||
# self.pooling_list.append([self.edge_num - 1])
|
||||
|
||||
self.description = 'SkeletonPool(in_edge_num={}, out_edge_num={})'.format(
|
||||
len(edges), len(self.pooling_list)
|
||||
)
|
||||
|
||||
self.weight = torch.zeros(len(self.pooling_list) * channels_per_edge, self.edge_num * channels_per_edge)
|
||||
|
||||
for i, pair in enumerate(self.pooling_list):
|
||||
for j in pair:
|
||||
for c in range(channels_per_edge):
|
||||
self.weight[i * channels_per_edge + c, j * channels_per_edge + c] = 1.0 / len(pair)
|
||||
|
||||
self.weight = nn.Parameter(self.weight, requires_grad=False)
|
||||
|
||||
def forward(self, input: torch.Tensor):
|
||||
# print('SkeletonPool')
|
||||
# print(f'input: {input.size()}')
|
||||
# print(f'self.weight: {self.weight.size()}')
|
||||
return torch.matmul(self.weight, input)
|
||||
|
||||
|
||||
class SkeletonUnpool(nn.Module):
|
||||
def __init__(self, pooling_list, channels_per_edge):
|
||||
super(SkeletonUnpool, self).__init__()
|
||||
self.pooling_list = pooling_list
|
||||
self.input_edge_num = len(pooling_list)
|
||||
self.output_edge_num = 0
|
||||
self.channels_per_edge = channels_per_edge
|
||||
for t in self.pooling_list:
|
||||
self.output_edge_num += len(t)
|
||||
|
||||
self.description = 'SkeletonUnpool(in_edge_num={}, out_edge_num={})'.format(
|
||||
self.input_edge_num, self.output_edge_num,
|
||||
)
|
||||
|
||||
self.weight = torch.zeros(self.output_edge_num * channels_per_edge, self.input_edge_num * channels_per_edge)
|
||||
|
||||
for i, pair in enumerate(self.pooling_list):
|
||||
for j in pair:
|
||||
for c in range(channels_per_edge):
|
||||
self.weight[j * channels_per_edge + c, i * channels_per_edge + c] = 1
|
||||
|
||||
self.weight = nn.Parameter(self.weight)
|
||||
self.weight.requires_grad_(False)
|
||||
|
||||
def forward(self, input: torch.Tensor):
|
||||
# print('SkeletonUnpool')
|
||||
# print(f'input: {input.size()}')
|
||||
# print(f'self.weight: {self.weight.size()}')
|
||||
return torch.matmul(self.weight, input)
|
||||
|
||||
|
||||
"""
|
||||
Helper functions for skeleton operation
|
||||
"""
|
||||
|
||||
|
||||
def dfs(x, fa, vis, dist):
|
||||
vis[x] = 1
|
||||
for y in range(len(fa)):
|
||||
if (fa[y] == x or fa[x] == y) and vis[y] == 0:
|
||||
dist[y] = dist[x] + 1
|
||||
dfs(y, fa, vis, dist)
|
||||
|
||||
|
||||
"""
|
||||
def find_neighbor_joint(fa, threshold):
|
||||
neighbor_list = [[]]
|
||||
for x in range(1, len(fa)):
|
||||
vis = [0 for _ in range(len(fa))]
|
||||
dist = [0 for _ in range(len(fa))]
|
||||
dist[0] = 10000
|
||||
dfs(x, fa, vis, dist)
|
||||
neighbor = []
|
||||
for j in range(1, len(fa)):
|
||||
if dist[j] <= threshold:
|
||||
neighbor.append(j)
|
||||
neighbor_list.append(neighbor)
|
||||
|
||||
neighbor = [0]
|
||||
for i, x in enumerate(neighbor_list):
|
||||
if i == 0: continue
|
||||
if 1 in x:
|
||||
neighbor.append(i)
|
||||
neighbor_list[i] = [0] + neighbor_list[i]
|
||||
neighbor_list[0] = neighbor
|
||||
return neighbor_list
|
||||
|
||||
|
||||
def build_edge_topology(topology, offset):
|
||||
# get all edges (pa, child, offset)
|
||||
edges = []
|
||||
joint_num = len(topology)
|
||||
for i in range(1, joint_num):
|
||||
edges.append((topology[i], i, offset[i]))
|
||||
return edges
|
||||
"""
|
||||
|
||||
|
||||
def build_edge_topology(topology):
|
||||
# get all edges (pa, child)
|
||||
edges = []
|
||||
joint_num = len(topology)
|
||||
edges.append((0, joint_num)) # add an edge between the root joint and a virtual joint
|
||||
for i in range(1, joint_num):
|
||||
edges.append((topology[i], i))
|
||||
return edges
|
||||
|
||||
|
||||
def build_joint_topology(edges, origin_names):
|
||||
parent = []
|
||||
offset = []
|
||||
names = []
|
||||
edge2joint = []
|
||||
joint_from_edge = [] # -1 means virtual joint
|
||||
joint_cnt = 0
|
||||
out_degree = [0] * (len(edges) + 10)
|
||||
for edge in edges:
|
||||
out_degree[edge[0]] += 1
|
||||
|
||||
# add root joint
|
||||
joint_from_edge.append(-1)
|
||||
parent.append(0)
|
||||
offset.append(np.array([0, 0, 0]))
|
||||
names.append(origin_names[0])
|
||||
joint_cnt += 1
|
||||
|
||||
def make_topology(edge_idx, pa):
|
||||
nonlocal edges, parent, offset, names, edge2joint, joint_from_edge, joint_cnt
|
||||
edge = edges[edge_idx]
|
||||
if out_degree[edge[0]] > 1:
|
||||
parent.append(pa)
|
||||
offset.append(np.array([0, 0, 0]))
|
||||
names.append(origin_names[edge[1]] + '_virtual')
|
||||
edge2joint.append(-1)
|
||||
pa = joint_cnt
|
||||
joint_cnt += 1
|
||||
|
||||
parent.append(pa)
|
||||
offset.append(edge[2])
|
||||
names.append(origin_names[edge[1]])
|
||||
edge2joint.append(edge_idx)
|
||||
pa = joint_cnt
|
||||
joint_cnt += 1
|
||||
|
||||
for idx, e in enumerate(edges):
|
||||
if e[0] == edge[1]:
|
||||
make_topology(idx, pa)
|
||||
|
||||
for idx, e in enumerate(edges):
|
||||
if e[0] == 0:
|
||||
make_topology(idx, 0)
|
||||
|
||||
return parent, offset, names, edge2joint
|
||||
|
||||
|
||||
def calc_edge_mat(edges):
|
||||
edge_num = len(edges)
|
||||
# edge_mat[i][j] = distance between edge(i) and edge(j)
|
||||
edge_mat = [[100000] * edge_num for _ in range(edge_num)]
|
||||
for i in range(edge_num):
|
||||
edge_mat[i][i] = 0
|
||||
|
||||
# initialize edge_mat with direct neighbor
|
||||
for i, a in enumerate(edges):
|
||||
for j, b in enumerate(edges):
|
||||
link = 0
|
||||
for x in range(2):
|
||||
for y in range(2):
|
||||
if a[x] == b[y]:
|
||||
link = 1
|
||||
if link:
|
||||
edge_mat[i][j] = 1
|
||||
|
||||
# calculate all the pairs distance
|
||||
for k in range(edge_num):
|
||||
for i in range(edge_num):
|
||||
for j in range(edge_num):
|
||||
edge_mat[i][j] = min(edge_mat[i][j], edge_mat[i][k] + edge_mat[k][j])
|
||||
return edge_mat
|
||||
|
||||
|
||||
def find_neighbor(edges, d):
|
||||
"""
|
||||
Args:
|
||||
edges: The list contains N elements, each element represents (parent, child).
|
||||
d: Distance between edges (the distance of the same edge is 0 and the distance of adjacent edges is 1).
|
||||
|
||||
Returns:
|
||||
The list contains N elements, each element is a list of edge indices whose distance <= d.
|
||||
"""
|
||||
edge_mat = calc_edge_mat(edges)
|
||||
neighbor_list = []
|
||||
edge_num = len(edge_mat)
|
||||
for i in range(edge_num):
|
||||
neighbor = []
|
||||
for j in range(edge_num):
|
||||
if edge_mat[i][j] <= d:
|
||||
neighbor.append(j)
|
||||
neighbor_list.append(neighbor)
|
||||
|
||||
# # add neighbor for global part
|
||||
# global_part_neighbor = neighbor_list[0].copy()
|
||||
# """
|
||||
# Line #373 is buggy. Thanks @crissallan!!
|
||||
# See issue #30 (https://github.com/DeepMotionEditing/deep-motion-editing/issues/30)
|
||||
# However, fixing this bug will make it unable to load the pretrained model and
|
||||
# affect the reproducibility of quantitative error reported in the paper.
|
||||
# It is not a fatal bug so we didn't touch it and we are looking for possible solutions.
|
||||
# """
|
||||
# for i in global_part_neighbor:
|
||||
# neighbor_list[i].append(edge_num)
|
||||
# neighbor_list.append(global_part_neighbor)
|
||||
|
||||
return neighbor_list
|
||||
|
||||
|
||||
def calc_node_depth(topology):
|
||||
def dfs(node, topology):
|
||||
if topology[node] < 0:
|
||||
return 0
|
||||
return 1 + dfs(topology[node], topology)
|
||||
depth = []
|
||||
for i in range(len(topology)):
|
||||
depth.append(dfs(i, topology))
|
||||
|
||||
return depth
|
||||
+652
@@ -0,0 +1,652 @@
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
||||
# All rights reserved.
|
||||
#
|
||||
# This source code is licensed under the BSD-style license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
import numpy as np
|
||||
from typing import Optional, Union
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
Device = Union[str, torch.device]
|
||||
|
||||
|
||||
"""
|
||||
The transformation matrices returned from the functions in this file assume
|
||||
the points on which the transformation will be applied are column vectors.
|
||||
i.e. the R matrix is structured as
|
||||
|
||||
R = [
|
||||
[Rxx, Rxy, Rxz],
|
||||
[Ryx, Ryy, Ryz],
|
||||
[Rzx, Rzy, Rzz],
|
||||
] # (3, 3)
|
||||
|
||||
This matrix can be applied to column vectors by post multiplication
|
||||
by the points e.g.
|
||||
|
||||
points = [[0], [1], [2]] # (3 x 1) xyz coordinates of a point
|
||||
transformed_points = R * points
|
||||
|
||||
To apply the same matrix to points which are row vectors, the R matrix
|
||||
can be transposed and pre multiplied by the points:
|
||||
|
||||
e.g.
|
||||
points = [[0, 1, 2]] # (1 x 3) xyz coordinates of a point
|
||||
transformed_points = points * R.transpose(1, 0)
|
||||
"""
|
||||
|
||||
|
||||
def quaternion_to_matrix(quaternions: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Convert rotations given as quaternions to rotation matrices.
|
||||
|
||||
Args:
|
||||
quaternions: quaternions with real part first,
|
||||
as tensor of shape (..., 4).
|
||||
|
||||
Returns:
|
||||
Rotation matrices as tensor of shape (..., 3, 3).
|
||||
"""
|
||||
r, i, j, k = torch.unbind(quaternions, -1)
|
||||
two_s = 2.0 / (quaternions * quaternions).sum(-1)
|
||||
|
||||
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 _copysign(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Return a tensor where each element has the absolute value taken from the,
|
||||
corresponding element of a, with sign taken from the corresponding
|
||||
element of b. This is like the standard copysign floating-point operation,
|
||||
but is not careful about negative 0 and NaN.
|
||||
|
||||
Args:
|
||||
a: source tensor.
|
||||
b: tensor whose signs will be used, of the same shape as a.
|
||||
|
||||
Returns:
|
||||
Tensor of the same shape as a with the signs of b.
|
||||
"""
|
||||
signs_differ = (a < 0) != (b < 0)
|
||||
return torch.where(signs_differ, -a, a)
|
||||
|
||||
|
||||
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
|
||||
ret[positive_mask] = torch.sqrt(x[positive_mask])
|
||||
return ret
|
||||
|
||||
|
||||
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(
|
||||
[
|
||||
torch.stack([q_abs[..., 0] ** 2, m21 - m12, m02 - m20, m10 - m01], dim=-1),
|
||||
torch.stack([m21 - m12, q_abs[..., 1] ** 2, m10 + m01, m02 + m20], dim=-1),
|
||||
torch.stack([m02 - m20, m10 + m01, q_abs[..., 2] ** 2, m12 + m21], dim=-1),
|
||||
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)
|
||||
|
||||
return quat_candidates[
|
||||
F.one_hot(q_abs.argmax(dim=-1), num_classes=4) > 0.5, : # pyre-ignore[16]
|
||||
].reshape(batch_dim + (4,))
|
||||
|
||||
|
||||
def _axis_angle_rotation(axis: str, angle: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Return the rotation matrices for one of the rotations about an axis
|
||||
of which Euler angles describe, for each value of the angle given.
|
||||
|
||||
Args:
|
||||
axis: Axis label "X" or "Y or "Z".
|
||||
angle: any shape tensor of Euler angles in radians
|
||||
|
||||
Returns:
|
||||
Rotation matrices as tensor of shape (..., 3, 3).
|
||||
"""
|
||||
|
||||
cos = torch.cos(angle)
|
||||
sin = torch.sin(angle)
|
||||
one = torch.ones_like(angle)
|
||||
zero = torch.zeros_like(angle)
|
||||
|
||||
if axis == "X":
|
||||
R_flat = (one, zero, zero, zero, cos, -sin, zero, sin, cos)
|
||||
elif axis == "Y":
|
||||
R_flat = (cos, zero, sin, zero, one, zero, -sin, zero, cos)
|
||||
elif axis == "Z":
|
||||
R_flat = (cos, -sin, zero, sin, cos, zero, zero, zero, one)
|
||||
else:
|
||||
raise ValueError("letter must be either X, Y or Z.")
|
||||
|
||||
return torch.stack(R_flat, -1).reshape(angle.shape + (3, 3))
|
||||
|
||||
|
||||
def euler_angles_to_matrix(euler_angles: torch.Tensor, convention: str) -> torch.Tensor:
|
||||
"""
|
||||
Convert rotations given as Euler angles in radians to rotation matrices.
|
||||
|
||||
Args:
|
||||
euler_angles: Euler angles in radians as tensor of shape (..., 3).
|
||||
convention: Convention string of three uppercase letters from
|
||||
{"X", "Y", and "Z"}.
|
||||
|
||||
Returns:
|
||||
Rotation matrices as tensor of shape (..., 3, 3).
|
||||
"""
|
||||
if euler_angles.dim() == 0 or euler_angles.shape[-1] != 3:
|
||||
raise ValueError("Invalid input euler angles.")
|
||||
if len(convention) != 3:
|
||||
raise ValueError("Convention must have 3 letters.")
|
||||
if convention[1] in (convention[0], convention[2]):
|
||||
raise ValueError(f"Invalid convention {convention}.")
|
||||
for letter in convention:
|
||||
if letter not in ("X", "Y", "Z"):
|
||||
raise ValueError(f"Invalid letter {letter} in convention string.")
|
||||
matrices = [
|
||||
_axis_angle_rotation(c, e)
|
||||
for c, e in zip(convention, torch.unbind(euler_angles, -1))
|
||||
]
|
||||
# return functools.reduce(torch.matmul, matrices)
|
||||
return torch.matmul(torch.matmul(matrices[0], matrices[1]), matrices[2])
|
||||
|
||||
|
||||
def _angle_from_tan(
|
||||
axis: str, other_axis: str, data, horizontal: bool, tait_bryan: bool
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Extract the first or third Euler angle from the two members of
|
||||
the matrix which are positive constant times its sine and cosine.
|
||||
|
||||
Args:
|
||||
axis: Axis label "X" or "Y or "Z" for the angle we are finding.
|
||||
other_axis: Axis label "X" or "Y or "Z" for the middle axis in the
|
||||
convention.
|
||||
data: Rotation matrices as tensor of shape (..., 3, 3).
|
||||
horizontal: Whether we are looking for the angle for the third axis,
|
||||
which means the relevant entries are in the same row of the
|
||||
rotation matrix. If not, they are in the same column.
|
||||
tait_bryan: Whether the first and third axes in the convention differ.
|
||||
|
||||
Returns:
|
||||
Euler Angles in radians for each matrix in data as a tensor
|
||||
of shape (...).
|
||||
"""
|
||||
|
||||
i1, i2 = {"X": (2, 1), "Y": (0, 2), "Z": (1, 0)}[axis]
|
||||
if horizontal:
|
||||
i2, i1 = i1, i2
|
||||
even = (axis + other_axis) in ["XY", "YZ", "ZX"]
|
||||
if horizontal == even:
|
||||
return torch.atan2(data[..., i1], data[..., i2])
|
||||
if tait_bryan:
|
||||
return torch.atan2(-data[..., i2], data[..., i1])
|
||||
return torch.atan2(data[..., i2], -data[..., i1])
|
||||
|
||||
|
||||
def _index_from_letter(letter: str) -> int:
|
||||
if letter == "X":
|
||||
return 0
|
||||
if letter == "Y":
|
||||
return 1
|
||||
if letter == "Z":
|
||||
return 2
|
||||
raise ValueError("letter must be either X, Y or Z.")
|
||||
|
||||
|
||||
def matrix_to_euler_angles(matrix: torch.Tensor, convention: str) -> torch.Tensor:
|
||||
"""
|
||||
Convert rotations given as rotation matrices to Euler angles in radians.
|
||||
|
||||
Args:
|
||||
matrix: Rotation matrices as tensor of shape (..., 3, 3).
|
||||
convention: Convention string of three uppercase letters.
|
||||
|
||||
Returns:
|
||||
Euler angles in radians as tensor of shape (..., 3).
|
||||
"""
|
||||
if len(convention) != 3:
|
||||
raise ValueError("Convention must have 3 letters.")
|
||||
if convention[1] in (convention[0], convention[2]):
|
||||
raise ValueError(f"Invalid convention {convention}.")
|
||||
for letter in convention:
|
||||
if letter not in ("X", "Y", "Z"):
|
||||
raise ValueError(f"Invalid letter {letter} in convention string.")
|
||||
if matrix.size(-1) != 3 or matrix.size(-2) != 3:
|
||||
raise ValueError(f"Invalid rotation matrix shape {matrix.shape}.")
|
||||
i0 = _index_from_letter(convention[0])
|
||||
i2 = _index_from_letter(convention[2])
|
||||
tait_bryan = i0 != i2
|
||||
if tait_bryan:
|
||||
central_angle = torch.asin(
|
||||
matrix[..., i0, i2] * (-1.0 if i0 - i2 in [-1, 2] else 1.0)
|
||||
)
|
||||
else:
|
||||
central_angle = torch.acos(matrix[..., i0, i0])
|
||||
|
||||
o = (
|
||||
_angle_from_tan(
|
||||
convention[0], convention[1], matrix[..., i2], False, tait_bryan
|
||||
),
|
||||
central_angle,
|
||||
_angle_from_tan(
|
||||
convention[2], convention[1], matrix[..., i0, :], True, tait_bryan
|
||||
),
|
||||
)
|
||||
return torch.stack(o, -1)
|
||||
|
||||
|
||||
def random_quaternions(
|
||||
n: int, dtype: Optional[torch.dtype] = None, device: Optional[Device] = None
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Generate random quaternions representing rotations,
|
||||
i.e. versors with nonnegative real part.
|
||||
|
||||
Args:
|
||||
n: Number of quaternions in a batch to return.
|
||||
dtype: Type to return.
|
||||
device: Desired device of returned tensor. Default:
|
||||
uses the current device for the default tensor type.
|
||||
|
||||
Returns:
|
||||
Quaternions as tensor of shape (N, 4).
|
||||
"""
|
||||
if isinstance(device, str):
|
||||
device = torch.device(device)
|
||||
o = torch.randn((n, 4), dtype=dtype, device=device)
|
||||
s = (o * o).sum(1)
|
||||
o = o / _copysign(torch.sqrt(s), o[:, 0])[:, None]
|
||||
return o
|
||||
|
||||
|
||||
def random_rotations(
|
||||
n: int, dtype: Optional[torch.dtype] = None, device: Optional[Device] = None
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Generate random rotations as 3x3 rotation matrices.
|
||||
|
||||
Args:
|
||||
n: Number of rotation matrices in a batch to return.
|
||||
dtype: Type to return.
|
||||
device: Device of returned tensor. Default: if None,
|
||||
uses the current device for the default tensor type.
|
||||
|
||||
Returns:
|
||||
Rotation matrices as tensor of shape (n, 3, 3).
|
||||
"""
|
||||
quaternions = random_quaternions(n, dtype=dtype, device=device)
|
||||
return quaternion_to_matrix(quaternions)
|
||||
|
||||
|
||||
def random_rotation(
|
||||
dtype: Optional[torch.dtype] = None, device: Optional[Device] = None
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Generate a single random 3x3 rotation matrix.
|
||||
|
||||
Args:
|
||||
dtype: Type to return
|
||||
device: Device of returned tensor. Default: if None,
|
||||
uses the current device for the default tensor type
|
||||
|
||||
Returns:
|
||||
Rotation matrix as tensor of shape (3, 3).
|
||||
"""
|
||||
return random_rotations(1, dtype, device)[0]
|
||||
|
||||
|
||||
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 quaternion_raw_multiply(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Multiply two quaternions.
|
||||
Usual torch rules for broadcasting apply.
|
||||
|
||||
Args:
|
||||
a: Quaternions as tensor of shape (..., 4), real part first.
|
||||
b: Quaternions as tensor of shape (..., 4), real part first.
|
||||
|
||||
Returns:
|
||||
The product of a and b, a tensor of quaternions shape (..., 4).
|
||||
"""
|
||||
aw, ax, ay, az = torch.unbind(a, -1)
|
||||
bw, bx, by, bz = torch.unbind(b, -1)
|
||||
ow = aw * bw - ax * bx - ay * by - az * bz
|
||||
ox = aw * bx + ax * bw + ay * bz - az * by
|
||||
oy = aw * by - ax * bz + ay * bw + az * bx
|
||||
oz = aw * bz + ax * by - ay * bx + az * bw
|
||||
return torch.stack((ow, ox, oy, oz), -1)
|
||||
|
||||
|
||||
def quaternion_multiply(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Multiply two quaternions representing rotations, returning the quaternion
|
||||
representing their composition, i.e. the versor with nonnegative real part.
|
||||
Usual torch rules for broadcasting apply.
|
||||
|
||||
Args:
|
||||
a: Quaternions as tensor of shape (..., 4), real part first.
|
||||
b: Quaternions as tensor of shape (..., 4), real part first.
|
||||
|
||||
Returns:
|
||||
The product of a and b, a tensor of quaternions of shape (..., 4).
|
||||
"""
|
||||
ab = quaternion_raw_multiply(a, b)
|
||||
return standardize_quaternion(ab)
|
||||
|
||||
|
||||
def quaternion_invert(quaternion: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Given a quaternion representing rotation, get the quaternion representing
|
||||
its inverse.
|
||||
|
||||
Args:
|
||||
quaternion: Quaternions as tensor of shape (..., 4), with real part
|
||||
first, which must be versors (unit quaternions).
|
||||
|
||||
Returns:
|
||||
The inverse, a tensor of quaternions of shape (..., 4).
|
||||
"""
|
||||
|
||||
scaling = torch.tensor([1, -1, -1, -1], device=quaternion.device)
|
||||
return quaternion * scaling
|
||||
|
||||
|
||||
def quaternion_apply(quaternion: torch.Tensor, point: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Apply the rotation given by a quaternion to a 3D point.
|
||||
Usual torch rules for broadcasting apply.
|
||||
|
||||
Args:
|
||||
quaternion: Tensor of quaternions, real part first, of shape (..., 4).
|
||||
point: Tensor of 3D points of shape (..., 3).
|
||||
|
||||
Returns:
|
||||
Tensor of rotated points of shape (..., 3).
|
||||
"""
|
||||
if point.size(-1) != 3:
|
||||
raise ValueError(f"Points are not in 3D, {point.shape}.")
|
||||
real_parts = point.new_zeros(point.shape[:-1] + (1,))
|
||||
point_as_quaternion = torch.cat((real_parts, point), -1)
|
||||
out = quaternion_raw_multiply(
|
||||
quaternion_raw_multiply(quaternion, point_as_quaternion),
|
||||
quaternion_invert(quaternion),
|
||||
)
|
||||
return out[..., 1:]
|
||||
|
||||
|
||||
def axis_angle_to_matrix(axis_angle: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Convert rotations given as axis/angle to rotation matrices.
|
||||
|
||||
Args:
|
||||
axis_angle: Rotations given as a vector in axis angle form,
|
||||
as a tensor of shape (..., 3), where the magnitude is
|
||||
the angle turned anticlockwise in radians around the
|
||||
vector's direction.
|
||||
|
||||
Returns:
|
||||
Rotation matrices as tensor of shape (..., 3, 3).
|
||||
"""
|
||||
return quaternion_to_matrix(axis_angle_to_quaternion(axis_angle))
|
||||
|
||||
|
||||
def matrix_to_axis_angle(matrix: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Convert rotations given as rotation matrices to axis/angle.
|
||||
|
||||
Args:
|
||||
matrix: Rotation matrices as tensor of shape (..., 3, 3).
|
||||
|
||||
Returns:
|
||||
Rotations given as a vector in axis angle form, as a tensor
|
||||
of shape (..., 3), where the magnitude is the angle
|
||||
turned anticlockwise in radians around the vector's
|
||||
direction.
|
||||
"""
|
||||
return quaternion_to_axis_angle(matrix_to_quaternion(matrix))
|
||||
|
||||
|
||||
def axis_angle_to_quaternion(axis_angle: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Convert rotations given as axis/angle to quaternions.
|
||||
|
||||
Args:
|
||||
axis_angle: Rotations given as a vector in axis angle form,
|
||||
as a tensor of shape (..., 3), where the magnitude is
|
||||
the angle turned anticlockwise in radians around the
|
||||
vector's direction.
|
||||
|
||||
Returns:
|
||||
quaternions with real part first, as tensor of shape (..., 4).
|
||||
"""
|
||||
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 quat_to_aa(quaternions: torch.Tensor) -> torch.Tensor:
|
||||
|
||||
norms = torch.norm(quaternions[..., 1:], p=2, dim=-1, keepdim=True)
|
||||
half_angles = torch.atan2(norms, quaternions[..., :1])
|
||||
angles = 2 * half_angles
|
||||
|
||||
# if angles.max() > torch.pi:
|
||||
# angles = angles % torch.pi
|
||||
# half_angles = angles / 2
|
||||
|
||||
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])
|
||||
|
||||
sin_half_angles_over_angles[~small_angles] = (
|
||||
torch.tensor(np.sin(half_angles[~small_angles].cpu()).to(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
|
||||
)
|
||||
return quaternions[..., 1:] / sin_half_angles_over_angles
|
||||
|
||||
|
||||
def quaternion_to_axis_angle(quaternions: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Convert rotations given as quaternions to axis/angle.
|
||||
|
||||
Args:
|
||||
quaternions: quaternions with real part first,
|
||||
as tensor of shape (..., 4).
|
||||
|
||||
Returns:
|
||||
Rotations given as a vector in axis angle form, as a tensor
|
||||
of shape (..., 3), where the magnitude is the angle
|
||||
turned anticlockwise in radians around the vector's
|
||||
direction.
|
||||
"""
|
||||
norms = torch.norm(quaternions[..., 1:], p=2, dim=-1, keepdim=True)
|
||||
half_angles = torch.atan2(norms, quaternions[..., :1])
|
||||
angles = 2 * half_angles
|
||||
|
||||
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
|
||||
)
|
||||
return quaternions[..., 1:] / sin_half_angles_over_angles
|
||||
|
||||
|
||||
def rotation_6d_to_matrix(d6: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Converts 6D rotation representation by Zhou et al. [1] to rotation matrix
|
||||
using Gram--Schmidt orthogonalization per Section B of [1].
|
||||
Args:
|
||||
d6: 6D rotation representation, of size (*, 6)
|
||||
|
||||
Returns:
|
||||
batch of rotation matrices of size (*, 3, 3)
|
||||
|
||||
[1] Zhou, Y., Barnes, C., Lu, J., Yang, J., & Li, H.
|
||||
On the Continuity of Rotation Representations in Neural Networks.
|
||||
IEEE Conference on Computer Vision and Pattern Recognition, 2019.
|
||||
Retrieved from http://arxiv.org/abs/1812.07035
|
||||
"""
|
||||
|
||||
a1, a2 = d6[..., :3], d6[..., 3:]
|
||||
b1 = F.normalize(a1, dim=-1)
|
||||
b2 = a2 - (b1 * a2).sum(-1, keepdim=True) * b1
|
||||
b2 = F.normalize(b2, dim=-1)
|
||||
b3 = torch.cross(b1, b2, dim=-1)
|
||||
return torch.stack((b1, b2, b3), dim=-2)
|
||||
|
||||
|
||||
def matrix_to_rotation_6d(matrix: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Converts rotation matrices to 6D rotation representation by Zhou et al. [1]
|
||||
by dropping the last row. Note that 6D representation is not unique.
|
||||
Args:
|
||||
matrix: batch of rotation matrices of size (*, 3, 3)
|
||||
|
||||
Returns:
|
||||
6D rotation representation, of size (*, 6)
|
||||
|
||||
[1] Zhou, Y., Barnes, C., Lu, J., Yang, J., & Li, H.
|
||||
On the Continuity of Rotation Representations in Neural Networks.
|
||||
IEEE Conference on Computer Vision and Pattern Recognition, 2019.
|
||||
Retrieved from http://arxiv.org/abs/1812.07035
|
||||
"""
|
||||
batch_dim = matrix.size()[:-2]
|
||||
return matrix[..., :2, :].clone().reshape(batch_dim + (6,))
|
||||
|
||||
|
||||
def batch_rodrigues(rot_vecs, epsilon=1e-8, dtype=torch.float32):
|
||||
''' Calculates the rotation matrices for a batch of rotation vectors
|
||||
Parameters
|
||||
----------
|
||||
rot_vecs: torch.tensor Nx3
|
||||
array of N axis-angle vectors
|
||||
Returns
|
||||
-------
|
||||
R: torch.tensor Nx3x3
|
||||
The rotation matrices for the given axis-angle parameters
|
||||
'''
|
||||
|
||||
batch_size = rot_vecs.shape[0]
|
||||
device = rot_vecs.device
|
||||
|
||||
angle = torch.norm(rot_vecs + 1e-8, dim=1, keepdim=True)
|
||||
rot_dir = rot_vecs / angle
|
||||
|
||||
cos = torch.unsqueeze(torch.cos(angle), dim=1)
|
||||
sin = torch.unsqueeze(torch.sin(angle), dim=1)
|
||||
|
||||
# Bx1 arrays
|
||||
rx, ry, rz = torch.split(rot_dir, 1, dim=1)
|
||||
K = torch.zeros((batch_size, 3, 3), dtype=dtype, device=device)
|
||||
|
||||
zeros = torch.zeros((batch_size, 1), dtype=dtype, device=device)
|
||||
K = torch.cat([zeros, -rz, ry, rz, zeros, -rx, -ry, rx, zeros], dim=1) \
|
||||
.view((batch_size, 3, 3))
|
||||
|
||||
ident = torch.eye(3, dtype=dtype, device=device).unsqueeze(dim=0)
|
||||
rot_mat = ident + sin * K + (1 - cos) * torch.bmm(K, K)
|
||||
return rot_mat
|
||||
+473
@@ -0,0 +1,473 @@
|
||||
import os
|
||||
import cv2
|
||||
import numpy as np
|
||||
import plyfile
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from human_body_prior.tools.omni_tools import copy2cpu as c2c
|
||||
from scipy.spatial.transform import Rotation as R
|
||||
from scipy.spatial.transform import Slerp
|
||||
|
||||
import holden.BVH as BVH
|
||||
from holden.Animation import Animation
|
||||
from holden.Quaternions import Quaternions
|
||||
from typing import List, Tuple, NewType
|
||||
import math
|
||||
Tensor = NewType('Tensor', torch.Tensor)
|
||||
|
||||
SMPL_JOINTS = {
|
||||
'Pelvis': 0,
|
||||
'L_Hip': 1, 'L_Knee': 4, 'L_Ankle': 7, 'L_Foot': 10,
|
||||
'R_Hip': 2, 'R_Knee': 5, 'R_Ankle': 8, 'R_Foot': 11,
|
||||
'Spine1': 3, 'Spine2': 6, 'Spine3': 9, 'Neck': 12, 'Head': 15,
|
||||
'L_Collar': 13, 'L_Shoulder': 16, 'L_Elbow': 18, 'L_Wrist': 20, 'L_Hand': 22,
|
||||
'R_Collar': 14, 'R_Shoulder': 17, 'R_Elbow': 19, 'R_Wrist': 21, 'R_Hand': 23
|
||||
}
|
||||
|
||||
MANO_RIGHT_JOINTS = {'R_Wrist':0, 'rindex0':1, 'rindex1':2, 'rindex2':3, 'rmiddle0':4, 'rmiddle1':5, 'rmiddle2':6,
|
||||
'rpinky0':7, 'rpinky1':8, 'rpinky2':9, 'rring0':10, 'rring1':11, 'rring2':12, 'rthumb0':13, 'rthumb1':14, 'rthumb2':15}
|
||||
|
||||
|
||||
FOOT_IDX = [SMPL_JOINTS['L_Ankle'], SMPL_JOINTS['R_Ankle'], SMPL_JOINTS['L_Foot'], SMPL_JOINTS['R_Foot']]
|
||||
|
||||
CONTACTS_IDX = [SMPL_JOINTS['L_Ankle'], SMPL_JOINTS['R_Ankle'],
|
||||
SMPL_JOINTS['L_Foot'], SMPL_JOINTS['R_Foot'],
|
||||
SMPL_JOINTS['L_Wrist'], SMPL_JOINTS['R_Wrist'],
|
||||
SMPL_JOINTS['L_Knee'], SMPL_JOINTS['R_Knee']]
|
||||
|
||||
|
||||
def align_joints(tensor, smpl_to_bvh=True):
|
||||
"""
|
||||
Args:
|
||||
tensor (T x J x D): The 3D torch tensor we need to process.
|
||||
|
||||
Returns:
|
||||
The 3D tensor whose joint order is compatible for bvh export.
|
||||
"""
|
||||
order = list(MANO_RIGHT_JOINTS.values())
|
||||
result = tensor.clone() if isinstance(tensor, torch.Tensor) else tensor.copy()
|
||||
for i in range(len(order)):
|
||||
if smpl_to_bvh:
|
||||
result[:, i] = tensor[:, order[i]]
|
||||
else:
|
||||
result[:, order[i]] = tensor[:, i]
|
||||
|
||||
return result
|
||||
|
||||
def ho3d_cam_extrinsics(cam_id):
|
||||
|
||||
# GPMF case
|
||||
extrinsic_0 = np.array([[-6.409640608401522277e-01, -2.650417602434713693e-01, 7.203595893984829912e-01, -5.675702013852904626e-01],
|
||||
[-7.601847604882365772e-01, 3.490763992765258950e-01, -5.479642300274587541e-01, 3.191066048126701693e-01],
|
||||
[-1.062271275746045074e-01, -8.988317600058360890e-01, -4.252261334538034454e-01, 3.074042501856224519e-01],
|
||||
[0, 0, 0, 1]])
|
||||
|
||||
extrinsic_1 = np.array([[-4.228051957326409149e-01, -2.078562115810418665e-01, -8.820609739517225600e-01, 4.637252480950287414e-01],
|
||||
[9.057441941962964815e-01, -6.536944939351106709e-02, -4.187532564239835331e-01, 1.535094236870509776e-01],
|
||||
[2.938062526878559150e-02, -9.759726586299259932e-01, 2.159030997129234852e-01, -1.448621926159473529e-02],
|
||||
[0, 0, 0, 1]])
|
||||
|
||||
extrinsic_2 = np.array([[9.846039213996782280e-01, 6.828612271512826681e-02, -1.609102961556882483e-01, 9.897043274267488394e-03],
|
||||
[1.747949086333730750e-01, -3.774480817251544273e-01, 9.093842342584639304e-01, -4.349111883283850455e-01],
|
||||
[1.363040801985337556e-03, -9.235095836246293155e-01, -3.835729279722913110e-01, 2.864983800867586528e-01],
|
||||
[0, 0, 0, 1]])
|
||||
|
||||
extrinsic_3 = np.array([[4.290309538248047727e-01, 8.850724776239779490e-02, 8.989432172021233347e-01, -4.754434780295317409e-01],
|
||||
[-9.018737295370956586e-01, 9.767471789067121157e-02, 4.208128152229064223e-01, -2.166588733650218201e-01],
|
||||
[-5.055904104142360661e-02, -9.912749954369395322e-01, 1.217278390117415843e-01, 6.127903398258782025e-02],
|
||||
[0, 0, 0, 1]])
|
||||
|
||||
extrinsic_4 = np.array([[-9.650817552456496529e-01, 5.435035576465601509e-02, 2.562484039369455346e-01, -9.846218864128720993e-02],
|
||||
[-2.604576318061128104e-01, -3.033325236743535935e-01, -9.165976227960112022e-01, 3.948193893157375123e-01],
|
||||
[2.791106816171780650e-02, -9.513334951054777111e-01, 3.068966493210200097e-01, -1.383363501223186694e-02],
|
||||
[0, 0, 0, 1]])
|
||||
|
||||
R, t = eval(f"extrinsic_{cam_id}")[:3, :3], eval(f"extrinsic_{cam_id}")[:3, 3]
|
||||
|
||||
return R, t
|
||||
|
||||
|
||||
def slerp(quat, trans, key_times, times, mask=True):
|
||||
"""
|
||||
Args:
|
||||
quat: (T x J x 4)
|
||||
trans: (T x 3)
|
||||
"""
|
||||
if mask:
|
||||
quat = c2c(quat[key_times])
|
||||
trans = c2c(trans[key_times])
|
||||
else:
|
||||
quat = c2c(quat)
|
||||
trans = c2c(trans)
|
||||
|
||||
quats = []
|
||||
for j in range(quat.shape[1]):
|
||||
key_rots = R.from_quat(quat[:, j])
|
||||
s = Slerp(key_times, key_rots)
|
||||
interp_rots = s(times)
|
||||
quats.append(interp_rots.as_quat())
|
||||
slerp_quat = np.stack(quats, axis=1)
|
||||
|
||||
lerp_trans = np.zeros((len(times), 3))
|
||||
for i in range(3):
|
||||
lerp_trans[:, i] = np.interp(times, key_times, trans[:, i])
|
||||
|
||||
return slerp_quat, lerp_trans
|
||||
|
||||
|
||||
def compute_orient_angle(matrix, traj):
|
||||
"""
|
||||
Args:
|
||||
matrix (N x T x 3 x 3): The 3D rotation matrix at the root joint
|
||||
traj (N x T x 3): The trajectory to align
|
||||
"""
|
||||
forward = matrix[:, :, :, 2].clone()
|
||||
forward[:, :, 2] = 0
|
||||
forward = F.normalize(forward, dim=-1) # normalized forward vector (N, T, 3)
|
||||
|
||||
traj[:, :, 2] = 0 # make sure the trajectory is projected to the plane
|
||||
|
||||
# first steps is forward diff
|
||||
init_tan = traj[:, 1:2] - traj[:, :1]
|
||||
# middle steps are second order
|
||||
middle_tan = (traj[:, 2:] - traj[:, 0:-2]) / 2
|
||||
# last step is backward diff
|
||||
final_tan = traj[:, -1:] - traj[:, -2:-1]
|
||||
|
||||
tangent = torch.cat([init_tan, middle_tan, final_tan], dim=1)
|
||||
tangent = F.normalize(tangent, dim=-1) # normalized tangent vector (N, T, 3)
|
||||
|
||||
cos = torch.sum(forward * tangent, dim=-1)
|
||||
|
||||
return cos
|
||||
|
||||
|
||||
def compute_trajectory(velocity, up, origin, dt, up_axis='z'):
|
||||
"""
|
||||
Args:
|
||||
velocity: (B, T, 3)
|
||||
up: (B, T)
|
||||
origin: (B, 3)
|
||||
up_axis: x, y, or z
|
||||
|
||||
Returns:
|
||||
trajectory: (B, T, 3)
|
||||
"""
|
||||
ordermap = {
|
||||
'x': 0,
|
||||
'y': 1,
|
||||
'z': 2,
|
||||
}
|
||||
v_axis = [x for x in ordermap.values() if x != ordermap[up_axis]]
|
||||
|
||||
origin = origin.unsqueeze(1) # (B, 3) => (B, 1, 3)
|
||||
trajectory = origin.repeat(1, up.shape[1], 1) # (B, 1, 3) => (B, T, 3)
|
||||
|
||||
for t in range(1, up.shape[1]):
|
||||
trajectory[:, t, v_axis[0]] = trajectory[:, t - 1, v_axis[0]] + velocity[:, t - 1, v_axis[0]] * dt
|
||||
trajectory[:, t, v_axis[1]] = trajectory[:, t - 1, v_axis[1]] + velocity[:, t - 1, v_axis[1]] * dt
|
||||
|
||||
trajectory[:, :, ordermap[up_axis]] = up
|
||||
|
||||
return trajectory
|
||||
|
||||
|
||||
def build_canonical_frame(forward, up_axis='z'):
|
||||
"""
|
||||
Args:
|
||||
forward: (..., 3)
|
||||
|
||||
Returns:
|
||||
frame: (..., 3, 3)
|
||||
"""
|
||||
forward[..., 'xyz'.index(up_axis)] = 0
|
||||
forward = F.normalize(forward, dim=-1) # normalized forward vector
|
||||
|
||||
up = torch.zeros_like(forward)
|
||||
up[..., 'xyz'.index(up_axis)] = 1 # normalized up vector
|
||||
right = torch.cross(up, forward)
|
||||
frame = torch.stack((right, up, forward), dim=-1) # canonical frame
|
||||
|
||||
return frame
|
||||
|
||||
|
||||
def estimate_linear_velocity(data_seq, dt):
|
||||
'''
|
||||
Given some batched data sequences of T timesteps in the shape (B, T, ...), estimates
|
||||
the velocity for the middle T-2 steps using a second order central difference scheme.
|
||||
The first and last frames are with forward and backward first-order
|
||||
differences, respectively
|
||||
- h : step size
|
||||
'''
|
||||
# first steps is forward diff (t+1 - t) / dt
|
||||
init_vel = (data_seq[:, 1:2] - data_seq[:, :1]) / dt
|
||||
# middle steps are second order (t+1 - t-1) / 2dt
|
||||
middle_vel = (data_seq[:, 2:] - data_seq[:, 0:-2]) / (2 * dt)
|
||||
# last step is backward diff (t - t-1) / dt
|
||||
final_vel = (data_seq[:, -1:] - data_seq[:, -2:-1]) / dt
|
||||
|
||||
vel_seq = torch.cat([init_vel, middle_vel, final_vel], dim=1)
|
||||
return vel_seq
|
||||
|
||||
|
||||
def estimate_angular_velocity(rot_seq, dt):
|
||||
'''
|
||||
Given a batch of sequences of T rotation matrices, estimates angular velocity at T-2 steps.
|
||||
Input sequence should be of shape (B, T, ..., 3, 3)
|
||||
'''
|
||||
# see https://en.wikipedia.org/wiki/Angular_velocity#Calculation_from_the_orientation_matrix
|
||||
dRdt = estimate_linear_velocity(rot_seq, dt)
|
||||
R = rot_seq
|
||||
RT = R.transpose(-1, -2)
|
||||
# compute skew-symmetric angular velocity tensor
|
||||
w_mat = torch.matmul(dRdt, RT)
|
||||
# pull out angular velocity vector by averaging symmetric entries
|
||||
w_x = (-w_mat[..., 1, 2] + w_mat[..., 2, 1]) / 2.0
|
||||
w_y = (w_mat[..., 0, 2] - w_mat[..., 2, 0]) / 2.0
|
||||
w_z = (-w_mat[..., 0, 1] + w_mat[..., 1, 0]) / 2.0
|
||||
w = torch.stack([w_x, w_y, w_z], axis=-1)
|
||||
return w
|
||||
|
||||
|
||||
def normalize(tensor, mean=None, std=None):
|
||||
"""
|
||||
Args:
|
||||
tensor: (B, T, ...)
|
||||
|
||||
Returns:
|
||||
normalized tensor with 0 mean and 1 standard deviation, std, mean
|
||||
"""
|
||||
if mean is None or std is None:
|
||||
# std, mean = torch.std_mean(tensor, dim=0, unbiased=False, keepdim=True)
|
||||
std, mean = torch.std_mean(tensor, dim=(0, 1), unbiased=False, keepdim=True)
|
||||
std[std == 0.0] = 1.0
|
||||
|
||||
return (tensor - mean) / std, mean, std
|
||||
|
||||
return (tensor - mean) / std
|
||||
|
||||
|
||||
def denormalize(tensor, mean, std):
|
||||
"""
|
||||
Args:
|
||||
tensor: B x T x D
|
||||
mean:
|
||||
std:
|
||||
"""
|
||||
return tensor * std + mean
|
||||
|
||||
|
||||
def export_bvh_animation(rotations, positions, offsets, parents, output_dir, prefix, joint_names, fps):
|
||||
"""
|
||||
Args:
|
||||
rotations: quaternions of the shape (B, T, J, 4)
|
||||
positions: global translations of the shape (B, T, 3)
|
||||
"""
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
for i in range(rotations.shape[0]):
|
||||
rotation = align_joints(rotations[i])
|
||||
position = positions[i]
|
||||
position = position.unsqueeze(1)
|
||||
anim = Animation(Quaternions(c2c(rotation)), c2c(position), None, offsets=offsets, parents=parents)
|
||||
BVH.save(os.path.join(output_dir, f'{prefix}_{i}.bvh'), anim, names=joint_names, frametime=1 / fps)
|
||||
|
||||
|
||||
def export_ply_trajectory(points, color, ply_fname):
|
||||
v = []
|
||||
for p in points:
|
||||
v += [(p[0], p[1], p[2], color[0], color[1], color[2])]
|
||||
v = np.array(v, dtype=[('x', 'f4'), ('y', 'f4'), ('z', 'f4'), ('red', 'u1'), ('green', 'u1'), ('blue', 'u1')])
|
||||
el_v = plyfile.PlyElement.describe(v, 'vertex')
|
||||
|
||||
e = np.empty(len(points) - 1, dtype=[('vertex1', 'i4'), ('vertex2', 'i4')])
|
||||
edge_data = np.array([[i, i + 1] for i in range(len(points) - 1)], dtype='i4')
|
||||
e['vertex1'] = edge_data[:, 0]
|
||||
e['vertex2'] = edge_data[:, 1]
|
||||
el_e = plyfile.PlyElement.describe(e, 'edge')
|
||||
|
||||
plyfile.PlyData([el_v, el_e], text=True).write(ply_fname)
|
||||
|
||||
|
||||
def get_keypoints_rectangle(keypoints: np.array, threshold: float) -> Tuple[float, float, float]:
|
||||
"""
|
||||
Compute rectangle enclosing keypoints above the threshold.
|
||||
Args:
|
||||
keypoints (np.array): Keypoint array of shape (N, 3).
|
||||
threshold (float): Confidence visualization threshold.
|
||||
Returns:
|
||||
Tuple[float, float, float]: Rectangle width, height and area.
|
||||
"""
|
||||
# print(keypoints.shape)
|
||||
valid_ind = keypoints[:, -1] > threshold
|
||||
if valid_ind.sum() > 0:
|
||||
valid_keypoints = keypoints[valid_ind][:, :-1]
|
||||
max_x = valid_keypoints[:,0].max()
|
||||
max_y = valid_keypoints[:,1].max()
|
||||
min_x = valid_keypoints[:,0].min()
|
||||
min_y = valid_keypoints[:,1].min()
|
||||
width = max_x - min_x
|
||||
height = max_y - min_y
|
||||
area = width * height
|
||||
return width, height, area
|
||||
else:
|
||||
return 0,0,0
|
||||
|
||||
def render_keypoints(img: np.array,
|
||||
keypoints: np.array,
|
||||
pairs: List,
|
||||
colors: List,
|
||||
thickness_circle_ratio: float,
|
||||
thickness_line_ratio_wrt_circle: float,
|
||||
pose_scales: List,
|
||||
threshold: float = 0.1,
|
||||
alpha: float = 1.0) -> np.array:
|
||||
"""
|
||||
Render keypoints on input image.
|
||||
Args:
|
||||
img (np.array): Input image of shape (H, W, 3) with pixel values in the [0,255] range.
|
||||
keypoints (np.array): Keypoint array of shape (N, 3).
|
||||
pairs (List): List of keypoint pairs per limb.
|
||||
colors: (List): List of colors per keypoint.
|
||||
thickness_circle_ratio (float): Circle thickness ratio.
|
||||
thickness_line_ratio_wrt_circle (float): Line thickness ratio wrt the circle.
|
||||
pose_scales (List): List of pose scales.
|
||||
threshold (float): Only visualize keypoints with confidence above the threshold.
|
||||
Returns:
|
||||
(np.array): Image of shape (H, W, 3) with keypoints drawn on top of the original image.
|
||||
"""
|
||||
img_orig = img.copy()
|
||||
width, height = img.shape[1], img.shape[2]
|
||||
area = width * height
|
||||
|
||||
lineType = 8
|
||||
shift = 0
|
||||
numberColors = len(colors)
|
||||
thresholdRectangle = 0.1
|
||||
|
||||
person_width, person_height, person_area = get_keypoints_rectangle(keypoints, thresholdRectangle)
|
||||
if person_area > 0:
|
||||
ratioAreas = min(1, max(person_width / width, person_height / height))
|
||||
thicknessRatio = np.maximum(np.round(math.sqrt(area) * thickness_circle_ratio * ratioAreas), 2)
|
||||
thicknessCircle = np.maximum(1, thicknessRatio if ratioAreas > 0.05 else -np.ones_like(thicknessRatio))
|
||||
thicknessLine = np.maximum(1, np.round(thicknessRatio * thickness_line_ratio_wrt_circle))
|
||||
radius = thicknessRatio / 2
|
||||
|
||||
img = np.ascontiguousarray(img.copy())
|
||||
for i, pair in enumerate(pairs):
|
||||
index1, index2 = pair
|
||||
if keypoints[index1, -1] > threshold and keypoints[index2, -1] > threshold:
|
||||
thicknessLineScaled = int(round(min(thicknessLine[index1], thicknessLine[index2]) * pose_scales[0]))
|
||||
colorIndex = index2
|
||||
color = colors[colorIndex % numberColors]
|
||||
keypoint1 = keypoints[index1, :-1].astype(np.int32)
|
||||
keypoint2 = keypoints[index2, :-1].astype(np.int32)
|
||||
cv2.line(img, tuple(keypoint1.tolist()), tuple(keypoint2.tolist()), tuple(color.tolist()), thicknessLineScaled, lineType, shift)
|
||||
for part in range(len(keypoints)):
|
||||
faceIndex = part
|
||||
if keypoints[faceIndex, -1] > threshold:
|
||||
radiusScaled = int(round(radius[faceIndex] * pose_scales[0]))
|
||||
thicknessCircleScaled = int(round(thicknessCircle[faceIndex] * pose_scales[0]))
|
||||
colorIndex = part
|
||||
color = colors[colorIndex % numberColors]
|
||||
center = keypoints[faceIndex, :-1].astype(np.int32)
|
||||
cv2.circle(img, tuple(center.tolist()), radiusScaled, tuple(color.tolist()), thicknessCircleScaled, lineType, shift)
|
||||
return img
|
||||
|
||||
def render_hand_keypoints(img, right_hand_keypoints, threshold=0.1, use_confidence=False, map_fn=lambda x: np.ones_like(x), alpha=1.0):
|
||||
if use_confidence and map_fn is not None:
|
||||
#thicknessCircleRatioLeft = 1./50 * map_fn(left_hand_keypoints[:, -1])
|
||||
thicknessCircleRatioRight = 1./50 * map_fn(right_hand_keypoints[:, -1])
|
||||
else:
|
||||
#thicknessCircleRatioLeft = 1./50 * np.ones(left_hand_keypoints.shape[0])
|
||||
thicknessCircleRatioRight = 1./50 * np.ones(right_hand_keypoints.shape[0])
|
||||
thicknessLineRatioWRTCircle = 0.75
|
||||
pairs = [0,1, 1,2, 2,3, 3,4, 0,5, 5,6, 6,7, 7,8, 0,9, 9,10, 10,11, 11,12, 0,13, 13,14, 14,15, 15,16, 0,17, 17,18, 18,19, 19,20]
|
||||
pairs = np.array(pairs).reshape(-1,2)
|
||||
|
||||
colors = [100., 100., 100.,
|
||||
100., 0., 0.,
|
||||
150., 0., 0.,
|
||||
200., 0., 0.,
|
||||
255., 0., 0.,
|
||||
100., 100., 0.,
|
||||
150., 150., 0.,
|
||||
200., 200., 0.,
|
||||
255., 255., 0.,
|
||||
0., 100., 50.,
|
||||
0., 150., 75.,
|
||||
0., 200., 100.,
|
||||
0., 255., 125.,
|
||||
0., 50., 100.,
|
||||
0., 75., 150.,
|
||||
0., 100., 200.,
|
||||
0., 125., 255.,
|
||||
100., 0., 100.,
|
||||
150., 0., 150.,
|
||||
200., 0., 200.,
|
||||
255., 0., 255.]
|
||||
colors = np.array(colors).reshape(-1,3)
|
||||
#colors = np.zeros_like(colors)
|
||||
poseScales = [1]
|
||||
#img = render_keypoints(img, left_hand_keypoints, pairs, colors, thicknessCircleRatioLeft, thicknessLineRatioWRTCircle, poseScales, threshold, alpha=alpha)
|
||||
img = render_keypoints(img, right_hand_keypoints, pairs, colors, thicknessCircleRatioRight, thicknessLineRatioWRTCircle, poseScales, threshold, alpha=alpha)
|
||||
#img = render_keypoints(img, right_hand_keypoints, pairs, colors, thickness_circle_ratio, thickness_line_ratio_wrt_circle, pose_scales, 0.1)
|
||||
return img
|
||||
|
||||
def render_body_keypoints(img: np.array,
|
||||
body_keypoints: np.array) -> np.array:
|
||||
"""
|
||||
Render OpenPose body keypoints on input image.
|
||||
Args:
|
||||
img (np.array): Input image of shape (H, W, 3) with pixel values in the [0,255] range.
|
||||
body_keypoints (np.array): Keypoint array of shape (N, 3); 3 <====> (x, y, confidence).
|
||||
Returns:
|
||||
(np.array): Image of shape (H, W, 3) with keypoints drawn on top of the original image.
|
||||
"""
|
||||
|
||||
thickness_circle_ratio = 1./75. * np.ones(body_keypoints.shape[0])
|
||||
thickness_line_ratio_wrt_circle = 0.75
|
||||
pairs = []
|
||||
pairs = [1,8,1,2,1,5,2,3,3,4,5,6,6,7,8,9,9,10,10,11,8,12,12,13,13,14,1,0,0,15,15,17,0,16,16,18,14,19,19,20,14,21,11,22,22,23,11,24]
|
||||
pairs = np.array(pairs).reshape(-1,2)
|
||||
colors = [255., 0., 85.,
|
||||
255., 0., 0.,
|
||||
255., 85., 0.,
|
||||
255., 170., 0.,
|
||||
255., 255., 0.,
|
||||
170., 255., 0.,
|
||||
85., 255., 0.,
|
||||
0., 255., 0.,
|
||||
255., 0., 0.,
|
||||
0., 255., 85.,
|
||||
0., 255., 170.,
|
||||
0., 255., 255.,
|
||||
0., 170., 255.,
|
||||
0., 85., 255.,
|
||||
0., 0., 255.,
|
||||
255., 0., 170.,
|
||||
170., 0., 255.,
|
||||
255., 0., 255.,
|
||||
85., 0., 255.,
|
||||
0., 0., 255.,
|
||||
0., 0., 255.,
|
||||
0., 0., 255.,
|
||||
0., 255., 255.,
|
||||
0., 255., 255.,
|
||||
0., 255., 255.]
|
||||
colors = np.array(colors).reshape(-1,3)
|
||||
pose_scales = [1]
|
||||
return render_keypoints(img, body_keypoints, pairs, colors, thickness_circle_ratio, thickness_line_ratio_wrt_circle, pose_scales, 0.1)
|
||||
|
||||
def render_openpose(img: np.array,
|
||||
hand_keypoints: np.array) -> np.array:
|
||||
"""
|
||||
Render keypoints in the OpenPose format on input image.
|
||||
Args:
|
||||
img (np.array): Input image of shape (H, W, 3) with pixel values in the [0,255] range.
|
||||
body_keypoints (np.array): Keypoint array of shape (N, 3); 3 <====> (x, y, confidence).
|
||||
Returns:
|
||||
(np.array): Image of shape (H, W, 3) with keypoints drawn on top of the original image.
|
||||
"""
|
||||
#img = render_body_keypoints(img, body_keypoints)
|
||||
img = render_hand_keypoints(img, hand_keypoints)
|
||||
return img
|
||||
+124
@@ -0,0 +1,124 @@
|
||||
"""Run the fixed 128-frame HMP on overlapping windows of a full sequence."""
|
||||
import copy
|
||||
import json
|
||||
from pathlib import Path
|
||||
import numpy as np
|
||||
import torch
|
||||
from scipy.spatial.transform import Rotation
|
||||
|
||||
|
||||
def window_starts(length, size=128, overlap=16):
|
||||
if length < 1 or not 0 <= overlap < size:
|
||||
raise ValueError('Invalid sequence length or overlap')
|
||||
starts = [0]
|
||||
while starts[-1] + size < length:
|
||||
starts.append(starts[-1] + size - overlap)
|
||||
return starts
|
||||
|
||||
|
||||
def slice_window(data, start, length, size=128):
|
||||
result = {}
|
||||
for key, value in data.items():
|
||||
if isinstance(value, (torch.Tensor, np.ndarray)) and value.ndim >= 2 and value.shape[1] == length:
|
||||
clip = value[:, start:min(start + size, length)].clone() if isinstance(value, torch.Tensor) else value[:, start:min(start + size, length)].copy()
|
||||
valid = clip.shape[1]
|
||||
if valid < size:
|
||||
if isinstance(clip, torch.Tensor):
|
||||
pad = clip[:, -1:].repeat_interleave(size-valid, dim=1)
|
||||
clip = torch.cat([clip, pad], dim=1)
|
||||
else:
|
||||
clip = np.concatenate([clip, np.repeat(clip[:, -1:], size-valid, axis=1)], axis=1)
|
||||
if key == 'vis_mask':
|
||||
clip[:, valid:] = 0
|
||||
elif key == 'joints2d':
|
||||
clip[:, valid:, :, -1] = 0
|
||||
result[key] = clip
|
||||
else:
|
||||
result[key] = copy.deepcopy(value)
|
||||
return result
|
||||
|
||||
|
||||
def blend_windows(original, clips, size=128, overlap=16):
|
||||
result = {k: np.array(v, copy=True) for k, v in original.items()}
|
||||
B, T = result['trans'].shape[:2]
|
||||
denom = np.zeros(T, dtype=np.float64)
|
||||
trans = np.zeros((B, T, 3), dtype=np.float64)
|
||||
rotation_keys = ['root_orient', 'pose_body']
|
||||
if all('decode_root' in clip for _, clip in clips):
|
||||
result['decode_root'] = np.zeros((B, T, 3), dtype=np.float32)
|
||||
rotation_keys.append('decode_root')
|
||||
accum, refs = {}, {}
|
||||
for key in rotation_keys:
|
||||
shape = result[key].shape
|
||||
refs[key] = Rotation.from_rotvec(result[key].reshape(-1, 3)).as_quat().reshape(*shape[:-1], 4)
|
||||
accum[key] = np.zeros_like(refs[key])
|
||||
for start, clip in clips:
|
||||
n = min(size, T-start)
|
||||
w = np.ones(n, dtype=np.float64)
|
||||
if start > 0:
|
||||
m = min(overlap, n)
|
||||
w[:m] *= np.arange(1, m+1)/(m+1)
|
||||
if start + size < T:
|
||||
m = min(overlap, n)
|
||||
w[-m:] *= np.arange(m, 0, -1)/(m+1)
|
||||
denom[start:start+n] += w
|
||||
trans[:, start:start+n] += clip['trans'][:, :n] * w[None, :, None]
|
||||
for key in rotation_keys:
|
||||
a = clip[key][:, :n].reshape(result[key][:, start:start+n].shape)
|
||||
q = Rotation.from_rotvec(a.reshape(-1, 3)).as_quat().reshape(*a.shape[:-1], 4)
|
||||
reference = refs[key][:, start:start+n]
|
||||
q *= np.where(np.sum(q*reference, axis=-1, keepdims=True) < 0, -1., 1.)
|
||||
weights = w.reshape((1,n) + (1,)*(q.ndim-2))
|
||||
accum[key][:, start:start+n] += q * weights
|
||||
if np.any(denom <= 0):
|
||||
raise ValueError('Uncovered frames in HMP windows')
|
||||
result['trans'] = (trans/denom[None,:,None]).astype(np.float32)
|
||||
for key, q in accum.items():
|
||||
if np.any(np.linalg.norm(q, axis=-1) < 1e-8):
|
||||
raise ValueError('Degenerate quaternion blend')
|
||||
result[key] = Rotation.from_quat(q.reshape(-1,4)).as_rotvec().reshape(result[key].shape).astype(np.float32)
|
||||
# HMP currently fixes shape. Reject unhandled per-window shape changes.
|
||||
for _, clip in clips:
|
||||
if not np.allclose(clip['betas'], result['betas'], atol=1e-5):
|
||||
raise ValueError('Windowed HMP requires fixed shape (opt_betas=false)')
|
||||
result['latent_pose'] = result['pose_body'].copy()
|
||||
return result
|
||||
|
||||
|
||||
def fit_windowed(obs_data, res_dicts, opt, data_args, out_dir, device, fit_fn):
|
||||
from body_model import MANO
|
||||
original = {k: v.detach().cpu().numpy() if isinstance(v, torch.Tensor) else np.asarray(v) for k,v in res_dicts[0].items()}
|
||||
T = original['trans'].shape[1]
|
||||
starts = window_starts(T)
|
||||
out = Path(out_dir)
|
||||
out.mkdir(parents=True, exist_ok=True)
|
||||
mano_cfg = {k.lower():v for k,v in dict(opt.MANO).items()}
|
||||
hand_model = MANO(batch_size=128, pose2rot=True, **mano_cfg).to(device)
|
||||
clips = []
|
||||
manifest = {'frames':T, 'window_size':128, 'overlap':16, 'windows':[]}
|
||||
for index, start in enumerate(starts):
|
||||
clip_dir = out / 'windows' / f'{start:06d}'
|
||||
path = clip_dir / f'{Path(opt.HMP.vid_path).stem}_000000_world_results.npz'
|
||||
print(f'HMP window {index+1}/{len(starts)}: frames {start}:{min(start+128,T)}', flush=True)
|
||||
if path.is_file():
|
||||
with np.load(path) as cached:
|
||||
compatible = np.allclose(cached['betas'], original['betas'], atol=1e-5)
|
||||
if not compatible:
|
||||
path.rename(path.with_suffix('.variable_shape_backup.npz'))
|
||||
if not path.is_file():
|
||||
obs = slice_window(obs_data, start, T)
|
||||
res = slice_window(res_dicts[0], start, T)
|
||||
fit_fn(obs, (res,), hand_model, opt, data_args, str(clip_dir), device)
|
||||
with np.load(path) as saved:
|
||||
clip = {k:saved[k].copy() for k in saved.files}
|
||||
if clip['trans'].shape[1] != 128 or not all(np.isfinite(v).all() for v in clip.values() if v.dtype.kind in 'fc'):
|
||||
raise ValueError(f'Invalid HMP window output: {path}')
|
||||
clips.append((start,clip))
|
||||
manifest['windows'].append({'start':start,'valid_frames':min(128,T-start),'path':str(path)})
|
||||
(out/'window_manifest.json').write_text(json.dumps(manifest,indent=2)+'\n')
|
||||
result = blend_windows(original,clips)
|
||||
target = out/f'{Path(opt.HMP.vid_path).stem}_000000_world_results.npz'
|
||||
temporary = target.with_suffix('.partial.npz')
|
||||
np.savez(temporary,**result)
|
||||
temporary.replace(target)
|
||||
print(f'Saved full-length windowed HMP: {target}', flush=True)
|
||||
@@ -0,0 +1,4 @@
|
||||
from .body_model import *
|
||||
from .specs import *
|
||||
from .utils import *
|
||||
from .mano_wrapper import *
|
||||
+2415
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,44 @@
|
||||
import torch
|
||||
import numpy as np
|
||||
import pickle
|
||||
from typing import Optional
|
||||
from body_model import MANOLayer
|
||||
from smplx.lbs import vertices2joints
|
||||
from smplx.utils import MANOOutput, to_tensor
|
||||
from smplx.vertex_ids import vertex_ids
|
||||
|
||||
|
||||
class MANO(MANOLayer):
|
||||
def __init__(self, *args, joint_regressor_extra: Optional[str] = None, **kwargs):
|
||||
"""
|
||||
Extension of the official MANO implementation to support more joints.
|
||||
Args:
|
||||
Same as MANOLayer.
|
||||
joint_regressor_extra (str): Path to extra joint regressor.
|
||||
"""
|
||||
super(MANO, self).__init__(*args, **kwargs)
|
||||
mano_to_openpose = [0, 13, 14, 15, 16, 1, 2, 3, 17, 4, 5, 6, 18, 10, 11, 12, 19, 7, 8, 9, 20]
|
||||
|
||||
#2, 3, 5, 4, 1
|
||||
if joint_regressor_extra is not None:
|
||||
self.register_buffer('joint_regressor_extra', torch.tensor(pickle.load(open(joint_regressor_extra, 'rb'), encoding='latin1'), dtype=torch.float32))
|
||||
self.register_buffer('extra_joints_idxs', to_tensor(list(vertex_ids['mano'].values()), dtype=torch.long))
|
||||
self.register_buffer('joint_map', torch.tensor(mano_to_openpose, dtype=torch.long))
|
||||
self.pose2rot = kwargs.get('pose2rot')
|
||||
# print([attr for attr in dir(self) if not callable(getattr(self, attr)) and not attr.startswith("__")])
|
||||
# print(self.hand_mean, self.flat_hand_mean, self.num_pca_comps, self.use_pca)
|
||||
# exit()
|
||||
|
||||
def forward(self, *args, **kwargs) -> MANOOutput:
|
||||
"""
|
||||
Run forward pass. Same as MANO and also append an extra set of joints if joint_regressor_extra is specified.
|
||||
"""
|
||||
mano_output = super(MANO, self).forward(*args, **kwargs)
|
||||
extra_joints = torch.index_select(mano_output.vertices, 1, self.extra_joints_idxs)
|
||||
joints = torch.cat([mano_output.joints, extra_joints], dim=1)
|
||||
joints = joints[:, self.joint_map, :]
|
||||
if hasattr(self, 'joint_regressor_extra'):
|
||||
extra_joints = vertices2joints(self.joint_regressor_extra, mano_output.vertices)
|
||||
joints = torch.cat([joints, extra_joints], dim=1)
|
||||
mano_output.joints = joints
|
||||
return mano_output
|
||||
+575
@@ -0,0 +1,575 @@
|
||||
import numpy as np
|
||||
|
||||
# SMPL_JOINTS = {
|
||||
# "hips": 0,
|
||||
# "leftUpLeg": 1,
|
||||
# "rightUpLeg": 2,
|
||||
# "spine": 3,
|
||||
# "leftLeg": 4,
|
||||
# "rightLeg": 5,
|
||||
# "spine1": 6,
|
||||
# "leftFoot": 7,
|
||||
# "rightFoot": 8,
|
||||
# "spine2": 9,
|
||||
# "leftToeBase": 10,
|
||||
# "rightToeBase": 11,
|
||||
# "neck": 12,
|
||||
# "leftShoulder": 13,
|
||||
# "rightShoulder": 14,
|
||||
# "head": 15,
|
||||
# "leftArm": 16,
|
||||
# "rightArm": 17,
|
||||
# "leftForeArm": 18,
|
||||
# "rightForeArm": 19,
|
||||
# "leftHand": 20,
|
||||
# "rightHand": 21,
|
||||
# }
|
||||
MANO_JOINTS = {
|
||||
'wrist': 0,
|
||||
'index1': 1,
|
||||
'index2': 2,
|
||||
'index3': 3,
|
||||
'middle1': 4,
|
||||
'middle2': 5,
|
||||
'middle3': 6,
|
||||
'pinky1': 7,
|
||||
'pinky2': 8,
|
||||
'pinky3': 9,
|
||||
'ring1': 10,
|
||||
'ring2': 11,
|
||||
'ring3': 12,
|
||||
'thumb1': 13,
|
||||
'thumb2': 14,
|
||||
'thumb3': 15,
|
||||
}
|
||||
OP_IGNORE_JOINTS = None
|
||||
OP_NUM_JOINTS = 21
|
||||
|
||||
# SMPL_PARENTS = [
|
||||
# -1,
|
||||
# 0,
|
||||
# 0,
|
||||
# 0,
|
||||
# 1,
|
||||
# 2,
|
||||
# 3,
|
||||
# 4,
|
||||
# 5,
|
||||
# 6,
|
||||
# 7,
|
||||
# 8,
|
||||
# 9,
|
||||
# 12,
|
||||
# 12,
|
||||
# 12,
|
||||
# 13,
|
||||
# 14,
|
||||
# 16,
|
||||
# 17,
|
||||
# 18,
|
||||
# 19,
|
||||
# ]
|
||||
|
||||
SMPLH_PATH = "./body_models/smplh"
|
||||
SMPLX_PATH = "./body_models/smplx"
|
||||
SMPL_PATH = "./body_models/smpl"
|
||||
VPOSER_PATH = "./body_models/vposer_v1_0"
|
||||
|
||||
# chosen virtual mocap markers that are "keypoints" to work with
|
||||
# KEYPT_VERTS = [
|
||||
# 4404,
|
||||
# 920,
|
||||
# 3076,
|
||||
# 3169,
|
||||
# 823,
|
||||
# 4310,
|
||||
# 1010,
|
||||
# 1085,
|
||||
# 4495,
|
||||
# 4569,
|
||||
# 6615,
|
||||
# 3217,
|
||||
# 3313,
|
||||
# 6713,
|
||||
# 6785,
|
||||
# 3383,
|
||||
# 6607,
|
||||
# 3207,
|
||||
# 1241,
|
||||
# 1508,
|
||||
# 4797,
|
||||
# 4122,
|
||||
# 1618,
|
||||
# 1569,
|
||||
# 5135,
|
||||
# 5040,
|
||||
# 5691,
|
||||
# 5636,
|
||||
# 5404,
|
||||
# 2230,
|
||||
# 2173,
|
||||
# 2108,
|
||||
# 134,
|
||||
# 3645,
|
||||
# 6543,
|
||||
# 3123,
|
||||
# 3024,
|
||||
# 4194,
|
||||
# 1306,
|
||||
# 182,
|
||||
# 3694,
|
||||
# 4294,
|
||||
# 744,
|
||||
# ]
|
||||
|
||||
|
||||
# """
|
||||
# Openpose
|
||||
# """
|
||||
# OP_NUM_JOINTS = 25
|
||||
# # OP_IGNORE_JOINTS = [1, 9, 12] # neck and left/right hip
|
||||
# OP_IGNORE_JOINTS = [1] # neck
|
||||
# OP_EDGE_LIST = [
|
||||
# [1, 8],
|
||||
# [1, 2],
|
||||
# [1, 5],
|
||||
# [2, 3],
|
||||
# [3, 4],
|
||||
# [5, 6],
|
||||
# [6, 7],
|
||||
# [8, 9],
|
||||
# [9, 10],
|
||||
# [10, 11],
|
||||
# [8, 12],
|
||||
# [12, 13],
|
||||
# [13, 14],
|
||||
# [1, 0],
|
||||
# [0, 15],
|
||||
# [15, 17],
|
||||
# [0, 16],
|
||||
# [16, 18],
|
||||
# [14, 19],
|
||||
# [19, 20],
|
||||
# [14, 21],
|
||||
# [11, 22],
|
||||
# [22, 23],
|
||||
# [11, 24],
|
||||
# ]
|
||||
# # indices to map an openpose detection to its flipped version
|
||||
# OP_FLIP_MAP = [
|
||||
# 0,
|
||||
# 1,
|
||||
# 5,
|
||||
# 6,
|
||||
# 7,
|
||||
# 2,
|
||||
# 3,
|
||||
# 4,
|
||||
# 8,
|
||||
# 12,
|
||||
# 13,
|
||||
# 14,
|
||||
# 9,
|
||||
# 10,
|
||||
# 11,
|
||||
# 16,
|
||||
# 15,
|
||||
# 18,
|
||||
# 17,
|
||||
# 22,
|
||||
# 23,
|
||||
# 24,
|
||||
# 19,
|
||||
# 20,
|
||||
# 21,
|
||||
# ]
|
||||
|
||||
|
||||
#
|
||||
# From https://github.com/vchoutas/smplify-x/blob/master/smplifyx/utils.py
|
||||
# Please see license for usage restrictions.
|
||||
#
|
||||
# def smpl_to_openpose(
|
||||
# model_type="smplx",
|
||||
# use_hands=True,
|
||||
# use_face=True,
|
||||
# use_face_contour=False,
|
||||
# openpose_format="coco25",
|
||||
# ):
|
||||
# """Returns the indices of the permutation that maps SMPL to OpenPose
|
||||
|
||||
# Parameters
|
||||
# ----------
|
||||
# model_type: str, optional
|
||||
# The type of SMPL-like model that is used. The default mapping
|
||||
# returned is for the SMPLX model
|
||||
# use_hands: bool, optional
|
||||
# Flag for adding to the returned permutation the mapping for the
|
||||
# hand keypoints. Defaults to True
|
||||
# use_face: bool, optional
|
||||
# Flag for adding to the returned permutation the mapping for the
|
||||
# face keypoints. Defaults to True
|
||||
# use_face_contour: bool, optional
|
||||
# Flag for appending the facial contour keypoints. Defaults to False
|
||||
# openpose_format: bool, optional
|
||||
# The output format of OpenPose. For now only COCO-25 and COCO-19 is
|
||||
# supported. Defaults to 'coco25'
|
||||
|
||||
# """
|
||||
# if openpose_format.lower() == "coco25":
|
||||
# if model_type == "smpl":
|
||||
# return np.array(
|
||||
# [
|
||||
# 24,
|
||||
# 12,
|
||||
# 17,
|
||||
# 19,
|
||||
# 21,
|
||||
# 16,
|
||||
# 18,
|
||||
# 20,
|
||||
# 0,
|
||||
# 2,
|
||||
# 5,
|
||||
# 8,
|
||||
# 1,
|
||||
# 4,
|
||||
# 7,
|
||||
# 25,
|
||||
# 26,
|
||||
# 27,
|
||||
# 28,
|
||||
# 29,
|
||||
# 30,
|
||||
# 31,
|
||||
# 32,
|
||||
# 33,
|
||||
# 34,
|
||||
# ],
|
||||
# dtype=np.int32,
|
||||
# )
|
||||
# elif model_type == "smplh":
|
||||
# body_mapping = np.array(
|
||||
# [
|
||||
# 52,
|
||||
# 12,
|
||||
# 17,
|
||||
# 19,
|
||||
# 21,
|
||||
# 16,
|
||||
# 18,
|
||||
# 20,
|
||||
# 0,
|
||||
# 2,
|
||||
# 5,
|
||||
# 8,
|
||||
# 1,
|
||||
# 4,
|
||||
# 7,
|
||||
# 53,
|
||||
# 54,
|
||||
# 55,
|
||||
# 56,
|
||||
# 57,
|
||||
# 58,
|
||||
# 59,
|
||||
# 60,
|
||||
# 61,
|
||||
# 62,
|
||||
# ],
|
||||
# dtype=np.int32,
|
||||
# )
|
||||
# mapping = [body_mapping]
|
||||
# if use_hands:
|
||||
# lhand_mapping = np.array(
|
||||
# [
|
||||
# 20,
|
||||
# 34,
|
||||
# 35,
|
||||
# 36,
|
||||
# 63,
|
||||
# 22,
|
||||
# 23,
|
||||
# 24,
|
||||
# 64,
|
||||
# 25,
|
||||
# 26,
|
||||
# 27,
|
||||
# 65,
|
||||
# 31,
|
||||
# 32,
|
||||
# 33,
|
||||
# 66,
|
||||
# 28,
|
||||
# 29,
|
||||
# 30,
|
||||
# 67,
|
||||
# ],
|
||||
# dtype=np.int32,
|
||||
# )
|
||||
# rhand_mapping = np.array(
|
||||
# [
|
||||
# 21,
|
||||
# 49,
|
||||
# 50,
|
||||
# 51,
|
||||
# 68,
|
||||
# 37,
|
||||
# 38,
|
||||
# 39,
|
||||
# 69,
|
||||
# 40,
|
||||
# 41,
|
||||
# 42,
|
||||
# 70,
|
||||
# 46,
|
||||
# 47,
|
||||
# 48,
|
||||
# 71,
|
||||
# 43,
|
||||
# 44,
|
||||
# 45,
|
||||
# 72,
|
||||
# ],
|
||||
# dtype=np.int32,
|
||||
# )
|
||||
# mapping += [lhand_mapping, rhand_mapping]
|
||||
# return np.concatenate(mapping)
|
||||
# # SMPLX
|
||||
# elif model_type == "smplx":
|
||||
# body_mapping = np.array(
|
||||
# [
|
||||
# 55,
|
||||
# 12,
|
||||
# 17,
|
||||
# 19,
|
||||
# 21,
|
||||
# 16,
|
||||
# 18,
|
||||
# 20,
|
||||
# 0,
|
||||
# 2,
|
||||
# 5,
|
||||
# 8,
|
||||
# 1,
|
||||
# 4,
|
||||
# 7,
|
||||
# 56,
|
||||
# 57,
|
||||
# 58,
|
||||
# 59,
|
||||
# 60,
|
||||
# 61,
|
||||
# 62,
|
||||
# 63,
|
||||
# 64,
|
||||
# 65,
|
||||
# ],
|
||||
# dtype=np.int32,
|
||||
# )
|
||||
# mapping = [body_mapping]
|
||||
# if use_hands:
|
||||
# lhand_mapping = np.array(
|
||||
# [
|
||||
# 20,
|
||||
# 37,
|
||||
# 38,
|
||||
# 39,
|
||||
# 66,
|
||||
# 25,
|
||||
# 26,
|
||||
# 27,
|
||||
# 67,
|
||||
# 28,
|
||||
# 29,
|
||||
# 30,
|
||||
# 68,
|
||||
# 34,
|
||||
# 35,
|
||||
# 36,
|
||||
# 69,
|
||||
# 31,
|
||||
# 32,
|
||||
# 33,
|
||||
# 70,
|
||||
# ],
|
||||
# dtype=np.int32,
|
||||
# )
|
||||
# rhand_mapping = np.array(
|
||||
# [
|
||||
# 21,
|
||||
# 52,
|
||||
# 53,
|
||||
# 54,
|
||||
# 71,
|
||||
# 40,
|
||||
# 41,
|
||||
# 42,
|
||||
# 72,
|
||||
# 43,
|
||||
# 44,
|
||||
# 45,
|
||||
# 73,
|
||||
# 49,
|
||||
# 50,
|
||||
# 51,
|
||||
# 74,
|
||||
# 46,
|
||||
# 47,
|
||||
# 48,
|
||||
# 75,
|
||||
# ],
|
||||
# dtype=np.int32,
|
||||
# )
|
||||
|
||||
# mapping += [lhand_mapping, rhand_mapping]
|
||||
# if use_face:
|
||||
# # end_idx = 127 + 17 * use_face_contour
|
||||
# face_mapping = np.arange(
|
||||
# 76, 127 + 17 * use_face_contour, dtype=np.int32
|
||||
# )
|
||||
# mapping += [face_mapping]
|
||||
|
||||
# return np.concatenate(mapping)
|
||||
# else:
|
||||
# raise ValueError("Unknown model type: {}".format(model_type))
|
||||
# elif openpose_format == "coco19":
|
||||
# if model_type == "smpl":
|
||||
# return np.array(
|
||||
# [24, 12, 17, 19, 21, 16, 18, 20, 0, 2, 5, 8, 1, 4, 7, 25, 26, 27, 28],
|
||||
# dtype=np.int32,
|
||||
# )
|
||||
# elif model_type == "smplh":
|
||||
# body_mapping = np.array(
|
||||
# [52, 12, 17, 19, 21, 16, 18, 20, 0, 2, 5, 8, 1, 4, 7, 53, 54, 55, 56],
|
||||
# dtype=np.int32,
|
||||
# )
|
||||
# mapping = [body_mapping]
|
||||
# if use_hands:
|
||||
# lhand_mapping = np.array(
|
||||
# [
|
||||
# 20,
|
||||
# 34,
|
||||
# 35,
|
||||
# 36,
|
||||
# 57,
|
||||
# 22,
|
||||
# 23,
|
||||
# 24,
|
||||
# 58,
|
||||
# 25,
|
||||
# 26,
|
||||
# 27,
|
||||
# 59,
|
||||
# 31,
|
||||
# 32,
|
||||
# 33,
|
||||
# 60,
|
||||
# 28,
|
||||
# 29,
|
||||
# 30,
|
||||
# 61,
|
||||
# ],
|
||||
# dtype=np.int32,
|
||||
# )
|
||||
# rhand_mapping = np.array(
|
||||
# [
|
||||
# 21,
|
||||
# 49,
|
||||
# 50,
|
||||
# 51,
|
||||
# 62,
|
||||
# 37,
|
||||
# 38,
|
||||
# 39,
|
||||
# 63,
|
||||
# 40,
|
||||
# 41,
|
||||
# 42,
|
||||
# 64,
|
||||
# 46,
|
||||
# 47,
|
||||
# 48,
|
||||
# 65,
|
||||
# 43,
|
||||
# 44,
|
||||
# 45,
|
||||
# 66,
|
||||
# ],
|
||||
# dtype=np.int32,
|
||||
# )
|
||||
# mapping += [lhand_mapping, rhand_mapping]
|
||||
# return np.concatenate(mapping)
|
||||
# # SMPLX
|
||||
# elif model_type == "smplx":
|
||||
# body_mapping = np.array(
|
||||
# [55, 12, 17, 19, 21, 16, 18, 20, 0, 2, 5, 8, 1, 4, 7, 56, 57, 58, 59],
|
||||
# dtype=np.int32,
|
||||
# )
|
||||
# mapping = [body_mapping]
|
||||
# if use_hands:
|
||||
# lhand_mapping = np.array(
|
||||
# [
|
||||
# 20,
|
||||
# 37,
|
||||
# 38,
|
||||
# 39,
|
||||
# 60,
|
||||
# 25,
|
||||
# 26,
|
||||
# 27,
|
||||
# 61,
|
||||
# 28,
|
||||
# 29,
|
||||
# 30,
|
||||
# 62,
|
||||
# 34,
|
||||
# 35,
|
||||
# 36,
|
||||
# 63,
|
||||
# 31,
|
||||
# 32,
|
||||
# 33,
|
||||
# 64,
|
||||
# ],
|
||||
# dtype=np.int32,
|
||||
# )
|
||||
# rhand_mapping = np.array(
|
||||
# [
|
||||
# 21,
|
||||
# 52,
|
||||
# 53,
|
||||
# 54,
|
||||
# 65,
|
||||
# 40,
|
||||
# 41,
|
||||
# 42,
|
||||
# 66,
|
||||
# 43,
|
||||
# 44,
|
||||
# 45,
|
||||
# 67,
|
||||
# 49,
|
||||
# 50,
|
||||
# 51,
|
||||
# 68,
|
||||
# 46,
|
||||
# 47,
|
||||
# 48,
|
||||
# 69,
|
||||
# ],
|
||||
# dtype=np.int32,
|
||||
# )
|
||||
|
||||
# mapping += [lhand_mapping, rhand_mapping]
|
||||
# if use_face:
|
||||
# face_mapping = np.arange(
|
||||
# 70, 70 + 51 + 17 * use_face_contour, dtype=np.int32
|
||||
# )
|
||||
# mapping += [face_mapping]
|
||||
|
||||
# return np.concatenate(mapping)
|
||||
# else:
|
||||
# raise ValueError("Unknown model type: {}".format(model_type))
|
||||
# else:
|
||||
# raise ValueError("Unknown joint format: {}".format(openpose_format))
|
||||
@@ -0,0 +1,76 @@
|
||||
import torch
|
||||
from .specs import MANO_JOINTS
|
||||
|
||||
|
||||
def run_mano(body_model, trans, root_orient, body_pose, is_right, betas=None):
|
||||
"""
|
||||
Forward pass of the MANO model and populates pred_data accordingly with
|
||||
joints3d, verts3d, points3d.
|
||||
|
||||
trans : B x T x 3
|
||||
root_orient : B x T x 3
|
||||
body_pose : B x T x J*3
|
||||
betas : (optional) B x D
|
||||
"""
|
||||
B, T, _ = trans.shape
|
||||
bm_batch_size = body_model.batch_size
|
||||
assert bm_batch_size % B == 0
|
||||
# assert (is_right[0]==is_right[0][0]).all(), f'{is_right}'
|
||||
# assert (is_right[1]==is_right[1][0]).all(), f'{is_right}'
|
||||
|
||||
seq_len = bm_batch_size // B
|
||||
bm_num_betas = body_model.num_betas
|
||||
J_BODY = len(MANO_JOINTS) - 1 # all joints except root
|
||||
# print('utils.py, seq_len, T: ', seq_len, T, bm_batch_size, B) # 1, 120, 1, 1
|
||||
# print('trans, root_orient, body_pose:', trans.shape, root_orient.shape, body_pose.shape)
|
||||
if T == 1:
|
||||
raise ValueError
|
||||
# must expand to use with body model
|
||||
trans = trans.expand(B, seq_len, 3)
|
||||
root_orient = root_orient.expand(B, seq_len, 3)
|
||||
body_pose = body_pose.expand(B, seq_len, J_BODY * 3)
|
||||
elif T != seq_len:
|
||||
trans, root_orient, body_pose = zero_pad_tensors(
|
||||
[trans, root_orient, body_pose], seq_len - T
|
||||
)
|
||||
if betas is None:
|
||||
betas = torch.zeros(B, bm_num_betas, device=trans.device)
|
||||
betas = betas.reshape((B, 1, bm_num_betas)).expand((B, seq_len, bm_num_betas))
|
||||
# print('body_pose: ', body_pose.reshape((B * seq_len, -1)).shape)
|
||||
|
||||
mano_output = body_model(
|
||||
hand_pose=body_pose.reshape((B * seq_len, -1)),
|
||||
betas=betas.reshape((B * seq_len, -1)),
|
||||
global_orient=root_orient.reshape((B * seq_len, -1)),
|
||||
transl=trans.reshape((B * seq_len, -1)),
|
||||
)
|
||||
joints = mano_output.joints
|
||||
verts = mano_output.vertices
|
||||
|
||||
joints = joints.reshape(B, seq_len, -1, 3)[:, :T]
|
||||
verts = verts.reshape(B, seq_len, -1, 3)[:, :T]
|
||||
|
||||
is_right = is_right.unsqueeze(-1)
|
||||
joints[:, :, :, 0] = (2*is_right-1)*joints[:, :, :, 0]
|
||||
verts[:, :, :, 0] = (2*is_right-1)*verts[:, :, :, 0]
|
||||
|
||||
return {
|
||||
"joints": joints,
|
||||
"vertices": verts,
|
||||
"l_faces": body_model.faces_tensor[:,[0,2,1]],
|
||||
"r_faces": body_model.faces_tensor,
|
||||
"is_right": is_right.squeeze(-1),
|
||||
'body_pose': body_pose.clone()
|
||||
}
|
||||
|
||||
|
||||
def zero_pad_tensors(pad_list, pad_size):
|
||||
"""
|
||||
Assumes tensors in pad_list are B x T x D and pad temporal dimension
|
||||
"""
|
||||
B = pad_list[0].size(0)
|
||||
new_pad_list = []
|
||||
for pad_idx, pad_tensor in enumerate(pad_list):
|
||||
padding = torch.zeros((B, pad_size, pad_tensor.size(2))).to(pad_tensor)
|
||||
new_pad_list.append(torch.cat([pad_tensor, padding], dim=1))
|
||||
return new_pad_list
|
||||
+64
@@ -0,0 +1,64 @@
|
||||
defaults:
|
||||
- data: HOT3D
|
||||
- optim
|
||||
- _self_
|
||||
|
||||
model:
|
||||
use_init: True
|
||||
opt_cams: False #True
|
||||
opt_scale: True
|
||||
async_tracks: True
|
||||
|
||||
overwrite: False
|
||||
run_opt: True
|
||||
run_vis: True
|
||||
run_prior: False
|
||||
datatype: ${data.video_dir}
|
||||
vis:
|
||||
phases:
|
||||
# - motion_chunks
|
||||
# - input
|
||||
# - root_fit
|
||||
- smooth_fit
|
||||
- prior
|
||||
# - init
|
||||
render_views:
|
||||
- above
|
||||
- side
|
||||
- front
|
||||
- src_cam
|
||||
make_grid: True
|
||||
overwrite: True
|
||||
render_keypoints: False # Overlay GT 2D keypoints on source camera view
|
||||
|
||||
MANO:
|
||||
MODEL_PATH: ${paths.DATA_DIR}/mano
|
||||
GENDER: neutral
|
||||
NUM_HAND_JOINTS: 15
|
||||
MEAN_PARAMS: ${paths.DATA_DIR}/mano_mean_params.npz
|
||||
CREATE_BODY_POSE: FALSE
|
||||
|
||||
paths:
|
||||
base_dir: None
|
||||
DATA_DIR: _DATA/data/
|
||||
vposer: VPoser/pretrained/Vposer_right_mirrored # _DATA/body_models/vposer_v1_0
|
||||
|
||||
HMP:
|
||||
config: hmp_config.yaml
|
||||
use_hposer: False
|
||||
vid_path: ${data.root}/video/${data.seq}.mp4
|
||||
exp_name: None
|
||||
|
||||
fps: 30
|
||||
log_root: ../outputs/logs
|
||||
log_dir: ${log_root}/${data.type}-${data.split}
|
||||
exp_name: ${now:%Y-%m-%d}
|
||||
temporal_smooth: True
|
||||
gpu: 0
|
||||
is_static: False
|
||||
|
||||
hydra:
|
||||
job:
|
||||
chdir: True
|
||||
run:
|
||||
dir: ${log_dir}/${exp_name}/${data.name}
|
||||
@@ -0,0 +1,24 @@
|
||||
type: video
|
||||
split: custom
|
||||
root: /data/home/zy3023/code/hand/Dyn-HaMR/test # put your videos in root/videos/vid.mp4
|
||||
video_dir: videos
|
||||
seq: demo1
|
||||
ext: mp4
|
||||
src_path: ${data.root}/${data.video_dir}/${data.seq}.${data.ext}
|
||||
frame_opts:
|
||||
ext: jpg
|
||||
fps: 25
|
||||
start_sec: 0
|
||||
end_sec: -1
|
||||
use_cams: True
|
||||
track_ids: "all"
|
||||
shot_idx: 0
|
||||
start_idx: 0
|
||||
end_idx: 500
|
||||
split_cameras: True
|
||||
name: ${data.seq}-${data.track_ids}-shot-${data.shot_idx}-${data.start_idx}-${data.end_idx}
|
||||
sources:
|
||||
images: ${data.root}/images/${data.seq}
|
||||
cameras: ${data.root}/dynhamr/cameras/${data.seq}/shot-${data.shot_idx}
|
||||
tracks: ${data.root}/dynhamr/track_preds/${data.seq}
|
||||
shots: ${data.root}/dynhamr/shot_idcs/${data.seq}.json
|
||||
@@ -0,0 +1,29 @@
|
||||
type: video
|
||||
split: custom
|
||||
root: /data/home/zy3023/code/hand/Dyn-HaMR/test # put your videos in root/videos/vid.mp4
|
||||
video_dir: videos
|
||||
seq: prod1
|
||||
ext: mp4
|
||||
src_path: ${data.root}/${data.video_dir}/${data.seq}.${data.ext}
|
||||
frame_opts:
|
||||
ext: jpg
|
||||
fps: 25
|
||||
start_sec: 0
|
||||
end_sec: -1
|
||||
use_cams: True
|
||||
track_ids: "all"
|
||||
shot_idx: 0
|
||||
start_idx: 0
|
||||
end_idx: -1
|
||||
split_cameras: True
|
||||
name: ${data.seq}-${data.track_ids}-shot-${data.shot_idx}-${data.start_idx}-${data.end_idx}
|
||||
sources:
|
||||
images: ${data.root}/images/${data.seq}
|
||||
cameras: ${data.root}/dynhamr/cameras/${data.seq}/shot-${data.shot_idx}
|
||||
tracks: ${data.root}/dynhamr/track_preds/${data.seq}
|
||||
shots: ${data.root}/dynhamr/shot_idcs/${data.seq}.json
|
||||
|
||||
# VIPE-specific options
|
||||
use_vipe: True
|
||||
vipe_dir: /data/home/zy3023/code/hand/Dyn-HaMR/third-party/vipe/vipe_results
|
||||
|
||||
+13
@@ -0,0 +1,13 @@
|
||||
defaults:
|
||||
- data: posetrack
|
||||
- _self_
|
||||
|
||||
gap: 1
|
||||
log_root: outputs
|
||||
save_per_frame: False
|
||||
print_err: False
|
||||
stride: 48
|
||||
|
||||
hydra:
|
||||
run:
|
||||
dir: ${log_root}/init/${data.type}-${data.seq}-${data.depth_dir}-${data.fov}fov-gap${gap}
|
||||
+52
@@ -0,0 +1,52 @@
|
||||
optim:
|
||||
options:
|
||||
robust_loss_type: "bisquare"
|
||||
robust_tuning_const: 4.6851
|
||||
joints2d_sigma: 100.0
|
||||
lr: 1.0
|
||||
lbfgs_max_iter: 20
|
||||
save_every: 20
|
||||
vis_every: -1
|
||||
max_chunk_steps: 20
|
||||
save_meshes: False
|
||||
|
||||
root:
|
||||
num_iters: 50
|
||||
|
||||
smpl:
|
||||
num_iters: 0
|
||||
|
||||
smooth:
|
||||
opt_scale: False
|
||||
num_iters: 300
|
||||
|
||||
motion_chunks:
|
||||
chunk_size: 10
|
||||
init_steps: 20
|
||||
chunk_steps: 20
|
||||
opt_cams: True
|
||||
|
||||
loss_weights:
|
||||
penetration: [0,0,0] # [0, 1.0, 10.0]
|
||||
# joints2d: [1, 1, 1] # Increased smooth_fit weight from 0.001 to 0.1 to better match 2D keypoints
|
||||
joints2d: [10000, 10000, 10000] # Increased smooth_fit weight from 0.001 to 0.1 to better match 2D keypoints
|
||||
bg2d: [0.0, 0.000, 0.000]
|
||||
cam_R_smooth : [0.0, 0.0, 0.0]
|
||||
cam_t_smooth : [0.0, 0.0, 0.0]
|
||||
bio : [0.0, 0.0, 0.0] # [10, 10, 10]
|
||||
depth_constraint: [100.0, 100.0, 0.0] # Prevent hand from being behind camera - MUST be stronger than 2D loss
|
||||
# bg2d: [0.0, 0.0001, 0.0001]
|
||||
# cam_R_smooth : [0.0, 1000.0, 1000.0]
|
||||
# cam_t_smooth : [0.0, 1000.0, 1000.0]
|
||||
joints3d: [0.0, 0.0, 0.0]
|
||||
joints3d_smooth: [1000, 10000.0, 0.0] # raw: [1.0, 10.0, 0.0]
|
||||
# joints3d_smooth: [0, 0, 0.0] # raw: [1.0, 10.0, 0.0]
|
||||
joints3d_rollout: [0.0, 0.0, 0.0]
|
||||
verts3d: [0.0, 0.0, 0.0]
|
||||
points3d: [0.0, 0.0, 0.0]
|
||||
pose_prior: [1, 1, 1]
|
||||
shape_prior: [0.05, 0.05, 0.05]
|
||||
motion_prior: [0.0, 0.0, 0.075]
|
||||
init_motion_prior: [0.0, 0.0, 0.075]
|
||||
joint_consistency: [0.0, 0.0, 100.0]
|
||||
bone_length: [0.0, 2000.0, 2000.0]
|
||||
@@ -0,0 +1,2 @@
|
||||
from .dataset import *
|
||||
from . import tools
|
||||
+482
@@ -0,0 +1,482 @@
|
||||
import os
|
||||
import glob
|
||||
import typing
|
||||
|
||||
import imageio
|
||||
import numpy as np
|
||||
import json
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch.utils.data import Dataset
|
||||
|
||||
from body_model import MANO_JOINTS
|
||||
assert len(MANO_JOINTS) == 16
|
||||
import time
|
||||
|
||||
from util.logger import Logger
|
||||
from geometry.camera import invert_camera
|
||||
|
||||
from .tools import read_keypoints, read_mask_path, load_mano_preds, load_keypoints_with_interp
|
||||
from .vidproc import preprocess_cameras, preprocess_frames, preprocess_tracks
|
||||
|
||||
|
||||
"""
|
||||
Define data-related constants
|
||||
"""
|
||||
# DEFAULT_GROUND = np.array([0.0, -1.0, 0.0, -0.5])
|
||||
|
||||
# XXX: TEMPORARY CONSTANTS
|
||||
SHOT_PAD = 0
|
||||
MIN_SEQ_LEN = 10
|
||||
MAX_NUM_TRACKS = 12
|
||||
MIN_TRACK_LEN = 60
|
||||
MIN_KEYP_CONF = 0.4
|
||||
|
||||
|
||||
def get_dataset_from_cfg(cfg):
|
||||
args = cfg.data
|
||||
if not args.use_cams:
|
||||
args.sources.cameras = ""
|
||||
|
||||
args.sources = expand_source_paths(args.sources)
|
||||
print("DATA SOURCES", args.sources)
|
||||
a = time.time()
|
||||
check_data_sources(args, cfg)
|
||||
print('end of check_data_sources, time: ', time.time() - a)
|
||||
|
||||
return MultiPeopleDataset(
|
||||
args.sources,
|
||||
args.seq,
|
||||
tid_spec=args.track_ids,
|
||||
shot_idx=args.shot_idx,
|
||||
start_idx=int(args.start_idx),
|
||||
end_idx=int(args.end_idx),
|
||||
is_static=cfg.is_static,
|
||||
split_cameras=args.get("split_cameras", True),
|
||||
)
|
||||
|
||||
|
||||
def expand_source_paths(data_sources):
|
||||
return {k: get_data_source(v) for k, v in data_sources.items()}
|
||||
|
||||
|
||||
def get_data_source(source):
|
||||
matches = glob.glob(source)
|
||||
if len(matches) < 1:
|
||||
print(f"{source} does not exist")
|
||||
return source # return anyway for default values
|
||||
if len(matches) > 1:
|
||||
raise ValueError(f"{source} is not unique")
|
||||
return matches[0]
|
||||
|
||||
|
||||
def check_data_sources(args, cfg):
|
||||
print()
|
||||
a = time.time()
|
||||
if args.type == "video":
|
||||
preprocess_frames(args.sources.images, args.src_path, **args.frame_opts)
|
||||
b = time.time()
|
||||
preprocess_tracks(cfg.datatype, args.sources.images, args.sources.tracks, args.sources.shots, gpu=cfg.gpu)
|
||||
c = time.time()
|
||||
preprocess_cameras(args, overwrite=args.get("overwrite_cams", False))
|
||||
d = time.time()
|
||||
print('frame, hand track, slam camera: ', b-a, c-b, d-c)
|
||||
|
||||
|
||||
class MultiPeopleDataset(Dataset):
|
||||
def __init__(
|
||||
self,
|
||||
data_sources: typing.Dict,
|
||||
seq_name,
|
||||
tid_spec="all",
|
||||
shot_idx=0,
|
||||
start_idx=0,
|
||||
end_idx=-1,
|
||||
is_static=False,
|
||||
pad_shot=False,
|
||||
split_cameras=True,
|
||||
):
|
||||
self.is_static = is_static
|
||||
self.seq_name = seq_name
|
||||
self.data_sources = data_sources
|
||||
self.split_cameras = split_cameras
|
||||
|
||||
# select only images in the desired shot
|
||||
img_files, _ = get_shot_img_files(
|
||||
self.data_sources["shots"], shot_idx, pad_shot
|
||||
)
|
||||
end_idx = end_idx if end_idx > 0 else len(img_files)
|
||||
self.data_start, self.data_end = start_idx, end_idx
|
||||
img_files = img_files[start_idx:end_idx]
|
||||
self.img_names = [get_name(f) for f in img_files]
|
||||
self.num_imgs = len(self.img_names)
|
||||
|
||||
img_dir = self.data_sources["images"]
|
||||
assert os.path.isdir(img_dir)
|
||||
self.img_paths = [os.path.join(img_dir, f) for f in img_files]
|
||||
img_h, img_w = imageio.imread(self.img_paths[0]).shape[:2]
|
||||
self.img_size = img_w, img_h
|
||||
print(f"USING TOTAL {self.num_imgs} {img_w}x{img_h} IMGS")
|
||||
|
||||
# find the tracks in the video
|
||||
track_root = self.data_sources["tracks"]
|
||||
if tid_spec == "all" or tid_spec.startswith("longest"):
|
||||
n_tracks = MAX_NUM_TRACKS
|
||||
if tid_spec.startswith("longest"):
|
||||
n_tracks = int(tid_spec.split("-")[1])
|
||||
# get the longest tracks in the selected shot
|
||||
track_ids = sorted(os.listdir(track_root))
|
||||
track_paths = [
|
||||
[f"{track_root}/{tid}/{name}_keypoints.json" for name in self.img_names]
|
||||
for tid in track_ids
|
||||
]
|
||||
track_lens = [
|
||||
len(list(filter(os.path.isfile, paths))) for paths in track_paths
|
||||
]
|
||||
print("raw TRACK IDS and LENGTHS", track_ids, track_lens)
|
||||
track_ids = [
|
||||
track_ids[i]
|
||||
for i in range(len(track_lens))# np.argsort(track_lens)[::-1]
|
||||
if track_lens[i] > MIN_TRACK_LEN
|
||||
]
|
||||
print("TRACK LENGTHS", track_ids, track_lens)
|
||||
track_ids = track_ids[:n_tracks]
|
||||
else:
|
||||
track_ids = [f"{int(tid):03d}" for tid in tid_spec.split("-")]
|
||||
|
||||
print("TRACK IDS", track_ids)
|
||||
|
||||
self.track_ids = track_ids
|
||||
self.n_tracks = len(track_ids)
|
||||
self.track_dirs = [os.path.join(track_root, tid) for tid in track_ids]
|
||||
|
||||
# keep a list of frame index masks of whether a track is available in a frame
|
||||
sidx = np.inf
|
||||
eidx = -1
|
||||
self.track_vis_masks = []
|
||||
for pred_dir in self.track_dirs:
|
||||
kp_paths = [f"{pred_dir}/{x}_keypoints.json" for x in self.img_names]
|
||||
has_kp = [os.path.isfile(x) for x in kp_paths]
|
||||
|
||||
# keep track of which frames this track is visible in
|
||||
vis_mask = np.array(has_kp)
|
||||
idcs = np.where(vis_mask)[0]
|
||||
if len(idcs) > 0:
|
||||
si, ei = min(idcs), max(idcs)
|
||||
sidx = min(sidx, si)
|
||||
eidx = max(eidx, ei)
|
||||
self.track_vis_masks.append(vis_mask)
|
||||
|
||||
eidx = max(eidx + 1, 0)
|
||||
sidx = min(sidx, eidx)
|
||||
print("START", sidx, "END", eidx)
|
||||
self.start_idx = sidx
|
||||
self.end_idx = eidx
|
||||
self.seq_len = eidx - sidx
|
||||
self.seq_intervals = [(sidx, eidx) for _ in track_ids]
|
||||
|
||||
self.sel_img_paths = self.img_paths[sidx:eidx]
|
||||
self.sel_img_names = self.img_names[sidx:eidx]
|
||||
|
||||
# used to cache data
|
||||
self.data_dict = {}
|
||||
self.cam_data = None
|
||||
|
||||
def __len__(self):
|
||||
return self.n_tracks
|
||||
|
||||
def load_data(self, interp_input=True):
|
||||
if len(self.data_dict) > 0:
|
||||
return
|
||||
|
||||
# load camera data
|
||||
self.load_camera_data()
|
||||
# get data for each track
|
||||
data_out = {
|
||||
"mask_paths": [],
|
||||
"floor_plane": [],
|
||||
"joints2d": [],
|
||||
"vis_mask": [],
|
||||
"track_interval": [],
|
||||
"init_body_pose": [],
|
||||
"init_body_shape": [],
|
||||
"init_root_orient": [],
|
||||
"init_trans": [],
|
||||
"is_right": []
|
||||
}
|
||||
|
||||
# create batches of sequences
|
||||
# each batch is a track for a person
|
||||
T = self.seq_len
|
||||
sidx, eidx = self.start_idx, self.end_idx
|
||||
for i, tid in enumerate(self.track_ids):
|
||||
if len(self.track_ids) == 2:
|
||||
assert int(i) == int(tid), f'{i}/{tid}/{self.track_ids}'
|
||||
# load mask of visible frames for this track
|
||||
vis_mask = self.track_vis_masks[i][sidx:eidx] # (T)
|
||||
vis_idcs = np.where(vis_mask)[0]
|
||||
track_s, track_e = min(vis_idcs), max(vis_idcs) + 1
|
||||
data_out["track_interval"].append([track_s, track_e])
|
||||
|
||||
vis_mask = get_ternary_mask(vis_mask)
|
||||
data_out["vis_mask"].append(vis_mask)
|
||||
|
||||
# load 2d keypoints for visible frames with interpolation
|
||||
kp_paths = [
|
||||
f"{self.track_dirs[i]}/{x}_keypoints.json" for x in self.sel_img_names
|
||||
]
|
||||
# (T, J, 3) (x, y, conf) - with interpolation for missing frames
|
||||
joints2d_data = load_keypoints_with_interp(kp_paths, interp=interp_input)
|
||||
# print('joints2d_data: ', joints2d_data.shape)
|
||||
# print('joints2d_data[:, :, [2]]: ', np.unique(joints2d_data[:, :, [2]]))
|
||||
# raise ValueError
|
||||
# Discard bad ViTPose detections
|
||||
joints2d_data[:, :, 2] = 1.0
|
||||
joints2d_data[
|
||||
np.repeat(joints2d_data[:, :, [2]] < MIN_KEYP_CONF, 3, axis=2)
|
||||
] = 0
|
||||
# Set all confidence values to 1.0 for optimization
|
||||
data_out["joints2d"].append(joints2d_data)
|
||||
|
||||
# load single image mano predictions
|
||||
pred_paths = [
|
||||
f"{self.track_dirs[i]}/{x}_mano.json" for x in self.sel_img_names
|
||||
]
|
||||
pose_init, orient_init, trans_init, betas_init, is_right = load_mano_preds(
|
||||
pred_paths, tid=tid, interp=interp_input
|
||||
)
|
||||
|
||||
n_joints = len(MANO_JOINTS) - 1 # 15
|
||||
data_out["init_body_pose"].append(pose_init)
|
||||
data_out["init_body_shape"].append(betas_init)
|
||||
data_out["init_root_orient"].append(orient_init)
|
||||
data_out["init_trans"].append(trans_init)
|
||||
data_out['is_right'].append(is_right)
|
||||
|
||||
# DEBUG: Print hand type for each track
|
||||
print(f"DEBUG: Track {i} (tid={tid}): is_right={is_right[0] if len(is_right) > 0 else 'empty'}, {len(pred_paths)} frames")
|
||||
|
||||
# data_out["floor_plane"].append(DEFAULT_GROUND[:3] * DEFAULT_GROUND[3:])
|
||||
|
||||
self.data_dict = data_out
|
||||
|
||||
def __getitem__(self, idx):
|
||||
if len(self.data_dict) < 1:
|
||||
self.load_data()
|
||||
|
||||
obs_data = dict()
|
||||
|
||||
# 2D keypoints
|
||||
joint2d_data = self.data_dict["joints2d"][idx]
|
||||
obs_data["joints2d"] = torch.Tensor(joint2d_data)
|
||||
|
||||
# single frame predictions
|
||||
obs_data["init_body_pose"] = torch.Tensor(self.data_dict["init_body_pose"][idx])
|
||||
obs_data["init_body_shape"] = torch.Tensor(self.data_dict["init_body_shape"][idx])
|
||||
obs_data["init_root_orient"] = torch.Tensor(
|
||||
self.data_dict["init_root_orient"][idx]
|
||||
)
|
||||
obs_data["init_trans"] = torch.Tensor(self.data_dict["init_trans"][idx])
|
||||
obs_data["is_right"] = torch.Tensor(self.data_dict["is_right"][idx])
|
||||
|
||||
# floor plane
|
||||
# obs_data["floor_plane"] = torch.Tensor(self.data_dict["floor_plane"][idx])
|
||||
|
||||
# the frames the track is visible in
|
||||
obs_data["vis_mask"] = torch.Tensor(self.data_dict["vis_mask"][idx])
|
||||
|
||||
# the frames used in this subsequence
|
||||
obs_data["seq_interval"] = torch.Tensor(list(self.seq_intervals[idx])).to(
|
||||
torch.int
|
||||
)
|
||||
# the start and end interval of available keypoints
|
||||
obs_data["track_interval"] = torch.Tensor(
|
||||
self.data_dict["track_interval"][idx]
|
||||
).int()
|
||||
|
||||
print(len(self.track_ids), self.track_ids[idx], obs_data["is_right"].shape, obs_data["seq_interval"].shape)
|
||||
|
||||
obs_data["track_id"] = int(self.track_ids[idx])
|
||||
# print(self.track_ids[idx], obs_data["is_right"])
|
||||
# print(self.track_ids[idx] == obs_data["is_right"])
|
||||
# is_right should be constant across time for each track
|
||||
assert torch.all(obs_data["is_right"] == obs_data["is_right"][0]), "is_right should be constant"
|
||||
assert int(self.track_ids[idx]) == int(obs_data["is_right"][0]), f"Track ID {self.track_ids[idx]} != is_right {obs_data['is_right'][0]}"
|
||||
obs_data["seq_name"] = self.seq_name
|
||||
return obs_data
|
||||
|
||||
def load_camera_data(self):
|
||||
cam_dir = self.data_sources["cameras"]
|
||||
data_interval = 0, -1
|
||||
if self.split_cameras:
|
||||
data_interval = self.data_start, self.data_end
|
||||
track_interval = self.start_idx, self.end_idx
|
||||
self.cam_data = CameraData(
|
||||
cam_dir, self.seq_len, self.img_size, self.is_static, data_interval, track_interval
|
||||
)
|
||||
|
||||
def get_camera_data(self):
|
||||
if self.cam_data is None:
|
||||
raise ValueError
|
||||
return self.cam_data.as_dict()
|
||||
|
||||
|
||||
class CameraData(object):
|
||||
def __init__(
|
||||
self, cam_dir, seq_len, img_size, is_static, data_interval=[0, -1], track_interval=[0, -1]
|
||||
):
|
||||
self.img_size = img_size
|
||||
self.cam_dir = cam_dir
|
||||
self.is_static = is_static
|
||||
|
||||
# inclusive exclusive
|
||||
data_start, data_end = data_interval
|
||||
if data_end < 0:
|
||||
data_end += seq_len + 1
|
||||
data_len = data_end - data_start
|
||||
|
||||
# start and end indices are with respect to the data interval
|
||||
sidx, eidx = track_interval
|
||||
if eidx < 0:
|
||||
eidx += data_len + 1
|
||||
|
||||
self.sidx, self.eidx = sidx + data_start, eidx + data_start
|
||||
print(self.sidx, self.eidx, data_start, data_end, data_interval, track_interval)
|
||||
#
|
||||
self.seq_len = self.eidx - self.sidx
|
||||
|
||||
self.load_data()
|
||||
|
||||
def load_data(self):
|
||||
# camera info
|
||||
sidx, eidx = self.sidx, self.eidx
|
||||
img_w, img_h = self.img_size
|
||||
fpath = os.path.join(self.cam_dir, "cameras.npz")
|
||||
if os.path.isfile(fpath) and not self.is_static:
|
||||
Logger.log(f"Loading cameras from {fpath}...")
|
||||
cam_R, cam_t, intrins, width, height = load_cameras_npz(fpath, self.seq_len)
|
||||
scale = img_w / width
|
||||
self.intrins = scale * intrins[sidx:eidx]
|
||||
# move first camera to origin
|
||||
# R0, t0 = invert_camera(cam_R[sidx], cam_t[sidx])
|
||||
# self.cam_R = torch.einsum("ij,...jk->...ik", R0, cam_R[sidx:eidx])
|
||||
# self.cam_t = t0 + torch.einsum("ij,...j->...i", R0, cam_t[sidx:eidx])
|
||||
# t0 = -cam_t[sidx:eidx].mean(dim=0) + torch.randn(3) * 0.1
|
||||
t0 = -cam_t[sidx:sidx+1] + torch.randn(3) * 0.1
|
||||
self.cam_R = cam_R[sidx:eidx]
|
||||
self.cam_t = cam_t[sidx:eidx] - t0
|
||||
else:
|
||||
# raise ValueError
|
||||
Logger.log(f"WARNING: {fpath} does not exist, using static cameras...")
|
||||
default_focal = 0.5 * (img_h + img_w)
|
||||
self.intrins = torch.tensor(
|
||||
[default_focal, default_focal, img_w / 2, img_h / 2]
|
||||
)[None].repeat(self.seq_len, 1)
|
||||
|
||||
self.cam_R = torch.eye(3)[None].repeat(self.seq_len, 1, 1)
|
||||
self.cam_t = torch.zeros(self.seq_len, 3)
|
||||
|
||||
Logger.log(f"Images have {img_w}x{img_h}, intrins {self.intrins[0]}")
|
||||
print("CAMERA DATA", self.cam_R.shape, self.cam_t.shape, self.intrins[0])
|
||||
|
||||
def world2cam(self):
|
||||
return self.cam_R, self.cam_t
|
||||
|
||||
def cam2world(self):
|
||||
print('dataset.py, cam2world: ', self.cam_R)
|
||||
R = self.cam_R.transpose(-1, -2)
|
||||
t = -torch.einsum("bij,bj->bi", R, self.cam_t)
|
||||
print('dataset.py, cam2world: ', R)
|
||||
return R, t
|
||||
|
||||
def as_dict(self):
|
||||
return {
|
||||
"cam_R": self.cam_R, # (T, 3, 3)
|
||||
"cam_t": self.cam_t, # (T, 3)
|
||||
"intrins": self.intrins, # (T, 4)
|
||||
"static": self.is_static, # bool
|
||||
}
|
||||
|
||||
|
||||
def get_ternary_mask(vis_mask):
|
||||
# get the track start and end idcs relative to the filtered interval
|
||||
vis_mask = torch.as_tensor(vis_mask)
|
||||
vis_idcs = torch.where(vis_mask)[0]
|
||||
track_s, track_e = min(vis_idcs), max(vis_idcs) + 1
|
||||
# -1 = track out of scene, 0 = occlusion, 1 = visible
|
||||
vis_mask = vis_mask.float()
|
||||
vis_mask[:track_s] = -1
|
||||
vis_mask[track_e:] = -1
|
||||
return vis_mask
|
||||
|
||||
|
||||
def get_shot_img_files(shots_path, shot_idx, shot_pad=SHOT_PAD):
|
||||
assert os.path.isfile(shots_path)
|
||||
with open(shots_path, "r") as f:
|
||||
shots_dict = json.load(f)
|
||||
img_names = sorted(shots_dict.keys())
|
||||
N = len(img_names)
|
||||
shot_mask = np.array([shots_dict[x] == shot_idx for x in img_names])
|
||||
|
||||
idcs = np.where(shot_mask)[0]
|
||||
if shot_pad > 0: # drop the frames before/after shot change
|
||||
raise ValueError
|
||||
if min(idcs) > 0:
|
||||
idcs = idcs[shot_pad:]
|
||||
if len(idcs) > 0 and max(idcs) < N - 1:
|
||||
idcs = idcs[:-shot_pad]
|
||||
if len(idcs) < MIN_SEQ_LEN:
|
||||
raise ValueError("shot is too short for optimization")
|
||||
|
||||
shot_mask = np.zeros(N, dtype=bool)
|
||||
shot_mask[idcs] = 1
|
||||
|
||||
sel_paths = [img_names[i] for i in idcs]
|
||||
print(f"FOUND {len(idcs)}/{len(shots_dict)} FRAMES FOR SHOT {shot_idx}")
|
||||
return sel_paths, idcs
|
||||
|
||||
|
||||
def load_cameras_npz(camera_path, v):
|
||||
assert os.path.splitext(camera_path)[-1] == ".npz"
|
||||
|
||||
cam_data = np.load(camera_path)
|
||||
height, width, focal = (
|
||||
int(cam_data["height"]),
|
||||
int(cam_data["width"]),
|
||||
float(cam_data["focal"]),
|
||||
)
|
||||
|
||||
w2c = torch.from_numpy(cam_data["w2c"]) # (N, 4, 4)
|
||||
cam_R = w2c[:, :3, :3] # (N, 3, 3)
|
||||
cam_t = w2c[:, :3, 3] # (N, 3)
|
||||
N = len(w2c)
|
||||
|
||||
if "intrins" in cam_data:
|
||||
intrins = torch.from_numpy(cam_data["intrins"].astype(np.float32))
|
||||
else:
|
||||
intrins = torch.tensor([focal, focal, width / 2, height / 2])[None].repeat(N, 1)
|
||||
|
||||
print(f"Loaded {N} cameras")
|
||||
return cam_R, cam_t, intrins, width, height
|
||||
|
||||
|
||||
def is_image(x):
|
||||
return (x.endswith(".png") or x.endswith(".jpg")) and not x.startswith(".")
|
||||
|
||||
|
||||
def get_name(x):
|
||||
return os.path.splitext(os.path.basename(x))[0]
|
||||
|
||||
|
||||
def split_name(x, suffix):
|
||||
return os.path.basename(x).split(suffix)[0]
|
||||
|
||||
|
||||
def get_names_in_dir(d, suffix):
|
||||
files = [split_name(x, suffix) for x in glob.glob(f"{d}/*{suffix}")]
|
||||
return sorted(files)
|
||||
|
||||
|
||||
def batch_join(parent, names, suffix=""):
|
||||
return [os.path.join(parent, f"{n}{suffix}") for n in names]
|
||||
+168
@@ -0,0 +1,168 @@
|
||||
import os
|
||||
import json
|
||||
import functools
|
||||
|
||||
import numpy as np
|
||||
|
||||
from body_model import OP_NUM_JOINTS
|
||||
from scipy.interpolate import interp1d
|
||||
from scipy.spatial.transform import Rotation, Slerp
|
||||
|
||||
|
||||
def read_keypoints(keypoint_fn):
|
||||
"""
|
||||
Only reads body keypoint data of first person.
|
||||
"""
|
||||
empty_kps = np.zeros((OP_NUM_JOINTS, 3), dtype=np.float32)
|
||||
if not os.path.isfile(keypoint_fn):
|
||||
return empty_kps
|
||||
|
||||
with open(keypoint_fn) as keypoint_file:
|
||||
data = json.load(keypoint_file)
|
||||
|
||||
if len(data["people"]) == 0:
|
||||
print("WARNING: Found no keypoints in %s! Returning zeros!" % (keypoint_fn))
|
||||
return empty_kps
|
||||
|
||||
person_data = data["people"][0]
|
||||
body_keypoints = np.array(person_data["pose_keypoints_2d"], dtype=np.float32)
|
||||
body_keypoints = body_keypoints.reshape([-1, 3])
|
||||
return body_keypoints
|
||||
|
||||
|
||||
def load_keypoints_with_interp(kp_paths, interp=True):
|
||||
"""
|
||||
Load 2D keypoints from paths and interpolate missing frames between tmin and tmax.
|
||||
|
||||
Args:
|
||||
kp_paths: List of paths to keypoint JSON files
|
||||
interp: Whether to interpolate missing frames
|
||||
|
||||
Returns:
|
||||
joints2d_data: (T, J, 3) array of keypoints with interpolation
|
||||
"""
|
||||
# Load all keypoints
|
||||
joints2d_data = np.stack([read_keypoints(p) for p in kp_paths], axis=0).astype(np.float32)
|
||||
|
||||
if not interp:
|
||||
return joints2d_data
|
||||
|
||||
# Find which frames have valid keypoints (file exists and has data)
|
||||
vis_mask = np.array([os.path.isfile(p) for p in kp_paths])
|
||||
# Also check if keypoints are not all zeros
|
||||
has_data = ~np.all(joints2d_data == 0, axis=(1, 2))
|
||||
vis_mask = vis_mask & has_data
|
||||
vis_idcs = np.where(vis_mask)[0]
|
||||
|
||||
if len(vis_idcs) == 0 or len(vis_idcs) == 1:
|
||||
# No interpolation needed
|
||||
return joints2d_data
|
||||
|
||||
# Interpolate only between tmin and tmax (same as MANO interpolation)
|
||||
T, J, _ = joints2d_data.shape
|
||||
tmin, tmax = min(vis_idcs), max(vis_idcs) + 1
|
||||
times = np.arange(tmin, tmax)
|
||||
|
||||
# Interpolate each joint's x, y coordinates separately
|
||||
for j in range(J):
|
||||
# Interpolate x coordinate
|
||||
x_interp = interp1d(vis_idcs, joints2d_data[vis_idcs, j, 0],
|
||||
kind='linear', bounds_error=False)
|
||||
# Interpolate y coordinate
|
||||
y_interp = interp1d(vis_idcs, joints2d_data[vis_idcs, j, 1],
|
||||
kind='linear', bounds_error=False)
|
||||
|
||||
# Interpolate for frames between tmin and tmax
|
||||
for t in times:
|
||||
if t not in vis_idcs:
|
||||
# Interpolate
|
||||
joints2d_data[t, j, 0] = x_interp(t)
|
||||
joints2d_data[t, j, 1] = y_interp(t)
|
||||
|
||||
# Set confidence as average of neighboring visible frames
|
||||
prev_vis = vis_idcs[vis_idcs < t]
|
||||
next_vis = vis_idcs[vis_idcs > t]
|
||||
if len(prev_vis) > 0 and len(next_vis) > 0:
|
||||
conf_prev = joints2d_data[prev_vis[-1], j, 2]
|
||||
conf_next = joints2d_data[next_vis[0], j, 2]
|
||||
joints2d_data[t, j, 2] = (conf_prev + conf_next) / 2.0 * 0.8 # Slightly lower confidence for interpolated
|
||||
elif len(prev_vis) > 0:
|
||||
joints2d_data[t, j, 2] = joints2d_data[prev_vis[-1], j, 2] * 0.8
|
||||
elif len(next_vis) > 0:
|
||||
joints2d_data[t, j, 2] = joints2d_data[next_vis[0], j, 2] * 0.8
|
||||
|
||||
return joints2d_data
|
||||
|
||||
|
||||
def read_mask_path(path):
|
||||
mask_path = None
|
||||
if not os.path.isfile(path):
|
||||
return mask_path
|
||||
|
||||
with open(path, "r") as f:
|
||||
data = json.load(path)
|
||||
|
||||
person_data = data["people"][0]
|
||||
if "mask_path" in person_data:
|
||||
mask_path = person_data["mask_path"]
|
||||
|
||||
return mask_path
|
||||
|
||||
|
||||
def read_mano_preds(pred_path, tid, num_betas=10):
|
||||
"""
|
||||
reads the betas, body_pose, global orientation and translation of a mano prediction
|
||||
exported from phalp outputs
|
||||
returns betas (10,), body_pose (23, 3), global_orientation (3,), translation (3,)
|
||||
"""
|
||||
pose = np.zeros((15, 3))
|
||||
rot = np.zeros(3)
|
||||
trans = np.zeros(3)
|
||||
betas = np.zeros(num_betas)
|
||||
if not os.path.isfile(pred_path):
|
||||
return pose, rot, trans, betas, int(tid)
|
||||
|
||||
with open(pred_path, "r") as f:
|
||||
data = json.load(f)
|
||||
|
||||
pose = np.array(data["body_pose"], dtype=np.float32)
|
||||
rot = np.array(data["global_orient"], dtype=np.float32)
|
||||
trans = np.array(data["cam_trans"], dtype=np.float32)
|
||||
betas = np.array(data["betas"], dtype=np.float32)
|
||||
is_right = np.array(data["is_right"], dtype=np.float32)
|
||||
|
||||
return pose, rot, trans, betas, is_right
|
||||
|
||||
|
||||
def load_mano_preds(pred_paths, tid, interp=True, num_betas=10):
|
||||
vis_mask = np.array([os.path.isfile(x) for x in pred_paths])
|
||||
vis_idcs = np.where(vis_mask)[0]
|
||||
|
||||
# load single image mano predictions
|
||||
stack_fnc = functools.partial(np.stack, axis=0)
|
||||
# (N, 16, 3), (N, 3), (N, 3), (N, 10)
|
||||
pose, orient, trans, betas, is_right = map(
|
||||
stack_fnc, zip(*[read_mano_preds(p, tid, num_betas=num_betas) for p in pred_paths])
|
||||
)
|
||||
|
||||
assert len(np.where(is_right!=int(tid))[0]) == 0
|
||||
if not interp:
|
||||
return pose, orient, trans, betas, is_right
|
||||
|
||||
# interpolate the occluded tracks
|
||||
orient_slerp = Slerp(vis_idcs, Rotation.from_rotvec(orient[vis_idcs]))
|
||||
trans_interp = interp1d(vis_idcs, trans[vis_idcs], axis=0)
|
||||
betas_interp = interp1d(vis_idcs, betas[vis_idcs], axis=0)
|
||||
|
||||
tmin, tmax = min(vis_idcs), max(vis_idcs) + 1
|
||||
times = np.arange(tmin, tmax)
|
||||
orient[times] = orient_slerp(times).as_rotvec()
|
||||
trans[times] = trans_interp(times)
|
||||
betas[times] = betas_interp(times)
|
||||
|
||||
# interpolate for each joint angle
|
||||
for i in range(pose.shape[1]):
|
||||
pose_slerp = Slerp(vis_idcs, Rotation.from_rotvec(pose[vis_idcs, i]))
|
||||
pose[times, i] = pose_slerp(times).as_rotvec()
|
||||
|
||||
return pose, orient, trans, betas, is_right
|
||||
+284
@@ -0,0 +1,284 @@
|
||||
import os
|
||||
import shlex
|
||||
import sys
|
||||
from pathlib import Path
|
||||
import numpy as np
|
||||
import subprocess
|
||||
import cv2
|
||||
|
||||
import preproc.launch_hamer as hamer
|
||||
from preproc.launch_slam import split_frames_shots, get_command, check_intrins
|
||||
from preproc.extract_frames import video_to_frames, split_frame
|
||||
|
||||
|
||||
def is_nonempty(d):
|
||||
return os.path.isdir(d) and len(os.listdir(d)) > 0
|
||||
|
||||
|
||||
def preprocess_frames(img_dir, src_path, overwrite=False, **kwargs):
|
||||
if not overwrite and is_nonempty(img_dir):
|
||||
print(f"FOUND {len(os.listdir(img_dir))} FRAMES in {img_dir}")
|
||||
return
|
||||
print(f"EXTRACTING FRAMES FROM {src_path} TO {img_dir}")
|
||||
print(kwargs)
|
||||
|
||||
# out = video_to_frames(src_path, img_dir, overwrite=overwrite, **kwargs)
|
||||
out = split_frame(src_path, img_dir, overwrite=overwrite, **kwargs)
|
||||
assert out == 0, "FAILED FRAME EXTRACTION"
|
||||
|
||||
|
||||
def preprocess_tracks(datatype, img_dir, track_dir, shot_dir, gpu, overwrite=False):
|
||||
"""
|
||||
:param img_dir
|
||||
:param track_dir, expected format: res_root/track_name/sequence
|
||||
:param shot_dir, expected format: res_root/shot_name/sequence
|
||||
"""
|
||||
if not overwrite and is_nonempty(track_dir):
|
||||
print(f"FOUND TRACKS IN {track_dir}")
|
||||
return
|
||||
|
||||
print(f"RUNNING HAMER ON {img_dir}")
|
||||
track_root, seq = os.path.split(track_dir.rstrip("/"))
|
||||
res_root, track_name = os.path.split(track_root)
|
||||
shot_name = shot_dir.rstrip("/").split("/")[-2]
|
||||
|
||||
hamer.process_seq(
|
||||
[gpu],
|
||||
res_root,
|
||||
seq,
|
||||
img_dir,
|
||||
track_name=track_name,
|
||||
shot_name=shot_name,
|
||||
datatype=datatype,
|
||||
overwrite=overwrite,
|
||||
)
|
||||
|
||||
|
||||
def load_vipe_cameras(vipe_dir, seq_name, img_dir, start=0, end=-1):
|
||||
"""
|
||||
Load VIPE camera outputs and convert to DROID-SLAM format
|
||||
|
||||
Args:
|
||||
vipe_dir: Path to VIPE results directory
|
||||
seq_name: Sequence name
|
||||
img_dir: Image directory to get image size
|
||||
start: Start frame index
|
||||
end: End frame index
|
||||
|
||||
Returns:
|
||||
w2c: (N, 4, 4) world-to-camera matrices
|
||||
intrins_full: (N, 6) intrinsics [fx, fy, cx, cy, W, H]
|
||||
"""
|
||||
# Load VIPE pose and intrinsics
|
||||
vipe_pose_path = os.path.join(vipe_dir, "pose", f"{seq_name}.npz")
|
||||
vipe_intrins_path = os.path.join(vipe_dir, "intrinsics", f"{seq_name}.npz")
|
||||
|
||||
if not os.path.exists(vipe_pose_path):
|
||||
raise FileNotFoundError(f"VIPE pose file not found: {vipe_pose_path}")
|
||||
if not os.path.exists(vipe_intrins_path):
|
||||
raise FileNotFoundError(f"VIPE intrinsics file not found: {vipe_intrins_path}")
|
||||
|
||||
print(f"Loading VIPE cameras from {vipe_dir}")
|
||||
pose_data = np.load(vipe_pose_path)
|
||||
c2w = pose_data['data'] # (N, 4, 4) camera-to-world
|
||||
pose_inds = pose_data['inds'] # (N,) frame indices
|
||||
|
||||
intrins_data = np.load(vipe_intrins_path)
|
||||
intrins = intrins_data['data'] # (N, 4) [fx, fy, cx, cy]
|
||||
intrins_inds = intrins_data['inds'] # (N,) frame indices
|
||||
|
||||
# Verify indices match
|
||||
assert np.array_equal(pose_inds, intrins_inds), "Pose and intrinsics indices don't match!"
|
||||
|
||||
# Get image size from first image
|
||||
image_files = sorted([f for f in os.listdir(img_dir)
|
||||
if f.endswith(('.png', '.jpg', '.jpeg'))])
|
||||
if not image_files:
|
||||
# Infer from intrinsics (cx, cy should be roughly at center)
|
||||
img_width = int(intrins[0, 2] * 2)
|
||||
img_height = int(intrins[0, 3] * 2)
|
||||
print(f"Inferred image size from intrinsics: {img_width}x{img_height}")
|
||||
else:
|
||||
img_path = os.path.join(img_dir, image_files[0])
|
||||
img = cv2.imread(img_path)
|
||||
img_height, img_width = img.shape[:2]
|
||||
print(f"Got image size from images: {img_width}x{img_height}")
|
||||
|
||||
# Select frames based on start/end
|
||||
if end < 0:
|
||||
end = len(c2w)
|
||||
c2w = c2w[start:end]
|
||||
intrins = intrins[start:end]
|
||||
|
||||
# Convert camera-to-world to world-to-camera
|
||||
w2c = np.linalg.inv(c2w)
|
||||
|
||||
# Add width and height to intrinsics
|
||||
N = len(w2c)
|
||||
intrins_full = np.zeros((N, 6), dtype=np.float32)
|
||||
intrins_full[:, :4] = intrins
|
||||
intrins_full[:, 4] = img_width
|
||||
intrins_full[:, 5] = img_height
|
||||
|
||||
print(f"Loaded {N} VIPE camera poses")
|
||||
return w2c, intrins_full
|
||||
|
||||
|
||||
def save_vipe_cameras_as_droid(output_dir, w2c, intrins_full):
|
||||
"""
|
||||
Save VIPE cameras in DROID-SLAM format
|
||||
|
||||
Args:
|
||||
output_dir: Output directory
|
||||
w2c: (N, 4, 4) world-to-camera matrices
|
||||
intrins_full: (N, 6) intrinsics [fx, fy, cx, cy, W, H]
|
||||
"""
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
|
||||
# Extract parameters
|
||||
W, H = intrins_full[0, 4], intrins_full[0, 5]
|
||||
focal = intrins_full[:, :2].mean()
|
||||
|
||||
print(f"Saving VIPE cameras to {output_dir}")
|
||||
print(f"Image size: {int(W)}x{int(H)}, focal: {focal:.2f}")
|
||||
|
||||
# Save cameras.npz (main format used by Dyn-HaMR)
|
||||
np.savez(
|
||||
f"{output_dir}/cameras.npz",
|
||||
height=H,
|
||||
width=W,
|
||||
focal=focal,
|
||||
intrins=intrins_full[:, :4],
|
||||
w2c=w2c,
|
||||
)
|
||||
print(f"Saved cameras.npz with {len(w2c)} frames")
|
||||
|
||||
|
||||
def run_vipe(video_path, vipe_dir, vipe_root=None, vipe_python=None, vipe_pipeline=None):
|
||||
"""Run VIPE without shell activation, using the project venv by default.
|
||||
|
||||
``VIPE_PYTHON`` or ``vipe_python`` can select a separate interpreter.
|
||||
Input/output paths are resolved before changing the child working directory.
|
||||
"""
|
||||
project_root = Path(__file__).resolve().parents[4]
|
||||
root = Path(vipe_root or project_root / "third_party" / "vipe").expanduser().resolve()
|
||||
video = Path(video_path).expanduser().absolute()
|
||||
output = Path(vipe_dir).expanduser().absolute()
|
||||
configured_python = vipe_python or os.environ.get("VIPE_PYTHON")
|
||||
project_python = project_root / ".venv" / "bin" / "python"
|
||||
python = Path(configured_python or (project_python if project_python.is_file() else sys.executable)).expanduser().absolute()
|
||||
if not video.is_file():
|
||||
raise FileNotFoundError(f"VIPE input video not found: {video}")
|
||||
if not (root / "vipe" / "cli" / "main.py").is_file():
|
||||
raise FileNotFoundError(f"VIPE source not found in {root}; set data.vipe_root")
|
||||
if not python.is_file():
|
||||
raise FileNotFoundError(f"VIPE interpreter not found: {python}; set VIPE_PYTHON")
|
||||
cmd = [str(python), "-m", "vipe.cli.main", "infer", str(video), "--output", str(output)]
|
||||
if vipe_pipeline:
|
||||
cmd.extend(["--pipeline", str(vipe_pipeline)])
|
||||
print(f"Running VIPE in {root}: {shlex.join(cmd)}", flush=True)
|
||||
out = subprocess.call(cmd, cwd=str(root))
|
||||
if out != 0:
|
||||
print(f"WARNING: VIPE failed with exit code {out}")
|
||||
else:
|
||||
print("VIPE completed successfully")
|
||||
return out
|
||||
|
||||
|
||||
def preprocess_cameras(cfg, overwrite=False):
|
||||
if not overwrite and is_nonempty(cfg.sources.cameras):
|
||||
print(f"FOUND CAMERAS IN {cfg.sources.cameras}")
|
||||
return
|
||||
|
||||
# Check if we should use VIPE instead of DROID-SLAM
|
||||
use_vipe = cfg.get("use_vipe", False)
|
||||
vipe_dir = cfg.get("vipe_dir", None)
|
||||
|
||||
if use_vipe and vipe_dir is not None:
|
||||
# Check if VIPE results exist for this sequence
|
||||
vipe_pose_path = os.path.join(vipe_dir, "pose", f"{cfg.seq}.npz")
|
||||
vipe_intrins_path = os.path.join(vipe_dir, "intrinsics", f"{cfg.seq}.npz")
|
||||
|
||||
# If VIPE results don't exist, run VIPE
|
||||
if not (os.path.exists(vipe_pose_path) and os.path.exists(vipe_intrins_path)):
|
||||
print(f"VIPE results not found for sequence '{cfg.seq}', running VIPE...")
|
||||
|
||||
# The source location is independent of the output directory.
|
||||
vipe_root = cfg.get("vipe_root", None)
|
||||
video_path = cfg.get("src_path", None)
|
||||
|
||||
if video_path is None or not os.path.exists(video_path):
|
||||
raise FileNotFoundError(
|
||||
f"Video path not found: {video_path}\n"
|
||||
f"Cannot run VIPE. Please provide a valid 'src_path' in your config."
|
||||
)
|
||||
|
||||
# Run VIPE
|
||||
out = run_vipe(video_path, vipe_dir, vipe_root, cfg.get("vipe_python", None), cfg.get("vipe_pipeline", None))
|
||||
if out != 0:
|
||||
raise RuntimeError(
|
||||
f"VIPE failed with exit code {out}\n"
|
||||
f"Check the command and error above. Select an installed Python "
|
||||
f"with VIPE_PYTHON or data.vipe_python; conda is not required."
|
||||
)
|
||||
|
||||
# Load VIPE results (after potentially running VIPE)
|
||||
if not (os.path.exists(vipe_pose_path) and os.path.exists(vipe_intrins_path)):
|
||||
raise FileNotFoundError(
|
||||
f"VIPE results not found after execution:\n"
|
||||
f" Pose: {vipe_pose_path}\n"
|
||||
f" Intrinsics: {vipe_intrins_path}\n"
|
||||
f"VIPE may have failed silently. Please check VIPE logs."
|
||||
)
|
||||
|
||||
print(f"USING VIPE CAMERAS FROM {vipe_dir}")
|
||||
img_dir = cfg.sources.images
|
||||
map_dir = cfg.sources.cameras
|
||||
|
||||
# Get frame range
|
||||
subseqs, shot_idcs = split_frames_shots(cfg.sources.images, cfg.sources.shots)
|
||||
shot_idx = np.where(shot_idcs == cfg.shot_idx)[0][0]
|
||||
start, end = subseqs[shot_idx]
|
||||
|
||||
if not cfg.split_cameras:
|
||||
# only run on specified segment within shot
|
||||
end = start + cfg.end_idx
|
||||
start = start + cfg.start_idx
|
||||
|
||||
# Load and convert VIPE cameras
|
||||
w2c, intrins_full = load_vipe_cameras(vipe_dir, cfg.seq, img_dir, start, end)
|
||||
save_vipe_cameras_as_droid(map_dir, w2c, intrins_full)
|
||||
return
|
||||
|
||||
# Default: use DROID-SLAM
|
||||
print(f"RUNNING SLAM ON {cfg.seq}")
|
||||
img_dir = cfg.sources.images
|
||||
map_dir = cfg.sources.cameras
|
||||
subseqs, shot_idcs = split_frames_shots(cfg.sources.images, cfg.sources.shots)
|
||||
print(shot_idcs, cfg.shot_idx, np.where(shot_idcs == cfg.shot_idx), cfg.sources.images, cfg.sources.shots)
|
||||
print(subseqs)
|
||||
shot_idx = np.where(shot_idcs == cfg.shot_idx)[0][0]
|
||||
# run on selected shot
|
||||
start, end = subseqs[shot_idx]
|
||||
if not cfg.split_cameras:
|
||||
# only run on specified segment within shot
|
||||
end = start + cfg.end_idx
|
||||
start = start + cfg.start_idx
|
||||
intrins_path = cfg.sources.get("intrins", None)
|
||||
if intrins_path is not None:
|
||||
intrins_path = check_intrins(cfg.type, cfg.root, intrins_path, cfg.seq, cfg.split)
|
||||
|
||||
print('img_dir, map_dir, start, end, intrins_path', img_dir, map_dir, start, end, intrins_path)
|
||||
# raise ValueERRROR
|
||||
cmd = get_command(
|
||||
img_dir,
|
||||
map_dir,
|
||||
start=start,
|
||||
end=end,
|
||||
intrins_path=intrins_path,
|
||||
overwrite=overwrite,
|
||||
)
|
||||
print(cmd)
|
||||
gpu = os.environ.get("CUDA_VISIBLE_DEVICES", 0)
|
||||
out = subprocess.call(f"CUDA_VISIBLE_DEVICES={gpu} {cmd}", shell=True)
|
||||
assert out == 0, "SLAM FAILED"
|
||||
@@ -0,0 +1,5 @@
|
||||
from . import camera
|
||||
from . import mesh
|
||||
from . import pcl
|
||||
from . import plane
|
||||
from . import rotation
|
||||
+351
@@ -0,0 +1,351 @@
|
||||
import torch
|
||||
import numpy as np
|
||||
|
||||
|
||||
def perspective_projection(
|
||||
points, focal_length, camera_center, rotation=None, translation=None
|
||||
):
|
||||
"""
|
||||
Adapted from https://github.com/mkocabas/VIBE/blob/master/lib/models/spin.py
|
||||
This function computes the perspective projection of a set of points.
|
||||
Input:
|
||||
points (bs, N, 3): 3D points
|
||||
focal_length (bs, 2): Focal length
|
||||
camera_center (bs, 2): Camera center
|
||||
rotation (bs, 3, 3): OPTIONAL Camera rotation
|
||||
translation (bs, 3): OPTIONAL Camera translation
|
||||
"""
|
||||
batch_size = points.shape[0]
|
||||
K = torch.zeros([batch_size, 3, 3], device=points.device)
|
||||
K[:, 0, 0] = focal_length[:, 0]
|
||||
K[:, 1, 1] = focal_length[:, 1]
|
||||
K[:, 2, 2] = 1.0
|
||||
K[:, :-1, -1] = camera_center
|
||||
|
||||
if rotation is not None and translation is not None:
|
||||
# Transform points
|
||||
points = torch.einsum("bij,bkj->bki", rotation, points)
|
||||
points = points + translation.unsqueeze(1)
|
||||
|
||||
# Apply perspective distortion
|
||||
projected_points = points / points[..., 2:3]
|
||||
|
||||
# Apply camera intrinsics
|
||||
projected_points = torch.einsum("bij,bkj->bki", K, projected_points)
|
||||
|
||||
return projected_points[:, :, :-1]
|
||||
|
||||
|
||||
def reproject(points3d, cam_R, cam_t, cam_f, cam_center):
|
||||
"""
|
||||
reproject points3d into the scene cameras
|
||||
:param points3d (B, T, N, 3)
|
||||
:param cam_R (B, T, 3, 3)
|
||||
:param cam_t (B, T, 3)
|
||||
:param cam_f (T, 2)
|
||||
:param cam_center (T, 2)
|
||||
"""
|
||||
# print(points3d.shape, cam_R.shape, cam_t.shape, cam_f.shape, cam_center.shape)
|
||||
# print(cam_f[None, :, None].shape, cam_center[None, :, None].shape)
|
||||
# raise ValueError
|
||||
B, T, N, _ = points3d.shape
|
||||
points3d = torch.einsum("btij,btnj->btni", cam_R, points3d)
|
||||
points3d = points3d + cam_t[..., None, :] # (B, T, N, 3)
|
||||
points2d = points3d[..., :2] / points3d[..., 2:3]
|
||||
points2d = cam_f[None, :, None] * points2d + cam_center[None, :, None]
|
||||
return points2d
|
||||
|
||||
|
||||
def focal2fov(focal, R):
|
||||
"""
|
||||
:param focal, focal length
|
||||
:param R, either W / 2 or H / 2
|
||||
"""
|
||||
return 2 * np.arctan(R / focal)
|
||||
|
||||
|
||||
def fov2focal(fov, R):
|
||||
"""
|
||||
:param fov, field of view in radians
|
||||
:param R, either W / 2 or H / 2
|
||||
"""
|
||||
return R / np.tan(fov / 2)
|
||||
|
||||
|
||||
def compute_lookat_box(bb_min, bb_max, intrins):
|
||||
"""
|
||||
The center and distance to a scene with bb_min, bb_max
|
||||
to place a camera with given intrinsics
|
||||
:param bb_min (3,)
|
||||
:param bb_max (3,)
|
||||
:param intrinsics, (fx, fy, cx, cy) of camera
|
||||
:param view_angle (optional) viewing angle in radians (elevation)
|
||||
"""
|
||||
fx, fy, cx, cy = intrins
|
||||
bb_min, bb_max = torch.tensor(bb_min), torch.tensor(bb_max)
|
||||
center = 0.5 * (bb_min + bb_max)
|
||||
size = torch.linalg.norm(bb_max - bb_min)
|
||||
cam_dist = np.sqrt(fx**2 + fy**2) / np.sqrt(cx**2 + cy**2)
|
||||
cam_dist = 0.75 * size * cam_dist
|
||||
return center, cam_dist
|
||||
|
||||
|
||||
def lookat_origin(cam_dist, view_angle=-np.pi / 6):
|
||||
"""
|
||||
:param cam_dist (float)
|
||||
:param view_angle (float)
|
||||
"""
|
||||
cam_dist = np.abs(cam_dist)
|
||||
view_angle = np.abs(view_angle)
|
||||
pos = cam_dist * torch.tensor([0, np.sin(view_angle), np.cos(view_angle)])
|
||||
rot = rotx(view_angle)
|
||||
return rot, pos
|
||||
|
||||
|
||||
def lookat_matrix(source_pos, target_pos, up):
|
||||
"""
|
||||
IMPORTANT: USES RIGHT UP BACK XYZ CONVENTION
|
||||
:param source_pos (*, 3)
|
||||
:param target_pos (*, 3)
|
||||
:param up (3,)
|
||||
"""
|
||||
*dims, _ = source_pos.shape
|
||||
up = up.reshape(*(1,) * len(dims), 3)
|
||||
up = up / torch.linalg.norm(up, dim=-1, keepdim=True)
|
||||
back = normalize(target_pos - source_pos)
|
||||
right = normalize(torch.linalg.cross(up, back))
|
||||
up = normalize(torch.linalg.cross(back, right))
|
||||
R = torch.stack([right, up, back], dim=-1)
|
||||
return make_4x4_pose(R, source_pos)
|
||||
|
||||
|
||||
def normalize(x):
|
||||
return x / torch.linalg.norm(x, dim=-1, keepdim=True)
|
||||
|
||||
|
||||
def invert_camera(R, t):
|
||||
"""
|
||||
:param R (*, 3, 3)
|
||||
:param t (*, 3)
|
||||
returns Ri (*, 3, 3), ti (*, 3)
|
||||
"""
|
||||
R, t = torch.tensor(R), torch.tensor(t)
|
||||
Ri = R.transpose(-1, -2)
|
||||
ti = -torch.einsum("...ij,...j->...i", Ri, t)
|
||||
return Ri, ti
|
||||
|
||||
|
||||
def compose_cameras(R1, t1, R2, t2):
|
||||
"""
|
||||
composes [R1, t1] and [R2, t2]
|
||||
:param R1 (*, 3, 3)
|
||||
:param t1 (*, 3)
|
||||
:param R2 (*, 3, 3)
|
||||
:param t2 (*, 3)
|
||||
"""
|
||||
R = torch.einsum("...ij,...jk->...ik", R1, R2)
|
||||
t = t1 + torch.einsum("...ij,...j->...i", R1, t2)
|
||||
return R, t
|
||||
|
||||
|
||||
def matmul_nd(A, x):
|
||||
"""
|
||||
multiply batch matrix A to batch nd tensors
|
||||
:param A (B, m, n)
|
||||
:param x (B, *dims, m)
|
||||
"""
|
||||
B, m, n = A.shape
|
||||
assert len(A) == len(x)
|
||||
assert x.shape[-1] == m
|
||||
B, *dims, _ = x.shape
|
||||
return torch.matmul(A.reshape(B, *(1,) * len(dims), m, n), x[..., None])[..., 0]
|
||||
|
||||
|
||||
def view_matrix(z, up, pos):
|
||||
"""
|
||||
:param z (*, 3) up (*, 3) pos (*, 3)
|
||||
returns (*, 4, 4)
|
||||
"""
|
||||
*dims, _ = z.shape
|
||||
x = normalize(torch.linalg.cross(up, z))
|
||||
y = normalize(torch.linalg.cross(z, x))
|
||||
bottom = (
|
||||
torch.tensor([0, 0, 0, 1], dtype=torch.float32)
|
||||
.reshape(*(1,) * len(dims), 1, 4)
|
||||
.expand(*dims, 1, 4)
|
||||
)
|
||||
|
||||
return torch.cat([torch.stack([x, y, z, pos], dim=-1), bottom], dim=-2)
|
||||
|
||||
|
||||
def average_pose(poses):
|
||||
"""
|
||||
:param poses (N, 4, 4)
|
||||
returns average pose (4, 4)
|
||||
"""
|
||||
center = poses[:, :3, 3].mean(0)
|
||||
up = normalize(poses[:, :3, 1].sum(0))
|
||||
z = normalize(poses[:, :3, 2].sum(0))
|
||||
return view_matrix(z, up, center)
|
||||
|
||||
|
||||
def project_so3(M, eps=1e-4):
|
||||
"""
|
||||
:param M (N, *, 3, 3)
|
||||
"""
|
||||
N, *dims, _, _ = M.shape
|
||||
M = M * (1 + torch.rand(N, *dims, 1, 3, device=M.device))
|
||||
U, D, Vt = torch.linalg.svd(M) # (N, *, 3, 3), (N, *, 3), (N, *, 3, 3)
|
||||
detuvt = torch.linalg.det(torch.matmul(U, Vt)) # (N, *)
|
||||
S = torch.cat(
|
||||
[torch.ones(N, *dims, 2, device=M.device), detuvt[..., None]], dim=-1
|
||||
) # (N, *, 3)
|
||||
return torch.matmul(U, torch.matmul(torch.diag_embed(S), Vt))
|
||||
|
||||
|
||||
def make_translation(t):
|
||||
return make_4x4_pose(torch.eye(3), t)
|
||||
|
||||
|
||||
def make_rotation(rx=0, ry=0, rz=0, order="xyz"):
|
||||
Rx = rotx(rx)
|
||||
Ry = roty(ry)
|
||||
Rz = rotz(rz)
|
||||
if order == "xyz":
|
||||
R = Rz @ Ry @ Rx
|
||||
elif order == "xzy":
|
||||
R = Ry @ Rz @ Rx
|
||||
elif order == "yxz":
|
||||
R = Rz @ Rx @ Ry
|
||||
elif order == "yzx":
|
||||
R = Rx @ Rz @ Ry
|
||||
elif order == "zyx":
|
||||
R = Rx @ Ry @ Rz
|
||||
elif order == "zxy":
|
||||
R = Ry @ Rx @ Rz
|
||||
return make_4x4_pose(R, torch.zeros(3))
|
||||
|
||||
|
||||
def make_4x4_pose(R, t):
|
||||
"""
|
||||
:param R (*, 3, 3)
|
||||
:param t (*, 3)
|
||||
return (*, 4, 4)
|
||||
"""
|
||||
dims = R.shape[:-2]
|
||||
pose_3x4 = torch.cat([R, t.view(*dims, 3, 1)], dim=-1)
|
||||
bottom = (
|
||||
torch.tensor([0, 0, 0, 1], device=R.device)
|
||||
.reshape(*(1,) * len(dims), 1, 4)
|
||||
.expand(*dims, 1, 4)
|
||||
)
|
||||
return torch.cat([pose_3x4, bottom], dim=-2)
|
||||
|
||||
|
||||
def normalize(x):
|
||||
return x / torch.sqrt(torch.sum(x**2, dim=-1, keepdim=True))
|
||||
|
||||
|
||||
def rotx(theta):
|
||||
return torch.tensor(
|
||||
[
|
||||
[1, 0, 0],
|
||||
[0, np.cos(theta), -np.sin(theta)],
|
||||
[0, np.sin(theta), np.cos(theta)],
|
||||
],
|
||||
dtype=torch.float32,
|
||||
)
|
||||
|
||||
|
||||
def roty(theta):
|
||||
return torch.tensor(
|
||||
[
|
||||
[np.cos(theta), 0, np.sin(theta)],
|
||||
[0, 1, 0],
|
||||
[-np.sin(theta), 0, np.cos(theta)],
|
||||
],
|
||||
dtype=torch.float32,
|
||||
)
|
||||
|
||||
|
||||
def rotz(theta):
|
||||
return torch.tensor(
|
||||
[
|
||||
[np.cos(theta), -np.sin(theta), 0],
|
||||
[np.sin(theta), np.cos(theta), 0],
|
||||
[0, 0, 1],
|
||||
],
|
||||
dtype=torch.float32,
|
||||
)
|
||||
|
||||
|
||||
def relative_pose_c2w(Rwc1, Rwc2, twc1, twc2):
|
||||
"""
|
||||
compute relative pose from cam 1 to cam 2 given c2w pose matrices
|
||||
:param Rwc1, Rwc2 (N, 3, 3) cam1, cam2 to world rotations
|
||||
:param twc1, twc2 (N, 3) cam1, cam2 to world translations
|
||||
returns R21 (N, 3, 3) t21 (N, 3)
|
||||
"""
|
||||
twc1 = twc1.view(-1, 3, 1)
|
||||
twc2 = twc2.view(-1, 3, 1)
|
||||
Rc2w = Rwc2.transpose(-1, -2) # world to c2
|
||||
tc2w = -torch.matmul(Rc2w, twc2)
|
||||
Rc2c1 = torch.matmul(Rc2w, Rwc1)
|
||||
tc2c1 = tc2w + torch.matmul(Rc2w, twc1)
|
||||
return Rc2c1, tc2c1[..., 0]
|
||||
|
||||
|
||||
def relative_pose_w2c(Rc1w, Rc2w, tc1w, tc2w):
|
||||
"""
|
||||
compute relative pose from cam 1 to cam 2 given w2c camera matrices
|
||||
:param Rc1w, Rc2w (N, 3, 3) world to cam1, cam2 rotations
|
||||
:param tc1w, tc2w (N, 3) world to cam1, cam2 translations
|
||||
"""
|
||||
tc1w = tc1w.view(-1, 3, 1)
|
||||
tc2w = tc2w.view(-1, 3, 1)
|
||||
# we keep the world to cam transforms
|
||||
Rwc1 = Rc1w.transpose(-1, -2) # c1 to world
|
||||
twc1 = -torch.matmul(Rwc1, tc1w)
|
||||
Rc2c1 = torch.matmul(Rc2w, Rwc1) # c1 to c2
|
||||
tc2c1 = tc2w + torch.matmul(Rc2w, twc1)
|
||||
return Rc2c1, tc2c1[..., 0]
|
||||
|
||||
|
||||
def project(xyz_c, center, focal, eps=1e-5):
|
||||
"""
|
||||
:param xyz_c (*, 3) 3d point in camera coordinates
|
||||
:param focal (1)
|
||||
:param center (*, 2)
|
||||
return (*, 2)
|
||||
"""
|
||||
return focal * xyz_c[..., :2] / (xyz_c[..., 2:3] + eps) + center # (N, *, 2)
|
||||
|
||||
|
||||
def convert_yup(xyz):
|
||||
"""
|
||||
converts points in x right y down z forward to x right y up z back
|
||||
:param xyz (*, 3)
|
||||
"""
|
||||
x, y, z = torch.split(xyz[..., :3], 1, dim=-1)
|
||||
return torch.cat([x, -y, -z], dim=-1)
|
||||
|
||||
|
||||
def inv_project(uv, z, center, focal, yup=True):
|
||||
"""
|
||||
:param uv (*, 2)
|
||||
:param z (*, 1)
|
||||
:param center (*, 2)
|
||||
:param focal (1)
|
||||
:returns (*, 3)
|
||||
"""
|
||||
uv = uv - center
|
||||
if yup:
|
||||
return z * torch.cat(
|
||||
[uv[..., :1] / focal, -uv[..., 1:2] / focal, -torch.ones_like(uv[..., :1])],
|
||||
dim=-1,
|
||||
) # (N, *, 3)
|
||||
|
||||
return z * torch.cat(
|
||||
[uv / focal, torch.ones_like(uv[..., :1])],
|
||||
dim=-1,
|
||||
) # (N, *, 3)
|
||||
+145
@@ -0,0 +1,145 @@
|
||||
import os
|
||||
import numpy as np
|
||||
import trimesh
|
||||
|
||||
|
||||
def get_mesh_bb(mesh):
|
||||
"""
|
||||
:param mesh - trimesh mesh object
|
||||
returns bb_min (3), bb_max (3)
|
||||
"""
|
||||
bb_min = mesh.vertices.max(axis=0)
|
||||
bb_max = mesh.vertices.min(axis=0)
|
||||
return bb_min, bb_max
|
||||
|
||||
|
||||
def get_scene_bb(meshes):
|
||||
"""
|
||||
:param mesh_seqs - (potentially nested) list of trimesh objects
|
||||
returns bb_min (3), bb_max (3)
|
||||
"""
|
||||
if isinstance(meshes, trimesh.Trimesh):
|
||||
return get_mesh_bb(meshes)
|
||||
|
||||
bb_mins, bb_maxs = zip(*[get_scene_bb(mesh) for mesh in meshes])
|
||||
bb_mins = np.stack(bb_mins, axis=0)
|
||||
bb_maxs = np.stack(bb_maxs, axis=0)
|
||||
return bb_mins.min(axis=0), bb_maxs.max(axis=0)
|
||||
|
||||
|
||||
def make_batch_mesh(verts, faces, colors):
|
||||
"""
|
||||
convenience function to make batch of meshes
|
||||
meshs have same faces in batch, verts have same color in mesh
|
||||
:param verts (B, V, 3)
|
||||
:param faces (F, 3)
|
||||
:param colors (B, 3)
|
||||
"""
|
||||
B, V, _ = verts.shape
|
||||
return [make_mesh(verts[b], faces, colors[b, None].expand(V, -1)) for b in range(B)]
|
||||
|
||||
|
||||
def make_mesh(verts, faces, colors=None, yup=True):
|
||||
"""
|
||||
create a trimesh object for the faces and vertices
|
||||
:param verts (V, 3) tensor
|
||||
:param faces (F, 3) tensor
|
||||
:param colors (optional) (V, 3) tensor
|
||||
:param yup (optional bool) whether or not to save with Y up
|
||||
"""
|
||||
verts = verts.detach().cpu().numpy()
|
||||
faces = faces.detach().cpu().numpy()
|
||||
if yup:
|
||||
verts = np.array([1, -1, -1])[None, :] * verts
|
||||
if colors is None:
|
||||
colors = np.ones_like(verts) * 0.5
|
||||
else:
|
||||
colors = colors.detach().cpu().numpy()
|
||||
return trimesh.Trimesh(
|
||||
vertices=verts, faces=faces, vertex_colors=colors, process=False
|
||||
)
|
||||
|
||||
|
||||
def save_mesh_scenes(out_dir, scenes):
|
||||
"""
|
||||
:param scenes, list of scenes (list of meshes)
|
||||
"""
|
||||
assert isinstance(scenes, list)
|
||||
assert isinstance(scenes[0], list)
|
||||
B = len(scenes[0])
|
||||
if B == 1:
|
||||
save_meshes_to_obj(out_dir, [x[0] for x in scenes])
|
||||
else:
|
||||
save_scenes_to_glb(out_dir, scenes)
|
||||
|
||||
|
||||
def save_scenes_to_glb(out_dir, scenes):
|
||||
"""
|
||||
Saves a list of scenes (list of meshes) each to glb files
|
||||
"""
|
||||
os.makedirs(out_dir, exist_ok=True)
|
||||
for t, meshes in enumerate(scenes):
|
||||
save_meshes_to_glb(f"{out_dir}/scene_{t:03d}.glb", meshes)
|
||||
|
||||
|
||||
def save_meshes_to_glb(path, meshes, names=None):
|
||||
"""
|
||||
put trimesh meshes in a scene and export to glb
|
||||
"""
|
||||
if names is not None:
|
||||
assert len(meshes) == len(names)
|
||||
|
||||
scene = trimesh.Scene()
|
||||
for i, mesh in enumerate(meshes):
|
||||
name = f"mesh_{i:03d}" if names is None else names[i]
|
||||
scene.add_geometry(mesh, node_name=name)
|
||||
|
||||
with open(path, "wb") as f:
|
||||
f.write(trimesh.exchange.gltf.export_glb(scene, include_normals=True))
|
||||
|
||||
|
||||
def save_meshes_to_obj(out_dir, meshes, names=None):
|
||||
if names is not None:
|
||||
assert len(meshes) == len(names)
|
||||
|
||||
os.makedirs(out_dir, exist_ok=True)
|
||||
for i, mesh in enumerate(meshes):
|
||||
name = f"mesh_{i:03d}" if names is None else names[i]
|
||||
path = os.path.join(out_dir, f"{name}.obj")
|
||||
with open(path, "w") as f:
|
||||
mesh.export(f, file_type="obj", include_color=False, include_normals=True)
|
||||
|
||||
def vertices_to_trimesh(vertices, faces, mesh_base_color=(1.0, 1.0, 0.9),
|
||||
rot_axis=[1,0,0], rot_angle=0, is_right=1):
|
||||
# material = pyrender.MetallicRoughnessMaterial(
|
||||
# metallicFactor=0.0,
|
||||
# alphaMode='OPAQUE',
|
||||
# baseColorFactor=(*mesh_base_color, 1.0))
|
||||
faces_new = np.array([[92, 38, 234],
|
||||
[234, 38, 239],
|
||||
[38, 122, 239],
|
||||
[239, 122, 279],
|
||||
[122, 118, 279],
|
||||
[279, 118, 215],
|
||||
[118, 117, 215],
|
||||
[215, 117, 214],
|
||||
[117, 119, 214],
|
||||
[214, 119, 121],
|
||||
[119, 120, 121],
|
||||
[121, 120, 78],
|
||||
[120, 108, 78],
|
||||
[78, 108, 79]])
|
||||
faces = np.concatenate([faces, faces_new], axis=0)
|
||||
vertex_colors = np.array([(*mesh_base_color, 1.0)] * vertices.shape[0])
|
||||
# print(vertices.shape, self.faces.shape)
|
||||
mesh = trimesh.Trimesh(vertices.copy(), faces.copy(), vertex_colors=vertex_colors)
|
||||
# mesh = trimesh.Trimesh(vertices.copy(), self.faces.copy())
|
||||
|
||||
rot = trimesh.transformations.rotation_matrix(
|
||||
np.radians(rot_angle), rot_axis)
|
||||
mesh.apply_transform(rot)
|
||||
|
||||
rot = trimesh.transformations.rotation_matrix(
|
||||
np.radians(180), [1, 0, 0])
|
||||
mesh.apply_transform(rot)
|
||||
return mesh
|
||||
+60
@@ -0,0 +1,60 @@
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
|
||||
def read_pcl_tensor(path):
|
||||
pcl_np = read_pcl(path)
|
||||
return torch.from_numpy(pcl_np)
|
||||
|
||||
|
||||
def align_pcl(Y, X, weight=None, fixed_scale=False):
|
||||
"""align similarity transform to align X with Y using umeyama method
|
||||
X' = s * R * X + t is aligned with Y
|
||||
:param Y (*, N, 3) first trajectory
|
||||
:param X (*, N, 3) second trajectory
|
||||
:param weight (*, N, 1) optional weight of valid correspondences
|
||||
:returns s (*, 1), R (*, 3, 3), t (*, 3)
|
||||
"""
|
||||
*dims, N, _ = Y.shape
|
||||
N = torch.ones(*dims, 1, 1) * N
|
||||
|
||||
if weight is not None:
|
||||
Y = Y * weight
|
||||
X = X * weight
|
||||
N = weight.sum(dim=-2, keepdim=True) # (*, 1, 1)
|
||||
|
||||
# subtract mean
|
||||
my = Y.sum(dim=-2) / N[..., 0] # (*, 3)
|
||||
mx = X.sum(dim=-2) / N[..., 0]
|
||||
y0 = Y - my[..., None, :] # (*, N, 3)
|
||||
x0 = X - mx[..., None, :]
|
||||
|
||||
if weight is not None:
|
||||
y0 = y0 * weight
|
||||
x0 = x0 * weight
|
||||
|
||||
# correlation
|
||||
C = torch.matmul(y0.transpose(-1, -2), x0) / N # (*, 3, 3)
|
||||
U, D, Vh = torch.linalg.svd(C) # (*, 3, 3), (*, 3), (*, 3, 3)
|
||||
|
||||
S = torch.eye(3).reshape(*(1,) * (len(dims)), 3, 3).repeat(*dims, 1, 1)
|
||||
neg = torch.det(U) * torch.det(Vh.transpose(-1, -2)) < 0
|
||||
S[neg, 2, 2] = -1
|
||||
|
||||
R = torch.matmul(U, torch.matmul(S, Vh)) # (*, 3, 3)
|
||||
|
||||
D = torch.diag_embed(D) # (*, 3, 3)
|
||||
if fixed_scale:
|
||||
s = torch.ones(*dims, 1, device=Y.device, dtype=torch.float32)
|
||||
else:
|
||||
var = torch.sum(torch.square(x0), dim=(-1, -2), keepdim=True) / N # (*, 1, 1)
|
||||
s = (
|
||||
torch.diagonal(torch.matmul(D, S), dim1=-2, dim2=-1).sum(
|
||||
dim=-1, keepdim=True
|
||||
)
|
||||
/ var[..., 0]
|
||||
) # (*, 1)
|
||||
|
||||
t = my - s * torch.matmul(R, mx[..., None])[..., 0] # (*, 3)
|
||||
|
||||
return s, R, t
|
||||
+101
@@ -0,0 +1,101 @@
|
||||
import torch
|
||||
|
||||
|
||||
def fit_plane(points):
|
||||
"""
|
||||
:param points (*, N, 3)
|
||||
returns (*, 3) plane parameters (returns in normal * offset format)
|
||||
"""
|
||||
*dims, N, D = points.shape
|
||||
mean = points.mean(dim=-2, keepdim=True)
|
||||
# (*, N, D), (*, D), (*, D, D)
|
||||
U, S, Vh = torch.linalg.svd(points - mean)
|
||||
normal = Vh[..., -1, :] # (*, D)
|
||||
offset = torch.einsum("...ij,...j->...i", points, normal) # (*, N)
|
||||
offset = offset.mean(dim=-1, keepdim=True)
|
||||
return torch.cat([normal, offset], dim=-1)
|
||||
|
||||
|
||||
def get_plane_transform(up, ground_plane=None, xyz_orig=None):
|
||||
"""
|
||||
get R, t rigid transform from plane and desired origin
|
||||
:param up (3,) up vector of coordinate frame
|
||||
:param ground_plane (4) (a, b, c, d) where a,b,c is the normal
|
||||
:param xyz_orig (3) desired origin
|
||||
"""
|
||||
R = torch.eye(3)
|
||||
t = torch.zeros(3)
|
||||
if ground_plane is None:
|
||||
return R, t
|
||||
|
||||
# compute transform between world up vector and passed in floor
|
||||
ground_plane = torch.as_tensor(ground_plane)
|
||||
ground_plane = torch.sign(ground_plane[3]) * ground_plane
|
||||
|
||||
normal = ground_plane[:3]
|
||||
normal = normal / torch.linalg.norm(normal)
|
||||
v = torch.linalg.cross(up, normal)
|
||||
ang_sin = torch.linalg.norm(v)
|
||||
ang_cos = up.dot(normal)
|
||||
skew_v = torch.as_tensor([[0.0, -v[2], v[1]], [v[2], 0.0, -v[0]], [-v[1], v[0], 0.0]])
|
||||
R = torch.eye(3) + skew_v + (skew_v @ skew_v) * ((1.0 - ang_cos) / (ang_sin**2))
|
||||
|
||||
# project origin onto plane
|
||||
if xyz_orig is None:
|
||||
xyz_orig = torch.zeros(3)
|
||||
t, _ = compute_plane_intersection(xyz_orig, -normal, ground_plane)
|
||||
|
||||
return R, t
|
||||
|
||||
|
||||
def parse_floor_plane(floor_plane):
|
||||
"""
|
||||
Takes floor plane in the optimization form (Bx3 with a,b,c * d) and parses into
|
||||
(a,b,c,d) from with (a,b,c) normal facing "up in the camera frame and d the offset.
|
||||
"""
|
||||
floor_offset = torch.norm(floor_plane, dim=-1, keepdim=True)
|
||||
floor_normal = floor_plane / floor_offset
|
||||
|
||||
# in camera system -y is up, so floor plane normal y component should never be positive
|
||||
# (assuming the camera is not sideways or upside down)
|
||||
neg_mask = floor_normal[..., 1:2] > 0.0
|
||||
floor_normal = torch.where(
|
||||
neg_mask.expand_as(floor_normal), -floor_normal, floor_normal
|
||||
)
|
||||
floor_offset = torch.where(neg_mask, -floor_offset, floor_offset)
|
||||
floor_plane_4d = torch.cat([floor_normal, floor_offset], dim=-1)
|
||||
|
||||
return floor_plane_4d
|
||||
|
||||
|
||||
def compute_plane_intersection(point, direction, plane):
|
||||
"""
|
||||
Given a ray defined by a point in space and a direction,
|
||||
compute the intersection point with the given plane.
|
||||
Detect intersection in either direction or -direction.
|
||||
Note, ray may not actually intersect with the plane.
|
||||
|
||||
Returns the intersection point and s where
|
||||
point + s * direction = intersection_point. if s < 0 it means
|
||||
-direction intersects.
|
||||
|
||||
- point : B x 3
|
||||
- direction : B x 3
|
||||
- plane : B x 4 (a, b, c, d) where (a, b, c) is the normal and (d) the offset.
|
||||
"""
|
||||
dims = point.shape[:-1]
|
||||
plane_normal = plane[..., :3]
|
||||
plane_off = plane[..., 3]
|
||||
s = (plane_off - bdot(plane_normal, point)) / (bdot(plane_normal, direction) + 1e-4)
|
||||
itsct_pt = point + s.reshape((-1, 1)) * direction
|
||||
return itsct_pt, s
|
||||
|
||||
|
||||
def bdot(A1, A2, keepdim=False):
|
||||
"""
|
||||
Batched dot product.
|
||||
- A1 : B x D
|
||||
- A2 : B x D.
|
||||
Returns B.
|
||||
"""
|
||||
return (A1 * A2).sum(dim=-1, keepdim=keepdim)
|
||||
+293
@@ -0,0 +1,293 @@
|
||||
import torch
|
||||
import numpy as np
|
||||
from torch.nn import functional as F
|
||||
|
||||
|
||||
def batch_rodrigues(rot_vecs, epsilon=1e-8, dtype=torch.float32):
|
||||
"""
|
||||
Taken from https://github.com/mkocabas/VIBE/blob/master/lib/utils/geometry.py
|
||||
Calculates the rotation matrices for a batch of rotation vectors
|
||||
- param rot_vecs: torch.tensor (N, 3) array of N axis-angle vectors
|
||||
- returns R: torch.tensor (N, 3, 3) rotation matrices
|
||||
"""
|
||||
batch_size = rot_vecs.shape[0]
|
||||
device = rot_vecs.device
|
||||
|
||||
angle = torch.norm(rot_vecs + 1e-8, dim=1, keepdim=True)
|
||||
rot_dir = rot_vecs / angle
|
||||
|
||||
cos = torch.unsqueeze(torch.cos(angle), dim=1)
|
||||
sin = torch.unsqueeze(torch.sin(angle), dim=1)
|
||||
|
||||
# Bx1 arrays
|
||||
rx, ry, rz = torch.split(rot_dir, 1, dim=1)
|
||||
K = torch.zeros((batch_size, 3, 3), dtype=dtype, device=device)
|
||||
|
||||
zeros = torch.zeros((batch_size, 1), dtype=dtype, device=device)
|
||||
K = torch.cat([zeros, -rz, ry, rz, zeros, -rx, -ry, rx, zeros], dim=1).view(
|
||||
(batch_size, 3, 3)
|
||||
)
|
||||
|
||||
ident = torch.eye(3, dtype=dtype, device=device).unsqueeze(dim=0)
|
||||
rot_mat = ident + sin * K + (1 - cos) * torch.bmm(K, K)
|
||||
return rot_mat
|
||||
|
||||
|
||||
def quaternion_mul(q0, q1):
|
||||
"""
|
||||
EXPECTS WXYZ
|
||||
:param q0 (*, 4)
|
||||
:param q1 (*, 4)
|
||||
"""
|
||||
r0, r1 = q0[..., :1], q1[..., :1]
|
||||
v0, v1 = q0[..., 1:], q1[..., 1:]
|
||||
r = r0 * r1 - (v0 * v1).sum(dim=-1, keepdim=True)
|
||||
v = r0 * v1 + r1 * v0 + torch.linalg.cross(v0, v1)
|
||||
return torch.cat([r, v], dim=-1)
|
||||
|
||||
|
||||
def quaternion_inverse(q, eps=1e-8):
|
||||
"""
|
||||
EXPECTS WXYZ
|
||||
:param q (*, 4)
|
||||
"""
|
||||
conj = torch.cat([q[..., :1], -q[..., 1:]], dim=-1)
|
||||
mag = torch.square(q).sum(dim=-1, keepdim=True) + eps
|
||||
return conj / mag
|
||||
|
||||
|
||||
def quaternion_slerp(t, q0, q1, eps=1e-8):
|
||||
"""
|
||||
:param t (*, 1) must be between 0 and 1
|
||||
:param q0 (*, 4)
|
||||
:param q1 (*, 4)
|
||||
"""
|
||||
dims = q0.shape[:-1]
|
||||
t = t.view(*dims, 1)
|
||||
|
||||
q0 = F.normalize(q0, p=2, dim=-1)
|
||||
q1 = F.normalize(q1, p=2, dim=-1)
|
||||
dot = (q0 * q1).sum(dim=-1, keepdim=True)
|
||||
|
||||
# make sure we give the shortest rotation path (< 180d)
|
||||
neg = dot < 0
|
||||
q1 = torch.where(neg, -q1, q1)
|
||||
dot = torch.where(neg, -dot, dot)
|
||||
angle = torch.acos(dot)
|
||||
|
||||
# if angle is too small, just do linear interpolation
|
||||
collin = torch.abs(dot) > 1 - eps
|
||||
fac = 1 / torch.sin(angle)
|
||||
w0 = torch.where(collin, 1 - t, torch.sin((1 - t) * angle) * fac)
|
||||
w1 = torch.where(collin, t, torch.sin(t * angle) * fac)
|
||||
slerp = q0 * w0 + q1 * w1
|
||||
return slerp
|
||||
|
||||
|
||||
def rotation_matrix_to_angle_axis(rotation_matrix):
|
||||
"""
|
||||
This function is borrowed from https://github.com/kornia/kornia
|
||||
|
||||
Convert rotation matrix to Rodrigues vector
|
||||
"""
|
||||
quaternion = rotation_matrix_to_quaternion(rotation_matrix)
|
||||
aa = quaternion_to_angle_axis(quaternion)
|
||||
aa[torch.isnan(aa)] = 0.0
|
||||
return aa
|
||||
|
||||
|
||||
def quaternion_to_angle_axis(quaternion):
|
||||
"""
|
||||
This function is borrowed from https://github.com/kornia/kornia
|
||||
|
||||
Convert quaternion vector to angle axis of rotation.
|
||||
Adapted from ceres C++ library: ceres-solver/include/ceres/rotation.h
|
||||
|
||||
:param quaternion (*, 4) expects WXYZ
|
||||
:returns angle_axis (*, 3)
|
||||
"""
|
||||
# unpack input and compute conversion
|
||||
q1 = quaternion[..., 1]
|
||||
q2 = quaternion[..., 2]
|
||||
q3 = quaternion[..., 3]
|
||||
sin_squared_theta = q1 * q1 + q2 * q2 + q3 * q3
|
||||
|
||||
sin_theta = torch.sqrt(sin_squared_theta)
|
||||
cos_theta = quaternion[..., 0]
|
||||
two_theta = 2.0 * torch.where(
|
||||
cos_theta < 0.0,
|
||||
torch.atan2(-sin_theta, -cos_theta),
|
||||
torch.atan2(sin_theta, cos_theta),
|
||||
)
|
||||
|
||||
k_pos = two_theta / sin_theta
|
||||
k_neg = 2.0 * torch.ones_like(sin_theta)
|
||||
k = torch.where(sin_squared_theta > 0.0, k_pos, k_neg)
|
||||
|
||||
angle_axis = torch.zeros_like(quaternion)[..., :3]
|
||||
angle_axis[..., 0] += q1 * k
|
||||
angle_axis[..., 1] += q2 * k
|
||||
angle_axis[..., 2] += q3 * k
|
||||
return angle_axis
|
||||
|
||||
|
||||
def angle_axis_to_rotation_matrix(angle_axis):
|
||||
"""
|
||||
:param angle_axis (*, 3)
|
||||
return (*, 3, 3)
|
||||
"""
|
||||
quat = angle_axis_to_quaternion(angle_axis)
|
||||
return quaternion_to_rotation_matrix(quat)
|
||||
|
||||
|
||||
def quaternion_to_rotation_matrix(quaternion):
|
||||
"""
|
||||
Convert a quaternion to a rotation matrix.
|
||||
Taken from https://github.com/kornia/kornia, based on
|
||||
https://github.com/matthew-brett/transforms3d/blob/8965c48401d9e8e66b6a8c37c65f2fc200a076fa/transforms3d/quaternions.py#L101
|
||||
https://github.com/tensorflow/graphics/blob/master/tensorflow_graphics/geometry/transformation/rotation_matrix_3d.py#L247
|
||||
:param quaternion (N, 4) expects WXYZ order
|
||||
returns rotation matrix (N, 3, 3)
|
||||
"""
|
||||
# normalize the input quaternion
|
||||
quaternion_norm = F.normalize(quaternion, p=2, dim=-1, eps=1e-12)
|
||||
*dims, _ = quaternion_norm.shape
|
||||
|
||||
# unpack the normalized quaternion components
|
||||
w, x, y, z = torch.chunk(quaternion_norm, chunks=4, dim=-1)
|
||||
|
||||
# compute the actual conversion
|
||||
tx = 2.0 * x
|
||||
ty = 2.0 * y
|
||||
tz = 2.0 * z
|
||||
twx = tx * w
|
||||
twy = ty * w
|
||||
twz = tz * w
|
||||
txx = tx * x
|
||||
txy = ty * x
|
||||
txz = tz * x
|
||||
tyy = ty * y
|
||||
tyz = tz * y
|
||||
tzz = tz * z
|
||||
one = torch.tensor(1.0)
|
||||
|
||||
matrix = torch.stack(
|
||||
(
|
||||
one - (tyy + tzz),
|
||||
txy - twz,
|
||||
txz + twy,
|
||||
txy + twz,
|
||||
one - (txx + tzz),
|
||||
tyz - twx,
|
||||
txz - twy,
|
||||
tyz + twx,
|
||||
one - (txx + tyy),
|
||||
),
|
||||
dim=-1,
|
||||
).view(*dims, 3, 3)
|
||||
return matrix
|
||||
|
||||
|
||||
def angle_axis_to_quaternion(angle_axis):
|
||||
"""
|
||||
This function is borrowed from https://github.com/kornia/kornia
|
||||
Convert angle axis to quaternion in WXYZ order
|
||||
:param angle_axis (*, 3)
|
||||
:returns quaternion (*, 4) WXYZ order
|
||||
"""
|
||||
theta_sq = torch.sum(angle_axis**2, dim=-1, keepdim=True) # (*, 1)
|
||||
# need to handle the zero rotation case
|
||||
valid = theta_sq > 0
|
||||
theta = torch.sqrt(theta_sq)
|
||||
half_theta = 0.5 * theta
|
||||
ones = torch.ones_like(half_theta)
|
||||
# fill zero with the limit of sin ax / x -> a
|
||||
k = torch.where(valid, torch.sin(half_theta) / theta, 0.5 * ones)
|
||||
w = torch.where(valid, torch.cos(half_theta), ones)
|
||||
quat = torch.cat([w, k * angle_axis], dim=-1)
|
||||
return quat
|
||||
|
||||
|
||||
def rotation_matrix_to_quaternion(rotation_matrix, eps=1e-6):
|
||||
"""
|
||||
This function is borrowed from https://github.com/kornia/kornia
|
||||
Convert rotation matrix to 4d quaternion vector
|
||||
This algorithm is based on algorithm described in
|
||||
https://github.com/KieranWynn/pyquaternion/blob/master/pyquaternion/quaternion.py#L201
|
||||
|
||||
:param rotation_matrix (N, 3, 3)
|
||||
"""
|
||||
*dims, m, n = rotation_matrix.shape
|
||||
rmat_t = torch.transpose(rotation_matrix.reshape(-1, m, n), -1, -2)
|
||||
|
||||
mask_d2 = rmat_t[:, 2, 2] < eps
|
||||
|
||||
mask_d0_d1 = rmat_t[:, 0, 0] > rmat_t[:, 1, 1]
|
||||
mask_d0_nd1 = rmat_t[:, 0, 0] < -rmat_t[:, 1, 1]
|
||||
|
||||
t0 = 1 + rmat_t[:, 0, 0] - rmat_t[:, 1, 1] - rmat_t[:, 2, 2]
|
||||
q0 = torch.stack(
|
||||
[
|
||||
rmat_t[:, 1, 2] - rmat_t[:, 2, 1],
|
||||
t0,
|
||||
rmat_t[:, 0, 1] + rmat_t[:, 1, 0],
|
||||
rmat_t[:, 2, 0] + rmat_t[:, 0, 2],
|
||||
],
|
||||
-1,
|
||||
)
|
||||
t0_rep = t0.repeat(4, 1).t()
|
||||
|
||||
t1 = 1 - rmat_t[:, 0, 0] + rmat_t[:, 1, 1] - rmat_t[:, 2, 2]
|
||||
q1 = torch.stack(
|
||||
[
|
||||
rmat_t[:, 2, 0] - rmat_t[:, 0, 2],
|
||||
rmat_t[:, 0, 1] + rmat_t[:, 1, 0],
|
||||
t1,
|
||||
rmat_t[:, 1, 2] + rmat_t[:, 2, 1],
|
||||
],
|
||||
-1,
|
||||
)
|
||||
t1_rep = t1.repeat(4, 1).t()
|
||||
|
||||
t2 = 1 - rmat_t[:, 0, 0] - rmat_t[:, 1, 1] + rmat_t[:, 2, 2]
|
||||
q2 = torch.stack(
|
||||
[
|
||||
rmat_t[:, 0, 1] - rmat_t[:, 1, 0],
|
||||
rmat_t[:, 2, 0] + rmat_t[:, 0, 2],
|
||||
rmat_t[:, 1, 2] + rmat_t[:, 2, 1],
|
||||
t2,
|
||||
],
|
||||
-1,
|
||||
)
|
||||
t2_rep = t2.repeat(4, 1).t()
|
||||
|
||||
t3 = 1 + rmat_t[:, 0, 0] + rmat_t[:, 1, 1] + rmat_t[:, 2, 2]
|
||||
q3 = torch.stack(
|
||||
[
|
||||
t3,
|
||||
rmat_t[:, 1, 2] - rmat_t[:, 2, 1],
|
||||
rmat_t[:, 2, 0] - rmat_t[:, 0, 2],
|
||||
rmat_t[:, 0, 1] - rmat_t[:, 1, 0],
|
||||
],
|
||||
-1,
|
||||
)
|
||||
t3_rep = t3.repeat(4, 1).t()
|
||||
|
||||
mask_c0 = mask_d2 * mask_d0_d1
|
||||
mask_c1 = mask_d2 * ~mask_d0_d1
|
||||
mask_c2 = ~mask_d2 * mask_d0_nd1
|
||||
mask_c3 = ~mask_d2 * ~mask_d0_nd1
|
||||
mask_c0 = mask_c0.view(-1, 1).type_as(q0)
|
||||
mask_c1 = mask_c1.view(-1, 1).type_as(q1)
|
||||
mask_c2 = mask_c2.view(-1, 1).type_as(q2)
|
||||
mask_c3 = mask_c3.view(-1, 1).type_as(q3)
|
||||
|
||||
q = q0 * mask_c0 + q1 * mask_c1 + q2 * mask_c2 + q3 * mask_c3
|
||||
q /= torch.sqrt(
|
||||
t0_rep * mask_c0
|
||||
+ t1_rep * mask_c1
|
||||
+ t2_rep * mask_c2 # noqa
|
||||
+ t3_rep * mask_c3
|
||||
) # noqa
|
||||
q *= 0.5
|
||||
return q.reshape(*dims, 4)
|
||||
+309
@@ -0,0 +1,309 @@
|
||||
import os
|
||||
import numpy as np
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from body_model import MANO_JOINTS, run_mano
|
||||
from geometry.rotation import (
|
||||
rotation_matrix_to_angle_axis,
|
||||
angle_axis_to_rotation_matrix,
|
||||
)
|
||||
from util.logger import Logger
|
||||
from util.tensor import move_to, detach_all
|
||||
|
||||
# from .helpers import estimate_initial_trans
|
||||
from .params import CameraParams
|
||||
|
||||
J_HAND = len(MANO_JOINTS) - 1 # no root
|
||||
|
||||
|
||||
class BaseSceneModel(nn.Module):
|
||||
"""
|
||||
Scene model of sequences of human poses.
|
||||
All poses are in their own INDEPENDENT camera reference frames.
|
||||
A basic class mostly for testing purposes.
|
||||
|
||||
Parameters:
|
||||
batch_size: number of sequences to optimize
|
||||
seq_len: length of the sequences
|
||||
body_model: MANO hand model
|
||||
pose_prior: VPoser model
|
||||
fit_gender: gender of model (optional)
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
batch_size,
|
||||
seq_len,
|
||||
body_model,
|
||||
pose_prior,
|
||||
# fit_gender="neutral",
|
||||
use_init=False,
|
||||
opt_cams=False,
|
||||
opt_scale=True,
|
||||
**kwargs,
|
||||
):
|
||||
super().__init__()
|
||||
B, T = batch_size, seq_len
|
||||
self.batch_size = batch_size
|
||||
self.seq_len = seq_len
|
||||
|
||||
self.body_model = body_model
|
||||
self.hand_mean = self.body_model.hand_mean
|
||||
|
||||
self.pose_prior = pose_prior
|
||||
if self.pose_prior is not None:
|
||||
self.latent_pose_dim = self.pose_prior.latentD
|
||||
|
||||
self.num_betas = body_model.num_betas
|
||||
|
||||
self.use_init = use_init
|
||||
print("USE INIT", use_init)
|
||||
self.opt_scale = opt_scale
|
||||
self.opt_cams = opt_cams
|
||||
print("OPT SCALE", self.opt_scale)
|
||||
print("OPT CAMERAS", self.opt_cams)
|
||||
print("Batch size: ", batch_size)
|
||||
self.params = CameraParams(batch_size)
|
||||
|
||||
def initialize(self, obs_data, cam_data):
|
||||
Logger.log("Initializing scene model with observed data")
|
||||
|
||||
# initialize cameras
|
||||
self.params.set_cameras(
|
||||
cam_data,
|
||||
opt_scale=self.opt_scale,
|
||||
opt_cams=self.opt_cams,
|
||||
opt_focal=self.opt_cams,
|
||||
)
|
||||
|
||||
# initialize body params
|
||||
B, T = self.batch_size, self.seq_len
|
||||
device = next(iter(cam_data.values())).device
|
||||
init_betas = torch.mean(obs_data["init_body_shape"], dim=1) # torch.zeros(B, self.num_betas, device=device)
|
||||
|
||||
if self.use_init and "init_body_pose" in obs_data:
|
||||
init_pose = obs_data["init_body_pose"][:, :, :J_HAND, :]
|
||||
init_pose_latent = self.pose2latent(init_pose)
|
||||
else:
|
||||
raise ValueError
|
||||
# init_pose = torch.zeros(B, T, J_HAND, 3, device=device)
|
||||
# init_pose_latent = torch.zeros(B, T, self.latent_pose_dim, device=device)
|
||||
|
||||
# transform into world frame (T, 3, 3), (T, 3)
|
||||
R_w2c, t_w2c = cam_data["cam_R"], cam_data["cam_t"]
|
||||
R_c2w = R_w2c.transpose(-1, -2)
|
||||
t_c2w = -torch.einsum("tij,tj->ti", R_c2w, t_w2c)
|
||||
|
||||
if self.use_init and "init_root_orient" in obs_data:
|
||||
init_rot = obs_data["init_root_orient"] # (B, T, 3)
|
||||
init_rot_mat = angle_axis_to_rotation_matrix(init_rot)
|
||||
print(R_c2w.shape, init_rot_mat.shape)
|
||||
init_rot_mat = torch.einsum("tij,btjk->btik", R_c2w, init_rot_mat)
|
||||
init_rot = rotation_matrix_to_angle_axis(init_rot_mat)
|
||||
else:
|
||||
raise ValueError
|
||||
# init_rot = (
|
||||
# torch.tensor([np.pi, 0, 0], dtype=torch.float32)
|
||||
# .reshape(1, 1, 3)
|
||||
# .repeat(B, T, 1)
|
||||
# )
|
||||
|
||||
init_trans = torch.zeros(B, T, 3, device=device)
|
||||
if self.use_init and "init_trans" in obs_data:
|
||||
# must offset by the root location before applying camera to world transform
|
||||
is_right = obs_data["is_right"]
|
||||
pred_data = self.pred_mano(init_trans, init_rot, init_pose, is_right, init_betas)
|
||||
root_loc = pred_data["joints3d"][..., 0, :] # (B, T, 3)
|
||||
init_trans = obs_data["init_trans"] # (B, T, 3)
|
||||
init_trans = (
|
||||
torch.einsum("tij,btj->bti", R_c2w, init_trans + root_loc)
|
||||
+ t_c2w[None]
|
||||
- root_loc
|
||||
)
|
||||
else:
|
||||
# initialize trans with reprojected joints
|
||||
# pred_data = self.pred_mano(init_trans, init_rot, init_pose, init_betas)
|
||||
# init_trans = estimate_initial_trans(
|
||||
# init_pose,
|
||||
# pred_data["joints3d_op"],
|
||||
# obs_data["joints2d"],
|
||||
# obs_data["intrins"][:, 0],
|
||||
# )
|
||||
raise ValueError
|
||||
|
||||
self.params.set_param("init_body_pose", init_pose)
|
||||
self.params.set_param("latent_pose", init_pose_latent)
|
||||
self.params.set_param("betas", init_betas)
|
||||
self.params.set_param("trans", init_trans)
|
||||
self.params.set_param("root_orient", init_rot)
|
||||
self.params.set_param("is_right", is_right, requires_grad=False)
|
||||
|
||||
# Store initial latent pose in obs_data for pose prior loss
|
||||
obs_data["init_latent_pose"] = init_pose_latent.detach()
|
||||
# print(init_pose.shape, init_pose_latent.shape, init_betas.shape, init_trans.shape, init_rot.shape, is_right.shape)
|
||||
# raise ValueError
|
||||
|
||||
def get_optim_result(self, **kwargs):
|
||||
"""
|
||||
Collect predicted outputs (latent_pose, trans, root_orient, betas, body pose) into dict
|
||||
"""
|
||||
print('running get_optim_result in base_scene.py...')
|
||||
res = self.params.get_dict()
|
||||
if "latent_pose" in res:
|
||||
res["pose_body"] = self.latent2pose(self.params.latent_pose).detach()
|
||||
|
||||
# add the cameras
|
||||
res["cam_R"], res["cam_t"], _, _ = self.params.get_cameras()
|
||||
res["intrins"] = self.params.intrins
|
||||
return {"world": res}
|
||||
|
||||
def latent2pose(self, latent_pose):
|
||||
"""
|
||||
Converts VPoser latent embedding to aa body pose.
|
||||
latent_pose : B x T x D
|
||||
body_pose : B x T x J*3
|
||||
"""
|
||||
if self.pose_prior is not None:
|
||||
B, T, _ = latent_pose.size()
|
||||
d_latent = self.pose_prior.latentD
|
||||
latent_pose = latent_pose.reshape((-1, d_latent))
|
||||
body_pose = self.pose_prior.decode(latent_pose, output_type="matrot")
|
||||
body_pose = rotation_matrix_to_angle_axis(
|
||||
body_pose.reshape((B * T * J_HAND, 3, 3))
|
||||
).reshape((B, T, J_HAND * 3))
|
||||
return body_pose + self.hand_mean
|
||||
else:
|
||||
return latent_pose
|
||||
|
||||
def pose2latent(self, body_pose):
|
||||
"""
|
||||
Encodes aa body pose to VPoser latent space.
|
||||
body_pose : B x T x J*3
|
||||
latent_pose : B x T x D
|
||||
"""
|
||||
if self.pose_prior is not None:
|
||||
B, T = body_pose.shape[:2]
|
||||
body_pose = body_pose.reshape((-1, J_HAND * 3))
|
||||
latent_pose_distrib = self.pose_prior.encode(body_pose - self.hand_mean)
|
||||
d_latent = self.pose_prior.latentD
|
||||
latent_pose = latent_pose_distrib.mean.reshape((B, T, d_latent))
|
||||
return latent_pose
|
||||
else:
|
||||
return body_pose
|
||||
|
||||
def pred_mano(self, trans, root_orient, body_pose, is_right, betas):
|
||||
"""
|
||||
Forward pass of the MANO model and populates pred_data accordingly with
|
||||
joints3d, verts3d, points3d.
|
||||
|
||||
trans : B x T x 3
|
||||
root_orient : B x T x 3
|
||||
body_pose : B x T x J*3
|
||||
betas : B x D
|
||||
"""
|
||||
|
||||
mano_out = run_mano(self.body_model, trans, root_orient, body_pose, is_right, betas=betas)
|
||||
joints3d, points3d = mano_out["joints"], mano_out["vertices"]
|
||||
|
||||
# select desired joints and vertices
|
||||
# joints3d_body = joints3d
|
||||
# joints3d_op = joints3d.clone()
|
||||
# joints3d_op = joints3d[:, :, self.smpl2op_map, :]
|
||||
# # hacky way to get hip joints that align with ViTPose keypoints
|
||||
# # this could be moved elsewhere in the future (and done properly)
|
||||
# joints3d_op[:, :, [9, 12]] = (
|
||||
# joints3d_op[:, :, [9, 12]]
|
||||
# + 0.25 * (joints3d_op[:, :, [9, 12]] - joints3d_op[:, :, [12, 9]])
|
||||
# + 0.5
|
||||
# * (
|
||||
# joints3d_op[:, :, [8]]
|
||||
# - 0.5 * (joints3d_op[:, :, [9, 12]] + joints3d_op[:, :, [12, 9]])
|
||||
# )
|
||||
# )
|
||||
verts3d = points3d #[:, :, KEYPT_VERTS, :]
|
||||
|
||||
return {
|
||||
"is_right": is_right,
|
||||
"points3d": points3d, # all vertices
|
||||
"verts3d": verts3d, # keypoint vertices
|
||||
"joints3d": joints3d, # smpl joints
|
||||
"joints3d_op": joints3d, # OP joints
|
||||
"l_faces": mano_out["l_faces"], # index array of faces
|
||||
"r_faces": mano_out["r_faces"], # index array of faces
|
||||
"body_pose": mano_out["body_pose"]
|
||||
}
|
||||
|
||||
def pred_params_mano(self, is_right, reproj=True):
|
||||
body_pose = self.latent2pose(self.params.latent_pose)
|
||||
pred_data = self.pred_mano(
|
||||
self.params.trans, self.params.root_orient, body_pose, is_right, self.params.betas
|
||||
)
|
||||
# from geometry.mesh import save_mesh_scenes, vertices_to_trimesh
|
||||
# LIGHT_BLUE=(0.65098039, 0.74117647, 0.85882353)
|
||||
# print(pred_data.keys())
|
||||
# verts, l_faces, r_faces, is_right = pred_data['verts3d'], pred_data['l_faces'], pred_data['l_faces'], pred_data['is_right']
|
||||
# T = len(verts[0])
|
||||
# print(f"{T} mesh frames")
|
||||
# print(verts.shape, l_faces.shape, is_right.shape, is_right)
|
||||
# times = list(range(0, T, 1))
|
||||
# os.makedirs('./mesh/')
|
||||
# for t in times:
|
||||
|
||||
# tmesh = vertices_to_trimesh(verts[0][t].detach().cpu().numpy(), l_faces.detach().cpu().numpy(), LIGHT_BLUE, is_right=0)
|
||||
# tmesh.export(os.path.join('./mesh/', f'{str(t).zfill(6)}_0.obj'))
|
||||
|
||||
# tmesh = vertices_to_trimesh(verts[1][t].detach().cpu().numpy(), r_faces.detach().cpu().numpy(), LIGHT_BLUE, is_right=1)
|
||||
# tmesh.export(os.path.join('./mesh/', f'{str(t).zfill(6)}_1.obj'))
|
||||
# print(os.path.join('./mesh/', f'{str(t).zfill(6)}_1.obj'))
|
||||
|
||||
#################################
|
||||
# print(self.params.root_orient.shape, self.params.init_body_pose.shape, body_pose.shape, self.params.trans.shape, self.params.betas.shape)
|
||||
# hand_model = mano.load(
|
||||
# model_path= '/vol/bitbucket/zy3023/code/hand/slam-hand/_DATA/data/mano',# '/vol/bitbucket/zy3023/code/hand/slam-hand/VPoser/pretrained/Vposer_right_mirrored',
|
||||
# is_right=True,
|
||||
# num_pca_comps=45,
|
||||
# batch_size=1,
|
||||
# flat_hand_mean=False,
|
||||
# ).cuda()
|
||||
|
||||
# output = hand_model(
|
||||
# betas=self.params.betas[0][None].repeat(180,1),
|
||||
# global_orient=self.params.root_orient[0],
|
||||
# hand_pose=self.params.init_body_pose[0].reshape(180, -1) - self.hand_mean,
|
||||
# transl=self.params.trans[0],
|
||||
# return_verts=True,
|
||||
# return_tips=False
|
||||
# )
|
||||
# np.save('x.npy', self.params.init_body_pose[0].reshape(180, -1).detach().cpu().numpy())
|
||||
|
||||
# out_path = './out'
|
||||
# h_meshes = hand_model.hand_meshes(output)
|
||||
# if not os.path.exists(out_path):
|
||||
# os.mkdir(out_path + '_after')
|
||||
# os.mkdir(out_path)
|
||||
|
||||
# for i in range(180):
|
||||
# #print(f"export to {out_path}/{str(i).zfill(6)}.ply")
|
||||
# h_meshes[i].export(f"{out_path}/{str(i).zfill(6)}.ply")
|
||||
|
||||
# # print(body_pose.shape, body_pose[0])
|
||||
# print(self.hand_mean)
|
||||
# output = hand_model(
|
||||
# betas=self.params.betas[0][None].repeat(180,1),
|
||||
# global_orient=self.params.root_orient[0],
|
||||
# hand_pose=body_pose[0].reshape(180, -1) - self.hand_mean,
|
||||
# transl=self.params.trans[0],
|
||||
# return_verts=True,
|
||||
# return_tips=False
|
||||
# )
|
||||
|
||||
# h_meshes = hand_model.hand_meshes(output)
|
||||
# for i in range(180):
|
||||
# #print(f"export to {out_path}_after/{str(i).zfill(6)}.ply")
|
||||
# h_meshes[i].export(f"{out_path}_after/{str(i).zfill(6)}.ply")
|
||||
# exit()
|
||||
|
||||
return pred_data
|
||||
+399
@@ -0,0 +1,399 @@
|
||||
import os
|
||||
|
||||
import matplotlib.pyplot as plt
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn.functional as torch_f
|
||||
|
||||
SNAP_PARENT = [
|
||||
0, # 0's parent
|
||||
0, # 1's parent
|
||||
1,
|
||||
2,
|
||||
3,
|
||||
0, # 5's parent
|
||||
5,
|
||||
6,
|
||||
7,
|
||||
0, # 9's parent
|
||||
9,
|
||||
10,
|
||||
11,
|
||||
0, # 13's parent
|
||||
13,
|
||||
14,
|
||||
15,
|
||||
0, # 17's parent
|
||||
17,
|
||||
18,
|
||||
19,
|
||||
]
|
||||
|
||||
JOINT_ROOT_IDX = 9
|
||||
|
||||
REF_BONE_LINK = (0, 9) # mid mcp
|
||||
|
||||
# bone indexes in 20 bones setting
|
||||
ID_ROOT_bone = np.array([0, 4, 8, 12, 16]) # ROOT_bone from wrist to MCP
|
||||
ID_PIP_bone = np.array([1, 5, 9, 13, 17]) # PIP_bone from MCP to PIP
|
||||
ID_DIP_bone = np.array([2, 6, 10, 14, 18]) # DIP_bone from PIP to DIP
|
||||
ID_TIP_bone = np.array([3, 7, 11, 15, 19]) # TIP_bone from DIP to TIP
|
||||
|
||||
def plot_hull(theta, hull):
|
||||
del_rdp_hull = hull.detach().cpu().numpy()
|
||||
theta = theta.detach().cpu().numpy()
|
||||
|
||||
fig = plt.figure()
|
||||
figManager = plt.get_current_fig_manager()
|
||||
figManager.window.showMaximized()
|
||||
|
||||
ax = fig.add_subplot(111)
|
||||
ax.scatter(theta[:, 0], theta[:, 1], s=10, c='r')
|
||||
ax.set_xlabel("flexion")
|
||||
ax.set_ylabel("abduction")
|
||||
|
||||
plt.plot(del_rdp_hull[:, 0],
|
||||
del_rdp_hull[:, 1],
|
||||
'-yo', linewidth=2)
|
||||
|
||||
plt.xticks(np.arange(-3, 4, 1))
|
||||
plt.yticks(np.arange(-2, 2, 0.5))
|
||||
plt.show()
|
||||
|
||||
|
||||
def two_norm(a):
|
||||
'''
|
||||
|
||||
Args:
|
||||
a: B*M*2 or B*M*3
|
||||
|
||||
Returns:
|
||||
|
||||
'''
|
||||
return torch.norm(a, dim=-1)
|
||||
|
||||
|
||||
def one_norm(a):
|
||||
'''
|
||||
|
||||
Args:
|
||||
a: B*M*2 or B*M*3
|
||||
|
||||
Returns:
|
||||
|
||||
'''
|
||||
return torch.norm(a, dim=-1, p=1)
|
||||
|
||||
|
||||
def calculate_joint_angle_loss(thetas, hulls):
|
||||
'''
|
||||
|
||||
Args:
|
||||
Theta: B*15*2
|
||||
hulls: list
|
||||
|
||||
Returns:
|
||||
|
||||
'''
|
||||
|
||||
loss = torch.Tensor([0]).cuda()
|
||||
for i in range(15):
|
||||
# print("i=",i)
|
||||
hull = hulls[i] # (M*2)
|
||||
theta = thetas[:, i] # (B*2)
|
||||
hull = torch.cat((hull, hull[0].unsqueeze(0)), dim=0)
|
||||
|
||||
v = (hull[1:] - hull[:-1]).unsqueeze(0) # (M-1)*2
|
||||
w = - hull[:-1].unsqueeze(0) + theta.unsqueeze(1).repeat(1, hull[:-1].shape[0], 1) # B*(M-1)*2
|
||||
|
||||
cross_product_2d = w[:, :, 0] * v[:, :, 1] - w[:, :, 1] * v[:, :, 0]
|
||||
tmp = torch.sum(cross_product_2d < 1e-6, dim=-1)
|
||||
|
||||
is_outside = (tmp != (hull.shape[0] - 1))
|
||||
if not torch.sum(is_outside):
|
||||
sub_loss = torch.Tensor([0]).cuda()
|
||||
else:
|
||||
outside_theta = theta[is_outside]
|
||||
outside_theta = outside_theta.unsqueeze(1).repeat(1, hull[:-1].shape[0], 1)
|
||||
w_outside = - hull[:-1].unsqueeze(0) + outside_theta # B*(M-1)*2
|
||||
t = torch.clamp(inner_product(w_outside, v) / (two_norm(v) ** 2), min=0, max=1).unsqueeze(2)
|
||||
p = hull[:-1] + t * v
|
||||
|
||||
D = one_norm(torch.cos(outside_theta) - torch.cos(p)) + one_norm(torch.sin(outside_theta) - torch.sin(p))
|
||||
sub_loss = torch.sum(torch.min(D, dim=-1)[0])
|
||||
|
||||
vis = 0
|
||||
if vis:
|
||||
print(theta)
|
||||
plot_hull(theta, hull)
|
||||
|
||||
loss += sub_loss
|
||||
|
||||
loss /= (15 * thetas.shape[0])
|
||||
|
||||
return loss
|
||||
|
||||
|
||||
def angle_between(v1, v2):
|
||||
epsilon = 1e-7
|
||||
cos = torch_f.cosine_similarity(v1, v2, dim=-1).clamp(-1 + epsilon, 1 - epsilon) # (B)
|
||||
theta = torch.acos(cos) # (B)
|
||||
return theta
|
||||
|
||||
|
||||
def normalize(vec):
|
||||
return torch_f.normalize(vec, p=2, dim=-1)
|
||||
|
||||
|
||||
def inner_product(x1, x2):
|
||||
return torch.sum(x1 * x2, dim=-1)
|
||||
|
||||
|
||||
def cross_product(x1, x2):
|
||||
return torch.cross(x1, x2, dim=-1)
|
||||
|
||||
|
||||
def axangle2mat_torch(axis, angle, is_normalized=False):
|
||||
""" Rotation matrix for rotation angle `angle` around `axis`
|
||||
Parameters
|
||||
----------
|
||||
axis : [B,M, 3] element sequence
|
||||
vector specifying axis for rotation.
|
||||
angle :[B,M, ] scalar
|
||||
angle of rotation in radians.
|
||||
is_normalized : bool, optional
|
||||
True if `axis` is already normalized (has norm of 1). Default False.
|
||||
Returns
|
||||
-------
|
||||
mat : array shape (B, M,3,3)
|
||||
rotation matrix for specified rotation
|
||||
Notes
|
||||
-----
|
||||
From: http://en.wikipedia.org/wiki/Rotation_matrix#Axis_and_angle
|
||||
"""
|
||||
B = axis.shape[0]
|
||||
M = axis.shape[1]
|
||||
|
||||
if not is_normalized:
|
||||
norm_axis = axis.norm(p=2, dim=-1, keepdim=True)
|
||||
normed_axis = axis / norm_axis
|
||||
else:
|
||||
normed_axis = axis
|
||||
x, y, z = normed_axis[:, :, 0], normed_axis[:, :, 1], normed_axis[:, :, 2]
|
||||
c = torch.cos(angle)
|
||||
s = torch.sin(angle)
|
||||
C = 1 - c
|
||||
|
||||
xs = x * s;
|
||||
ys = y * s;
|
||||
zs = z * s # noqa
|
||||
xC = x * C;
|
||||
yC = y * C;
|
||||
zC = z * C # noqa
|
||||
xyC = x * yC;
|
||||
yzC = y * zC;
|
||||
zxC = z * xC # noqa
|
||||
|
||||
TMP = torch.stack([x * xC + c, xyC - zs, zxC + ys, xyC + zs, y * yC + c, yzC - xs, zxC - ys, yzC + xs, z * zC + c],
|
||||
dim=-1)
|
||||
return TMP.reshape(B, M, 3, 3)
|
||||
|
||||
|
||||
def interval_loss(value, min, max):
|
||||
'''
|
||||
calculate interval loss
|
||||
Args:
|
||||
value: B*M
|
||||
max: M
|
||||
min: M
|
||||
|
||||
Returns:
|
||||
|
||||
'''
|
||||
|
||||
batch_3d_size = value.shape[0]
|
||||
|
||||
min = min.repeat(value.shape[0], 1)
|
||||
max = max.repeat(value.shape[0], 1)
|
||||
|
||||
loss1 = torch.max(min - value, torch.Tensor([0]).cuda())
|
||||
loss2 = torch.max(value - max, torch.Tensor([0]).cuda())
|
||||
|
||||
loss = (loss1 + loss2).sum()
|
||||
|
||||
loss /= (batch_3d_size * value.shape[1])
|
||||
|
||||
return loss
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
class BMCLoss:
|
||||
def __init__(
|
||||
self,
|
||||
lambda_bl=0.,
|
||||
lambda_rb=0.,
|
||||
lambda_a=0.,
|
||||
):
|
||||
self.lambda_bl = lambda_bl
|
||||
self.lambda_rb = lambda_rb
|
||||
self.lambda_a = lambda_a
|
||||
|
||||
self.lp = os.path.join(os.path.dirname(__file__), "../../_DATA/BMC")
|
||||
# resolved_path = os.path.abspath(self.lp)
|
||||
# print(f"Resolved path: {resolved_path}")
|
||||
# print(f"Current working directory: {os.getcwd()}")
|
||||
# exit()
|
||||
|
||||
self.bone_len_max = np.load(os.path.join(self.lp, "bone_len_max.npy"))
|
||||
self.bone_len_min = np.load(os.path.join(self.lp, "bone_len_min.npy"))
|
||||
self.rb_curvatures_max = np.load(os.path.join(self.lp, "curvatures_max.npy"))
|
||||
self.rb_curvatures_min = np.load(os.path.join(self.lp, "curvatures_min.npy"))
|
||||
self.rb_PHI_max = np.load(os.path.join(self.lp, "PHI_max.npy"))
|
||||
self.rb_PHI_min = np.load(os.path.join(self.lp, "PHI_min.npy"))
|
||||
|
||||
self.joint_angle_limit = np.load(os.path.join(self.lp, "CONVEX_HULLS.npy"),
|
||||
allow_pickle=True)
|
||||
LEN_joint_angle_limit = len(self.joint_angle_limit)
|
||||
|
||||
self.bl_max = torch.from_numpy(self.bone_len_max).float().cuda()
|
||||
self.bl_min = torch.from_numpy(self.bone_len_min).float().cuda()
|
||||
self.rb_curvatures_max = torch.from_numpy(self.rb_curvatures_max).float().cuda()
|
||||
self.rb_curvatures_min = torch.from_numpy(self.rb_curvatures_min).float().cuda()
|
||||
self.rb_PHI_max = torch.from_numpy(self.rb_PHI_max).float().cuda()
|
||||
self.rb_PHI_min = torch.from_numpy(self.rb_PHI_min).float().cuda()
|
||||
|
||||
self.joint_angle_limit = [torch.from_numpy(self.joint_angle_limit[i]).float().cuda() for i in
|
||||
range(LEN_joint_angle_limit)]
|
||||
|
||||
def compute_loss(self, ori_joints, valid_mask):
|
||||
'''
|
||||
|
||||
Args:
|
||||
joints: B*21*3
|
||||
|
||||
Returns:
|
||||
|
||||
'''
|
||||
|
||||
joints = ori_joints[valid_mask]
|
||||
print(joints.shape, ori_joints.shape)
|
||||
batch_size = joints.shape[0]
|
||||
final_loss = torch.Tensor([0]).cuda()
|
||||
|
||||
BMC_losses = {"bmc_bl": torch.Tensor([0]).cuda(), "bmc_rb": torch.Tensor([0]).cuda(), "bmc_a": torch.Tensor([0]).cuda()}
|
||||
|
||||
if (self.lambda_bl < 1e-6) and (self.lambda_rb < 1e-6) and (self.lambda_a < 1e-6):
|
||||
return final_loss, BMC_losses
|
||||
|
||||
ALL_bones = [
|
||||
(
|
||||
joints[:, i, :] -
|
||||
joints[:, SNAP_PARENT[i], :]
|
||||
) for i in range(21)
|
||||
]
|
||||
ALL_bones = torch.stack(ALL_bones[1:], dim=1) # (B,20,3)
|
||||
ROOT_bones = ALL_bones[:, ID_ROOT_bone] # (B,5,3)
|
||||
PIP_bones = ALL_bones[:, ID_PIP_bone]
|
||||
DIP_bones = ALL_bones[:, ID_DIP_bone]
|
||||
TIP_bones = ALL_bones[:, ID_TIP_bone]
|
||||
|
||||
ALL_Z_axis = normalize(ALL_bones)
|
||||
PIP_Z_axis = ALL_Z_axis[:, ID_ROOT_bone]
|
||||
DIP_Z_axis = ALL_Z_axis[:, ID_PIP_bone]
|
||||
TIP_Z_axis = ALL_Z_axis[:, ID_DIP_bone]
|
||||
|
||||
normals = normalize(cross_product(ROOT_bones[:, 1:], ROOT_bones[:, :-1]))
|
||||
|
||||
# compute loss of bone length
|
||||
bl_loss = torch.Tensor([0]).cuda()
|
||||
if self.lambda_bl:
|
||||
bls = two_norm(ALL_bones) # (B,20,1)
|
||||
bl_loss = interval_loss(value=bls, min=self.bl_min, max=self.bl_max)
|
||||
final_loss += self.lambda_bl * bl_loss
|
||||
BMC_losses["bmc_bl"] = bl_loss
|
||||
|
||||
# compute loss of Root bones
|
||||
rb_loss = torch.Tensor([0]).cuda()
|
||||
if self.lambda_rb:
|
||||
edge_normals = torch.zeros_like(ROOT_bones).cuda() # (B,5,3)
|
||||
edge_normals[:, [0, 4]] = normals[:, [0, 3]]
|
||||
edge_normals[:, 1:4] = normalize(normals[:, 1:4] + normals[:, :3])
|
||||
|
||||
curvatures = inner_product(edge_normals[:, 1:] - edge_normals[:, :4],
|
||||
ROOT_bones[:, 1:] - ROOT_bones[:, :4]) / \
|
||||
(two_norm(ROOT_bones[:, 1:] - ROOT_bones[:, :4]) ** 2)
|
||||
PHI = angle_between(ROOT_bones[:, :4], ROOT_bones[:, 1:]) # (B)
|
||||
|
||||
rb_loss = interval_loss(value=curvatures, min=self.rb_curvatures_min, max=self.rb_curvatures_max) + \
|
||||
interval_loss(value=PHI, min=self.rb_PHI_min, max=self.rb_PHI_max)
|
||||
final_loss += self.lambda_rb * rb_loss
|
||||
BMC_losses["bmc_rb"] = rb_loss
|
||||
|
||||
# compute loss of Joint angles
|
||||
a_loss = torch.Tensor([0]).cuda()
|
||||
if self.lambda_a:
|
||||
# PIP bones
|
||||
PIP_X_axis = torch.zeros([batch_size, 5, 3]).cuda() # (B,5,3)
|
||||
PIP_X_axis[:, [0, 1, 4], :] = -normals[:, [0, 1, 3]]
|
||||
PIP_X_axis[:, 2:4] = -normalize(normals[:, 2:4] + normals[:, 1:3]) # (B,2,3)
|
||||
PIP_Y_axis = normalize(cross_product(PIP_Z_axis, PIP_X_axis)) # (B,5,3)
|
||||
|
||||
PIP_bones_xz = PIP_bones - inner_product(PIP_bones, PIP_Y_axis).unsqueeze(2) * PIP_Y_axis
|
||||
PIP_theta_flexion = angle_between(PIP_bones_xz, PIP_Z_axis) # in global coordinate (B)
|
||||
PIP_theta_abduction = angle_between(PIP_bones_xz, PIP_bones) # in global coordinate (B)
|
||||
# x-component of the bone vector
|
||||
tmp = inner_product(PIP_bones, PIP_X_axis)
|
||||
PIP_theta_flexion = torch.where(tmp < 1e-6, -PIP_theta_flexion, PIP_theta_flexion)
|
||||
# y-component of the bone vector
|
||||
tmp = torch.sum((PIP_bones * PIP_Y_axis), dim=-1)
|
||||
PIP_theta_abduction = torch.where(tmp < 1e-6, -PIP_theta_abduction, PIP_theta_abduction)
|
||||
|
||||
temp_axis = normalize(cross_product(PIP_Z_axis, PIP_bones))
|
||||
temp_alpha = angle_between(PIP_Z_axis, PIP_bones) # alpha belongs to [pi/2, pi]
|
||||
temp_R = axangle2mat_torch(axis=temp_axis, angle=temp_alpha, is_normalized=True)
|
||||
|
||||
# DIP bones
|
||||
DIP_X_axis = torch.matmul(temp_R, PIP_X_axis.unsqueeze(3)).squeeze()
|
||||
DIP_Y_axis = torch.matmul(temp_R, PIP_Y_axis.unsqueeze(3)).squeeze()
|
||||
|
||||
DIP_bones_xz = DIP_bones - inner_product(DIP_bones, DIP_Y_axis).unsqueeze(2) * DIP_Y_axis
|
||||
DIP_theta_flexion = angle_between(DIP_bones_xz, DIP_Z_axis) # in global coordinate
|
||||
DIP_theta_abduction = angle_between(DIP_bones_xz, DIP_bones) # in global coordinate
|
||||
# x-component of the bone vector
|
||||
tmp = inner_product(DIP_bones, DIP_X_axis)
|
||||
DIP_theta_flexion = torch.where(tmp < 0, -DIP_theta_flexion, DIP_theta_flexion)
|
||||
# y-component of the bone vector
|
||||
tmp = inner_product(DIP_bones, DIP_Y_axis)
|
||||
DIP_theta_abduction = torch.where(tmp < 0, -DIP_theta_abduction, DIP_theta_abduction)
|
||||
|
||||
temp_axis = normalize(cross_product(DIP_Z_axis, DIP_bones))
|
||||
temp_alpha = angle_between(DIP_Z_axis, DIP_bones) # alpha belongs to [pi/2, pi]
|
||||
temp_R = axangle2mat_torch(axis=temp_axis, angle=temp_alpha, is_normalized=True)
|
||||
|
||||
# TIP bones
|
||||
TIP_X_axis = torch.matmul(temp_R, DIP_X_axis.unsqueeze(3)).squeeze()
|
||||
TIP_Y_axis = torch.matmul(temp_R, DIP_Y_axis.unsqueeze(3)).squeeze()
|
||||
TIP_bones_xz = TIP_bones - inner_product(TIP_bones, TIP_Y_axis).unsqueeze(2) * TIP_Y_axis
|
||||
|
||||
TIP_theta_flexion = angle_between(TIP_bones_xz, TIP_Z_axis) # in global coordinate
|
||||
TIP_theta_abduction = angle_between(TIP_bones_xz, TIP_bones) # in global coordinate
|
||||
# x-component of the bone vector
|
||||
tmp = inner_product(TIP_bones, TIP_X_axis)
|
||||
TIP_theta_flexion = torch.where(tmp < 1e-6, -TIP_theta_flexion, TIP_theta_flexion)
|
||||
# y-component of the bone vector
|
||||
tmp = inner_product(TIP_bones, TIP_Y_axis)
|
||||
TIP_theta_abduction = torch.where(tmp < 1e-6, -TIP_theta_abduction, TIP_theta_abduction)
|
||||
|
||||
# ALL
|
||||
ALL_theta_flexion = torch.cat((PIP_theta_flexion, DIP_theta_flexion, TIP_theta_flexion), dim=-1)
|
||||
ALL_theta_abduction = torch.cat((PIP_theta_abduction, DIP_theta_abduction, TIP_theta_abduction), dim=-1)
|
||||
ALL_theta = torch.stack((ALL_theta_flexion, ALL_theta_abduction), dim=-1)
|
||||
|
||||
a_loss = calculate_joint_angle_loss(ALL_theta, self.joint_angle_limit)
|
||||
final_loss += self.lambda_a * a_loss
|
||||
|
||||
BMC_losses["bmc_a"] = a_loss
|
||||
|
||||
return final_loss, BMC_losses
|
||||
+180
@@ -0,0 +1,180 @@
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from body_model import OP_EDGE_LIST
|
||||
from geometry.rotation import batch_rodrigues
|
||||
from geometry.plane import (
|
||||
parse_floor_plane,
|
||||
compute_plane_intersection,
|
||||
fit_plane,
|
||||
)
|
||||
|
||||
|
||||
def compute_world2prior(floor_plane, trans, root_orient, origin):
|
||||
"""
|
||||
Computes rotation and translation from the camera frame to the canonical coordinate system
|
||||
used by the motion and initial state priors.
|
||||
- floor_plane : B x 3
|
||||
- trans : B x 3
|
||||
- root_orient : B x 3
|
||||
- origin: B x 3 desired origin (first joint of human model)
|
||||
- returns R (B, 3, 3), t (B, 3), root_height (B)
|
||||
"""
|
||||
B = trans.size(0)
|
||||
if floor_plane.size(1) == 3:
|
||||
floor_plane_4d = parse_floor_plane(floor_plane)
|
||||
else:
|
||||
floor_plane_4d = floor_plane
|
||||
floor_normal = floor_plane_4d[:, :3]
|
||||
floor_trans, _ = compute_plane_intersection(trans, -floor_normal, floor_plane_4d)
|
||||
|
||||
# compute prior frame axes within the camera frame
|
||||
# up is the floor_plane normal
|
||||
up_axis = floor_normal
|
||||
# right is body -x direction projected to floor plane
|
||||
root_orient_mat = batch_rodrigues(root_orient)
|
||||
body_right = -root_orient_mat[:, :, 0]
|
||||
floor_body_right, s = compute_plane_intersection(trans, body_right, floor_plane_4d)
|
||||
right_axis = floor_body_right - floor_trans
|
||||
# body right may not actually intersect - in this case must negate axis because we have the -x
|
||||
right_axis = torch.where(s.reshape((B, 1)) < 0, -right_axis, right_axis)
|
||||
right_axis = right_axis / torch.norm(right_axis, dim=1, keepdim=True)
|
||||
# forward is their cross product
|
||||
fwd_axis = torch.linalg.cross(up_axis, right_axis)
|
||||
fwd_axis = fwd_axis / torch.norm(fwd_axis, dim=1, keepdim=True)
|
||||
|
||||
# prior frame is right, fwd, up
|
||||
prior_R = torch.stack([right_axis, fwd_axis, up_axis], dim=2)
|
||||
world2prior_R = prior_R.transpose(-1, -2)
|
||||
|
||||
# translation takes translation to origin plus offset to the floor
|
||||
world2prior_t = -trans
|
||||
|
||||
# compute the distance from root to ground plane
|
||||
_, s_root = compute_plane_intersection(origin, -floor_normal, floor_plane_4d)
|
||||
root_height = s_root.reshape(B, 1)
|
||||
|
||||
if False:
|
||||
# apply transform to origin
|
||||
origin_prior = torch.matmul(world2prior_R, (origin - trans)[..., None])[..., 0]
|
||||
root_height = s_root.reshape(B, 1) - origin_prior[:, 2:3]
|
||||
|
||||
return world2prior_R, world2prior_t, root_height
|
||||
|
||||
|
||||
def find_cliques(edges):
|
||||
cliques = []
|
||||
N = edges.shape[0]
|
||||
empty = np.ones(N, dtype=bool)
|
||||
while empty.sum() > 0:
|
||||
idcs = np.where(empty)[0]
|
||||
i = idcs[0]
|
||||
cur_clique = [i]
|
||||
for j in idcs[1:]:
|
||||
if all(edges[c, j] for c in cur_clique):
|
||||
cur_clique.append(j)
|
||||
empty[cur_clique] = False
|
||||
cliques.append(cur_clique)
|
||||
return cliques
|
||||
|
||||
|
||||
def estimate_floor_planes(
|
||||
smpl_joints, valid_mask, thresh=0.5, group=False, flatten=False
|
||||
):
|
||||
"""
|
||||
:param smpl_joints (B, T, J, 3)
|
||||
:param valid_mask (B, T)
|
||||
:param thresh (optional, default 0.5) threshold for separating
|
||||
:param group (optional, default False) whether to group the planes
|
||||
:param flatten (optional, default False) whether to estimate one plane for everyone
|
||||
returns floor_plane (C, 3), floor_idcs (B,)
|
||||
"""
|
||||
B, T = valid_mask.shape
|
||||
device = valid_mask.device
|
||||
feet_joints = smpl_joints[..., 10:12, :] # (B, T, 2, 3)
|
||||
points = [feet_joints[b, valid_mask[b]].reshape(-1, 3) for b in range(B)]
|
||||
|
||||
if flatten:
|
||||
planes = fit_plane(torch.cat(points, dim=0))[None]
|
||||
labels = torch.zeros(B, device=device, dtype=torch.long)
|
||||
return planes[..., :3] * planes[..., 3:], labels
|
||||
|
||||
# estimate floors independently
|
||||
planes = torch.stack([fit_plane(pts) for pts in points], dim=0) # (B, 4)
|
||||
normals, offsets = planes[..., :3], planes[..., 3:] # (B, 3), (B, 1)
|
||||
sgn = torch.sign(offsets)
|
||||
normals, offsets = sgn * normals, sgn * offsets
|
||||
planes = normals * offsets # (B, 3)
|
||||
labels = torch.arange(B, device=device)
|
||||
|
||||
if not group:
|
||||
return planes, labels # (B, 3), (B,)
|
||||
|
||||
# asymetric, project each plane vector onto the other normals
|
||||
diffs = offsets - (planes[None] * normals[:, None]).sum(dim=-1)
|
||||
print(diffs)
|
||||
edges = diffs < thresh
|
||||
edges = edges.T & edges
|
||||
clusters = find_cliques(edges) # list of indices of clusters
|
||||
planes = []
|
||||
labels = torch.zeros(B, dtype=torch.long, device=device)
|
||||
for c, idcs in enumerate(clusters):
|
||||
plane = fit_plane(torch.cat([points[i] for i in idcs], dim=0))
|
||||
planes.append(plane)
|
||||
labels[idcs] = c
|
||||
planes = torch.stack(planes, dim=0) # (C, 4)
|
||||
return planes[..., :3] * planes[..., 3:], labels
|
||||
|
||||
|
||||
def estimate_initial_trans(body_pose, joints3d_op, joints2d_op, focal):
|
||||
"""
|
||||
use focal length and bone lengths to approximate distance from camera
|
||||
(based on PROX https://github.com/mohamedhassanmus/prox/blob/master/prox/fitting.py)
|
||||
:param body_pose
|
||||
:param joints3d_op OP keypoints in 3d
|
||||
:param joints2d_op OP keypoints in 2d
|
||||
"""
|
||||
B, T, N, _ = joints2d_op.shape
|
||||
joints2d_obs = joints2d_op[:, :, :, :2] # (B, T, N, 3)
|
||||
joints2d_conf = joints2d_op[:, :, :, 2]
|
||||
|
||||
# find least-occluded 2d frame
|
||||
num_2d_vis = torch.sum(joints2d_conf > 0.0, dim=2)
|
||||
best_2d_idx = torch.max(num_2d_vis, dim=1)[1]
|
||||
|
||||
# calculate bone lengths and confidence in each bone length
|
||||
bone3d = []
|
||||
bone2d = []
|
||||
conf2d = []
|
||||
for pair in OP_EDGE_LIST:
|
||||
diff3d = torch.norm(
|
||||
joints3d_op[:, 0, pair[0], :] - joints3d_op[:, 0, pair[1], :],
|
||||
dim=1,
|
||||
) # does not change over time
|
||||
diff2d = torch.norm(
|
||||
joints2d_obs[:, :, pair[0], :] - joints2d_obs[:, :, pair[1], :],
|
||||
dim=2,
|
||||
)
|
||||
minconf2d = torch.min(
|
||||
joints2d_conf[:, :, pair[0]], joints2d_conf[:, :, pair[1]]
|
||||
)
|
||||
bone3d.append(diff3d)
|
||||
bone2d.append(diff2d)
|
||||
conf2d.append(minconf2d)
|
||||
|
||||
bone3d = torch.stack(bone3d, dim=1)
|
||||
bone2d = torch.stack(bone2d, dim=2)
|
||||
bone2d = bone2d[np.arange(B), best_2d_idx, :]
|
||||
conf2d = torch.stack(conf2d, dim=2)
|
||||
conf2d = conf2d[np.arange(B), best_2d_idx, :]
|
||||
|
||||
# mean over all bones
|
||||
mean_bone3d = torch.mean(bone3d, dim=1)
|
||||
mean_bone2d = torch.mean(bone2d * (conf2d > 0.0), dim=1)
|
||||
|
||||
# approx z based on ratio
|
||||
init_z = focal * (mean_bone3d / mean_bone2d)
|
||||
init_trans = torch.zeros(B, T, 3)
|
||||
init_trans[:, :, 2] = init_z.unsqueeze(1).expand(B, T).detach()
|
||||
|
||||
return init_trans
|
||||
+1456
File diff suppressed because it is too large
Load Diff
+727
@@ -0,0 +1,727 @@
|
||||
import os
|
||||
import numpy as np
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from body_model import MANO_JOINTS, KEYPT_VERTS, smpl_to_openpose
|
||||
from geometry.rotation import batch_rodrigues, rotation_matrix_to_angle_axis
|
||||
from util.logger import Logger
|
||||
from util.tensor import (
|
||||
detach_all,
|
||||
scatter_intervals,
|
||||
select_intervals,
|
||||
get_scatter_mask,
|
||||
)
|
||||
|
||||
from .base_scene import BaseSceneModel
|
||||
from .helpers import compute_world2prior, estimate_floor_planes
|
||||
|
||||
|
||||
J_BODY = len(MANO_JOINTS) - 1 # no root
|
||||
CONTACT_ORDERING = [
|
||||
"hips",
|
||||
"leftLeg",
|
||||
"rightLeg",
|
||||
"leftFoot",
|
||||
"rightFoot",
|
||||
"leftToeBase",
|
||||
"rightToeBase",
|
||||
"leftHand",
|
||||
"rightHand",
|
||||
]
|
||||
CONTACT_INDS = [MANO_JOINTS[jname] for jname in CONTACT_ORDERING]
|
||||
CONTACT_THRESH = 0.5
|
||||
|
||||
|
||||
class MovingSceneModel(BaseSceneModel):
|
||||
"""
|
||||
Scene model of moving people in a shared global reference frame
|
||||
|
||||
Parameters:
|
||||
batch_size: number of sequences to optimize
|
||||
seq_len: length of the sequences
|
||||
pose_prior: VPoser model
|
||||
motion_prior: humor model
|
||||
init_motion_prior: dict of GMM params to use for prior on initial motion state
|
||||
shared_floor: (default True) if true, sequences are in the same reference frame
|
||||
fit_gender: (optional) gender of SMPL model
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
batch_size,
|
||||
seq_len,
|
||||
body_model,
|
||||
pose_prior,
|
||||
motion_prior,
|
||||
init_motion_prior=None,
|
||||
fit_gender="neutral",
|
||||
use_init=False,
|
||||
opt_cams=False,
|
||||
opt_scale=True,
|
||||
cam_graph=False,
|
||||
est_floor=True,
|
||||
floor_type="shared",
|
||||
async_tracks=True,
|
||||
**kwargs,
|
||||
):
|
||||
super().__init__(
|
||||
batch_size,
|
||||
seq_len,
|
||||
body_model,
|
||||
pose_prior,
|
||||
fit_gender=fit_gender,
|
||||
use_init=use_init,
|
||||
opt_cams=opt_cams,
|
||||
opt_scale=opt_scale,
|
||||
cam_graph=cam_graph,
|
||||
)
|
||||
assert motion_prior is not None
|
||||
assert motion_prior.model_data_config in [
|
||||
"mano+joints",
|
||||
"mano+joints+contacts",
|
||||
], "Only mano+joints motion prior configuration is supported!"
|
||||
|
||||
self.motion_prior = motion_prior
|
||||
self.init_motion_prior = init_motion_prior
|
||||
|
||||
# need latent dynamics sequence as well
|
||||
self.latent_motion_dim = self.motion_prior.latent_size
|
||||
self.cond_prior = self.motion_prior.use_conditional_prior
|
||||
assert self.cond_prior
|
||||
|
||||
# the frame chosen to use for the initial state (first frame by default)
|
||||
self.async_tracks = async_tracks
|
||||
self.register_buffer("track_start", torch.zeros(self.batch_size))
|
||||
self.register_buffer("track_end", torch.ones(self.batch_size) * self.seq_len)
|
||||
|
||||
self.shared_floor = floor_type == "shared"
|
||||
self.group_floor = floor_type == "group"
|
||||
self.floor_type = floor_type
|
||||
print("FLOOR TYPE", floor_type)
|
||||
self.est_floor = est_floor
|
||||
|
||||
@property
|
||||
def is_motion_active(self):
|
||||
return hasattr(self.params, "latent_motion")
|
||||
|
||||
def initialize(self, obs_data, cam_data, param_dict, data_fps):
|
||||
"""
|
||||
we need to also optimize for floor and world scale
|
||||
"""
|
||||
Logger.log("Initializing moving scene model with observed data")
|
||||
|
||||
self.params.set_cameras(
|
||||
cam_data,
|
||||
opt_scale=self.opt_scale,
|
||||
opt_cams=self.opt_cams,
|
||||
opt_focal=self.opt_cams,
|
||||
**param_dict,
|
||||
)
|
||||
# self.init_floor(obs_data, param_dict)
|
||||
self.init_first_state(obs_data, param_dict, data_fps)
|
||||
|
||||
# def init_floor(self, obs_data, param_dict):
|
||||
# if self.est_floor or self.group_floor:
|
||||
# with torch.no_grad():
|
||||
# smpl_preds = self.pred_mano(
|
||||
# param_dict["trans"],
|
||||
# param_dict["root_orient"],
|
||||
# self.latent2pose(param_dict["latent_pose"]),
|
||||
# param_dict["betas"],
|
||||
# )
|
||||
# floor_plane, floor_idcs = estimate_floor_planes(
|
||||
# smpl_preds["joints3d"].detach(),
|
||||
# obs_data["vis_mask"] > 0,
|
||||
# group=self.group_floor,
|
||||
# flatten=self.shared_floor,
|
||||
# )
|
||||
# if self.group_floor and not self.est_floor:
|
||||
# # don't use the estimated floor as initial plane
|
||||
# floor_plane = obs_data["floor_plane"][: len(floor_plane)]
|
||||
# else: # fixed shared or separate floors
|
||||
# num_floors = 1 if self.shared_floor else self.batch_size
|
||||
# floor_plane = obs_data["floor_plane"][:num_floors]
|
||||
# if self.shared_floor:
|
||||
# floor_idcs = torch.zeros(
|
||||
# self.batch_size, dtype=torch.long, device=floor_plane.device
|
||||
# )
|
||||
# else:
|
||||
# floor_idcs = torch.arange(
|
||||
# self.batch_size, dtype=torch.long, device=floor_plane.device
|
||||
# )
|
||||
|
||||
# Logger.log(f"ESTIMATED FLOORS: {str(floor_plane.detach().cpu())}")
|
||||
# Logger.log(f"FLOOR IDCS: {str(floor_idcs.detach().cpu())}")
|
||||
# self.params.set_param("floor_plane", floor_plane.float().detach())
|
||||
# self.params.set_param(
|
||||
# "floor_idcs", floor_idcs.long().detach(), requires_grad=False
|
||||
# )
|
||||
|
||||
def init_first_state(self, obs_data, param_dict, data_fps):
|
||||
"""
|
||||
initialize the latent motion and first state of each track
|
||||
using per-frame trajectories of trans, rot, latent_pose, betas
|
||||
and observed data
|
||||
"""
|
||||
B, T = self.batch_size, self.seq_len
|
||||
param_dict = detach_all(param_dict)
|
||||
|
||||
# select the valid segments of each track; pad the shorter tracks with final element
|
||||
if self.async_tracks and "track_interval" in obs_data:
|
||||
print("asynchronous tracks")
|
||||
interval = obs_data["track_interval"] # (B, 2)
|
||||
start, end = interval[:, 0], interval[:, 1]
|
||||
self.track_start, self.track_end = start, end
|
||||
param_dict = select_dict_segments(
|
||||
param_dict, start, end, names=["trans", "root_orient", "latent_pose"]
|
||||
)
|
||||
|
||||
trans = param_dict["trans"] # (B, T, 3)
|
||||
root_orient = param_dict["root_orient"] # (B, T, 3)
|
||||
latent_pose = param_dict["latent_pose"] # (B, T, D)
|
||||
betas = param_dict["betas"] # (B, b)
|
||||
|
||||
Logger.log(f"SEL TRACKS {trans.shape}, {root_orient.shape}")
|
||||
|
||||
# save each track's first appearance
|
||||
self.params.set_param("trans", trans[:, :1])
|
||||
Logger.log(f"INITIAL TRANS {trans[:, :1].detach().cpu()}")
|
||||
self.params.set_param("root_orient", root_orient[:, :1])
|
||||
self.params.set_param("latent_pose", latent_pose[:, :1])
|
||||
self.params.set_param("betas", betas)
|
||||
|
||||
# pass the current pose estimates through the motion prior
|
||||
self.data_fps = data_fps
|
||||
body_pose = self.latent2pose(latent_pose)
|
||||
init_latent = self.infer_latent_motion(
|
||||
trans, root_orient, body_pose, betas, data_fps
|
||||
).detach() # (B, T-1, D)
|
||||
self.params.set_param("latent_motion", init_latent)
|
||||
|
||||
# estimate velocities in prior frame and save initial state
|
||||
trans_vel, joints_vel, root_orient_vel = self.estimate_prior_velocities(
|
||||
trans, root_orient, body_pose, betas, data_fps
|
||||
)
|
||||
self.params.set_param("trans_vel", trans_vel[:, :1].detach())
|
||||
self.params.set_param("joints_vel", joints_vel[:, :1].detach())
|
||||
self.params.set_param("root_orient_vel", root_orient_vel[:, :1].detach())
|
||||
|
||||
def get_optim_result(self, num_steps=-1):
|
||||
res = super().get_optim_result()
|
||||
if not self.is_motion_active:
|
||||
return res
|
||||
|
||||
num_steps = self.seq_len if num_steps < 0 else num_steps
|
||||
prior_res, world_res = self.rollout_latent_motion(
|
||||
self.params.latent_motion[:, : num_steps - 1]
|
||||
)
|
||||
world_res = self.synchronize_preds(world_res, num_steps)
|
||||
res["world"].update(world_res)
|
||||
|
||||
prior_res = self.synchronize_preds(prior_res, num_steps)
|
||||
res["prior"] = prior_res
|
||||
return res
|
||||
|
||||
def estimate_prior_velocities(self, trans, root_orient, body_pose, betas, data_fps):
|
||||
"""
|
||||
compute velocities of trajectory in prior frame (not world frame)
|
||||
Transforms the first of each track into the prior frame
|
||||
:param trans (B, T, 3)
|
||||
:param root_orient (B, T, 3)
|
||||
:param body_pose (B, T, Dp)
|
||||
:param betas (B, Db)
|
||||
"""
|
||||
with torch.no_grad():
|
||||
self.update_world2prior(trans, root_orient, body_pose, betas)
|
||||
trans, root_orient = self.apply_world2prior(
|
||||
trans, root_orient, body_pose, betas
|
||||
)
|
||||
mano_results = self.pred_mano(trans, root_orient, body_pose, betas)
|
||||
return estimate_velocities(
|
||||
trans, root_orient, mano_results["joints3d"], data_fps
|
||||
)
|
||||
|
||||
def update_world2prior(self, trans, root_orient, body_pose, betas):
|
||||
B = trans.shape[0]
|
||||
mano_data = self.pred_mano(
|
||||
trans[:, :1], root_orient[:, :1], body_pose[:, :1], betas
|
||||
)
|
||||
floor_plane = self.params.floor_plane[self.params.floor_idcs]
|
||||
R, t, height = compute_world2prior(
|
||||
floor_plane,
|
||||
trans[:, 0],
|
||||
root_orient[:, 0],
|
||||
mano_data["joints3d"][:, 0, 0],
|
||||
)
|
||||
self.world2prior_R = R # (B, 3, 3)
|
||||
self.world2prior_t = t # (B, 3)
|
||||
self.world2prior_root_height = height # (B)
|
||||
|
||||
def apply_world2prior(self, trans, root_orient, body_pose, betas, inverse=False):
|
||||
"""
|
||||
Applies the world2prior tranformation to trans, root_orient
|
||||
If T=1, this function assumes they are at key_frame_idx,
|
||||
which we need to compute the offset from the origin
|
||||
:param trans (B, T, 3)
|
||||
:param root_orient (B, T, 3)
|
||||
:param body_pose (B, T, J, 3)
|
||||
:param betas (B, b)
|
||||
:param inverse (bool) optional, default False
|
||||
"""
|
||||
B, T, _ = trans.size()
|
||||
# (B, 3, 3), (B, 3), (B)
|
||||
R, t, root_height = (
|
||||
self.world2prior_R,
|
||||
self.world2prior_t,
|
||||
self.world2prior_root_height,
|
||||
)
|
||||
R_time = R.unsqueeze(1).expand((B, T, 3, 3))
|
||||
t_time = t.unsqueeze(1).expand((B, T, 3))
|
||||
root_orient_mat = batch_rodrigues(root_orient.reshape(-1, 3)).reshape(
|
||||
B, T, 3, 3
|
||||
)
|
||||
|
||||
if inverse:
|
||||
R_time = R_time.transpose(-1, -2)
|
||||
|
||||
root_orient_mat = torch.matmul(R_time, root_orient_mat)
|
||||
root_orient = rotation_matrix_to_angle_axis(
|
||||
root_orient_mat.reshape(B * T, 3, 3)
|
||||
).reshape(B, T, 3)
|
||||
|
||||
if inverse:
|
||||
# transform so first frame is at origin
|
||||
trans = trans - trans[:, :1, :]
|
||||
|
||||
# rotates to world frame
|
||||
trans = torch.matmul(R_time, trans[..., None])[..., 0]
|
||||
# translate to world frame
|
||||
trans = trans - t_time
|
||||
|
||||
return trans, root_orient
|
||||
|
||||
# first transform so the trans of key frame is at origin
|
||||
trans = trans + t_time
|
||||
# then rotate to canonical frame
|
||||
trans = torch.matmul(R_time, trans[..., None])[..., 0]
|
||||
# compute the root height after transforming
|
||||
cur_smpl_data = self.pred_mano(trans, root_orient, body_pose, betas)
|
||||
cur_root_height = cur_smpl_data["joints3d"][:, 0, 0, 2:3]
|
||||
# then apply floor offset so the root joint is at root_height
|
||||
height_diff = root_height - cur_root_height
|
||||
trans_offset = torch.cat(
|
||||
[torch.zeros((B, 2)).to(height_diff), height_diff], axis=1
|
||||
)
|
||||
trans = trans + trans_offset.reshape((B, 1, 3))
|
||||
|
||||
return trans, root_orient
|
||||
|
||||
def convert_prior_rot_inputs(self, root_orient, body_pose):
|
||||
# convert rots
|
||||
# body pose and root orient are both in aa
|
||||
B, T = root_orient.shape[:2]
|
||||
if (
|
||||
self.motion_prior.in_rot_rep == "mat"
|
||||
or self.motion_prior.in_rot_rep == "6d"
|
||||
):
|
||||
root_orient_in = batch_rodrigues(root_orient.reshape(-1, 3)).reshape(
|
||||
(B, T, 9)
|
||||
)
|
||||
body_pose_in = batch_rodrigues(body_pose.reshape(-1, 3)).reshape(
|
||||
(B, T, J_BODY * 9)
|
||||
)
|
||||
if self.motion_prior.in_rot_rep == "6d":
|
||||
root_orient_in = root_orient[:, :, :6]
|
||||
body_pose_in = body_pose.reshape((B, T, J_BODY, 9))[:, :, :, :6].reshape(
|
||||
(B, T, J_BODY * 6)
|
||||
)
|
||||
return root_orient_in, body_pose_in
|
||||
|
||||
def convert_prior_rot_outputs(self, out_dict):
|
||||
keys = ["root_orient", "pose_body"]
|
||||
for k in keys:
|
||||
out = out_dict[k]
|
||||
B, T = out.shape[:2]
|
||||
out_dict[k] = rotation_matrix_to_angle_axis(
|
||||
out.reshape((-1, 3, 3))
|
||||
).reshape((B, T, -1))
|
||||
return out_dict
|
||||
|
||||
def infer_latent_motion(
|
||||
self, trans, root_orient, body_pose, betas, data_fps, full_forward_pass=False
|
||||
):
|
||||
"""
|
||||
By default, gets a sequence of z's from the current SMPL optim params.
|
||||
|
||||
If full_forward_pass is true, also samples from the posterior and feeds
|
||||
through the motion prior decoder to get all terms needed to calculate the ELBO.
|
||||
"""
|
||||
B, T, _ = trans.size()
|
||||
h = 1.0 / data_fps
|
||||
|
||||
# need to first transform into canonical coordinate frame
|
||||
self.update_world2prior(trans, root_orient, body_pose, betas)
|
||||
trans, root_orient = self.apply_world2prior(
|
||||
trans, root_orient, body_pose, betas
|
||||
)
|
||||
|
||||
mano_results = self.pred_mano(trans, root_orient, body_pose, betas)
|
||||
joints = mano_results["joints3d"] # (B, T, len(MANO_JOINTS), 3)
|
||||
trans_vel, joints_vel, root_orient_vel = estimate_velocities(
|
||||
trans, root_orient, joints, data_fps
|
||||
)
|
||||
|
||||
root_orient_in, body_pose_in = self.convert_prior_rot_inputs(
|
||||
root_orient, body_pose
|
||||
)
|
||||
joints_in = joints.reshape((B, T, -1))
|
||||
joints_vel_in = joints_vel.reshape((B, T, -1))
|
||||
|
||||
seq_dict = {
|
||||
"trans": trans,
|
||||
"trans_vel": trans_vel,
|
||||
"root_orient": root_orient_in,
|
||||
"root_orient_vel": root_orient_vel,
|
||||
"pose_body": body_pose_in,
|
||||
"joints": joints_in,
|
||||
"joints_vel": joints_vel_in,
|
||||
}
|
||||
|
||||
infer_results = self.motion_prior.infer_global_seq(
|
||||
seq_dict, full_forward_pass=full_forward_pass
|
||||
)
|
||||
if full_forward_pass:
|
||||
# return both the given motion and the one from the forward pass
|
||||
# make sure rotations are matrix
|
||||
# NOTE: assumes seq_dict is same thing we want to compute loss on
|
||||
# Need to change if multiple future steps.
|
||||
if self.motion_prior.in_rot_rep != "mat":
|
||||
seq_dict["trans"] = batch_rodrigues(root_orient.reshape(-1, 3)).reshape(
|
||||
(B, T, 9)
|
||||
)
|
||||
seq_dict["pose_body"] = batch_rodrigues(
|
||||
body_pose.reshape(-1, 3)
|
||||
).reshape((B, T, J_BODY * 9))
|
||||
# do not need initial step anymore since output will be T-1
|
||||
for k, v in seq_dict.items():
|
||||
seq_dict[k] = v[:, 1:]
|
||||
for k in infer_results.keys():
|
||||
if k != "posterior_distrib" and k != "prior_distrib":
|
||||
infer_results[k] = infer_results[k][
|
||||
:, :, 0
|
||||
] # only want first output step
|
||||
infer_results = (seq_dict, infer_results)
|
||||
else:
|
||||
prior_z, posterior_z = infer_results
|
||||
infer_results = posterior_z[0] # mean of the approximate posterior
|
||||
|
||||
return infer_results
|
||||
|
||||
def rollout_mano_steps(self, num_steps=-1):
|
||||
# rollout motion given initial state
|
||||
num_steps = self.seq_len if num_steps < 0 else num_steps
|
||||
latent_motion = self.params.latent_motion[:, : num_steps - 1]
|
||||
|
||||
# roll out tracks from their first appearances
|
||||
# NOTE rolled out tracks are unsynced in time
|
||||
res, world_res = self.rollout_latent_motion(
|
||||
latent_motion, return_prior=self.cond_prior
|
||||
)
|
||||
|
||||
# get the smpl predictions for the unsynced tracks
|
||||
preds = self.pred_mano(
|
||||
res["trans"], res["root_orient"], res["pose_body"], res["betas"]
|
||||
)
|
||||
world_preds = self.pred_mano(
|
||||
world_res["trans"],
|
||||
world_res["root_orient"],
|
||||
world_res["pose_body"],
|
||||
world_res["betas"],
|
||||
)
|
||||
|
||||
# pass along relevant results
|
||||
preds["joints3d_rollout"] = res["joints"]
|
||||
preds["latent_pose"] = self.pose2latent(res["pose_body"])
|
||||
preds["joints3d_init"] = preds["joints3d"][:, :1]
|
||||
|
||||
if "contacts" in res:
|
||||
preds["contacts"] = res["contacts"]
|
||||
if "contacts_conf" in res:
|
||||
preds["contacts_conf"] = res["contacts_conf"]
|
||||
|
||||
# synchronize the tracklets to the same timesteps
|
||||
preds = self.synchronize_preds(preds, num_steps)
|
||||
world_preds = self.synchronize_preds(world_preds, num_steps)
|
||||
|
||||
# pass along unsynchronized motion latents
|
||||
preds["latent_motion"] = latent_motion
|
||||
if "cond_prior" in res:
|
||||
preds["cond_prior"] = res["cond_prior"]
|
||||
|
||||
return preds, world_preds
|
||||
|
||||
def rollout_latent_motion(
|
||||
self,
|
||||
latent_motion,
|
||||
return_prior=False,
|
||||
return_vel=False,
|
||||
num_steps=-1,
|
||||
canonicalize_input=False,
|
||||
):
|
||||
"""
|
||||
From the stored initial state, rolls out a sequence from the latent_motion vector
|
||||
NOTE: the stored initial states are not at the same time step, but we return all
|
||||
predicted sequences starting from time 0.
|
||||
|
||||
If latent_motion is None, samples num_steps into the future sequence from the prior.
|
||||
using the mean of the prior rather than random samples.
|
||||
|
||||
If canonicalize_input is True, transform the initial state into the local canonical
|
||||
frame before roll out
|
||||
"""
|
||||
# get the first frame state
|
||||
# NOTE: first states do not occur at same time step
|
||||
trans, root_orient, betas = (
|
||||
self.params.trans,
|
||||
self.params.root_orient,
|
||||
self.params.betas,
|
||||
)
|
||||
trans_vel, root_orient_vel, joints_vel = (
|
||||
self.params.trans_vel,
|
||||
self.params.root_orient_vel,
|
||||
self.params.joints_vel,
|
||||
)
|
||||
body_pose = self.latent2pose(self.params.latent_pose)
|
||||
|
||||
B = trans.size(0)
|
||||
is_sampling = latent_motion is None
|
||||
Tm1 = num_steps if latent_motion is None else latent_motion.size(1)
|
||||
if is_sampling and Tm1 <= 0:
|
||||
Logger.log("num_steps must be positive to sample!")
|
||||
exit()
|
||||
|
||||
# need to first transform initial state into canonical coordinate frame
|
||||
self.update_world2prior(trans, root_orient, body_pose, betas)
|
||||
trans, root_orient = self.apply_world2prior(
|
||||
trans, root_orient, body_pose, betas
|
||||
)
|
||||
|
||||
smpl_results = self.pred_mano(trans, root_orient, body_pose, betas)
|
||||
joints = smpl_results["joints3d"] # (B, T, len(MANO_JOINTS), 3)
|
||||
|
||||
# update to correct rotations for input
|
||||
root_orient_in, body_pose_in = self.convert_prior_rot_inputs(
|
||||
root_orient, body_pose
|
||||
)
|
||||
joints_in = joints.reshape((B, 1, -1))
|
||||
joints_vel_in = joints_vel.reshape((B, 1, -1))
|
||||
|
||||
rollout_in_dict = {
|
||||
"trans": trans,
|
||||
"root_orient": root_orient_in,
|
||||
"pose_body": body_pose_in,
|
||||
"joints": joints_in,
|
||||
"trans_vel": trans_vel,
|
||||
"root_orient_vel": root_orient_vel,
|
||||
"joints_vel": joints_vel_in,
|
||||
}
|
||||
|
||||
roll_output = self.motion_prior.roll_out(
|
||||
None,
|
||||
rollout_in_dict,
|
||||
Tm1,
|
||||
z_seq=latent_motion,
|
||||
return_prior=return_prior,
|
||||
return_z=is_sampling,
|
||||
canonicalize_input=canonicalize_input,
|
||||
gender=[self.fit_gender] * B,
|
||||
betas=betas.reshape((B, 1, -1)),
|
||||
)
|
||||
|
||||
pred_dict = prior_out = None
|
||||
if return_prior:
|
||||
pred_dict, prior_out = roll_output
|
||||
else:
|
||||
pred_dict = roll_output
|
||||
|
||||
pred_dict = self.convert_prior_rot_outputs(pred_dict)
|
||||
|
||||
# concat with initial state
|
||||
trans_out = torch.cat([trans, pred_dict["trans"]], dim=1)
|
||||
root_orient_out = torch.cat([root_orient, pred_dict["root_orient"]], dim=1)
|
||||
body_pose_out = torch.cat([body_pose, pred_dict["pose_body"]], dim=1)
|
||||
joints_out = torch.cat(
|
||||
[joints, pred_dict["joints"].reshape((B, Tm1, -1, 3))], dim=1
|
||||
)
|
||||
out_dict = {
|
||||
"trans": trans_out,
|
||||
"root_orient": root_orient_out,
|
||||
"pose_body": body_pose_out,
|
||||
"joints": joints_out,
|
||||
"betas": betas,
|
||||
}
|
||||
if return_vel:
|
||||
out_dict["trans_vel"] = torch.cat(
|
||||
[trans_vel, pred_dict["trans_vel"]], dim=1
|
||||
)
|
||||
out_dict["root_orient_vel"] = torch.cat(
|
||||
[root_orient_vel, pred_dict["root_orient_vel"]], dim=1
|
||||
)
|
||||
out_dict["joints_vel"] = torch.cat(
|
||||
[joints_vel, pred_dict["joints_vel"].reshape((B, Tm1, -1, 3))],
|
||||
dim=1,
|
||||
)
|
||||
|
||||
if return_prior: # return the mean and var of distribution stacked
|
||||
pm, pv = prior_out
|
||||
out_dict["cond_prior"] = torch.stack([pm, pv], dim=-1)
|
||||
|
||||
if self.motion_prior.model_data_config == "smpl+joints+contacts":
|
||||
pred_contacts = pred_dict["contacts"]
|
||||
# get binary classification
|
||||
contact_conf = torch.sigmoid(pred_contacts)
|
||||
pred_contacts = (contact_conf > CONTACT_THRESH).to(torch.float)
|
||||
# expand to full body
|
||||
full_contact_conf = torch.zeros((B, Tm1, len(MANO_JOINTS))).to(contact_conf)
|
||||
full_contact_conf[:, :, CONTACT_INDS] = (
|
||||
full_contact_conf[:, :, CONTACT_INDS] + contact_conf
|
||||
)
|
||||
full_contacts = torch.zeros((B, Tm1, len(MANO_JOINTS))).to(pred_contacts)
|
||||
full_contacts[:, :, CONTACT_INDS] = (
|
||||
full_contacts[:, :, CONTACT_INDS] + pred_contacts
|
||||
)
|
||||
# repeat first entry for t0
|
||||
full_contact_conf = torch.cat(
|
||||
[full_contact_conf[:, 0:1], full_contact_conf], dim=1
|
||||
)
|
||||
full_contacts = torch.cat([full_contacts[:, 0:1], full_contacts], dim=1)
|
||||
out_dict["contacts_conf"] = full_contact_conf
|
||||
out_dict["contacts"] = full_contacts
|
||||
|
||||
if is_sampling:
|
||||
out_dict["z"] = pred_dict["z"]
|
||||
|
||||
cam_dict = {
|
||||
"trans": out_dict["trans"],
|
||||
"root_orient": out_dict["root_orient"],
|
||||
"pose_body": out_dict["pose_body"],
|
||||
"betas": out_dict["betas"],
|
||||
}
|
||||
|
||||
# also must return trans and root orient in camera frame
|
||||
cam_dict["trans"], cam_dict["root_orient"] = self.apply_world2prior(
|
||||
out_dict["trans"],
|
||||
out_dict["root_orient"],
|
||||
out_dict["pose_body"],
|
||||
out_dict["betas"],
|
||||
inverse=True,
|
||||
)
|
||||
return out_dict, cam_dict
|
||||
|
||||
def synchronize_preds(self, pred_dict, seg_len):
|
||||
"""
|
||||
synchronize predictions in time, scatter predictions into [0, seq_len)
|
||||
"""
|
||||
if not self.async_tracks:
|
||||
# predictions are already synchronized
|
||||
return pred_dict, None
|
||||
|
||||
# return time-synchronized predictions
|
||||
start = self.track_start
|
||||
end = torch.clip(start + seg_len, max=self.track_end)
|
||||
T = end.max()
|
||||
pred_dict = scatter_dict_segments(pred_dict, start, end, T)
|
||||
pred_mask = get_scatter_mask(start, end, T)
|
||||
pred_dict["track_mask"] = pred_mask
|
||||
return pred_dict
|
||||
|
||||
|
||||
def scatter_dict_segments(data_dict, start, end, T=None, names=None):
|
||||
"""
|
||||
the tracks as they are stored are synchronized by time step
|
||||
Uses the start and end to synchronize output predictions
|
||||
"""
|
||||
sync_dict = data_dict.copy()
|
||||
min_len = (end - start).min()
|
||||
if names is None:
|
||||
names = data_dict.keys()
|
||||
for name in names:
|
||||
val = data_dict[name]
|
||||
if not isinstance(val, torch.Tensor) or val.ndim < 3 or val.shape[1] < min_len:
|
||||
continue
|
||||
sync_dict[name] = scatter_intervals(val, start, end, T)
|
||||
return sync_dict
|
||||
|
||||
|
||||
def select_dict_segments(data_dict, start, end, names=None):
|
||||
out_data_dict = data_dict.copy()
|
||||
min_len = (end - start).min()
|
||||
if names is None:
|
||||
names = data_dict.keys()
|
||||
for name in names:
|
||||
val = data_dict[name]
|
||||
if not isinstance(val, torch.Tensor) or val.ndim < 3 or val.shape[1] < min_len:
|
||||
continue
|
||||
# only select for time-series observations
|
||||
out_data_dict[name] = select_intervals(val, start, end)
|
||||
return out_data_dict
|
||||
|
||||
|
||||
def estimate_velocities(trans, root_orient, joints3d, data_fps):
|
||||
"""
|
||||
Estimates velocity inputs to the motion prior.
|
||||
- trans (B, T, 3) root translation
|
||||
- root_orient (B, T, 3) aa root orientation
|
||||
- joints3d (B, T, len(MANO_JOINTS), 3) joints3d of SMPL prediction
|
||||
"""
|
||||
B, T, _ = trans.size()
|
||||
h = 1.0 / data_fps
|
||||
trans_vel = estimate_linear_velocity(trans, h)
|
||||
joints_vel = estimate_linear_velocity(joints3d, h)
|
||||
root_orient_mat = batch_rodrigues(root_orient.reshape((-1, 3))).reshape(
|
||||
(B, T, 3, 3)
|
||||
)
|
||||
root_orient_vel = estimate_angular_velocity(root_orient_mat, h)
|
||||
return trans_vel, joints_vel, root_orient_vel
|
||||
|
||||
|
||||
def estimate_linear_velocity(data_seq, h):
|
||||
"""
|
||||
Given some batched data sequences of T timesteps in the shape (B, T, ...), estimates
|
||||
the velocity for the middle T-2 steps using a second order central difference scheme.
|
||||
The first and last frames are with forward and backward first-order
|
||||
differences, respectively
|
||||
- h : step size
|
||||
"""
|
||||
# first steps is forward diff (t+1 - t) / h
|
||||
init_vel = (data_seq[:, 1:2] - data_seq[:, :1]) / h
|
||||
# middle steps are second order (t+1 - t-1) / 2h
|
||||
middle_vel = (data_seq[:, 2:] - data_seq[:, 0:-2]) / (2 * h)
|
||||
# last step is backward diff (t - t-1) / h
|
||||
final_vel = (data_seq[:, -1:] - data_seq[:, -2:-1]) / h
|
||||
|
||||
vel_seq = torch.cat([init_vel, middle_vel, final_vel], dim=1)
|
||||
return vel_seq
|
||||
|
||||
|
||||
def estimate_angular_velocity(rot_seq, h):
|
||||
"""
|
||||
Given a batch of sequences of T rotation matrices, estimates angular velocity at T-2 steps.
|
||||
Input sequence should be of shape (B, T, ..., 3, 3)
|
||||
"""
|
||||
# see https://en.wikipedia.org/wiki/Angular_velocity#Calculation_from_the_orientation_matrix
|
||||
dRdt = estimate_linear_velocity(rot_seq, h)
|
||||
R = rot_seq
|
||||
RT = R.transpose(-1, -2)
|
||||
# compute skew-symmetric angular velocity tensor
|
||||
w_mat = torch.matmul(dRdt, RT)
|
||||
# pull out angular velocity vector
|
||||
# average symmetric entries
|
||||
w_x = (-w_mat[..., 1, 2] + w_mat[..., 2, 1]) / 2.0
|
||||
w_y = (w_mat[..., 0, 2] - w_mat[..., 2, 0]) / 2.0
|
||||
w_z = (-w_mat[..., 0, 1] + w_mat[..., 1, 0]) / 2.0
|
||||
w = torch.stack([w_x, w_y, w_z], axis=-1)
|
||||
return w
|
||||
+643
@@ -0,0 +1,643 @@
|
||||
import os
|
||||
import glob
|
||||
import numpy as np
|
||||
import matplotlib.pyplot as plt
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import torch.optim as optim
|
||||
|
||||
from body_model import OP_IGNORE_JOINTS
|
||||
from geometry.mesh import save_mesh_scenes, vertices_to_trimesh
|
||||
from util.logger import Logger, log_cur_stats
|
||||
from util.tensor import move_to, detach_all
|
||||
from vis.output import prep_result_vis, animate_scene
|
||||
|
||||
from .losses import RootLoss, SMPLLoss #, MotionLoss
|
||||
from .output import save_camera_json
|
||||
import time
|
||||
|
||||
LINE_SEARCH = "strong_wolfe"
|
||||
LIGHT_BLUE=(0.65098039, 0.74117647, 0.85882353)
|
||||
|
||||
"""
|
||||
Optimization happens in multiple stages
|
||||
Stage 1: fit root orients and trans of each frame indendpently
|
||||
Stage 2: fit SMPL poses and betas of each frame independently
|
||||
Stage 3: fit poses and roots in same global coordinate frame
|
||||
"""
|
||||
|
||||
|
||||
class StageOptimizer(object):
|
||||
def __init__(
|
||||
self,
|
||||
name,
|
||||
model,
|
||||
param_names,
|
||||
lr=1.0,
|
||||
lbfgs_max_iter=20,
|
||||
save_every=10,
|
||||
vis_every=-1,
|
||||
save_meshes=True,
|
||||
max_chunk_steps=10,
|
||||
**kwargs,
|
||||
):
|
||||
Logger.log(f"INITIALIZING OPTIMIZER {name} for {param_names}")
|
||||
self.name = name
|
||||
self.model = model
|
||||
|
||||
self.set_opt_vars(param_names)
|
||||
|
||||
self.optim = torch.optim.LBFGS(
|
||||
self.opt_params, max_iter=lbfgs_max_iter, lr=lr, line_search_fn=LINE_SEARCH
|
||||
)
|
||||
# LBFGS computes losses multiple times per iteration,
|
||||
# save a dict mapping iteration to list of stats_dicts
|
||||
self.loss_dicts = {}
|
||||
|
||||
self.save_every = save_every
|
||||
self.vis_every = vis_every
|
||||
self.save_meshes = save_meshes
|
||||
self.max_chunk_steps = max_chunk_steps
|
||||
|
||||
self.cur_step = 0
|
||||
|
||||
self.add_chunk = 0
|
||||
self.cur_loss = 0
|
||||
self.prev_loss = np.inf
|
||||
self.last_updated = 0
|
||||
self.reached_max = False
|
||||
self.reached_max_iter = -1
|
||||
|
||||
def set_opt_vars(self, param_names):
|
||||
Logger.log("Set param names:")
|
||||
Logger.log(param_names)
|
||||
|
||||
self.param_names = param_names
|
||||
self.model.params.set_require_grads(self.param_names)
|
||||
self.opt_params = [
|
||||
getattr(self.model.params, name) for name in self.param_names
|
||||
]
|
||||
|
||||
def forward_pass(self, obs_data):
|
||||
raise NotImplementedError
|
||||
|
||||
def load_checkpoint(self, out_dir, device=None):
|
||||
if device is None:
|
||||
device = torch.device("cpu")
|
||||
|
||||
param_path = os.path.join(out_dir, f"{self.name}_params.pth")
|
||||
if os.path.isfile(param_path):
|
||||
param_dict = torch.load(param_path, map_location=device)
|
||||
self.model.params.load_dict(param_dict)
|
||||
Logger.log(f"Params loaded from {param_path}")
|
||||
|
||||
optim_path = os.path.join(out_dir, f"{self.name}_optim.pth")
|
||||
if os.path.isfile(optim_path):
|
||||
optim_dict = torch.load(optim_path)
|
||||
self.optim.load_state_dict(optim_dict["optim"])
|
||||
self.cur_step = optim_dict["cur_step"]
|
||||
Logger.log(f"Optimizer loaded from {optim_path} at iter {self.cur_step}")
|
||||
|
||||
def save_checkpoint(self, out_dir):
|
||||
param_path = os.path.join(out_dir, f"{self.name}_params.pth")
|
||||
param_dict = self.model.params.get_dict()
|
||||
if "world_scale" in param_dict:
|
||||
print("WORLD_SCALE", param_dict["world_scale"].detach().cpu())
|
||||
if "floor_plane" in param_dict:
|
||||
print("FLOOR PLANE", param_dict["floor_plane"].detach().cpu())
|
||||
if "cam_f" in param_dict:
|
||||
print("CAM_F", param_dict["cam_f"].detach().cpu())
|
||||
torch.save(param_dict, param_path)
|
||||
Logger.log(f"Model saved at {param_path}")
|
||||
|
||||
optim_path = os.path.join(out_dir, f"{self.name}_optim.pth")
|
||||
torch.save(
|
||||
{"optim": self.optim.state_dict(), "cur_step": self.cur_step},
|
||||
optim_path,
|
||||
)
|
||||
Logger.log(f"Optimizer saved at {optim_path}")
|
||||
|
||||
def save_results(self, out_dir, seq_name):
|
||||
"""
|
||||
pred dict will be a dictionary of trajectories.
|
||||
each trajectory will have params and lists of trimesh sequences
|
||||
"""
|
||||
os.makedirs(out_dir, exist_ok=True)
|
||||
|
||||
with torch.no_grad():
|
||||
print('go into get_optim_result from optimizer.py save_results')
|
||||
pred_dict = self.model.get_optim_result()
|
||||
pred_dict = move_to(detach_all(pred_dict), "cpu")
|
||||
|
||||
i = self.cur_step
|
||||
for name, results in pred_dict.items():
|
||||
print('save_results, name: ', name)
|
||||
# save parameters of trajectory
|
||||
out_path = f"{out_dir}/{seq_name}_{i:06d}_{name}_results.npz"
|
||||
Logger.log(f"saving params to {out_path}")
|
||||
np.savez(out_path, **results)
|
||||
print('cam_R in save_results and its .npz', results['cam_R'])
|
||||
|
||||
# also save the cameras
|
||||
# print('save the cameras in ptimizer.py save_results')
|
||||
# print(self.model.params)
|
||||
# with torch.no_grad():
|
||||
# cam_R, cam_t = self.model.params.get_extrinsics()
|
||||
# intrins = self.model.params.get_intrinsics()
|
||||
# print('cam_R in 0000.json:', cam_R)
|
||||
# save_camera_json(
|
||||
# f"{out_dir}/{seq_name}_cameras_{self.cur_step:06d}.json",
|
||||
# cam_R.detach().cpu(),
|
||||
# cam_t.detach().cpu(),
|
||||
# intrins.detach().cpu(),
|
||||
# )
|
||||
|
||||
# plot losses
|
||||
self.plot_losses(out_dir)
|
||||
|
||||
# def save_meshes_all(self, res_dir, obs_data, seq_name, num_steps=-1):
|
||||
|
||||
# # check which results are saved
|
||||
# seq_name = obs_data["seq_name"][0]
|
||||
# res_pre = f"{res_dir}/{seq_name}_opt_{self.cur_step:06d}"
|
||||
# with torch.no_grad():
|
||||
# print('go into get_optim_result from optimizer.py vis_result')
|
||||
# pred_dict = self.model.get_optim_result(num_steps=num_steps)
|
||||
|
||||
# res_dict = detach_all(pred_dict["world"])
|
||||
# scene_dict = move_to(
|
||||
# prep_result_vis(
|
||||
# res_dict,
|
||||
# obs_data["vis_mask"],
|
||||
# obs_data["track_id"],
|
||||
# self.model.body_model,
|
||||
# temporal_smooth=True
|
||||
# ),
|
||||
# "cpu",
|
||||
# )
|
||||
|
||||
# # save the meshes
|
||||
# os.makedirs(res_dir, exist_ok=True)
|
||||
# i = self.cur_step
|
||||
# scene_dir = f"{res_dir}/{seq_name}_{i:06d}_meshes"
|
||||
# os.makedirs(scene_dir, exist_ok=True)
|
||||
|
||||
# verts, joints, colors, l_faces, r_faces, is_right, bounds = scene_dict["geometry"]
|
||||
# T = len(verts)
|
||||
# print(f"{T} mesh frames")
|
||||
# times = list(range(0, T, 1))
|
||||
# for t in times:
|
||||
# if len(is_right[t]) > 1:
|
||||
# assert (is_right[t].cpu().numpy().tolist() == [0,1])
|
||||
# # l_meshes = make_batch_mesh(verts[t][0][None], l_faces[t], colors[t][0][None])
|
||||
# # r_meshes = make_batch_mesh(verts[t][1][None], r_faces[t], colors[t][1][None])
|
||||
# # assert len(l_meshes) == 1
|
||||
# # assert len(r_meshes) == 1
|
||||
# tmesh = vertices_to_trimesh(verts[t][0].detach().cpu().numpy(), l_faces[t].detach().cpu().numpy(), LIGHT_BLUE, is_right=0)
|
||||
# tmesh.export(os.path.join(scene_dir, f'{str(t).zfill(6)}_0.obj'))
|
||||
|
||||
# tmesh = vertices_to_trimesh(verts[t][1].detach().cpu().numpy(), r_faces[t].detach().cpu().numpy(), LIGHT_BLUE, is_right=1)
|
||||
# tmesh.export(os.path.join(scene_dir, f'{str(t).zfill(6)}_1.obj'))
|
||||
|
||||
# else:
|
||||
# assert len(is_right[t]) == 1
|
||||
# if is_right[t] == 0:
|
||||
# # meshes = make_batch_mesh(verts[t][0][None], l_faces[t], colors[t][0][None])
|
||||
# tmesh = vertices_to_trimesh(verts[t][0].detach().cpu().numpy(), l_faces[t].detach().cpu().numpy(), LIGHT_BLUE, is_right=0)
|
||||
# tmesh.export(os.path.join(scene_dir, f'{str(t).zfill(6)}_0.obj'))
|
||||
# elif is_right[t] == 1:
|
||||
# # meshes = make_batch_mesh(verts[t][0][None], r_faces[t], colors[t][0][None])
|
||||
# tmesh = vertices_to_trimesh(verts[t][0].detach().cpu().numpy(), r_faces[t].detach().cpu().numpy(), LIGHT_BLUE, is_right=1)
|
||||
# tmesh.export(os.path.join(scene_dir, f'{str(t).zfill(6)}_1.obj'))
|
||||
|
||||
def vis_result(self, res_dir, obs_data, vis=None, num_steps=-1):
|
||||
if vis is None or self.vis_every < 0:
|
||||
return
|
||||
|
||||
# check which results are saved
|
||||
seq_name = obs_data["seq_name"][0]
|
||||
res_pre = f"{res_dir}/{seq_name}_opt_{self.cur_step:06d}"
|
||||
with torch.no_grad():
|
||||
print('go into get_optim_result from optimizer.py vis_result')
|
||||
pred_dict = self.model.get_optim_result(num_steps=num_steps)
|
||||
|
||||
res_dict = detach_all(pred_dict["world"])
|
||||
scene_dict = move_to(
|
||||
prep_result_vis(
|
||||
res_dict,
|
||||
obs_data["vis_mask"],
|
||||
obs_data["track_id"],
|
||||
self.model.body_model,
|
||||
),
|
||||
"cpu",
|
||||
)
|
||||
animate_scene(vis, scene_dict, res_pre, render_views=["src_cam", "above"])
|
||||
|
||||
def log_losses(self, stats_dict):
|
||||
stats_dict = move_to(detach_all(stats_dict), "cpu")
|
||||
log_cur_stats(
|
||||
stats_dict,
|
||||
iter=self.cur_step,
|
||||
to_stdout=(self.cur_step % self.save_every == 0),
|
||||
)
|
||||
for loss_name, loss_val in stats_dict.items():
|
||||
loss_dict = self.loss_dicts.get(loss_name, {})
|
||||
loss_series = loss_dict.get(self.cur_step, [])
|
||||
loss_series.append(loss_val)
|
||||
loss_dict[self.cur_step] = loss_series
|
||||
self.loss_dicts[loss_name] = loss_dict
|
||||
|
||||
def record_current_losses(self, writer):
|
||||
"""
|
||||
record the mean of current step's loss values in tensorboard
|
||||
"""
|
||||
if len(self.loss_dicts) < 1:
|
||||
return
|
||||
|
||||
for loss_name, loss_dict in self.loss_dicts.items():
|
||||
loss_mean = np.mean(loss_dict[self.cur_step])
|
||||
writer.add_scalar(f"{self.name}/{loss_name}", loss_mean, self.cur_step)
|
||||
|
||||
def plot_losses(self, res_dir):
|
||||
"""
|
||||
plot a box plot for each BFGS iteration
|
||||
"""
|
||||
if len(self.loss_dicts) < 1:
|
||||
return
|
||||
for loss_name, loss_dict in self.loss_dicts.items():
|
||||
# times (list len T)
|
||||
# loss vals (list len T of loss value lists)
|
||||
times, loss_vals = zip(*loss_dict.items())
|
||||
plt.figure()
|
||||
plt.boxplot(loss_vals, labels=times, showfliers=False)
|
||||
plt.savefig(f"{res_dir}/{loss_name}.png")
|
||||
|
||||
def run(self, obs_data, num_iters, out_dir, vis=None, writer=None):
|
||||
self.cur_step = 0
|
||||
self.loss.cur_step = 0
|
||||
res_dir = os.path.join(out_dir, self.name)
|
||||
os.makedirs(res_dir, exist_ok=True)
|
||||
seq_name = obs_data["seq_name"][0]
|
||||
print("SEQ NAME", seq_name, "num_iters", num_iters)
|
||||
|
||||
# try to load from checkpoint if exists
|
||||
device = obs_data["joints2d"].device
|
||||
self.load_checkpoint(out_dir, device=device)
|
||||
|
||||
if self.cur_step >= num_iters:
|
||||
Logger.log(f"Current optimizer {self}")
|
||||
Logger.log(f"Checkpoint at {self.cur_step} >= {num_iters}, skipping")
|
||||
return
|
||||
|
||||
Logger.log(f"OPTIMIZING {self.name} FOR {num_iters} ITERATIONS")
|
||||
# time.sleep(5)
|
||||
|
||||
# save initial results and vis
|
||||
print('go into save_results in optimizers.py run()')
|
||||
self.save_results(res_dir, seq_name)
|
||||
|
||||
for i in range(self.cur_step, num_iters):
|
||||
Logger.log("ITER: %d" % (i))
|
||||
|
||||
if (i + 1) % self.save_every == 0: # save before
|
||||
self.save_checkpoint(out_dir)
|
||||
self.save_results(res_dir, seq_name)
|
||||
else:
|
||||
self.save_checkpoint(out_dir)
|
||||
|
||||
if (i + 1) % self.vis_every == 0: # render
|
||||
self.vis_result(res_dir, obs_data, vis)
|
||||
|
||||
self.cur_step = i
|
||||
self.loss.cur_step = i
|
||||
|
||||
#print("ITER: %d" % (i))
|
||||
self.optim_step(obs_data, i, writer)
|
||||
|
||||
# early termination in case of nans
|
||||
if np.isnan(self.cur_loss):
|
||||
# we need to backtrack
|
||||
self.load_checkpoint(out_dir, device=device)
|
||||
print(self.cur_loss, self.kkk, self.ppp)
|
||||
raise ValueError
|
||||
|
||||
# termination for the last chunk
|
||||
if self.reached_max and self.reached_max_iter < 0:
|
||||
self.reached_max_iter = i - 1
|
||||
if self.reached_max and i - self.reached_max_iter >= self.max_chunk_steps:
|
||||
break
|
||||
|
||||
# termination for middle chunks
|
||||
loss_change = self.prev_loss - self.cur_loss
|
||||
if self.last_updated == i - 1 and loss_change == 0:
|
||||
break
|
||||
if (
|
||||
(self.cur_loss < 0 and loss_change < 100)
|
||||
or (i - self.last_updated >= self.max_chunk_steps)
|
||||
or (loss_change < 20 and i - self.last_updated > 5)
|
||||
):
|
||||
self.add_chunk = self.add_chunk + 1
|
||||
self.last_updated = i
|
||||
self.prev_loss = self.cur_loss
|
||||
|
||||
# final save and vis step
|
||||
self.cur_step = num_iters
|
||||
self.save_checkpoint(out_dir)
|
||||
self.save_results(res_dir, seq_name)
|
||||
self.vis_result(res_dir, obs_data, vis)
|
||||
# self.save_meshes_all(res_dir, obs_data, seq_name)
|
||||
|
||||
def optim_step(self, obs_data, i, writer=None):
|
||||
def closure():
|
||||
self.optim.zero_grad()
|
||||
#print('******************************', i)
|
||||
loss, stats_dict, preds = self.forward_pass(obs_data)
|
||||
#print('******************************', i)
|
||||
stats_dict["total"] = loss
|
||||
self.log_losses(move_to(detach_all(stats_dict), "cpu"))
|
||||
self.cur_loss = stats_dict["total"].detach().cpu().item()
|
||||
self.kkk = stats_dict
|
||||
self.ppp = preds
|
||||
loss_keys = stats_dict.keys()
|
||||
if "motion_prior" in loss_keys and "pose_prior" in loss_keys:
|
||||
self.cur_loss = (
|
||||
self.cur_loss
|
||||
- 0.04 * stats_dict["pose_prior"].detach().cpu().item()
|
||||
)
|
||||
print("optimizer print", loss, loss.shape, stats_dict, loss.requires_grad)
|
||||
loss.backward()
|
||||
# print('end of loss backward')
|
||||
return loss
|
||||
|
||||
#print('!!!!!!!!!!!!!!!!!!!!!!!!!!', i)
|
||||
self.optim.step(closure)
|
||||
#print('!!!!!!!!!!!!!!!!!!!!!!!!!!', i)
|
||||
if writer is not None:
|
||||
self.record_current_losses(writer)
|
||||
|
||||
|
||||
class RootOptimizer(StageOptimizer):
|
||||
name = "root_fit"
|
||||
stage = 0
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model,
|
||||
all_loss_weights,
|
||||
use_chamfer=False,
|
||||
robust_loss_type="none",
|
||||
robust_tuning_const=4.6851,
|
||||
joints2d_sigma=100,
|
||||
**kwargs,
|
||||
):
|
||||
param_names = ["trans", "root_orient"]
|
||||
super().__init__(self.name, model, param_names, **kwargs)
|
||||
|
||||
self.loss = RootLoss(
|
||||
all_loss_weights[self.stage],
|
||||
ignore_op_joints=OP_IGNORE_JOINTS,
|
||||
joints2d_sigma=joints2d_sigma,
|
||||
use_chamfer=use_chamfer,
|
||||
robust_loss=robust_loss_type,
|
||||
robust_tuning_const=robust_tuning_const,
|
||||
faces=model.body_model.faces_tensor
|
||||
)
|
||||
|
||||
def forward_pass(self, obs_data):
|
||||
"""
|
||||
Takes in observed data, predicts the smpl parameters and returns loss
|
||||
"""
|
||||
# raise ValueError
|
||||
# Use current params to go through SMPL and get joints3d, verts3d, points3d
|
||||
pred_data = self.model.pred_params_mano(obs_data["is_right"])
|
||||
pred_data["cameras"] = self.model.params.get_cameras()
|
||||
|
||||
# compute data losses only
|
||||
vis_mask = obs_data["vis_mask"] >= 0
|
||||
# print('entering loss from RootOptimizer')
|
||||
loss, stats_dict = self.loss(obs_data, pred_data, vis_mask)
|
||||
return loss, stats_dict, pred_data
|
||||
|
||||
|
||||
class SmoothOptimizer(StageOptimizer):
|
||||
name = "smooth_fit"
|
||||
stage = 1
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model,
|
||||
all_loss_weights,
|
||||
use_chamfer=False,
|
||||
robust_loss_type="none",
|
||||
robust_tuning_const=4.6851,
|
||||
joints2d_sigma=100,
|
||||
**kwargs,
|
||||
):
|
||||
param_names = ["trans", "root_orient", "betas", "latent_pose"]
|
||||
if model.opt_scale:
|
||||
param_names += ["world_scale"]
|
||||
if model.opt_cams:
|
||||
param_names += ["cam_f", "delta_cam_R"]
|
||||
|
||||
super().__init__(self.name, model, param_names, **kwargs)
|
||||
|
||||
self.loss = SMPLLoss(
|
||||
all_loss_weights[self.stage],
|
||||
ignore_op_joints=OP_IGNORE_JOINTS,
|
||||
joints2d_sigma=joints2d_sigma,
|
||||
use_chamfer=use_chamfer,
|
||||
robust_loss=robust_loss_type,
|
||||
robust_tuning_const=robust_tuning_const,
|
||||
faces=model.body_model.faces_tensor
|
||||
)
|
||||
|
||||
def forward_pass(self, obs_data):
|
||||
"""
|
||||
Takes in observed data, predicts the smpl parameters and returns loss
|
||||
"""
|
||||
# raise ValueError
|
||||
# Use current params to go through SMPL and get joints3d, verts3d, points3d
|
||||
pred_data = self.model.pred_params_mano(obs_data["is_right"])
|
||||
pred_data["cameras"] = self.model.params.get_cameras()
|
||||
pred_data.update(self.model.params.get_vars())
|
||||
|
||||
# camera predictions
|
||||
# pts_ij, target, weight = self.model.params.get_reprojected_points()
|
||||
# pred_data["bg2d_err"] = (
|
||||
# (weight * (pts_ij - target) ** 2).sum(dim=(-1, -2)).mean()
|
||||
# )
|
||||
pred_data["cam_R"], pred_data["cam_t"] = self.model.params.get_extrinsics()
|
||||
|
||||
# compute data losses only
|
||||
vis_mask = obs_data["vis_mask"] >= 0
|
||||
loss, stats_dict = self.loss(obs_data, pred_data, self.model.seq_len, vis_mask)
|
||||
return loss, stats_dict, pred_data
|
||||
|
||||
|
||||
# class MotionOptimizer(StageOptimizer):
|
||||
# name = "motion_fit"
|
||||
# stage = 2
|
||||
|
||||
# def __init__(self, model, all_loss_weights, opt_cams=False, **kwargs):
|
||||
# self.opt_cams = opt_cams and model.opt_cams
|
||||
# param_names = [
|
||||
# "trans",
|
||||
# "root_orient",
|
||||
# "latent_pose",
|
||||
# "trans_vel",
|
||||
# "root_orient_vel",
|
||||
# "joints_vel",
|
||||
# "betas",
|
||||
# "latent_motion",
|
||||
# "floor_plane",
|
||||
# ]
|
||||
# if model.opt_scale:
|
||||
# param_names += ["world_scale"]
|
||||
|
||||
# if self.opt_cams:
|
||||
# Logger.log(f"{self.name} OPTIMIZING CAMERAS")
|
||||
# param_names += ["delta_cam_R", "cam_f"]
|
||||
|
||||
# super().__init__(self.name, model, param_names, **kwargs)
|
||||
# self.set_loss(model, all_loss_weights[self.stage], **kwargs)
|
||||
|
||||
# def set_loss(
|
||||
# self,
|
||||
# model,
|
||||
# loss_weights,
|
||||
# use_chamfer=False,
|
||||
# robust_loss_type="none",
|
||||
# robust_tuning_const=4.6851,
|
||||
# joints2d_sigma=100,
|
||||
# **kwargs,
|
||||
# ):
|
||||
# self.loss = MotionLoss(
|
||||
# loss_weights,
|
||||
# init_motion_prior=model.init_motion_prior,
|
||||
# ignore_op_joints=OP_IGNORE_JOINTS,
|
||||
# joints2d_sigma=joints2d_sigma,
|
||||
# use_chamfer=use_chamfer,
|
||||
# robust_loss=robust_loss_type,
|
||||
# robust_tuning_const=robust_tuning_const,
|
||||
# )
|
||||
|
||||
# def get_motion_scale(self):
|
||||
# return 1.0
|
||||
|
||||
# def forward_pass(self, obs_data, num_steps=-1):
|
||||
# p = self.model.params
|
||||
# param_names = [
|
||||
# "betas",
|
||||
# "joints_vel",
|
||||
# "trans_vel",
|
||||
# "root_orient_vel",
|
||||
# ]
|
||||
# param_dict = p.get_vars(param_names)
|
||||
# # param_dict["floor_plane"] = p.floor_plane[p.floor_idcs]
|
||||
|
||||
# preds, world_preds = self.model.rollout_mano_steps(num_steps)
|
||||
# preds.update(param_dict)
|
||||
# world_preds.update(param_dict)
|
||||
|
||||
# # track mask is the length of the sequence that contains num_steps of each track
|
||||
# track_mask = world_preds.get("track_mask", None)
|
||||
# T = track_mask.shape[1] if track_mask is not None else num_steps
|
||||
# world_preds["cameras"] = p.get_cameras(np.arange(T))
|
||||
# if self.opt_cams:
|
||||
# cam_R, cam_t = p.get_extrinsics()
|
||||
# world_preds["cam_R"], world_preds["cam_t"] = cam_R[:T], cam_t[:T]
|
||||
|
||||
# obs_data = slice_dict(obs_data, 0, T)
|
||||
# motion_scale = self.get_motion_scale()
|
||||
# loss, stats_dict = self.loss(
|
||||
# obs_data,
|
||||
# preds,
|
||||
# world_preds,
|
||||
# self.model.seq_len,
|
||||
# valid_mask=track_mask,
|
||||
# init_motion_scale=motion_scale,
|
||||
# )
|
||||
# return loss, stats_dict, (preds, world_preds)
|
||||
|
||||
|
||||
# def slice_dict(d, start, end):
|
||||
# out = d.copy()
|
||||
# for k, v in d.items():
|
||||
# if not isinstance(v, torch.Tensor) or v.ndim < 3 or v.shape[1] < end:
|
||||
# continue
|
||||
# out[k] = v[:, start:end]
|
||||
# return out
|
||||
|
||||
|
||||
# class MotionOptimizerChunks(MotionOptimizer):
|
||||
# """
|
||||
# Incrementally optimize sequence in chunks (with all variables)
|
||||
# :param chunk_size (int) number of frames per chunk
|
||||
# :param init_steps (int) number of optimization steps to spend on first chunk
|
||||
# :param chunk_steps (int) number of opt steps to spend on subsequent chunks
|
||||
# """
|
||||
|
||||
# name = "motion_chunks"
|
||||
# stage = 2
|
||||
|
||||
# def __init__(self, *args, chunk_size=20, init_steps=20, chunk_steps=10, **kwargs):
|
||||
# self.chunk_size = chunk_size
|
||||
# self.init_steps = init_steps
|
||||
# self.chunk_steps = chunk_steps
|
||||
# super().__init__(*args, **kwargs)
|
||||
|
||||
# @property
|
||||
# def num_iters(self):
|
||||
# num_chunks = int(np.ceil(self.model.seq_len / self.chunk_size)) + 1
|
||||
# return self.init_steps + self.chunk_steps * num_chunks
|
||||
|
||||
# @property
|
||||
# def end_idx(self):
|
||||
# if self.cur_step < self.init_steps and self.add_chunk == 0:
|
||||
# return self.chunk_size
|
||||
# chunk_idx = max(
|
||||
# (self.cur_step - self.init_steps) // self.chunk_steps + 1, self.add_chunk
|
||||
# )
|
||||
# if int(np.ceil(self.model.seq_len / self.chunk_size)) == self.add_chunk:
|
||||
# self.reached_max = True
|
||||
# return min(self.chunk_size * (chunk_idx + 1), self.model.seq_len)
|
||||
|
||||
# @property
|
||||
# def start_idx(self):
|
||||
# return 0
|
||||
|
||||
# def get_motion_scale(self):
|
||||
# num_frames = self.end_idx - self.start_idx
|
||||
# return max(1.0, float(self.model.seq_len) / num_frames)
|
||||
|
||||
# def forward_pass(self, obs_data):
|
||||
# return super().forward_pass(obs_data, num_steps=self.end_idx)
|
||||
|
||||
# def vis_result(self, res_dir, obs_data, vis=None, **kwargs):
|
||||
# print("start, end", self.start_idx, self.end_idx)
|
||||
# return super().vis_result(res_dir, obs_data, vis=vis, num_steps=self.end_idx)
|
||||
|
||||
|
||||
# class MotionOptimizerFreeze(MotionOptimizer):
|
||||
# """
|
||||
# Only optimizing a subset of parameters
|
||||
# """
|
||||
|
||||
# name = "motion_freeze"
|
||||
# stage = 2
|
||||
|
||||
# def __init__(self, model, all_loss_weights, no_contacts=True, **kwargs):
|
||||
# loss_weights = all_loss_weights[self.stage].copy()
|
||||
|
||||
# if no_contacts:
|
||||
# loss_weights["contact_height"] = 0.0
|
||||
# loss_weights["contact_vel"] = 0.0
|
||||
|
||||
# param_names = ["latent_motion", "betas", "trans", "floor_plane"]
|
||||
# if model.optim_scale:
|
||||
# param_names += ["world_scale"]
|
||||
|
||||
# StageOptimizer.__init__(self, self.name, model, param_names, **kwargs)
|
||||
# self.set_loss(model, loss_weights, **kwargs)
|
||||
+188
@@ -0,0 +1,188 @@
|
||||
import os
|
||||
import imageio
|
||||
import glob
|
||||
import json
|
||||
import subprocess
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from util.tensor import move_to, detach_all, to_torch
|
||||
from util.logger import Logger
|
||||
|
||||
|
||||
def get_results_paths(res_dir):
|
||||
"""
|
||||
get the iterations of all saved results in res_dir
|
||||
:param res_dir (str) result dir
|
||||
returns a dict of iter to result path
|
||||
"""
|
||||
res_files = sorted(glob.glob(f"{res_dir}/*_results.npz"))
|
||||
print(f"found {len(res_files)} results in {res_dir}")
|
||||
|
||||
path_dict = {}
|
||||
for res_file in res_files:
|
||||
it, name, _ = os.path.basename(res_file).split("_")[-3:]
|
||||
assert name in ["world", "prior"]
|
||||
if it not in path_dict:
|
||||
path_dict[it] = {}
|
||||
path_dict[it][name] = res_file
|
||||
return path_dict
|
||||
|
||||
|
||||
def load_result(res_path_dict):
|
||||
"""
|
||||
load all saved results for a given iteration
|
||||
:param res_path_dict (dict) paths to relevant results
|
||||
returns dict of results
|
||||
"""
|
||||
res_dict = {}
|
||||
for name, path in res_path_dict.items():
|
||||
res = np.load(path)
|
||||
res_dict[name] = to_torch({k: res[k] for k in res.files})
|
||||
return res_dict
|
||||
|
||||
|
||||
def save_initial_predictions(model, out_dir, seq_name):
|
||||
os.makedirs(out_dir, exist_ok=True)
|
||||
with torch.no_grad():
|
||||
print('go into get_optim_result from output.py save_initial_predictions')
|
||||
pred_dict = model.get_optim_result()
|
||||
pred_dict = move_to(detach_all(pred_dict), "cpu")
|
||||
|
||||
for name, results in pred_dict.items():
|
||||
out_path = f"{out_dir}/{seq_name}_{0:06d}_init_{name}_results.npz"
|
||||
Logger.log(f"saving initial results to {out_path}")
|
||||
np.savez(out_path, **results)
|
||||
# print('cam_R in save_initial_predictions: ', results['cam_R'].shape, results['cam_R'])
|
||||
# print((results['cam_R'][0] == results['cam_R'][1]).all())
|
||||
# print((results['cam_R'][1] == results['cam_R'][2]).all())
|
||||
# print((results['cam_R'][2] == results['cam_R'][3]).all())
|
||||
# print((results['cam_R'][3] == results['cam_R'][4]).all())
|
||||
# print((results['cam_R'][4] == results['cam_R'][5]).all())
|
||||
# print((results['cam_R'][5] == results['cam_R'][6]).all())
|
||||
|
||||
|
||||
def save_input_poses(dataset, out_dir, seq_name, name="phalp"):
|
||||
os.makedirs(out_dir, exist_ok=True)
|
||||
dataset.load_data(interp_input=False)
|
||||
d = dataset.data_dict
|
||||
res = {
|
||||
"pose_body": np.stack(d["init_body_pose"], axis=0),
|
||||
"trans": np.stack(d["init_trans"], axis=0),
|
||||
"root_orient": np.stack(d["init_root_orient"], axis=0),
|
||||
}
|
||||
print({k: v.shape for k, v in res.items()})
|
||||
out_path = f"{out_dir}/{seq_name}_{0:06d}_{name}_world_results.npz"
|
||||
Logger.log(f"saving inputs to {out_path}")
|
||||
np.savez(out_path, **res)
|
||||
|
||||
|
||||
def save_input_frames_ffmpeg(dataset, out_dir, name="input", fps=30, overwrite=False):
|
||||
vid_path = f"{out_dir}/{name}.mp4"
|
||||
if not overwrite and os.path.isfile(vid_path):
|
||||
return
|
||||
|
||||
list_path = f"{out_dir}/src_paths.txt"
|
||||
with open(list_path, "w") as f:
|
||||
f.write("\n".join([f"file '{p}'" for p in dataset.sel_img_paths]))
|
||||
|
||||
filter_str = f"-framerate {fps} -vf 'format=yuv420p'"
|
||||
cmd = (
|
||||
f"ffmpeg -loglevel error -f concat -safe 0 "
|
||||
f"-i {list_path} -c:v libx264 {filter_str} {vid_path} -y"
|
||||
)
|
||||
print(cmd)
|
||||
subprocess.call(cmd, shell=True, stdin=subprocess.PIPE)
|
||||
print(f"SAVED {len(dataset.sel_img_paths)} INPUT FRAMES TO {vid_path}")
|
||||
return vid_path
|
||||
|
||||
|
||||
def save_input_frames(dataset, vid_path, fps=30, overwrite=False):
|
||||
if not overwrite and os.path.isfile(vid_path):
|
||||
return
|
||||
|
||||
writer = imageio.get_writer(vid_path, fps=fps)
|
||||
for path in dataset.sel_img_paths:
|
||||
writer.append_data(imageio.imread(path))
|
||||
writer.close()
|
||||
print(f"SAVED {len(dataset.sel_img_paths)} INPUT FRAMES TO {vid_path}")
|
||||
return vid_path
|
||||
|
||||
|
||||
def load_track_info(path):
|
||||
assert os.path.isfile(path)
|
||||
|
||||
with open(path, "r") as f:
|
||||
track_info = json.load(f)
|
||||
|
||||
meta = track_info["meta"]
|
||||
tracks = track_info["tracks"]
|
||||
track_ids, idcs, vis_masks = map(
|
||||
torch.as_tensor,
|
||||
zip(
|
||||
*[
|
||||
(int(tid), info["index"], info["vis_mask"])
|
||||
for tid, info in tracks.items()
|
||||
]
|
||||
),
|
||||
)
|
||||
track_ids = track_ids[idcs]
|
||||
vis_masks = vis_masks[idcs]
|
||||
data_interval = meta["data_interval"]
|
||||
seq_interval = meta["seq_interval"]
|
||||
return track_ids, vis_masks, data_interval, seq_interval
|
||||
|
||||
|
||||
def save_track_info(dataset, out_dir):
|
||||
# track indices are relative to the entire sequence
|
||||
track_ids = dataset.track_ids
|
||||
track_dict = {}
|
||||
for i, (tid, mask) in enumerate(zip(dataset.track_ids, dataset.track_vis_masks)):
|
||||
track_dict[int(tid)] = {"index": i, "vis_mask": mask.tolist()}
|
||||
|
||||
out_dict = {
|
||||
"tracks": track_dict,
|
||||
"meta": {
|
||||
"seq_interval": (int(dataset.start_idx), int(dataset.end_idx)),
|
||||
"data_interval": (int(dataset.data_start), int(dataset.data_end)),
|
||||
},
|
||||
}
|
||||
|
||||
with open(f"{out_dir}/track_info.json", "w") as f:
|
||||
json.dump(out_dict, f)
|
||||
print("SAVED TRACK INFO")
|
||||
|
||||
|
||||
def load_camera_json(path):
|
||||
assert os.path.isfile(path)
|
||||
|
||||
with open(path, "r") as f:
|
||||
cam_data = json.load(f)
|
||||
R = torch.as_tensor(cam_data["rotation"]).reshape(-1, 3, 3)
|
||||
t = torch.as_tensor(cam_data["translation"])
|
||||
intrins = torch.as_tensor(cam_data["intrinsics"])
|
||||
return R.float(), t.float(), intrins.float()
|
||||
|
||||
|
||||
def save_camera_json(path, cam_R, cam_t, intrins):
|
||||
"""
|
||||
:param path
|
||||
:param cam_R (N, 3, 3)
|
||||
:param cam_t (N, 3)
|
||||
:param intrins (N, 4)
|
||||
"""
|
||||
N = len(cam_R)
|
||||
T = torch.tensor([[1, 0, 0], [0, -1, 0], [0, 0, -1]], dtype=torch.float32)
|
||||
cam_R = torch.einsum("ij,bjk->bik", T, cam_R)
|
||||
cam_t = torch.einsum("ij,bj->bi", T, cam_t)
|
||||
with open(path, "w") as f:
|
||||
json.dump(
|
||||
{
|
||||
"rotation": cam_R.reshape(N, 9).tolist(),
|
||||
"translation": cam_t.tolist(),
|
||||
"intrinsics": intrins.tolist(),
|
||||
},
|
||||
f,
|
||||
indent=1,
|
||||
)
|
||||
+174
@@ -0,0 +1,174 @@
|
||||
import numpy as np
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
# from body_model import MANO_JOINTS
|
||||
from geometry.camera import invert_camera, compose_cameras
|
||||
from geometry.rotation import batch_rodrigues
|
||||
from util.logger import Logger
|
||||
from util.tensor import detach_all
|
||||
|
||||
# J_BODY = len(MANO_JOINTS) - 1 # no root
|
||||
# assert J_BODY == 15
|
||||
|
||||
class Params(nn.Module):
|
||||
def __init__(self, batch_size):
|
||||
super().__init__()
|
||||
self.batch_size = batch_size
|
||||
self.param_names = set()
|
||||
|
||||
def set_param(self, name, val, requires_grad=False):
|
||||
print("SETTING PARAM", name, val.shape)
|
||||
with torch.no_grad():
|
||||
setattr(
|
||||
self, name, nn.Parameter(val.contiguous(), requires_grad=requires_grad)
|
||||
)
|
||||
self.param_names.add(name)
|
||||
|
||||
def get_param(self, name):
|
||||
if name not in self.param_names:
|
||||
raise ValueError(f"{name} not stored as opt param")
|
||||
return getattr(self, name)
|
||||
|
||||
def load_dict(self, param_dict):
|
||||
for name, val in param_dict.items():
|
||||
self.set_param(name, val, requires_grad=False)
|
||||
|
||||
def get_dict(self):
|
||||
return {name: self.get_param_item(name) for name in self.param_names}
|
||||
|
||||
def get_vars(self, names=None):
|
||||
if names is None:
|
||||
names = self.param_names
|
||||
return {name: self.get_param(name) for name in names}
|
||||
|
||||
def get_param_item(self, name):
|
||||
with torch.no_grad():
|
||||
param = self.get_param(name)
|
||||
return param.detach()
|
||||
|
||||
def _set_param_grad(self, name, val: bool):
|
||||
if name not in self.param_names:
|
||||
raise ValueError(f"{name} not stored as param")
|
||||
param = getattr(self, name)
|
||||
assert isinstance(param, torch.Tensor)
|
||||
param.requires_grad = val
|
||||
|
||||
def set_require_grads(self, names):
|
||||
"""
|
||||
set parameters in names to True, set all others to False
|
||||
"""
|
||||
for name in self.param_names:
|
||||
self._set_param_grad(name, False)
|
||||
|
||||
for name in names:
|
||||
self._set_param_grad(name, True)
|
||||
|
||||
Logger.log("Set parameter grads:")
|
||||
Logger.log(
|
||||
{name: getattr(self, name).requires_grad for name in self.param_names}
|
||||
)
|
||||
|
||||
|
||||
class CameraParams(Params):
|
||||
"""
|
||||
Parameter container with cameras
|
||||
"""
|
||||
|
||||
def set_cameras(
|
||||
self, cam_data, opt_scale=True, opt_cams=False, opt_focal=False, **kwargs
|
||||
):
|
||||
self.opt_scale = opt_scale
|
||||
self.opt_cams = opt_cams
|
||||
self.opt_focal = opt_focal
|
||||
|
||||
# (T, 3, 3), (T, 3)
|
||||
cam_R, cam_t = cam_data["cam_R"], cam_data["cam_t"]
|
||||
intrins = cam_data["intrins"] # (T, 4)
|
||||
|
||||
T = cam_R.shape[0]
|
||||
device = cam_R.device
|
||||
|
||||
# assume focal length and center are same for all timesteps
|
||||
# print(intrins.shape, cam_R.shape, cam_t.shape)
|
||||
# torch.Size([128, 4]) torch.Size([128, 3, 3]) torch.Size([128, 3])
|
||||
|
||||
self.cam_center = intrins[:, 2:] # (T, 2)
|
||||
cam_f = intrins[:, :2] # (T, 2)
|
||||
if self.opt_focal:
|
||||
self.set_param("cam_f", cam_f)
|
||||
else:
|
||||
self.cam_f = cam_f
|
||||
|
||||
self._cam_R = cam_R # (T, 3, 3)
|
||||
self._cam_t = cam_t # (T, 3)
|
||||
|
||||
world_scale = torch.ones(1, 1, device=device)
|
||||
if self.opt_scale:
|
||||
if "world_scale" in kwargs:
|
||||
world_scale = kwargs["world_scale"]
|
||||
self.set_param("world_scale", world_scale)
|
||||
else:
|
||||
self.world_scale = world_scale
|
||||
|
||||
if self.opt_cams:
|
||||
init_delta_R = torch.zeros(T, 3, device=device)
|
||||
init_delta_t = torch.zeros(T, 3, device=device)
|
||||
if "delta_cam_R" in kwargs:
|
||||
print("setting delta_cam_R from kwargs")
|
||||
init_delta_R = kwargs["delta_cam_R"]
|
||||
if "delta_cam_t" in kwargs:
|
||||
print("setting delta_cam_t from kwargs")
|
||||
init_delta_t = kwargs["delta_cam_t"]
|
||||
|
||||
self.set_param("delta_cam_R", init_delta_R)
|
||||
self.set_param("delta_cam_t", init_delta_t)
|
||||
|
||||
@property
|
||||
def intrins(self): # (4,)
|
||||
return torch.cat([self.cam_f[0], self.cam_center[0]], dim=-1).detach().cpu()
|
||||
|
||||
def get_extrinsics(self):
|
||||
"""
|
||||
returns (T, 3, 3), (T, 3)
|
||||
"""
|
||||
cam_R, cam_t = self._cam_R, self._cam_t
|
||||
print('get_extrinsics')
|
||||
if self.opt_cams:
|
||||
raise ValueError
|
||||
print('self.opt_cams:', self.opt_cams)
|
||||
# raise ValueError
|
||||
dR = batch_rodrigues(self.delta_cam_R)
|
||||
cam_R = torch.matmul(cam_R, dR)
|
||||
cam_t = cam_t + self.delta_cam_t
|
||||
return cam_R, cam_t * self.world_scale
|
||||
|
||||
def get_intrinsics(self):
|
||||
"""
|
||||
returns (T, 4)
|
||||
"""
|
||||
return torch.cat([self.cam_f, self.cam_center], axis=-1)
|
||||
|
||||
def get_cameras(self, idcs=None):
|
||||
"""
|
||||
returns cam_R (B, T, 3, 3) cam_t (B, T, 3), cam_f (T, 2), cam_center (T, 2)
|
||||
"""
|
||||
if idcs is None:
|
||||
idcs = np.arange(self._cam_R.shape[0])
|
||||
|
||||
# print('idcs:', idcs, self._cam_R.shape[0], self._cam_R.shape)
|
||||
cam_R, cam_t = self.get_extrinsics()
|
||||
cam_R, cam_t = cam_R[None, idcs], cam_t[None, idcs]
|
||||
cam_f, cam_center = self.cam_f[idcs], self.cam_center[idcs]
|
||||
|
||||
B = self.batch_size
|
||||
cam_R = cam_R.repeat(B, 1, 1, 1)
|
||||
cam_t = cam_t.repeat(B, 1, 1)
|
||||
# print(cam_R.shape, cam_t.shape, cam_f.shape, cam_center.shape)
|
||||
# print(B, self.batch_size)
|
||||
# raise ValueError
|
||||
# torch.Size([2, 128, 3, 3]) torch.Size([2, 128, 3]) torch.Size([128, 2]) torch.Size([128, 2])
|
||||
# 2 2
|
||||
return cam_R, cam_t, cam_f, cam_center
|
||||
+137
@@ -0,0 +1,137 @@
|
||||
import os
|
||||
import glob
|
||||
import pandas as pd
|
||||
from copy import deepcopy
|
||||
import warnings
|
||||
|
||||
|
||||
def update_args(args):
|
||||
args = deepcopy(args)
|
||||
if args.type == "egobody":
|
||||
if args.root is None:
|
||||
warnings.warn(
|
||||
"args.root flag is set to None, using root directory in source"
|
||||
)
|
||||
args.root = "/path/to/egobody"
|
||||
if args.seqs is None:
|
||||
args.seqs = get_egobody_seqs(args.root, args.split)
|
||||
args.img_dirs = [
|
||||
glob.glob(f"{args.root}/egocentric_color/{seq}/**/PV")[0]
|
||||
for seq in args.seqs
|
||||
]
|
||||
return args
|
||||
|
||||
if args.type == "3dpw":
|
||||
if args.root is None:
|
||||
warnings.warn(
|
||||
"args.root flag is set to None, using root directory in source"
|
||||
)
|
||||
args.root = "/path/to/3DPW"
|
||||
if args.split == "val":
|
||||
args.split = "validation"
|
||||
if args.seqs is None:
|
||||
args.seqs = get_3dpw_seqs(args.root, args.split)
|
||||
args.img_name = "imageFiles"
|
||||
|
||||
elif args.type == "posetrack":
|
||||
if args.root is None:
|
||||
warnings.warn(
|
||||
"args.root flag is set to None, using root directory in source"
|
||||
)
|
||||
args.root = "/path/to/posetrack/posetrack2018/posetrack_data"
|
||||
args.img_name = f"images/{args.split}"
|
||||
|
||||
elif args.type == "davis":
|
||||
if args.root is None:
|
||||
warnings.warn(
|
||||
"args.root flag is set to None, using root directory in source"
|
||||
)
|
||||
args.root = "/path/to/DAVIS"
|
||||
args.img_name = "JPEGImages/Full-Resolution"
|
||||
args.split = ""
|
||||
|
||||
else:
|
||||
assert args.root is not None
|
||||
args.img_name = args.img_name if args.img_name is not None else "images"
|
||||
args.split = ""
|
||||
|
||||
args.img_dirs, args.seqs = get_all_img_dirs(
|
||||
f"{args.root}/{args.img_name}", args.seqs
|
||||
)
|
||||
return args
|
||||
|
||||
|
||||
def isimage(f):
|
||||
ext = os.path.splitext(f)[-1].lower()
|
||||
return ext in [".png", ".jpg"]
|
||||
|
||||
|
||||
def get_all_img_dirs(data_root, seqs=None):
|
||||
"""
|
||||
returns all image directories in root, recursively searching, and the sequence names
|
||||
"""
|
||||
seqs = seqs if seqs is not None else []
|
||||
img_dirs = []
|
||||
|
||||
def get_name(p):
|
||||
return p.removeprefix(f"{data_root}/")
|
||||
|
||||
for root, _, files in os.walk(data_root):
|
||||
num_imgs = len([f for f in files if isimage(f)])
|
||||
if num_imgs > 0:
|
||||
name = get_name(root)
|
||||
if len(seqs) == 0 or name in seqs:
|
||||
img_dirs.append(root)
|
||||
img_dirs = sorted(img_dirs)
|
||||
seqs = [get_name(d) for d in img_dirs]
|
||||
return img_dirs, seqs
|
||||
|
||||
|
||||
def get_custom_seqs(data_root):
|
||||
img_dir = f"{data_root}/images"
|
||||
if not os.path.isdir(img_dir):
|
||||
return []
|
||||
return sorted(os.listdir(img_dir))
|
||||
|
||||
|
||||
def get_egobody_seqs(data_root, split):
|
||||
split_file = f"{data_root}/data_splits.csv"
|
||||
df = pd.read_csv(split_file)
|
||||
if split not in df.columns:
|
||||
print(f"{split} not in {split_file}")
|
||||
return []
|
||||
return sorted(df[split].dropna().tolist())
|
||||
|
||||
|
||||
def get_3dpw_seqs(data_root, split):
|
||||
split_dir = f"{data_root}/sequenceFiles/{split}"
|
||||
if not os.path.isdir(split_dir):
|
||||
return []
|
||||
seq_files = sorted(os.listdir(split_dir))
|
||||
return [os.path.splitext(f)[0] for f in seq_files]
|
||||
|
||||
|
||||
def get_posetrack_seqs(data_root, split):
|
||||
split_dir = f"{data_root}/images/{split}"
|
||||
if not os.path.isdir(split_dir):
|
||||
return []
|
||||
return sorted(os.listdir(split_dir))
|
||||
|
||||
|
||||
def get_davis_seqs(data_root):
|
||||
img_root = f"{data_root}/JPEGImages/Full-Resolution"
|
||||
if not os.path.isdir(img_root):
|
||||
return []
|
||||
return sorted(os.listdir(img_root))
|
||||
|
||||
|
||||
def get_img_dir(data_type, data_root, seq, img_name):
|
||||
if data_type == "posetrack":
|
||||
return f"{data_root}/{img_name}/{seq}"
|
||||
if data_type == "3dpw":
|
||||
return f"{data_root}/imageFiles/{seq}"
|
||||
if data_type == "davis":
|
||||
return f"{data_root}/JPEGImages/Full-Resolution/{seq}"
|
||||
if data_type == "custom":
|
||||
return f"{data_root}/images/{seq}" # custom sequence
|
||||
return None
|
||||
@@ -0,0 +1,124 @@
|
||||
import os
|
||||
import joblib
|
||||
import json
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
|
||||
def export_hamer_predictions(res_path, target_dir):
|
||||
"""
|
||||
exports hamer output results into a directory for each track
|
||||
of JSON files with each frame's mano prediction
|
||||
"""
|
||||
os.makedirs(target_dir, exist_ok=True)
|
||||
print(f"exporting hamer predictions from {res_path} to {target_dir}")
|
||||
|
||||
tracklet_data = joblib.load(res_path)
|
||||
for frame, track_data in tracklet_data.items():
|
||||
name = os.path.splitext(frame.split("/")[-1])[0]
|
||||
print(track_data['tid'])
|
||||
track_dicts = unpack_frame(track_data)
|
||||
for tid, pred_dict in track_dicts.items():
|
||||
track_dir = os.path.join(target_dir, f"{tid:03d}")
|
||||
os.makedirs(track_dir, exist_ok=True)
|
||||
pred_path = os.path.join(track_dir, f"{name}_mano.json")
|
||||
with open(pred_path, "w") as f:
|
||||
json.dump(pred_dict, f, indent=1)
|
||||
|
||||
|
||||
def export_vitpose_keypoints(res_path, target_dir):
|
||||
"""
|
||||
exports each track of each frame into its own json of keypoints
|
||||
"""
|
||||
os.makedirs(target_dir, exist_ok=True)
|
||||
print(f"exporting keypoints from {res_path} to {target_dir}")
|
||||
|
||||
tracklet_data = joblib.load(res_path)
|
||||
for frame, track_data in tracklet_data.items():
|
||||
name = os.path.splitext(frame.split("/")[-1])[0]
|
||||
tids = track_data["tid"]
|
||||
valid_tids = track_data["tracked_ids"]
|
||||
kps_all = track_data["extra_data"]
|
||||
# mask_paths = track_data["mask_name"]
|
||||
for i, tid in enumerate(tids):
|
||||
if tid not in valid_tids:
|
||||
continue
|
||||
track_dir = os.path.join(target_dir, f"{tid:03d}")
|
||||
os.makedirs(track_dir, exist_ok=True)
|
||||
kp_path = os.path.join(track_dir, f"{name}_keypoints.json")
|
||||
kp_dict = {
|
||||
"people": [
|
||||
{
|
||||
"pose_keypoints_2d": kps_all[i],#.tolist(),
|
||||
# "mask_path": mask_paths[i],
|
||||
}
|
||||
],
|
||||
}
|
||||
with open(kp_path, "w") as f:
|
||||
json.dump(kp_dict, f, indent=1)
|
||||
|
||||
|
||||
def export_shot_changes(res_path, target_path):
|
||||
"""
|
||||
Exports the phalp output of shot changes into a JSON file
|
||||
dict with the shot index for each frame in the sequence
|
||||
"""
|
||||
os.makedirs(os.path.dirname(target_path), exist_ok=True)
|
||||
print(f"exporting shot changes to {target_path}")
|
||||
|
||||
tracklet_data = joblib.load(res_path)
|
||||
frames = sorted(tracklet_data.keys())
|
||||
# print(frames)
|
||||
# print(tracklet_data[frames[0]].keys())
|
||||
# for frame in frames:
|
||||
# print(frame, tracklet_data[frame]["shot"])
|
||||
shots = np.cumsum([tracklet_data[frame]["shot"] for frame in frames])
|
||||
shots = shots.astype(int).tolist()
|
||||
shot_dict = {frame.split("/")[-1]: shot for frame, shot in zip(frames, shots)}
|
||||
with open(target_path, "w") as f:
|
||||
json.dump(shot_dict, f, indent=1)
|
||||
|
||||
|
||||
def unpack_frame(track_data):
|
||||
"""
|
||||
Unpack one frame of track data from the phalp output into JSON files for each track.
|
||||
Track data contains body pose, global translation and rotation,
|
||||
and keypoints for all tracks in the frame.
|
||||
Each track is unpacked into folders with a JSON for each frame.
|
||||
"""
|
||||
out_dict = {}
|
||||
for i, tid in enumerate(track_data["tid"]):
|
||||
if tid not in track_data["tracked_ids"]:
|
||||
continue
|
||||
|
||||
cam_trans = track_data["cam_trans"][i].squeeze() # (3,)
|
||||
|
||||
mano = track_data["mano"][i]
|
||||
betas = mano["betas"].squeeze() # (10,)
|
||||
hand_pose = mano["hand_pose"].squeeze() # (23, 3, 3)
|
||||
is_right = mano["is_right"]
|
||||
print(is_right, tid, track_data["tid"], track_data["tracked_ids"])
|
||||
assert is_right == tid, f'{is_right}, {tid}'
|
||||
hand_pose_aa = np.stack(
|
||||
[cv2.Rodrigues(x)[0].squeeze() for x in hand_pose], axis=0
|
||||
) # (23, 3)
|
||||
global_orient = mano["global_orient"].squeeze() # (3, 3)
|
||||
global_orient_aa = cv2.Rodrigues(global_orient)[0].squeeze() # (3,)
|
||||
|
||||
out_dict[tid] = {
|
||||
"betas": betas.tolist(),
|
||||
"body_pose": hand_pose_aa.tolist(),
|
||||
"global_orient": global_orient_aa.tolist(),
|
||||
"cam_trans": cam_trans.tolist(),
|
||||
"is_right": is_right
|
||||
}
|
||||
|
||||
return out_dict
|
||||
|
||||
|
||||
def export_sequence_results(hamer_res_path, track_dir, shot_path):
|
||||
print('exporting sequence results in export_hamer.py...')
|
||||
export_hamer_predictions(hamer_res_path, track_dir)
|
||||
export_vitpose_keypoints(hamer_res_path, track_dir)
|
||||
export_shot_changes(hamer_res_path, shot_path)
|
||||
@@ -0,0 +1,206 @@
|
||||
import argparse
|
||||
import imageio
|
||||
import os
|
||||
import subprocess
|
||||
|
||||
|
||||
def split_frame(videopath,
|
||||
out_dir,
|
||||
fps=30,
|
||||
ext="jpg",
|
||||
down_scale=1,
|
||||
start_sec=0,
|
||||
end_sec=-1,
|
||||
overwrite=False,
|
||||
**kwargs,):
|
||||
import cv2
|
||||
base_name = videopath.replace('.mp4', '').replace('.MP4', '').replace('.avi', '')
|
||||
if os.path.isfile(videopath):
|
||||
print(f'processing {videopath} .....')
|
||||
# filepath = os.path.dirname(videopath)
|
||||
# videoname = os.path.basename(videopath)[:-4]
|
||||
if not os.path.exists(videopath):
|
||||
raise('{} not exists!'.format(videopath))
|
||||
|
||||
vc = cv2.VideoCapture(videopath)
|
||||
|
||||
path = out_dir
|
||||
if not os.path.exists(path):
|
||||
os.makedirs(path)
|
||||
|
||||
if vc.isOpened():
|
||||
rval , frame = vc.read()
|
||||
if '.MOV' not in videopath:
|
||||
cv2.imencode('.jpg',frame)[1].tofile(path +'/' + str(1).zfill(6) + '.jpg')
|
||||
else:
|
||||
cv2.imencode('.jpg',frame[::-1, ::-1])[1].tofile(path +'/' + str(0).zfill(6) + '.jpg')
|
||||
else:
|
||||
rval = False
|
||||
|
||||
Frame = 2
|
||||
timeF = 1
|
||||
while rval:
|
||||
rval,frame = vc.read()
|
||||
if rval==False:
|
||||
break
|
||||
|
||||
if '.MOV' not in videopath:
|
||||
cv2.imencode('.jpg',frame)[1].tofile(path +'/' + str(Frame).zfill(6) + '.jpg')
|
||||
else:
|
||||
cv2.imencode('.jpg',frame[::-1, ::-1])[1].tofile(path +'/' + str(Frame).zfill(6) + '.jpg')
|
||||
|
||||
Frame += 1
|
||||
elif os.path.isdir(videopath):
|
||||
print()
|
||||
print('moving ', base_name , 'to ', out_dir)
|
||||
os.system(f"cp -r {videopath} {out_dir}")
|
||||
raise ValueError
|
||||
return 0
|
||||
|
||||
else:
|
||||
print(os.path.abspath(videopath), os.path.exists(videopath))
|
||||
raise ValueError(f"{videopath} does not exist or is not a valid file or directory.")
|
||||
|
||||
return 0
|
||||
|
||||
def video_to_frames(
|
||||
path,
|
||||
out_dir,
|
||||
fps=30,
|
||||
ext="jpg",
|
||||
down_scale=1,
|
||||
start_sec=0,
|
||||
end_sec=-1,
|
||||
overwrite=False,
|
||||
**kwargs,
|
||||
):
|
||||
"""
|
||||
:param path
|
||||
:param out_dir
|
||||
:param fps
|
||||
:param down_scale (optional int)
|
||||
"""
|
||||
path = str(path)
|
||||
base_name = path.replace('.mp4', '').replace('.MP4', '').replace('.avi', '')
|
||||
print(path, base_name)
|
||||
print(os.path.isfile(path), os.path.isdir(base_name))
|
||||
|
||||
if os.path.isfile(path):
|
||||
os.makedirs(out_dir, exist_ok=True)
|
||||
|
||||
arg_str = f"-copyts -qscale:v 2 -vf fps={fps}"
|
||||
if down_scale != 1:
|
||||
arg_str = f"{arg_str},scale='iw/{down_scale}:ih/{down_scale}'"
|
||||
if start_sec > 0:
|
||||
arg_str = f"{arg_str} -ss {start_sec}"
|
||||
if end_sec > start_sec:
|
||||
arg_str = f"{arg_str} -to {end_sec}"
|
||||
|
||||
yn = "-y" if overwrite else "-n"
|
||||
cmd = f"ffmpeg -i {path} {arg_str} {out_dir}/%06d.{ext} {yn}"
|
||||
print(cmd)
|
||||
elif os.path.isdir(base_name):
|
||||
os.system(f"cp -r {base_name} {out_dir}")
|
||||
return 0
|
||||
else:
|
||||
print(path, os.path.exists(path))
|
||||
raise ValueError
|
||||
|
||||
return subprocess.call(cmd, shell=True, stdin=subprocess.PIPE)
|
||||
|
||||
# def video_to_frames(
|
||||
# path,
|
||||
# out_dir,
|
||||
# fps=30,
|
||||
# ext="jpg",
|
||||
# down_scale=1,
|
||||
# start_sec=0,
|
||||
# end_sec=-1,
|
||||
# overwrite=False,
|
||||
# **kwargs,
|
||||
# ):
|
||||
# """
|
||||
# :param path
|
||||
# :param out_dir
|
||||
# :param fps
|
||||
# :param down_scale (optional int)
|
||||
# """
|
||||
# path = str(path)
|
||||
# base_name = path.replace('.mp4', '').replace('.MP4', '').replace('.avi', '')
|
||||
# print(path, base_name)
|
||||
# print(os.path.isfile(path), os.path.isdir(base_name))
|
||||
|
||||
# if os.path.isfile(path):
|
||||
# os.makedirs(out_dir, exist_ok=True)
|
||||
|
||||
# arg_str = f"-copyts -qscale:v 2 -vf fps={fps}"
|
||||
# if down_scale != 1:
|
||||
# arg_str = f"{arg_str},scale='iw/{down_scale}:ih/{down_scale}'"
|
||||
# if start_sec > 0:
|
||||
# arg_str = f"{arg_str} -ss {start_sec}"
|
||||
# if end_sec > start_sec:
|
||||
# arg_str = f"{arg_str} -to {end_sec}"
|
||||
|
||||
# yn = "-y" if overwrite else "-n"
|
||||
# # Changed %06d to start from 0 instead of 1
|
||||
# cmd = f"ffmpeg -i {path} {arg_str} -start_number 0 {out_dir}/%06d.{ext} {yn}"
|
||||
# print(cmd)
|
||||
# elif os.path.isdir(base_name):
|
||||
# os.system(f"cp -r {base_name} {out_dir}")
|
||||
# return 0
|
||||
# else:
|
||||
# print(path, os.path.exists(path))
|
||||
# raise ValueError
|
||||
|
||||
# return subprocess.call(cmd, shell=True, stdin=subprocess.PIPE)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--data_root", type=str, required=True, help="path to video")
|
||||
parser.add_argument(
|
||||
"--out_root", type=str, required=True, help="output dir for frames"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--seqs",
|
||||
nargs="*",
|
||||
default=None,
|
||||
help="[optional] sequences to run, default runs all available",
|
||||
)
|
||||
parser.add_argument("--fps", type=int, default=30, help="fps to extract frames")
|
||||
parser.add_argument(
|
||||
"--ext", type=str, default="jpg", help="output filetype for frames"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--down_scale", type=int, default=1, help="scale to extract frames"
|
||||
)
|
||||
parser.add_argument(
|
||||
"-ss", "--start_sec", type=float, default=0, help="seconds to start_sec"
|
||||
)
|
||||
parser.add_argument(
|
||||
"-es", "--end_sec", type=float, default=-1, help="seconds to end_sec"
|
||||
)
|
||||
parser.add_argument(
|
||||
"-y", "--overwrite", action="store_true", help="overwrite if already exist"
|
||||
)
|
||||
args = parser.parse_args()
|
||||
seqs_all = os.listdir(args.data_root)
|
||||
if args.seqs is None:
|
||||
args.seqs = seqs_all
|
||||
|
||||
for seq in args.seqs:
|
||||
path = os.path.join(args.data_root, seq)
|
||||
print(f"EXTRACTING FRAMES FROM {path}")
|
||||
assert os.path.isfile(path)
|
||||
seq_name = os.path.splitext(os.path.basename(path.rstrip("/")))[0]
|
||||
out_dir = os.path.join(args.out_root, seq_name)
|
||||
video_to_frames(
|
||||
path,
|
||||
out_dir,
|
||||
args.fps,
|
||||
args.ext,
|
||||
args.down_scale,
|
||||
args.start_sec,
|
||||
args.end_sec,
|
||||
args.overwrite,
|
||||
)
|
||||
@@ -0,0 +1,122 @@
|
||||
import os
|
||||
import subprocess
|
||||
import multiprocessing as mp
|
||||
from concurrent import futures
|
||||
|
||||
from preproc.datasets import update_args
|
||||
from preproc.export_hamer import export_sequence_results
|
||||
|
||||
ROOT_DIR = os.path.abspath(f"{__file__}/../../../")
|
||||
SRC_DIR = os.path.join(ROOT_DIR, "third-party/hamer")
|
||||
|
||||
def launch_hamer(gpus, seq, img_dir, res_dir, name, datatype, overwrite=False):
|
||||
"""
|
||||
run hamer using GPU pool
|
||||
"""
|
||||
cur_proc = mp.current_process()
|
||||
print("PROCESS", cur_proc.name, cur_proc._identity)
|
||||
# 1-indexed processes
|
||||
# worker_id = cur_proc._identity[0] - 1 if len(cur_proc._identity) > 0 else 0
|
||||
# gpu = gpus[worker_id % len(gpus)]
|
||||
gpu = gpus[0]
|
||||
|
||||
HAMER_DIR = SRC_DIR
|
||||
print("HAMER DIR", HAMER_DIR)
|
||||
|
||||
cmd_args = [
|
||||
f"cd {HAMER_DIR};",
|
||||
f"CUDA_VISIBLE_DEVICES={gpu}",
|
||||
"python -u run.py",
|
||||
f"--img_folder {img_dir} ",
|
||||
f"--res_folder {res_dir}/demo_{name}.pkl ",
|
||||
f"--batch_size=48 --side_view --save_mesh --full_frame",
|
||||
f"--type {datatype}",
|
||||
f"--checkpoint {ROOT_DIR}",
|
||||
"--render"
|
||||
]
|
||||
|
||||
cmd = " ".join(cmd_args)
|
||||
print(cmd)
|
||||
return subprocess.call(cmd, shell=True)
|
||||
|
||||
|
||||
def process_seq(
|
||||
gpus,
|
||||
out_root,
|
||||
seq,
|
||||
img_dir,
|
||||
out_name="hamer_out",
|
||||
datatype=None,
|
||||
track_name="track_preds",
|
||||
shot_name="shot_idcs",
|
||||
overwrite=False,
|
||||
):
|
||||
"""
|
||||
Run and export HAMER results
|
||||
"""
|
||||
name = os.path.basename(seq)
|
||||
res_root = f"{out_root}/{out_name}/{seq}"
|
||||
os.makedirs(res_root, exist_ok=True)
|
||||
res_dir = os.path.join(res_root, "results")
|
||||
res_path = f"{res_root}/{name}.pkl"
|
||||
|
||||
if overwrite or not os.path.isfile(res_path):
|
||||
res = launch_hamer(gpus, seq, img_dir, res_dir, name, datatype, overwrite)
|
||||
print(f'rename {res_dir}/demo_{name}.pkl into ', res_path)
|
||||
os.rename(f"{res_dir}/demo_{name}.pkl", res_path)
|
||||
assert res == 0, "HAMER FAILED"
|
||||
|
||||
# export the HAMER predictions
|
||||
track_dir = f"{out_root}/{track_name}/{seq}"
|
||||
shot_path = f"{out_root}/{shot_name}/{seq}.json"
|
||||
|
||||
export_sequence_results(res_path, track_dir, shot_path)
|
||||
return 0
|
||||
|
||||
|
||||
def get_out_dir(src_root, src_dir, src_token, out_token):
|
||||
"""
|
||||
:param src_root (str) root of all data
|
||||
:param src_dir (str) img input dir
|
||||
:param src_token (str) parent name of image input dir
|
||||
:param out_token (str) name of output dir
|
||||
"""
|
||||
src_suffix = src_dir.removeprefix(src_root)
|
||||
out_dir = f"{out_root}/{src_suffix}"
|
||||
return out_dir.replace(src_token, out_token)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import argparse
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--type", default="posetrack", help="dataset to process")
|
||||
parser.add_argument("--root", default=None, help="root dir of data, default None")
|
||||
parser.add_argument("--split", default="val", help="split of dataset, default val")
|
||||
parser.add_argument(
|
||||
"--img_name", default=None, help="input image directory name, default None"
|
||||
)
|
||||
parser.add_argument("--seqs", nargs="*", default=None)
|
||||
parser.add_argument("--gpus", nargs="*", default=[0])
|
||||
parser.add_argument("-y", "--overwrite", action="store_true")
|
||||
|
||||
args = parser.parse_args()
|
||||
args = update_args(args)
|
||||
|
||||
out_root = f"{args.root}/slahmr/{args.split}"
|
||||
|
||||
print(f"running phalp on {len(args.img_dirs)} image directories")
|
||||
if len(args.gpus) > 1:
|
||||
with futures.ProcessPoolExecutor(max_workers=len(args.gpus)) as exe:
|
||||
for img_dir, seq in zip(args.img_dirs, args.seqs):
|
||||
exe.submit(
|
||||
process_seq,
|
||||
args.gpus,
|
||||
out_root,
|
||||
seq,
|
||||
img_dir,
|
||||
overwrite=args.overwrite,
|
||||
)
|
||||
else:
|
||||
for img_dir, seq in zip(args.img_dirs, args.seqs):
|
||||
process_seq(args.gpus, out_root, seq, img_dir, overwrite=args.overwrite)
|
||||
@@ -0,0 +1,186 @@
|
||||
import os
|
||||
import glob
|
||||
import json
|
||||
import numpy as np
|
||||
|
||||
import subprocess
|
||||
from concurrent import futures
|
||||
import multiprocessing as mp
|
||||
|
||||
from preproc.datasets import update_args
|
||||
|
||||
|
||||
SRC_DIR = os.path.abspath(f"{__file__}/../")
|
||||
|
||||
|
||||
def isimage(path):
|
||||
ext = os.path.splitext(path)[-1].lower()
|
||||
return ext == ".png" or ext == ".jpg" or ext == ".jpeg"
|
||||
|
||||
|
||||
def split_frames_equal(img_dir, seq_len=-1):
|
||||
"""
|
||||
split the sequence into subsequences of seq_len
|
||||
returns start and end indices for each subsequence
|
||||
"""
|
||||
if seq_len == -1: # don't split
|
||||
return [(0, -1)]
|
||||
|
||||
image_list = sorted(list(filter(isimage, os.listdir(img_dir))))
|
||||
num_imgs = len(image_list)
|
||||
splits = list(range(0, num_imgs, seq_len))
|
||||
splits[-1] = num_imgs # add remainder to last subseq
|
||||
return list(zip(splits[:-1], splits[1:]))
|
||||
|
||||
|
||||
def split_frames_shots(img_dir, shot_path, pad_shot=0, min_len=0):
|
||||
"""
|
||||
split the sequence into subsequences according by shots in the video
|
||||
returns start and end indices for each subsequence
|
||||
"""
|
||||
if shot_path is None or not os.path.isfile(shot_path):
|
||||
print(f"{shot_path} DOES NOT EXIST, USING WHOLE SEQUENCE")
|
||||
return [(0, -1)], [0]
|
||||
|
||||
image_list = sorted(list(filter(isimage, os.listdir(img_dir))))#[:500]
|
||||
num_frames = len(image_list)
|
||||
|
||||
with open(shot_path, "r") as f:
|
||||
shot_dict = json.load(f)
|
||||
|
||||
# print(shot_dict.keys())
|
||||
shot_idcs = np.array([shot_dict[name] for name in image_list])
|
||||
subseqs = []
|
||||
for i in np.unique(shot_idcs):
|
||||
idcs = np.where(shot_idcs == i)[0]
|
||||
start, end = idcs.min(), idcs.max() + 1
|
||||
if pad_shot > 0:
|
||||
if start > 0:
|
||||
start = start + pad_shot
|
||||
if end < num_frames:
|
||||
end = end - pad_shot
|
||||
if (end - start) < min_len: # this shot is too short
|
||||
continue
|
||||
subseqs.append((start, end))
|
||||
return subseqs, np.unique(shot_idcs)
|
||||
|
||||
|
||||
def split_sequence(args, img_dir, seq):
|
||||
subseqs = split_frames_equal(img_dir, args.seq_len)
|
||||
return subseqs, [None for _ in subseqs]
|
||||
|
||||
|
||||
def get_out_dir(args, seq, shot_idx=None, start=0, end=-1):
|
||||
if "egobody" in args.type or "3dpw" in args.type:
|
||||
if args.use_intrins:
|
||||
out_name = "cameras_intrins"
|
||||
else:
|
||||
out_name = "cameras_default"
|
||||
else:
|
||||
out_name = "cameras"
|
||||
|
||||
if args.seq_len > 0:
|
||||
out_name = f"{out_name}_split"
|
||||
|
||||
out_dir = f"{args.root}/slahmr/{args.split}/{out_name}/{seq}"
|
||||
|
||||
if args.seq_len > 0:
|
||||
return f"{out_dir}/{start}-{end}"
|
||||
|
||||
if args.type == "posetrack" and shot_idx is not None:
|
||||
return f"{out_dir}/shot-{shot_idx}"
|
||||
|
||||
return out_dir
|
||||
|
||||
|
||||
def get_intrins_path(data_type, data_root, seq):
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
def get_command(img_dir, out_dir, start=0, end=-1, intrins_path=None, overwrite=False):
|
||||
cmd_args = [
|
||||
f"python {SRC_DIR}/run_slam.py",
|
||||
"-i",
|
||||
img_dir,
|
||||
"--map_dir",
|
||||
out_dir,
|
||||
"--disable_vis",
|
||||
]
|
||||
if start != 0:
|
||||
cmd_args += ["--start", str(start)]
|
||||
if end != -1:
|
||||
cmd_args += ["--end", str(end)]
|
||||
|
||||
if intrins_path is not None:
|
||||
cmd_args += ["--intrins_path", intrins_path]
|
||||
|
||||
if overwrite:
|
||||
cmd_args.append("--overwrite")
|
||||
|
||||
cmd = " ".join(cmd_args)
|
||||
return cmd
|
||||
|
||||
|
||||
def check_intrins(data_type, data_root, intrins_path, img_dir):
|
||||
assert intrins_path is not None
|
||||
out_name = intrins_path.split(data_root)[1].strip("/").split("/")[0]
|
||||
out_root = f"{data_root}/{out_name}"
|
||||
return intrins_path
|
||||
|
||||
|
||||
def get_slam_command(args, img_dir, seq, shot_idx=None, start=0, end=-1):
|
||||
out_dir = get_out_dir(args, seq, shot_idx, start, end)
|
||||
intrins_path = None
|
||||
if args.use_intrins:
|
||||
intrins_path = get_intrins_path(args.type, args.root, seq)
|
||||
if intrins_path is not None:
|
||||
intrins_path = check_intrins(args.type, args.root, intrins_path, img_dir)
|
||||
return get_command(
|
||||
img_dir,
|
||||
out_dir,
|
||||
start=start,
|
||||
end=end,
|
||||
intrins_path=intrins_path,
|
||||
overwrite=args.overwrite,
|
||||
)
|
||||
|
||||
|
||||
def launch_job(gpus, cmd):
|
||||
cur_proc = mp.current_process()
|
||||
print(cur_proc.name, cur_proc._identity)
|
||||
worker_id = cur_proc._identity[0] - 1
|
||||
gpu = gpus[worker_id % len(gpus)]
|
||||
cmd = f"CUDA_VISIBLE_DEVICES={gpu} {cmd}"
|
||||
print(cmd)
|
||||
subprocess.call(cmd, shell=True)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import argparse
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--type", default="posetrack", help="dataset to process")
|
||||
parser.add_argument(
|
||||
"--root", default=None, help="root of data to process, default None"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--split", default="val", help="split of dataset to process, default val"
|
||||
)
|
||||
parser.add_argument("--gpus", nargs="*", default=[0])
|
||||
parser.add_argument("--seqs", nargs="*", default=None)
|
||||
parser.add_argument("--seq_len", type=int, default=-1)
|
||||
parser.add_argument(
|
||||
"--use_intrins", action="store_true", help="use GT intrinsics if available"
|
||||
)
|
||||
parser.add_argument("-y", "--overwrite", action="store_true")
|
||||
args = parser.parse_args()
|
||||
args = update_args(args)
|
||||
|
||||
print(f"Running SLAM on {len(args.seqs)} sequences")
|
||||
with futures.ProcessPoolExecutor(max_workers=len(args.gpus)) as ex:
|
||||
for img_dir, seq in zip(args.img_dirs, args.seqs):
|
||||
subseqs, shot_idcs = split_sequence(args, img_dir, seq)
|
||||
for shot, (start, end) in zip(shot_idcs, subseqs):
|
||||
cmd = get_slam_command(args, img_dir, seq, shot, start, end)
|
||||
print(cmd)
|
||||
ex.submit(launch_job, args.gpus, cmd)
|
||||
+414
@@ -0,0 +1,414 @@
|
||||
import os
|
||||
|
||||
ROOT_DIR = os.path.abspath(f"{__file__}/../../../")
|
||||
SRC_DIR = os.path.join(ROOT_DIR, "third-party/DROID-SLAM")
|
||||
print("PROJ SRC", ROOT_DIR)
|
||||
print("DROID SRC", SRC_DIR)
|
||||
|
||||
import sys
|
||||
|
||||
sys.path.append(SRC_DIR)
|
||||
sys.path.append(f"{SRC_DIR}/droid_slam")
|
||||
|
||||
import glob
|
||||
import shutil
|
||||
import argparse
|
||||
from tqdm import tqdm
|
||||
|
||||
import json
|
||||
|
||||
import cv2
|
||||
import trimesh
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from lietorch import SE3
|
||||
from droid import Droid
|
||||
import droid_backends
|
||||
|
||||
|
||||
def get_image(image):
|
||||
return image.permute(1, 2, 0).cpu().numpy()[..., ::-1]
|
||||
|
||||
|
||||
def show_image(image):
|
||||
image = image.permute(1, 2, 0).cpu().numpy()
|
||||
cv2.imshow("image", image / 255.0)
|
||||
cv2.waitKey(1)
|
||||
|
||||
|
||||
def get_hwf(test_path):
|
||||
test_img = cv2.imread(test_path)
|
||||
H, W, _ = test_img.shape
|
||||
F = 0.5 * (H + W)
|
||||
return H, W, F
|
||||
|
||||
|
||||
def isimage(path):
|
||||
ext = os.path.splitext(path)[-1].lower()
|
||||
return ext == ".png" or ext == ".jpg" or ext == ".jpeg"
|
||||
|
||||
|
||||
def get_image_files(img_dir, stride, start=0, end=-1):
|
||||
image_list = sorted(list(filter(isimage, os.listdir(img_dir))))
|
||||
end = len(image_list) + 1 + end if end < 0 else end
|
||||
return [os.path.join(img_dir, name) for name in image_list[start:end:stride]]
|
||||
|
||||
|
||||
def load_intrins(intrins_path, image_files, start=0, end=-1):
|
||||
N = len(image_files)
|
||||
H, W, F = get_hwf(image_files[0])
|
||||
# intrinsics for default FOV
|
||||
default = torch.tensor([F, F, W / 2, H / 2, W, H])[None].repeat(N, 1)
|
||||
if intrins_path is None:
|
||||
return default
|
||||
assert os.path.isfile(intrins_path), f"{intrins_path} DOES NOT EXIST"
|
||||
assert intrins_path.endswith(".txt"), f"{intrins_path} INCORRECT EXT"
|
||||
|
||||
print(f"LOADING INTRINSICS FROM {intrins_path}")
|
||||
intrins = torch.from_numpy(np.loadtxt(intrins_path).astype(np.float32))
|
||||
print(intrins.shape, N)
|
||||
if intrins.ndim < 2: # assume intrins is same for all frames
|
||||
intrins = intrins[None].repeat(N, 1)
|
||||
else:
|
||||
end = N + 1 + end if end < 0 else end
|
||||
intrins = intrins[start:end]
|
||||
print(intrins.shape)
|
||||
if intrins.shape[1] == 4: # fx, fy, cx, cy, adding on W and H
|
||||
ones = torch.ones_like(intrins[:, :1])
|
||||
intrins = torch.cat([intrins, ones * W, ones * H], dim=-1)
|
||||
return intrins
|
||||
|
||||
|
||||
def image_stream(image_files, intrins_all):
|
||||
"""image generator"""
|
||||
|
||||
N = len(image_files)
|
||||
|
||||
assert intrins_all.shape[0] == N
|
||||
assert intrins_all.shape[1] >= 4
|
||||
|
||||
for t, imfile in enumerate(image_files):
|
||||
image = cv2.imread(imfile)
|
||||
if image is None:
|
||||
print(imfile, "is none, exiting")
|
||||
sys.exit(1)
|
||||
|
||||
h0, w0, _ = image.shape
|
||||
h1 = int(h0 * np.sqrt((384 * 512) / (h0 * w0)))
|
||||
w1 = int(w0 * np.sqrt((384 * 512) / (h0 * w0)))
|
||||
|
||||
image = cv2.resize(image, (w1, h1))
|
||||
image = image[: h1 - h1 % 8, : w1 - w1 % 8]
|
||||
image = torch.as_tensor(image).permute(2, 0, 1)
|
||||
|
||||
fx, fy, cx, cy = intrins_all[t, :4]
|
||||
sx, sy = w1 / w0, h1 / h0
|
||||
frame_intrins = torch.as_tensor([fx * sx, fy * sy, cx * sx, cy * sy])
|
||||
yield t, image[None], frame_intrins
|
||||
|
||||
|
||||
def unpack_video(video, depth_filter_thresh=0.005):
|
||||
"""
|
||||
extract the video of keyframe disps and poses
|
||||
returns:
|
||||
(T) timestamps
|
||||
(T, 7) world2cam pose for each keyframe
|
||||
(T, H, W) disps
|
||||
(T, H, W, 3) images
|
||||
(T, H, W) mask of valid pixels
|
||||
(4) intrinsics (fx, fy, cx, cy)
|
||||
"""
|
||||
with torch.no_grad():
|
||||
with video.get_lock():
|
||||
# dirty frames have disparity normalized over all previous frames
|
||||
# should probably change this in visualizer
|
||||
t = video.counter.value
|
||||
|
||||
print(f"{t} keyframes in map")
|
||||
|
||||
# normalization happens after all frames have been processed
|
||||
if t < 1:
|
||||
print("no keyframes")
|
||||
return
|
||||
|
||||
# taken from visualization.py#80
|
||||
tstamps = video.tstamp[:t]
|
||||
dirty_index = torch.arange(t, device=video.poses.device)
|
||||
|
||||
poses = torch.index_select(video.poses, 0, dirty_index) # (T, 7)
|
||||
disps = torch.index_select(video.disps, 0, dirty_index) # (T, H, W)
|
||||
images = torch.index_select(video.images, 0, dirty_index) # (T, 3, H, W)
|
||||
|
||||
images = images[:, [2, 1, 0], 3::8, 3::8].permute(0, 2, 3, 1) / 255.0
|
||||
intrins = video.intrinsics[0] # (4,)
|
||||
|
||||
# get mask of static points
|
||||
thresh = depth_filter_thresh * torch.ones_like(disps.mean(dim=[1, 2]))
|
||||
count = droid_backends.depth_filter(
|
||||
video.poses, video.disps, intrins, dirty_index, thresh
|
||||
)
|
||||
disp_mean = disps.mean(dim=(1, 2), keepdim=True).clamp_min(0.2)
|
||||
valid = (count >= 2) & (disps > 0.5 * disp_mean) # (T, H, W)
|
||||
return {
|
||||
"tstamps": tstamps.cpu(),
|
||||
"w2c": poses.cpu(),
|
||||
"disps": disps.cpu(),
|
||||
"images": images.cpu(),
|
||||
"valid": valid.cpu(),
|
||||
"intrins": intrins.cpu(),
|
||||
}
|
||||
|
||||
|
||||
def get_keyframe_map(video_dict):
|
||||
"""
|
||||
:param video_dict (dict) unpacked droid slam video
|
||||
:returns
|
||||
(T, 4, 4) camera to world matrix
|
||||
(N, 3) valid points in world
|
||||
(N, 3) colors of valid points
|
||||
"""
|
||||
poses = video_dict["w2c"]
|
||||
images = video_dict["images"]
|
||||
valid = video_dict["valid"]
|
||||
t = len(poses)
|
||||
c2w = SE3(poses).inv()
|
||||
|
||||
# `iproj` only has a cuda backend; if we pass in CPU tensors, it will silently
|
||||
# return zeros
|
||||
points = droid_backends.iproj(
|
||||
c2w.data.cuda(),
|
||||
video_dict["disps"].cuda(),
|
||||
video_dict["intrins"].cuda(),
|
||||
)
|
||||
valid_pts = torch.cat([points[i, valid[i]] for i in range(t)], dim=0)
|
||||
valid_rgb = torch.cat([images[i, valid[i]] for i in range(t)], dim=0)
|
||||
return (
|
||||
c2w.matrix().detach().cpu(),
|
||||
valid_pts.detach().cpu(),
|
||||
valid_rgb.detach().cpu(),
|
||||
)
|
||||
|
||||
|
||||
def get_frame_cameras(droid, img_paths, intrins_all):
|
||||
N = len(img_paths)
|
||||
|
||||
t = droid.video.counter.value
|
||||
print(f"{t} keyframes in map")
|
||||
|
||||
if t > 1:
|
||||
with torch.no_grad():
|
||||
# localize all frames and get edges into keyframe graph
|
||||
# returns 7D tensor (3D trans, 4D quat)
|
||||
c2w = droid.terminate(image_stream(img_paths, intrins_all))
|
||||
c2w = torch.from_numpy(c2w.astype(np.float32))
|
||||
return SE3(c2w).inv().matrix()
|
||||
|
||||
return torch.eye(4)[None].repeat(N, 1, 1)
|
||||
|
||||
|
||||
def save_camera_json(path, extrins, intrins, yup=False):
|
||||
"""
|
||||
:param path
|
||||
:param extrins (N, 4, 4)
|
||||
:param intrins (N, 4)
|
||||
"""
|
||||
N = extrins.shape[0]
|
||||
if yup:
|
||||
T = torch.tensor(
|
||||
[[1, 0, 0, 0], [0, -1, 0, 0], [0, 0, -1, 0], [0, 0, 0, 1]],
|
||||
dtype=torch.float32,
|
||||
)
|
||||
extrins = torch.matmul(T[None], extrins)
|
||||
with open(path, "w") as f:
|
||||
json.dump(
|
||||
{
|
||||
"rotation": extrins[:, :3, :3].reshape((N, 9)).tolist(),
|
||||
"translation": extrins[:, :3, 3].tolist(),
|
||||
"intrinsics": intrins.tolist(),
|
||||
},
|
||||
f,
|
||||
indent=1,
|
||||
)
|
||||
|
||||
|
||||
def save_pcl(path, points, colors, yup=False):
|
||||
assert len(points) == len(colors)
|
||||
if yup:
|
||||
T = torch.tensor(
|
||||
[[1, 0, 0, 0], [0, -1, 0, 0], [0, 0, -1, 0], [0, 0, 0, 1]],
|
||||
dtype=torch.float32,
|
||||
)
|
||||
points = torch.einsum("ij,...j->...i", T, points) # (*, 3)
|
||||
points = points.reshape(-1, 3).numpy()
|
||||
colors = colors.reshape(-1, 3).numpy()
|
||||
obj = trimesh.Trimesh(vertices=points, vertex_colors=colors)
|
||||
with open(path, "wb") as f:
|
||||
f.write(trimesh.exchange.ply.export_ply(obj))
|
||||
|
||||
|
||||
def save_keyframe_map(out_dir, kf_nodes):
|
||||
"""
|
||||
Saves everything with x right y up z back
|
||||
(converts from x right y down z forward)
|
||||
:param out_dir
|
||||
:param kf_nodes (dict)
|
||||
"""
|
||||
c2w, points, colors = get_keyframe_map(kf_nodes)
|
||||
intrins = kf_nodes["intrins"].detach().cpu()
|
||||
|
||||
print("c2w", c2w.shape)
|
||||
print("points", points.shape)
|
||||
print("colors", colors.shape)
|
||||
|
||||
os.makedirs(out_dir, exist_ok=True)
|
||||
save_pcl(f"{out_dir}/map_points.ply", points, colors)
|
||||
save_camera_json(f"{out_dir}/map_cameras.json", c2w, intrins)
|
||||
print(f"saved {len(c2w)} keyframes to {out_dir}")
|
||||
|
||||
|
||||
def save_cameras(map_dir, frame_w2c, intrins):
|
||||
os.makedirs(map_dir, exist_ok=True)
|
||||
print("SAVING CAMERA INTRINS AND EXTRINS IN", map_dir)
|
||||
|
||||
# save all camera extrinsics and intrinsic
|
||||
W, H = intrins[0, 4], intrins[0, 5]
|
||||
focal = intrins[:, :2].mean()
|
||||
print("INTRINS", intrins[0])
|
||||
np.savez(
|
||||
f"{map_dir}/cameras.npz",
|
||||
height=H,
|
||||
width=W,
|
||||
focal=focal,
|
||||
intrins=intrins[:, :4],
|
||||
w2c=frame_w2c,
|
||||
)
|
||||
|
||||
# save visualization
|
||||
frame_intrins = torch.tensor([focal, focal, W / 2, H / 2, W, H])
|
||||
frame_c2w = torch.linalg.inv(frame_w2c)
|
||||
save_camera_json(f"{map_dir}/frame_cameras.json", frame_c2w, intrins)
|
||||
|
||||
|
||||
def main(args):
|
||||
args.stereo = False
|
||||
torch.multiprocessing.set_start_method("spawn")
|
||||
|
||||
droid = None
|
||||
|
||||
img_paths = get_image_files(
|
||||
args.img_dir, args.stride, start=args.start, end=args.end
|
||||
)
|
||||
if args.map_dir is not None:
|
||||
os.makedirs(args.map_dir, exist_ok=True)
|
||||
camera_file = os.path.join(args.map_dir, "cameras.npz")
|
||||
if os.path.isfile(camera_file) and not args.overwrite:
|
||||
print(glob.glob(f"{args.map_dir}/*.npz"), "already exist, skipping")
|
||||
return
|
||||
|
||||
intrins_all = load_intrins(
|
||||
args.intrins_path, img_paths, start=args.start, end=args.end
|
||||
)
|
||||
print("intrins shape", intrins_all.shape, "num images", len(img_paths))
|
||||
|
||||
for t, image, intrinsics in tqdm(image_stream(img_paths, intrins_all)):
|
||||
if t < args.t0:
|
||||
continue
|
||||
|
||||
if not args.disable_vis:
|
||||
show_image(image[0])
|
||||
|
||||
if droid is None:
|
||||
args.image_size = [image.shape[2], image.shape[3]]
|
||||
droid = Droid(args)
|
||||
|
||||
droid.track(t, image, intrinsics=intrinsics)
|
||||
|
||||
if args.map_dir is None:
|
||||
return
|
||||
|
||||
# save cameras
|
||||
frame_w2c = get_frame_cameras(droid, img_paths, intrins_all)
|
||||
save_cameras(args.map_dir, frame_w2c, intrins_all)
|
||||
|
||||
# save keyframe cameras and points
|
||||
kf_nodes = unpack_video(droid.video)
|
||||
save_keyframe_map(args.map_dir, kf_nodes)
|
||||
|
||||
|
||||
def get_slam_parser():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--stride", default=1, type=int, help="frame stride")
|
||||
parser.add_argument("--t0", default=0, type=int, help="starting frame")
|
||||
parser.add_argument("--weights", default=f"{ROOT_DIR}/_DATA/droid.pth")
|
||||
parser.add_argument("--buffer", type=int, default=512)
|
||||
parser.add_argument("--image_size", default=[240, 320])
|
||||
parser.add_argument("--disable_vis", action="store_true")
|
||||
|
||||
parser.add_argument(
|
||||
"--beta",
|
||||
type=float,
|
||||
default=0.3,
|
||||
help="weight for translation / rotation components of flow",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--filter_thresh",
|
||||
type=float,
|
||||
default=1.0,
|
||||
help="how much motion before considering new keyframe",
|
||||
)
|
||||
parser.add_argument("--warmup", type=int, default=8, help="number of warmup frames")
|
||||
parser.add_argument(
|
||||
"--keyframe_thresh",
|
||||
type=float,
|
||||
default=4.0,
|
||||
help="threshold to create a new keyframe",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--frontend_thresh",
|
||||
type=float,
|
||||
default=16.0,
|
||||
help="add edges between frames whithin this distance",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--frontend_window", type=int, default=25, help="frontend optimization window"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--frontend_radius",
|
||||
type=int,
|
||||
default=2,
|
||||
help="force edges between frames within radius",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--frontend_nms", type=int, default=1, help="non-maximal supression of edges"
|
||||
)
|
||||
|
||||
parser.add_argument("--backend_thresh", type=float, default=22.0)
|
||||
parser.add_argument("--backend_radius", type=int, default=2)
|
||||
parser.add_argument("--backend_nms", type=int, default=3)
|
||||
parser.add_argument("--upsample", action="store_true")
|
||||
return parser
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = get_slam_parser()
|
||||
parser.add_argument(
|
||||
"-i", "--img_dir", required=True, type=str, help="path to image directory"
|
||||
)
|
||||
parser.add_argument(
|
||||
"-o", "--map_dir", type=str, default=None, help="path to output directory"
|
||||
)
|
||||
parser.add_argument(
|
||||
"-y", "--overwrite", action="store_true", help="overwrite existing cameras"
|
||||
)
|
||||
parser.add_argument(
|
||||
"-g", "--save_graph", action="store_true", help="save the frame graph"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--intrins_path", default=None, help="path to camera intrinsics"
|
||||
)
|
||||
parser.add_argument("--start", default=0, type=int, help="start frame")
|
||||
parser.add_argument("--end", default=-1, type=int, help="end frame")
|
||||
|
||||
args = parser.parse_args()
|
||||
main(args)
|
||||
+204
@@ -0,0 +1,204 @@
|
||||
import os
|
||||
import glob
|
||||
import json
|
||||
import subprocess
|
||||
import random
|
||||
import numpy as np
|
||||
|
||||
import torch
|
||||
from torch.utils.data import DataLoader
|
||||
from torch.utils.tensorboard import SummaryWriter
|
||||
|
||||
from data import get_dataset_from_cfg, expand_source_paths
|
||||
|
||||
# from humor.humor_model import HumorModel
|
||||
from optim.base_scene import BaseSceneModel
|
||||
|
||||
from optim.optimizers import (
|
||||
RootOptimizer,
|
||||
SmoothOptimizer
|
||||
)
|
||||
from optim.output import (
|
||||
save_track_info,
|
||||
save_camera_json,
|
||||
save_input_poses,
|
||||
save_initial_predictions,
|
||||
)
|
||||
from vis.viewer import init_viewer
|
||||
from body_model import MANO
|
||||
from util.loaders import (
|
||||
# load_vposer,
|
||||
# load_state,
|
||||
# load_gmm,
|
||||
resolve_cfg_paths,
|
||||
)
|
||||
from util.logger import Logger
|
||||
from util.tensor import get_device, move_to, detach_all, to_torch
|
||||
|
||||
from run_vis import run_vis
|
||||
|
||||
import hydra
|
||||
from omegaconf import DictConfig, OmegaConf
|
||||
import time
|
||||
|
||||
N_STAGES = 3
|
||||
|
||||
import sys
|
||||
sys.path.append('src/human_body_prior')
|
||||
sys.path.append('HMP/')
|
||||
# print(sys.path)
|
||||
# print(torch.cuda.memory_summary())
|
||||
from HMP.fitting import run_prior
|
||||
|
||||
|
||||
def set_seed(seed=42):
|
||||
"""
|
||||
Set random seed for reproducibility
|
||||
"""
|
||||
random.seed(seed)
|
||||
np.random.seed(seed)
|
||||
torch.manual_seed(seed)
|
||||
torch.cuda.manual_seed(seed)
|
||||
torch.cuda.manual_seed_all(seed)
|
||||
torch.backends.cudnn.benchmark = False
|
||||
torch.backends.cudnn.deterministic = True
|
||||
os.environ['PYTHONHASHSEED'] = str(seed)
|
||||
|
||||
def run_opt(cfg, dataset, out_dir, device):
|
||||
a = time.time()
|
||||
args = cfg.data
|
||||
B = len(dataset)
|
||||
T = dataset.seq_len
|
||||
loader = DataLoader(dataset, batch_size=B, shuffle=False)
|
||||
|
||||
obs_data = move_to(next(iter(loader)), device)
|
||||
cam_data = move_to(dataset.get_camera_data(), device)
|
||||
print("Batch size (dataset_length), T (dataset.seq_len): ", B, len(dataset), T)
|
||||
print("OBS DATA", obs_data.keys())
|
||||
print("CAM DATA", cam_data.keys(), 'cam_R', cam_data['cam_R'])
|
||||
|
||||
# save cameras
|
||||
cam_R, cam_t = dataset.cam_data.cam2world()
|
||||
intrins = dataset.cam_data.intrins
|
||||
save_camera_json(f"cameras.json", cam_R, cam_t, intrins)
|
||||
|
||||
# check whether the cameras are static
|
||||
# if static, cannot optimize scale
|
||||
cfg.model.opt_scale &= not dataset.cam_data.is_static
|
||||
Logger.log(f"OPT SCALE {cfg.model.opt_scale}")
|
||||
|
||||
# loss weights for all stages
|
||||
all_loss_weights = cfg.optim.loss_weights
|
||||
assert all(len(wts) == N_STAGES for wts in all_loss_weights.values())
|
||||
stage_loss_weights = [
|
||||
{k: wts[i] for k, wts in all_loss_weights.items()} for i in range(N_STAGES)
|
||||
]
|
||||
max_loss_weights = {k: max(wts) for k, wts in all_loss_weights.items()}
|
||||
|
||||
# load models
|
||||
cfg = resolve_cfg_paths(cfg)
|
||||
cfg.paths.base_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), "../"))
|
||||
paths = cfg.paths
|
||||
Logger.log(f"Loading pose prior from {paths.vposer}")
|
||||
# pose_prior, _ = load_model(paths.vposer, model_code=VPoser, remove_words_in_model_weights='vp_model.', disable_grad=True)
|
||||
# pose_prior = pose_prior.to(device)
|
||||
pose_prior = None
|
||||
|
||||
# load models
|
||||
cfg = resolve_cfg_paths(cfg)
|
||||
paths = cfg.paths
|
||||
Logger.log(f"Loading hand model")
|
||||
# Instantiate MANO model
|
||||
mano_cfg = {k.lower(): v for k,v in dict(cfg.MANO).items()}
|
||||
print('initializing MANO model with cfgs:', mano_cfg)
|
||||
hand_model = MANO(batch_size=B*T, pose2rot=True, **mano_cfg).to(device)
|
||||
|
||||
################################################################
|
||||
######################## optimization ##########################
|
||||
################################################################
|
||||
margs = cfg.model
|
||||
base_model = BaseSceneModel(
|
||||
B, T, hand_model, pose_prior, **margs
|
||||
)
|
||||
|
||||
base_model.initialize(obs_data, cam_data)
|
||||
base_model.to(device)
|
||||
|
||||
# save initial results for later visualization
|
||||
save_input_poses(dataset, os.path.join(out_dir, "hamer"), args.seq)
|
||||
save_initial_predictions(base_model, os.path.join(out_dir, "init"), args.seq)
|
||||
|
||||
opts = cfg.optim.options
|
||||
vis_scale = 0.25
|
||||
vis = None
|
||||
if opts.vis_every > 0:
|
||||
vis = init_viewer(
|
||||
dataset.img_size,
|
||||
cam_data["intrins"][0],
|
||||
vis_scale=vis_scale,
|
||||
bg_paths=dataset.sel_img_paths,
|
||||
fps=cfg.fps,
|
||||
)
|
||||
print("OPTIMIZER OPTIONS:", opts)
|
||||
|
||||
writer = SummaryWriter(out_dir)
|
||||
|
||||
print('start optimization')
|
||||
a = time.time()
|
||||
optim = RootOptimizer(base_model, stage_loss_weights, **opts)
|
||||
optim.run(obs_data, cfg.optim.root.num_iters, out_dir, vis, writer)
|
||||
|
||||
args = cfg.optim.smooth
|
||||
print(args)
|
||||
b = time.time()
|
||||
print('root optimization time: ', b - a)
|
||||
optim = SmoothOptimizer(
|
||||
base_model, stage_loss_weights, opt_scale=args.opt_scale, **opts
|
||||
)
|
||||
optim.run(obs_data, args.num_iters, out_dir, vis, writer)
|
||||
c = time.time()
|
||||
print('Smooth optimization time: ', c - b)
|
||||
|
||||
# HMP
|
||||
if cfg.run_prior and not glob.glob(os.path.join(out_dir, 'prior', '*_world_results.npz')):
|
||||
run_prior(cfg, dataset, out_dir, device, ['smooth_fit'], \
|
||||
obs_data, hand_model, cfg, cfg.data, os.path.join(out_dir, 'prior'))
|
||||
d = time.time()
|
||||
print('prior optimization time: ', d-c)
|
||||
|
||||
|
||||
@hydra.main(version_base=None, config_path="confs", config_name="config.yaml")
|
||||
def main(cfg: DictConfig):
|
||||
OmegaConf.register_new_resolver("eval", eval)
|
||||
print('run_opt.py: ', cfg)
|
||||
|
||||
# Set random seed
|
||||
set_seed(cfg.get('seed', 42))
|
||||
|
||||
out_dir = os.getcwd()
|
||||
print("out_dir", out_dir)
|
||||
Logger.init(f"{out_dir}/opt_log.txt")
|
||||
|
||||
# make sure we get all necessary inputs
|
||||
print("init SOURCES", cfg.data.sources)
|
||||
cfg.data.sources = expand_source_paths(cfg.data.sources)
|
||||
print("SOURCES", cfg.data.sources)
|
||||
|
||||
dataset = get_dataset_from_cfg(cfg)
|
||||
save_track_info(dataset, out_dir)
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = str(cfg.get("gpu"))
|
||||
print("CUDA_VISIBLE_DEVICES", os.environ["CUDA_VISIBLE_DEVICES"])
|
||||
device_id = cfg.get("gpu")
|
||||
|
||||
if cfg.run_opt:
|
||||
device = get_device(device_id)
|
||||
run_opt(cfg, dataset, out_dir, device)
|
||||
|
||||
if cfg.run_vis:
|
||||
run_vis(
|
||||
cfg, dataset, out_dir, device_id, **cfg.get("vis", dict())
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
+475
@@ -0,0 +1,475 @@
|
||||
import os
|
||||
import glob
|
||||
|
||||
import imageio
|
||||
import numpy as np
|
||||
import torch
|
||||
from torch.utils.data import DataLoader
|
||||
from body_model import MANO
|
||||
from omegaconf import DictConfig, OmegaConf
|
||||
|
||||
from data import get_dataset_from_cfg, expand_source_paths
|
||||
from optim.output import (
|
||||
get_results_paths,
|
||||
load_result,
|
||||
save_input_frames,
|
||||
save_input_poses,
|
||||
)
|
||||
from util.loaders import load_config_from_log, resolve_cfg_paths
|
||||
from util.tensor import get_device, move_to, detach_all, to_torch
|
||||
from vis.output import prep_result_vis, animate_scene, make_video_grid_2x2
|
||||
from vis.tools import vis_keypoints
|
||||
from vis.viewer import init_viewer
|
||||
import time
|
||||
from geometry.mesh import save_mesh_scenes, vertices_to_trimesh
|
||||
LIGHT_BLUE=(0.65098039, 0.74117647, 0.85882353)
|
||||
|
||||
def save_meshes_all(cfg, dataset, res_dicts, dev_id, mesh_dirs, num_steps=-1):
|
||||
B = len(dataset)
|
||||
T = dataset.seq_len
|
||||
loader = DataLoader(dataset, batch_size=B, shuffle=False)
|
||||
device = get_device(dev_id)
|
||||
obs_data = move_to(next(iter(loader)), device)
|
||||
|
||||
# load models
|
||||
cfg = resolve_cfg_paths(cfg)
|
||||
# Instantiate MANO model
|
||||
mano_cfg = {k.lower(): v for k,v in dict(cfg.MANO).items()}
|
||||
print('initializing MANO model with cfgs:', mano_cfg, 'B*T', B*T)
|
||||
hand_model = MANO(batch_size=B*T, pose2rot=True, **mano_cfg).to(device)
|
||||
|
||||
for res_dict, mesh_dir in zip(res_dicts, mesh_dirs):
|
||||
res_dict = move_to(res_dict, device)
|
||||
scene_dict = move_to(
|
||||
prep_result_vis(
|
||||
res_dict,
|
||||
obs_data["vis_mask"],
|
||||
obs_data["track_id"],
|
||||
hand_model,
|
||||
temporal_smooth=cfg.temporal_smooth,
|
||||
smooth_trans=True # For mesh export, smooth everything
|
||||
),
|
||||
"cpu",
|
||||
)
|
||||
|
||||
scene_dir = mesh_dir
|
||||
verts, joints, colors, l_faces, r_faces, is_right, bounds = scene_dict["geometry"]
|
||||
T = len(verts)
|
||||
print(f"{T} mesh frames for ", mesh_dir)
|
||||
times = list(range(0, T, 1))
|
||||
flag = False
|
||||
for t in times:
|
||||
if len(is_right[t]) > 1:
|
||||
flag = True
|
||||
vv = t
|
||||
|
||||
if flag:
|
||||
init_trans = (joints[vv][0][9].clone() + joints[vv][1][9].clone()) / 2
|
||||
else:
|
||||
init_trans = joints[0][0][9].clone()
|
||||
|
||||
for t in times:
|
||||
if len(is_right[t]) > 1:
|
||||
assert (is_right[t].cpu().numpy().tolist() == [0,1])
|
||||
# l_meshes = make_batch_mesh(verts[t][0][None], l_faces[t], colors[t][0][None])
|
||||
# r_meshes = make_batch_mesh(verts[t][1][None], r_faces[t], colors[t][1][None])
|
||||
# assert len(l_meshes) == 1
|
||||
# assert len(r_meshes) == 1
|
||||
|
||||
verts[t][0] -= init_trans
|
||||
joints[t][0] -= init_trans
|
||||
tmesh = vertices_to_trimesh(verts[t][0].detach().cpu().numpy(), l_faces[t].detach().cpu().numpy(), LIGHT_BLUE, is_right=0)
|
||||
tmesh.export(os.path.join(scene_dir, f'{str(t).zfill(6)}_0.obj'))
|
||||
# np.save(f'{str(t).zfill(6)}_0.npy', joints[t][0].detach().cpu().numpy())
|
||||
|
||||
verts[t][1] -= init_trans
|
||||
joints[t][1] -= init_trans
|
||||
tmesh = vertices_to_trimesh(verts[t][1].detach().cpu().numpy(), r_faces[t].detach().cpu().numpy(), LIGHT_BLUE, is_right=1)
|
||||
tmesh.export(os.path.join(scene_dir, f'{str(t).zfill(6)}_1.obj'))
|
||||
# np.save(f'{str(t).zfill(6)}_1.npy', joints[t][1].detach().cpu().numpy())
|
||||
|
||||
else:
|
||||
assert len(is_right[t]) == 1
|
||||
if is_right[t] == 0:
|
||||
verts[t][0] -= init_trans
|
||||
joints[t][0] -= init_trans
|
||||
tmesh = vertices_to_trimesh(verts[t][0].detach().cpu().numpy(), l_faces[t].detach().cpu().numpy(), LIGHT_BLUE, is_right=0)
|
||||
tmesh.export(os.path.join(scene_dir, f'{str(t).zfill(6)}_0.obj'))
|
||||
# np.save(f'{str(t).zfill(6)}_0.npy', joints[t][0].detach().cpu().numpy())
|
||||
|
||||
elif is_right[t] == 1:
|
||||
verts[t][0] -= init_trans
|
||||
joints[t][0] -= init_trans
|
||||
tmesh = vertices_to_trimesh(verts[t][0].detach().cpu().numpy(), r_faces[t].detach().cpu().numpy(), LIGHT_BLUE, is_right=1)
|
||||
tmesh.export(os.path.join(scene_dir, f'{str(t).zfill(6)}_1.obj'))
|
||||
# np.save(f'{str(t).zfill(6)}_1.npy', joints[t][0].detach().cpu().numpy())
|
||||
|
||||
|
||||
def run_vis(
|
||||
cfg,
|
||||
dataset,
|
||||
out_dir,
|
||||
dev_id,
|
||||
phases=["smooth_fit"],
|
||||
render_views=["src_cam", "above", "side"],
|
||||
make_grid=True,
|
||||
overwrite=False,
|
||||
save_dir=None,
|
||||
render_kps=False,
|
||||
render_layers=False,
|
||||
save_frames=False,
|
||||
**kwargs
|
||||
):
|
||||
save_dir = out_dir if save_dir is None else save_dir
|
||||
print("OUT_DIR", out_dir)
|
||||
print("SAVE_DIR", save_dir)
|
||||
print("VISUALIZING PHASES", phases)
|
||||
print("RENDERING VIEWS", render_views)
|
||||
print("RENDER_KPS", render_kps)
|
||||
print("OVERWRITE", overwrite)
|
||||
|
||||
# save input frames
|
||||
inp_vid_path = save_input_frames(
|
||||
dataset,
|
||||
f"{save_dir}/{dataset.seq_name}_input.mp4",
|
||||
fps=cfg.fps,
|
||||
overwrite=True,
|
||||
)
|
||||
|
||||
if render_kps:
|
||||
render_keypoints_2d(dataset, save_dir, overwrite=overwrite)
|
||||
|
||||
if len(render_views) < 1:
|
||||
return
|
||||
|
||||
out_ext = "/" if render_layers or save_frames else ".mp4"
|
||||
phase_results = {}
|
||||
phase_max_iters = {}
|
||||
for phase in phases:
|
||||
res_dir = os.path.join(out_dir, phase)
|
||||
if phase == "input":
|
||||
res = get_input_dict(dataset)
|
||||
it = f"{0:06d}"
|
||||
|
||||
elif os.path.isdir(res_dir):
|
||||
res_path_dict = get_results_paths(res_dir)
|
||||
print(f"FOUND {len(res_path_dict)} results in {res_dir}")
|
||||
it = sorted(res_path_dict.keys())[-1]
|
||||
res = load_result(res_path_dict[it])["world"]
|
||||
|
||||
else:
|
||||
print(f"{res_dir} does not exist, skipping")
|
||||
continue
|
||||
|
||||
out_name = f"{save_dir}/{dataset.seq_name}_{phase}_final_{it}"
|
||||
mesh_dir = f"{save_dir}/{phase}/{dataset.seq_name}_{it}_meshes"
|
||||
os.makedirs(mesh_dir, exist_ok=True)
|
||||
phase_max_iters[phase] = it
|
||||
|
||||
out_paths = [f"{out_name}_{view}{out_ext}" for view in render_views]
|
||||
if not overwrite and all(os.path.exists(p) for p in out_paths):
|
||||
print("FOUND OUT PATHS", out_paths)
|
||||
continue
|
||||
|
||||
phase_results[phase] = out_name, mesh_dir, res
|
||||
|
||||
if len(phase_results) > 0:
|
||||
out_names, mesh_dir, res_dicts = zip(*phase_results.values())
|
||||
render_results(
|
||||
cfg,
|
||||
dataset,
|
||||
dev_id,
|
||||
res_dicts,
|
||||
out_names,
|
||||
render_views=render_views,
|
||||
render_layers=render_layers,
|
||||
save_frames=save_frames,
|
||||
**kwargs,
|
||||
)
|
||||
save_meshes_all(cfg, dataset, res_dicts, dev_id, mesh_dir)
|
||||
|
||||
if make_grid:
|
||||
for phase, it in phase_max_iters.items():
|
||||
grid_path = f"{save_dir}/{dataset.seq_name}_{phase}_grid.mp4"
|
||||
vid_paths = [
|
||||
f"{save_dir}/{dataset.seq_name}_{phase}_final_{it}_src_cam.mp4",
|
||||
f"{save_dir}/{dataset.seq_name}_{phase}_final_{it}_front.mp4",
|
||||
f"{save_dir}/{dataset.seq_name}_{phase}_final_{it}_above.mp4",
|
||||
f"{save_dir}/{dataset.seq_name}_{phase}_final_{it}_side.mp4",
|
||||
]
|
||||
make_video_grid_2x2(
|
||||
grid_path,
|
||||
vid_paths,
|
||||
fps=cfg.fps,
|
||||
overwrite=True,
|
||||
)
|
||||
|
||||
|
||||
def get_input_dict(dataset):
|
||||
dataset.load_data(interp_input=False)
|
||||
d = dataset.data_dict
|
||||
input_params = {
|
||||
"pose_body": np.stack(d["init_body_pose"], axis=0),
|
||||
"trans": np.stack(d["init_trans"], axis=0),
|
||||
"root_orient": np.stack(d["init_root_orient"], axis=0),
|
||||
}
|
||||
input_params = to_torch(input_params)
|
||||
print({k: v.shape for k, v in input_params.items()})
|
||||
return input_params
|
||||
|
||||
|
||||
def render_keypoints_2d(dataset, save_dir, overwrite=False):
|
||||
"""
|
||||
render 2d keypoints for each track
|
||||
"""
|
||||
dataset.load_data()
|
||||
out_dir = f"{save_dir}/{dataset.seq_name}_joints2d"
|
||||
B, T = dataset.n_tracks, dataset.seq_len
|
||||
if not overwrite and (os.path.isdir(out_dir) and len(os.listdir(out_dir)) >= B * T):
|
||||
print(f"Keypoints already rendered in {out_dir}")
|
||||
return
|
||||
|
||||
os.makedirs(out_dir, exist_ok=True)
|
||||
for i, tid in enumerate(dataset.track_ids):
|
||||
joints2d = dataset.data_dict["joints2d"][i] # (T, J, 3)
|
||||
for t, sel_img_name in enumerate(dataset.sel_img_names):
|
||||
img = vis_keypoints(joints2d[t : t + 1], dataset.img_size)
|
||||
out_path = f"{out_dir}/{sel_img_name}_{tid}.png"
|
||||
imageio.imwrite(out_path, img)
|
||||
|
||||
|
||||
|
||||
def render_results(cfg, dataset, dev_id, res_dicts, out_names, **kwargs):
|
||||
"""
|
||||
render results for all selected phases
|
||||
"""
|
||||
assert len(res_dicts) == len(out_names)
|
||||
if len(res_dicts) < 1:
|
||||
print("no results to render, skipping")
|
||||
return
|
||||
|
||||
B = len(dataset)
|
||||
T = dataset.seq_len
|
||||
loader = DataLoader(dataset, batch_size=B, shuffle=False)
|
||||
|
||||
device = get_device(dev_id)
|
||||
obs_data = move_to(next(iter(loader)), device)
|
||||
cam_data = dataset.get_camera_data()
|
||||
|
||||
# load models
|
||||
cfg = resolve_cfg_paths(cfg)
|
||||
# Instantiate MANO model
|
||||
mano_cfg = {k.lower(): v for k,v in dict(cfg.MANO).items()}
|
||||
print('initializing MANO model with cfgs:', mano_cfg, 'B*T', B*T)
|
||||
hand_model = MANO(batch_size=B*T, pose2rot=True, **mano_cfg).to(device)
|
||||
vis = init_viewer(
|
||||
dataset.img_size,
|
||||
cam_data["intrins"][0],
|
||||
vis_scale=1.0,
|
||||
bg_paths=dataset.sel_img_paths,
|
||||
fps=cfg.fps,
|
||||
)
|
||||
|
||||
# Set 2D keypoints for overlay visualization if render_keypoints is enabled
|
||||
if kwargs.get('render_keypoints', False):
|
||||
# Load 2D keypoints from dataset - assume single track for now
|
||||
# joints2d is list of (T, J, 3) arrays, one per track
|
||||
dataset.load_data()
|
||||
if len(dataset.data_dict["joints2d"]) > 0:
|
||||
joints2d = dataset.data_dict["joints2d"][0] # First track: (T, J, 3)
|
||||
keypoints_seq = [joints2d[t] for t in range(T)] # List of (J, 3) arrays
|
||||
vis.set_keypoints_seq(keypoints_seq)
|
||||
print(f"Loaded {T} frames of 2D keypoints with {joints2d.shape[1]} joints")
|
||||
|
||||
save_paths_all = []
|
||||
render_views = kwargs.get('render_views', ['src_cam', 'above', 'side'])
|
||||
|
||||
# Separate src_cam from other views
|
||||
src_cam_views = [v for v in render_views if v == 'src_cam']
|
||||
other_views = [v for v in render_views if v != 'src_cam']
|
||||
|
||||
for res_dict, out_name in zip(res_dicts, out_names):
|
||||
print(f'preparing results for rendering {out_name}')
|
||||
res_dict = move_to(res_dict, device)
|
||||
|
||||
# Render src_cam WITH temporal smoothing but WITHOUT trans smoothing
|
||||
if src_cam_views:
|
||||
print(f"Rendering src_cam view WITH temporal smoothing (excluding trans)")
|
||||
scene_dict_no_smooth = prep_result_vis(
|
||||
res_dict,
|
||||
obs_data["vis_mask"],
|
||||
obs_data["track_id"],
|
||||
hand_model,
|
||||
temporal_smooth=cfg.temporal_smooth, # Apply smoothing
|
||||
smooth_trans=False # But don't smooth translation
|
||||
)
|
||||
|
||||
# Compute predicted 2D keypoints using THE SAME function as optimization
|
||||
if kwargs.get('render_keypoints', False):
|
||||
from geometry import camera as cam_util
|
||||
from body_model import run_mano
|
||||
|
||||
# Get the same data as optimization
|
||||
joints3d_op = run_mano(
|
||||
hand_model,
|
||||
res_dict["trans"],
|
||||
res_dict["root_orient"],
|
||||
res_dict["pose_body"],
|
||||
res_dict["is_right"],
|
||||
res_dict.get("betas", None),
|
||||
)["joints"] # (B, T, J, 3)
|
||||
|
||||
# Get cameras (same as optimization)
|
||||
cam_R = res_dict["cam_R"] # (B, T, 3, 3)
|
||||
cam_t = res_dict["cam_t"] # (B, T, 3)
|
||||
intrins = res_dict["intrins"] # (4,) or (T, 4) or (B, T, 4)
|
||||
|
||||
# Ensure intrins is (B, T, 4)
|
||||
if intrins.ndim == 1: # (4,)
|
||||
intrins = intrins[None, None].expand(cam_R.shape[0], cam_R.shape[1], -1) # (B, T, 4)
|
||||
elif intrins.ndim == 2: # (T, 4)
|
||||
intrins = intrins[None].expand(cam_R.shape[0], -1, -1) # (B, T, 4)
|
||||
|
||||
cam_f = intrins[:, :, :2] # (B, T, 2)
|
||||
cam_center = intrins[:, :, 2:] # (B, T, 2)
|
||||
|
||||
# Reproject using EXACTLY the same function as optimization
|
||||
joints2d_pred = cam_util.reproject(joints3d_op, cam_R, cam_t, cam_f[0], cam_center[0]) # (B, T, J, 2)
|
||||
|
||||
# Convert to list for first track
|
||||
pred_keypoints_seq = [joints2d_pred[0, t].cpu().numpy() for t in range(joints2d_pred.shape[1])]
|
||||
vis.set_pred_keypoints_seq(pred_keypoints_seq)
|
||||
print(f"Computed {len(pred_keypoints_seq)} frames of predicted 2D keypoints using reproject()")
|
||||
|
||||
# Render only src_cam view
|
||||
kwargs_src = kwargs.copy()
|
||||
kwargs_src['render_views'] = src_cam_views
|
||||
save_paths_src = animate_scene(
|
||||
vis, scene_dict_no_smooth, out_name, seq_name=dataset.seq_name, **kwargs_src
|
||||
)
|
||||
save_paths_all.append(save_paths_src)
|
||||
|
||||
# Render other views WITH temporal smoothing (including trans)
|
||||
if other_views:
|
||||
print(f"Rendering {other_views} views WITH temporal smoothing (including trans)")
|
||||
scene_dict_smooth = prep_result_vis(
|
||||
res_dict,
|
||||
obs_data["vis_mask"],
|
||||
obs_data["track_id"],
|
||||
hand_model,
|
||||
temporal_smooth=cfg.temporal_smooth, # Apply smoothing for other views
|
||||
smooth_trans=True # Smooth translation for other views
|
||||
)
|
||||
|
||||
# Render other views
|
||||
kwargs_other = kwargs.copy()
|
||||
kwargs_other['render_views'] = other_views
|
||||
kwargs_other['render_keypoints'] = False # No keypoints on other views
|
||||
save_paths_other = animate_scene(
|
||||
vis, scene_dict_smooth, out_name, seq_name=dataset.seq_name, **kwargs_other
|
||||
)
|
||||
if not src_cam_views:
|
||||
save_paths_all.append(save_paths_other)
|
||||
|
||||
vis.close()
|
||||
|
||||
return save_paths_all
|
||||
|
||||
|
||||
def visualize_log(log_dir, dev_id, phases, save_dir=None, **kwargs):
|
||||
print(log_dir)
|
||||
cfg = load_config_from_log(log_dir)
|
||||
print(cfg)
|
||||
print(cfg.data)
|
||||
print(cfg.data.sources)
|
||||
|
||||
# make sure we get all necessary inputs
|
||||
cfg.data.sources = expand_source_paths(cfg.data.sources)
|
||||
print("SOURCES", cfg.data.sources)
|
||||
# cfg.data.track_ids = "001"
|
||||
dataset = get_dataset_from_cfg(cfg)
|
||||
if len(dataset) < 1:
|
||||
print(f"No tracks in dataset, skipping")
|
||||
return
|
||||
|
||||
run_vis(cfg, dataset, log_dir, dev_id, phases=phases, save_dir=save_dir, **kwargs)
|
||||
|
||||
|
||||
def launch_vis(i, args):
|
||||
log_dir = args.log_dirs[i]
|
||||
dev_id = args.gpus[i % len(args.gpus)]
|
||||
os.environ["EGL_DEVICE_ID"] = str(dev_id)
|
||||
os.environ["PYOPENGL_PLATFORM"] = "egl"
|
||||
|
||||
if args.save_root is not None:
|
||||
path_name = log_dir.split(args.log_root)[-1].strip("/")
|
||||
exp_name = "-".join(path_name.split("/")[:2])
|
||||
save_dir = f"{args.save_root}/{exp_name}"
|
||||
os.makedirs(save_dir, exist_ok=True)
|
||||
|
||||
print('args.save_root: ', args.save_root)
|
||||
visualize_log(
|
||||
log_dir,
|
||||
dev_id,
|
||||
phases=args.phases,
|
||||
save_dir=save_dir,
|
||||
overwrite=args.overwrite,
|
||||
accumulate=args.accumulate,
|
||||
render_kps=args.render_kps,
|
||||
render_layers=args.render_layers,
|
||||
render_views=args.render_views,
|
||||
save_frames=args.save_frames,
|
||||
make_grid=args.grid,
|
||||
)
|
||||
|
||||
|
||||
def main(args):
|
||||
"""
|
||||
visualize all runs in root
|
||||
"""
|
||||
OmegaConf.register_new_resolver("eval", eval)
|
||||
log_dirs = []
|
||||
for root, subd, files in os.walk(args.log_root):
|
||||
if ".hydra" in subd:
|
||||
log_dirs.append(root)
|
||||
args.log_dirs = log_dirs
|
||||
print(f"FOUND {len(args.log_dirs)} TO RENDER")
|
||||
|
||||
if len(args.gpus) > 1:
|
||||
from torch.multiprocessing import Pool
|
||||
|
||||
torch.multiprocessing.set_start_method("spawn")
|
||||
|
||||
with Pool(processes=len(args.gpus)) as pool:
|
||||
res = pool.starmap(
|
||||
launch_vis, [(i, args) for i in range(len(args.log_dirs))]
|
||||
)
|
||||
return
|
||||
|
||||
for i in range(len(args.log_dirs)):
|
||||
launch_vis(i, args)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import argparse
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--log_root", required=True)
|
||||
parser.add_argument("--save_root", default='./')
|
||||
parser.add_argument("--phases", nargs="*", default=["smooth_fit"])
|
||||
parser.add_argument("--gpus", nargs="*", default=[0])
|
||||
parser.add_argument(
|
||||
"-rv",
|
||||
"--render_views",
|
||||
nargs="*",
|
||||
default=["src_cam", "front", "above", "side"],
|
||||
)
|
||||
parser.add_argument("-g", "--grid", action="store_true")
|
||||
parser.add_argument("-rl", "--render_layers", action="store_true")
|
||||
parser.add_argument("-kp", "--render_kps", action="store_true")
|
||||
parser.add_argument("-sf", "--save_frames", action="store_true")
|
||||
parser.add_argument("-ra", "--accumulate", action="store_true")
|
||||
parser.add_argument("-y", "--overwrite", action="store_true")
|
||||
args = parser.parse_args()
|
||||
|
||||
main(args)
|
||||
+54
@@ -0,0 +1,54 @@
|
||||
#!/bin/bash
|
||||
|
||||
# Point BASE_PATH directly to subject1
|
||||
BASE_PATH="/data/home/zy3023/code/hand/slam-hand/test/HOT3D"
|
||||
VIDEO_JSON="/data/home/zy3023/code/hand/slam-hand/slahmr-eccv/confs/data/video.yaml"
|
||||
|
||||
# Directory for log files
|
||||
LOG_DIR="logs"
|
||||
mkdir -p "$LOG_DIR"
|
||||
|
||||
# How many Python processes to run in parallel at most
|
||||
MAX_JOBS=2
|
||||
|
||||
count=0 # Tracks how many jobs have been launched
|
||||
|
||||
for mp4file in "$BASE_PATH"/*.mp4; do
|
||||
[ -f "$mp4file" ] || continue
|
||||
|
||||
# basename "a.mp4" => "a.mp4"
|
||||
mp4basename=$(basename "$mp4file")
|
||||
# strip the ".mp4" extension => "a"
|
||||
vidname="${mp4basename%.mp4}"
|
||||
|
||||
logfile="${LOG_DIR}/${vidname}.log"
|
||||
|
||||
# Skip this video if the log file exists and is non-empty
|
||||
if [ -s "$logfile" ]; then
|
||||
echo "Skipping $vidname, log file exists and is non-empty: $logfile"
|
||||
continue
|
||||
fi
|
||||
|
||||
# Update 'seq' in the YAML to "a" instead of "a.mp4"
|
||||
sed -i "s|seq: *\"[^\"]*\"|seq: \"$vidname\"|" "$VIDEO_JSON"
|
||||
|
||||
echo "Launching job for: $vidname => seq: \"$vidname\". Logs: $logfile"
|
||||
|
||||
# Run Python in the background (&).
|
||||
# Output goes to "$logfile"
|
||||
python -u run_opt.py data=video run_opt=True run_vis=True data.seq=$vidname \
|
||||
> "$logfile" 2>&1 &
|
||||
|
||||
((count++))
|
||||
# If we've launched MAX_JOBS in parallel, wait for them to complete
|
||||
if (( count % MAX_JOBS == 0 )); then
|
||||
echo "Reached $count jobs. Waiting for them to finish..."
|
||||
wait
|
||||
echo "Continuing..."
|
||||
fi
|
||||
done
|
||||
|
||||
# After the loop, wait for any remaining jobs
|
||||
wait
|
||||
|
||||
echo "All done."
|
||||
@@ -0,0 +1,3 @@
|
||||
from . import loaders
|
||||
from . import logger
|
||||
from . import tensor
|
||||
+207
@@ -0,0 +1,207 @@
|
||||
"""
|
||||
Utilities for loading input and models (reading pre-processed data)
|
||||
"""
|
||||
|
||||
import os
|
||||
import importlib
|
||||
import glob
|
||||
from omegaconf import OmegaConf
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
# from body_model import HandModel
|
||||
|
||||
|
||||
ROOT_DIR = os.path.abspath(f"{__file__}/../../../")
|
||||
|
||||
|
||||
def load_config_from_log(log_dir):
|
||||
hydra_dir = f"{log_dir}/.hydra"
|
||||
cfg_path = f"{hydra_dir}/config.yaml"
|
||||
assert os.path.isdir(hydra_dir), f"{hydra_dir} does not exist"
|
||||
assert os.path.isfile(cfg_path), f"{cfg_path} does not exist"
|
||||
return OmegaConf.load(cfg_path)
|
||||
|
||||
|
||||
def resolve_cfg_paths(cfg):
|
||||
paths = cfg.paths
|
||||
for name, rel_path in paths.items():
|
||||
if rel_path.startswith("/"): # absolute path
|
||||
continue
|
||||
paths[name] = os.path.join(ROOT_DIR, rel_path)
|
||||
print("RESOLVED PATHS", paths)
|
||||
return cfg
|
||||
|
||||
|
||||
# def load_mano_body_model(cfg):
|
||||
# """
|
||||
# Load SMPL model
|
||||
# """
|
||||
# if device is None:
|
||||
# device = torch.device("cpu")
|
||||
# return (
|
||||
# HandModel(
|
||||
# bm_path=path,
|
||||
# num_betas=num_betas,
|
||||
# batch_size=batch_size,
|
||||
# use_vtx_selector=use_vtx_selector,
|
||||
# model_type=model_type,
|
||||
# ).to(device),
|
||||
# )
|
||||
|
||||
|
||||
def expid2model(expr_dir):
|
||||
""" "
|
||||
Reading VPoser models (https://github.com/nghorbani/human_body_prior).
|
||||
"""
|
||||
from configer import Configer
|
||||
|
||||
if not os.path.exists(expr_dir):
|
||||
raise ValueError("Could not find the experiment directory: %s" % expr_dir)
|
||||
|
||||
best_model_fname = sorted(
|
||||
glob.glob(os.path.join(expr_dir, "snapshots", "*.pt")), key=os.path.getmtime
|
||||
)[-1]
|
||||
try_num = os.path.basename(best_model_fname).split("_")[0]
|
||||
|
||||
print(("Found Trained Model: %s" % best_model_fname))
|
||||
|
||||
default_ps_fname = glob.glob(os.path.join(expr_dir, "*.ini"))[0]
|
||||
if not os.path.exists(default_ps_fname):
|
||||
raise ValueError(
|
||||
"Could not find the appropriate vposer_settings: %s" % default_ps_fname
|
||||
)
|
||||
ps = Configer(
|
||||
default_ps_fname=default_ps_fname,
|
||||
work_dir=expr_dir,
|
||||
best_model_fname=best_model_fname,
|
||||
)
|
||||
|
||||
return ps, best_model_fname
|
||||
|
||||
|
||||
def load_vposer(expr_dir, vp_model="snapshot"):
|
||||
"""
|
||||
:param expr_dir:
|
||||
:param vp_model: either 'snapshot' to use the experiment folder's code or a VPoser imported module, e.g.
|
||||
from human_body_prior.train.vposer_smpl import VPoser, then pass VPoser to this function
|
||||
:param if True will load the model definition used for training, and not the one in current repository
|
||||
:return:
|
||||
"""
|
||||
ps, trained_model_fname = expid2model(expr_dir)
|
||||
if vp_model == "snapshot":
|
||||
|
||||
vposer_path = sorted(
|
||||
glob.glob(os.path.join(expr_dir, "vposer_*.py")), key=os.path.getmtime
|
||||
)[-1]
|
||||
|
||||
spec = importlib.util.spec_from_file_location("VPoser", vposer_path)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(module)
|
||||
|
||||
vposer_pt = getattr(module, "VPoser")(
|
||||
num_neurons=ps.num_neurons, latentD=ps.latentD, data_shape=ps.data_shape
|
||||
)
|
||||
else:
|
||||
vposer_pt = vp_model(
|
||||
num_neurons=ps.num_neurons, latentD=ps.latentD, data_shape=ps.data_shape
|
||||
)
|
||||
|
||||
vposer_pt.load_state_dict(torch.load(trained_model_fname, map_location="cpu"))
|
||||
vposer_pt.eval()
|
||||
|
||||
return vposer_pt, ps
|
||||
|
||||
|
||||
def load_state(
|
||||
load_path,
|
||||
model,
|
||||
optimizer=None,
|
||||
is_parallel=False,
|
||||
map_location=None,
|
||||
ignore_keys=None,
|
||||
):
|
||||
"""
|
||||
Load Humor model checkpoint
|
||||
"""
|
||||
if not os.path.exists(load_path):
|
||||
print("Could not find checkpoint at path " + load_path)
|
||||
|
||||
full_checkpoint_dict = torch.load(load_path, map_location=map_location)
|
||||
model_state_dict = full_checkpoint_dict["model"]
|
||||
optim_state_dict = full_checkpoint_dict["optim"]
|
||||
|
||||
# load model weights
|
||||
for k, v in model_state_dict.items():
|
||||
if k.split(".")[0] == "module" and not is_parallel:
|
||||
# then it was trained with Data parallel
|
||||
print("Loading weights trained with DataParallel...")
|
||||
model_state_dict = {
|
||||
".".join(k.split(".")[1:]): v
|
||||
for k, v in model_state_dict.items()
|
||||
if k.split(".")[0] == "module"
|
||||
}
|
||||
break
|
||||
|
||||
if ignore_keys is not None:
|
||||
model_state_dict = {
|
||||
k: v
|
||||
for k, v in model_state_dict.items()
|
||||
if k.split(".")[0] not in ignore_keys
|
||||
}
|
||||
|
||||
# overwrite entries in the existing state dict
|
||||
missing_keys, unexpected_keys = model.load_state_dict(
|
||||
model_state_dict, strict=False
|
||||
)
|
||||
if ignore_keys is not None:
|
||||
missing_keys = [k for k in missing_keys if k.split(".")[0] not in ignore_keys]
|
||||
unexpected_keys = [
|
||||
k for k in unexpected_keys if k.split(".")[0] not in ignore_keys
|
||||
]
|
||||
if len(missing_keys) > 0:
|
||||
print(
|
||||
"WARNING: The following keys could not be found in the given state dict - ignoring..."
|
||||
)
|
||||
print(missing_keys)
|
||||
if len(unexpected_keys) > 0:
|
||||
print(
|
||||
"WARNING: The following keys were found in the given state dict but not in the current model - ignoring..."
|
||||
)
|
||||
print(unexpected_keys)
|
||||
|
||||
# load optimizer weights
|
||||
if optimizer is not None:
|
||||
optimizer.load_state_dict(optim_state_dict)
|
||||
|
||||
min_train_loss = float("Inf")
|
||||
if "min_train_loss" in full_checkpoint_dict.keys():
|
||||
min_train_loss = full_checkpoint_dict["min_train_loss"]
|
||||
|
||||
return (
|
||||
full_checkpoint_dict["epoch"],
|
||||
full_checkpoint_dict["min_val_loss"],
|
||||
min_train_loss,
|
||||
)
|
||||
|
||||
|
||||
def load_gmm(path, device=None):
|
||||
"""
|
||||
Load motion prior GMM for first frame motion prior
|
||||
"""
|
||||
if device is None:
|
||||
device = torch.device("cpu")
|
||||
|
||||
gmm_path = os.path.join(path, "prior_gmm.npz")
|
||||
if not os.path.isfile(gmm_path):
|
||||
raise ValueError(f"Init motion prior path {gmm_path} does not exist")
|
||||
|
||||
res = np.load(gmm_path)
|
||||
return {
|
||||
"gmm": (
|
||||
torch.from_numpy(res["weights"]).to(device),
|
||||
torch.from_numpy(res["means"]).to(device),
|
||||
torch.from_numpy(res["covariances"]).to(device),
|
||||
)
|
||||
}
|
||||
+37
@@ -0,0 +1,37 @@
|
||||
import datetime
|
||||
|
||||
|
||||
class Logger(object):
|
||||
"""
|
||||
Static logging class
|
||||
"""
|
||||
|
||||
log_file = None
|
||||
|
||||
@staticmethod
|
||||
def init(log_path):
|
||||
Logger.log_file = log_path
|
||||
|
||||
@staticmethod
|
||||
def log(write_str, to_stdout=True):
|
||||
if to_stdout:
|
||||
print(write_str)
|
||||
if not Logger.log_file:
|
||||
print("Logger must be initialized before logging!")
|
||||
return
|
||||
time_str = datetime.datetime.now().strftime("%Y-%m-%d_%H:%M:%S")
|
||||
with open(Logger.log_file, "a") as f:
|
||||
f.write(time_str + " ")
|
||||
f.write(str(write_str) + "\n")
|
||||
|
||||
|
||||
def log_cur_stats(stats_dict, iter=None, to_stdout=True):
|
||||
loss = stats_dict.pop("total", 0)
|
||||
Logger.log(f"LOSS: {loss:.04f}", to_stdout=to_stdout)
|
||||
for k, v in stats_dict.items():
|
||||
Logger.log(f"{k}: {v:.04f}", to_stdout=to_stdout)
|
||||
if to_stdout:
|
||||
if iter is not None:
|
||||
print("======= iter %d =======" % iter)
|
||||
else:
|
||||
print("========")
|
||||
+139
@@ -0,0 +1,139 @@
|
||||
import torch
|
||||
import numpy as np
|
||||
|
||||
|
||||
def move_to(obj, device):
|
||||
if isinstance(obj, torch.Tensor):
|
||||
return obj.to(device)
|
||||
if isinstance(obj, dict):
|
||||
return {k: move_to(v, device) for k, v in obj.items()}
|
||||
if isinstance(obj, (list, tuple)):
|
||||
return [move_to(x, device) for x in obj]
|
||||
return obj # otherwise do nothing
|
||||
|
||||
|
||||
def detach_all(obj):
|
||||
if isinstance(obj, torch.Tensor):
|
||||
return obj.detach()
|
||||
if isinstance(obj, dict):
|
||||
return {k: detach_all(v) for k, v in obj.items()}
|
||||
if isinstance(obj, (list, tuple)):
|
||||
return [detach_all(x) for x in obj]
|
||||
return obj # otherwise do nothing
|
||||
|
||||
|
||||
def to_torch(obj):
|
||||
if isinstance(obj, np.ndarray):
|
||||
return torch.from_numpy(obj).float()
|
||||
if isinstance(obj, dict):
|
||||
return {k: to_torch(v) for k, v in obj.items()}
|
||||
if isinstance(obj, (list, tuple)):
|
||||
return [to_torch(x) for x in obj]
|
||||
return obj
|
||||
|
||||
|
||||
def to_np(obj):
|
||||
if isinstance(obj, torch.Tensor):
|
||||
return obj.numpy()
|
||||
if isinstance(obj, dict):
|
||||
return {k: to_np(v) for k, v in obj.items()}
|
||||
if isinstance(obj, (list, tuple)):
|
||||
return [to_np(x) for x in obj]
|
||||
return obj
|
||||
|
||||
|
||||
def get_device(i=0):
|
||||
device = f"cuda:{i}" if torch.cuda.is_available() else "cpu"
|
||||
return torch.device(device)
|
||||
|
||||
|
||||
def scatter_intervals(tensor, start, end, T=None):
|
||||
"""
|
||||
Scatter the tensor contents into intervals from start to end
|
||||
output tensor indexed from 0 to end.max()
|
||||
:param tensor (B, S, *)
|
||||
:param start (B) start indices
|
||||
:param end (B) end indices
|
||||
:param T (int, optional) max length
|
||||
returns (B, T, *) scattered tensor
|
||||
"""
|
||||
assert isinstance(tensor, torch.Tensor) and tensor.ndim >= 2
|
||||
if T is not None:
|
||||
end = torch.where(end < 0, T + 1 + end, end)
|
||||
assert torch.all(end <= T)
|
||||
else:
|
||||
T = end.max()
|
||||
|
||||
B, S, *dims = tensor.shape
|
||||
start, end = start.long(), end.long()
|
||||
# get idcs that go past the last time step so we don't have repeat indices in scatter
|
||||
idcs = time_segment_idcs(start, end, clip=False) # (B, S)
|
||||
# mask out the extra padding
|
||||
mask = idcs >= end[:, None]
|
||||
tensor[mask] = 0
|
||||
|
||||
idcs = idcs.reshape(B, S, *(1,) * len(dims)).repeat(1, 1, *dims)
|
||||
output = torch.zeros(
|
||||
B, idcs.max() + 1, *dims, device=tensor.device, dtype=tensor.dtype
|
||||
)
|
||||
output.scatter_(1, idcs, tensor)
|
||||
# slice out the extra segments
|
||||
return output[:, :T]
|
||||
|
||||
|
||||
def get_scatter_mask(start, end, T):
|
||||
"""
|
||||
get the mask of selected intervals
|
||||
"""
|
||||
B = start.shape[0]
|
||||
start, end = start.long(), end.long()
|
||||
assert torch.all(end <= T)
|
||||
idcs = time_segment_idcs(start, end, clip=True)
|
||||
mask = torch.zeros(B, T, device=start.device, dtype=torch.bool)
|
||||
mask.scatter_(1, idcs, 1)
|
||||
return mask
|
||||
|
||||
|
||||
def select_intervals(series, start, end):
|
||||
"""
|
||||
Select slices of a tensor from start to end
|
||||
will pad uneven sequences to all the max segment length
|
||||
:param series (B, T, *)
|
||||
:param start (B)
|
||||
:param end (B)
|
||||
returns (B, S, *) selected segments, S = max(end - start)
|
||||
"""
|
||||
B, T, *dims = series.shape
|
||||
assert torch.all(end <= T)
|
||||
sel = time_segment_idcs(start, end, clip=True)
|
||||
S = sel.shape[1]
|
||||
sel = sel.reshape(B, S, *(1,) * len(dims)).repeat(1, 1, *dims)
|
||||
return torch.gather(series, 1, sel)
|
||||
|
||||
|
||||
def get_select_mask(start, end):
|
||||
"""
|
||||
get the mask of unpadded elementes for the selected time segments
|
||||
e.g. sel[mask] are the unpadded elements
|
||||
:param start (B)
|
||||
:param end (B)
|
||||
"""
|
||||
idcs = time_segment_idcs(start, end, clip=False)
|
||||
return idcs < end[:, None] # (B, S)
|
||||
|
||||
|
||||
def time_segment_idcs(start, end, clip=True):
|
||||
"""
|
||||
:param start (B)
|
||||
:param end (B)
|
||||
returns (B, S) long tensor of indices, where S = max(end - start)
|
||||
"""
|
||||
start, end = start.long(), end.long()
|
||||
S = (end - start).max()
|
||||
seg = torch.arange(S, dtype=torch.int64, device=start.device)
|
||||
idcs = start[:, None] + seg[None, :] # (B, S)
|
||||
if clip:
|
||||
# clip at the lengths of each track
|
||||
imax = end[:, None] - 1
|
||||
idcs = idcs.clamp(max=imax)
|
||||
return idcs
|
||||
@@ -0,0 +1,4 @@
|
||||
from . import viewer
|
||||
|
||||
# from . import render
|
||||
from . import tools
|
||||
+98
@@ -0,0 +1,98 @@
|
||||
241 91 181
|
||||
0 187 249
|
||||
254 228 64
|
||||
0 245 212
|
||||
222 158 54
|
||||
255 133 141
|
||||
93 169 233
|
||||
155 93 229
|
||||
24 242 178
|
||||
252 116 186
|
||||
138 225 252
|
||||
178 171 242
|
||||
24 231 98
|
||||
232 93 117
|
||||
207 250 231
|
||||
253 208 242
|
||||
109 157 197
|
||||
166 189 219
|
||||
253 146 207
|
||||
167 201 87
|
||||
117 68 177
|
||||
255 229 50
|
||||
251 202 239
|
||||
58 134 255
|
||||
255 0 110
|
||||
251 86 7
|
||||
188 51 209
|
||||
122 229 130
|
||||
0 48 73
|
||||
214 40 40
|
||||
229 179 179
|
||||
0 187 249
|
||||
255 190 11
|
||||
204 213 174
|
||||
0 245 212
|
||||
255 153 200
|
||||
144 251 146
|
||||
189 211 147
|
||||
230 0 86
|
||||
0 95 57
|
||||
0 174 126
|
||||
255 116 163
|
||||
189 198 255
|
||||
90 219 255
|
||||
158 0 142
|
||||
255 147 126
|
||||
164 36 0
|
||||
0 21 68
|
||||
145 208 203
|
||||
95 173 78
|
||||
107 104 130
|
||||
0 125 181
|
||||
106 130 108
|
||||
252 246 189
|
||||
208 244 222
|
||||
169 222 249
|
||||
228 193 249
|
||||
122 204 174
|
||||
194 140 159
|
||||
0 143 156
|
||||
235 0 0
|
||||
255 2 157
|
||||
104 61 59
|
||||
150 138 232
|
||||
152 255 82
|
||||
167 87 64
|
||||
1 255 254
|
||||
255 238 232
|
||||
254 137 0
|
||||
1 208 255
|
||||
187 136 0
|
||||
165 255 210
|
||||
255 166 254
|
||||
119 77 0
|
||||
122 71 130
|
||||
38 52 0
|
||||
0 71 84
|
||||
67 0 44
|
||||
181 0 255
|
||||
255 177 103
|
||||
255 219 102
|
||||
126 45 210
|
||||
229 111 254
|
||||
222 255 116
|
||||
0 255 120
|
||||
0 155 255
|
||||
0 100 1
|
||||
0 118 255
|
||||
133 169 0
|
||||
0 185 23
|
||||
120 130 49
|
||||
0 255 198
|
||||
255 110 65
|
||||
232 94 190
|
||||
1 0 103
|
||||
149 0 58
|
||||
98 14 0
|
||||
0 0 0
|
||||
+245
@@ -0,0 +1,245 @@
|
||||
import torch
|
||||
import numpy as np
|
||||
|
||||
|
||||
def get_seq_figure_skip(seq_name=None):
|
||||
if seq_name == "017437_mpii_test":
|
||||
return 15
|
||||
if seq_name == "012968_mpii_test":
|
||||
return 15
|
||||
if seq_name == "recording_20211002_S03_S18_04-all-100-200":
|
||||
return 11
|
||||
return 10
|
||||
|
||||
|
||||
def get_seq_static_lookat_points(seq_name=None, bounds=None):
|
||||
if seq_name == "002276_mpii_test":
|
||||
top_source = torch.tensor([-0.5, -3.0, -2.0])
|
||||
top_target = torch.tensor([0.0, 0.0, 4.0])
|
||||
|
||||
side_source = torch.tensor([6.0, -2.0, 2.0])
|
||||
side_target = torch.tensor([2.0, 0.0, 6.0])
|
||||
return (top_source, top_target), (side_source, side_target)
|
||||
|
||||
if seq_name == "002374_mpii_test":
|
||||
top_source = torch.tensor([-4.0, 0.0, 0.0])
|
||||
top_target = torch.tensor([-2.0, 2.0, 3.0])
|
||||
|
||||
side_source = torch.tensor([-7.0, 1.0, 7.0])
|
||||
side_target = torch.tensor([0.0, 1.0, 8.0])
|
||||
return (top_source, top_target), (side_source, side_target)
|
||||
|
||||
if seq_name == "003742_mpii_test":
|
||||
top_source = torch.tensor([3.0, -2.0, -6])
|
||||
top_target = torch.tensor([3.0, 1.0, 2.0])
|
||||
|
||||
side_source = torch.tensor([-5.0, -1.0, -2])
|
||||
side_target = torch.tensor([2.0, 0.0, 5])
|
||||
return (top_source, top_target), (side_source, side_target)
|
||||
|
||||
if seq_name == "003747_mpii_test":
|
||||
top_source = torch.tensor([-5.0, -2.0, -5])
|
||||
top_target = torch.tensor([-3.0, 2.0, 1.0])
|
||||
|
||||
side_source = torch.tensor([-9.0, -1.0, -4])
|
||||
side_target = torch.tensor([-3.0, 2.0, 3])
|
||||
return (top_source, top_target), (side_source, side_target)
|
||||
|
||||
if seq_name == "003943_mpii_test":
|
||||
top_source = torch.tensor([0, -1, -2])
|
||||
top_target = torch.tensor([0, 1, 3])
|
||||
|
||||
side_source = torch.tensor([6, 0, 1])
|
||||
side_target = torch.tensor([0, 0, 5])
|
||||
return (top_source, top_target), (side_source, side_target)
|
||||
|
||||
if seq_name == "006537_mpii_test":
|
||||
top_source = torch.tensor([0, -4, -1])
|
||||
top_target = torch.tensor([1.0, 0.0, 5.0])
|
||||
|
||||
side_source = torch.tensor([5, -2, 0])
|
||||
side_target = torch.tensor([1.0, 1.0, 7.0])
|
||||
return (top_source, top_target), (side_source, side_target)
|
||||
|
||||
if seq_name == "007684_mpii_test":
|
||||
top_source = torch.tensor([-2, -4, -6])
|
||||
top_target = torch.tensor([0.5, 0.0, 4.0])
|
||||
|
||||
side_source = torch.tensor([-8, -3, -2])
|
||||
side_target = torch.tensor([1.0, 0.0, 5.0])
|
||||
return (top_source, top_target), (side_source, side_target)
|
||||
|
||||
if seq_name == "009039_mpii_test":
|
||||
top_source = torch.tensor([-2.0, -3.0, -1.5])
|
||||
top_target = torch.tensor([1.0, 0.0, 3.0])
|
||||
|
||||
side_source = torch.tensor([6, -2, -1.5])
|
||||
side_target = torch.tensor([2.0, 0.0, 6.0])
|
||||
return (top_source, top_target), (side_source, side_target)
|
||||
|
||||
if seq_name == "009607_mpii_test":
|
||||
top_source = torch.tensor([0, -3, -1])
|
||||
top_target = torch.tensor([0, 0, 5])
|
||||
|
||||
side_source = torch.tensor([-6, -0.5, 2])
|
||||
side_target = torch.tensor([-1.0, 0.5, 7])
|
||||
return (top_source, top_target), (side_source, side_target)
|
||||
|
||||
if seq_name == "012968_mpii_test":
|
||||
top_source = torch.tensor([-0.5, -3, -4])
|
||||
top_target = torch.tensor([0.0, 0.5, 2.0])
|
||||
|
||||
side_source = torch.tensor([7, -4, 0])
|
||||
side_target = torch.tensor([0, 0.5, 2.0])
|
||||
return (top_source, top_target), (side_source, side_target)
|
||||
|
||||
if seq_name == "014531_mpii_test":
|
||||
top_source = torch.tensor([0, -2, -2])
|
||||
top_target = torch.tensor([0, 0, 2])
|
||||
|
||||
side_source = torch.tensor([-2, 0, -2])
|
||||
side_target = torch.tensor([0, 0, 2])
|
||||
return (top_source, top_target), (side_source, side_target)
|
||||
|
||||
if seq_name == "015933_mpii_test":
|
||||
top_source = torch.tensor([0.0, 0.5, -5.0])
|
||||
top_target = torch.tensor([0.0, 2.0, -2.0])
|
||||
|
||||
side_source = torch.tensor([5.0, 1.5, -3.0])
|
||||
side_target = torch.tensor([0.0, 2.5, 1.5])
|
||||
return (top_source, top_target), (side_source, side_target)
|
||||
|
||||
if seq_name == "016195_mpii_test":
|
||||
top_source = torch.tensor([0, -1, -2])
|
||||
top_target = torch.tensor([0, 1, 1])
|
||||
|
||||
side_source = torch.tensor([-3, 0, 1])
|
||||
side_target = torch.tensor([0, 1, 2])
|
||||
return (top_source, top_target), (side_source, side_target)
|
||||
|
||||
if seq_name == "018061_mpii_test":
|
||||
top_source = torch.tensor([4, -1, 2])
|
||||
top_target = torch.tensor([1, 1, 3])
|
||||
|
||||
side_source = torch.tensor([-3, 0, 4])
|
||||
side_target = torch.tensor([1, 1, 3])
|
||||
return (top_source, top_target), (side_source, side_target)
|
||||
|
||||
if seq_name == "018713_mpii_test":
|
||||
top_source = torch.tensor([0, -1, -1])
|
||||
top_target = torch.tensor([0, 0.5, 2])
|
||||
|
||||
side_source = torch.tensor([3, 0, 1])
|
||||
side_target = torch.tensor([0, 0, 3])
|
||||
return (top_source, top_target), (side_source, side_target)
|
||||
|
||||
if seq_name == "022691_mpii_test":
|
||||
top_source = torch.tensor([0, -4, -4])
|
||||
top_target = torch.tensor([-2.0, 0.0, 1.0])
|
||||
|
||||
side_source = torch.tensor([-5, -2, -5])
|
||||
side_target = torch.tensor([-3.5, 0.0, 1.0])
|
||||
return (top_source, top_target), (side_source, side_target)
|
||||
|
||||
if seq_name == "023390_mpii_test":
|
||||
top_source = torch.tensor([0.0, -1.5, -1.0])
|
||||
top_target = torch.tensor([0.0, 0.0, 2.0])
|
||||
|
||||
side_source = torch.tensor([4.0, -0.5, 0])
|
||||
side_target = torch.tensor([1.0, 0.5, 4])
|
||||
# top_source = torch.tensor([0.0, -3.0, -2.0])
|
||||
# top_target = torch.tensor([0.0, 0.5, 2.0])
|
||||
|
||||
# side_source = torch.tensor([5.0, -0.5, -3])
|
||||
# side_target = torch.tensor([1.0, 0.5, 4])
|
||||
return (top_source, top_target), (side_source, side_target)
|
||||
|
||||
if seq_name=="024165_mpii_test":
|
||||
top_source = torch.tensor([0, -3, -3])
|
||||
top_target = torch.tensor([0, 0, 3])
|
||||
|
||||
side_source = torch.tensor([-4, -1, -2])
|
||||
side_target = torch.tensor([0, 0, 3])
|
||||
return (top_source, top_target), (side_source, side_target)
|
||||
|
||||
if seq_name == "024154_mpii_test":
|
||||
top_source = torch.tensor([-1, -0.5, 1])
|
||||
top_target = torch.tensor([-1, 1, 8])
|
||||
|
||||
side_source = torch.tensor([-4, 1, 3])
|
||||
side_target = torch.tensor([-1, 1, 8])
|
||||
return (top_source, top_target), (side_source, side_target)
|
||||
|
||||
# if seq_name == "024159_mpii_test":
|
||||
# top_source = torch.tensor([-4.0, -2.0, -6.0])
|
||||
# top_target = torch.tensor([-4.0, 0.0, 0.0])
|
||||
|
||||
# side_source = torch.tensor([-12, -1, -2])
|
||||
# side_target = torch.tensor([-10.0, 0.0, 0.0])
|
||||
# return (top_source, top_target), (side_source, side_target)
|
||||
|
||||
if seq_name == "023962_mpii_test":
|
||||
top_source = torch.tensor([0, -0.5, -2])
|
||||
top_target = torch.tensor([0, 0.5, 2])
|
||||
|
||||
side_source = torch.tensor([3, 0.5, 0])
|
||||
side_target = torch.tensor([0, 0, 2])
|
||||
return (top_source, top_target), (side_source, side_target)
|
||||
|
||||
if seq_name == "017437_mpii_test":
|
||||
top_source = torch.tensor([0, -8, -6])
|
||||
top_target = torch.tensor([-2.0, -1.0, 2.0])
|
||||
|
||||
side_source = torch.tensor([-5, -1, -7])
|
||||
side_target = torch.tensor([-3.0, -1.0, 1.0])
|
||||
return (top_source, top_target), (side_source, side_target)
|
||||
|
||||
if seq_name == "recording_20211002_S03_S18_04-all-100-200":
|
||||
top_source = torch.tensor([1.0, -1, -4.5])
|
||||
top_target = torch.tensor([1.0, 0.0, 1])
|
||||
|
||||
side_source = torch.tensor([1.0, 0.5, -4.5])
|
||||
side_target = torch.tensor([1.0, 0.5, 1])
|
||||
return (top_source, top_target), (side_source, side_target)
|
||||
|
||||
if seq_name == "recording_20211002_S03_S18_04-all-100-500":
|
||||
top_source = torch.tensor([0.0, -0.5, -3])
|
||||
top_target = torch.tensor([0.0, 1.0, 1])
|
||||
|
||||
side_source = torch.tensor([-2.0, 0.5, -3])
|
||||
side_target = torch.tensor([0.0, 1.0, 2])
|
||||
return (top_source, top_target), (side_source, side_target)
|
||||
|
||||
if seq_name == "recording_20211002_S03_S18_01-all-700-862":
|
||||
top_source = torch.tensor([0.5, -2.0, -2])
|
||||
top_target = torch.tensor([0.5, -0.5, 2])
|
||||
|
||||
side_source = torch.tensor([2.0, -0.5, -2])
|
||||
side_target = torch.tensor([0.5, -0.5, 1])
|
||||
return (top_source, top_target), (side_source, side_target)
|
||||
|
||||
if seq_name == "recording_20220315_S21_S30_02-all-0-100":
|
||||
top_source = torch.tensor([0.0, -1, -2])
|
||||
top_target = torch.tensor([0.0, 0.5, 1])
|
||||
|
||||
side_source = torch.tensor([-2.0, 0.0, -1])
|
||||
side_target = torch.tensor([0.0, 0.5, 1])
|
||||
return (top_source, top_target), (side_source, side_target)
|
||||
|
||||
if bounds is not None:
|
||||
# (3), (3), (3)
|
||||
bb_min, bb_max, center = bounds
|
||||
print("SCENE BOUNDS", bb_min, bb_max, center)
|
||||
length = torch.abs(bb_max - bb_min).max()
|
||||
print(length)
|
||||
top_source = center + torch.tensor([0.0, -2.0, -0.9 * length])
|
||||
|
||||
side_source = center + torch.tensor([0.5 * length, -0.5, -0.7 * length])
|
||||
return (top_source, center), (side_source, center)
|
||||
|
||||
top_source = torch.tensor([0.0, -2.0, -3.0])
|
||||
top_target = torch.tensor([0.0, 0.0, 1.0])
|
||||
|
||||
side_source = torch.tensor([3.0, -1.0, -1.0])
|
||||
side_target = torch.tensor([0.0, 0.0, 1.0])
|
||||
return (top_source, top_target), (side_source, side_target)
|
||||
+379
@@ -0,0 +1,379 @@
|
||||
import os
|
||||
import cv2
|
||||
import imageio
|
||||
import numpy as np
|
||||
import torch
|
||||
import subprocess
|
||||
|
||||
from body_model import run_mano
|
||||
from geometry import camera as cam_util
|
||||
from geometry.mesh import make_batch_mesh, save_mesh_scenes
|
||||
from geometry.plane import parse_floor_plane, get_plane_transform
|
||||
|
||||
from util.tensor import detach_all, to_torch, move_to
|
||||
|
||||
from .fig_specs import get_seq_figure_skip, get_seq_static_lookat_points
|
||||
from .tools import mesh_to_geometry, smooth_results
|
||||
|
||||
|
||||
def prep_result_vis(res, vis_mask, track_ids, body_model, temporal_smooth, smooth_trans=True):
|
||||
"""
|
||||
:param res (dict) with (B, T, *) tensor elements, B tracks and T frames
|
||||
:param vis_mask (B, T) with visibility of each track in each frame
|
||||
:param track_ids (B,) index of each track
|
||||
:param temporal_smooth (bool) whether to apply temporal smoothing
|
||||
:param smooth_trans (bool) whether to smooth translation (only used if temporal_smooth=True)
|
||||
"""
|
||||
print("RESULT FIELDS", res.keys())
|
||||
res = detach_all(res)
|
||||
# print(res.keys(), res["is_right"])
|
||||
with torch.no_grad():
|
||||
if temporal_smooth:
|
||||
if smooth_trans:
|
||||
print('running temporal smooth (including trans)')
|
||||
# Smooth everything including trans
|
||||
res["root_orient"], res["pose_body"], res['betas'], res["trans"] = smooth_results(res["root_orient"], res["pose_body"], res['betas'], res["is_right"], res["trans"])
|
||||
else:
|
||||
print('running temporal smooth (excluding trans)')
|
||||
# Smooth only pose and shape, not trans
|
||||
res["root_orient"], res["pose_body"], res['betas'], _ = smooth_results(res["root_orient"], res["pose_body"], res['betas'], res["is_right"])
|
||||
|
||||
world_smpl = run_mano(
|
||||
body_model,
|
||||
res["trans"],
|
||||
res["root_orient"],
|
||||
res["pose_body"],
|
||||
res["is_right"],
|
||||
res.get("betas", None),
|
||||
)
|
||||
|
||||
T_w2c = None
|
||||
floor_plane = None
|
||||
if "cam_R" in res and "cam_t" in res:
|
||||
T_w2c = cam_util.make_4x4_pose(res["cam_R"][0], res["cam_t"][0])
|
||||
if "floor_plane" in res:
|
||||
floor_plane = res["floor_plane"][0]
|
||||
return build_scene_dict(
|
||||
world_smpl,
|
||||
vis_mask,
|
||||
track_ids,
|
||||
T_w2c=T_w2c,
|
||||
floor_plane=floor_plane,
|
||||
)
|
||||
|
||||
|
||||
def build_scene_dict(
|
||||
world_smpl, vis_mask, track_ids, T_w2c=None, floor_plane=None, **kwargs
|
||||
):
|
||||
scene_dict = {}
|
||||
|
||||
# first get the geometry of the people
|
||||
# lists of length T with (B, V, 3), (F, 3), (B, 3)
|
||||
scene_dict["geometry"] = mesh_to_geometry(
|
||||
world_smpl["vertices"], world_smpl["joints"], world_smpl["l_faces"], world_smpl["r_faces"], world_smpl["is_right"], vis_mask, track_ids
|
||||
)
|
||||
|
||||
if T_w2c is None:
|
||||
T_w2c = torch.eye(4)[None]
|
||||
|
||||
T_c2w = torch.linalg.inv(T_w2c)
|
||||
# rotate the camera slightly down and translate back and up
|
||||
T = cam_util.make_4x4_pose(
|
||||
cam_util.rotx(-np.pi / 10), torch.tensor([0, -1, -2])
|
||||
).to(T_c2w.device)
|
||||
|
||||
scene_dict["cameras"] = {
|
||||
"src_cam": T_c2w,
|
||||
# "front": torch.einsum("ij,...jk->...ik", T, T_c2w),
|
||||
}
|
||||
|
||||
# Create ground plane based on hand mesh height
|
||||
# DISABLED: Ground plane visualization disabled to better see coordinate system
|
||||
# verts = scene_dict["geometry"][0] # Get vertices from geometry
|
||||
# if len(verts) > 0:
|
||||
# # Find minimum height across all frames and vertices
|
||||
# min_height = float('inf')
|
||||
# for frame_verts in verts:
|
||||
# if len(frame_verts) > 0:
|
||||
# frame_min = frame_verts[..., 1].min().item() # y-coordinate is height
|
||||
# min_height = min(min_height, frame_min)
|
||||
#
|
||||
# # Set ground plane further below minimum height
|
||||
# ground_offset = -0.5 # 20cm below minimum height
|
||||
# ground_height = min_height - ground_offset
|
||||
#
|
||||
# # Create ground plane transform
|
||||
# R = torch.eye(3) # Identity rotation (flat ground)
|
||||
# t = torch.tensor([0.0, ground_height, 0.0]) # Translate to ground height
|
||||
# scene_dict["ground"] = cam_util.make_4x4_pose(R, t)
|
||||
#
|
||||
# # Save ground mesh for Blender debugging
|
||||
# import trimesh
|
||||
# from vis.viewer import make_checkerboard
|
||||
# ground_mesh = make_checkerboard(color0=[0.9, 0.95, 1.0], color1=[0.7, 0.8, 0.85], up="y", alpha=1.0)
|
||||
# ground_mesh.apply_translation([0.0, ground_height, 0.0])
|
||||
# ground_mesh.export("ground_debug.obj")
|
||||
|
||||
# if floor_plane is not None:
|
||||
# # compute the ground transform
|
||||
# # use the first appearance of a track as the reference point
|
||||
# tid, sid = torch.where(vis_mask > 0)
|
||||
# idx = tid[torch.argmin(sid)]
|
||||
# root = world_smpl["joints"][idx, 0, 0].detach().cpu()
|
||||
# floor = parse_floor_plane(floor_plane.detach().cpu())
|
||||
# R, t = get_plane_transform(torch.tensor([0.0, 1.0, 0.0]), floor, root)
|
||||
# scene_dict["ground"] = cam_util.make_4x4_pose(R, t)
|
||||
|
||||
return scene_dict
|
||||
|
||||
|
||||
# def render_scene_dict(renderer, scene_dict, out_name, fps=30, **kwargs):
|
||||
# # lists of T (B, V, 3), (B, 3), (F, 3)
|
||||
# verts, colors, faces, bounds = scene_dict["geometry"]
|
||||
# print("NUM VERTS", len(verts))
|
||||
|
||||
# # add a top view
|
||||
# scene_dict["cameras"]["above"] = cam_util.make_4x4_pose(
|
||||
# torch.eye(3), torch.tensor([0, 0, -10])
|
||||
# )[None]
|
||||
|
||||
# for cam_name, cam_poses in scene_dict["cameras"].items():
|
||||
# print("rendering scene for", cam_name)
|
||||
# # cam_poses are (T, 4, 4)
|
||||
# render_bg = cam_name == "src_cam"
|
||||
# ground_pose = scene_dict.get("ground", None)
|
||||
# print(cam_name, ground_pose)
|
||||
# frames = renderer.render_video(
|
||||
# cam_poses[None], verts, faces, colors, render_bg, ground_pose=ground_pose
|
||||
# )
|
||||
# os.makedirs(f"{out_name}_{cam_name}/", exist_ok=True)
|
||||
# for idx, i in enumerate(frames):
|
||||
# imageio.imwrite(f"{out_name}_{cam_name}/" + f'{str(idx).zfill(6)}.jpg', i)
|
||||
# imageio.mimwrite(f"{out_name}_{cam_name}.mp4", frames, fps=fps)
|
||||
|
||||
|
||||
def animate_scene(
|
||||
vis,
|
||||
scene,
|
||||
out_name,
|
||||
seq_name=None,
|
||||
accumulate=False,
|
||||
render_views=["src_cam", "front", "above", "side"],
|
||||
render_bg=True,
|
||||
render_cam=True,
|
||||
render_ground=True,
|
||||
debug=False,
|
||||
render_keypoints=False,
|
||||
**kwargs,
|
||||
):
|
||||
if len(render_views) < 1:
|
||||
return
|
||||
|
||||
scene = build_pyrender_scene(
|
||||
vis,
|
||||
scene,
|
||||
seq_name,
|
||||
render_views=render_views,
|
||||
render_cam=render_cam,
|
||||
accumulate=accumulate,
|
||||
debug=debug,
|
||||
)
|
||||
|
||||
print("RENDERING VIEWS", scene["cameras"].keys())
|
||||
render_ground = render_ground and "ground" in scene
|
||||
save_paths = []
|
||||
for cam_name, cam_poses in scene["cameras"].items():
|
||||
is_src = cam_name == "src_cam"
|
||||
show_bg = is_src and render_bg
|
||||
show_ground = render_ground and not is_src
|
||||
show_cam = render_cam and not is_src
|
||||
show_keypoints = is_src and render_keypoints # Only show keypoints on source view
|
||||
vis_name = f"{out_name}_{cam_name}"
|
||||
print(f"{cam_name} has {len(cam_poses)} poses")
|
||||
skip = 10 if debug else 1
|
||||
vis.set_camera_seq(cam_poses[::skip])
|
||||
save_path = vis.animate(
|
||||
vis_name,
|
||||
render_bg=show_bg,
|
||||
render_ground=show_ground,
|
||||
render_cam=show_cam,
|
||||
render_keypoints=show_keypoints,
|
||||
**kwargs,
|
||||
)
|
||||
save_paths.append(save_path)
|
||||
|
||||
return save_paths
|
||||
|
||||
|
||||
def build_pyrender_scene(
|
||||
vis,
|
||||
scene,
|
||||
seq_name,
|
||||
render_views=["src_cam", "front", "above", "side"],
|
||||
render_cam=True,
|
||||
accumulate=False,
|
||||
debug=False,
|
||||
):
|
||||
"""
|
||||
:param vis (viewer object)
|
||||
:param scene (scene dict with geometry, cameras, etc)
|
||||
:param accumulate (optional bool, default False) whether to render entire trajectory together
|
||||
:param render_views (list str) camera views to render
|
||||
"""
|
||||
if len(render_views) < 1:
|
||||
return
|
||||
|
||||
assert all(view in ["src_cam", "front", "above", "side"] for view in render_views)
|
||||
|
||||
scene = move_to(detach_all(scene), "cpu")
|
||||
src_cams = scene["cameras"]["src_cam"]
|
||||
verts, _, colors, l_faces, r_faces, is_right, bounds = scene["geometry"]
|
||||
T = len(verts)
|
||||
print(f"{T} mesh frames for {seq_name}, {len(verts)}")
|
||||
|
||||
# set camera views
|
||||
if not "cameras" in scene:
|
||||
scene["cameras"] = {}
|
||||
|
||||
# remove default views from source camera perspective if desired
|
||||
if "src_cam" not in render_views:
|
||||
scene["cameras"].pop("src_cam", None)
|
||||
if "front" not in render_views:
|
||||
scene["cameras"].pop("front", None)
|
||||
|
||||
# add static viewpoints if desired
|
||||
top_pose, side_pose, _skip = get_static_views(seq_name, bounds)
|
||||
if "above" in render_views:
|
||||
scene["cameras"]["above"] = top_pose[None]
|
||||
if "side" in render_views:
|
||||
scene["cameras"]["side"] = side_pose[None]
|
||||
if "front" in render_views:
|
||||
# Use a static front camera: place it in front of the hand mesh, looking at the center
|
||||
# Place the camera at a fixed distance along the -z axis from the center
|
||||
if bounds is not None:
|
||||
bb_min, bb_max, center = bounds
|
||||
length = torch.abs(bb_max - bb_min).max()
|
||||
front_offset = 1.0 # 1 closer, 2 farther
|
||||
front_source = center + torch.tensor([0.0, -0.5, -front_offset])
|
||||
front_target = center
|
||||
up = torch.tensor([0.0, 1, 0.0])
|
||||
front_pose = cam_util.lookat_matrix(front_source, front_target, up)
|
||||
scene["cameras"]["front"] = front_pose[None]
|
||||
else:
|
||||
# fallback to previous behavior if bounds not available
|
||||
pass
|
||||
|
||||
# accumulate meshes if possible (can only accumulate for static camera)
|
||||
moving_cam = "src_cam" in render_views or "front" in render_views
|
||||
accumulate = accumulate and not moving_cam
|
||||
skip = _skip if accumulate else 1
|
||||
|
||||
vis.clear_meshes()
|
||||
|
||||
if "ground" in scene:
|
||||
vis.set_ground(scene["ground"])
|
||||
|
||||
if debug:
|
||||
skip = 10
|
||||
times = list(range(0, T, skip))
|
||||
# print(times, verts[0].shape)
|
||||
# raise ValueError
|
||||
for t in times:
|
||||
if len(is_right[t]) > 1:
|
||||
assert (is_right[t].cpu().numpy().tolist() == [0,1])
|
||||
l_meshes = make_batch_mesh(verts[t][0][None], l_faces[t], colors[t][0][None])
|
||||
r_meshes = make_batch_mesh(verts[t][1][None], r_faces[t], colors[t][1][None])
|
||||
assert len(l_meshes) == 1
|
||||
assert len(r_meshes) == 1
|
||||
meshes = [l_meshes[0], r_meshes[0]]
|
||||
else:
|
||||
assert len(is_right[t]) == 1
|
||||
if is_right[t] == 0:
|
||||
meshes = make_batch_mesh(verts[t][0][None], l_faces[t], colors[t][0][None])
|
||||
elif is_right[t] == 1:
|
||||
meshes = make_batch_mesh(verts[t][0][None], r_faces[t], colors[t][0][None])
|
||||
|
||||
if accumulate:
|
||||
vis.add_static_meshes(meshes)
|
||||
else:
|
||||
vis.add_mesh_frame(meshes, debug=debug)
|
||||
|
||||
# add camera markers
|
||||
if render_cam:
|
||||
if accumulate:
|
||||
vis.add_camera_markers_static(src_cams[::skip])
|
||||
else:
|
||||
vis.add_camera_markers(src_cams[::skip])
|
||||
|
||||
return scene
|
||||
|
||||
|
||||
def get_static_views(seq_name=None, bounds=None):
|
||||
print("STATIC VIEWS FOR SEQ NAME", seq_name)
|
||||
up = torch.tensor([0.0, 1.0, 0.0])
|
||||
|
||||
skip = get_seq_figure_skip(seq_name)
|
||||
top_vp, side_vp = get_seq_static_lookat_points(seq_name, bounds)
|
||||
top_source, top_target = top_vp
|
||||
side_source, side_target = side_vp
|
||||
top_pose = cam_util.lookat_matrix(top_source, top_target, up)
|
||||
side_pose = cam_util.lookat_matrix(side_source, side_target, up)
|
||||
return top_pose, side_pose, skip
|
||||
|
||||
|
||||
def make_video_grid_2x2(out_path, vid_paths, fps, overwrite=False):
|
||||
if os.path.isfile(out_path) and not overwrite:
|
||||
print(f"{out_path} already exists, skipping.")
|
||||
return
|
||||
|
||||
if any(not os.path.isfile(v) for v in vid_paths):
|
||||
print("not all inputs exist!", vid_paths)
|
||||
return
|
||||
|
||||
# Open all videos for reading
|
||||
readers = [imageio.get_reader(v, 'ffmpeg') for v in vid_paths]
|
||||
try:
|
||||
# Check min frame count to avoid index errors; also get shape/fps
|
||||
lengths = [reader.count_frames() for reader in readers]
|
||||
min_len = min(lengths)
|
||||
if min_len == 0:
|
||||
print("Some videos have 0 frames!")
|
||||
for reader in readers:
|
||||
reader.close()
|
||||
return
|
||||
|
||||
# We assume the first video dictates the output fps and (half)size
|
||||
meta = readers[0].get_meta_data()
|
||||
# fps = meta.get('fps', 30)
|
||||
orig_size = meta['size'] # (width, height)
|
||||
width = orig_size[0] // 2
|
||||
height = orig_size[1] // 2
|
||||
out_size = (width * 2, height * 2)
|
||||
|
||||
writer = imageio.get_writer(out_path, fps=fps, codec='libx264', ffmpeg_log_level="quiet", macro_block_size=None)
|
||||
|
||||
for t in range(min_len):
|
||||
frames = []
|
||||
for r in readers:
|
||||
try:
|
||||
frame = r.get_data(t)
|
||||
except Exception:
|
||||
frame = None
|
||||
frames.append(frame)
|
||||
if any(f is None for f in frames):
|
||||
break
|
||||
# Resize all to (width, height)
|
||||
frames = [cv2.resize(f, (width, height)) for f in frames]
|
||||
row1 = np.hstack(frames[:2])
|
||||
row2 = np.hstack(frames[2:])
|
||||
grid = np.vstack([row1, row2])
|
||||
# Ensure uint8
|
||||
if grid.dtype != np.uint8:
|
||||
grid = np.clip(grid, 0, 255).astype(np.uint8)
|
||||
writer.append_data(grid)
|
||||
|
||||
writer.close()
|
||||
print(f"Video grid created: {out_path}")
|
||||
finally:
|
||||
for reader in readers:
|
||||
reader.close()
|
||||
+254
@@ -0,0 +1,254 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import numpy as np
|
||||
|
||||
from pytorch3d.renderer import (
|
||||
RasterizationSettings,
|
||||
MeshRenderer,
|
||||
MeshRasterizer,
|
||||
HardPhongShader,
|
||||
PointLights,
|
||||
PerspectiveCameras,
|
||||
)
|
||||
from pytorch3d.structures import Meshes
|
||||
from pytorch3d.structures.meshes import join_meshes_as_scene
|
||||
from pytorch3d.renderer import TexturesAtlas, TexturesVertex
|
||||
from pytorch3d.transforms import euler_angles_to_matrix
|
||||
|
||||
from .tools import get_colors, read_image, transform_torch3d, checkerboard_geometry
|
||||
|
||||
|
||||
def init_renderer(img_size, intrins, device, vis_scale=1.0, bg_paths=None):
|
||||
img_size = int(vis_scale * img_size[0]), int(vis_scale * img_size[1])
|
||||
renderer = RenderBase(device, img_size, intrins=vis_scale * intrins)
|
||||
if bg_paths is not None:
|
||||
bg_imgs = [
|
||||
torch.from_numpy(read_image(p, scale=vis_scale)).float() / 255
|
||||
for p in bg_paths
|
||||
]
|
||||
print(bg_imgs[0].shape)
|
||||
renderer.set_bg_seq(bg_imgs)
|
||||
print("RENDERER", renderer)
|
||||
return renderer
|
||||
|
||||
|
||||
class RenderBase(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
device,
|
||||
img_size,
|
||||
intrins=None,
|
||||
blur_radius: float = 0.0, # rasterizer settings
|
||||
faces_per_pixel: int = 3,
|
||||
):
|
||||
"""
|
||||
Uses Pytorch3D to render results.
|
||||
We use the PyTorch3D coordinate system, which assumes:
|
||||
+X:left, +Y: up and +Z: from us to scene (right-handed)
|
||||
https://pytorch3d.org/docs/cameras
|
||||
"""
|
||||
super().__init__()
|
||||
self.device = device
|
||||
W, H = img_size
|
||||
self.img_size = (H, W)
|
||||
if intrins is not None:
|
||||
intrins = intrins.to(device)
|
||||
self.cam_f, self.cam_center = intrins[:2], intrins[2:]
|
||||
cx, cy = self.cam_center.detach().cpu()
|
||||
else:
|
||||
self.cam_f = torch.tensor([0.5 * (H + W), 0.5 * (H + W)], device=device)
|
||||
self.cam_center = torch.tensor([0.5 * W, 0.5 * H], device=device)
|
||||
print("cam_f", self.cam_f, "cam_center", self.cam_center)
|
||||
print("img_size", self.img_size)
|
||||
self.blur_radius = blur_radius
|
||||
self.faces_per_pixel = faces_per_pixel
|
||||
|
||||
self.light_loc = [[0, 0, -0.5]]
|
||||
self.cameras = self.make_cameras()
|
||||
self.bg_imgs = None
|
||||
self.set_ground()
|
||||
self.build_renderer()
|
||||
|
||||
def build_renderer(self):
|
||||
self.lights = PointLights(device=self.device, location=self.light_loc)
|
||||
|
||||
self.rasterizer = RasterizationSettings(
|
||||
image_size=self.img_size,
|
||||
blur_radius=self.blur_radius,
|
||||
faces_per_pixel=self.faces_per_pixel,
|
||||
max_faces_per_bin=100000,
|
||||
)
|
||||
self.renderer = MeshRenderer(
|
||||
rasterizer=MeshRasterizer(
|
||||
cameras=self.cameras, raster_settings=self.rasterizer
|
||||
),
|
||||
shader=HardPhongShader(
|
||||
cameras=self.cameras, lights=self.lights, device=self.device
|
||||
),
|
||||
)
|
||||
|
||||
def make_cameras(self, cam_poses=None):
|
||||
device = self.device
|
||||
if cam_poses is None:
|
||||
cam_poses = torch.eye(4, device=device)[None]
|
||||
cam_poses = cam_poses.to(device)
|
||||
cam_R, cam_t = transform_torch3d(cam_poses)
|
||||
Nc = len(cam_poses)
|
||||
focal = self.cam_f[None].repeat(Nc, 1)
|
||||
ppt = self.cam_center[None].repeat(Nc, 1)
|
||||
img_size = [self.img_size for _ in range(Nc)]
|
||||
return PerspectiveCameras(
|
||||
device=self.device,
|
||||
in_ndc=False,
|
||||
focal_length=focal,
|
||||
principal_point=ppt,
|
||||
image_size=img_size,
|
||||
R=cam_R,
|
||||
T=cam_t,
|
||||
)
|
||||
|
||||
def render_frame(self, cam_poses, verts, faces, colors):
|
||||
"""
|
||||
:param cam_poses (C, 4, 4)
|
||||
:param verts (B, V, 3)
|
||||
:param faces (F, 3)
|
||||
:param colors (B, 4)
|
||||
"""
|
||||
with torch.no_grad():
|
||||
cameras = self.make_cameras(cam_poses)
|
||||
mesh = build_meshes(verts, faces, colors)
|
||||
image = self.renderer(mesh, cameras=cameras)
|
||||
return image
|
||||
|
||||
def render_with_ground(self, cam_poses, verts, faces, colors, ground_pose):
|
||||
"""
|
||||
:param cam_poses (C, 4, 4)
|
||||
:param verts (B, V, 3)
|
||||
:param faces (F, 3)
|
||||
:param colors (B, 3)
|
||||
:param ground_pose (4, 4)
|
||||
"""
|
||||
with torch.no_grad():
|
||||
self.set_ground_pose(ground_pose)
|
||||
cameras = self.make_cameras(cam_poses)
|
||||
# (B, V, 3), (B, F, 3), (B, V, 3)
|
||||
verts, faces, colors = prep_shared_geometry(verts, faces, colors)
|
||||
# (V, 3), (F, 3), (V, 3)
|
||||
gv, gf, gc = self.ground_geometry
|
||||
verts = list(torch.unbind(verts, dim=0)) + [gv]
|
||||
faces = list(torch.unbind(faces, dim=0)) + [gf]
|
||||
colors = list(torch.unbind(colors, dim=0)) + [gc[..., :3]]
|
||||
mesh = create_meshes(verts, faces, colors)
|
||||
|
||||
image = self.renderer(mesh, cameras=cameras)
|
||||
return image
|
||||
|
||||
def set_bg_seq(self, bg_imgs):
|
||||
self.bg_imgs = bg_imgs
|
||||
|
||||
def set_ground(self, pose=None):
|
||||
device = self.device
|
||||
v, f, vc, fc = map(torch.from_numpy, checkerboard_geometry(up="y"))
|
||||
v, f, vc = v.to(device), f.to(device), vc.to(device)
|
||||
self.ground_geometry = [v, f, vc]
|
||||
if pose is not None:
|
||||
self.set_ground_pose(pose)
|
||||
|
||||
def set_ground_pose(self, pose):
|
||||
"""
|
||||
:param pose (4, 4)
|
||||
"""
|
||||
if self.ground_geometry is None:
|
||||
return
|
||||
v = self.ground_geometry[0]
|
||||
pose = pose.to(self.device)
|
||||
# R, t = pose[:3, :3], pose[:3, 3]
|
||||
R, t = transform_torch3d(pose)
|
||||
self.ground_geometry[0] = torch.einsum("ij,vj->vi", R, v) + t[None]
|
||||
|
||||
def composite_bg(self, frames, fac=1.0):
|
||||
T = len(frames)
|
||||
if self.bg_imgs is None:
|
||||
return frames
|
||||
|
||||
out_frames = []
|
||||
for bg, frame in zip(self.bg_imgs, frames):
|
||||
alpha = fac * frame[..., 3:]
|
||||
out = alpha * frame[..., :3] + (1 - alpha) * bg
|
||||
out_frames.append(out)
|
||||
return out_frames
|
||||
|
||||
def render_video(
|
||||
self,
|
||||
cam_poses,
|
||||
verts,
|
||||
faces,
|
||||
colors,
|
||||
render_bg=False,
|
||||
fac=1.0,
|
||||
ground_pose=None,
|
||||
):
|
||||
"""
|
||||
:param cam_poses (C, T, 4, 4)
|
||||
:param verts list of T (B, V, 3)
|
||||
:param faces list of T (F, 3)
|
||||
:param colors list of T (B, 4)
|
||||
:param render_bg (optional default False)
|
||||
:param fac (optional float, default 1)
|
||||
:param ground_pose (optional, default None) (4, 4)
|
||||
returns list of T (H, W, 4)
|
||||
"""
|
||||
T = len(verts)
|
||||
print(f"rendering video with {T} frames")
|
||||
if cam_poses.shape[1] == 1:
|
||||
cam_poses = cam_poses.expand(-1, T, -1, -1)
|
||||
|
||||
if render_bg: # don't render ground on source frame
|
||||
ground_pose = None
|
||||
|
||||
frames = []
|
||||
for t in range(T):
|
||||
if ground_pose is None:
|
||||
# (H, W, 4)
|
||||
frame = self.render_frame(
|
||||
cam_poses[:, t], verts[t], faces[t], colors[t]
|
||||
)
|
||||
else:
|
||||
# (H, W, 4)
|
||||
frame = self.render_with_ground(
|
||||
cam_poses[:, t], verts[t], faces[t], colors[t], ground_pose
|
||||
)
|
||||
frames.append(frame.cpu()[0])
|
||||
|
||||
if render_bg and self.bg_imgs is not None:
|
||||
frames = self.composite_bg(frames, fac=fac)
|
||||
|
||||
return [(255 * f).byte() for f in frames]
|
||||
|
||||
|
||||
def build_meshes(verts, faces, colors):
|
||||
return create_meshes(*prep_shared_geometry(verts, faces, colors))
|
||||
|
||||
|
||||
def prep_shared_geometry(verts, faces, colors):
|
||||
"""
|
||||
:param verts (B, V, 3)
|
||||
:param faces (F, 3)
|
||||
:param colors (B, 4)
|
||||
"""
|
||||
B, V, _ = verts.shape
|
||||
F, _ = faces.shape
|
||||
colors = colors.unsqueeze(1).expand(B, V, -1)[..., :3]
|
||||
faces = faces.unsqueeze(0).expand(B, F, -1)
|
||||
return verts, faces, colors
|
||||
|
||||
|
||||
def create_meshes(verts, faces, colors):
|
||||
"""
|
||||
:param verts (B, V, 3)
|
||||
:param faces (B, F, 3)
|
||||
:param colors (B, V, 3)
|
||||
"""
|
||||
textures = TexturesVertex(verts_features=colors)
|
||||
meshes = Meshes(verts=verts, faces=faces, textures=textures)
|
||||
return join_meshes_as_scene(meshes)
|
||||
+1109
File diff suppressed because it is too large
Load Diff
+672
@@ -0,0 +1,672 @@
|
||||
import os
|
||||
import imageio
|
||||
import numpy as np
|
||||
|
||||
import time
|
||||
import torch
|
||||
import trimesh
|
||||
|
||||
os.environ["PYOPENGL_PLATFORM"] = "egl"
|
||||
|
||||
import pyrender
|
||||
from pyrender.constants import RenderFlags
|
||||
from pyrender.light import DirectionalLight
|
||||
from pyrender.node import Node
|
||||
|
||||
from geometry.camera import make_rotation, make_translation
|
||||
from .tools import (
|
||||
read_image,
|
||||
checkerboard_geometry,
|
||||
camera_marker_geometry,
|
||||
transform_pyrender,
|
||||
)
|
||||
|
||||
|
||||
def init_viewer(
|
||||
img_size, intrins, vis_scale=1.0, bg_paths=None, fps=30,
|
||||
):
|
||||
img_size = int(vis_scale * img_size[0]), int(vis_scale * img_size[1])
|
||||
intrins = vis_scale * intrins
|
||||
|
||||
platform = os.environ.get("PYOPENGL_PLATFORM", "pyglet")
|
||||
if platform == "pyglet":
|
||||
vis = AnimationViewer(img_size, intrins=intrins, fps=fps)
|
||||
print("VIS", vis)
|
||||
return vis
|
||||
|
||||
if platform != "egl":
|
||||
raise NotImplementedError
|
||||
|
||||
vis = OffscreenAnimation(img_size, intrins=intrins, fps=fps)
|
||||
if bg_paths is not None:
|
||||
bg_imgs = [
|
||||
read_image(p, scale=vis_scale).astype(np.float32) / 255 for p in bg_paths
|
||||
]
|
||||
vis.set_bg_seq(bg_imgs)
|
||||
print("VIS", vis)
|
||||
return vis
|
||||
|
||||
|
||||
class AnimationBase(object):
|
||||
"""
|
||||
Render an animation
|
||||
IMPORTANT: pyrender uses coordinates with Y UP (i.e. xyz = right up back),
|
||||
we must keep this in mind because the human prior frame has Z UP (i.e.
|
||||
xyz = right forward up)
|
||||
"""
|
||||
|
||||
def __init__(self, img_size, intrins=None):
|
||||
self.scene = pyrender.Scene(
|
||||
ambient_light=[0.3, 0.3, 0.3], bg_color=[1.0, 1.0, 1.0, 0.0]
|
||||
)
|
||||
self.scene.bg_color = np.array([1.0, 1.0, 1.0, 0.0])
|
||||
|
||||
self.viewport_size = img_size
|
||||
self.camera = make_pyrender_camera(img_size, intrins)
|
||||
|
||||
# set the ground with Z up
|
||||
self.ground_pose = np.eye(4)
|
||||
ground = pyrender.Mesh.from_trimesh(
|
||||
make_checkerboard(color0=[0.9, 0.95, 1.0], color1=[0.7, 0.8, 0.85], up="y", alpha=1.0), smooth=False
|
||||
)
|
||||
self.ground_node = self.scene.add(ground, name="ground", pose=self.ground_pose)
|
||||
self.cam_node = self.scene.add(self.camera, name="camera")
|
||||
|
||||
self.cam_marker_node = None
|
||||
self.cam_marker_poses = []
|
||||
|
||||
self.light_nodes = []
|
||||
self.light_poses = []
|
||||
|
||||
self.bg_seq = []
|
||||
|
||||
# pyrender nodes, reused between frames
|
||||
self.anim_nodes = []
|
||||
# list of meshes per frame of animation
|
||||
self.anim_meshes = []
|
||||
# list of camera poses per frame of animation
|
||||
self.anim_cameras = None
|
||||
# current timestep of animation
|
||||
self.anim_idx = 0
|
||||
|
||||
def acquire_lock(self):
|
||||
pass
|
||||
|
||||
def release_lock(self):
|
||||
pass
|
||||
|
||||
def close(self):
|
||||
pass
|
||||
|
||||
def delete(self):
|
||||
pass
|
||||
|
||||
def clear_meshes(self):
|
||||
self.anim_meshes = []
|
||||
self.anim_idx = 0
|
||||
|
||||
@property
|
||||
def anim_len(self):
|
||||
return len(self.anim_meshes)
|
||||
|
||||
def add_lighting(self, color=np.ones(3), intensity=1.0):
|
||||
self.light_poses = get_light_poses()
|
||||
self.light_poses.append(np.eye(4))
|
||||
cam_pose = self.scene.get_pose(self.cam_node)
|
||||
for i, pose in enumerate(self.light_poses):
|
||||
matrix = cam_pose @ pose
|
||||
node = Node(
|
||||
name=f"light-{i:02d}",
|
||||
light=DirectionalLight(color=color, intensity=intensity),
|
||||
matrix=matrix,
|
||||
)
|
||||
if self.scene.has_node(node):
|
||||
continue
|
||||
self.scene.add_node(node)
|
||||
self.light_nodes.append(node)
|
||||
|
||||
def add_mesh_frame(self, meshes, debug=False):
|
||||
"""
|
||||
:param meshes (list) trimesh meshes for this frame timestep
|
||||
"""
|
||||
meshes = [pyrender.Mesh.from_trimesh(m) for m in meshes]
|
||||
num_nodes = len(meshes)
|
||||
self.anim_meshes.append(meshes)
|
||||
|
||||
if num_nodes > len(self.anim_nodes):
|
||||
# need to add a node in the pyrender scene
|
||||
self.acquire_lock()
|
||||
for i in range(len(self.anim_nodes), num_nodes):
|
||||
self.anim_nodes.append(self.scene.add(meshes[i], name=f"mesh_{i:03d}"))
|
||||
self.release_lock()
|
||||
|
||||
def add_static_meshes(self, meshes, smooth=True):
|
||||
"""
|
||||
add all meshes to a single frame (timestep)
|
||||
"""
|
||||
meshes = [pyrender.Mesh.from_trimesh(m, smooth=smooth) for m in meshes]
|
||||
if len(self.anim_meshes) < 1:
|
||||
self.anim_meshes.append([])
|
||||
self.anim_meshes[0].extend(meshes)
|
||||
num_meshes = len(self.anim_meshes[0])
|
||||
num_nodes = len(self.anim_nodes)
|
||||
if num_meshes > num_nodes:
|
||||
# need to add a node in the pyrender scene
|
||||
self.acquire_lock()
|
||||
self.anim_nodes.extend(
|
||||
[
|
||||
self.scene.add(self.anim_meshes[0][i])
|
||||
for i in range(num_nodes, num_meshes)
|
||||
]
|
||||
)
|
||||
self.release_lock()
|
||||
|
||||
def set_camera_seq(self, poses):
|
||||
"""
|
||||
:param poses (*, 4, 4)
|
||||
"""
|
||||
poses = transform_pyrender(poses)
|
||||
poses = np.asarray(poses)
|
||||
assert poses.ndim >= 2 and poses.shape[-2:] == (4, 4)
|
||||
if poses.ndim < 3:
|
||||
poses = poses[None]
|
||||
print(f"Adding camera sequence length {len(poses)}")
|
||||
self.anim_cameras = poses
|
||||
|
||||
def add_camera_markers(self, poses):
|
||||
"""
|
||||
Add a single camera marker node that we move around when we animate
|
||||
"""
|
||||
print("ADDING CAMERA MARKERS")
|
||||
if self.cam_marker_node is None:
|
||||
cam_marker = make_camera_marker(up="y")
|
||||
self.cam_marker_node = self.scene.add(
|
||||
pyrender.Mesh.from_trimesh(cam_marker, smooth=False)
|
||||
)
|
||||
poses = transform_pyrender(poses)
|
||||
poses = np.asarray(poses)
|
||||
self.cam_marker_poses = poses
|
||||
|
||||
def add_camera_markers_static(self, poses):
|
||||
"""
|
||||
Add all camera markers as separate mesh nodes to render in one pass
|
||||
"""
|
||||
print("ADDING STATIC CAMERA MARKERS")
|
||||
poses = transform_pyrender(poses)
|
||||
poses = np.asarray(poses)
|
||||
cam_meshes = [make_camera_marker(up="y", transform=pose) for pose in poses]
|
||||
self.add_static_meshes(cam_meshes, smooth=False)
|
||||
|
||||
def update_camera(self, cam_pose):
|
||||
self.acquire_lock()
|
||||
self.scene.set_pose(self.cam_node, pose=cam_pose)
|
||||
self.release_lock()
|
||||
|
||||
def set_ground(self, pose):
|
||||
pose = transform_pyrender(pose)
|
||||
print("Setting ground pose to be", pose)
|
||||
self.ground_pose = np.asarray(pose)
|
||||
self.scene.set_pose(self.ground_node, pose=self.ground_pose)
|
||||
|
||||
def set_bg_seq(self, bg_imgs):
|
||||
pass
|
||||
|
||||
def set_mesh_visibility(self, vis):
|
||||
self.acquire_lock()
|
||||
for node in self.anim_nodes:
|
||||
node.mesh.is_visible = vis
|
||||
self.release_lock()
|
||||
|
||||
def check_mesh_visibility(self):
|
||||
for node in self.anim_nodes:
|
||||
if node.mesh.is_visible:
|
||||
return True
|
||||
|
||||
if self.ground_node.mesh.is_visible:
|
||||
return True
|
||||
|
||||
if self.cam_marker_node is not None and self.cam_marker_node.mesh.is_visible:
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
def update_frame(self):
|
||||
t = self.anim_idx % self.anim_len
|
||||
|
||||
# update camera
|
||||
if self.anim_cameras is not None:
|
||||
cam_t = min(t, len(self.anim_cameras) - 1)
|
||||
cam_pose = self.anim_cameras[cam_t]
|
||||
self.update_camera(cam_pose)
|
||||
if self.cam_marker_node is not None and t < len(self.cam_marker_poses):
|
||||
self.acquire_lock()
|
||||
self.scene.set_pose(self.cam_marker_node, self.cam_marker_poses[t])
|
||||
self.release_lock()
|
||||
|
||||
# update meshes
|
||||
meshes = self.anim_meshes[t]
|
||||
self.acquire_lock()
|
||||
for i, node in enumerate(self.anim_nodes):
|
||||
if i < len(meshes):
|
||||
node.mesh = meshes[i]
|
||||
node.mesh.is_visible = True
|
||||
else:
|
||||
node.mesh.is_visible = False
|
||||
self.release_lock()
|
||||
|
||||
|
||||
class OffscreenAnimation(AnimationBase):
|
||||
def __init__(self, img_size, intrins=None, fps=30, ext="mp4"):
|
||||
super().__init__(img_size, intrins=intrins)
|
||||
self.add_lighting(0.9)
|
||||
|
||||
self.fps = fps
|
||||
self.ext = ext
|
||||
self.viewer = pyrender.OffscreenRenderer(*self.viewport_size)
|
||||
|
||||
self.bg_seq = []
|
||||
self.keypoints_seq = None # Store GT 2D keypoints for overlay
|
||||
self.pred_keypoints_seq = None # Store predicted 2D keypoints
|
||||
|
||||
def set_bg_seq(self, bg_imgs):
|
||||
assert isinstance(bg_imgs, list)
|
||||
print(f"setting length {len(bg_imgs)} bg sequence")
|
||||
self.bg_seq = bg_imgs
|
||||
|
||||
def set_keypoints_seq(self, keypoints_seq):
|
||||
"""
|
||||
Set GT 2D keypoints sequence for overlay visualization
|
||||
Args:
|
||||
keypoints_seq: List of (J, 3) arrays with [x, y, confidence] for each frame
|
||||
"""
|
||||
self.keypoints_seq = keypoints_seq
|
||||
print(f"setting length {len(keypoints_seq)} GT keypoints sequence")
|
||||
|
||||
def set_pred_keypoints_seq(self, pred_keypoints_seq):
|
||||
"""
|
||||
Set predicted 2D keypoints sequence
|
||||
Args:
|
||||
pred_keypoints_seq: List of (J, 2) arrays with [x, y] for each frame
|
||||
"""
|
||||
self.pred_keypoints_seq = pred_keypoints_seq
|
||||
print(f"setting length {len(pred_keypoints_seq)} predicted keypoints sequence")
|
||||
|
||||
def delete(self):
|
||||
self.viewer.delete()
|
||||
|
||||
def close(self):
|
||||
self.delete()
|
||||
|
||||
def draw_keypoints_on_image(self, img, keypoints, radius=4, color=(0, 255, 0), thickness=-1, conf_threshold=0.3):
|
||||
"""
|
||||
Draw 2D keypoints on an image
|
||||
Args:
|
||||
img: numpy array (H, W, 3)
|
||||
keypoints: (J, 3) array with [x, y, confidence]
|
||||
radius: circle radius for keypoints
|
||||
color: BGR color tuple
|
||||
thickness: -1 for filled circles
|
||||
conf_threshold: minimum confidence to draw keypoint
|
||||
Returns:
|
||||
img with keypoints drawn
|
||||
"""
|
||||
import cv2
|
||||
img = img.copy()
|
||||
for j in range(keypoints.shape[0]):
|
||||
x, y, conf = keypoints[j]
|
||||
if conf > conf_threshold:
|
||||
# Draw circle at keypoint location
|
||||
cv2.circle(img, (int(x), int(y)), radius, color, thickness)
|
||||
return img
|
||||
|
||||
def render(
|
||||
self, render_bg=True, render_ground=True, render_cam=True, fac=1.0, render_keypoints=False, **kwargs
|
||||
):
|
||||
flags = RenderFlags.RGBA | RenderFlags.SHADOWS_DIRECTIONAL
|
||||
self.ground_node.mesh.is_visible = render_ground and (not render_bg)
|
||||
if self.cam_marker_node is not None:
|
||||
self.cam_marker_node.mesh.is_visible = render_cam
|
||||
|
||||
rgba = np.zeros((*self.viewport_size[::-1], 4))
|
||||
if self.check_mesh_visibility(): # if any meshes are visible
|
||||
rgba, _ = self.viewer.render(self.scene, flags=flags)
|
||||
|
||||
img = rgba
|
||||
if render_bg and len(self.bg_seq) > 0:
|
||||
t = min(self.anim_idx, len(self.bg_seq) - 1)
|
||||
bg = self.bg_seq[t]
|
||||
rgba = rgba.astype(np.float32) / 255
|
||||
alpha = fac * rgba[..., 3:]
|
||||
img = alpha * rgba[..., :3] + (1 - alpha) * bg
|
||||
img = (255 * img).astype(np.uint8)
|
||||
|
||||
# Overlay 2D keypoints if available and requested
|
||||
if render_keypoints and render_bg:
|
||||
# Draw GT keypoints in GREEN (larger)
|
||||
if self.keypoints_seq is not None:
|
||||
t = min(self.anim_idx, len(self.keypoints_seq) - 1)
|
||||
keypoints = self.keypoints_seq[t]
|
||||
if keypoints is not None and len(keypoints) > 0:
|
||||
img = self.draw_keypoints_on_image(img, keypoints, color=(0, 255, 0), radius=5)
|
||||
|
||||
# Draw predicted keypoints in RED (smaller)
|
||||
if self.pred_keypoints_seq is not None:
|
||||
t = min(self.anim_idx, len(self.pred_keypoints_seq) - 1)
|
||||
pred_kp = self.pred_keypoints_seq[t]
|
||||
if pred_kp is not None and len(pred_kp) > 0:
|
||||
# Add confidence=1.0 if only (J, 2)
|
||||
if pred_kp.ndim == 2 and pred_kp.shape[1] == 2:
|
||||
pred_kp = np.concatenate([pred_kp, np.ones((len(pred_kp), 1))], axis=1)
|
||||
img = self.draw_keypoints_on_image(img, pred_kp, color=(0, 0, 255), radius=3, conf_threshold=0.0)
|
||||
|
||||
return img
|
||||
|
||||
def render_mesh_layers(self, render_ground=True, render_cam=True, **kwargs):
|
||||
# render each mesh in its own layer
|
||||
flags = RenderFlags.RGBA | RenderFlags.SHADOWS_DIRECTIONAL
|
||||
render_cam = render_cam and self.cam_marker_node is not None
|
||||
|
||||
# current timestep
|
||||
t = self.anim_idx % self.anim_len
|
||||
meshes = self.anim_meshes[t]
|
||||
layers = []
|
||||
|
||||
# toggle ground and camera
|
||||
self.ground_node.mesh.is_visible = False
|
||||
if self.cam_marker_node is not None:
|
||||
self.cam_marker_node.mesh.is_visible = False
|
||||
|
||||
# render animation meshes for current timestep
|
||||
for i in range(len(meshes)):
|
||||
for j, node in enumerate(self.anim_nodes):
|
||||
if i == j:
|
||||
node.mesh.is_visible = True
|
||||
else:
|
||||
node.mesh.is_visible = False
|
||||
rgba, _ = self.viewer.render(self.scene, flags=flags)
|
||||
layers.append(rgba)
|
||||
|
||||
if not render_ground and not render_cam:
|
||||
return layers
|
||||
|
||||
# render the ground and camera
|
||||
for node in self.anim_nodes:
|
||||
node.mesh.is_visible = False
|
||||
|
||||
if render_cam:
|
||||
self.cam_marker_node.mesh.is_visible = render_cam
|
||||
rgba, _ = self.viewer.render(self.scene, flags=flags)
|
||||
layers.append(rgba)
|
||||
self.cam_marker_node.mesh.is_visible = False
|
||||
|
||||
if render_ground:
|
||||
self.ground_node.mesh.is_visible = render_ground
|
||||
rgba, _ = self.viewer.render(self.scene, flags=flags)
|
||||
layers.append(rgba)
|
||||
|
||||
return layers
|
||||
|
||||
def render_frames_layers(self):
|
||||
frames_layers = []
|
||||
for t in range(self.anim_len):
|
||||
self.anim_idx = t
|
||||
self.update_frame()
|
||||
layers = self.render_mesh_layers()
|
||||
frames_layers.append(layers)
|
||||
return frames_layers
|
||||
|
||||
def animate(self, save_name, save_frames=False, render_layers=False, **kwargs):
|
||||
if render_layers:
|
||||
frames_layers = self.render_frames_layers()
|
||||
os.makedirs(save_name, exist_ok=True)
|
||||
for t, layers in enumerate(frames_layers):
|
||||
for l, layer in enumerate(layers):
|
||||
path = f"{save_name}/{t:06d}_{l:03d}.png"
|
||||
imageio.imwrite(path, layer)
|
||||
print(f"saved frame layers to {save_name}")
|
||||
return save_name
|
||||
|
||||
frames = self.render_frames(**kwargs)
|
||||
if len(frames) > 1:
|
||||
if save_frames:
|
||||
os.makedirs(save_name, exist_ok=True)
|
||||
for i, frame in enumerate(frames):
|
||||
path = f"{save_name}/{i:06d}.png"
|
||||
imageio.imwrite(path, frame)
|
||||
return save_name
|
||||
save_path = f"{save_name}.{self.ext}"
|
||||
os.makedirs(save_name, exist_ok=True)
|
||||
for idx, i in enumerate(frames):
|
||||
imageio.imwrite(f"{save_name}/" + f'{str(idx).zfill(6)}.jpg', i[:, :, :3])
|
||||
imageio.mimwrite(save_path, frames, fps=self.fps)
|
||||
print("wrote video to", save_path)
|
||||
return save_path
|
||||
|
||||
save_path = f"{save_name}.png"
|
||||
imageio.imwrite(save_path, frames[0])
|
||||
print("wrote image to", save_path)
|
||||
return save_path
|
||||
|
||||
def render_frames(self, **kwargs):
|
||||
print("ANIMATION LENGTH", self.anim_len)
|
||||
frames = []
|
||||
for t in range(self.anim_len):
|
||||
self.anim_idx = t
|
||||
self.update_frame()
|
||||
img = self.render(**kwargs)
|
||||
frames.append(img)
|
||||
return frames
|
||||
|
||||
def render_layers(self, save_dir, composite=False, fac=0.4):
|
||||
os.makedirs(save_dir, exist_ok=True)
|
||||
|
||||
# render all frames
|
||||
frames = self.render_frames(render_bg=False, render_ground=False)
|
||||
|
||||
# render background
|
||||
self.set_mesh_visibility(False)
|
||||
bg_img = self.render(render_bg=False, render_ground=True) / 255
|
||||
|
||||
# save frames
|
||||
for t, frame in enumerate(frames):
|
||||
imageio.imwrite(f"{save_dir}/{t:06d}.png", frame)
|
||||
imageio.imwrite(f"{save_dir}/background.png", bg_img)
|
||||
|
||||
if composite:
|
||||
rgba = composite_layers(frames, bg_img, fac=fac)
|
||||
imageio.imwrite(f"{save_dir}/composite.png", (255 * rgba).astype(np.uint8))
|
||||
|
||||
|
||||
def composite_layers(frames, bg_img, fac=0.4):
|
||||
# composite front to back
|
||||
H, W = frames[0].shape[:2]
|
||||
comp = np.zeros((H, W, 3))
|
||||
vis = np.ones((H, W, 1))
|
||||
for frame in frames:
|
||||
frame = frame / 255
|
||||
alpha = fac * frame[..., 3:]
|
||||
comp += vis * alpha * frame[..., :3]
|
||||
vis *= 1 - alpha
|
||||
|
||||
comp += vis * bg_img[..., 3:] * bg_img[..., :3]
|
||||
vis *= 1 - bg_img[..., 3:]
|
||||
return comp + vis
|
||||
rgba = np.concatenate([comp, 1 - vis], axis=-1)
|
||||
return rgba
|
||||
|
||||
|
||||
class AnimationViewer(AnimationBase):
|
||||
def __init__(self, img_size, intrins=None, fps=30, num_repeats=1):
|
||||
super().__init__(img_size, intrins=intrins)
|
||||
|
||||
self.fps = fps
|
||||
self.ext = "gif"
|
||||
|
||||
def pause_play_callback(_, viewer):
|
||||
viewer.is_paused = not viewer.is_paused
|
||||
|
||||
def step_callback(_, viewer, step):
|
||||
viewer.anim_idx = (viewer.anim_idx + step) % viewer.anim_len
|
||||
|
||||
def stop_callback(_, viewer):
|
||||
viewer.do_animate = False
|
||||
|
||||
def loop_callback(_, viewer):
|
||||
viewer.num_repeats *= -1
|
||||
|
||||
registered_keys = {
|
||||
"p": (pause_play_callback, [self]),
|
||||
".": (step_callback, [self, 1]),
|
||||
",": (step_callback, [self, -1]),
|
||||
"l": (loop_callback, [self]),
|
||||
"s": (stop_callback, [self]),
|
||||
}
|
||||
|
||||
self.viewer = pyrender.Viewer(
|
||||
self.scene,
|
||||
run_in_thread=True,
|
||||
viewport_size=self.viewport_size,
|
||||
use_raymond_lighting=True,
|
||||
registered_keys=registered_keys,
|
||||
)
|
||||
|
||||
self.num_repeats = num_repeats
|
||||
self.do_animate = False
|
||||
self.is_paused = False
|
||||
|
||||
def acquire_lock(self):
|
||||
self.viewer.render_lock.acquire()
|
||||
|
||||
def release_lock(self):
|
||||
self.viewer.render_lock.release()
|
||||
|
||||
def close(self):
|
||||
self.viewer.close_external()
|
||||
|
||||
def animate(self, save_name=None, **kwargs):
|
||||
print("===================")
|
||||
print(f"PLAYING ANIMATION {self.num_repeats} TIMES")
|
||||
print("VIEWER CONTROLS:")
|
||||
print("p\tpause/play")
|
||||
print("l\tloop/unloop")
|
||||
print(".\tstep forward")
|
||||
print(",\tstep backward")
|
||||
print("s\tstop")
|
||||
print("===================")
|
||||
|
||||
self.anim_idx = 0
|
||||
self.do_animate = True
|
||||
frame_dur = 1 / self.fps
|
||||
if save_name is not None:
|
||||
self.viewer.viewer_flags["record"] = True
|
||||
|
||||
# set up initial frame
|
||||
self.update_frame()
|
||||
|
||||
while self.do_animate:
|
||||
if self.is_paused:
|
||||
self.update_frame()
|
||||
continue
|
||||
|
||||
if (
|
||||
self.num_repeats > 0
|
||||
and self.anim_idx >= self.num_repeats * self.anim_len
|
||||
):
|
||||
break
|
||||
|
||||
if save_name is not None and self.anim_idx >= self.anim_len:
|
||||
self.viewer.viewer_flags["record"] = False
|
||||
|
||||
start_time = time.time()
|
||||
self.update_frame()
|
||||
render_time = time.time() - start_time
|
||||
|
||||
self.anim_idx += 1
|
||||
sleep_time = frame_dur - render_time
|
||||
if sleep_time > 0:
|
||||
time.sleep(sleep_time)
|
||||
|
||||
if save_name is not None:
|
||||
print(f"SAVING RECORDING TO {save_name}")
|
||||
self.viewer.save_gif(f"{save_name}.{self.ext}")
|
||||
|
||||
|
||||
def make_checkerboard(
|
||||
length=25.0,
|
||||
color0=[0.8, 0.9, 0.9],
|
||||
color1=[0.6, 0.7, 0.7],
|
||||
tile_width=0.5,
|
||||
alpha=1.0,
|
||||
up="y",
|
||||
):
|
||||
v, f, _, fc = checkerboard_geometry(length, color0, color1, tile_width, alpha, up)
|
||||
return trimesh.Trimesh(v, f, face_colors=fc, process=False)
|
||||
|
||||
|
||||
def make_pyrender_camera(img_size, intrins=None):
|
||||
if intrins is not None:
|
||||
print("USING INTRINSICS CAMERA", intrins)
|
||||
fx, fy, cx, cy = intrins
|
||||
return pyrender.IntrinsicsCamera(fx, fy, cx, cy)
|
||||
|
||||
W, H = img_size
|
||||
focal = 0.5 * (H + W)
|
||||
yfov = 2 * np.arctan(0.5 * H / focal)
|
||||
print("USING PERSPECTIVE CAMERA", H, W, focal)
|
||||
return pyrender.PerspectiveCamera(yfov=yfov, aspectRatio=W / H)
|
||||
|
||||
|
||||
def make_camera_marker(radius=0.1, height=0.2, up="y", transform=None):
|
||||
"""
|
||||
:param radius (default 0.1) radius of pyramid base, diagonal of image plane
|
||||
:param height (default 0.2) height of pyramid, focal length
|
||||
:param up (default y) camera up vector
|
||||
:param transform (default None) (4, 4) cam to world transform
|
||||
"""
|
||||
verts, faces, face_colors = camera_marker_geometry(radius, height, up)
|
||||
if transform is not None:
|
||||
assert transform.shape == (4, 4)
|
||||
verts = (
|
||||
np.einsum("ij,nj->ni", transform[:3, :3], verts) + transform[None, :3, 3]
|
||||
)
|
||||
return trimesh.Trimesh(verts, faces, face_colors=face_colors, process=False)
|
||||
|
||||
|
||||
def get_light_poses(n_lights=5, elevation=np.pi / 3, dist=12):
|
||||
# get lights in a circle around origin at elevation
|
||||
thetas = elevation * np.ones(n_lights)
|
||||
phis = 2 * np.pi * np.arange(n_lights) / n_lights
|
||||
poses = []
|
||||
trans = make_translation(torch.tensor([0, 0, dist]))
|
||||
for phi, theta in zip(phis, thetas):
|
||||
rot = make_rotation(rx=-theta, ry=phi, order="xyz")
|
||||
poses.append((rot @ trans).numpy())
|
||||
return poses
|
||||
|
||||
|
||||
def get_raymond_light_poses(up="z"):
|
||||
thetas = np.pi * np.ones(3) / 6
|
||||
phis = 2 * np.pi * np.arange(3) / 3
|
||||
|
||||
poses = []
|
||||
|
||||
for phi, theta in zip(phis, thetas):
|
||||
xp = np.sin(theta) * np.cos(phi)
|
||||
yp = np.sin(theta) * np.sin(phi)
|
||||
zp = np.cos(theta)
|
||||
|
||||
if up == "y":
|
||||
z = np.array([xp, zp, -yp])
|
||||
else:
|
||||
z = np.array([xp, yp, zp])
|
||||
z = z / np.linalg.norm(z)
|
||||
x = np.array([-z[1], z[0], 0.0])
|
||||
if np.linalg.norm(x) == 0:
|
||||
x = np.array([1.0, 0.0, 0.0])
|
||||
x = x / np.linalg.norm(x)
|
||||
y = np.cross(z, x)
|
||||
|
||||
matrix = np.eye(4)
|
||||
matrix[:3, :3] = np.c_[x, y, z]
|
||||
poses.append(matrix)
|
||||
return poses
|
||||
+28
@@ -0,0 +1,28 @@
|
||||
git+https://github.com/nghorbani/configer
|
||||
git+https://github.com/mattloper/chumpy
|
||||
setuptools==59.5.0
|
||||
torchgeometry==0.1.2
|
||||
tensorboard
|
||||
numpy
|
||||
smplx
|
||||
pyrender
|
||||
open3d
|
||||
imageio-ffmpeg
|
||||
matplotlib
|
||||
opencv-python
|
||||
scipy
|
||||
scikit-image
|
||||
joblib
|
||||
cython
|
||||
tqdm
|
||||
hydra-core
|
||||
pyyaml
|
||||
gdown
|
||||
dill
|
||||
motmetrics
|
||||
einops
|
||||
mmcv==1.3.9
|
||||
timm==0.4.9
|
||||
xtcocotools
|
||||
pandas==1.4.0
|
||||
ultralytics==8.1.34
|
||||
@@ -0,0 +1,94 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Convert MP4 videos to optimized GIFs for GitHub README display.
|
||||
This script creates smaller, optimized GIFs suitable for web display.
|
||||
"""
|
||||
import os
|
||||
import sys
|
||||
import imageio
|
||||
from pathlib import Path
|
||||
|
||||
def convert_mp4_to_gif(mp4_path, gif_path, max_frames=150, fps=10, scale=0.5):
|
||||
"""
|
||||
Convert MP4 to optimized GIF.
|
||||
|
||||
Args:
|
||||
mp4_path: Path to input MP4 file
|
||||
gif_path: Path to output GIF file
|
||||
max_frames: Maximum number of frames to include (for file size)
|
||||
fps: Frames per second for output GIF
|
||||
scale: Scale factor (0.5 = half size)
|
||||
"""
|
||||
print(f"Converting {mp4_path} to {gif_path}...")
|
||||
|
||||
try:
|
||||
# Read video
|
||||
reader = imageio.get_reader(mp4_path)
|
||||
fps_original = reader.get_meta_data()['fps']
|
||||
|
||||
# Calculate frame skip to get desired fps
|
||||
frame_skip = max(1, int(fps_original / fps))
|
||||
|
||||
frames = []
|
||||
frame_count = 0
|
||||
|
||||
for i, frame in enumerate(reader):
|
||||
if i % frame_skip == 0 and frame_count < max_frames:
|
||||
# Resize frame
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
img = Image.fromarray(frame)
|
||||
new_size = (int(img.width * scale), int(img.height * scale))
|
||||
img_resized = img.resize(new_size, Image.Resampling.LANCZOS)
|
||||
frames.append(np.array(img_resized))
|
||||
frame_count += 1
|
||||
|
||||
reader.close()
|
||||
|
||||
# Write GIF
|
||||
imageio.mimsave(gif_path, frames, fps=fps, loop=0)
|
||||
print(f" ✓ Created {gif_path} ({frame_count} frames, {fps} fps)")
|
||||
|
||||
# Get file sizes
|
||||
mp4_size = os.path.getsize(mp4_path) / (1024 * 1024) # MB
|
||||
gif_size = os.path.getsize(gif_path) / (1024 * 1024) # MB
|
||||
print(f" Size: {mp4_size:.1f} MB → {gif_size:.1f} MB")
|
||||
|
||||
except Exception as e:
|
||||
print(f" ✗ Error: {e}")
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def main():
|
||||
assets_dir = Path(__file__).parent.parent / "assets"
|
||||
|
||||
videos = [
|
||||
"droid_result.mp4",
|
||||
"vipe_result.mp4",
|
||||
"handedness1.mp4",
|
||||
"handedness2.mp4",
|
||||
"handedness3.mp4",
|
||||
"handedness4.mp4",
|
||||
]
|
||||
|
||||
print("Converting videos to GIFs for GitHub README...")
|
||||
print("=" * 60)
|
||||
|
||||
for video in videos:
|
||||
mp4_path = assets_dir / video
|
||||
gif_path = assets_dir / video.replace(".mp4", ".gif")
|
||||
|
||||
if not mp4_path.exists():
|
||||
print(f"⚠ Skipping {video} (not found)")
|
||||
continue
|
||||
|
||||
convert_mp4_to_gif(mp4_path, gif_path, max_frames=1000, fps=20, scale=0.6)
|
||||
print()
|
||||
|
||||
print("=" * 60)
|
||||
print("Done! Update README.md to use .gif files instead of .mp4")
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
||||
+34
@@ -0,0 +1,34 @@
|
||||
#!/usr/bin/env bash
|
||||
set -e
|
||||
|
||||
export CONDA_ENV_NAME=dynhamr
|
||||
|
||||
conda create -n $CONDA_ENV_NAME python=3.10 -y
|
||||
|
||||
conda activate $CONDA_ENV_NAME
|
||||
|
||||
# install pytorch using pip, update with appropriate cuda drivers if necessary
|
||||
pip install torch==1.13.0 torchvision==0.14.0 --index-url https://download.pytorch.org/whl/cu117
|
||||
# uncomment if pip installation isn't working
|
||||
# conda install pytorch=1.13.0 torchvision=0.14.0 pytorch-cuda=11.7 -c pytorch -c nvidia -y
|
||||
|
||||
# install pytorch scatter using pip, update with appropriate cuda drivers if necessary
|
||||
pip install torch-scatter -f https://data.pyg.org/whl/torch-1.13.0+cu117.html
|
||||
# uncomment if pip installation isn't working
|
||||
# conda install pytorch-scatter -c pyg -y
|
||||
|
||||
# install remaining requirements
|
||||
pip install -r requirements.txt
|
||||
|
||||
# install source
|
||||
pip install -e .
|
||||
|
||||
# install DROID-SLAM/DPVO
|
||||
cd third-party/DROID-SLAM
|
||||
python setup.py install
|
||||
cd ../..
|
||||
|
||||
# install HaMeR
|
||||
cd third-party/HaMeR
|
||||
pip install -e .[all]
|
||||
pip install -v -e third-party/ViTPose
|
||||
+30
@@ -0,0 +1,30 @@
|
||||
#!/usr/bin/env bash
|
||||
set -e
|
||||
|
||||
echo "Creating virtual environment"
|
||||
python3.10 -m venv .dynhamr
|
||||
echo "Activating virtual environment"
|
||||
|
||||
source $PWD/.dynhamr/bin/activate
|
||||
|
||||
# install pytorch
|
||||
$PWD/.dynhamr/bin/pip install torch==1.13.0 torchvision==0.14.0 --index-url https://download.pytorch.org/whl/cu117
|
||||
|
||||
# torch-scatter
|
||||
$PWD/.dynhamr/bin/pip install torch-scatter -f https://data.pyg.org/whl/torch-1.13.0+cu117.html
|
||||
|
||||
# install source
|
||||
$PWD/.dynhamr/bin/pip install -e .
|
||||
|
||||
# install remaining requirements
|
||||
$PWD/.dynhamr/bin/pip install -r requirements.txt
|
||||
|
||||
# install DROID-SLAM/DPVO
|
||||
cd third-party/DROID-SLAM
|
||||
python setup.py install
|
||||
cd ../..
|
||||
|
||||
# install HaMeR
|
||||
cd third-party/HaMeR
|
||||
pip install -e .[all]
|
||||
pip install -v -e third-party/ViTPose
|
||||
+21
@@ -0,0 +1,21 @@
|
||||
#!/usr/bin/env bash
|
||||
|
||||
|
||||
# Download for HaMeR
|
||||
gdown https://drive.google.com/uc?id=1mv7CUAnm73oKsEEG1xE3xH2C_oqcFSzT
|
||||
# Alternatively, you can use wget
|
||||
#wget https://www.cs.utexas.edu/~pavlakos/hamer/data/hamer_demo_data.tar.gz
|
||||
tar --warning=no-unknown-keyword --exclude=".*" -xvf hamer_demo_data.tar.gz
|
||||
rm -rf hamer_demo_data.tar.gz
|
||||
|
||||
# Download for SLAM
|
||||
gdown https://drive.google.com/uc?id=1VD1vGhl_NPzy8mza4Fx6vvqFpnlzZ86L
|
||||
mv droid.pth ./_DATA/
|
||||
|
||||
# Download for HMP
|
||||
gdown https://drive.google.com/uc?id=1_rp_SMxZA6Gjm6Z49M6hXam9kjMD6-GQ
|
||||
unzip hmp_model.zip
|
||||
mv hmp_model ./_DATA/
|
||||
rm -rf hmp_model.zip
|
||||
|
||||
wget https://huggingface.co/spaces/rolpotamias/WiLoR/resolve/main/pretrained_models/detector.pt -P ./third-party/hamer/pretrained_models/
|
||||
Vendored
+9
@@ -0,0 +1,9 @@
|
||||
from setuptools import find_packages, setup
|
||||
|
||||
setup(
|
||||
name="dyn-hamr",
|
||||
packages=find_packages(
|
||||
where="dyn-hamr",
|
||||
),
|
||||
package_dir={"": "dyn-hamr"},
|
||||
)
|
||||
@@ -0,0 +1,29 @@
|
||||
BSD 3-Clause License
|
||||
|
||||
Copyright (c) 2021, Princeton Vision & Learning Lab
|
||||
All rights reserved.
|
||||
|
||||
Redistribution and use in source and binary forms, with or without
|
||||
modification, are permitted provided that the following conditions are met:
|
||||
|
||||
1. Redistributions of source code must retain the above copyright notice, this
|
||||
list of conditions and the following disclaimer.
|
||||
|
||||
2. Redistributions in binary form must reproduce the above copyright notice,
|
||||
this list of conditions and the following disclaimer in the documentation
|
||||
and/or other materials provided with the distribution.
|
||||
|
||||
3. Neither the name of the copyright holder nor the names of its
|
||||
contributors may be used to endorse or promote products derived from
|
||||
this software without specific prior written permission.
|
||||
|
||||
THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
@@ -0,0 +1,139 @@
|
||||
# DROID-SLAM
|
||||
|
||||
|
||||
<!-- <center><img src="misc/DROID.png" width="640" style="center"></center> -->
|
||||
|
||||
|
||||
[](https://www.youtube.com/watch?v=GG78CSlSHSA)
|
||||
|
||||
|
||||
|
||||
[DROID-SLAM: Deep Visual SLAM for Monocular, Stereo, and RGB-D Cameras](https://arxiv.org/abs/2108.10869)
|
||||
Zachary Teed and Jia Deng
|
||||
|
||||
```
|
||||
@article{teed2021droid,
|
||||
title={{DROID-SLAM: Deep Visual SLAM for Monocular, Stereo, and RGB-D Cameras}},
|
||||
author={Teed, Zachary and Deng, Jia},
|
||||
journal={Advances in neural information processing systems},
|
||||
year={2021}
|
||||
}
|
||||
```
|
||||
|
||||
**Initial Code Release:** This repo currently provides a single GPU implementation of our monocular, stereo, and RGB-D SLAM systems. It currently contains demos, training, and evaluation scripts.
|
||||
|
||||
|
||||
## Requirements
|
||||
|
||||
To run the code you will need ...
|
||||
* **Inference:** Running the demos will require a GPU with at least 11G of memory.
|
||||
|
||||
* **Training:** Training requires a GPU with at least 24G of memory. We train on 4 x RTX-3090 GPUs.
|
||||
|
||||
## Getting Started
|
||||
1. Clone the repo using the `--recursive` flag
|
||||
```Bash
|
||||
git clone --recursive https://github.com/princeton-vl/DROID-SLAM.git
|
||||
```
|
||||
|
||||
2. Creating a new anaconda environment using the provided .yaml file. Use `environment_novis.yaml` to if you do not want to use the visualization
|
||||
```Bash
|
||||
conda env create -f environment.yaml
|
||||
pip install evo --upgrade --no-binary evo
|
||||
pip install gdown
|
||||
```
|
||||
|
||||
3. Compile the extensions (takes about 10 minutes)
|
||||
```Bash
|
||||
python setup.py install
|
||||
```
|
||||
|
||||
|
||||
## Demos
|
||||
|
||||
1. Download the model from google drive: [droid.pth](https://drive.google.com/file/d/1PpqVt1H4maBa_GbPJp4NwxRsd9jk-elh/view?usp=sharing)
|
||||
|
||||
2. Download some sample videos using the provided script.
|
||||
```Bash
|
||||
./tools/download_sample_data.sh
|
||||
```
|
||||
|
||||
Run the demo on any of the samples (all demos can be run on a GPU with 11G of memory). While running, press the "s" key to increase the filtering threshold (= more points) and "a" to decrease the filtering threshold (= fewer points). To save the reconstruction with full resolution depth maps use the `--reconstruction_path` flag.
|
||||
|
||||
|
||||
```Python
|
||||
python demo.py --imagedir=data/abandonedfactory --calib=calib/tartan.txt --stride=2
|
||||
```
|
||||
|
||||
```Python
|
||||
python demo.py --imagedir=data/sfm_bench/rgb --calib=calib/eth.txt
|
||||
```
|
||||
|
||||
```Python
|
||||
python demo.py --imagedir=data/Barn --calib=calib/barn.txt --stride=1 --backend_nms=4
|
||||
```
|
||||
|
||||
```Python
|
||||
python demo.py --imagedir=data/mav0/cam0/data --calib=calib/euroc.txt --t0=150
|
||||
```
|
||||
|
||||
```Python
|
||||
python demo.py --imagedir=data/rgbd_dataset_freiburg3_cabinet/rgb --calib=calib/tum3.txt
|
||||
```
|
||||
|
||||
|
||||
**Running on your own data:** All you need is a calibration file. Calibration files are in the form
|
||||
```
|
||||
fx fy cx cy [k1 k2 p1 p2 [ k3 [ k4 k5 k6 ]]]
|
||||
```
|
||||
with parameters in brackets optional.
|
||||
|
||||
## Evaluation
|
||||
We provide evaluation scripts for TartanAir, EuRoC, and TUM. EuRoC and TUM can be run on a 1080Ti. The TartanAir and ETH will require 24G of memory.
|
||||
|
||||
### TartanAir (Mono + Stereo)
|
||||
Download the [TartanAir](https://theairlab.org/tartanair-dataset/) dataset using the script `thirdparty/tartanair_tools/download_training.py` and put them in `datasets/TartanAir`
|
||||
```Bash
|
||||
./tools/validate_tartanair.sh --plot_curve # monocular eval
|
||||
./tools/validate_tartanair.sh --plot_curve --stereo # stereo eval
|
||||
```
|
||||
|
||||
### EuRoC (Mono + Stereo)
|
||||
Download the [EuRoC](https://projects.asl.ethz.ch/datasets/doku.php?id=kmavvisualinertialdatasets) sequences (ASL format) and put them in `datasets/EuRoC`
|
||||
```Bash
|
||||
./tools/evaluate_euroc.sh # monocular eval
|
||||
./tools/evaluate_euroc.sh --stereo # stereo eval
|
||||
```
|
||||
|
||||
### TUM-RGBD (Mono)
|
||||
Download the fr1 sequences from [TUM-RGBD](https://vision.in.tum.de/data/datasets/rgbd-dataset/download) and put them in `datasets/TUM-RGBD`
|
||||
```Bash
|
||||
./tools/evaluate_tum.sh # monocular eval
|
||||
```
|
||||
|
||||
### ETH3D (RGB-D)
|
||||
Download the [ETH3D](https://www.eth3d.net/slam_datasets) dataset
|
||||
```Bash
|
||||
./tools/evaluate_eth3d.sh # RGB-D eval
|
||||
```
|
||||
|
||||
## Training
|
||||
|
||||
First download the TartanAir dataset. The download script can be found in `thirdparty/tartanair_tools/download_training.py`. You will only need the `rgb` and `depth` data.
|
||||
|
||||
```
|
||||
python download_training.py --rgb --depth
|
||||
```
|
||||
|
||||
You can then run the training script. We use 4x3090 RTX GPUs for training which takes approximatly 1 week. If you use a different number of GPUs, adjust the learning rate accordingly.
|
||||
|
||||
**Note:** On the first training run, covisibility is computed between all pairs of frames. This can take several hours, but the results are cached so that future training runs will start immediately.
|
||||
|
||||
|
||||
```
|
||||
python train.py --datapath=<path to tartanair> --gpus=4 --lr=0.00025
|
||||
```
|
||||
|
||||
|
||||
## Acknowledgements
|
||||
Data from [TartanAir](https://theairlab.org/tartanair-dataset/) was used to train our model. We additionally use evaluation tools from [evo](https://github.com/MichaelGrupp/evo) and [tartanair_tools](https://github.com/castacks/tartanair_tools).
|
||||
@@ -0,0 +1 @@
|
||||
1161.545689 1161.545689 960.000000 540.000000 -0.025158 0.0 0.0 0.0
|
||||
@@ -0,0 +1 @@
|
||||
726.21081542969 726.21081542969 359.2048034668 202.47247314453
|
||||
@@ -0,0 +1 @@
|
||||
458.654 457.296 367.215 248.375 -0.28340811 0.07395907 0.00019359 1.76187114e-05
|
||||
@@ -0,0 +1 @@
|
||||
320.0 320.0 320.0 240.0
|
||||
@@ -0,0 +1 @@
|
||||
535.4 539.2 320.1 247.6
|
||||
+134
@@ -0,0 +1,134 @@
|
||||
import sys
|
||||
sys.path.append('droid_slam')
|
||||
|
||||
from tqdm import tqdm
|
||||
import numpy as np
|
||||
import torch
|
||||
import lietorch
|
||||
import cv2
|
||||
import os
|
||||
import glob
|
||||
import time
|
||||
import argparse
|
||||
|
||||
from torch.multiprocessing import Process
|
||||
from droid import Droid
|
||||
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
def show_image(image):
|
||||
image = image.permute(1, 2, 0).cpu().numpy()
|
||||
cv2.imshow('image', image / 255.0)
|
||||
cv2.waitKey(1)
|
||||
|
||||
def image_stream(imagedir, calib, stride):
|
||||
""" image generator """
|
||||
|
||||
calib = np.loadtxt(calib, delimiter=" ")
|
||||
fx, fy, cx, cy = calib[:4]
|
||||
|
||||
K = np.eye(3)
|
||||
K[0,0] = fx
|
||||
K[0,2] = cx
|
||||
K[1,1] = fy
|
||||
K[1,2] = cy
|
||||
|
||||
image_list = sorted(os.listdir(imagedir))[::stride]
|
||||
|
||||
for t, imfile in enumerate(image_list):
|
||||
image = cv2.imread(os.path.join(imagedir, imfile))
|
||||
if len(calib) > 4:
|
||||
image = cv2.undistort(image, K, calib[4:])
|
||||
|
||||
h0, w0, _ = image.shape
|
||||
h1 = int(h0 * np.sqrt((384 * 512) / (h0 * w0)))
|
||||
w1 = int(w0 * np.sqrt((384 * 512) / (h0 * w0)))
|
||||
|
||||
image = cv2.resize(image, (w1, h1))
|
||||
image = image[:h1-h1%8, :w1-w1%8]
|
||||
image = torch.as_tensor(image).permute(2, 0, 1)
|
||||
|
||||
intrinsics = torch.as_tensor([fx, fy, cx, cy])
|
||||
intrinsics[0::2] *= (w1 / w0)
|
||||
intrinsics[1::2] *= (h1 / h0)
|
||||
|
||||
yield t, image[None], intrinsics
|
||||
|
||||
|
||||
def save_reconstruction(droid, reconstruction_path):
|
||||
|
||||
from pathlib import Path
|
||||
import random
|
||||
import string
|
||||
|
||||
t = droid.video.counter.value
|
||||
tstamps = droid.video.tstamp[:t].cpu().numpy()
|
||||
images = droid.video.images[:t].cpu().numpy()
|
||||
disps = droid.video.disps_up[:t].cpu().numpy()
|
||||
poses = droid.video.poses[:t].cpu().numpy()
|
||||
intrinsics = droid.video.intrinsics[:t].cpu().numpy()
|
||||
|
||||
Path("reconstructions/{}".format(reconstruction_path)).mkdir(parents=True, exist_ok=True)
|
||||
np.save("reconstructions/{}/tstamps.npy".format(reconstruction_path), tstamps)
|
||||
np.save("reconstructions/{}/images.npy".format(reconstruction_path), images)
|
||||
np.save("reconstructions/{}/disps.npy".format(reconstruction_path), disps)
|
||||
np.save("reconstructions/{}/poses.npy".format(reconstruction_path), poses)
|
||||
np.save("reconstructions/{}/intrinsics.npy".format(reconstruction_path), intrinsics)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--imagedir", type=str, help="path to image directory")
|
||||
parser.add_argument("--calib", type=str, help="path to calibration file")
|
||||
parser.add_argument("--t0", default=0, type=int, help="starting frame")
|
||||
parser.add_argument("--stride", default=3, type=int, help="frame stride")
|
||||
|
||||
parser.add_argument("--weights", default="droid.pth")
|
||||
parser.add_argument("--buffer", type=int, default=512)
|
||||
parser.add_argument("--image_size", default=[240, 320])
|
||||
parser.add_argument("--disable_vis", action="store_true")
|
||||
|
||||
parser.add_argument("--beta", type=float, default=0.3, help="weight for translation / rotation components of flow")
|
||||
parser.add_argument("--filter_thresh", type=float, default=2.4, help="how much motion before considering new keyframe")
|
||||
parser.add_argument("--warmup", type=int, default=8, help="number of warmup frames")
|
||||
parser.add_argument("--keyframe_thresh", type=float, default=4.0, help="threshold to create a new keyframe")
|
||||
parser.add_argument("--frontend_thresh", type=float, default=16.0, help="add edges between frames whithin this distance")
|
||||
parser.add_argument("--frontend_window", type=int, default=25, help="frontend optimization window")
|
||||
parser.add_argument("--frontend_radius", type=int, default=2, help="force edges between frames within radius")
|
||||
parser.add_argument("--frontend_nms", type=int, default=1, help="non-maximal supression of edges")
|
||||
|
||||
parser.add_argument("--backend_thresh", type=float, default=22.0)
|
||||
parser.add_argument("--backend_radius", type=int, default=2)
|
||||
parser.add_argument("--backend_nms", type=int, default=3)
|
||||
parser.add_argument("--upsample", action="store_true")
|
||||
parser.add_argument("--reconstruction_path", help="path to saved reconstruction")
|
||||
args = parser.parse_args()
|
||||
|
||||
args.stereo = False
|
||||
torch.multiprocessing.set_start_method('spawn')
|
||||
|
||||
droid = None
|
||||
|
||||
# need high resolution depths
|
||||
if args.reconstruction_path is not None:
|
||||
args.upsample = True
|
||||
|
||||
tstamps = []
|
||||
for (t, image, intrinsics) in tqdm(image_stream(args.imagedir, args.calib, args.stride)):
|
||||
if t < args.t0:
|
||||
continue
|
||||
|
||||
if not args.disable_vis:
|
||||
show_image(image[0])
|
||||
|
||||
if droid is None:
|
||||
args.image_size = [image.shape[2], image.shape[3]]
|
||||
droid = Droid(args)
|
||||
|
||||
droid.track(t, image, intrinsics=intrinsics)
|
||||
|
||||
if args.reconstruction_path is not None:
|
||||
save_reconstruction(droid, args.reconstruction_path)
|
||||
|
||||
traj_est = droid.terminate(image_stream(args.imagedir, args.calib, args.stride))
|
||||
+1
@@ -0,0 +1 @@
|
||||
|
||||
+58
@@ -0,0 +1,58 @@
|
||||
import torch
|
||||
import torchvision.transforms as transforms
|
||||
import numpy as np
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
class RGBDAugmentor:
|
||||
""" perform augmentation on RGB-D video """
|
||||
|
||||
def __init__(self, crop_size):
|
||||
self.crop_size = crop_size
|
||||
self.augcolor = transforms.Compose([
|
||||
transforms.ToPILImage(),
|
||||
transforms.ColorJitter(brightness=0.25, contrast=0.25, saturation=0.25, hue=0.4/3.14),
|
||||
transforms.RandomGrayscale(p=0.1),
|
||||
transforms.ToTensor()])
|
||||
|
||||
self.max_scale = 0.25
|
||||
|
||||
def spatial_transform(self, images, depths, poses, intrinsics):
|
||||
""" cropping and resizing """
|
||||
ht, wd = images.shape[2:]
|
||||
|
||||
max_scale = self.max_scale
|
||||
min_scale = np.log2(np.maximum(
|
||||
(self.crop_size[0] + 1) / float(ht),
|
||||
(self.crop_size[1] + 1) / float(wd)))
|
||||
|
||||
scale = 2 ** np.random.uniform(min_scale, max_scale)
|
||||
intrinsics = scale * intrinsics
|
||||
depths = depths.unsqueeze(dim=1)
|
||||
|
||||
images = F.interpolate(images, scale_factor=scale, mode='bilinear',
|
||||
align_corners=False, recompute_scale_factor=True)
|
||||
|
||||
depths = F.interpolate(depths, scale_factor=scale, recompute_scale_factor=True)
|
||||
|
||||
# always perform center crop (TODO: try non-center crops)
|
||||
y0 = (images.shape[2] - self.crop_size[0]) // 2
|
||||
x0 = (images.shape[3] - self.crop_size[1]) // 2
|
||||
|
||||
intrinsics = intrinsics - torch.tensor([0.0, 0.0, x0, y0])
|
||||
images = images[:, :, y0:y0+self.crop_size[0], x0:x0+self.crop_size[1]]
|
||||
depths = depths[:, :, y0:y0+self.crop_size[0], x0:x0+self.crop_size[1]]
|
||||
|
||||
depths = depths.squeeze(dim=1)
|
||||
return images, poses, depths, intrinsics
|
||||
|
||||
def color_transform(self, images):
|
||||
""" color jittering """
|
||||
num, ch, ht, wd = images.shape
|
||||
images = images.permute(1, 2, 3, 0).reshape(ch, ht, wd*num)
|
||||
images = 255 * self.augcolor(images[[2,1,0]] / 255.0)
|
||||
return images[[2,1,0]].reshape(ch, ht, wd, num).permute(3,0,1,2).contiguous()
|
||||
|
||||
def __call__(self, images, poses, depths, intrinsics):
|
||||
images = self.color_transform(images)
|
||||
return self.spatial_transform(images, depths, poses, intrinsics)
|
||||
@@ -0,0 +1,157 @@
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.utils.data as data
|
||||
import torch.nn.functional as F
|
||||
|
||||
import csv
|
||||
import os
|
||||
import cv2
|
||||
import math
|
||||
import random
|
||||
import json
|
||||
import pickle
|
||||
import os.path as osp
|
||||
|
||||
from .augmentation import RGBDAugmentor
|
||||
from .rgbd_utils import *
|
||||
|
||||
class RGBDDataset(data.Dataset):
|
||||
def __init__(self, name, datapath, n_frames=4, crop_size=[384,512], fmin=8.0, fmax=75.0, do_aug=True):
|
||||
""" Base class for RGBD dataset """
|
||||
self.aug = None
|
||||
self.root = datapath
|
||||
self.name = name
|
||||
|
||||
self.n_frames = n_frames
|
||||
self.fmin = fmin # exclude very easy examples
|
||||
self.fmax = fmax # exclude very hard examples
|
||||
|
||||
if do_aug:
|
||||
self.aug = RGBDAugmentor(crop_size=crop_size)
|
||||
|
||||
# building dataset is expensive, cache so only needs to be performed once
|
||||
cur_path = osp.dirname(osp.abspath(__file__))
|
||||
if not os.path.isdir(osp.join(cur_path, 'cache')):
|
||||
os.mkdir(osp.join(cur_path, 'cache'))
|
||||
|
||||
cache_path = osp.join(cur_path, 'cache', '{}.pickle'.format(self.name))
|
||||
|
||||
if osp.isfile(cache_path):
|
||||
scene_info = pickle.load(open(cache_path, 'rb'))[0]
|
||||
else:
|
||||
scene_info = self._build_dataset()
|
||||
with open(cache_path, 'wb') as cachefile:
|
||||
pickle.dump((scene_info,), cachefile)
|
||||
|
||||
self.scene_info = scene_info
|
||||
self._build_dataset_index()
|
||||
|
||||
def _build_dataset_index(self):
|
||||
self.dataset_index = []
|
||||
for scene in self.scene_info:
|
||||
if not self.__class__.is_test_scene(scene):
|
||||
graph = self.scene_info[scene]['graph']
|
||||
for i in graph:
|
||||
if len(graph[i][0]) > self.n_frames:
|
||||
self.dataset_index.append((scene, i))
|
||||
else:
|
||||
print("Reserving {} for validation".format(scene))
|
||||
|
||||
@staticmethod
|
||||
def image_read(image_file):
|
||||
return cv2.imread(image_file)
|
||||
|
||||
@staticmethod
|
||||
def depth_read(depth_file):
|
||||
return np.load(depth_file)
|
||||
|
||||
def build_frame_graph(self, poses, depths, intrinsics, f=16, max_flow=256):
|
||||
""" compute optical flow distance between all pairs of frames """
|
||||
def read_disp(fn):
|
||||
depth = self.__class__.depth_read(fn)[f//2::f, f//2::f]
|
||||
depth[depth < 0.01] = np.mean(depth)
|
||||
return 1.0 / depth
|
||||
|
||||
poses = np.array(poses)
|
||||
intrinsics = np.array(intrinsics) / f
|
||||
|
||||
disps = np.stack(list(map(read_disp, depths)), 0)
|
||||
d = f * compute_distance_matrix_flow(poses, disps, intrinsics)
|
||||
|
||||
# uncomment for nice visualization
|
||||
# import matplotlib.pyplot as plt
|
||||
# plt.imshow(d)
|
||||
# plt.show()
|
||||
|
||||
graph = {}
|
||||
for i in range(d.shape[0]):
|
||||
j, = np.where(d[i] < max_flow)
|
||||
graph[i] = (j, d[i,j])
|
||||
|
||||
return graph
|
||||
|
||||
def __getitem__(self, index):
|
||||
""" return training video """
|
||||
|
||||
index = index % len(self.dataset_index)
|
||||
scene_id, ix = self.dataset_index[index]
|
||||
|
||||
frame_graph = self.scene_info[scene_id]['graph']
|
||||
images_list = self.scene_info[scene_id]['images']
|
||||
depths_list = self.scene_info[scene_id]['depths']
|
||||
poses_list = self.scene_info[scene_id]['poses']
|
||||
intrinsics_list = self.scene_info[scene_id]['intrinsics']
|
||||
|
||||
inds = [ ix ]
|
||||
while len(inds) < self.n_frames:
|
||||
# get other frames within flow threshold
|
||||
k = (frame_graph[ix][1] > self.fmin) & (frame_graph[ix][1] < self.fmax)
|
||||
frames = frame_graph[ix][0][k]
|
||||
|
||||
# prefer frames forward in time
|
||||
if np.count_nonzero(frames[frames > ix]):
|
||||
ix = np.random.choice(frames[frames > ix])
|
||||
|
||||
elif np.count_nonzero(frames):
|
||||
ix = np.random.choice(frames)
|
||||
|
||||
inds += [ ix ]
|
||||
|
||||
images, depths, poses, intrinsics = [], [], [], []
|
||||
for i in inds:
|
||||
images.append(self.__class__.image_read(images_list[i]))
|
||||
depths.append(self.__class__.depth_read(depths_list[i]))
|
||||
poses.append(poses_list[i])
|
||||
intrinsics.append(intrinsics_list[i])
|
||||
|
||||
images = np.stack(images).astype(np.float32)
|
||||
depths = np.stack(depths).astype(np.float32)
|
||||
poses = np.stack(poses).astype(np.float32)
|
||||
intrinsics = np.stack(intrinsics).astype(np.float32)
|
||||
|
||||
images = torch.from_numpy(images).float()
|
||||
images = images.permute(0, 3, 1, 2)
|
||||
|
||||
disps = torch.from_numpy(1.0 / depths)
|
||||
poses = torch.from_numpy(poses)
|
||||
intrinsics = torch.from_numpy(intrinsics)
|
||||
|
||||
if self.aug is not None:
|
||||
images, poses, disps, intrinsics = \
|
||||
self.aug(images, poses, disps, intrinsics)
|
||||
|
||||
# scale scene
|
||||
if len(disps[disps>0.01]) > 0:
|
||||
s = disps[disps>0.01].mean()
|
||||
disps = disps / s
|
||||
poses[...,:3] *= s
|
||||
|
||||
return images, poses, disps, intrinsics
|
||||
|
||||
def __len__(self):
|
||||
return len(self.dataset_index)
|
||||
|
||||
def __imul__(self, x):
|
||||
self.dataset_index *= x
|
||||
return self
|
||||
+82
@@ -0,0 +1,82 @@
|
||||
|
||||
import pickle
|
||||
import os
|
||||
import os.path as osp
|
||||
|
||||
# RGBD-Dataset
|
||||
from .tartan import TartanAir
|
||||
|
||||
from .stream import ImageStream
|
||||
from .stream import StereoStream
|
||||
from .stream import RGBDStream
|
||||
|
||||
# streaming datasets for inference
|
||||
from .tartan import TartanAirStream
|
||||
from .tartan import TartanAirTestStream
|
||||
|
||||
def dataset_factory(dataset_list, **kwargs):
|
||||
""" create a combined dataset """
|
||||
|
||||
from torch.utils.data import ConcatDataset
|
||||
|
||||
dataset_map = { 'tartan': (TartanAir, ) }
|
||||
db_list = []
|
||||
for key in dataset_list:
|
||||
# cache datasets for faster future loading
|
||||
db = dataset_map[key][0](**kwargs)
|
||||
|
||||
print("Dataset {} has {} images".format(key, len(db)))
|
||||
db_list.append(db)
|
||||
|
||||
return ConcatDataset(db_list)
|
||||
|
||||
|
||||
def create_datastream(dataset_path, **kwargs):
|
||||
""" create data_loader to stream images 1 by 1 """
|
||||
|
||||
from torch.utils.data import DataLoader
|
||||
|
||||
if osp.isfile(osp.join(dataset_path, 'calibration.txt')):
|
||||
db = ETH3DStream(dataset_path, **kwargs)
|
||||
|
||||
elif osp.isdir(osp.join(dataset_path, 'image_left')):
|
||||
db = TartanAirStream(dataset_path, **kwargs)
|
||||
|
||||
elif osp.isfile(osp.join(dataset_path, 'rgb.txt')):
|
||||
db = TUMStream(dataset_path, **kwargs)
|
||||
|
||||
elif osp.isdir(osp.join(dataset_path, 'mav0')):
|
||||
db = EurocStream(dataset_path, **kwargs)
|
||||
|
||||
elif osp.isfile(osp.join(dataset_path, 'calib.txt')):
|
||||
db = KITTIStream(dataset_path, **kwargs)
|
||||
|
||||
else:
|
||||
# db = TartanAirStream(dataset_path, **kwargs)
|
||||
db = TartanAirTestStream(dataset_path, **kwargs)
|
||||
|
||||
stream = DataLoader(db, shuffle=False, batch_size=1, num_workers=4)
|
||||
return stream
|
||||
|
||||
|
||||
def create_imagestream(dataset_path, **kwargs):
|
||||
""" create data_loader to stream images 1 by 1 """
|
||||
from torch.utils.data import DataLoader
|
||||
|
||||
db = ImageStream(dataset_path, **kwargs)
|
||||
return DataLoader(db, shuffle=False, batch_size=1, num_workers=4)
|
||||
|
||||
def create_stereostream(dataset_path, **kwargs):
|
||||
""" create data_loader to stream images 1 by 1 """
|
||||
from torch.utils.data import DataLoader
|
||||
|
||||
db = StereoStream(dataset_path, **kwargs)
|
||||
return DataLoader(db, shuffle=False, batch_size=1, num_workers=4)
|
||||
|
||||
def create_rgbdstream(dataset_path, **kwargs):
|
||||
""" create data_loader to stream images 1 by 1 """
|
||||
from torch.utils.data import DataLoader
|
||||
|
||||
db = RGBDStream(dataset_path, **kwargs)
|
||||
return DataLoader(db, shuffle=False, batch_size=1, num_workers=4)
|
||||
|
||||
+190
@@ -0,0 +1,190 @@
|
||||
import numpy as np
|
||||
import os.path as osp
|
||||
|
||||
import torch
|
||||
from lietorch import SE3
|
||||
|
||||
import geom.projective_ops as pops
|
||||
from scipy.spatial.transform import Rotation
|
||||
|
||||
|
||||
def parse_list(filepath, skiprows=0):
|
||||
""" read list data """
|
||||
data = np.loadtxt(filepath, delimiter=' ', dtype=np.unicode_, skiprows=skiprows)
|
||||
return data
|
||||
|
||||
def associate_frames(tstamp_image, tstamp_depth, tstamp_pose, max_dt=1.0):
|
||||
""" pair images, depths, and poses """
|
||||
associations = []
|
||||
for i, t in enumerate(tstamp_image):
|
||||
if tstamp_pose is None:
|
||||
j = np.argmin(np.abs(tstamp_depth - t))
|
||||
if (np.abs(tstamp_depth[j] - t) < max_dt):
|
||||
associations.append((i, j))
|
||||
|
||||
else:
|
||||
j = np.argmin(np.abs(tstamp_depth - t))
|
||||
k = np.argmin(np.abs(tstamp_pose - t))
|
||||
|
||||
if (np.abs(tstamp_depth[j] - t) < max_dt) and \
|
||||
(np.abs(tstamp_pose[k] - t) < max_dt):
|
||||
associations.append((i, j, k))
|
||||
|
||||
return associations
|
||||
|
||||
def loadtum(datapath, frame_rate=-1):
|
||||
""" read video data in tum-rgbd format """
|
||||
if osp.isfile(osp.join(datapath, 'groundtruth.txt')):
|
||||
pose_list = osp.join(datapath, 'groundtruth.txt')
|
||||
|
||||
elif osp.isfile(osp.join(datapath, 'pose.txt')):
|
||||
pose_list = osp.join(datapath, 'pose.txt')
|
||||
|
||||
else:
|
||||
return None, None, None, None
|
||||
|
||||
image_list = osp.join(datapath, 'rgb.txt')
|
||||
depth_list = osp.join(datapath, 'depth.txt')
|
||||
|
||||
calib_path = osp.join(datapath, 'calibration.txt')
|
||||
intrinsic = None
|
||||
if osp.isfile(calib_path):
|
||||
intrinsic = np.loadtxt(calib_path, delimiter=' ')
|
||||
intrinsic = intrinsic.astype(np.float64)
|
||||
|
||||
image_data = parse_list(image_list)
|
||||
depth_data = parse_list(depth_list)
|
||||
pose_data = parse_list(pose_list, skiprows=1)
|
||||
pose_vecs = pose_data[:,1:].astype(np.float64)
|
||||
|
||||
tstamp_image = image_data[:,0].astype(np.float64)
|
||||
tstamp_depth = depth_data[:,0].astype(np.float64)
|
||||
tstamp_pose = pose_data[:,0].astype(np.float64)
|
||||
associations = associate_frames(tstamp_image, tstamp_depth, tstamp_pose)
|
||||
|
||||
# print(len(tstamp_image))
|
||||
# print(len(associations))
|
||||
|
||||
indicies = range(len(associations))[::5]
|
||||
|
||||
# indicies = [ 0 ]
|
||||
# for i in range(1, len(associations)):
|
||||
# t0 = tstamp_image[associations[indicies[-1]][0]]
|
||||
# t1 = tstamp_image[associations[i][0]]
|
||||
# if t1 - t0 > 1.0 / frame_rate:
|
||||
# indicies += [ i ]
|
||||
|
||||
images, poses, depths, intrinsics, tstamps = [], [], [], [], []
|
||||
for ix in indicies:
|
||||
(i, j, k) = associations[ix]
|
||||
images += [ osp.join(datapath, image_data[i,1]) ]
|
||||
depths += [ osp.join(datapath, depth_data[j,1]) ]
|
||||
poses += [ pose_vecs[k] ]
|
||||
tstamps += [ tstamp_image[i] ]
|
||||
|
||||
if intrinsic is not None:
|
||||
intrinsics += [ intrinsic ]
|
||||
|
||||
return images, depths, poses, intrinsics, tstamps
|
||||
|
||||
|
||||
def all_pairs_distance_matrix(poses, beta=2.5):
|
||||
""" compute distance matrix between all pairs of poses """
|
||||
poses = np.array(poses, dtype=np.float32)
|
||||
poses[:,:3] *= beta # scale to balence rot + trans
|
||||
poses = SE3(torch.from_numpy(poses))
|
||||
|
||||
r = (poses[:,None].inv() * poses[None,:]).log()
|
||||
return r.norm(dim=-1).cpu().numpy()
|
||||
|
||||
def pose_matrix_to_quaternion(pose):
|
||||
""" convert 4x4 pose matrix to (t, q) """
|
||||
q = Rotation.from_matrix(pose[:3, :3]).as_quat()
|
||||
return np.concatenate([pose[:3, 3], q], axis=0)
|
||||
|
||||
def compute_distance_matrix_flow(poses, disps, intrinsics):
|
||||
""" compute flow magnitude between all pairs of frames """
|
||||
if not isinstance(poses, SE3):
|
||||
poses = torch.from_numpy(poses).float().cuda()[None]
|
||||
poses = SE3(poses).inv()
|
||||
|
||||
disps = torch.from_numpy(disps).float().cuda()[None]
|
||||
intrinsics = torch.from_numpy(intrinsics).float().cuda()[None]
|
||||
|
||||
N = poses.shape[1]
|
||||
|
||||
ii, jj = torch.meshgrid(torch.arange(N), torch.arange(N))
|
||||
ii = ii.reshape(-1).cuda()
|
||||
jj = jj.reshape(-1).cuda()
|
||||
|
||||
MAX_FLOW = 100.0
|
||||
matrix = np.zeros((N, N), dtype=np.float32)
|
||||
|
||||
s = 2048
|
||||
for i in range(0, ii.shape[0], s):
|
||||
flow1, val1 = pops.induced_flow(poses, disps, intrinsics, ii[i:i+s], jj[i:i+s])
|
||||
flow2, val2 = pops.induced_flow(poses, disps, intrinsics, jj[i:i+s], ii[i:i+s])
|
||||
|
||||
flow = torch.stack([flow1, flow2], dim=2)
|
||||
val = torch.stack([val1, val2], dim=2)
|
||||
|
||||
mag = flow.norm(dim=-1).clamp(max=MAX_FLOW)
|
||||
mag = mag.view(mag.shape[1], -1)
|
||||
val = val.view(val.shape[1], -1)
|
||||
|
||||
mag = (mag * val).mean(-1) / val.mean(-1)
|
||||
mag[val.mean(-1) < 0.7] = np.inf
|
||||
|
||||
i1 = ii[i:i+s].cpu().numpy()
|
||||
j1 = jj[i:i+s].cpu().numpy()
|
||||
matrix[i1, j1] = mag.cpu().numpy()
|
||||
|
||||
return matrix
|
||||
|
||||
|
||||
def compute_distance_matrix_flow2(poses, disps, intrinsics, beta=0.4):
|
||||
""" compute flow magnitude between all pairs of frames """
|
||||
# if not isinstance(poses, SE3):
|
||||
# poses = torch.from_numpy(poses).float().cuda()[None]
|
||||
# poses = SE3(poses).inv()
|
||||
|
||||
# disps = torch.from_numpy(disps).float().cuda()[None]
|
||||
# intrinsics = torch.from_numpy(intrinsics).float().cuda()[None]
|
||||
|
||||
N = poses.shape[1]
|
||||
|
||||
ii, jj = torch.meshgrid(torch.arange(N), torch.arange(N))
|
||||
ii = ii.reshape(-1)
|
||||
jj = jj.reshape(-1)
|
||||
|
||||
MAX_FLOW = 128.0
|
||||
matrix = np.zeros((N, N), dtype=np.float32)
|
||||
|
||||
s = 2048
|
||||
for i in range(0, ii.shape[0], s):
|
||||
flow1a, val1a = pops.induced_flow(poses, disps, intrinsics, ii[i:i+s], jj[i:i+s], tonly=True)
|
||||
flow1b, val1b = pops.induced_flow(poses, disps, intrinsics, ii[i:i+s], jj[i:i+s])
|
||||
flow2a, val2a = pops.induced_flow(poses, disps, intrinsics, jj[i:i+s], ii[i:i+s], tonly=True)
|
||||
flow2b, val2b = pops.induced_flow(poses, disps, intrinsics, ii[i:i+s], jj[i:i+s])
|
||||
|
||||
flow1 = flow1a + beta * flow1b
|
||||
val1 = val1a * val2b
|
||||
|
||||
flow2 = flow2a + beta * flow2b
|
||||
val2 = val2a * val2b
|
||||
|
||||
flow = torch.stack([flow1, flow2], dim=2)
|
||||
val = torch.stack([val1, val2], dim=2)
|
||||
|
||||
mag = flow.norm(dim=-1).clamp(max=MAX_FLOW)
|
||||
mag = mag.view(mag.shape[1], -1)
|
||||
val = val.view(val.shape[1], -1)
|
||||
|
||||
mag = (mag * val).mean(-1) / val.mean(-1)
|
||||
mag[val.mean(-1) < 0.8] = np.inf
|
||||
|
||||
i1 = ii[i:i+s].cpu().numpy()
|
||||
j1 = jj[i:i+s].cpu().numpy()
|
||||
matrix[i1, j1] = mag.cpu().numpy()
|
||||
|
||||
return matrix
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user