smpl_core/common/
animation.rs

1use super::{
2    expression::Expression,
3    metadata::smpl_metadata,
4    pose::Pose,
5    types::{AngleType, SmplType, UpAxis},
6};
7use crate::codec::codec::SmplCodec;
8use core::time::Duration;
9use gloss_utils::nshare::{RefNdarray1, ToNalgebra};
10use log::debug;
11use log::warn;
12use nalgebra as na;
13use nd::concatenate;
14use ndarray as nd;
15use ndarray_npy::NpzReader;
16use serde_json::Value;
17use smpl_utils::{
18    io::FileLoader,
19    numerical::{euler2angleaxis, map},
20};
21use std::io::{Read, Seek};
22/// Animation Wrap mode
23#[derive(PartialEq, PartialOrd, Clone, Default)]
24pub enum AnimWrap {
25    Clamp,
26    #[default]
27    Loop,
28    Reverse,
29}
30/// Animation config
31#[derive(Clone)]
32pub struct AnimationConfig {
33    pub fps: f32,
34    pub wrap_behaviour: AnimWrap,
35    pub angle_type: AngleType,
36    pub up_axis: UpAxis,
37    pub smpl_type: SmplType,
38}
39impl Default for AnimationConfig {
40    fn default() -> Self {
41        Self {
42            fps: 60.0,
43            wrap_behaviour: AnimWrap::Loop,
44            angle_type: AngleType::AxisAngle,
45            up_axis: UpAxis::Y,
46            smpl_type: SmplType::SmplX,
47        }
48    }
49}
50/// The runner for animations
51#[derive(Clone)]
52#[allow(clippy::struct_excessive_bools)]
53pub struct AnimationRunner {
54    pub anim_current_time: Duration,
55    pub anim_reversed: bool,
56    pub nr_repetitions: u32,
57    pub paused: bool,
58    pub temporary_pause: bool,
59}
60impl Default for AnimationRunner {
61    fn default() -> Self {
62        Self {
63            anim_current_time: Duration::ZERO,
64            anim_reversed: false,
65            nr_repetitions: 0,
66            paused: false,
67            temporary_pause: false,
68        }
69    }
70}
71/// Animation struct for all data regarding a certain animation
72#[derive(Clone)]
73pub struct Animation {
74    pub per_frame_joint_poses: nd::Array3<f32>,
75    pub per_frame_root_trans: nd::Array2<f32>,
76    pub per_frame_expression_coeffs: Option<nd::Array2<f32>>,
77    pub start_offset: usize,
78    pub runner: AnimationRunner,
79    pub config: AnimationConfig,
80}
81impl Animation {
82    /// # Panics
83    /// Will panic if the translation and rotation do not cover the same number
84    /// of timesteps
85    #[allow(clippy::cast_possible_truncation)]
86    pub fn new_from_matrices(
87        per_frame_joint_poses: nd::Array3<f32>,
88        per_frame_global_trans: nd::Array2<f32>,
89        per_frame_expression_coeffs: Option<nd::Array2<f32>>,
90        config: AnimationConfig,
91    ) -> Self {
92        assert!(
93            per_frame_joint_poses.dim().0 == per_frame_global_trans.dim().0,
94            "The translation and rotation should cover the same number of timesteps"
95        );
96        let mut per_frame_joint_poses = per_frame_joint_poses;
97        let per_frame_global_trans = per_frame_global_trans;
98        if config.smpl_type == SmplType::SmplPP && config.angle_type == AngleType::Euler {
99            warn!("Angle type Euler is not allowed with SMPL++");
100        }
101        if config.smpl_type != SmplType::SmplPP && config.angle_type == AngleType::Euler {
102            let animation_frames = per_frame_joint_poses.dim().0;
103            let num_active_joints = per_frame_joint_poses.dim().1;
104            let mut new_per_frame_joint_poses: nd::Array3<f32> = nd::Array3::<f32>::zeros((animation_frames, num_active_joints, 3));
105            for (idx_timestep, poses_for_timestep) in per_frame_joint_poses.axis_iter(nd::Axis(0)).enumerate() {
106                for (idx_joint, joint_pose) in poses_for_timestep.axis_iter(nd::Axis(0)).enumerate() {
107                    let angle_axis = euler2angleaxis(joint_pose[0], joint_pose[1], joint_pose[2]);
108                    new_per_frame_joint_poses[(idx_timestep, idx_joint, 0)] = angle_axis.x;
109                    new_per_frame_joint_poses[(idx_timestep, idx_joint, 1)] = angle_axis.y;
110                    new_per_frame_joint_poses[(idx_timestep, idx_joint, 2)] = angle_axis.z;
111                }
112            }
113            per_frame_joint_poses = new_per_frame_joint_poses;
114        }
115        Self {
116            per_frame_joint_poses,
117            per_frame_root_trans: per_frame_global_trans,
118            per_frame_expression_coeffs,
119            start_offset: 0,
120            runner: AnimationRunner::default(),
121            config,
122        }
123    }
124    #[allow(clippy::cast_possible_truncation)]
125    fn new_from_npz_reader<R: Read + Seek>(npz: &mut NpzReader<R>, config: AnimationConfig) -> Self {
126        debug!("npz names is {:?}", npz.names().unwrap());
127        let per_frame_joint_poses: nd::Array2<f64> = npz.by_name("poses").unwrap();
128        let animation_frames = per_frame_joint_poses.nrows();
129        let num_joints_3 = per_frame_joint_poses.ncols();
130        let per_frame_joint_poses = per_frame_joint_poses.mapv(|x| x as f32);
131        let per_frame_joint_poses = per_frame_joint_poses
132            .into_shape_with_order((animation_frames, num_joints_3 / 3, 3))
133            .unwrap();
134        let per_frame_global_trans: nd::Array2<f64> = npz.by_name("trans").unwrap();
135        let per_frame_global_trans = per_frame_global_trans.mapv(|x| x as f32);
136        let per_frame_expression_coeffs: Option<nd::Array2<f64>> = npz.by_name("expressionParameters").ok();
137        let per_frame_expression_coeffs = per_frame_expression_coeffs.map(|x| x.mapv(|x| x as f32));
138        Self::new_from_matrices(per_frame_joint_poses, per_frame_global_trans, per_frame_expression_coeffs, config)
139    }
140    /// # Panics
141    /// Will panic if the path cannot be opened
142    /// Will panic if the translation and rotation do not cover the same number
143    /// of timesteps
144    #[cfg(not(target_arch = "wasm32"))]
145    #[allow(clippy::cast_possible_truncation)]
146    pub fn new_from_npz(anim_npz_path: &str, config: AnimationConfig) -> Self {
147        let mut npz =
148            NpzReader::new(std::fs::File::open(anim_npz_path).unwrap_or_else(|_| panic!("Could not find/open file: {anim_npz_path}"))).unwrap();
149        Self::new_from_npz_reader(&mut npz, config)
150    }
151    /// # Panics
152    /// Will panic if the path cannot be opened
153    /// Will panic if the translation and rotation do not cover the same number
154    /// of timesteps
155    #[allow(clippy::cast_possible_truncation)]
156    pub async fn new_from_npz_async(anim_npz_path: &str, config: AnimationConfig) -> Self {
157        let reader = FileLoader::open(anim_npz_path).await;
158        let mut npz = NpzReader::new(reader).unwrap();
159        Self::new_from_npz_reader(&mut npz, config)
160    }
161    /// # Panics
162    /// Will panic if the path cannot be opened
163    /// Will panic if the translation and rotation do not cover the same number
164    /// of timesteps
165    #[allow(clippy::cast_possible_truncation)]
166    #[allow(clippy::identity_op)]
167    pub fn new_from_json(path: &str, _anim_fps: f32, config: AnimationConfig) -> Self {
168        let file = std::fs::File::open(path).unwrap();
169        let reader = std::io::BufReader::new(file);
170        let v: Value = serde_json::from_reader(reader).unwrap();
171        let poses = &v["poses"];
172        let animation_frames = poses.as_array().unwrap().len();
173        let num_active_joints = poses[0].as_array().unwrap().len() / 3;
174        let per_frame_global_trans: nd::Array2<f32> = nd::Array2::<f32>::zeros((animation_frames, 3));
175        let poses_json_vec: Vec<f32> = poses.as_array().unwrap().iter().map(|x| x.as_f64().unwrap() as f32).collect();
176        let per_frame_joint_poses = nd::Array3::from_shape_vec((animation_frames, num_active_joints, 3), poses_json_vec).unwrap();
177        Self::new_from_matrices(per_frame_joint_poses, per_frame_global_trans, None, config)
178    }
179    /// Create an ``Animation`` component from a ``SmplCodec``
180    /// # Panics
181    /// Will panic if the individual body part poses in the codec don't have the
182    /// correct shape to be concatenated together into a full pose for the whole
183    /// body
184    #[allow(clippy::cast_sign_loss)]
185    pub fn new_from_smpl_codec(codec: &SmplCodec, wrap_behaviour: AnimWrap) -> Option<Self> {
186        let nr_frames = codec.frame_count as usize;
187        let metadata = smpl_metadata(&codec.smpl_type());
188        let body_translation = codec
189            .body_translation
190            .as_ref()
191            .unwrap_or(&ndarray::Array2::<f32>::zeros((nr_frames, 3)))
192            .clone();
193        let fps = codec.frame_rate?;
194        if codec.smpl_type() == SmplType::SmplPP {
195            let body_pose = codec
196                .body_pose
197                .as_ref()
198                .unwrap_or(&ndarray::Array3::<f32>::zeros((nr_frames, metadata.pose_dim, 1)))
199                .clone();
200            let config = AnimationConfig {
201                smpl_type: SmplType::SmplPP,
202                wrap_behaviour,
203                fps,
204                ..Default::default()
205            };
206            Some(Self::new_from_matrices(body_pose, body_translation, None, config))
207        } else {
208            let body_pose = codec
209                .body_pose
210                .as_ref()
211                .unwrap_or(&ndarray::Array3::<f32>::zeros((nr_frames, 1 + metadata.num_body_joints, 3)))
212                .clone();
213            let head_pose = codec
214                .head_pose
215                .as_ref()
216                .unwrap_or(&ndarray::Array3::<f32>::zeros((nr_frames, metadata.num_face_joints, 3)))
217                .clone();
218            let left_hand_pose = codec
219                .left_hand_pose
220                .as_ref()
221                .unwrap_or(&ndarray::Array3::<f32>::zeros((nr_frames, metadata.num_hand_joints, 3)))
222                .clone();
223            let right_hand_pose = codec
224                .right_hand_pose
225                .as_ref()
226                .unwrap_or(&ndarray::Array3::<f32>::zeros((nr_frames, metadata.num_hand_joints, 3)))
227                .clone();
228            let per_frame_joint_poses = concatenate(
229                nd::Axis(1),
230                &[body_pose.view(), head_pose.view(), left_hand_pose.view(), right_hand_pose.view()],
231            )
232            .unwrap();
233            let per_frame_expression_coeffs = codec.expression_parameters.clone();
234            let config = AnimationConfig {
235                smpl_type: codec.smpl_type(),
236                wrap_behaviour,
237                fps,
238                ..Default::default()
239            };
240            Some(Self::new_from_matrices(
241                per_frame_joint_poses,
242                body_translation,
243                per_frame_expression_coeffs,
244                config,
245            ))
246        }
247    }
248    /// Create an ``Animation`` component from a ``.smpl`` file
249    #[cfg(not(target_arch = "wasm32"))]
250    #[allow(clippy::cast_possible_truncation)]
251    pub fn new_from_smpl_file(path: &str, wrap_behaviour: AnimWrap) -> Option<Self> {
252        let codec = SmplCodec::from_file(path);
253        Self::new_from_smpl_codec(&codec, wrap_behaviour)
254    }
255    /// Create an ``Animation`` component from a ``.smpl`` buffer
256    #[allow(clippy::cast_possible_truncation)]
257    pub fn new_from_smpl_buf(buf: &[u8], wrap_behaviour: AnimWrap) -> Option<Self> {
258        let codec = SmplCodec::from_buf(buf);
259        Self::new_from_smpl_codec(&codec, wrap_behaviour)
260    }
261    pub fn num_active_joints(&self) -> usize {
262        self.per_frame_joint_poses.dim().1
263    }
264    pub fn num_animation_frames(&self) -> usize {
265        self.per_frame_joint_poses.dim().0
266    }
267    /// Advances the animation by the amount of time elapsed since last time we
268    /// got the current pose
269    pub fn advance(&mut self, dt_raw: Duration, first_time: bool) {
270        let duration = self.duration();
271        let runner = &mut self.runner;
272        let config = &self.config;
273        let mut dt = dt_raw;
274        if first_time {
275            dt = Duration::ZERO;
276        }
277        let will_overflow = runner.anim_current_time + dt > duration;
278        let will_underflow = runner.anim_current_time < dt && runner.anim_reversed;
279        if will_overflow || will_underflow {
280            if will_overflow {
281                match config.wrap_behaviour {
282                    AnimWrap::Clamp => {
283                        dt = Duration::ZERO;
284                        runner.anim_current_time = duration;
285                    }
286                    AnimWrap::Loop => {
287                        dt = Duration::from_secs_f64(dt.as_secs_f64() % duration.as_secs_f64());
288                        runner.anim_current_time = Duration::ZERO;
289                        runner.nr_repetitions += 1;
290                    }
291                    AnimWrap::Reverse => {
292                        dt = Duration::from_secs_f64(dt.as_secs_f64() % duration.as_secs_f64());
293                        runner.anim_current_time = duration;
294                        runner.anim_reversed = !runner.anim_reversed;
295                        runner.nr_repetitions += 1;
296                    }
297                }
298            } else {
299                match config.wrap_behaviour {
300                    AnimWrap::Clamp => {
301                        dt = Duration::ZERO;
302                        runner.anim_current_time = Duration::ZERO;
303                    }
304                    AnimWrap::Loop => {
305                        dt = Duration::from_secs_f64(dt.as_secs_f64() % duration.as_secs_f64());
306                        runner.anim_current_time = duration;
307                        runner.nr_repetitions += 1;
308                    }
309                    AnimWrap::Reverse => {
310                        dt = Duration::from_secs_f64(dt.as_secs_f64() % duration.as_secs_f64());
311                        runner.anim_reversed = !runner.anim_reversed;
312                        runner.nr_repetitions += 1;
313                    }
314                }
315            }
316        }
317        if runner.anim_reversed {
318            runner.anim_current_time = runner.anim_current_time.saturating_sub(dt);
319        } else {
320            runner.anim_current_time = runner.anim_current_time.saturating_add(dt);
321        }
322    }
323    pub fn has_expression(&self) -> bool {
324        self.per_frame_expression_coeffs.is_some()
325    }
326    #[allow(clippy::cast_precision_loss)]
327    #[allow(clippy::cast_possible_truncation)]
328    #[allow(clippy::cast_sign_loss)]
329    pub fn get_smooth_time_indices(&self) -> (usize, usize, f32) {
330        let frame_time = map(
331            self.runner.anim_current_time.as_secs_f32(),
332            0.0,
333            self.duration().as_secs_f32(),
334            0.0,
335            (self.num_animation_frames() - 1) as f32,
336        );
337        let frame_ceil = frame_time.ceil();
338        let frame_ceil = frame_ceil.clamp(0.0, (self.num_animation_frames() - 1) as f32);
339        let frame_floor = frame_time.floor();
340        let frame_floor = frame_floor.clamp(0.0, (self.num_animation_frames() - 1) as f32);
341        let w_ceil = frame_ceil - frame_time;
342        let w_ceil = 1.0 - w_ceil;
343        (frame_floor as usize, frame_ceil as usize, w_ceil)
344    }
345    /// Get the pose and translation at the current time, interpolates if
346    /// necessary
347    #[allow(clippy::cast_precision_loss)]
348    #[allow(clippy::cast_possible_truncation)]
349    #[allow(clippy::cast_sign_loss)]
350    pub fn get_current_pose(&mut self) -> Pose {
351        let (frame_floor, frame_ceil, w_ceil) = self.get_smooth_time_indices();
352        let anim_frame_ceil = self.get_pose_at_idx(frame_ceil);
353        let anim_frame_floor = self.get_pose_at_idx(frame_floor);
354        anim_frame_floor.interpolate(&anim_frame_ceil, w_ceil)
355    }
356    /// Get pose at a certain frame ID
357    #[allow(clippy::cast_precision_loss)]
358    #[allow(clippy::cast_possible_truncation)]
359    #[allow(clippy::cast_sign_loss)]
360    pub fn get_pose_at_idx(&self, idx: usize) -> Pose {
361        let joint_poses = self.per_frame_joint_poses.index_axis(nd::Axis(0), idx).to_owned();
362        let global_trans = self.per_frame_root_trans.index_axis(nd::Axis(0), idx).to_owned();
363        Pose::new(joint_poses, global_trans, self.config.up_axis, self.config.smpl_type)
364    }
365    /// Get expression at current time
366    pub fn get_current_expression(&mut self) -> Option<Expression> {
367        let (frame_floor, frame_ceil, w_ceil) = self.get_smooth_time_indices();
368        let expression_ceil = self.get_expression_at_idx(frame_ceil);
369        let expresion_floor = self.get_expression_at_idx(frame_floor);
370        expresion_floor.map(|expresion_floor| expresion_floor.interpolate(&expression_ceil.unwrap(), w_ceil))
371    }
372    /// Get expression at a given frame ID
373    #[allow(clippy::cast_precision_loss)]
374    #[allow(clippy::cast_possible_truncation)]
375    #[allow(clippy::cast_sign_loss)]
376    pub fn get_expression_at_idx(&self, idx: usize) -> Option<Expression> {
377        if let Some(ref per_frame_expression_coeffs) = self.per_frame_expression_coeffs {
378            let expr_coeffs = per_frame_expression_coeffs.index_axis(nd::Axis(0), idx).to_owned();
379            Some(Expression::new(expr_coeffs))
380        } else {
381            None
382        }
383    }
384    #[must_use]
385    pub fn slice_time_range(&self, start_sec: f32, end_sec: f32) -> Animation {
386        let mut cur_anim = self.clone();
387        cur_anim.set_cur_time_as_sec(start_sec);
388        let (start_idx, _, _) = cur_anim.get_smooth_time_indices();
389        cur_anim.set_cur_time_as_sec(end_sec);
390        let (_, end_idx, _) = cur_anim.get_smooth_time_indices();
391        let nr_frames = end_idx - start_idx + 1;
392        let nr_joints = cur_anim.per_frame_joint_poses.shape()[1];
393        let mut new_per_frame_joint_poses = nd::Array3::<f32>::zeros((nr_frames, nr_joints, 3));
394        let mut new_per_frame_root_trans = nd::Array2::<f32>::zeros((nr_frames, 3));
395        for (idx_insert_to, idx_extract_from) in (start_idx..=end_idx).enumerate() {
396            let joint_poses = cur_anim.per_frame_joint_poses.index_axis(nd::Axis(0), idx_extract_from);
397            let trans = cur_anim.per_frame_root_trans.index_axis(nd::Axis(0), idx_extract_from);
398            new_per_frame_joint_poses.index_axis_mut(nd::Axis(0), idx_insert_to).assign(&joint_poses);
399            new_per_frame_root_trans.index_axis_mut(nd::Axis(0), idx_insert_to).assign(&trans);
400        }
401        let _new_per_frame_expression_coeffs = if let Some(ref per_frame_expression_coeffs) = cur_anim.per_frame_expression_coeffs {
402            let nr_expr_coeffs = per_frame_expression_coeffs.shape()[1];
403            let mut new_per_frame_expression_coeffs = nd::Array2::<f32>::zeros((nr_frames, nr_expr_coeffs));
404            for (idx_insert_to, idx_extract_from) in (start_idx..=end_idx).enumerate() {
405                let expr = per_frame_expression_coeffs.index_axis(nd::Axis(0), idx_extract_from);
406                new_per_frame_expression_coeffs.index_axis_mut(nd::Axis(0), idx_insert_to).assign(&expr);
407            }
408            Some(new_per_frame_expression_coeffs)
409        } else {
410            None
411        };
412        let new_per_frame_expression_coeffs = if let Some(ref per_frame_expression_coeffs) = cur_anim.per_frame_expression_coeffs {
413            let nr_expr_coeffs = per_frame_expression_coeffs.shape()[1];
414            let mut new_per_frame_expression_coeffs = nd::Array2::<f32>::zeros((nr_frames, nr_expr_coeffs));
415            for (idx_insert_to, idx_extract_from) in (start_idx..=end_idx).enumerate() {
416                let expr = per_frame_expression_coeffs.index_axis(nd::Axis(0), idx_extract_from);
417                new_per_frame_expression_coeffs.index_axis_mut(nd::Axis(0), idx_insert_to).assign(&expr);
418            }
419            Some(new_per_frame_expression_coeffs)
420        } else {
421            None
422        };
423        Animation::new_from_matrices(
424            new_per_frame_joint_poses,
425            new_per_frame_root_trans,
426            new_per_frame_expression_coeffs,
427            cur_anim.config.clone(),
428        )
429    }
430    /// Rotates multiple of 90 until the axis of the body is aligned with some
431    /// arbitrary vector
432    pub fn align_y_axis_quadrant(&mut self, current_axis: &nd::Array1<f32>, desired_axis: &nd::Array1<f32>) {
433        let mut cur_axis_xz = na::Vector2::new(current_axis[0], -current_axis[2]).normalize();
434        let desired_axis_xz = na::Vector2::new(desired_axis[0], -desired_axis[2]).normalize();
435        let mut best_dot = f32::MIN;
436        let mut best_angle: f32 = 0.0;
437        let mut cur_angle = 0.0;
438        let rot_90 = na::Rotation2::new(std::f32::consts::FRAC_PI_2);
439        for _iters in 0..4 {
440            cur_axis_xz = rot_90 * cur_axis_xz;
441            cur_angle += 90.0;
442            let dot = cur_axis_xz.dot(&desired_axis_xz);
443            if dot > best_dot {
444                best_dot = dot;
445                best_angle = cur_angle;
446            }
447        }
448        let alignment_rot = na::Rotation3::from_euler_angles(0.0, best_angle.to_radians(), 0.0);
449        ndarray::Zip::from(self.per_frame_joint_poses.outer_iter_mut())
450            .and(self.per_frame_root_trans.outer_iter_mut())
451            .for_each(|mut poses_for_timestep, mut trans_for_timestep| {
452                let pelvis_axis_angle = poses_for_timestep.row(0).into_nalgebra();
453                let pelvis_axis_angle = pelvis_axis_angle.fixed_rows::<3>(0);
454                let pelvis_rot = na::Rotation3::from_scaled_axis(pelvis_axis_angle);
455                let new_pelvis_rot = alignment_rot * pelvis_rot;
456                poses_for_timestep.row_mut(0).assign(&new_pelvis_rot.scaled_axis().ref_ndarray1());
457                let trans_for_timestep_na = trans_for_timestep.to_owned().into_nalgebra();
458                let new_trans_for_timestep_na = alignment_rot * trans_for_timestep_na;
459                trans_for_timestep.assign(&new_trans_for_timestep_na.ref_ndarray1());
460            });
461    }
462    pub fn get_cur_time(&self) -> Duration {
463        self.runner.anim_current_time
464    }
465    pub fn set_cur_time_as_sec(&mut self, time_sec: f32) {
466        self.runner.anim_current_time = Duration::from_secs_f32(time_sec);
467    }
468    pub fn pause(&mut self) {
469        self.runner.paused = true;
470    }
471    pub fn play(&mut self) {
472        self.runner.paused = false;
473    }
474    /// Duration of the animation
475    #[allow(clippy::cast_precision_loss)]
476    pub fn duration(&self) -> Duration {
477        Duration::from_secs_f32(self.num_animation_frames() as f32 / self.config.fps)
478    }
479    pub fn is_finished(&self) -> bool {
480        self.config.wrap_behaviour == AnimWrap::Clamp && self.runner.anim_current_time >= self.duration()
481    }
482    pub fn nr_repetitions(&self) -> u32 {
483        self.runner.nr_repetitions
484    }
485    /// Shift each frame of the animation by the given translation vector.
486    ///
487    /// # Errors
488    ///
489    /// Will return `Err` for array size mismatch.
490    pub fn translate(&mut self, translation: &nd::Array1<f32>) -> Result<(), String> {
491        if translation.len() != self.per_frame_root_trans.ncols() {
492            return Err("Translation vector should be length-3 array".to_owned());
493        }
494        for mut row in self.per_frame_root_trans.rows_mut() {
495            row += translation;
496        }
497        Ok(())
498    }
499}
smpl_core/common/animation.rs

smpl_core/common/
animation.rs
