1use crate::scene::SceneAnimation;
2use burn::backend::{Candle, NdArray, Wgpu};
3use burn::prelude::Backend;
4use gloss_img::dynamic_image::DynImage;
5use gloss_renderer::{
6 components::{DiffuseImg, MetalnessImg, Name, NormalImg, RoughnessImg},
7 scene::Scene,
8};
9use gloss_utils::{
10 bshare::{ToNalgebraFloat, ToNalgebraInt, ToNdArray},
11 nshare::ToNalgebra,
12};
13use image::imageops::FilterType;
14use log::info;
15use nalgebra::DMatrix;
16use ndarray::{self as nd, s};
17use smpl_core::common::types::SmplType;
18use smpl_core::{
19 codec::gltf::GltfCodec,
20 common::{metadata::smpl_metadata, pose::Pose, smpl_model::SmplCache, smpl_params::SmplParams},
21 conversions::pose_remap::PoseRemap,
22};
23use smpl_core::{
24 codec::{gltf::PerBodyData, scene::CameraTrack},
25 common::{
26 animation::Animation, betas::Betas, expression::Expression, pose_override::PoseOverride, pose_retarget::RetargetPoseYShift,
27 smpl_model::SmplCacheDynamic, smpl_options::SmplOptions, types::UpAxis,
28 },
29};
30use smpl_utils::array::{Gather2D, Gather3D};
31use std::f32::consts::PI;
32pub trait GltfCodecGloss {
34 fn from_scene(scene: &Scene, max_texture_size: Option<u32>, ignore_ent: Option<String>) -> GltfCodec;
35}
36fn get_image(image: &DynImage, to_gray: bool, max_texture_size: Option<u32>) -> DynImage {
37 let mut image = image.clone();
38 if to_gray {
39 image = image.grayscale();
40 }
41 if let Some(force_image_size) = max_texture_size {
42 if image.width() > force_image_size {
43 image.resize(force_image_size, force_image_size, FilterType::Gaussian)
44 } else {
45 image
46 }
47 } else {
48 image
49 }
50}
51impl GltfCodecGloss for GltfCodec {
53 fn from_scene(scene: &Scene, max_texture_size: Option<u32>, ignore_ent: Option<String>) -> GltfCodec {
55 let smpl_models = scene.get_resource::<&SmplCacheDynamic>().unwrap();
56 match &*smpl_models {
57 SmplCacheDynamic::NdArray(models) => from_scene_on_backend::<NdArray>(scene, models, max_texture_size, &ignore_ent),
58 SmplCacheDynamic::Wgpu(models) => from_scene_on_backend::<Wgpu>(scene, models, max_texture_size, &ignore_ent),
59 SmplCacheDynamic::Candle(models) => from_scene_on_backend::<Candle>(scene, models, max_texture_size, &ignore_ent),
60 }
61 }
62}
63#[allow(clippy::too_many_lines)]
66#[allow(clippy::trivially_copy_pass_by_ref)]
67fn from_scene_on_backend<B: Backend>(
68 scene: &Scene,
69 smpl_models: &SmplCache<B>,
70 max_texture_size: Option<u32>,
71 ignore_ent: &Option<String>,
72) -> GltfCodec
73where
74 <B as Backend>::FloatTensorPrimitive<2>: Sync,
75 <B as Backend>::IntTensorPrimitive<2>: Sync,
76 B::QuantizedTensorPrimitive<1>: std::marker::Sync,
77 B::QuantizedTensorPrimitive<2>: std::marker::Sync,
78 B::QuantizedTensorPrimitive<3>: std::marker::Sync,
79{
80 let now = wasm_timer::Instant::now();
81 let mut gltf_codec = GltfCodec::default();
82 let scene_anim = scene.get_resource::<&SceneAnimation>().unwrap();
83 let nr_frames = scene_anim.num_frames;
84 let fps = scene_anim.config.fps;
85 let mut cameras_query = scene.world.query::<&CameraTrack>();
86 for (_, camera_track) in cameras_query.iter() {
87 gltf_codec.camera_track = Some(camera_track.clone());
88 }
89 let mut query = scene.world.query::<(&SmplParams, &Name)>();
90 let num_bodies = query.iter().len();
91 gltf_codec.num_bodies = num_bodies;
92 let mut should_export_posedirs = false;
93 let mut should_export_exprdirs = false;
94 let mut num_expression_blend_shapes = 0;
95 for (entity, (smpl_params, _name)) in query.iter() {
96 if scene.world.has::<Animation>(entity).unwrap() && smpl_params.enable_pose_corrective {
97 should_export_posedirs = true;
98 }
99 let smpl_model = smpl_models.get_model_ref(smpl_params.smpl_type, smpl_params.gender).unwrap();
100 if let Ok(anim) = scene.get_comp::<&Animation>(&entity) {
101 if anim.has_expression() && smpl_model.get_expression_dirs().is_some() {
102 should_export_exprdirs = true;
103 num_expression_blend_shapes = smpl_model.get_expression_dirs().unwrap().shape().dims[1];
104 }
105 }
106 }
107 for (body_idx, (entity, (smpl_params, name))) in query.iter().enumerate() {
108 if let Some(ent_to_ignore) = ignore_ent {
109 if *ent_to_ignore == name.0 {
110 continue;
111 }
112 }
113 let smpl_version = smpl_params.smpl_type;
114 let gender = smpl_params.gender as i32;
115 let mut current_body = PerBodyData::default();
116 assert!(smpl_version != SmplType::SmplPP, "GLTF export for SMPL++ is not supported yet!");
117 let smpl_model = smpl_models.get_model_ref(smpl_params.smpl_type, smpl_params.gender).unwrap();
118 let Ok(betas) = scene.get_comp::<&Betas>(&entity) else {
119 panic!("Betas component does not exist!");
120 };
121 let default_pose = Pose::new_empty(UpAxis::Y, smpl_params.smpl_type);
122 let default_expression = Expression::new_empty(10);
123 let mut smpl_output = smpl_model.forward(&SmplOptions::default(), &betas, &default_pose, Some(&default_expression));
124 smpl_output.compute_normals();
125 smpl_output = smpl_model.create_body_with_uv(&smpl_output);
126 let metadata = smpl_metadata(&smpl_params.smpl_type);
127 let mut num_total_blendshapes = 0;
128 if should_export_posedirs {
129 num_total_blendshapes += metadata.num_pose_blend_shapes + 1;
130 }
131 if should_export_exprdirs {
132 num_total_blendshapes += num_expression_blend_shapes;
133 }
134 gltf_codec.smpl_type = smpl_version;
135 gltf_codec.gender = gender;
136 current_body.pose = Some(default_pose.clone());
137 gltf_codec.default_joint_poses = Some(default_pose.clone().joint_poses);
138 current_body.body_translation = Some(default_pose.clone().global_trans.to_shape((1, 3)).unwrap().to_owned());
139 let verts_na = smpl_output.verts.to_nalgebra();
140 let normals_na = smpl_output.normals.as_ref().expect("SMPL Output is missing normals!").to_nalgebra();
141 let faces_na = smpl_output.faces.to_nalgebra();
142 let uvs_na = smpl_output.uvs.as_ref().expect("SMPL Output is missing UVs!").to_nalgebra();
143 current_body.positions = Some(verts_na);
144 current_body.normals = Some(normals_na);
145 gltf_codec.faces = Some(faces_na);
146 gltf_codec.uvs = Some(uvs_na);
147 let smpl_joints = smpl_output.joints.clone().to_ndarray();
148 let joint_count = smpl_joints.shape()[0];
149 let lbs_weights = smpl_model.lbs_weights_split().to_ndarray();
150 let vertex_count = smpl_output.verts.dims()[0];
151 let mut skin_vertex_index = DMatrix::<u32>::zeros(vertex_count, 4);
152 let mut skin_vertex_weight = DMatrix::<f32>::zeros(vertex_count, 4);
153 for (vertex_id, row) in lbs_weights.outer_iter().enumerate() {
154 let mut vertex_weights: Vec<(usize, f32)> = row.iter().enumerate().map(|(index, &weight)| (index, weight)).collect();
155 vertex_weights.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
156 assert_eq!(vertex_weights.len().min(4), 4, "Illegal vertex weights");
157 for (i, (index, weight)) in vertex_weights.iter().take(4).enumerate() {
158 skin_vertex_index[(vertex_id, i)] = u32::try_from(*index).expect("Cannot convert to u32!");
159 skin_vertex_weight[(vertex_id, i)] = *weight;
160 }
161 }
162 gltf_codec.joint_index = Some(skin_vertex_index);
163 gltf_codec.joint_weight = Some(skin_vertex_weight);
164 let diffuse_img = scene.get_comp::<&DiffuseImg>(&entity);
165 if let Ok(diffuse_img) = diffuse_img {
166 if let Some(img) = &diffuse_img.generic_img.cpu_img {
167 current_body.diffuse_textures = Some(get_image(img, false, max_texture_size));
168 }
169 }
170 let normals_img = scene.get_comp::<&NormalImg>(&entity);
171 if let Ok(normals_img) = normals_img {
172 if let Some(img) = &normals_img.generic_img.cpu_img {
173 current_body.normals_textures = Some(get_image(img, false, max_texture_size));
174 }
175 }
176 let metalness_img = scene.get_comp::<&MetalnessImg>(&entity);
177 if let Ok(metalness_img) = metalness_img {
178 if let Some(img) = &metalness_img.generic_img.cpu_img {
179 current_body.metalness_textures = Some(get_image(img, true, max_texture_size));
180 }
181 }
182 let roughness_img = scene.get_comp::<&RoughnessImg>(&entity);
183 if let Ok(roughness_img) = roughness_img {
184 if let Some(img) = &roughness_img.generic_img.cpu_img {
185 current_body.roughness_textures = Some(get_image(img, true, max_texture_size));
186 }
187 }
188 if scene.world.has::<Pose>(entity).unwrap() && !scene.world.has::<Animation>(entity).unwrap() {
189 let Ok(pose_ref) = scene.get_comp::<&Pose>(&entity) else {
190 panic!("Pose component doesn't exist");
191 };
192 let current_pose: &Pose = &pose_ref;
193 let current_body_translation = current_pose.global_trans.to_shape((1, 3)).unwrap().to_owned();
194 current_body.pose = Some(current_pose.clone());
195 current_body.body_translation = Some(current_body_translation);
196 if smpl_params.enable_pose_corrective {
197 let vertex_offsets_merged = smpl_model.compute_pose_correctives(current_pose).to_ndarray();
198 let mapping = &smpl_model.idx_split_2_merged_vec();
199 let cols = vec![0, 1, 2];
200 let vertex_offsets = vertex_offsets_merged.gather(mapping, &cols).into_nalgebra();
201 current_body.positions = Some(current_body.positions.as_ref().unwrap() + vertex_offsets);
202 }
203 }
204 #[allow(clippy::cast_precision_loss)]
205 if scene.world.has::<Animation>(entity).unwrap() {
206 info!("Processing Animation for body {:?}", body_idx);
207 let anim = scene.get_comp::<&Animation>(&entity).unwrap();
208 gltf_codec.frame_count = Some(nr_frames);
209 let mut keyframe_times: Vec<f32> = Vec::new();
210 let mut current_body_rotations = nd::Array3::<f32>::zeros((joint_count, nr_frames, 3));
211 let mut current_body_translations = nd::Array2::<f32>::zeros((nr_frames, 3));
212 let mut current_body_scales = nd::Array2::<f32>::zeros((nr_frames, 3));
213 let mut current_per_frame_blend_weights = nd::Array2::<f32>::zeros((nr_frames, num_total_blendshapes));
214 if should_export_posedirs || should_export_exprdirs {
215 let mut full_morph_targets = nd::Array3::<f32>::zeros((num_total_blendshapes, vertex_count, 3));
216 let mut running_idx_morph_target = 0;
217 if should_export_posedirs {
218 let mut pose_morph_targets = nd::Array3::<f32>::zeros((metadata.num_pose_blend_shapes + 1, vertex_count, 3));
219 let nr_elem_merged = smpl_model.get_pose_dirs().dims()[0] / 3;
220 let pose_dirs_merged = smpl_model
221 .get_pose_dirs()
222 .to_ndarray()
223 .into_shape_with_order((nr_elem_merged, 3, metadata.num_pose_blend_shapes))
224 .unwrap();
225 let mapping = smpl_model.idx_split_2_merged_vec();
226 let cols = vec![0, 1, 2];
227 let depth = (0..metadata.num_pose_blend_shapes).collect::<Vec<_>>().into_boxed_slice();
228 let pose_blend_shapes = pose_dirs_merged
229 .gather(mapping, &cols, &depth)
230 .into_shape_with_order((vertex_count, 3, metadata.num_pose_blend_shapes))
231 .unwrap()
232 .permuted_axes([2, 0, 1]);
233 let morph_targets = (2.0 * PI) * pose_blend_shapes.clone();
234 let pi = nd::Array1::<f32>::from_elem(metadata.num_pose_blend_shapes, -PI);
235 let pi_array = pi.insert_axis(nd::Axis(1)).insert_axis(nd::Axis(2));
236 assert_eq!(pose_blend_shapes.shape()[0], pi_array.len());
237 let template_offset = (pose_blend_shapes * &pi_array).sum_axis(nd::Axis(0));
238 pose_morph_targets
239 .slice_mut(s![0..metadata.num_pose_blend_shapes, .., ..])
240 .assign(&morph_targets);
241 pose_morph_targets
242 .slice_mut(s![metadata.num_pose_blend_shapes, .., ..])
243 .assign(&template_offset);
244 #[allow(clippy::range_plus_one)]
245 full_morph_targets
246 .slice_mut(s![
247 running_idx_morph_target..running_idx_morph_target + metadata.num_pose_blend_shapes + 1,
248 ..,
249 ..
250 ])
251 .assign(&pose_morph_targets);
252 running_idx_morph_target += metadata.num_pose_blend_shapes + 1;
253 }
254 #[allow(unused_assignments)]
255 if should_export_exprdirs {
256 if let Some(expr_dirs) = smpl_model.get_expression_dirs() {
257 let nr_elem_merged = expr_dirs.dims()[0] / 3;
258 let expression_dirs_merged = expr_dirs
259 .to_ndarray()
260 .into_shape_with_order((nr_elem_merged, 3, num_expression_blend_shapes))
261 .unwrap();
262 let mapping = smpl_model.idx_split_2_merged_vec();
263 let cols = vec![0, 1, 2];
264 let depth = (0..metadata.expression_space_dim).collect::<Vec<_>>().into_boxed_slice();
265 let expression_dirs_split = expression_dirs_merged
266 .gather(mapping, &cols, &depth)
267 .into_shape_with_order((vertex_count, 3, num_expression_blend_shapes))
268 .unwrap()
269 .permuted_axes([2, 0, 1]);
270 full_morph_targets
271 .slice_mut(s![
272 running_idx_morph_target..running_idx_morph_target + num_expression_blend_shapes,
273 ..,
274 ..
275 ])
276 .assign(&expression_dirs_split);
277 running_idx_morph_target += num_expression_blend_shapes;
278 }
279 }
280 gltf_codec.morph_targets = Some(full_morph_targets);
281 }
282 for global_frame_idx in 0..nr_frames {
283 keyframe_times.push((global_frame_idx as f32) / fps);
284 if global_frame_idx < anim.start_offset || global_frame_idx > anim.start_offset + anim.num_animation_frames() {
285 continue;
286 }
287 let mut local_frame_idx = global_frame_idx - anim.start_offset;
288 if global_frame_idx == (anim.start_offset + anim.num_animation_frames()) {
289 local_frame_idx -= 1;
290 }
291 let mut pose = anim.get_pose_at_idx(local_frame_idx);
292 let pose_remap = PoseRemap::new(pose.smpl_type, smpl_params.smpl_type);
293 pose = pose_remap.remap(&pose);
294 if let Ok(ref pose_mask) = scene.get_comp::<&PoseOverride>(&entity) {
295 let mut new_pose_mask = PoseOverride::clone(pose_mask);
296 pose.apply_mask(&mut new_pose_mask);
297 }
298 if let Ok(ref pose_retarget) = scene.get_comp::<&RetargetPoseYShift>(&entity) {
299 let mut pose_retarget_local = RetargetPoseYShift::clone(pose_retarget);
300 pose_retarget_local.apply(&mut pose);
301 }
302 current_body_rotations.slice_mut(s![.., global_frame_idx, ..]).assign(&pose.joint_poses);
303 let mut skeleton_root_translation = pose.global_trans.to_owned();
304 let root_translation = smpl_output.joints.to_ndarray().slice(s![0, ..]).to_owned();
305 skeleton_root_translation = skeleton_root_translation + root_translation;
306 current_body_translations
307 .slice_mut(s![global_frame_idx, ..])
308 .assign(&skeleton_root_translation);
309 if global_frame_idx < (anim.start_offset + anim.num_animation_frames()) {
310 current_body_scales.slice_mut(s![global_frame_idx, ..]).assign(&nd::Array1::ones(3));
311 }
312 let mut running_idx_morph_target = 0;
313 if should_export_posedirs {
314 let pose_blend_weights = &smpl_model.compute_pose_feature(&pose);
315 let rescaled_pose_blend_weights = pose_blend_weights.map(|&elem| (elem + PI) / (2.0 * PI));
316 current_per_frame_blend_weights
317 .slice_mut(s![global_frame_idx, 0..metadata.num_pose_blend_shapes])
318 .assign(&rescaled_pose_blend_weights);
319 if global_frame_idx == (anim.start_offset + anim.num_animation_frames()) {
320 current_per_frame_blend_weights
321 .slice_mut(s![global_frame_idx..nr_frames, 0..metadata.num_pose_blend_shapes])
322 .assign(&rescaled_pose_blend_weights);
323 }
324 running_idx_morph_target += metadata.num_pose_blend_shapes + 1;
325 }
326 #[allow(unused_assignments)]
327 if should_export_exprdirs {
328 let expr_opt = anim.get_expression_at_idx(local_frame_idx);
329 if let Some(expr) = expr_opt.as_ref() {
330 let max_nr_expr_coeffs = num_expression_blend_shapes.min(expr.expr_coeffs.len());
331 let expr_coeffs = expr.expr_coeffs.slice(s![0..max_nr_expr_coeffs]);
332 current_per_frame_blend_weights
333 .slice_mut(s![
334 global_frame_idx,
335 running_idx_morph_target..running_idx_morph_target + max_nr_expr_coeffs
336 ])
337 .assign(&expr_coeffs);
338 }
339 running_idx_morph_target += num_expression_blend_shapes;
340 }
341 }
342 gltf_codec.keyframe_times = Some(keyframe_times);
343 current_body.body_scales = Some(current_body_scales);
344 current_body.body_translations = Some(current_body_translations);
345 current_body.body_rotations = Some(current_body_rotations);
346 if should_export_posedirs {
347 current_per_frame_blend_weights
348 .slice_mut(s![.., metadata.num_pose_blend_shapes])
349 .assign(&nd::Array1::<f32>::from_elem(nr_frames, 1.0));
350 }
351 if should_export_posedirs || should_export_exprdirs {
352 current_body.per_frame_blend_weights = Some(current_per_frame_blend_weights);
353 }
354 }
355 current_body.default_joint_translations = Some(smpl_joints);
356 gltf_codec.per_body_data.push(current_body);
357 }
358 info!(
359 "Writing {} body scene to GltfCodec: Took {} seconds for {} frames",
360 num_bodies,
361 now.elapsed().as_secs(),
362 nr_frames
363 );
364 gltf_codec
365}