smpl_core/common/
transform_sequence.rs

1use gloss_renderer::components::ModelMatrix;
2use nalgebra as na;
3use ndarray as nd;
4use ndarray::s;
5use ndarray_npy::NpzReader;
6use std::time::Duration;
7fn map(value: f32, from_min: f32, from_max: f32, to_min: f32, to_max: f32) -> f32 {
8    to_min + (value - from_min) * (to_max - to_min) / (from_max - from_min)
9}
10/// The ``TransformSequence`` contains the rigid pose sequence of a model over time.
11/// It can contain per-frame translations, rotations (in angle-axis format) and scales.
12#[derive(Debug, Clone)]
13pub struct TransformSequence {
14    pub translations: nd::Array2<f32>,
15    pub rotations: nd::Array2<f32>,
16    pub scales: nd::Array1<f32>,
17}
18impl TransformSequence {
19    /// Create a new `TransformSequence` from given npz file path.
20    /// # Panics
21    /// Will panic if the number of frames in translations and rotations do not match.
22    pub fn new_from_npz(path: &str) -> Self {
23        let mut npz = NpzReader::new(std::fs::File::open(path).unwrap()).unwrap();
24        let translations: nd::Array2<f32> = npz.by_name("translation.npy").unwrap();
25        let rotations: nd::Array2<f32> = npz.by_name("rotation.npy").unwrap();
26        let scales: nd::Array1<f32> = npz.by_name("objectSize.npy").unwrap();
27        assert_eq!(
28            translations.shape()[0],
29            rotations.shape()[0],
30            "Number of frames in translations and rotations must match"
31        );
32        Self {
33            translations,
34            rotations,
35            scales,
36        }
37    }
38    /// Create an empty `TransformSequence` with given number of frames.
39    pub fn new_empty(num_frames: usize) -> Self {
40        let translations = nd::Array2::<f32>::zeros((num_frames, 3));
41        let rotations = nd::Array2::<f32>::zeros((num_frames, 3));
42        let scales = nd::Array1::<f32>::zeros(num_frames);
43        Self {
44            translations,
45            rotations,
46            scales,
47        }
48    }
49    /// Create a `TransformSequence` from a translations and quaternion rotations
50    pub fn new_from_quat_rot_trans(rot: &nd::Array2<f32>, trans: &nd::Array2<f32>) -> Self {
51        assert_eq!(
52            rot.shape()[1],
53            4,
54            "Rotations must be in quaternion format (nr_frames x 4); found rotations of shape {:?}",
55            rot.shape()
56        );
57        let num_frames = rot.shape()[0];
58        let mut rotations = nd::Array2::<f32>::zeros((num_frames, 3));
59        for i in 0..num_frames {
60            let quat_slice = rot.slice(s![i, ..]);
61            let quat = na::Quaternion::new(quat_slice[3], quat_slice[0], quat_slice[1], quat_slice[2]);
62            let unit_quat = na::UnitQuaternion::from_quaternion(quat);
63            let axis_angle = unit_quat.scaled_axis();
64            rotations[[i, 0]] = axis_angle.x;
65            rotations[[i, 1]] = axis_angle.y;
66            rotations[[i, 2]] = axis_angle.z;
67        }
68        let scales = nd::Array1::<f32>::ones(trans.shape()[0]);
69        Self {
70            translations: trans.clone(),
71            rotations,
72            scales,
73        }
74    }
75    /// Create a `TransformSequence` from a translations and axis-angle rotations
76    pub fn new_from_axisangle_rot_trans(rot: &nd::Array2<f32>, trans: &nd::Array2<f32>) -> Self {
77        assert_eq!(
78            rot.shape()[1],
79            3,
80            "Rotations must be in axis-angle format (nr_frames x 3); found rotations of shape {:?}",
81            rot.shape()
82        );
83        let scales = nd::Array1::<f32>::ones(trans.shape()[0]);
84        Self {
85            translations: trans.clone(),
86            rotations: rot.clone(),
87            scales,
88        }
89    }
90    /// Duration of the animation
91    #[allow(clippy::cast_precision_loss)]
92    pub fn duration(&self, fps: f32) -> Duration {
93        Duration::from_secs_f32(self.num_frames() as f32 / fps)
94    }
95    pub fn get_rotations_as_quaternions(&self) -> nd::Array2<f32> {
96        let num_frames = self.rotations.shape()[0];
97        let mut quaternions = nd::Array2::<f32>::zeros((num_frames, 4));
98        for i in 0..num_frames {
99            let rotation = self.rotations.slice(s![i, ..]);
100            let aa_rotation = na::Vector3::from_row_slice(rotation.as_slice().unwrap());
101            let quaternion = na::UnitQuaternion::from_scaled_axis(aa_rotation);
102            quaternions[[i, 0]] = quaternion.i;
103            quaternions[[i, 1]] = quaternion.j;
104            quaternions[[i, 2]] = quaternion.k;
105            quaternions[[i, 3]] = quaternion.w;
106        }
107        quaternions
108    }
109    #[allow(clippy::cast_precision_loss)]
110    #[allow(clippy::cast_possible_truncation)]
111    #[allow(clippy::cast_sign_loss)]
112    pub fn get_smooth_time_indices(&self, time_sec: f32, fps: f32) -> (usize, usize, f32) {
113        let frame_time = map(time_sec, 0.0, self.duration(fps).as_secs_f32(), 0.0, (self.num_frames() - 1) as f32);
114        let frame_ceil = frame_time.ceil();
115        let frame_ceil = frame_ceil.clamp(0.0, (self.num_frames() - 1) as f32);
116        let frame_floor = frame_time.floor();
117        let frame_floor = frame_floor.clamp(0.0, (self.num_frames() - 1) as f32);
118        let w_ceil = frame_ceil - frame_time;
119        let w_ceil = 1.0 - w_ceil;
120        (frame_floor as usize, frame_ceil as usize, w_ceil)
121    }
122    pub fn get_transform_at_idx(&self, idx: u32) -> ModelMatrix {
123        let translation = self.translations.slice(s![idx as usize, ..]);
124        let rotation = self.rotations.slice(s![idx as usize, ..]);
125        let scale = self.scales[idx as usize];
126        let translation_vec = na::Vector3::from_row_slice(translation.as_slice().unwrap());
127        let rotation_vec = na::Vector3::from_row_slice(rotation.as_slice().unwrap());
128        let rotation_matrix = if rotation_vec.norm() > 0.0 {
129            na::Rotation3::from_axis_angle(&na::UnitVector3::new_normalize(rotation_vec), rotation_vec.norm())
130        } else {
131            na::Rotation3::identity()
132        };
133        let similarity = na::SimilarityMatrix3::from_parts(na::Translation3::from(translation_vec), rotation_matrix, scale);
134        ModelMatrix(similarity)
135    }
136    #[allow(clippy::cast_possible_truncation)]
137    pub fn get_transform_at_time(&self, time_sec: f32, fps: f32) -> ModelMatrix {
138        let (frame_floor, frame_ceil, w_ceil) = self.get_smooth_time_indices(time_sec, fps);
139        let mm_floor = self.get_transform_at_idx(frame_floor as u32);
140        let mm_ceil = self.get_transform_at_idx(frame_ceil as u32);
141        mm_floor.interpolate(&mm_ceil, w_ceil)
142    }
143    /// Get the number of frames in the `TransformSequence`.
144    pub fn num_frames(&self) -> usize {
145        self.translations.shape()[0]
146    }
147    /// Returns a new `TransformSequence` cropped from the start frame to the end frame.
148    #[must_use]
149    pub fn crop(&self, start_frame: usize, end_frame: usize) -> Self {
150        assert!(start_frame < end_frame, "Start frame must be less than end frame");
151        assert!(start_frame < self.num_frames(), "Start frame must be less than number of frames");
152        assert!(end_frame <= self.num_frames(), "End frame must be less than or equal to number of frames");
153        Self {
154            translations: self.translations.slice(s![start_frame..end_frame, ..]).to_owned(),
155            rotations: self.rotations.slice(s![start_frame..end_frame, ..]).to_owned(),
156            scales: self.scales.slice(s![start_frame..end_frame]).to_owned(),
157        }
158    }
159}
smpl_core/common/transform_sequence.rs

smpl_core/common/
transform_sequence.rs
