Skip to main content

smpl_core/codec/
scene.rs

1use super::codec::SmplCodec;
2use crate::common::transform_sequence::TransformSequence;
3use base64;
4use gloss_renderer::components::ProjectionWithFov;
5use gltf_json::{validation::Checked::Valid, Root, Value};
6use log::info;
7use ndarray as nd;
8use smpl_utils::log;
9use std::{
10    fs::{self, File},
11    io::Read,
12    path::Path,
13};
14/// The ``McsCodec`` contains all of the contents of an ``.mcs`` file
15#[derive(Clone)]
16pub struct McsCodec {
17    pub num_frames: usize,
18    pub frame_rate: Option<f32>,
19    pub smpl_bodies: Vec<SmplBody>,
20    pub smpl_camera: Option<SmplCamera>,
21}
22/// ``SmplBody`` holds the contents of the ``.smpl`` file along with the frame presence
23#[derive(Clone)]
24pub struct SmplBody {
25    pub frame_presence: Vec<usize>,
26    pub codec: SmplCodec,
27}
28/// ``SmplCamera`` holds the information regarding the camera track
29#[derive(Clone)]
30pub struct SmplCamera {
31    pub projection: ProjectionWithFov,
32    pub transform_sequence: TransformSequence,
33}
34/// ``McsCodec`` for ``.mcs`` files
35#[allow(clippy::cast_possible_truncation)]
36impl McsCodec {
37    /// Load `McsCodec` from a GLTF file.
38    pub fn from_file(path: &str) -> Self {
39        let mut file = File::open(path).expect("Failed to open GLTF file");
40        let mut json_data = String::new();
41        file.read_to_string(&mut json_data).expect("Failed to read GLTF file");
42        let gltf: Root = serde_json::from_str(&json_data).expect("Failed to parse GLTF JSON");
43        Self::from_gltf(&gltf)
44    }
45    /// Parse a scene GLTF into ``McsCodec``
46    pub fn from_gltf(gltf: &Root) -> Self {
47        let mut num_frames = 0;
48        let mut smpl_bodies = Vec::new();
49        if let Some(scene) = gltf.scenes.first() {
50            if let Some(extensions) = &scene.extensions {
51                if let Some(extension_value) = extensions.others.get("MC_scene_description") {
52                    let extension = extension_value.as_object().expect("Expected extension to be an object");
53                    if let Some(nf) = extension.get("num_frames").and_then(gltf_json::Value::as_u64) {
54                        num_frames = nf as usize;
55                    }
56                    if let Some(smpl_bodies_data) = extension.get("smpl_bodies").and_then(|v| v.as_array()) {
57                        smpl_bodies = Self::extract_smpl_bodies(gltf, smpl_bodies_data);
58                    }
59                }
60            }
61            Self {
62                num_frames,
63                frame_rate: smpl_bodies.first().and_then(|b| b.codec.frame_rate),
64                smpl_bodies,
65                smpl_camera: Self::extract_gltf_camera(gltf),
66            }
67        } else {
68            panic!("Not able to find GLTF root! Check the GLTF file format!")
69        }
70    }
71    /// Extract SMPL bodies and their `frame_presence` from the GLTF extension.
72    fn extract_smpl_bodies(gltf: &Root, smpl_bodies_data: &[serde_json::Value]) -> Vec<SmplBody> {
73        smpl_bodies_data
74            .iter()
75            .filter_map(|smpl_body_data| {
76                smpl_body_data
77                    .get("bufferView")
78                    .and_then(gltf_json::Value::as_u64)
79                    .map(|buffer_view_index| {
80                        let buffer = Self::read_smpl_buffer(gltf, buffer_view_index as usize);
81                        let frame_presence = smpl_body_data
82                            .get("frame_presence")
83                            .and_then(|v| v.as_array())
84                            .map_or_else(Vec::new, |arr| arr.iter().filter_map(|v| v.as_u64().map(|n| n as usize)).collect());
85                        let codec = SmplCodec::from_buf(&buffer);
86                        SmplBody { frame_presence, codec }
87                    })
88            })
89            .collect()
90    }
91    /// Reads buffer data based on buffer index.
92    fn read_smpl_buffer(gltf: &Root, buffer_index: usize) -> Vec<u8> {
93        let buffer = &gltf.buffers[buffer_index];
94        buffer
95            .uri
96            .as_ref()
97            .and_then(|uri| {
98                if uri.starts_with("data:") {
99                    uri.split(',')
100                        .nth(1)
101                        .map(|encoded_data| base64::decode(encoded_data).expect("Failed to decode Base64 data"))
102                } else {
103                    panic!("The data buffers must not be separate files!")
104                }
105            })
106            .unwrap_or_default()
107    }
108    /// Extract camera track from `.mcs` file
109    fn extract_gltf_camera(gltf: &Root) -> Option<SmplCamera> {
110        if let Some(camera) = gltf.cameras.first() {
111            let (yfov, znear, zfar, aspect_ratio) = match camera.type_.unwrap() {
112                gltf_json::camera::Type::Perspective => (
113                    camera.perspective.as_ref().map_or(std::f32::consts::FRAC_PI_2, |p| p.yfov),
114                    camera.perspective.as_ref().map_or(0.1, |p| p.znear),
115                    camera.perspective.as_ref().map_or(1000.0, |p| p.zfar.unwrap_or(1000.0)),
116                    camera.perspective.as_ref().map_or(1.0, |p| p.aspect_ratio.unwrap_or(1.0)),
117                ),
118                gltf_json::camera::Type::Orthographic => {
119                    panic!("Orthographic camera not supported!")
120                }
121            };
122            if gltf.animations.is_empty() {
123                let mut static_translation = nd::Array2::<f32>::zeros((1, 3));
124                let mut static_rotation = nd::array![[0.0, 0.0, 0.0, 1.0]];
125                for node in &gltf.nodes {
126                    if node.camera.is_some() {
127                        if let Some(translation) = &node.translation {
128                            static_translation = nd::Array2::from_shape_vec((1, 3), vec![translation[0], translation[1], translation[2]]).unwrap();
129                        }
130                        if let Some(rotation) = &node.rotation {
131                            let quat = rotation.0;
132                            static_rotation = nd::Array2::from_shape_vec((1, 4), vec![quat[0], quat[1], quat[2], quat[3]]).unwrap();
133                        }
134                        break;
135                    }
136                }
137                return Some(SmplCamera {
138                    projection: ProjectionWithFov {
139                        fovy: yfov,
140                        near: znear,
141                        far: zfar,
142                        aspect_ratio,
143                    },
144                    transform_sequence: TransformSequence::new_from_quat_rot_trans(&static_rotation, &static_translation),
145                });
146            }
147            let mut per_frame_translations: Option<nd::Array2<f32>> = None;
148            let mut per_frame_rotations: Option<nd::Array2<f32>> = None;
149            for animation in &gltf.animations {
150                for channel in &animation.channels {
151                    let target = &channel.target;
152                    let Some(node) = gltf.nodes.get(target.node.value()) else { continue };
153                    if node.camera.is_none() {
154                        continue;
155                    }
156                    let Some(sampler) = animation.samplers.get(channel.sampler.value()) else {
157                        continue;
158                    };
159                    let Some(output_accessor) = gltf.accessors.get(sampler.output.value()) else {
160                        continue;
161                    };
162                    let Some(buffer_view) = output_accessor.buffer_view.as_ref().and_then(|bv| gltf.buffer_views.get(bv.value())) else {
163                        continue;
164                    };
165                    let Some(buffer) = gltf.buffers.get(buffer_view.buffer.value()) else {
166                        continue;
167                    };
168                    let Some(uri) = &buffer.uri else { continue };
169                    if !uri.starts_with("data:") {
170                        continue;
171                    }
172                    let encoded_data = uri.split(',').nth(1).expect("Invalid data URI");
173                    let buffer_data = base64::decode(encoded_data).expect("Failed to decode Base64 data");
174                    let start = buffer_view.byte_offset.map_or(0, |x| x.0 as usize);
175                    let length = buffer_view.byte_length.0 as usize;
176                    let data = &buffer_data[start..start + length];
177                    let floats: Vec<f32> = data.chunks(4).map(|b| f32::from_le_bytes([b[0], b[1], b[2], b[3]])).collect();
178                    if let Valid(path) = &target.path {
179                        match path {
180                            gltf_json::animation::Property::Translation => {
181                                let num_frames = floats.len() / 3;
182                                per_frame_translations = Some(nd::Array2::from_shape_vec((num_frames, 3), floats).unwrap());
183                            }
184                            gltf_json::animation::Property::Rotation => {
185                                let num_frames = floats.len() / 4;
186                                per_frame_rotations = Some(nd::Array2::from_shape_vec((num_frames, 4), floats).unwrap());
187                            }
188                            _ => {}
189                        }
190                    }
191                }
192            }
193            Some(SmplCamera {
194                projection: ProjectionWithFov {
195                    fovy: yfov,
196                    near: znear,
197                    far: zfar,
198                    aspect_ratio,
199                },
200                transform_sequence: TransformSequence::new_from_quat_rot_trans(
201                    &per_frame_rotations.unwrap_or(nd::array![[0.0, 0.0, 0.0, 1.0]]),
202                    &per_frame_translations.unwrap_or(nd::Array2::<f32>::zeros((1, 3))),
203                ),
204            })
205        } else {
206            None
207        }
208    }
209    /// Export `McsCodec` to an MCS file
210    pub fn to_file(&self, path: &str) {
211        let parent_path = Path::new(path).parent();
212        let file_name = Path::new(path).file_name();
213        let Some(parent_path) = parent_path else {
214            log!("Error: Exporting MCS - Something wrong with the path: {}", path);
215            return;
216        };
217        if !parent_path.exists() {
218            let _ = fs::create_dir_all(parent_path);
219        }
220        let Some(_) = file_name else {
221            log!("Error: Exporting MCS - no file name found: {}", path);
222            return;
223        };
224        let gltf_json = self.to_gltf_json();
225        std::fs::write(path, gltf_json).expect("Failed to write MCS file");
226        info!("Exported MCS file to: {path}");
227    }
228    /// Create the base empty Mcs structure
229    pub fn create_gltf_structure(&self) -> Value {
230        let mut gltf_json = serde_json::json!(
231            { "asset" : { "version" : "2.0", "generator" : "smpl-rs McsCodec Exporter" },
232            "scene" : 0, "scenes" : [{ "nodes" : [0], "extensions" : {
233            "MC_scene_description" : { "num_frames" : self.num_frames, "smpl_bodies" : []
234            } } }], "buffers" : [], "bufferViews" : [], "accessors" : [], "animations" :
235            [], "extensionsUsed" : ["MC_scene_description"] }
236        );
237        if let Some(smpl_camera) = &self.smpl_camera {
238            gltf_json["nodes"] = serde_json::json!(
239                [{ "name" : "RootNode", "children" : [1] }, { "name" : "AnimatedCamera",
240                "camera" : 0, "translation" : [0.0, 0.0, 0.0], "rotation" : [0.0, 0.0,
241                0.0, 1.0] }]
242            );
243            gltf_json["cameras"] = serde_json::json!(
244                [{ "type" : "perspective", "perspective" : { "yfov" : smpl_camera
245                .projection.fovy, "znear" : smpl_camera.projection.near, "zfar" :
246                smpl_camera.projection.far, "aspectRatio" : smpl_camera.projection
247                .aspect_ratio } }]
248            );
249        } else {
250            gltf_json["nodes"] = serde_json::json!([{ "name" : "RootNode" }]);
251        }
252        gltf_json
253    }
254    /// Add SMPL buffers to the Mcs GLTF JSON
255    pub fn add_smpl_buffers_to_gltf(&self, gltf_json: &mut Value) {
256        for (body_index, smpl_body) in self.smpl_bodies.iter().enumerate() {
257            let buffer_data = smpl_body.codec.to_buf();
258            let buffer_base64 = base64::encode(&buffer_data);
259            gltf_json["buffers"].as_array_mut().unwrap().push(serde_json::json!(
260                { "byteLength" : buffer_data.len(), "uri" :
261                format!("data:application/octet-stream;base64,{}", buffer_base64)
262                }
263            ));
264            gltf_json["bufferViews"].as_array_mut().unwrap().push(serde_json::json!(
265                { "buffer" : body_index, "byteOffset" : 0, "byteLength" :
266                buffer_data.len() }
267            ));
268            gltf_json["scenes"][0]["extensions"]["MC_scene_description"]["smpl_bodies"]
269                .as_array_mut()
270                .unwrap()
271                .push(serde_json::json!(
272                    { "frame_presence" : smpl_body.frame_presence, "bufferView" :
273                    body_index }
274                ));
275        }
276    }
277    /// Add camera animation to the Mcs GLTF JSON
278    #[allow(clippy::too_many_lines)]
279    pub fn add_camera_animation(&self, gltf_json: &mut Value) {
280        let buffers_start_idx = self.smpl_bodies.len();
281        let num_frames = self.num_frames;
282        let fps = self.frame_rate.unwrap_or(30.0);
283        #[allow(clippy::cast_precision_loss)]
284        let times: Vec<f32> = (0..num_frames).map(|i| i as f32 / fps).collect();
285        let time_bytes = times.iter().flat_map(|f| f.to_le_bytes()).collect::<Vec<u8>>();
286        let camera_positions = self.smpl_camera.as_ref().unwrap().transform_sequence.translations.clone();
287        let camera_rotations = self.smpl_camera.as_ref().unwrap().transform_sequence.get_rotations_as_quaternions();
288        if camera_positions.dim().0 == 1 {
289            if let Some(node) = gltf_json["nodes"].as_array_mut().and_then(|nodes| nodes.get_mut(1)) {
290                node["translation"] = serde_json::json!([camera_positions[[0, 0]], camera_positions[[0, 1]], camera_positions[[0, 2]]]);
291                node["rotation"] = serde_json::json!([
292                    camera_rotations[[0, 0]],
293                    camera_rotations[[0, 1]],
294                    camera_rotations[[0, 2]],
295                    camera_rotations[[0, 3]]
296                ]);
297            }
298            return;
299        }
300        let translation_bytes = camera_positions.iter().flat_map(|f| f.to_le_bytes()).collect::<Vec<u8>>();
301        let rotation_bytes = camera_rotations.iter().flat_map(|f| f.to_le_bytes()).collect::<Vec<u8>>();
302        gltf_json["buffers"].as_array_mut().unwrap().extend([
303            serde_json::json!(
304                { "byteLength" : time_bytes.len(), "uri" :
305                format!("data:application/octet-stream;base64,{}", base64::encode(&
306                time_bytes)) }
307            ),
308            serde_json::json!(
309                { "byteLength" : translation_bytes.len(), "uri" :
310                format!("data:application/octet-stream;base64,{}", base64::encode(&
311                translation_bytes)) }
312            ),
313            serde_json::json!(
314                { "byteLength" : rotation_bytes.len(), "uri" :
315                format!("data:application/octet-stream;base64,{}", base64::encode(&
316                rotation_bytes)) }
317            ),
318        ]);
319        gltf_json["bufferViews"].as_array_mut().unwrap().extend([
320            serde_json::json!(
321                { "name" : "TimeBufferView", "buffer" : buffers_start_idx,
322                "byteOffset" : 0, "byteLength" : time_bytes.len() }
323            ),
324            serde_json::json!(
325                { "name" : "camera_track_translations_buffer_view", "buffer" :
326                buffers_start_idx + 1, "byteOffset" : 0, "byteLength" :
327                translation_bytes.len() }
328            ),
329            serde_json::json!(
330                { "name" : "camera_track_rotations_buffer_view", "buffer" :
331                buffers_start_idx + 2, "byteOffset" : 0, "byteLength" :
332                rotation_bytes.len() }
333            ),
334        ]);
335        let buffer_views_len = gltf_json["bufferViews"].as_array().unwrap().len();
336        gltf_json["accessors"].as_array_mut().unwrap().extend([
337            serde_json::json!(
338                { "name" : "TimeAccessor", "bufferView" : buffer_views_len - 3,
339                "componentType" : 5126, "count" : num_frames, "type" : "SCALAR",
340                "min" : [times[0]], "max" : [times[num_frames - 1]] }
341            ),
342            serde_json::json!(
343                { "name" : "camera_track_translations_accessor", "bufferView" :
344                buffer_views_len - 2, "componentType" : 5126, "count" : num_frames,
345                "type" : "VEC3" }
346            ),
347            serde_json::json!(
348                { "name" : "camera_track_rotations_accessor", "bufferView" :
349                buffer_views_len - 1, "componentType" : 5126, "count" : num_frames,
350                "type" : "VEC4" }
351            ),
352        ]);
353        let accessors_len = gltf_json["accessors"].as_array().unwrap().len();
354        gltf_json["animations"].as_array_mut().unwrap().push(serde_json::json!(
355            { "channels" : [{ "sampler" : 0, "target" : { "node" : 1, "path" :
356            "translation" } }, { "sampler" : 1, "target" : { "node" : 1, "path" :
357            "rotation" } }], "samplers" : [{ "input" : accessors_len - 3,
358            "interpolation" : "LINEAR", "output" : accessors_len - 2 }, { "input"
359            : accessors_len - 3, "interpolation" : "LINEAR", "output" :
360            accessors_len - 1 }] }
361        ));
362    }
363    /// Convert `McsCodec` to GLTF JSON string
364    pub fn to_gltf_json(&self) -> String {
365        let mut gltf_json = self.create_gltf_structure();
366        self.add_smpl_buffers_to_gltf(&mut gltf_json);
367        if self.smpl_camera.is_some() {
368            self.add_camera_animation(&mut gltf_json);
369        }
370        serde_json::to_string_pretty(&gltf_json).unwrap()
371    }
372}