smpl_core/common/
smpl_model.rs

1use super::{
2    expression::ExpressionG,
3    pose::PoseG,
4    smpl_options::SmplOptions,
5    types::{Gender, SmplType},
6};
7use crate::{
8    common::{betas::BetasG, outputs::SmplOutputG, vertex_offsets::VertexOffsetsG},
9    AppBackend,
10};
11use burn::{
12    prelude::Backend,
13    tensor::{Float, Int, Tensor},
14};
15use dyn_clone::DynClone;
16use enum_map::EnumMap;
17use gloss_geometry::csr::VertexFaceCSRBurn;
18use std::any::Any;
19pub trait FaceModel<B: Backend>: Send + Sync + 'static + Any + DynClone {
20    fn expression2offsets(&self, expression: &ExpressionG<B>) -> Tensor<B, 2, Float>;
21    fn get_face_model(&self) -> &dyn FaceModel<B>;
22}
23impl<B: Backend> Clone for Box<dyn FaceModel<B>> {
24    #[allow(unconditional_recursion)]
25    fn clone(&self) -> Box<dyn FaceModel<B>> {
26        self.clone()
27    }
28}
29/// Trait for a Smpl based model. Smpl-rs expects all Smpl models to implement
30/// this.
31pub trait SmplModel<B: Backend>: Send + Sync + 'static + Any + DynClone {
32    fn smpl_type(&self) -> SmplType;
33    fn gender(&self) -> Gender;
34    fn device(&self) -> B::Device;
35    fn forward(
36        &self,
37        options: &SmplOptions,
38        betas: &BetasG<B>,
39        pose_raw: &PoseG<B>,
40        expression: Option<&ExpressionG<B>>,
41        vertex_offsets: Option<&VertexOffsetsG<B>>,
42    ) -> SmplOutputG<B>;
43    fn create_body_with_uv(&self, smpl_output: &SmplOutputG<B>) -> SmplOutputG<B>;
44    fn get_face_model(&self) -> &dyn FaceModel<B>;
45    fn betas2verts(&self, betas: &BetasG<B>) -> Tensor<B, 2, Float>;
46    fn verts2joints(&self, verts_t_pose: Tensor<B, 2, Float>) -> Tensor<B, 2, Float>;
47    fn compute_pose_correctives(&self, pose: &PoseG<B>) -> Tensor<B, 2, Float>;
48    fn compute_pose_feature(&self, pose: &PoseG<B>) -> Tensor<B, 1, Float>;
49    #[allow(clippy::type_complexity)]
50    fn apply_pose(
51        &self,
52        verts_t_pose: &Tensor<B, 2, Float>,
53        joints: &Tensor<B, 2, Float>,
54        lbs_weights: &Tensor<B, 2, Float>,
55        pose: &PoseG<B>,
56    ) -> (Tensor<B, 2, Float>, Tensor<B, 2, Float>);
57    fn faces(&self) -> &Tensor<B, 2, Int>;
58    fn faces_uv(&self) -> &Tensor<B, 2, Int>;
59    fn uv(&self) -> &Tensor<B, 2, Float>;
60    fn lbs_weights(&self) -> Tensor<B, 2, Float>;
61    fn lbs_weights_split(&self) -> Tensor<B, 2, Float>;
62    fn idx_split_2_merged(&self) -> Tensor<B, 1, Int>;
63    fn idx_split_2_merged_vec(&self) -> &Vec<usize>;
64    fn set_pose_dirs(&mut self, posedirs: Tensor<B, 2, Float>);
65    fn get_pose_dirs(&self) -> Tensor<B, 2, Float>;
66    fn get_expression_dirs(&self) -> Option<Tensor<B, 2, Float>>;
67    fn vertex_face_csr(&self) -> Option<VertexFaceCSRBurn<B>>;
68    fn vertex_face_uv_csr(&self) -> Option<VertexFaceCSRBurn<B>>;
69    fn kinematic_tree_depth(&self) -> usize;
70    fn clone_dyn(&self) -> Box<dyn SmplModel<B>>;
71    fn as_any(&self) -> &dyn Any;
72}
73impl<B: Backend> Clone for Box<dyn SmplModel<B>> {
74    #[allow(unconditional_recursion)]
75    fn clone(&self) -> Box<dyn SmplModel<B>> {
76        self.clone()
77    }
78}
79/// A mapping from ``Gender`` to ``SmplModel``
80#[derive(Default, Clone)]
81pub struct Gender2Model<B: Backend> {
82    gender_to_model: EnumMap<Gender, Option<Box<dyn SmplModel<B>>>>,
83}
84#[derive(Default, Clone)]
85pub struct Gender2Path {
86    gender_to_path: EnumMap<Gender, Option<String>>,
87}
88/// A Cache for storing and easy access to ``SmplModels`` which is generic over
89/// Burn Backend
90#[derive(Default, Clone)]
91pub struct SmplCacheG<B: Backend> {
92    type_to_model: EnumMap<SmplType, Gender2Model<B>>,
93    type_to_path: EnumMap<SmplType, Gender2Path>,
94}
95impl<B: Backend> SmplCacheG<B> {
96    pub fn add_model<T: SmplModel<B> + FaceModel<B>>(&mut self, model: T, cache_models: bool) {
97        let smpl_type = model.smpl_type();
98        self.add_model_under_type(smpl_type, model, cache_models);
99    }
100    pub fn add_model_under_type<T: SmplModel<B> + FaceModel<B>>(&mut self, smpl_type: SmplType, model: T, cache_models: bool) {
101        let gender = model.gender();
102        if !cache_models {
103            self.type_to_model = EnumMap::default();
104        }
105        self.type_to_model[smpl_type].gender_to_model[gender] = Some(Box::new(model));
106    }
107    pub fn remove_all_models(&mut self) {
108        self.type_to_model = EnumMap::default();
109    }
110    #[allow(clippy::borrowed_box)]
111    pub fn get_model_box_ref(&self, smpl_type: SmplType, gender: Gender) -> Option<&Box<dyn SmplModel<B>>> {
112        self.type_to_model[smpl_type].gender_to_model[gender].as_ref()
113    }
114    #[allow(clippy::redundant_closure_for_method_calls)]
115    pub fn get_model_ref(&self, smpl_type: SmplType, gender: Gender) -> Option<&dyn SmplModel<B>> {
116        let opt = &self.type_to_model[smpl_type].gender_to_model[gender];
117        let model = opt.as_ref().map(|x| x.as_ref());
118        model
119    }
120    #[allow(clippy::redundant_closure_for_method_calls)]
121    pub fn get_face_model_ref(&self, smpl_type: SmplType, gender: Gender) -> Option<&dyn FaceModel<B>> {
122        let opt = &self.type_to_model[smpl_type].gender_to_model[gender];
123        opt.as_ref().map(|model| model.get_face_model())
124    }
125    #[allow(clippy::redundant_closure_for_method_calls)]
126    pub fn get_model_mut(&mut self, smpl_type: SmplType, gender: Gender) -> Option<&mut dyn SmplModel<B>> {
127        let opt = &mut self.type_to_model[smpl_type].gender_to_model[gender];
128        let model = opt.as_mut().map(|x| x.as_mut());
129        model
130    }
131    pub fn has_model(&self, smpl_type: SmplType, gender: Gender) -> bool {
132        self.type_to_model[smpl_type].gender_to_model[gender].is_some()
133    }
134    pub fn has_lazy_loading(&self, smpl_type: SmplType, gender: Gender) -> bool {
135        self.type_to_path[smpl_type].gender_to_path[gender].is_some()
136    }
137    pub fn get_lazy_loading(&self, smpl_type: SmplType, gender: Gender) -> Option<String> {
138        self.type_to_path[smpl_type].gender_to_path[gender].clone()
139    }
140    pub fn lazy_load_defaults(&mut self) {
141        self.set_lazy_loading(SmplType::SmplX, Gender::Neutral, "./data/smplx/SMPLX_neutral_array_f32_slim.npz");
142        self.set_lazy_loading(SmplType::SmplX, Gender::Male, "./data/smplx/SMPLX_male_array_f32_slim.npz");
143        self.set_lazy_loading(SmplType::SmplX, Gender::Female, "./data/smplx/SMPLX_female_array_f32_slim.npz");
144    }
145    pub fn set_lazy_loading(&mut self, smpl_type: SmplType, gender: Gender, path: &str) {
146        self.type_to_path[smpl_type].gender_to_path[gender] = Some(path.to_string());
147        #[cfg(not(target_arch = "wasm32"))]
148        assert!(
149            std::path::Path::new(&path).exists(),
150            "File at path {path} does not exist. Please follow the data download instructions in the README."
151        );
152    }
153    #[cfg(not(target_arch = "wasm32"))]
154    pub fn add_model_from_type(&mut self, smpl_type: SmplType, path: &str, gender: Gender, max_num_betas: usize, num_expression_components: usize) {
155        match smpl_type {
156            SmplType::SmplX | SmplType::SmplXS => {
157                use crate::smpl_x::smpl_x_gpu::SmplXGPUG;
158                let new_model = SmplXGPUG::new_from_npz(path, gender, max_num_betas, num_expression_components);
159                self.add_model_under_type(smpl_type, new_model, true);
160            }
161            _ => panic!("Model loading for {smpl_type:?} if not supported yet!"),
162        }
163    }
164}
165pub type SmplCache = SmplCacheG<AppBackend>;
smpl_core/common/smpl_model.rs

smpl_core/common/
smpl_model.rs
