Kieran/ffmpeg-rs-raw
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: decoder options

Browse Source
This commit is contained in:
kieran
2024-10-24 12:26:52 +01:00
parent 3e74159ade
commit cbaf4069c7
6 changed files with 405 additions and 215 deletions
Download Patch File Download Diff File Expand all files Collapse all files

109
src/decode.rs
View File

@ -1,3 +1,4 @@
use crate::{options_to_dict, StreamInfoChannel};
use std::collections::HashMap;
use std::ffi::CStr;
use std::ptr;
@ -12,12 +13,23 @@ use ffmpeg_sys_the_third::{
};
use libc::memcpy;
struct CodecContext {
pub struct DecoderCodecContext {
pub context: *mut AVCodecContext,
pub codec: *const AVCodec,
}
impl Drop for CodecContext {
impl DecoderCodecContext {
/// Set [AVCodecContext] options
pub fn set_opt(&mut self, options: HashMap<String, String>) -> Result<(), Error> {
crate::set_opts(self.context as *mut libc::c_void, options)
}
pub fn list_opts(&self) -> Result<Vec<String>, Error> {
crate::list_opts(self.context as *mut libc::c_void)
}
}
impl Drop for DecoderCodecContext {
fn drop(&mut self) {
unsafe {
avcodec_free_context(&mut self.context);
@ -27,14 +39,13 @@ impl Drop for CodecContext {
}
}
unsafe impl Send for DecoderCodecContext {}
unsafe impl Sync for DecoderCodecContext {}
pub struct Decoder {
codecs: HashMap<i32, CodecContext>,
codecs: HashMap<i32, DecoderCodecContext>,
}
unsafe impl Send for Decoder {}
unsafe impl Sync for Decoder {}
impl Decoder {
pub fn new() -> Self {
Self {
@ -42,6 +53,62 @@ impl Decoder {
}
}
/// Set up a decoder for a given channel
pub fn setup_decoder(
&mut self,
channel: &StreamInfoChannel,
options: Option<HashMap<String, String>>,
) -> Result<&mut DecoderCodecContext, Error> {
unsafe { self.setup_decoder_for_stream(channel.stream, options) }
}
/// Set up a decoder from an [AVStream]
pub unsafe fn setup_decoder_for_stream(
&mut self,
stream: *mut AVStream,
options: Option<HashMap<String, String>>,
) -> Result<&mut DecoderCodecContext, Error> {
if stream.is_null() {
anyhow::bail!("stream is null");
}
let codec_par = (*stream).codecpar;
assert_ne!(
codec_par,
ptr::null_mut(),
"Codec parameters are missing from stream"
);
if let std::collections::hash_map::Entry::Vacant(e) = self.codecs.entry((*stream).index) {
let codec = avcodec_find_decoder((*codec_par).codec_id);
if codec.is_null() {
anyhow::bail!(
"Failed to find codec: {}",
CStr::from_ptr(avcodec_get_name((*codec_par).codec_id)).to_str()?
)
}
let context = avcodec_alloc_context3(codec);
if context.is_null() {
anyhow::bail!("Failed to alloc context")
}
if avcodec_parameters_to_context(context, (*stream).codecpar) != 0 {
anyhow::bail!("Failed to copy codec parameters to context")
}
let mut dict = if let Some(options) = options {
options_to_dict(options)?
} else {
ptr::null_mut()
};
if avcodec_open2(context, codec, &mut dict) < 0 {
anyhow::bail!("Failed to open codec")
}
Ok(e.insert(DecoderCodecContext { context, codec }))
} else {
anyhow::bail!("Decoder already setup");
}
}
pub unsafe fn decode_pkt(
&mut self,
pkt: *mut AVPacket,
@ -54,34 +121,6 @@ impl Decoder {
"Passed stream reference does not match stream_index of packet"
);
let codec_par = (*stream).codecpar;
assert_ne!(
codec_par,
ptr::null_mut(),
"Codec parameters are missing from stream"
);
if let std::collections::hash_map::Entry::Vacant(e) = self.codecs.entry(stream_index) {
let codec = avcodec_find_decoder((*codec_par).codec_id);
if codec.is_null() {
return Err(Error::msg(format!(
"Failed to find codec: {}",
CStr::from_ptr(avcodec_get_name((*codec_par).codec_id)).to_str()?
)));
}
let context = avcodec_alloc_context3(ptr::null());
if context.is_null() {
return Err(Error::msg("Failed to alloc context"));
}
if avcodec_parameters_to_context(context, (*stream).codecpar) != 0 {
return Err(Error::msg("Failed to copy codec parameters to context"));
}
if avcodec_open2(context, codec, ptr::null_mut()) < 0 {
return Err(Error::msg("Failed to open codec"));
}
e.insert(CodecContext { context, codec });
}
if let Some(ctx) = self.codecs.get_mut(&stream_index) {
// subtitles don't need decoding, create a frame from the pkt data
if (*ctx.codec).type_ == AVMediaType::AVMEDIA_TYPE_SUBTITLE {

180
src/demux.rs
View File

@ -1,14 +1,11 @@
use crate::return_ffmpeg_error;
use crate::{get_ffmpeg_error_msg, DemuxerInfo, StreamChannelType, StreamInfoChannel};
use anyhow::Error;
use ffmpeg_sys_the_third::*;
use std::ffi::CStr;
use std::{ptr, slice};
use crate::get_ffmpeg_error_msg;
use crate::return_ffmpeg_error;
use slimbox::{slimbox_unsize, SlimBox, SlimMut};
use std::fmt::{Display, Formatter};
use std::collections::HashMap;
use std::io::Read;
use std::mem::transmute;
use std::{ptr, slice};
#[no_mangle]
unsafe extern "C" fn read_data(
@ -27,131 +24,9 @@ unsafe extern "C" fn read_data(
}
}
#[derive(Clone, Debug, PartialEq)]
pub struct DemuxerInfo {
pub bitrate: usize,
pub duration: f32,
pub channels: Vec<StreamInfoChannel>,
}
unsafe impl Send for DemuxerInfo {}
unsafe impl Sync for DemuxerInfo {}
impl DemuxerInfo {
pub fn best_stream(&self, t: StreamChannelType) -> Option<&StreamInfoChannel> {
self.channels
.iter()
.filter(|a| a.channel_type == t)
.reduce(|acc, channel| {
if channel.best_metric() > acc.best_metric() {
channel
} else {
acc
}
})
}
pub fn best_video(&self) -> Option<&StreamInfoChannel> {
self.best_stream(StreamChannelType::Video)
}
pub fn best_audio(&self) -> Option<&StreamInfoChannel> {
self.best_stream(StreamChannelType::Audio)
}
pub fn best_subtitle(&self) -> Option<&StreamInfoChannel> {
self.best_stream(StreamChannelType::Subtitle)
}
pub unsafe fn is_best_stream(&self, stream: *mut AVStream) -> bool {
match (*(*stream).codecpar).codec_type {
AVMediaType::AVMEDIA_TYPE_VIDEO => {
(*stream).index == self.best_video().map_or(usize::MAX, |r| r.index) as libc::c_int
}
AVMediaType::AVMEDIA_TYPE_AUDIO => {
(*stream).index == self.best_audio().map_or(usize::MAX, |r| r.index) as libc::c_int
}
AVMediaType::AVMEDIA_TYPE_SUBTITLE => {
(*stream).index
== self.best_subtitle().map_or(usize::MAX, |r| r.index) as libc::c_int
}
_ => false,
}
}
}
impl Display for DemuxerInfo {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
write!(f, "Demuxer Info:")?;
for c in &self.channels {
write!(f, "\n{}", c)?;
}
Ok(())
}
}
#[derive(Clone, Debug, PartialEq)]
pub enum StreamChannelType {
Video,
Audio,
Subtitle,
}
impl Display for StreamChannelType {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
write!(
f,
"{}",
match self {
StreamChannelType::Video => "video",
StreamChannelType::Audio => "audio",
StreamChannelType::Subtitle => "subtitle",
}
)
}
}
#[derive(Clone, Debug, PartialEq)]
pub struct StreamInfoChannel {
pub index: usize,
pub channel_type: StreamChannelType,
pub codec: usize,
pub width: usize,
pub height: usize,
pub fps: f32,
pub sample_rate: usize,
pub format: usize,
}
impl StreamInfoChannel {
pub fn best_metric(&self) -> f32 {
match self.channel_type {
StreamChannelType::Video => self.width as f32 * self.height as f32 * self.fps,
StreamChannelType::Audio => self.sample_rate as f32,
StreamChannelType::Subtitle => 999. - self.index as f32,
}
}
}
impl Display for StreamInfoChannel {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
let codec_name = unsafe { CStr::from_ptr(avcodec_get_name(transmute(self.codec as i32))) };
write!(
f,
"{} #{}: codec={},size={}x{},fps={}",
self.channel_type,
self.index,
codec_name.to_str().unwrap(),
self.width,
self.height,
self.fps
)
}
}
pub enum DemuxerInput {
Url(String),
Reader(Option<SlimBox<dyn Read + 'static>>),
Reader(Option<SlimBox<dyn Read + 'static>>, Option<String>),
}
pub struct Demuxer {
@ -160,28 +35,35 @@ pub struct Demuxer {
}
impl Demuxer {
/// Create a new [Demuxer] from a file path or url
pub fn new(input: &str) -> Self {
unsafe {
let ps = avformat_alloc_context();
let ctx = avformat_alloc_context();
Self {
ctx: ps,
ctx,
input: DemuxerInput::Url(input.to_string()),
}
}
}
pub fn new_custom_io<R: Read + 'static>(reader: R) -> Self {
/// Create a new [Demuxer] from an object that implements [Read]
pub fn new_custom_io<R: Read + 'static>(reader: R, url: Option<String>) -> Self {
unsafe {
let ps = avformat_alloc_context();
(*ps).flags |= AVFMT_FLAG_CUSTOM_IO;
let ctx = avformat_alloc_context();
(*ctx).flags |= AVFMT_FLAG_CUSTOM_IO;
Self {
ctx: ps,
input: DemuxerInput::Reader(Some(slimbox_unsize!(reader))),
ctx,
input: DemuxerInput::Reader(Some(slimbox_unsize!(reader)), url),
}
}
}
/// Set [AVFormatContext] options
pub fn set_opt(&mut self, options: HashMap<String, String>) -> Result<(), Error> {
crate::set_opts(self.ctx as *mut libc::c_void, options)
}
unsafe fn open_input(&mut self) -> libc::c_int {
match &mut self.input {
DemuxerInput::Url(input) => avformat_open_input(
@ -190,7 +72,7 @@ impl Demuxer {
ptr::null_mut(),
ptr::null_mut(),
),
DemuxerInput::Reader(input) => {
DemuxerInput::Reader(input, url) => {
let input = input.take().expect("input stream already taken");
const BUFFER_SIZE: usize = 4096;
let pb = avio_alloc_context(
@ -206,7 +88,11 @@ impl Demuxer {
(*self.ctx).pb = pb;
avformat_open_input(
&mut self.ctx,
ptr::null_mut(),
if let Some(url) = url {
format!("{}\0", url).as_ptr() as *const libc::c_char
} else {
ptr::null_mut()
},
ptr::null_mut(),
ptr::null_mut(),
)
@ -221,7 +107,6 @@ impl Demuxer {
if avformat_find_stream_info(self.ctx, ptr::null_mut()) < 0 {
return Err(Error::msg("Could not find stream info"));
}
//av_dump_format(self.ctx, 0, ptr::null_mut(), 0);
let mut channel_infos = vec![];
@ -230,6 +115,7 @@ impl Demuxer {
match (*(*stream).codecpar).codec_type {
AVMediaType::AVMEDIA_TYPE_VIDEO => {
channel_infos.push(StreamInfoChannel {
stream,
index: (*stream).index as usize,
codec: (*(*stream).codecpar).codec_id as usize,
channel_type: StreamChannelType::Video,
@ -242,6 +128,7 @@ impl Demuxer {
}
AVMediaType::AVMEDIA_TYPE_AUDIO => {
channel_infos.push(StreamInfoChannel {
stream,
index: (*stream).index as usize,
codec: (*(*stream).codecpar).codec_id as usize,
channel_type: StreamChannelType::Audio,
@ -254,6 +141,7 @@ impl Demuxer {
}
AVMediaType::AVMEDIA_TYPE_SUBTITLE => {
channel_infos.push(StreamInfoChannel {
stream,
index: (*stream).index as usize,
codec: (*(*stream).codecpar).codec_id as usize,
channel_type: StreamChannelType::Subtitle,
@ -296,12 +184,14 @@ impl Demuxer {
impl Drop for Demuxer {
fn drop(&mut self) {
unsafe {
if let DemuxerInput::Reader(_) = self.input {
drop(SlimBox::<dyn Read>::from_raw((*(*self.ctx).pb).opaque));
if !self.ctx.is_null() {
unsafe {
if let DemuxerInput::Reader(_, _) = self.input {
drop(SlimBox::<dyn Read>::from_raw((*(*self.ctx).pb).opaque));
}
avformat_free_context(self.ctx);
self.ctx = ptr::null_mut();
}
avformat_free_context(self.ctx);
self.ctx = ptr::null_mut();
}
}
}

60
src/lib.rs
View File

@ -1,10 +1,17 @@
use ffmpeg_sys_the_third::av_make_error_string;
use anyhow::Error;
use ffmpeg_sys_the_third::{
av_dict_set, av_make_error_string, av_opt_next, av_opt_set, AVDictionary, AVOption,
AV_OPT_SEARCH_CHILDREN,
};
use std::collections::HashMap;
use std::ffi::CStr;
use std::ptr;
mod decode;
mod demux;
mod resample;
mod scale;
mod stream_info;
#[macro_export]
macro_rules! return_ffmpeg_error {
@ -24,8 +31,59 @@ fn get_ffmpeg_error_msg(ret: libc::c_int) -> String {
}
}
unsafe fn options_to_dict(options: HashMap<String, String>) -> Result<*mut AVDictionary, Error> {
let mut dict = ptr::null_mut();
for (key, value) in options {
let ret = av_dict_set(
&mut dict,
format!("{}\0", key).as_ptr() as *const libc::c_char,
format!("{}\0", value).as_ptr() as *const libc::c_char,
0,
);
return_ffmpeg_error!(ret);
}
Ok(dict)
}
fn list_opts(ctx: *mut libc::c_void) -> Result<Vec<String>, Error> {
let mut opt_ptr: *const AVOption = ptr::null_mut();
let mut ret = vec![];
unsafe {
loop {
opt_ptr = av_opt_next(ctx, opt_ptr);
if opt_ptr.is_null() {
break;
}
ret.push(
CStr::from_ptr((*opt_ptr).name)
.to_string_lossy()
.to_string(),
);
}
}
Ok(ret)
}
fn set_opts(ctx: *mut libc::c_void, options: HashMap<String, String>) -> Result<(), Error> {
unsafe {
for (key, value) in options {
let ret = av_opt_set(
ctx,
format!("{}\0", key).as_ptr() as *const libc::c_char,
format!("{}\0", value).as_ptr() as *const libc::c_char,
AV_OPT_SEARCH_CHILDREN,
);
return_ffmpeg_error!(ret);
}
}
Ok(())
}
pub use decode::*;
pub use demux::*;
pub use ffmpeg_sys_the_third;
pub use resample::*;
pub use scale::*;
pub use stream_info::*;

129
src/stream_info.rs Normal file
View File

@ -0,0 +1,129 @@
use ffmpeg_sys_the_third::{avcodec_get_name, AVMediaType, AVStream};
use std::ffi::CStr;
use std::fmt::{Display, Formatter};
use std::intrinsics::transmute;
#[derive(Clone, Debug, PartialEq)]
pub struct DemuxerInfo {
pub bitrate: usize,
pub duration: f32,
pub channels: Vec<StreamInfoChannel>,
}
impl DemuxerInfo {
pub fn best_stream(&self, t: StreamChannelType) -> Option<&StreamInfoChannel> {
self.channels
.iter()
.filter(|a| a.channel_type == t)
.reduce(|acc, channel| {
if channel.best_metric() > acc.best_metric() {
channel
} else {
acc
}
})
}
pub fn best_video(&self) -> Option<&StreamInfoChannel> {
self.best_stream(StreamChannelType::Video)
}
pub fn best_audio(&self) -> Option<&StreamInfoChannel> {
self.best_stream(StreamChannelType::Audio)
}
pub fn best_subtitle(&self) -> Option<&StreamInfoChannel> {
self.best_stream(StreamChannelType::Subtitle)
}
pub unsafe fn is_best_stream(&self, stream: *mut AVStream) -> bool {
match (*(*stream).codecpar).codec_type {
AVMediaType::AVMEDIA_TYPE_VIDEO => {
(*stream).index == self.best_video().map_or(usize::MAX, |r| r.index) as libc::c_int
}
AVMediaType::AVMEDIA_TYPE_AUDIO => {
(*stream).index == self.best_audio().map_or(usize::MAX, |r| r.index) as libc::c_int
}
AVMediaType::AVMEDIA_TYPE_SUBTITLE => {
(*stream).index
== self.best_subtitle().map_or(usize::MAX, |r| r.index) as libc::c_int
}
_ => false,
}
}
}
impl Display for DemuxerInfo {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
write!(f, "Demuxer Info:")?;
for c in &self.channels {
write!(f, "\n{}", c)?;
}
Ok(())
}
}
#[derive(Clone, Debug, PartialEq)]
pub enum StreamChannelType {
Video,
Audio,
Subtitle,
}
impl Display for StreamChannelType {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
write!(
f,
"{}",
match self {
StreamChannelType::Video => "video",
StreamChannelType::Audio => "audio",
StreamChannelType::Subtitle => "subtitle",
}
)
}
}
#[derive(Clone, Debug, PartialEq)]
pub struct StreamInfoChannel {
pub index: usize,
pub channel_type: StreamChannelType,
pub codec: usize,
pub width: usize,
pub height: usize,
pub fps: f32,
pub sample_rate: usize,
pub format: usize,
// private stream pointer
pub(crate) stream: *mut AVStream,
}
unsafe impl Send for StreamInfoChannel {}
unsafe impl Sync for StreamInfoChannel {}
impl StreamInfoChannel {
pub fn best_metric(&self) -> f32 {
match self.channel_type {
StreamChannelType::Video => self.width as f32 * self.height as f32 * self.fps,
StreamChannelType::Audio => self.sample_rate as f32,
StreamChannelType::Subtitle => 999. - self.index as f32,
}
}
}
impl Display for StreamInfoChannel {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
let codec_name = unsafe { CStr::from_ptr(avcodec_get_name(transmute(self.codec as i32))) };
write!(
f,
"{} #{}: codec={},size={}x{},fps={}",
self.channel_type,
self.index,
codec_name.to_str().unwrap(),
self.width,
self.height,
self.fps
)
}
}