Revise extract_mvs example (#142)

* Revise `extract_mvs` example

* Remove `AVFrame::get_motion_vectors`

* Add `get_media_type_string`

src/avutil/frame.rs

@@ -161,12 +161,6 @@ impl<'frame> AVFrame {
             .upgrade()
             .map(|side_data_ptr| unsafe { AVFrameSideDataRef::from_raw(side_data_ptr) })
     }
-
-    pub fn get_motion_vectors(&'frame self) -> Option<&'frame [AVMotionVector]> {
-        let side_data =
-            self.get_side_data(ffi::AVFrameSideDataType_AV_FRAME_DATA_MOTION_VECTORS)?;
-        Some(unsafe { side_data.as_motion_vectors() })
-    }
 }
 
 impl Drop for AVFrame {
@@ -229,7 +223,10 @@ impl AVFrameWithImage {
 wrap_ref!(AVFrameSideData: ffi::AVFrameSideData);
 
 impl<'frame> AVFrameSideDataRef<'frame> {
-    unsafe fn as_motion_vectors(&self) -> &'frame [AVMotionVector] {
+    /// # Safety
+    ///
+    /// You should only call this function when you ensure side data is motion vector.
+    pub unsafe fn as_motion_vectors(&self) -> &'frame [AVMotionVector] {
         unsafe {
             slice::from_raw_parts(
                 self.data as *const _ as *const ffi::AVMotionVector,

src/avutil/mod.rs

@@ -15,6 +15,7 @@ mod pixfmt;
 mod rational;
 mod samplefmt;
 mod timestamp;
+mod utils;
 
 pub use audio_fifo::*;
 pub use channel_layout::*;
@@ -32,3 +33,4 @@ pub use pixfmt::*;
 pub use rational::*;
 pub use samplefmt::*;
 pub use timestamp::*;
+pub use utils::*;

src/avutil/utils.rs

@@ -0,0 +1,9 @@
+use crate::{ffi, shared::PointerUpgrade};
+use std::ffi::CStr;
+
+/// Return a string describing the media_type enum, NULL if media_type is unknown.
+pub fn get_media_type_string(media_type: i32) -> Option<&'static CStr> {
+    unsafe { ffi::av_get_media_type_string(media_type) }
+        .upgrade()
+        .map(|str| unsafe { CStr::from_ptr(str.as_ptr()) })
+}

tests/extract_mvs.rs

@@ -1,24 +1,19 @@
-use anyhow::{Context, Result};
+//! RIIR: https://github.com/FFmpeg/FFmpeg/blob/master/doc/examples/extract_mvs.c
+use anyhow::{anyhow, Context, Result};
 use cstr::cstr;
 use rsmpeg::{
     avcodec::{AVCodecContext, AVPacket},
     avformat::AVFormatContextInput,
-    avutil::{AVDictionary, AVMotionVector},
+    avutil::{get_media_type_string, AVDictionary},
     error::RsmpegError,
     ffi,
 };
 use std::ffi::{CStr, CString};
 
-struct MotionVector {
-    motion_vector: AVMotionVector,
-    frame_index: usize,
-}
-
 fn decode_packet(
     decode_context: &mut AVCodecContext,
     packet: Option<&AVPacket>,
-    frame_index: &mut usize,
-    motion_vectors: &mut Vec<MotionVector>,
+    video_frame_count: &mut usize,
 ) -> Result<()> {
     decode_context
         .send_packet(packet)
@@ -31,43 +26,66 @@ fn decode_packet(
             Err(e) => return Err(e.into()),
         };
 
-        *frame_index += 1;
+        *video_frame_count += 1;
 
-        if let Some(raw_motion_vectors) = frame.get_motion_vectors() {
+        if let Some(side_data) =
+            frame.get_side_data(ffi::AVFrameSideDataType_AV_FRAME_DATA_MOTION_VECTORS)
+        {
+            let raw_motion_vectors = unsafe { side_data.as_motion_vectors() };
             for &motion_vector in raw_motion_vectors {
-                // framenum,source,blockw,blockh,srcx,srcy,dstx,dsty,flags
-                motion_vectors.push(MotionVector {
-                    motion_vector,
-                    frame_index: *frame_index,
-                });
+                println!(
+                    "{},{:2},{:2},{:2},{:4},{:4},{:4},{:4},{:#x},{:4},{:4},{:4}",
+                    video_frame_count,
+                    motion_vector.source,
+                    motion_vector.w,
+                    motion_vector.h,
+                    motion_vector.src_x,
+                    motion_vector.src_y,
+                    motion_vector.dst_x,
+                    motion_vector.dst_y,
+                    motion_vector.flags,
+                    motion_vector.motion_x,
+                    motion_vector.motion_y,
+                    motion_vector.motion_scale,
+                );
             }
-        }
+        };
     }
     Ok(())
 }
 
 /// Extract motion vectors from a video.
-fn extract_mvs(video_path: &CStr) -> Result<Vec<MotionVector>> {
+fn extract_mvs(video_path: &CStr) -> Result<()> {
     let mut input_format_context = AVFormatContextInput::open(video_path, None, &mut None)?;
+    let media_type = ffi::AVMediaType_AVMEDIA_TYPE_VIDEO;
 
     let (stream_index, mut decode_context) = {
         let (stream_index, decoder) = input_format_context
-            .find_best_stream(ffi::AVMediaType_AVMEDIA_TYPE_VIDEO)?
-            .context("Cannot find best stream in this file.")?;
+            .find_best_stream(media_type)?
+            .with_context(|| {
+                anyhow!(
+                    "Could not find {} stream in input file '{}'",
+                    get_media_type_string(media_type).unwrap().to_string_lossy(),
+                    video_path.to_string_lossy()
+                )
+            })?;
 
-        let stream = input_format_context.streams().get(stream_index).unwrap();
+        let stream = &input_format_context.streams()[stream_index];
 
         let mut decode_context = AVCodecContext::new(&decoder);
 
-        decode_context.apply_codecpar(&stream.codecpar())?;
+        decode_context
+            .apply_codecpar(&stream.codecpar())
+            .context("Failed to copy codec parameters to codec context")?;
 
-        let key = cstr!("flags2");
-        let value = cstr!("+export_mvs");
-        let opts = AVDictionary::new(key, value, 0);
+        let opts = AVDictionary::new(cstr!("flags2"), cstr!("+export_mvs"), 0);
 
-        decode_context
-            .open(Some(opts))
-            .context("failed to open decode codec")?;
+        decode_context.open(Some(opts)).with_context(|| {
+            anyhow!(
+                "Failed to open {} codec",
+                get_media_type_string(media_type).unwrap().to_string_lossy()
+            )
+        })?;
 
         (stream_index, decode_context)
     };
@@ -76,51 +94,25 @@ fn extract_mvs(video_path: &CStr) -> Result<Vec<MotionVector>> {
         .dump(0, video_path)
         .context("Input format context dump failed.")?;
 
+    println!(
+        "framenum,source,blockw,blockh,srcx,srcy,dstx,dsty,flags,motion_x,motion_y,motion_scale"
+    );
+
     let mut video_frame_count = 0;
-    let mut motion_vectors = vec![];
 
     while let Some(packet) = input_format_context.read_packet().unwrap() {
         if packet.stream_index == stream_index as i32 {
-            decode_packet(
-                &mut decode_context,
-                Some(&packet),
-                &mut video_frame_count,
-                &mut motion_vectors,
-            )?;
+            decode_packet(&mut decode_context, Some(&packet), &mut video_frame_count)?;
         }
     }
 
-    decode_packet(
-        &mut decode_context,
-        None,
-        &mut video_frame_count,
-        &mut motion_vectors,
-    )?;
+    decode_packet(&mut decode_context, None, &mut video_frame_count)?;
 
-    Ok(motion_vectors)
+    Ok(())
 }
 
 #[test]
 fn extract_mvs_test() {
-    fn to_string(motion_vector: &MotionVector) -> String {
-        format!(
-            "{},{:2},{:2},{:2},{:4},{:4},{:4},{:4},{:#x}",
-            motion_vector.frame_index,
-            motion_vector.motion_vector.source,
-            motion_vector.motion_vector.w,
-            motion_vector.motion_vector.h,
-            motion_vector.motion_vector.src_x,
-            motion_vector.motion_vector.src_y,
-            motion_vector.motion_vector.dst_x,
-            motion_vector.motion_vector.dst_y,
-            motion_vector.motion_vector.flags,
-        )
-    }
     let video_path = &CString::new("tests/assets/vids/bear.mp4").unwrap();
-    let mvs = extract_mvs(video_path).unwrap();
-    assert_eq!(10783, mvs.len());
-    assert_eq!("2, 1,16,16, 264,  56, 264,  56,0x0", to_string(&mvs[114]));
-    assert_eq!("3,-1, 8, 8, 220,  52, 220,  52,0x0", to_string(&mvs[514]));
-    assert_eq!("7,-1,16,16,  87,   8,  88,   8,0x0", to_string(&mvs[1919]));
-    assert_eq!("4, 1,16,16, 232,  24, 232,  24,0x0", to_string(&mvs[810]));
+    extract_mvs(video_path).unwrap();
 }