diff --git a/src/avutil/frame.rs b/src/avutil/frame.rs index e6313c5..7bc25ae 100644 --- a/src/avutil/frame.rs +++ b/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, diff --git a/src/avutil/mod.rs b/src/avutil/mod.rs index c042fea..ffc2250 100644 --- a/src/avutil/mod.rs +++ b/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::*; diff --git a/src/avutil/utils.rs b/src/avutil/utils.rs new file mode 100644 index 0000000..5176544 --- /dev/null +++ b/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()) }) +} diff --git a/tests/extract_mvs.rs b/tests/extract_mvs.rs index 860103f..d6fc96a 100644 --- a/tests/extract_mvs.rs +++ b/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, + 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> { +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> { .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(); }