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