smpl_core/conversions/
pose_chunked.rs

1use crate::common::{
2    metadata::SmplMetadata,
3    pose::Pose,
4    pose_parts::PosePart,
5    types::{SmplType, UpAxis},
6};
7use ndarray as nd;
8use ndarray::prelude::*;
9use std::ops::Range;
10/// Chunk ``Pose`` into various pose parts
11#[derive(Debug)]
12pub struct PoseChunked {
13    pub global_trans: nd::Array2<f32>,
14    pub global_orient: Option<nd::Array2<f32>>,
15    pub body_pose: Option<nd::Array2<f32>>,
16    pub left_hand_pose: Option<nd::Array2<f32>>,
17    pub right_hand_pose: Option<nd::Array2<f32>>,
18    pub jaw_pose: Option<nd::Array2<f32>>,
19    pub left_eye_pose: Option<nd::Array2<f32>>,
20    pub right_eye_pose: Option<nd::Array2<f32>>,
21    pub up_axis: UpAxis,
22    pub smpl_type: SmplType,
23}
24impl Default for PoseChunked {
25    fn default() -> Self {
26        let global_trans = ndarray::Array2::<f32>::zeros((1, 3));
27        Self {
28            global_trans,
29            global_orient: None,
30            body_pose: None,
31            left_hand_pose: None,
32            right_hand_pose: None,
33            jaw_pose: None,
34            left_eye_pose: None,
35            right_eye_pose: None,
36            up_axis: UpAxis::Y,
37            smpl_type: SmplType::SmplX,
38        }
39    }
40}
41impl PoseChunked {
42    #[allow(clippy::missing_panics_doc)]
43    pub fn new(pose: &Pose, metadata: &SmplMetadata) -> Self {
44        if pose.smpl_type == SmplType::SmplPP {
45            return Self {
46                global_trans: pose.global_trans.to_shape((1, 3)).unwrap().to_owned(),
47                global_orient: None,
48                body_pose: Some(pose.joint_poses.clone()),
49                left_hand_pose: None,
50                right_hand_pose: None,
51                jaw_pose: None,
52                left_eye_pose: None,
53                right_eye_pose: None,
54                up_axis: pose.up_axis,
55                smpl_type: pose.smpl_type,
56            };
57        }
58        let p2r = &metadata.parts2jointranges;
59        let joint_poses = &pose.joint_poses;
60        let max_range = 0..joint_poses.dim().0;
61        let clamp_closure = |lhs: &Range<usize>, rhs: &Range<usize>| -> Range<usize> {
62            if lhs.end > rhs.end {
63                0..0
64            } else {
65                lhs.clone()
66            }
67        };
68        let global_orient_clamped = clamp_closure(&p2r[PosePart::RootRotation], &max_range);
69        let mut global_orient = Some(joint_poses.slice(s![global_orient_clamped, ..]).to_owned());
70        let body_clamped = clamp_closure(&p2r[PosePart::Body], &max_range);
71        let mut body_pose = Some(joint_poses.slice(s![body_clamped, ..]).to_owned());
72        let left_hand_clamped = clamp_closure(&p2r[PosePart::LeftHand], &max_range);
73        let mut left_hand_pose = Some(joint_poses.slice(s![left_hand_clamped, ..]).to_owned());
74        let right_hand_clamped = clamp_closure(&p2r[PosePart::RightHand], &max_range);
75        let mut right_hand_pose = Some(joint_poses.slice(s![right_hand_clamped, ..]).to_owned());
76        let jaw_clamped = clamp_closure(&p2r[PosePart::Jaw], &max_range);
77        let mut jaw_pose = Some(joint_poses.slice(s![jaw_clamped, ..]).to_owned());
78        let left_eye_clamped = clamp_closure(&p2r[PosePart::LeftEye], &max_range);
79        let mut left_eye_pose = Some(joint_poses.slice(s![left_eye_clamped, ..]).to_owned());
80        let right_eye_clamped = clamp_closure(&p2r[PosePart::RightEye], &max_range);
81        let mut right_eye_pose = Some(joint_poses.slice(s![right_eye_clamped, ..]).to_owned());
82        if global_orient.as_ref().unwrap().dim().0 == 0 {
83            global_orient = None;
84        }
85        if body_pose.as_ref().unwrap().dim().0 == 0 {
86            body_pose = None;
87        }
88        if left_hand_pose.as_ref().unwrap().dim().0 == 0 {
89            left_hand_pose = None;
90        }
91        if right_hand_pose.as_ref().unwrap().dim().0 == 0 {
92            right_hand_pose = None;
93        }
94        if jaw_pose.as_ref().unwrap().dim().0 == 0 {
95            jaw_pose = None;
96        }
97        if left_eye_pose.as_ref().unwrap().dim().0 == 0 {
98            left_eye_pose = None;
99        }
100        if right_eye_pose.as_ref().unwrap().dim().0 == 0 {
101            right_eye_pose = None;
102        }
103        Self {
104            global_trans: pose.global_trans.to_shape((1, 3)).unwrap().to_owned(),
105            global_orient,
106            body_pose,
107            left_hand_pose,
108            right_hand_pose,
109            jaw_pose,
110            left_eye_pose,
111            right_eye_pose,
112            up_axis: pose.up_axis,
113            smpl_type: pose.smpl_type,
114        }
115    }
116    #[allow(clippy::missing_panics_doc)]
117    pub fn to_pose(&self, metadata: &SmplMetadata, smpl_type: SmplType) -> Pose {
118        if smpl_type == SmplType::SmplPP {
119            let zeros = nd::Array2::<f32>::zeros((46, 1));
120            let mut pose = Pose::new_empty(self.up_axis, smpl_type);
121            pose.joint_poses.assign(self.body_pose.as_ref().unwrap_or(&zeros));
122            pose.global_trans.assign(&self.global_trans.to_shape(3).unwrap().to_owned());
123            return pose;
124        }
125        let mut pose = Pose::new_empty(self.up_axis, smpl_type);
126        let zeros = nd::Array2::<f32>::zeros((1, 3));
127        pose.global_trans = self.global_trans.to_shape(3).unwrap().to_owned();
128        pose.joint_poses
129            .slice_mut(s![metadata.parts2jointranges[PosePart::RootRotation].clone(), ..])
130            .assign(self.global_orient.as_ref().unwrap_or(&zeros));
131        pose.joint_poses
132            .slice_mut(s![metadata.parts2jointranges[PosePart::Body].clone(), ..])
133            .assign(self.body_pose.as_ref().unwrap_or(&zeros));
134        pose.joint_poses
135            .slice_mut(s![metadata.parts2jointranges[PosePart::LeftHand].clone(), ..])
136            .assign(self.left_hand_pose.as_ref().unwrap_or(&zeros));
137        pose.joint_poses
138            .slice_mut(s![metadata.parts2jointranges[PosePart::RightHand].clone(), ..])
139            .assign(self.right_hand_pose.as_ref().unwrap_or(&zeros));
140        pose.joint_poses
141            .slice_mut(s![metadata.parts2jointranges[PosePart::Jaw].clone(), ..])
142            .assign(self.jaw_pose.as_ref().unwrap_or(&zeros));
143        pose.joint_poses
144            .slice_mut(s![metadata.parts2jointranges[PosePart::LeftEye].clone(), ..])
145            .assign(self.left_eye_pose.as_ref().unwrap_or(&zeros));
146        pose.joint_poses
147            .slice_mut(s![metadata.parts2jointranges[PosePart::RightEye].clone(), ..])
148            .assign(self.right_eye_pose.as_ref().unwrap_or(&zeros));
149        pose
150    }
151}
smpl_core/conversions/pose_chunked.rs

smpl_core/conversions/
pose_chunked.rs
