smpl_core/codec/
scene.rs

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

smpl_core/codec/
scene.rs
