Skip to main content

oxihuman_export/
pose_export.rs

1// Copyright (C) 2026 COOLJAPAN OU (Team KitaSan)
2// SPDX-License-Identifier: Apache-2.0
3
4//! Pose and animation-frame export with keyframe support.
5//!
6//! Provides structs and functions for building, sampling, and serializing
7//! animation clips composed of pose frames.
8
9// ── Types ──────────────────────────────────────────────────────────────────
10
11/// A single pose frame in a clip.
12#[allow(dead_code)]
13#[derive(Debug, Clone)]
14pub struct ExportPoseFrame {
15    /// Time of this keyframe in seconds.
16    pub time: f32,
17    /// Per-bone transforms: `(position [x,y,z], rotation [x,y,z,w])`.
18    pub bone_transforms: Vec<([f32; 3], [f32; 4])>,
19}
20
21/// An ordered sequence of pose frames forming an animation clip.
22#[allow(dead_code)]
23#[derive(Debug, Clone)]
24pub struct ExportPoseClip {
25    /// Clip identifier.
26    pub name: String,
27    /// Ordered keyframes.
28    pub frames: Vec<ExportPoseFrame>,
29    /// Playback rate in frames per second.
30    pub fps: f32,
31    /// Whether the clip loops.
32    pub looping: bool,
33}
34
35/// Configuration for pose/animation export.
36#[allow(dead_code)]
37#[derive(Debug, Clone)]
38pub struct PoseExportConfig {
39    /// Target FPS for export baking.
40    pub fps: f32,
41    /// Include rotation channels.
42    pub include_rotation: bool,
43    /// Include position channels.
44    pub include_position: bool,
45    /// Floating-point decimal precision.
46    pub precision: u32,
47}
48
49// ── Type aliases ───────────────────────────────────────────────────────────
50
51/// Pair of frame references used for interpolation.
52pub type FramePair<'a> = (&'a ExportPoseFrame, &'a ExportPoseFrame);
53
54// ── Config ─────────────────────────────────────────────────────────────────
55
56/// Return a sensible default [`PoseExportConfig`].
57#[allow(dead_code)]
58pub fn default_pose_export_config() -> PoseExportConfig {
59    PoseExportConfig {
60        fps: 30.0,
61        include_rotation: true,
62        include_position: true,
63        precision: 6,
64    }
65}
66
67// ── Clip construction ──────────────────────────────────────────────────────
68
69/// Create a new, empty [`ExportPoseClip`].
70#[allow(dead_code)]
71pub fn new_pose_clip(name: &str, fps: f32) -> ExportPoseClip {
72    ExportPoseClip {
73        name: name.to_string(),
74        frames: Vec::new(),
75        fps: fps.max(f32::EPSILON),
76        looping: false,
77    }
78}
79
80/// Append a frame to the clip (kept in insertion order; call `sort` if needed).
81#[allow(dead_code)]
82pub fn add_frame(clip: &mut ExportPoseClip, frame: ExportPoseFrame) {
83    clip.frames.push(frame);
84}
85
86/// Return the number of frames in the clip.
87#[allow(dead_code)]
88pub fn frame_count(clip: &ExportPoseClip) -> usize {
89    clip.frames.len()
90}
91
92/// Return the clip duration in seconds (time of last frame minus first, or 0).
93#[allow(dead_code)]
94pub fn pose_clip_duration(clip: &ExportPoseClip) -> f32 {
95    if clip.frames.len() < 2 {
96        return 0.0;
97    }
98    clip.frames.last().map_or(0.0, |f| f.time) - clip.frames.first().map_or(0.0, |f| f.time)
99}
100
101/// Return the configured FPS of the clip.
102#[allow(dead_code)]
103pub fn pose_clip_fps(clip: &ExportPoseClip) -> f32 {
104    clip.fps
105}
106
107/// Update the FPS of the clip.
108#[allow(dead_code)]
109pub fn set_clip_fps(clip: &mut ExportPoseClip, fps: f32) {
110    clip.fps = fps.max(f32::EPSILON);
111}
112
113// ── Clip operations ────────────────────────────────────────────────────────
114
115/// Remove frames with timestamps outside `[start, end]` (inclusive).
116#[allow(dead_code)]
117pub fn trim_clip(clip: &mut ExportPoseClip, start: f32, end: f32) {
118    clip.frames.retain(|f| f.time >= start && f.time <= end);
119}
120
121/// Reverse the time order of all frames in-place, re-normalizing timestamps
122/// so the first frame starts at 0.
123#[allow(dead_code)]
124pub fn reverse_clip(clip: &mut ExportPoseClip) {
125    clip.frames.reverse();
126    if let Some(first_t) = clip.frames.first().map(|f| f.time) {
127        let total = clip.frames.last().map_or(0.0, |f| f.time) - first_t;
128        for frame in &mut clip.frames {
129            frame.time = total - (frame.time - first_t);
130        }
131        clip.frames.sort_by(|a, b| {
132            a.time
133                .partial_cmp(&b.time)
134                .unwrap_or(std::cmp::Ordering::Equal)
135        });
136    }
137}
138
139/// Scale all frame timestamps by `factor` (e.g. 0.5 = double speed).
140#[allow(dead_code)]
141pub fn scale_clip_timing(clip: &mut ExportPoseClip, factor: f32) {
142    let f = factor.max(f32::EPSILON);
143    for frame in &mut clip.frames {
144        frame.time *= f;
145    }
146}
147
148/// Merge `other` into `clip`, shifting `other`'s timestamps so they start
149/// immediately after the last frame of `clip`.
150#[allow(dead_code)]
151pub fn merge_clips(clip: &mut ExportPoseClip, other: &ExportPoseClip) {
152    let offset = clip.frames.last().map_or(0.0, |f| f.time);
153    let first_other = other.frames.first().map_or(0.0, |f| f.time);
154    for frame in &other.frames {
155        clip.frames.push(ExportPoseFrame {
156            time: offset + (frame.time - first_other),
157            bone_transforms: frame.bone_transforms.clone(),
158        });
159    }
160}
161
162/// Linearly interpolate between two `[f32; 3]` arrays.
163fn lerp3(a: [f32; 3], b: [f32; 3], t: f32) -> [f32; 3] {
164    [
165        a[0] + (b[0] - a[0]) * t,
166        a[1] + (b[1] - a[1]) * t,
167        a[2] + (b[2] - a[2]) * t,
168    ]
169}
170
171/// Linearly interpolate (NLERP) between two quaternions `[f32; 4]`.
172fn nlerp4(a: [f32; 4], b: [f32; 4], t: f32) -> [f32; 4] {
173    let raw = [
174        a[0] + (b[0] - a[0]) * t,
175        a[1] + (b[1] - a[1]) * t,
176        a[2] + (b[2] - a[2]) * t,
177        a[3] + (b[3] - a[3]) * t,
178    ];
179    let len = (raw[0] * raw[0] + raw[1] * raw[1] + raw[2] * raw[2] + raw[3] * raw[3])
180        .sqrt()
181        .max(f32::EPSILON);
182    [raw[0] / len, raw[1] / len, raw[2] / len, raw[3] / len]
183}
184
185/// Sample the clip at `time_sec` by linearly interpolating between the two
186/// nearest frames.  Returns `None` if the clip has no frames.
187#[allow(dead_code)]
188pub fn sample_clip_at(clip: &ExportPoseClip, time_sec: f32) -> Option<ExportPoseFrame> {
189    if clip.frames.is_empty() {
190        return None;
191    }
192    if clip.frames.len() == 1 {
193        return Some(clip.frames[0].clone());
194    }
195    // Clamp to range
196    let t_min = clip.frames.first().map_or(0.0, |f| f.time);
197    let t_max = clip.frames.last().map_or(0.0, |f| f.time);
198    let t = time_sec.clamp(t_min, t_max);
199
200    // Find bracketing pair
201    let idx = clip
202        .frames
203        .partition_point(|f| f.time <= t)
204        .saturating_sub(1)
205        .min(clip.frames.len() - 2);
206
207    let fa = &clip.frames[idx];
208    let fb = &clip.frames[idx + 1];
209    let span = fb.time - fa.time;
210    let alpha = if span.abs() < f32::EPSILON {
211        0.0
212    } else {
213        (t - fa.time) / span
214    };
215
216    let bone_count = fa.bone_transforms.len().min(fb.bone_transforms.len());
217    let bone_transforms = (0..bone_count)
218        .map(|i| {
219            let (pa, ra) = fa.bone_transforms[i];
220            let (pb, rb) = fb.bone_transforms[i];
221            (lerp3(pa, pb, alpha), nlerp4(ra, rb, alpha))
222        })
223        .collect();
224
225    Some(ExportPoseFrame {
226        time: t,
227        bone_transforms,
228    })
229}
230
231// ── Serialization ──────────────────────────────────────────────────────────
232
233/// Serialize the clip to a compact JSON string.
234#[allow(dead_code)]
235pub fn clip_to_json(clip: &ExportPoseClip) -> String {
236    let frame_strs: Vec<String> = clip
237        .frames
238        .iter()
239        .map(|f| {
240            let bt_strs: Vec<String> = f
241                .bone_transforms
242                .iter()
243                .map(|(p, r)| {
244                    format!(
245                        r#"{{"pos":[{},{},{}],"rot":[{},{},{},{}]}}"#,
246                        p[0], p[1], p[2], r[0], r[1], r[2], r[3]
247                    )
248                })
249                .collect();
250            format!(r#"{{"time":{},"bones":[{}]}}"#, f.time, bt_strs.join(","))
251        })
252        .collect();
253    format!(
254        r#"{{"name":"{}","fps":{},"looping":{},"frames":[{}]}}"#,
255        clip.name,
256        clip.fps,
257        clip.looping,
258        frame_strs.join(",")
259    )
260}
261
262/// Serialize the clip to CSV format (one bone-transform per row).
263#[allow(dead_code)]
264pub fn clip_to_csv(clip: &ExportPoseClip) -> String {
265    let mut out = String::from("frame_time,bone_idx,pos_x,pos_y,pos_z,rot_x,rot_y,rot_z,rot_w\n");
266    for f in &clip.frames {
267        for (i, (pos, rot)) in f.bone_transforms.iter().enumerate() {
268            out.push_str(&format!(
269                "{},{},{},{},{},{},{},{},{}\n",
270                f.time, i, pos[0], pos[1], pos[2], rot[0], rot[1], rot[2], rot[3]
271            ));
272        }
273    }
274    out
275}
276
277// ── Tests ──────────────────────────────────────────────────────────────────
278
279#[cfg(test)]
280mod tests {
281    use super::*;
282
283    fn make_frame(time: f32, n: usize) -> ExportPoseFrame {
284        ExportPoseFrame {
285            time,
286            bone_transforms: (0..n)
287                .map(|i| ([i as f32, 0.0, 0.0], [0.0, 0.0, 0.0, 1.0]))
288                .collect(),
289        }
290    }
291
292    #[test]
293    fn test_default_pose_export_config() {
294        let cfg = default_pose_export_config();
295        assert!((cfg.fps - 30.0).abs() < 1e-5);
296        assert!(cfg.include_rotation);
297        assert!(cfg.include_position);
298    }
299
300    #[test]
301    fn test_new_pose_clip() {
302        let clip = new_pose_clip("run", 24.0);
303        assert_eq!(clip.name, "run");
304        assert!((clip.fps - 24.0).abs() < 1e-5);
305        assert!(clip.frames.is_empty());
306    }
307
308    #[test]
309    fn test_add_frame() {
310        let mut clip = new_pose_clip("c", 30.0);
311        add_frame(&mut clip, make_frame(0.0, 2));
312        assert_eq!(frame_count(&clip), 1);
313    }
314
315    #[test]
316    fn test_frame_count_empty() {
317        let clip = new_pose_clip("c", 30.0);
318        assert_eq!(frame_count(&clip), 0);
319    }
320
321    #[test]
322    fn test_pose_clip_duration_two_frames() {
323        let mut clip = new_pose_clip("c", 30.0);
324        add_frame(&mut clip, make_frame(0.0, 1));
325        add_frame(&mut clip, make_frame(1.0, 1));
326        assert!((pose_clip_duration(&clip) - 1.0).abs() < 1e-5);
327    }
328
329    #[test]
330    fn test_pose_clip_duration_single_frame() {
331        let mut clip = new_pose_clip("c", 30.0);
332        add_frame(&mut clip, make_frame(0.5, 1));
333        assert!((pose_clip_duration(&clip) - 0.0).abs() < 1e-5);
334    }
335
336    #[test]
337    fn test_pose_clip_fps() {
338        let clip = new_pose_clip("c", 60.0);
339        assert!((pose_clip_fps(&clip) - 60.0).abs() < 1e-5);
340    }
341
342    #[test]
343    fn test_set_clip_fps() {
344        let mut clip = new_pose_clip("c", 30.0);
345        set_clip_fps(&mut clip, 24.0);
346        assert!((clip.fps - 24.0).abs() < 1e-5);
347    }
348
349    #[test]
350    fn test_trim_clip() {
351        let mut clip = new_pose_clip("c", 30.0);
352        add_frame(&mut clip, make_frame(0.0, 1));
353        add_frame(&mut clip, make_frame(0.5, 1));
354        add_frame(&mut clip, make_frame(1.0, 1));
355        add_frame(&mut clip, make_frame(2.0, 1));
356        trim_clip(&mut clip, 0.4, 1.1);
357        assert_eq!(frame_count(&clip), 2);
358    }
359
360    #[test]
361    fn test_reverse_clip() {
362        let mut clip = new_pose_clip("c", 30.0);
363        add_frame(&mut clip, make_frame(0.0, 1));
364        add_frame(&mut clip, make_frame(1.0, 1));
365        reverse_clip(&mut clip);
366        assert_eq!(frame_count(&clip), 2);
367        // After reversing, first frame time should be <= last frame time
368        assert!(clip.frames[0].time <= clip.frames[1].time);
369    }
370
371    #[test]
372    fn test_scale_clip_timing() {
373        let mut clip = new_pose_clip("c", 30.0);
374        add_frame(&mut clip, make_frame(0.0, 1));
375        add_frame(&mut clip, make_frame(2.0, 1));
376        scale_clip_timing(&mut clip, 0.5);
377        assert!((clip.frames[1].time - 1.0).abs() < 1e-5);
378    }
379
380    #[test]
381    fn test_merge_clips() {
382        let mut clip1 = new_pose_clip("c1", 30.0);
383        add_frame(&mut clip1, make_frame(0.0, 1));
384        add_frame(&mut clip1, make_frame(1.0, 1));
385        let mut clip2 = new_pose_clip("c2", 30.0);
386        add_frame(&mut clip2, make_frame(0.0, 1));
387        add_frame(&mut clip2, make_frame(0.5, 1));
388        merge_clips(&mut clip1, &clip2);
389        assert_eq!(frame_count(&clip1), 4);
390    }
391
392    #[test]
393    fn test_sample_clip_at_empty() {
394        let clip = new_pose_clip("c", 30.0);
395        assert!(sample_clip_at(&clip, 0.5).is_none());
396    }
397
398    #[test]
399    fn test_sample_clip_at_single_frame() {
400        let mut clip = new_pose_clip("c", 30.0);
401        add_frame(&mut clip, make_frame(0.0, 2));
402        let s = sample_clip_at(&clip, 0.5).expect("should succeed");
403        assert_eq!(s.bone_transforms.len(), 2);
404    }
405
406    #[test]
407    fn test_sample_clip_at_midpoint() {
408        let mut clip = new_pose_clip("c", 30.0);
409        add_frame(
410            &mut clip,
411            ExportPoseFrame {
412                time: 0.0,
413                bone_transforms: vec![([0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 1.0])],
414            },
415        );
416        add_frame(
417            &mut clip,
418            ExportPoseFrame {
419                time: 1.0,
420                bone_transforms: vec![([2.0, 0.0, 0.0], [0.0, 0.0, 0.0, 1.0])],
421            },
422        );
423        let s = sample_clip_at(&clip, 0.5).expect("should succeed");
424        assert!((s.bone_transforms[0].0[0] - 1.0).abs() < 1e-4);
425    }
426
427    #[test]
428    fn test_clip_to_json_contains_name() {
429        let mut clip = new_pose_clip("idle", 30.0);
430        add_frame(&mut clip, make_frame(0.0, 1));
431        let json = clip_to_json(&clip);
432        assert!(json.contains("idle"));
433    }
434
435    #[test]
436    fn test_clip_to_csv_has_header() {
437        let mut clip = new_pose_clip("c", 30.0);
438        add_frame(&mut clip, make_frame(0.0, 2));
439        let csv = clip_to_csv(&clip);
440        assert!(csv.starts_with("frame_time"));
441    }
442}