1use crate::scene::SceneAnimation;
2use gloss_img::dynamic_image::DynImage;
3use gloss_renderer::{
4 components::{DiffuseImg, Faces, MetalnessImg, Name, NormalImg, RoughnessImg, UVs, Verts},
5 scene::Scene,
6};
7use gloss_utils::{
8 bshare::{ToNalgebraFloat, ToNalgebraInt, ToNdArray},
9 nshare::ToNalgebra,
10 tensor::DynamicMatrixOps,
11};
12use image::imageops::FilterType;
13use log::info;
14use nalgebra::DMatrix;
15use ndarray::{self as nd, s};
16use smpl_core::common::types::SmplType;
17use smpl_core::{
18 codec::gltf::GltfCodec,
19 common::{metadata::smpl_metadata, pose::Pose, smpl_params::SmplParams},
20 conversions::pose_remap::PoseRemap,
21};
22use smpl_core::{codec::gltf::PropData, common::transform_sequence::TransformSequence};
23use smpl_core::{
24 codec::{gltf::PerBodyData, scene::CameraTrack},
25 common::{
26 animation::Animation,
27 betas::Betas,
28 expression::Expression,
29 pose_override::PoseOverride,
30 pose_retarget::RetargetPoseYShift,
31 smpl_model::SmplCache,
32 smpl_options::SmplOptions,
33 types::{FaceType, UpAxis},
34 },
35};
36use smpl_utils::array::{Gather2D, Gather3D};
37use std::f32::consts::PI;
38use std::ops::Deref;
39pub trait GltfCodecGloss {
41 fn from_scene(scene: &Scene, options: &GltfInteropOptions) -> GltfCodec;
42 fn from_entities(scene: &Scene, options: &GltfInteropOptions, entities: Vec<String>) -> GltfCodec;
43 fn from_scene_with_body_indices(scene: &Scene, options: &GltfInteropOptions, body_idxs: Vec<usize>) -> GltfCodec;
44}
45fn get_image(image: &DynImage, to_gray: bool, max_texture_size: Option<u32>) -> DynImage {
46 let mut image = image.clone();
47 if to_gray {
48 image = image.grayscale();
49 }
50 if let Some(force_image_size) = max_texture_size {
51 if image.width() > force_image_size {
52 image.resize(force_image_size, force_image_size, FilterType::Gaussian)
53 } else {
54 image
55 }
56 } else {
57 image
58 }
59}
60pub struct GltfInteropOptions {
61 pub max_texture_size: Option<u32>,
62 pub export_camera: bool,
63 pub export_shape: bool,
64}
65impl Default for GltfInteropOptions {
66 fn default() -> Self {
67 Self {
68 max_texture_size: None,
69 export_camera: true,
70 export_shape: true,
71 }
72 }
73}
74impl GltfCodecGloss for GltfCodec {
76 fn from_scene(scene: &Scene, options: &GltfInteropOptions) -> GltfCodec {
78 let smpl_models = scene.get_resource::<&SmplCache>().unwrap();
79 gltfcodec_from_scene(scene, &smpl_models, options, None)
80 }
81 fn from_entities(scene: &Scene, options: &GltfInteropOptions, entities: Vec<String>) -> GltfCodec {
83 let smpl_models = scene.get_resource::<&SmplCache>().unwrap();
84 gltfcodec_from_scene(scene, &smpl_models, options, Some(&entities))
85 }
86 fn from_scene_with_body_indices(scene: &Scene, options: &GltfInteropOptions, body_idxs: Vec<usize>) -> GltfCodec {
88 assert!(!body_idxs.is_empty(), "At least one index must be specified");
89 let smpl_models = scene.get_resource::<&SmplCache>().unwrap();
90 let num_smpl_bodies = scene.world.query::<&SmplParams>().iter().len();
91 if *body_idxs.iter().max().unwrap() >= num_smpl_bodies {
92 info!("Some of the indices are out of range. Ignoring these indices.");
93 }
94 let entities_to_export: Vec<String> = body_idxs.into_iter().map(|idx| format!("avatar-{:02}", idx + 1)).collect();
95 gltfcodec_from_scene(scene, &smpl_models, options, Some(&entities_to_export))
96 }
97}
98#[allow(clippy::too_many_lines)]
100#[allow(clippy::trivially_copy_pass_by_ref)]
101fn gltfcodec_from_scene(scene: &Scene, smpl_models: &SmplCache, options: &GltfInteropOptions, entities: Option<&Vec<String>>) -> GltfCodec {
102 let now = wasm_timer::Instant::now();
103 let mut gltf_codec = GltfCodec::default();
104 let mut nr_frames = 0;
105 if options.export_camera {
106 let mut cameras_query = scene.world.query::<&CameraTrack>();
107 for (_, camera_track) in cameras_query.iter() {
108 gltf_codec.camera_track = Some(camera_track.clone());
109 }
110 }
111 let mut query = scene.world.query::<(&SmplParams, &Name)>();
112 let num_bodies = query.iter().len();
113 gltf_codec.num_bodies = num_bodies;
114 let mut should_export_posedirs = false;
115 let mut should_export_exprdirs = false;
116
117 let mut num_expression_blend_shapes = 0;
118 for (entity, (smpl_params, _name)) in query.iter() {
119 if scene.world.has::<Animation>(entity).unwrap() && smpl_params.enable_pose_corrective {
120 should_export_posedirs = true;
121 }
122 let smpl_model = smpl_models.get_model_ref(smpl_params.smpl_type, smpl_params.gender).unwrap();
123 if let Ok(anim) = scene.get_comp::<&Animation>(&entity) {
124 if anim.has_expression() && smpl_model.get_expression_dirs().is_some() {
125 should_export_exprdirs = true;
126 num_expression_blend_shapes = smpl_model.get_expression_dirs().unwrap().shape().dims[1];
127 }
128 }
129 }
130
131 if let Ok(scene_anim) = scene.get_resource::<&SceneAnimation>() {
132 nr_frames = scene_anim.num_frames;
133 let mut keyframe_times: Vec<f32> = Vec::new();
134 let fps = scene_anim.config.fps;
135 #[allow(clippy::cast_precision_loss)]
136 for global_frame_idx in 0..nr_frames {
137 keyframe_times.push((global_frame_idx as f32) / fps);
138 }
139 gltf_codec.keyframe_times = Some(keyframe_times);
140 gltf_codec.frame_count = Some(nr_frames);
141 }
142 for (body_idx, (entity, (smpl_params, name))) in query.iter().enumerate() {
143 if let Some(entities) = entities {
144 if !entities.contains(&name.0) {
145 continue;
146 }
147 }
148 let smpl_version = smpl_params.smpl_type;
149 let gender = smpl_params.gender as i32;
150 let mut current_body = PerBodyData::default();
151 assert!(smpl_version != SmplType::SmplPP, "GLTF export for SMPL++ is not supported yet!");
152 let smpl_model = smpl_models.get_model_ref(smpl_params.smpl_type, smpl_params.gender).unwrap();
153 let betas = if options.export_shape {
154 scene
155 .get_comp::<&Betas>(&entity)
156 .expect("Betas component does not exist!")
157 .deref()
158 .clone()
159 } else {
160 Betas::default()
161 };
162 let default_pose = Pose::new_empty(UpAxis::Y, smpl_params.smpl_type);
163 let default_expression = Expression::new_empty(10, FaceType::SmplX);
164 let mut smpl_output = smpl_model.forward(&SmplOptions::default(), &betas, &default_pose, Some(&default_expression));
165 smpl_output.compute_normals();
166 smpl_output = smpl_model.create_body_with_uv(&smpl_output);
167 let metadata = smpl_metadata(&smpl_params.smpl_type);
168 let mut num_total_blendshapes = 0;
169 if should_export_posedirs {
170 num_total_blendshapes += metadata.num_pose_blend_shapes + 1;
171 }
172 if should_export_exprdirs {
173 num_total_blendshapes += num_expression_blend_shapes + 1;
174 }
175 gltf_codec.smpl_type = smpl_version;
176 gltf_codec.gender = gender;
177 current_body.joint_poses = Some(default_pose.joint_poses.to_ndarray().clone());
178 gltf_codec.default_joint_poses = Some(default_pose.clone().joint_poses.to_ndarray());
179 current_body.body_translation = Some(default_pose.clone().global_trans.to_ndarray().to_shape((1, 3)).unwrap().to_owned());
180 let verts_na = smpl_output.verts.to_nalgebra();
181 let normals_na = smpl_output.normals.as_ref().expect("SMPL Output is missing normals!").to_nalgebra();
182 let faces_na = smpl_output.faces.to_nalgebra();
183 let uvs_na = smpl_output.uvs.as_ref().expect("SMPL Output is missing UVs!").to_nalgebra();
184 current_body.positions = Some(verts_na);
185 current_body.normals = Some(normals_na);
186 gltf_codec.faces = Some(faces_na);
187 gltf_codec.uvs = Some(uvs_na);
188 let smpl_joints = smpl_output.joints.clone().to_ndarray();
189 let joint_count = smpl_joints.shape()[0];
190 let lbs_weights = smpl_model.lbs_weights_split().to_ndarray();
191 let vertex_count = smpl_output.verts.dims()[0];
192 let mut skin_vertex_index = DMatrix::<u32>::zeros(vertex_count, 4);
193 let mut skin_vertex_weight = DMatrix::<f32>::zeros(vertex_count, 4);
194 for (vertex_id, row) in lbs_weights.outer_iter().enumerate() {
195 let mut vertex_weights: Vec<(usize, f32)> = row.iter().enumerate().map(|(index, &weight)| (index, weight)).collect();
196 vertex_weights.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
197 assert_eq!(vertex_weights.len().min(4), 4, "Illegal vertex weights");
198 for (i, (index, weight)) in vertex_weights.iter().take(4).enumerate() {
199 skin_vertex_index[(vertex_id, i)] = u32::try_from(*index).expect("Cannot convert to u32!");
200 skin_vertex_weight[(vertex_id, i)] = *weight;
201 }
202 }
203 gltf_codec.joint_index = Some(skin_vertex_index);
204 gltf_codec.joint_weight = Some(skin_vertex_weight);
205 let diffuse_img = scene.get_comp::<&DiffuseImg>(&entity);
206 if let Ok(diffuse_img) = diffuse_img {
207 if let Some(img) = &diffuse_img.generic_img.cpu_img {
208 current_body.diffuse_textures = Some(get_image(img, false, options.max_texture_size));
209 }
210 }
211 let normals_img = scene.get_comp::<&NormalImg>(&entity);
212 if let Ok(normals_img) = normals_img {
213 if let Some(img) = &normals_img.generic_img.cpu_img {
214 current_body.normals_textures = Some(get_image(img, false, options.max_texture_size));
215 }
216 }
217 let metalness_img = scene.get_comp::<&MetalnessImg>(&entity);
218 if let Ok(metalness_img) = metalness_img {
219 if let Some(img) = &metalness_img.generic_img.cpu_img {
220 current_body.metalness_textures = Some(get_image(img, true, options.max_texture_size));
221 }
222 }
223 let roughness_img = scene.get_comp::<&RoughnessImg>(&entity);
224 if let Ok(roughness_img) = roughness_img {
225 if let Some(img) = &roughness_img.generic_img.cpu_img {
226 current_body.roughness_textures = Some(get_image(img, true, options.max_texture_size));
227 }
228 }
229 if scene.world.has::<Pose>(entity).unwrap() && !scene.world.has::<Animation>(entity).unwrap() {
230 let Ok(pose_ref) = scene.get_comp::<&Pose>(&entity) else {
231 panic!("Pose component doesn't exist");
232 };
233 let current_pose: &Pose = &pose_ref;
234 let current_body_translation = current_pose.global_trans.to_ndarray().to_shape((1, 3)).unwrap().to_owned();
235 current_body.joint_poses = Some(current_pose.joint_poses.clone().to_ndarray());
236 current_body.body_translation = Some(current_body_translation);
237 if smpl_params.enable_pose_corrective {
238 let vertex_offsets_merged = smpl_model.compute_pose_correctives(current_pose).to_ndarray();
239 let mapping = &smpl_model.idx_split_2_merged_vec();
240 let cols = vec![0, 1, 2];
241 let vertex_offsets = vertex_offsets_merged.gather(mapping, &cols).into_nalgebra();
242 current_body.positions = Some(current_body.positions.as_ref().unwrap() + vertex_offsets);
243 }
244 }
245 #[allow(clippy::cast_precision_loss)]
246 if scene.world.has::<Animation>(entity).unwrap() {
247 let scene_anim = scene.get_resource::<&SceneAnimation>().unwrap();
248 nr_frames = scene_anim.num_frames;
249 info!("Processing Animation for body {body_idx:?}");
250 let anim = scene.get_comp::<&Animation>(&entity).unwrap();
251 let mut current_body_rotations = nd::Array3::<f32>::zeros((joint_count, nr_frames, 3));
252 let mut current_body_translations = nd::Array2::<f32>::zeros((nr_frames, 3));
253 let mut current_body_scales = nd::Array2::<f32>::zeros((nr_frames, 3));
254 let mut current_per_frame_blend_weights = nd::Array2::<f32>::zeros((nr_frames, num_total_blendshapes));
255 if should_export_posedirs || should_export_exprdirs {
256 let mut full_morph_targets = nd::Array3::<f32>::zeros((num_total_blendshapes, vertex_count, 3));
257 let mut running_idx_morph_target = 0;
258 if should_export_posedirs {
259 let mut pose_morph_targets = nd::Array3::<f32>::zeros((metadata.num_pose_blend_shapes + 1, vertex_count, 3));
260 let nr_elem_merged = smpl_model.get_pose_dirs().dims()[0] / 3;
261 let pose_dirs_merged = smpl_model
262 .get_pose_dirs()
263 .to_ndarray()
264 .into_shape_with_order((nr_elem_merged, 3, metadata.num_pose_blend_shapes))
265 .unwrap();
266 let mapping = smpl_model.idx_split_2_merged_vec();
267 let cols = vec![0, 1, 2];
268 let depth = (0..metadata.num_pose_blend_shapes).collect::<Vec<_>>().into_boxed_slice();
269 let pose_blend_shapes = pose_dirs_merged
270 .gather(mapping, &cols, &depth)
271 .into_shape_with_order((vertex_count, 3, metadata.num_pose_blend_shapes))
272 .unwrap()
273 .permuted_axes([2, 0, 1]);
274 let pi = nd::Array1::<f32>::from_elem(metadata.num_pose_blend_shapes, -PI);
275 let pi_array = pi.insert_axis(nd::Axis(1)).insert_axis(nd::Axis(2));
276 assert_eq!(pose_blend_shapes.shape()[0], pi_array.len());
277 let template_offset = (pose_blend_shapes.clone() * &pi_array).sum_axis(nd::Axis(0));
278 pose_morph_targets
279 .slice_mut(s![0..metadata.num_pose_blend_shapes, .., ..])
280 .assign(&(pose_blend_shapes));
281 pose_morph_targets
282 .slice_mut(s![metadata.num_pose_blend_shapes, .., ..])
283 .assign(&template_offset);
284 #[allow(clippy::range_plus_one)]
285 full_morph_targets
286 .slice_mut(s![
287 running_idx_morph_target..running_idx_morph_target + metadata.num_pose_blend_shapes + 1,
288 ..,
289 ..
290 ])
291 .assign(&pose_morph_targets);
292 running_idx_morph_target += metadata.num_pose_blend_shapes + 1;
293 gltf_codec.num_pose_morph_targets = metadata.num_pose_blend_shapes + 1;
294 }
295 #[allow(unused_assignments)]
296 if should_export_exprdirs {
297 if let Some(expr_dirs) = smpl_model.get_expression_dirs() {
298 let mut expression_morph_targets = nd::Array3::<f32>::zeros((num_expression_blend_shapes + 1, vertex_count, 3));
299 let nr_elem_merged = expr_dirs.dims()[0] / 3;
300 let expression_dirs_merged = expr_dirs
301 .to_ndarray()
302 .into_shape_with_order((nr_elem_merged, 3, num_expression_blend_shapes))
303 .unwrap();
304 let mapping = smpl_model.idx_split_2_merged_vec();
305 let cols = vec![0, 1, 2];
306 let depth = (0..num_expression_blend_shapes).collect::<Vec<_>>().into_boxed_slice();
307 let expression_dirs_split = expression_dirs_merged
308 .gather(mapping, &cols, &depth)
309 .into_shape_with_order((vertex_count, 3, num_expression_blend_shapes))
310 .unwrap()
311 .permuted_axes([2, 0, 1]);
312 let expression_bounds = nd::Array1::<f32>::from_elem(num_expression_blend_shapes, -7.0);
313 let expression_bounds_array = expression_bounds.insert_axis(nd::Axis(1)).insert_axis(nd::Axis(2));
314 assert_eq!(expression_dirs_split.shape()[0], expression_bounds_array.len());
315 let template_offset = (expression_dirs_split.clone() * &expression_bounds_array).sum_axis(nd::Axis(0));
316 expression_morph_targets
317 .slice_mut(s![0..num_expression_blend_shapes, .., ..])
318 .assign(&(expression_dirs_split));
319 expression_morph_targets
320 .slice_mut(s![num_expression_blend_shapes, .., ..])
321 .assign(&template_offset);
322 #[allow(clippy::range_plus_one)]
323 full_morph_targets
324 .slice_mut(s![
325 running_idx_morph_target..running_idx_morph_target + num_expression_blend_shapes + 1,
326 ..,
327 ..
328 ])
329 .assign(&expression_morph_targets);
330 running_idx_morph_target += num_expression_blend_shapes + 1;
331 gltf_codec.num_expression_morph_targets = num_expression_blend_shapes + 1;
332 }
333 }
334 gltf_codec.morph_targets = Some(full_morph_targets);
335 }
336 for global_frame_idx in 0..nr_frames {
337 if global_frame_idx < anim.start_offset || global_frame_idx > anim.start_offset + anim.num_animation_frames() {
338 continue;
339 }
340 let mut local_frame_idx = global_frame_idx - anim.start_offset;
341 if global_frame_idx == (anim.start_offset + anim.num_animation_frames()) {
342 local_frame_idx -= 1;
343 }
344 let mut pose = anim.get_pose_at_idx(local_frame_idx);
345 let pose_remap = PoseRemap::new(pose.smpl_type, smpl_params.smpl_type);
346 pose = pose_remap.remap(&pose);
347 if let Ok(ref pose_mask) = scene.get_comp::<&PoseOverride>(&entity) {
348 let mut new_pose_mask = PoseOverride::clone(pose_mask);
349 pose.apply_mask(&mut new_pose_mask);
350 }
351 if let Ok(ref pose_retarget) = scene.get_comp::<&RetargetPoseYShift>(&entity) {
352 let mut pose_retarget_local = RetargetPoseYShift::clone(pose_retarget);
353 pose_retarget_local.apply(&mut pose);
354 }
355 current_body_rotations
356 .slice_mut(s![.., global_frame_idx, ..])
357 .assign(&pose.joint_poses.to_ndarray());
358 let mut skeleton_root_translation = pose.global_trans.to_ndarray().to_owned();
359 let root_translation = smpl_output.joints.to_ndarray().slice(s![0, ..]).to_owned();
360 skeleton_root_translation = skeleton_root_translation + root_translation;
361 current_body_translations
362 .slice_mut(s![global_frame_idx, ..])
363 .assign(&skeleton_root_translation);
364 if global_frame_idx < (anim.start_offset + anim.num_animation_frames()) {
365 current_body_scales.slice_mut(s![global_frame_idx, ..]).assign(&nd::Array1::ones(3));
366 }
367 let mut running_idx_morph_target = 0;
368 if should_export_posedirs {
369 let pose_blend_weights = &smpl_model.compute_pose_feature(&pose);
370 current_per_frame_blend_weights
371 .slice_mut(s![global_frame_idx, 0..metadata.num_pose_blend_shapes])
372 .assign(&pose_blend_weights.to_ndarray());
373 if global_frame_idx == (anim.start_offset + anim.num_animation_frames()) {
374 current_per_frame_blend_weights
375 .slice_mut(s![global_frame_idx..nr_frames, 0..metadata.num_pose_blend_shapes])
376 .assign(&pose_blend_weights.to_ndarray());
377 }
378 running_idx_morph_target += metadata.num_pose_blend_shapes + 1;
379 }
380 #[allow(unused_assignments)]
381 if should_export_exprdirs {
382 let expr_opt = anim.get_expression_at_idx(local_frame_idx);
383 if let Some(expr) = expr_opt.as_ref() {
384 let expr_nd = expr.expr_coeffs.to_ndarray();
385 let max_nr_expr_coeffs = num_expression_blend_shapes.min(expr_nd.len());
386 let expr_coeffs = expr_nd.slice(s![0..max_nr_expr_coeffs]);
387 current_per_frame_blend_weights
388 .slice_mut(s![
389 global_frame_idx,
390 running_idx_morph_target..running_idx_morph_target + max_nr_expr_coeffs
391 ])
392 .assign(&expr_coeffs);
393 }
394 running_idx_morph_target += num_expression_blend_shapes + 1;
395 }
396 }
397 current_body.body_scales = Some(current_body_scales);
398 current_body.body_translations = Some(current_body_translations);
399 current_body.body_rotations = Some(current_body_rotations);
400 if should_export_posedirs {
401 current_per_frame_blend_weights
402 .slice_mut(s![.., metadata.num_pose_blend_shapes])
403 .assign(&nd::Array1::<f32>::from_elem(nr_frames, 0.0));
404 }
405 if should_export_exprdirs {
406 current_per_frame_blend_weights
407 .slice_mut(s![.., metadata.num_pose_blend_shapes + num_expression_blend_shapes + 1])
408 .assign(&nd::Array1::<f32>::from_elem(nr_frames, 0.0));
409 }
410 if should_export_posedirs || should_export_exprdirs {
411 current_body.per_frame_blend_weights = Some(current_per_frame_blend_weights);
412 }
413 }
414 current_body.default_joint_translations = Some(smpl_joints);
415 gltf_codec.per_body_data.push(current_body);
416 }
417 let mut query_props = scene
418 .world
419 .query::<(&Verts, &Faces, &TransformSequence, Option<&UVs>, Option<&DiffuseImg>, &Name)>();
420 for (_entity, (verts, faces, transform_sequence, uv, diffuse_img, _name)) in query_props.iter() {
421 let scales_vec3 = transform_sequence.scales.iter().flat_map(|s| vec![*s, *s, *s]).collect::<Vec<f32>>();
422 let scales_ndarray3 = nd::Array2::from_shape_vec((transform_sequence.num_frames(), 3), scales_vec3).unwrap();
423 let prop_data = PropData {
424 positions: verts.0.to_ndarray().into_nalgebra(),
425 faces: faces.0.to_ndarray().into_nalgebra(),
426 normals: None,
427 uvs: uv.map(|uv| uv.0.to_ndarray().into_nalgebra()),
428 diffuse_texture: diffuse_img.and_then(|diffuse_img| {
429 diffuse_img
430 .generic_img
431 .cpu_img
432 .as_ref()
433 .map(|img| get_image(img, false, options.max_texture_size))
434 }),
435 metalness_texture: None,
436 roughness_texture: None,
437 normals_texture: None,
438 translations: transform_sequence.translations.clone(),
439 rotations: transform_sequence.rotations.clone().insert_axis(nd::Axis(0)),
440 scales: scales_ndarray3.clone(),
441 default_translation: nd::Array2::<f32>::zeros([1, 3]),
442 default_joint_poses: nd::Array2::<f32>::zeros([1, 3]),
443 };
444 gltf_codec.props.push(prop_data);
445 }
446 info!(
447 "Writing {} body scene to GltfCodec: Took {} seconds for {} frames",
448 num_bodies,
449 now.elapsed().as_secs(),
450 nr_frames
451 );
452 gltf_codec
453}