smpl_core/common/
pose_retarget.rs1use super::pose::PoseG;
2use burn::{prelude::Backend, tensor::Tensor};
3#[derive(Clone)]
6pub struct RetargetPoseYShift {
7 pub y_shift: f32,
8 pub dist_chest_to_feet: f32,
9 pub currently_applied_y_shift: f32,
10}
11impl RetargetPoseYShift {
12 pub fn new(y_shift: f32, dist_chest_to_feet: f32) -> Self {
13 Self {
14 y_shift,
15 dist_chest_to_feet,
16 currently_applied_y_shift: 0.0,
17 }
18 }
19 pub fn update(&mut self, new_height: f32) {
20 let old_height = self.dist_chest_to_feet;
21 let diff = new_height - old_height;
22 self.dist_chest_to_feet = new_height;
23 self.y_shift += -diff;
24 }
25 #[allow(clippy::missing_panics_doc)]
26 pub fn apply<B: Backend>(&mut self, pose: &mut PoseG<B>) {
27 if pose.non_retargeted_pose.is_none() {
28 pose.non_retargeted_pose = Some(Box::new(pose.clone()));
29 }
30 let original_trans = &pose.non_retargeted_pose.as_ref().unwrap().global_trans;
31 let device = original_trans.device();
32 let shift_tensor = Tensor::from_floats([0.0, self.y_shift, 0.0], &device);
33 pose.global_trans = original_trans.clone() + shift_tensor;
34 self.currently_applied_y_shift = self.y_shift;
35 pose.retargeted = true;
36 }
37 #[allow(clippy::missing_panics_doc)]
38 pub fn remove_retarget<B: Backend>(&self, pose: &mut PoseG<B>) {
39 if pose.non_retargeted_pose.is_some() {
40 let original = *pose.non_retargeted_pose.take().unwrap();
41 pose.clone_from(&original);
42 }
43 pose.retargeted = false;
44 }
45}