mkv_strand_reader/
lib.rs

1// Copyright (C) The Strand-Braid Authors
2// SPDX-License-Identifier: MIT OR Apache-2.0
3
4use mkv_parser_kit::{BoxData, EbmlElement, Tag, ebml_parse};
5
6const STRAND_MKV_FILENAME_TEMPLATE: &str = "movie%Y%m%d_%H%M%S.%f";
7
8#[derive(thiserror::Error, Debug)]
9pub enum Error {
10    #[error("parse error: {source}")]
11    Parser {
12        #[from]
13        source: mkv_parser_kit::Error,
14    },
15    #[error("missing data: {missing}")]
16    MissingData { missing: String },
17    #[error("no cluster timestamp")]
18    NoClusterTimestamp,
19    #[error("could not determine filename")]
20    CouldNotDetermineFilename,
21    #[error("filename is not valid UTF-8")]
22    FilenameNotUtf8,
23    #[error("filename is could not be parsed to datetime")]
24    DatetimeInFilenameNotParsedCorrectly,
25}
26
27fn missing(what: &str) -> Error {
28    Error::MissingData {
29        missing: what.to_string(),
30    }
31}
32
33pub type Result<T> = std::result::Result<T, Error>;
34
35#[derive(Debug, Clone)]
36pub struct BlockData {
37    pub pts: std::time::Duration,
38    pub start_idx: u64,
39    pub size: usize,
40    pub is_keyframe: bool,
41}
42
43#[derive(Debug)]
44pub struct StrandCamMkvMetadata {
45    pub creation_time: chrono::DateTime<chrono::FixedOffset>,
46    pub camera_name: Option<String>,
47    pub gamma: Option<f32>,
48    pub writing_app: String,
49}
50
51#[derive(Debug)]
52pub struct ParsedStrandCamMkv {
53    pub width: u32,
54    pub height: u32,
55    pub timestep_nanos: u32,
56    pub metadata: StrandCamMkvMetadata,
57    pub block_data: Vec<BlockData>,
58    pub uncompressed_fourcc: Option<String>,
59    pub codec: String,
60}
61
62impl TryFrom<Accum> for ParsedStrandCamMkv {
63    type Error = Error;
64    fn try_from(a: Accum) -> Result<Self> {
65        let width = a.width.ok_or(missing("width"))?;
66        let height = a.height.ok_or(missing("height"))?;
67        let metadata = StrandCamMkvMetadata {
68            creation_time: a.creation_time.ok_or(missing("creation time"))?,
69            gamma: a.gamma,
70            // Strand Cam only ever saved the camera name to the MKV title field.
71            camera_name: a.title,
72            writing_app: a.writing_app.ok_or(missing("writing app"))?,
73        };
74        Ok(Self {
75            width,
76            height,
77            timestep_nanos: a.timestep_nanos.ok_or(missing("timestep"))?,
78            block_data: a.block_data,
79            codec: a.codec.ok_or(missing("codec"))?,
80            metadata,
81            uncompressed_fourcc: a.uncompressed_fourcc,
82        })
83    }
84}
85
86#[derive(Debug, Default)]
87struct Accum {
88    n_segments: u8,
89    width: Option<u32>,
90    height: Option<u32>,
91    timestep_nanos: Option<u32>,
92    creation_time: Option<chrono::DateTime<chrono::FixedOffset>>,
93    title: Option<String>,
94    gamma: Option<f32>,
95    writing_app: Option<String>,
96    block_data: Vec<BlockData>,
97    /// The timestamp for the entire cluster.
98    // segment_cluster_timestamp: Option<chrono::DateTime<chrono::Utc>>,
99    pts: Option<std::time::Duration>,
100    uncompressed_fourcc: Option<String>,
101    codec: Option<String>,
102}
103
104fn line_summary(e: &EbmlElement) -> String {
105    let name = format!("{:?}", e.tag());
106    let position = e.position();
107    let full_size = e.full_size();
108    let data_size = e.data_size();
109    format!("{name} at {position} size {full_size} data size {data_size}")
110}
111
112const PRINT_ALL: bool = false;
113
114fn do_parse(
115    element: &EbmlElement,
116    depth: u8,
117    accum: &mut Accum,
118    tag_path: &[Tag],
119    verbose: bool,
120    filename: Option<&str>,
121) -> Result<()> {
122    let verbose_prefix = if verbose {
123        let mut prefix = String::new();
124        for idx in 0..depth {
125            if idx == 0 {
126                prefix.push('|');
127            } else {
128                prefix.push(' ');
129            }
130        }
131        Some(prefix)
132    } else {
133        None
134    };
135    if let Some(prefix) = &verbose_prefix
136        && PRINT_ALL
137    {
138        println!("{}+ {}", prefix, line_summary(element));
139    }
140    for child in element.children().iter() {
141        let mut child_tag_path = tag_path.to_vec();
142        child_tag_path.push(child.tag());
143        do_parse(child, depth + 1, accum, &child_tag_path, verbose, filename)?;
144    }
145
146    if let Some(prefix) = &verbose_prefix
147        && let Some(bd) = &element.box_data()
148    {
149        if !PRINT_ALL {
150            println!("{}+ {}", prefix, line_summary(element));
151        }
152        println!("{prefix}+           {bd:?}");
153    }
154
155    match tag_path {
156        [Tag::Segment] => {
157            accum.n_segments += 1;
158            assert_eq!(accum.n_segments, 1); // no support for > 1 segment
159        }
160        [Tag::Segment, Tag::Info, Tag::DateUTC] => {
161            if let Some(BoxData::DateTime(creation_time_utc)) = element.box_data() {
162                let creation_time = infer_timezone(creation_time_utc, filename)?;
163                accum.creation_time = Some(creation_time);
164            } else {
165                panic!("need DateUTC");
166            }
167        }
168        [Tag::Segment, Tag::Info, Tag::TimestampScale] => {
169            accum.timestep_nanos = Some(get_uint(element));
170        }
171        [Tag::Segment, Tag::Info, Tag::WritingApp] => {
172            accum.writing_app = Some(get_string(element));
173        }
174        [Tag::Segment, Tag::Info, Tag::Title] => {
175            accum.title = Some(get_string(element));
176        }
177        [Tag::Segment, Tag::Info, Tag::GammaValue] => {
178            accum.gamma = Some(get_float(element));
179        }
180        [Tag::Segment, Tag::Cluster] => {
181            // accum.segment_cluster_timestamp = None;
182            accum.pts = None;
183        }
184        [Tag::Segment, Tag::Cluster, Tag::Timestamp] => {
185            // assert!(accum.segment_cluster_timestamp.is_none());
186            assert!(accum.pts.is_none());
187            // convert to u64 because these numbers can get big.
188            let n_timesteps: u64 = get_uint(element).into();
189            let timestep_nanos: u64 = accum.timestep_nanos.unwrap().into();
190            let pts_total_nanos = n_timesteps * timestep_nanos;
191
192            let pts = std::time::Duration::from_nanos(pts_total_nanos);
193            // let timestamp = accum.creation_time.unwrap() + pts;
194            // accum.segment_cluster_timestamp = Some(timestamp);
195            accum.pts = Some(pts);
196        }
197        [Tag::Segment, Tag::Cluster, Tag::SimpleBlock] => {
198            // if let Some(cluster_timestamp) = accum.segment_cluster_timestamp {
199            if let Some(cluster_pts) = accum.pts {
200                let x = if let Some(BoxData::SimpleBlockData(block_data)) = &element.box_data() {
201                    let n_timesteps: u64 = block_data.timestamp.try_into().unwrap();
202                    let timestep_nanos: u64 = accum.timestep_nanos.unwrap().into();
203                    let cluster_offset_nanos = n_timesteps * timestep_nanos;
204                    let cluster_offset = std::time::Duration::from_nanos(cluster_offset_nanos);
205                    let pts = cluster_pts + cluster_offset;
206                    BlockData {
207                        pts,
208                        start_idx: block_data.start,
209                        is_keyframe: block_data.is_keyframe,
210                        size: block_data.size.try_into().unwrap(),
211                    }
212                } else {
213                    panic!("expected UncompressedFourCC in {:?}", element.tag());
214                };
215                accum.block_data.push(x);
216            } else {
217                return Err(Error::NoClusterTimestamp);
218            }
219        }
220        [Tag::Segment, Tag::Tracks, Tag::TrackEntry, Tag::CodecID] => {
221            accum.codec = Some(get_ascii_string(element));
222        }
223        [
224            Tag::Segment,
225            Tag::Tracks,
226            Tag::TrackEntry,
227            Tag::Video,
228            Tag::UncompressedFourCC,
229        ] => {
230            accum.uncompressed_fourcc =
231                if let Some(BoxData::UncompressedFourCC(s)) = &element.box_data() {
232                    Some(s.clone())
233                } else {
234                    panic!("expected UncompressedFourCC in {:?}", element.tag());
235                }
236        }
237        [
238            Tag::Segment,
239            Tag::Tracks,
240            Tag::TrackEntry,
241            Tag::Video,
242            Tag::PixelWidth,
243        ] => {
244            accum.width = Some(get_uint(element));
245        }
246        [
247            Tag::Segment,
248            Tag::Tracks,
249            Tag::TrackEntry,
250            Tag::Video,
251            Tag::PixelHeight,
252        ] => {
253            accum.height = Some(get_uint(element));
254        }
255        _ => {
256            // dbg!(&tag_path);
257        }
258    }
259    Ok(())
260}
261
262fn get_ascii_string(element: &EbmlElement) -> String {
263    if let Some(BoxData::AsciiString(s)) = &element.box_data() {
264        s.clone()
265    } else {
266        panic!("expected ascii string in {:?}", element.tag());
267    }
268}
269
270fn get_string(element: &EbmlElement) -> String {
271    if let Some(BoxData::String(s)) = &element.box_data() {
272        s.clone()
273    } else {
274        panic!("expected string in {:?}", element.tag());
275    }
276}
277
278fn get_float(element: &EbmlElement) -> f32 {
279    if let Some(BoxData::Float(f)) = &element.box_data() {
280        *f
281    } else {
282        panic!("expected string");
283    }
284}
285
286fn get_uint(element: &EbmlElement) -> u32 {
287    if let Some(BoxData::UnsignedInt(v)) = &element.box_data() {
288        *v
289    } else {
290        panic!("expected string");
291    }
292}
293
294pub fn parse_strand_cam_mkv<R, P>(
295    rdr: R,
296    verbose: bool,
297    path: Option<P>,
298) -> Result<(ParsedStrandCamMkv, R)>
299where
300    R: std::io::Read + std::io::Seek,
301    P: AsRef<std::path::Path>,
302{
303    let filename = path.map(|p| format!("{}", p.as_ref().display()));
304    let (parsed, rdr) = ebml_parse(rdr)?;
305    // println!("parsing done ---------- ");
306    let mut accum = Accum::default();
307    for element in parsed.iter() {
308        do_parse(
309            element,
310            0,
311            &mut accum,
312            &[element.tag()],
313            verbose,
314            filename.as_deref(),
315        )?;
316    }
317    Ok((accum.try_into()?, rdr))
318}
319
320/// Attempt to parse a filename to set the correct timezone.
321///
322/// Strand Camera (and its predecessors) have saved times in UTC but set
323/// filenames based on the local time. This allows us to infer the timezone
324/// offset from UTC at the time the file was recorded without losing the UTC
325/// time. Here attempt to do this.
326pub fn infer_timezone(
327    creation_time_utc: &chrono::DateTime<chrono::Utc>,
328    filename: Option<&str>,
329) -> Result<chrono::DateTime<chrono::FixedOffset>> {
330    let zero_offset = chrono::FixedOffset::east_opt(0).unwrap();
331    let mut creation_time = creation_time_utc.with_timezone(&zero_offset);
332    if let Some(filename) = filename {
333        let path_buf = std::path::PathBuf::from(filename);
334        let filename = path_buf
335            .file_name()
336            .ok_or(Error::CouldNotDetermineFilename)?;
337        let filename = filename.to_str().ok_or(Error::FilenameNotUtf8)?;
338
339        // Optimisitcally parse filename according to default for
340        // `StrandCamArgs::mkv_filename_template`.
341        let underscores: Vec<_> = filename.split('_').collect();
342        if underscores.len() > 2 {
343            let joined = format!("{}_{}", underscores[0], underscores[1]);
344
345            match chrono::NaiveDateTime::parse_from_str(&joined, STRAND_MKV_FILENAME_TEMPLATE) {
346                Ok(naive) => {
347                    let offset_dur = naive - creation_time_utc.naive_utc();
348                    let offset =
349                        chrono::FixedOffset::east_opt(offset_dur.num_seconds().try_into().unwrap())
350                            .unwrap();
351                    creation_time = creation_time_utc.with_timezone(&offset);
352                    let test_str = creation_time
353                        .format(STRAND_MKV_FILENAME_TEMPLATE)
354                        .to_string();
355                    // There was a datetime, but we could not parse it
356                    // correctly. It's not clear to me if this could ever
357                    // happen.
358                    if joined != test_str {
359                        return Err(Error::DatetimeInFilenameNotParsedCorrectly);
360                    }
361                }
362                Err(_e) => {}
363            }
364        }
365    };
366    Ok(creation_time)
367}
mkv_strand_reader/lib.rs

mkv_strand_reader/
lib.rs
