burn_human 0.3.0

parametric human mesh pipeline
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162

use burn_human::data::reference::load_reference_bundle;
use burn_human::{AnnyBody, AnnyInput, AnnyReference};
use std::error::Error;

// Reference artifact is stored in mixed f16/f32; allow a looser tolerance.
const TOLERANCE: f64 = 1e-3;

fn load_model() -> Result<AnnyReference, Box<dyn Error>> {
    Ok(AnnyReference::from_paths(
        "assets/model/fullbody_default.safetensors",
        "assets/model/fullbody_default.meta.json",
    )?)
}

#[test]
fn reference_cases_roundtrip() -> Result<(), Box<dyn Error>> {
    let model = load_model()?;
    let bundle = load_reference_bundle(
        "assets/model/fullbody_default.safetensors",
        "assets/model/fullbody_default.meta.json",
    )?;

    for name in model.case_names() {
        let expected = bundle
            .cases
            .iter()
            .find(|c| c.name == name)
            .expect("missing case in bundle");
        let out = model.forward_case(name)?;

        assert_eq!(out.rest_vertices.shape, expected.rest_vertices.shape);
        assert_eq!(out.posed_vertices.shape, expected.posed_vertices.shape);
        assert_eq!(out.rest_bone_poses.shape, expected.rest_bone_poses.shape);
        assert_eq!(out.bone_poses.shape, expected.bone_poses.shape);
        assert_eq!(out.bone_heads.shape, expected.bone_heads.shape);
        assert_eq!(out.bone_tails.shape, expected.bone_tails.shape);

        let rest_err = max_abs_diff(&out.rest_vertices.data, &expected.rest_vertices.data);
        let posed_err = max_abs_diff(&out.posed_vertices.data, &expected.posed_vertices.data);
        let rest_bone_err = max_abs_diff(&out.rest_bone_poses.data, &expected.rest_bone_poses.data);
        let bone_err = max_abs_diff(&out.bone_poses.data, &expected.bone_poses.data);
        let head_err = max_abs_diff(&out.bone_heads.data, &expected.bone_heads.data);
        let tail_err = max_abs_diff(&out.bone_tails.data, &expected.bone_tails.data);

        assert!(
            rest_err < TOLERANCE,
            "rest vertices err {} for case {}",
            rest_err,
            name
        );
        assert!(
            posed_err < TOLERANCE,
            "posed vertices err {} for case {}",
            posed_err,
            name
        );
        assert!(
            rest_bone_err < TOLERANCE,
            "rest bone pose err {} for case {}",
            rest_bone_err,
            name
        );
        assert!(
            bone_err < TOLERANCE,
            "bone pose err {} for case {}",
            bone_err,
            name
        );
        assert!(
            head_err < TOLERANCE,
            "bone heads err {} for case {}",
            head_err,
            name
        );
        assert!(
            tail_err < TOLERANCE,
            "bone tails err {} for case {}",
            tail_err,
            name
        );
    }

    Ok(())
}

#[test]
fn reference_metadata_present() -> Result<(), Box<dyn Error>> {
    let model = load_model()?;
    let meta = &model.metadata().metadata;
    assert!(!meta.case_names.is_empty());
    assert!(!meta.bone_labels.is_empty());
    assert!(!meta.phenotype_labels.is_empty());
    assert!(!meta.rig.is_empty());
    assert!(!meta.topology.is_empty());
    assert!(!meta.pose_parameterization.is_empty());
    assert!(!meta.bone_parents.is_empty());
    assert!(!meta.macrodetail_keys.is_empty());
    assert!(!meta.phenotype_variations.is_empty());
    assert!(!meta.phenotype_anchors.is_empty());
    Ok(())
}

#[test]
fn anny_body_forward_returns_reference_outputs() -> Result<(), Box<dyn Error>> {
    let body = AnnyBody::from_reference_paths(
        "assets/model/fullbody_default.safetensors",
        "assets/model/fullbody_default.meta.json",
    )?;
    let bundle = load_reference_bundle(
        "assets/model/fullbody_default.safetensors",
        "assets/model/fullbody_default.meta.json",
    )?;

    for case in &bundle.cases {
        let out = body.forward(AnnyInput {
            case_name: Some(&case.name),
            ..Default::default()
        })?;
        // Compare rest vertices and posed vertices.
        assert_eq!(out.posed_vertices.shape, case.posed_vertices.shape);
        assert_eq!(out.rest_vertices.shape, case.rest_vertices.shape);
        assert_eq!(out.bone_poses.shape, case.bone_poses.shape);
        assert_eq!(out.rest_bone_poses.shape, case.rest_bone_poses.shape);

        let max_abs = max_abs_diff(&out.posed_vertices.data, &case.posed_vertices.data);
        assert!(
            max_abs < TOLERANCE,
            "posed vertices mismatch for case {} (max abs diff {})",
            case.name,
            max_abs
        );
        let max_rest = max_abs_diff(&out.rest_vertices.data, &case.rest_vertices.data);
        assert!(
            max_rest < TOLERANCE,
            "rest vertices mismatch for case {} (max abs diff {})",
            case.name,
            max_rest
        );

        let max_bone = max_abs_diff(&out.bone_poses.data, &case.bone_poses.data);
        assert!(
            max_bone < TOLERANCE,
            "bone poses mismatch for case {} (max abs diff {})",
            case.name,
            max_bone
        );
        assert_eq!(out.posed_vertices.shape, case.posed_vertices.shape);
    }
    Ok(())
}

fn max_abs_diff(a: &[f64], b: &[f64]) -> f64 {
    assert_eq!(a.len(), b.len());
    let mut max = 0.0;
    for (x, y) in a.iter().zip(b.iter()) {
        let d = (x - y).abs();
        if d > max {
            max = d;
        }
    }
    max
}