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

smpl_core/common/
animation.rs
