1use gloss_utils::nshare::ToNalgebra;
2use nalgebra as na;
3use nalgebra::clamp;
4use ndarray as nd;
5use ndarray::prelude::*;
6use std::{
7 f32::consts::PI,
8 ops::{Div, SubAssign},
9};
10pub fn interpolate_angle(cur_angle: f32, other_angle: f32, _cur_w: f32, other_w: f32) -> f32 {
11 let mut diff = other_angle - cur_angle;
12 if diff.abs() > PI {
13 if diff > 0.0 {
14 diff -= 2.0 * PI;
15 } else {
16 diff += 2.0 * PI;
17 }
18 }
19 cur_angle + other_w * diff
20}
21pub fn map(value: f32, in_min: f32, in_max: f32, out_min: f32, out_max: f32) -> f32 {
22 let value_clamped = clamp(value, in_min, in_max);
23 out_min + (out_max - out_min) * (value_clamped - in_min) / (in_max - in_min)
24}
25pub fn smootherstep(low: f32, high: f32, val: f32) -> f32 {
26 let t = map(val, low, high, 0.0, 1.0);
27 t * t * t * (t * (t * 6.0 - 15.0) + 10.0)
28}
29pub fn batch_rodrigues(full_pose: &nd::Array2<f32>) -> nd::Array3<f32> {
30 let mut rotations_per_join = ndarray::Array3::<f32>::zeros((full_pose.shape()[0], 3, 3));
31 for (idx, v) in full_pose.axis_iter(nd::Axis(0)).enumerate() {
32 let angle = v.iter().map(|x| x * x).sum::<f32>().sqrt();
33 let rot_dir = full_pose.row(idx).to_owned().div(angle + 1e-6);
34 let cos = angle.cos();
35 let sin = angle.sin();
36 let (rx, ry, rz) = (rot_dir[0], rot_dir[1], rot_dir[2]);
37 let k = array![[0.0, -rz, ry], [rz, 0.0, -rx], [-ry, rx, 0.0]];
38 let identity = ndarray::Array2::<f32>::eye(3);
39 let rot_mat = identity + sin * k.clone() + (1.0 - cos) * k.dot(&k);
40 rotations_per_join.slice_mut(s![idx, .., ..]).assign(&rot_mat);
41 }
42 rotations_per_join
43}
44pub fn euler2angleaxis(euler_x: f32, euler_y: f32, euler_z: f32) -> na::Vector3<f32> {
45 let c1 = f32::cos(euler_x / 2.0);
46 let c2 = f32::cos(euler_y / 2.0);
47 let c3 = f32::cos(euler_z / 2.0);
48 let s1 = f32::sin(euler_x / 2.0);
49 let s2 = f32::sin(euler_y / 2.0);
50 let s3 = f32::sin(euler_z / 2.0);
51 let rot = na::Quaternion::new(
52 c1 * c2 * c3 - s1 * s2 * s3,
53 s1 * c2 * c3 + c1 * s2 * s3,
54 c1 * s2 * c3 - s1 * c2 * s3,
55 c1 * c2 * s3 + s1 * s2 * c3,
56 );
57 let rot = na::UnitQuaternion::new_normalize(rot);
58 rot.scaled_axis()
59}
60pub fn interpolate_axis_angle(this_axis: &nd::Array1<f32>, other_axis: &nd::Array1<f32>, other_weight: f32) -> nd::Array1<f32> {
62 let this_axis_na = this_axis.clone().into_nalgebra();
63 let other_axis_na = other_axis.clone().into_nalgebra();
64 let cur_r = na::Rotation3::new(this_axis_na.fixed_rows(0));
65 let other_r = na::Rotation3::new(other_axis_na.fixed_rows(0));
66 let new_r = cur_r.slerp(&other_r, other_weight);
67 let axis_angle = new_r.scaled_axis();
68 let new_axis_angle_nd = array![axis_angle.x, axis_angle.y, axis_angle.z];
69 new_axis_angle_nd
70}
71pub fn interpolate_axis_angle_batch(this_axis: &nd::Array2<f32>, other_axis: &nd::Array2<f32>, other_weight: f32) -> nd::Array2<f32> {
74 let this_axis_na = this_axis.clone().into_nalgebra();
75 let other_axis_na = other_axis.clone().into_nalgebra();
76 let mut new_axis_angles = nd::Array2::<f32>::zeros(this_axis_na.shape());
77 for ((this_axis, other_axis), mut new_joint) in this_axis_na
78 .row_iter()
79 .zip(other_axis_na.row_iter())
80 .zip(new_axis_angles.axis_iter_mut(nd::Axis(0)))
81 {
82 let cur_r = na::Rotation3::new(this_axis.transpose().fixed_rows(0));
83 let other_r = na::Rotation3::new(other_axis.transpose().fixed_rows(0));
84 let new_r = cur_r.slerp(&other_r, other_weight);
85 let axis_angle = new_r.scaled_axis();
86 new_joint.assign(&array![axis_angle.x, axis_angle.y, axis_angle.z]);
87 }
88 new_axis_angles
89}
90#[allow(clippy::missing_panics_doc)]
91#[allow(clippy::similar_names)]
92#[allow(clippy::cast_sign_loss)]
93pub fn batch_rigid_transform(
94 parent_idx_per_joint: Vec<u32>,
95 rot_mats: &nd::Array3<f32>,
96 joints: &nd::Array2<f32>,
97 num_joints: usize,
98) -> (nd::Array2<f32>, nd::Array3<f32>) {
99 let mut rel_joints = joints.clone();
100 let parent_idx_data_u32 = parent_idx_per_joint;
101 let parent_idx_per_joint = nd::Array1::from_vec(parent_idx_data_u32);
102 for (idx_cur, idx_parent) in parent_idx_per_joint.iter().enumerate().skip(1) {
103 let parent_joint_position = joints.row(*idx_parent as usize);
104 rel_joints.row_mut(idx_cur).sub_assign(&parent_joint_position);
105 }
106 let mut transforms_mat = ndarray::Array3::<f32>::zeros((num_joints + 1, 4, 4));
107 for idx in 0..=num_joints {
108 let rot = rot_mats.slice(s![idx, .., ..]).to_owned();
109 let t = rel_joints.row(idx).to_owned();
110 transforms_mat.slice_mut(s![idx, 0..3, 0..3]).assign(&rot);
111 transforms_mat.slice_mut(s![idx, 0..3, 3]).assign(&t);
112 transforms_mat.slice_mut(s![idx, 3, 0..4]).assign(&array![0.0, 0.0, 0.0, 1.0]);
113 }
114 let mut transform_chain = Vec::new();
115 transform_chain.push(transforms_mat.slice(s![0, 0..4, 0..4]).to_owned().into_shape_with_order((4, 4)).unwrap());
116 for i in 1..=num_joints {
117 let mat_1 = &transform_chain[parent_idx_per_joint[[i]] as usize];
118 let mat_2 = transforms_mat.slice(s![i, 0..4, 0..4]);
119 let curr_res = mat_1.dot(&mat_2);
120 transform_chain.push(curr_res);
121 }
122 let mut posed_joints = joints.clone();
123 for (i, tf) in transform_chain.iter().enumerate() {
124 let t = tf.slice(s![0..3, 3]);
125 posed_joints.row_mut(i).assign(&t);
126 }
127 let mut rel_transforms = ndarray::Array3::<f32>::zeros((num_joints + 1, 4, 4));
128 for (i, transform) in transform_chain.iter().enumerate() {
129 let (jx, jy, jz) = (joints.row(i)[0], joints.row(i)[1], joints.row(i)[2]);
130 let joint_homogen = array![jx, jy, jz, 0.0];
131 let transformed_joint = transform.dot(&joint_homogen);
132 let mut transformed_joint_4 = nd::Array2::<f32>::zeros((4, 4));
133 transformed_joint_4.slice_mut(s![0..4, 3]).assign(&transformed_joint);
134 transformed_joint_4 = transform - transformed_joint_4;
135 rel_transforms.slice_mut(s![i, .., ..]).assign(&transformed_joint_4);
136 }
137 (posed_joints, rel_transforms)
138}