use std::collections::HashMap;
use std::ffi::{CStr, CString};
use std::ptr::{null, null_mut};
#[cfg(not(feature = "docs-rs"))]
use ffmpeg_sys_next::AVChannelOrder;
use ffmpeg_sys_next::AVMediaType::{
AVMEDIA_TYPE_ATTACHMENT, AVMEDIA_TYPE_AUDIO, AVMEDIA_TYPE_DATA, AVMEDIA_TYPE_SUBTITLE,
AVMEDIA_TYPE_UNKNOWN, AVMEDIA_TYPE_VIDEO,
};
use ffmpeg_sys_next::{
av_dict_free, av_dict_get, av_dict_iterate, av_find_best_stream, avcodec_get_name,
avformat_find_stream_info, AVCodecID, AVDictionary, AVDictionaryEntry, AVRational,
};
use ffmpeg_sys_next::{avformat_alloc_context, avformat_close_input, avformat_open_input};
use crate::core::context::AVFormatContextBox;
use crate::error::{FindStreamError, OpenInputError, Result};
#[derive(Debug, Clone)]
pub enum StreamInfo {
Video {
index: i32,
time_base: AVRational,
start_time: i64,
duration: i64,
nb_frames: i64,
r_frame_rate: AVRational,
sample_aspect_ratio: AVRational,
metadata: HashMap<String, String>,
avg_frame_rate: AVRational,
codec_id: AVCodecID,
codec_name: String,
width: i32,
height: i32,
bit_rate: i64,
pixel_format: i32,
video_delay: i32,
fps: f64,
rotate: i32,
},
Audio {
index: i32,
time_base: AVRational,
start_time: i64,
duration: i64,
nb_frames: i64,
metadata: HashMap<String, String>,
avg_frame_rate: AVRational,
codec_id: AVCodecID,
codec_name: String,
sample_rate: i32,
#[cfg(not(feature = "docs-rs"))]
order: AVChannelOrder,
nb_channels: i32,
bit_rate: i64,
sample_format: i32,
frame_size: i32,
},
Subtitle {
index: i32,
time_base: AVRational,
start_time: i64,
duration: i64,
nb_frames: i64,
metadata: HashMap<String, String>,
codec_id: AVCodecID,
codec_name: String,
},
Data {
index: i32,
time_base: AVRational,
start_time: i64,
duration: i64,
metadata: HashMap<String, String>,
},
Attachment {
index: i32,
metadata: HashMap<String, String>,
codec_id: AVCodecID,
codec_name: String,
},
Unknown {
index: i32,
metadata: HashMap<String, String>,
},
}
impl StreamInfo {
pub fn stream_type(&self) -> &'static str {
match self {
StreamInfo::Video { .. } => "Video",
StreamInfo::Audio { .. } => "Audio",
StreamInfo::Subtitle { .. } => "Subtitle",
StreamInfo::Data { .. } => "Data",
StreamInfo::Attachment { .. } => "Attachment",
StreamInfo::Unknown { .. } => "Unknown",
}
}
pub fn is_video(&self) -> bool {
matches!(self, StreamInfo::Video { .. })
}
pub fn is_audio(&self) -> bool {
matches!(self, StreamInfo::Audio { .. })
}
pub fn index(&self) -> i32 {
match self {
StreamInfo::Video { index, .. }
| StreamInfo::Audio { index, .. }
| StreamInfo::Subtitle { index, .. }
| StreamInfo::Data { index, .. }
| StreamInfo::Attachment { index, .. }
| StreamInfo::Unknown { index, .. } => *index,
}
}
}
unsafe fn extract_stream_info_from_stream(raw_stream: *mut ffmpeg_sys_next::AVStream) -> StreamInfo {
let stream = &*raw_stream;
let metadata = dict_to_hashmap(stream.metadata);
if stream.codecpar.is_null() {
return StreamInfo::Unknown {
index: stream.index,
metadata,
};
}
let codecpar = &*stream.codecpar;
let codec_id = codecpar.codec_id;
let codec_name = codec_name(codec_id);
let index = stream.index;
let time_base = stream.time_base;
let start_time = stream.start_time;
let duration = stream.duration;
let nb_frames = stream.nb_frames;
let avg_frame_rate = stream.avg_frame_rate;
match codecpar.codec_type {
AVMEDIA_TYPE_VIDEO => {
let width = codecpar.width;
let height = codecpar.height;
let bit_rate = codecpar.bit_rate;
let pixel_format = codecpar.format;
let video_delay = codecpar.video_delay;
let r_frame_rate = stream.r_frame_rate;
let sample_aspect_ratio = stream.sample_aspect_ratio;
let fps = if avg_frame_rate.den == 0 {
0.0
} else {
avg_frame_rate.num as f64 / avg_frame_rate.den as f64
};
let rotate = metadata
.get("rotate")
.and_then(|rotate| rotate.parse::<i32>().ok())
.unwrap_or(0);
StreamInfo::Video {
index,
time_base,
start_time,
duration,
nb_frames,
r_frame_rate,
sample_aspect_ratio,
metadata,
avg_frame_rate,
codec_id,
codec_name,
width,
height,
bit_rate,
pixel_format,
video_delay,
fps,
rotate,
}
}
AVMEDIA_TYPE_AUDIO => {
let sample_rate = codecpar.sample_rate;
#[cfg(not(feature = "docs-rs"))]
let ch_layout = codecpar.ch_layout;
let sample_format = codecpar.format;
let frame_size = codecpar.frame_size;
let bit_rate = codecpar.bit_rate;
StreamInfo::Audio {
index,
time_base,
start_time,
duration,
nb_frames,
metadata,
avg_frame_rate,
codec_id,
codec_name,
sample_rate,
#[cfg(not(feature = "docs-rs"))]
order: ch_layout.order,
#[cfg(feature = "docs-rs")]
nb_channels: 0,
#[cfg(not(feature = "docs-rs"))]
nb_channels: ch_layout.nb_channels,
bit_rate,
sample_format,
frame_size,
}
}
AVMEDIA_TYPE_SUBTITLE => StreamInfo::Subtitle {
index,
time_base,
start_time,
duration,
nb_frames,
metadata,
codec_id,
codec_name,
},
AVMEDIA_TYPE_DATA => StreamInfo::Data {
index,
time_base,
start_time,
duration,
metadata,
},
AVMEDIA_TYPE_ATTACHMENT => StreamInfo::Attachment {
index,
metadata,
codec_id,
codec_name,
},
_ => StreamInfo::Unknown { index, metadata },
}
}
pub(crate) unsafe fn extract_stream_infos(fmt_ctx_box: &AVFormatContextBox) -> Result<Vec<StreamInfo>> {
let fmt_ctx = fmt_ctx_box.fmt_ctx;
if fmt_ctx.is_null() {
return Err(OpenInputError::OutOfMemory.into());
}
let nb_streams = (*fmt_ctx).nb_streams as usize;
let streams_ptr = (*fmt_ctx).streams;
if nb_streams > 0 && streams_ptr.is_null() {
return Err(FindStreamError::NoStreamFound.into());
}
let mut infos = Vec::with_capacity(nb_streams);
for i in 0..nb_streams {
let raw_stream = *streams_ptr.add(i);
if raw_stream.is_null() {
infos.push(StreamInfo::Unknown {
index: i as i32,
metadata: HashMap::new(),
});
continue;
}
infos.push(extract_stream_info_from_stream(raw_stream));
}
if !infos.is_empty() && infos.iter().all(|i| matches!(i, StreamInfo::Unknown { .. })) {
return Err(FindStreamError::NoStreamFound.into());
}
Ok(infos)
}
fn find_best_stream_info(
url: impl Into<String>,
media_type: ffmpeg_sys_next::AVMediaType,
) -> Result<Option<StreamInfo>> {
let in_fmt_ctx_box = init_format_context(url)?;
unsafe {
let best_index = av_find_best_stream(
in_fmt_ctx_box.fmt_ctx,
media_type,
-1,
-1,
null_mut(),
0,
);
if best_index < 0 {
return Ok(None);
}
let nb_streams = (*in_fmt_ctx_box.fmt_ctx).nb_streams as usize;
let index = best_index as usize;
if index >= nb_streams {
return Err(FindStreamError::NoStreamFound.into());
}
let streams_ptr = (*in_fmt_ctx_box.fmt_ctx).streams;
if streams_ptr.is_null() {
return Err(FindStreamError::NoStreamFound.into());
}
let raw_stream = *streams_ptr.add(index);
if raw_stream.is_null() {
return Err(FindStreamError::NoStreamFound.into());
}
let info = extract_stream_info_from_stream(raw_stream);
if media_type != AVMEDIA_TYPE_UNKNOWN && matches!(info, StreamInfo::Unknown { .. }) {
return Ok(None);
}
Ok(Some(info))
}
}
pub fn find_video_stream_info(url: impl Into<String>) -> Result<Option<StreamInfo>> {
find_best_stream_info(url, AVMEDIA_TYPE_VIDEO)
}
pub fn find_audio_stream_info(url: impl Into<String>) -> Result<Option<StreamInfo>> {
find_best_stream_info(url, AVMEDIA_TYPE_AUDIO)
}
pub fn find_subtitle_stream_info(url: impl Into<String>) -> Result<Option<StreamInfo>> {
find_best_stream_info(url, AVMEDIA_TYPE_SUBTITLE)
}
pub fn find_data_stream_info(url: impl Into<String>) -> Result<Option<StreamInfo>> {
find_best_stream_info(url, AVMEDIA_TYPE_DATA)
}
pub fn find_attachment_stream_info(url: impl Into<String>) -> Result<Option<StreamInfo>> {
find_best_stream_info(url, AVMEDIA_TYPE_ATTACHMENT)
}
pub fn find_unknown_stream_info(url: impl Into<String>) -> Result<Option<StreamInfo>> {
find_best_stream_info(url, AVMEDIA_TYPE_UNKNOWN)
}
pub fn find_all_stream_infos(url: impl Into<String>) -> Result<Vec<StreamInfo>> {
let in_fmt_ctx_box = init_format_context(url)?;
unsafe { extract_stream_infos(&in_fmt_ctx_box) }
}
#[inline]
fn codec_name(id: AVCodecID) -> String {
unsafe {
let ptr = avcodec_get_name(id);
if ptr.is_null() {
"Unknown codec".into()
} else {
CStr::from_ptr(ptr).to_string_lossy().into_owned()
}
}
}
pub(crate) fn init_format_context(url: impl Into<String>) -> Result<AVFormatContextBox> {
crate::core::initialize_ffmpeg();
let url_cstr = CString::new(url.into())?;
unsafe {
let mut in_fmt_ctx = avformat_alloc_context();
if in_fmt_ctx.is_null() {
return Err(OpenInputError::OutOfMemory.into());
}
let mut format_opts = null_mut();
let scan_all_pmts_key = CString::new("scan_all_pmts")?;
if av_dict_get(
format_opts,
scan_all_pmts_key.as_ptr(),
null(),
ffmpeg_sys_next::AV_DICT_MATCH_CASE,
)
.is_null()
{
let scan_all_pmts_value = CString::new("1")?;
ffmpeg_sys_next::av_dict_set(
&mut format_opts,
scan_all_pmts_key.as_ptr(),
scan_all_pmts_value.as_ptr(),
ffmpeg_sys_next::AV_DICT_DONT_OVERWRITE,
);
};
#[cfg(not(feature = "docs-rs"))]
let mut ret =
{ avformat_open_input(&mut in_fmt_ctx, url_cstr.as_ptr(), null(), &mut format_opts) };
#[cfg(feature = "docs-rs")]
let mut ret = 0;
av_dict_free(&mut format_opts);
if ret < 0 {
avformat_close_input(&mut in_fmt_ctx);
return Err(OpenInputError::from(ret).into());
}
ret = avformat_find_stream_info(in_fmt_ctx, null_mut());
if ret < 0 {
avformat_close_input(&mut in_fmt_ctx);
return Err(FindStreamError::from(ret).into());
}
Ok(AVFormatContextBox::new(in_fmt_ctx, true, false))
}
}
fn dict_to_hashmap(dict: *mut AVDictionary) -> HashMap<String, String> {
if dict.is_null() {
return HashMap::new();
}
let mut map = HashMap::new();
unsafe {
let mut e: *const AVDictionaryEntry = null_mut();
while {
e = av_dict_iterate(dict, e);
!e.is_null()
} {
let k = CStr::from_ptr((*e).key).to_string_lossy().into_owned();
let v = CStr::from_ptr((*e).value).to_string_lossy().into_owned();
map.insert(k, v);
}
}
map
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_not_found() {
let result = find_all_stream_infos("not_found.mp4");
assert!(result.is_err());
let error = result.err().unwrap();
println!("{error}");
assert!(matches!(
error,
crate::error::Error::OpenInputStream(OpenInputError::NotFound)
))
}
#[test]
fn test_find_all_stream_infos() {
let stream_infos = find_all_stream_infos("test.mp4").unwrap();
assert_eq!(2, stream_infos.len());
for stream_info in stream_infos {
println!("{:?}", stream_info);
}
}
#[test]
fn test_find_video_stream_info() {
let option = find_video_stream_info("test.mp4").unwrap();
assert!(option.is_some());
let video_stream_info = option.unwrap();
println!("video_stream_info:{:?}", video_stream_info);
}
#[test]
fn test_find_audio_stream_info() {
let option = find_audio_stream_info("test.mp4").unwrap();
assert!(option.is_some());
let audio_stream_info = option.unwrap();
println!("audio_stream_info:{:?}", audio_stream_info);
}
#[test]
fn test_find_subtitle_stream_info() {
let option = find_subtitle_stream_info("test.mp4").unwrap();
assert!(option.is_none())
}
#[test]
fn test_find_data_stream_info() {
let option = find_data_stream_info("test.mp4").unwrap();
assert!(option.is_none());
}
#[test]
fn test_find_attachment_stream_info() {
let option = find_attachment_stream_info("test.mp4").unwrap();
assert!(option.is_none())
}
#[test]
fn test_find_unknown_stream_info() {
let option = find_unknown_stream_info("test.mp4").unwrap();
assert!(option.is_none())
}
#[test]
fn test_is_video() {
let video = StreamInfo::Video {
index: 0, time_base: AVRational { num: 1, den: 30 },
start_time: 0, duration: 100, nb_frames: 100,
r_frame_rate: AVRational { num: 30, den: 1 },
sample_aspect_ratio: AVRational { num: 1, den: 1 },
avg_frame_rate: AVRational { num: 30, den: 1 },
width: 1920, height: 1080, bit_rate: 0, pixel_format: 0,
video_delay: 0, fps: 30.0, rotate: 0,
codec_id: AVCodecID::AV_CODEC_ID_H264,
codec_name: "h264".to_string(), metadata: HashMap::new(),
};
let unknown = StreamInfo::Unknown { index: 1, metadata: HashMap::new() };
assert!(video.is_video());
assert!(!video.is_audio());
assert!(!unknown.is_video());
}
#[test]
fn test_is_audio() {
let audio = StreamInfo::Audio {
index: 1, time_base: AVRational { num: 1, den: 44100 },
start_time: 0, duration: 100, nb_frames: 0,
avg_frame_rate: AVRational { num: 0, den: 1 },
sample_rate: 44100,
#[cfg(not(feature = "docs-rs"))]
order: AVChannelOrder::AV_CHANNEL_ORDER_UNSPEC,
nb_channels: 2, bit_rate: 128000, sample_format: 0, frame_size: 1024,
codec_id: AVCodecID::AV_CODEC_ID_AAC,
codec_name: "aac".to_string(), metadata: HashMap::new(),
};
assert!(audio.is_audio());
assert!(!audio.is_video());
}
#[test]
fn test_index() {
let video = StreamInfo::Video {
index: 5, time_base: AVRational { num: 1, den: 30 },
start_time: 0, duration: 100, nb_frames: 100,
r_frame_rate: AVRational { num: 30, den: 1 },
sample_aspect_ratio: AVRational { num: 1, den: 1 },
avg_frame_rate: AVRational { num: 30, den: 1 },
width: 1920, height: 1080, bit_rate: 0, pixel_format: 0,
video_delay: 0, fps: 30.0, rotate: 0,
codec_id: AVCodecID::AV_CODEC_ID_H264,
codec_name: "h264".to_string(), metadata: HashMap::new(),
};
let unknown = StreamInfo::Unknown { index: 42, metadata: HashMap::new() };
assert_eq!(video.index(), 5);
assert_eq!(unknown.index(), 42);
}
}