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