frame_source/
strand_cam_mkv_source.rs

1// Copyright (C) The Strand-Braid Authors
2// SPDX-License-Identifier: MIT OR Apache-2.0
3
4use std::{
5    io::{BufReader, Read, Seek},
6    path::Path,
7};
8
9use mkv_strand_reader::ParsedStrandCamMkv;
10
11#[cfg(feature = "openh264")]
12use machine_vision_formats::owned::OImage;
13
14use super::*;
15
16#[derive(thiserror::Error, Debug)]
17pub enum StrandMkvSourceError {
18    #[error("cannot skip frames without decoding H264")]
19    CannotSkipWithoutDecodingH264,
20    #[error("could not decode single frame with openh264")]
21    CouldNotDecodeSingleFrameWithOpenH264,
22    #[error("Uncompressed MKV with fourcc '{0}' unsupported")]
23    UnsupportedFourcc(String),
24    #[error("Unsupported codec '{0}'")]
25    UnsupportedCodec(String),
26    #[error(
27        "Support for {timestamp_source:?} timestamp source not (yet) implemented for Strand Cam MKV files."
28    )]
29    UnimplTsSource { timestamp_source: TimestampSource },
30    #[error("unexpected image data")]
31    UnexpectedImageData,
32}
33
34// NAL unit start for b"MISPmicrosectime":
35const PRECISION_TIME_NALU_START: &[u8] = &[
36    0x00, 0x00, 0x00, 0x01, 0x06, 0x05, 28, b'M', b'I', b'S', b'P', b'm', b'i', b'c', b'r', b'o',
37    b's', b'e', b'c', b't', b'i', b'm', b'e',
38];
39
40/// An MKV file saved by Strand Camera.
41///
42/// Note that this is not a general purpose MKV file converter but is specific
43/// to MKV files which have been saved by Strand Camera.
44pub struct StrandCamMkvSource<R: Read + Seek> {
45    rdr: R,
46    pub parsed: ParsedStrandCamMkv,
47    src_format: Format,
48    is_uncompressed: bool,
49    h264_decoder_state: Option<crate::opt_openh264_decoder::DecoderType>,
50    keyframes_cache: Option<Vec<usize>>,
51}
52
53impl<R: Read + Seek> FrameDataSource for StrandCamMkvSource<R> {
54    fn width(&self) -> u32 {
55        self.parsed.width
56    }
57    fn height(&self) -> u32 {
58        self.parsed.height
59    }
60    fn camera_name(&self) -> Option<&str> {
61        self.parsed
62            .metadata
63            .camera_name
64            .as_ref()
65            .map(|x| x.as_ref())
66    }
67    fn gamma(&self) -> Option<f32> {
68        self.parsed.metadata.gamma
69    }
70    fn frame0_time(&self) -> Option<chrono::DateTime<chrono::FixedOffset>> {
71        Some(self.parsed.metadata.creation_time)
72    }
73    fn average_framerate(&self) -> Option<f64> {
74        None
75    }
76    fn skip_n_frames(&mut self, n_frames: usize) -> Result<()> {
77        if n_frames > 0 && self.src_format == Format::H264 {
78            if self.keyframes_cache.is_none() {
79                self.read_keyframes();
80            }
81            let keyframes = self.keyframes_cache.as_ref().unwrap();
82
83            let decoder = match self.h264_decoder_state.as_mut() {
84                Some(decoder) => decoder,
85                None => {
86                    return Err(StrandMkvSourceError::CannotSkipWithoutDecodingH264.into());
87                }
88            };
89
90            let mut best_keyframe = keyframes[0];
91            let target_frame = n_frames + 1; // we skip N so want N+1.
92            for keyframe in keyframes.iter() {
93                if *keyframe <= target_frame {
94                    best_keyframe = *keyframe;
95                }
96            }
97
98            for (idx, bd) in self.parsed.block_data.iter().take(n_frames).enumerate() {
99                // always decode keyframe and subsequent before target
100                if idx < best_keyframe {
101                    // always decode first frame with SPS and PPS
102                    if idx != 0 {
103                        // skip decoding this frame
104                        continue;
105                    }
106                }
107                self.rdr.seek(std::io::SeekFrom::Start(bd.start_idx))?;
108                let mut h264_raw_buf = vec![0u8; bd.size];
109                self.rdr.read_exact(&mut h264_raw_buf)?;
110
111                if let Some(decoded_yuv) = decoder.decode(&h264_raw_buf)? {
112                    decoded_yuv
113                } else {
114                    return Err(StrandMkvSourceError::CouldNotDecodeSingleFrameWithOpenH264.into());
115                };
116            }
117
118            self.keyframes_cache = None; // reset this.
119        }
120
121        let block_data = self.parsed.block_data.split_off(n_frames);
122
123        let timeshift = block_data[0].pts;
124        self.parsed.block_data = block_data
125            .into_iter()
126            .map(|mut el| {
127                el.pts -= timeshift;
128                el
129            })
130            .collect();
131        self.parsed.metadata.creation_time += chrono::Duration::from_std(timeshift).unwrap();
132        self.keyframes_cache = None;
133        Ok(())
134    }
135    fn estimate_luminance_range(&mut self) -> Result<(u16, u16)> {
136        Err(Error::UnsupportedForEsimatingLuminangeRange)
137    }
138    fn decode_order_iter<'a>(&'a mut self) -> Box<dyn Iterator<Item = Result<FrameData>> + 'a> {
139        Box::new(StrandCamMkvSourceIter {
140            parent: self,
141            idx: 0,
142        })
143    }
144    fn timestamp_source(&self) -> &str {
145        "MKV creation time + PTS"
146    }
147    fn has_timestamps(&self) -> bool {
148        true
149    }
150}
151
152struct StrandCamMkvSourceIter<'a, R: Read + Seek> {
153    parent: &'a mut StrandCamMkvSource<R>,
154    idx: usize,
155}
156
157#[derive(PartialEq)]
158enum Format {
159    UncompressedMono,
160    H264,
161}
162
163impl<R: Read + Seek> StrandCamMkvSource<R> {
164    fn new<P>(
165        rdr: R,
166        path: Option<P>,
167        do_decode_h264: bool,
168        timestamp_source: crate::TimestampSource,
169    ) -> Result<Self>
170    where
171        P: AsRef<std::path::Path>,
172    {
173        let (parsed, rdr) = mkv_strand_reader::parse_strand_cam_mkv(rdr, false, path)?;
174        if let Some(uncompressed_fourcc) = &parsed.uncompressed_fourcc
175            && uncompressed_fourcc.as_str() != "Y800"
176        {
177            return Err(
178                StrandMkvSourceError::UnsupportedFourcc(uncompressed_fourcc.clone()).into(),
179            );
180        }
181
182        let is_uncompressed = parsed.uncompressed_fourcc.is_some();
183
184        let src_format = if is_uncompressed {
185            Format::UncompressedMono
186        } else if &parsed.codec == "V_MPEG4/ISO/AVC" {
187            Format::H264
188        } else {
189            return Err(StrandMkvSourceError::UnsupportedCodec(parsed.codec).into());
190        };
191
192        let h264_decoder_state = if do_decode_h264 {
193            Some(crate::opt_openh264_decoder::DecoderType::new()?)
194        } else {
195            None
196        };
197
198        match timestamp_source {
199            crate::TimestampSource::BestGuess => {}
200            _ => {
201                return Err(StrandMkvSourceError::UnimplTsSource { timestamp_source }.into());
202            }
203        }
204
205        Ok(Self {
206            rdr,
207            parsed,
208            src_format,
209            is_uncompressed,
210            h264_decoder_state,
211            keyframes_cache: None,
212        })
213    }
214
215    pub fn is_uncompressed(&self) -> bool {
216        self.is_uncompressed
217    }
218
219    /// Return all indices of keyframes (I frames)
220    fn read_keyframes(&mut self) {
221        let mut keyframes_cache = vec![];
222        for (idx, bd) in self.parsed.block_data.iter().enumerate() {
223            if bd.is_keyframe {
224                keyframes_cache.push(idx);
225            }
226        }
227        self.keyframes_cache = Some(keyframes_cache);
228    }
229
230    fn get_frame(&mut self, idx: usize) -> Option<Result<FrameData>> {
231        let bd = self.parsed.block_data.get(idx);
232        bd?;
233        Some(self.get_frame_inner(idx))
234    }
235    fn get_frame_inner(&mut self, idx: usize) -> Result<FrameData> {
236        let bd = &self.parsed.block_data[idx];
237
238        let width = self.parsed.width;
239        let height = self.parsed.height;
240        let stride = usize::try_from(self.parsed.width).unwrap();
241
242        self.rdr.seek(std::io::SeekFrom::Start(bd.start_idx))?;
243        let mut image_data = vec![0u8; bd.size];
244        self.rdr.read_exact(&mut image_data)?;
245
246        let pts = bd.pts;
247
248        let image = match self.src_format {
249            Format::UncompressedMono => super::ImageData::Decoded(
250                DynamicFrameOwned::from_buf(
251                    width,
252                    height,
253                    stride,
254                    image_data,
255                    machine_vision_formats::PixFmt::Mono8,
256                )
257                .unwrap(),
258            ),
259            Format::H264 => {
260                // This is a hacky and imperfect way to check if the h264 stream
261                // has a timestamp. It is hacky because:
262                //  1) it assumes the timestamp NAL unit will be the first NAL
263                // unit (or that there will only be one NAL unit). That said, I
264                // think this is actually what the MISB standard specifies.
265                //  2) it does not really parse the NAL unit structure and
266                // assumes, for example, that the start bytes are `[0x00, 0x00,
267                // 0x00, 0x01]` whereas `[0x00, 0x00, 0x01]` would also be
268                // theoretically valid start bytes. Still, we write the full 4
269                // start bytes, so this should be OK.
270                if !image_data.starts_with(&[0, 0, 0, 1]) {
271                    return Err(StrandMkvSourceError::UnexpectedImageData.into());
272                }
273                let has_precision_timestamp = image_data.starts_with(PRECISION_TIME_NALU_START);
274                if let Some(decoder) = self.h264_decoder_state.as_mut() {
275                    let dynamic_frame = if let Some(decoded_yuv) = decoder.decode(&image_data)? {
276                        my_decode(decoded_yuv, width, height)?
277                    } else {
278                        return Err(
279                            StrandMkvSourceError::CouldNotDecodeSingleFrameWithOpenH264.into()
280                        );
281                    };
282                    super::ImageData::Decoded(dynamic_frame)
283                } else {
284                    super::ImageData::EncodedH264(super::EncodedH264 {
285                        data: H264EncodingVariant::AnnexB(image_data),
286                        has_precision_timestamp,
287                    })
288                }
289            }
290        };
291
292        Ok(FrameData {
293            timestamp: Timestamp::Duration(pts),
294            image,
295            buf_len: bd.size,
296            idx,
297            // The MKV reader does not reconstruct picture order count, so this
298            // source relies on the default (identity) `presentation_order_iter`.
299            poc: None,
300        })
301    }
302}
303
304#[cfg(not(feature = "openh264"))]
305fn my_decode(_decoded_yuv: (), _width: u32, _height: u32) -> Result<DynamicFrameOwned> {
306    Err(Error::H264Error("No H264 decoder support at compile time"))
307}
308
309#[cfg(feature = "openh264")]
310fn my_decode(
311    decoded_yuv: openh264::decoder::DecodedYUV<'_>,
312    width: u32,
313    height: u32,
314) -> Result<DynamicFrameOwned> {
315    use openh264::formats::YUVSource;
316    let dim = decoded_yuv.dimensions();
317
318    let stride = dim.0 * 3;
319    let mut image_data = vec![0u8; stride * dim.1];
320    decoded_yuv.write_rgb8(&mut image_data);
321    Ok(strand_dynamic_frame::DynamicFrameOwned::from_static(
322        OImage::<machine_vision_formats::pixel_format::RGB8>::new(
323            width, height, stride, image_data,
324        )
325        .unwrap(),
326    ))
327}
328
329impl<R: Read + Seek> Iterator for StrandCamMkvSourceIter<'_, R> {
330    type Item = Result<FrameData>;
331    fn next(&mut self) -> Option<Self::Item> {
332        let result = self.parent.get_frame(self.idx);
333        self.idx += 1;
334        result
335    }
336    fn size_hint(&self) -> (usize, Option<usize>) {
337        let remaining = self.parent.parsed.block_data.len() - self.idx;
338        (remaining, Some(remaining))
339    }
340}
341
342pub(crate) fn mkv_source_from_path_with_timestamp_source<P: AsRef<Path>>(
343    path: P,
344    do_decode_h264: bool,
345    timestamp_source: crate::TimestampSource,
346) -> Result<StrandCamMkvSource<BufReader<std::fs::File>>> {
347    let rdr = std::fs::File::open(path.as_ref())?;
348    let buf_reader = BufReader::new(rdr);
349    StrandCamMkvSource::new(
350        buf_reader,
351        Some(path.as_ref().to_path_buf()),
352        do_decode_h264,
353        timestamp_source,
354    )
355}
frame_source/strand_cam_mkv_source.rs

frame_source/
strand_cam_mkv_source.rs
