smpl_core/common/
pose_retarget.rs

1use super::pose::PoseG;
2use burn::{prelude::Backend, tensor::Tensor};
3/// After adding a pose or a animation, we align the avatar to the floor and
4/// store the tallness that it has; this is the most naive way of retargetting
5#[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}