From f8a085af09497e174c4d7c074f28cd738faf668b Mon Sep 17 00:00:00 2001 From: kieran Date: Wed, 6 Nov 2024 14:12:45 +0000 Subject: [PATCH] feat: scale tests --- examples/main.rs | 7 ++-- src/decode.rs | 20 ++++------- src/demux.rs | 6 ++-- src/filter.rs | 12 +++---- src/lib.rs | 29 +++++++--------- src/scale.rs | 82 +++++++++++++++++++++++++++++++++++++++++----- src/stream_info.rs | 24 +++++--------- 7 files changed, 110 insertions(+), 70 deletions(-) diff --git a/examples/main.rs b/examples/main.rs index f11cb43..2938376 100644 --- a/examples/main.rs +++ b/examples/main.rs @@ -1,8 +1,5 @@ -use ffmpeg_rs_raw::{Decoder, Demuxer, DemuxerInfo, Filter, Scaler}; -use ffmpeg_sys_the_third::AVHWDeviceType::{ - AV_HWDEVICE_TYPE_CUDA, AV_HWDEVICE_TYPE_D3D12VA, AV_HWDEVICE_TYPE_MEDIACODEC, - AV_HWDEVICE_TYPE_OPENCL, AV_HWDEVICE_TYPE_VDPAU, AV_HWDEVICE_TYPE_VULKAN, -}; +use ffmpeg_rs_raw::{Decoder, Demuxer, DemuxerInfo, Filter}; +use ffmpeg_sys_the_third::AVHWDeviceType::AV_HWDEVICE_TYPE_CUDA; use ffmpeg_sys_the_third::{av_frame_free, av_packet_free, AVMediaType}; use log::{error, info}; use std::env::args; diff --git a/src/decode.rs b/src/decode.rs index e3d2347..e549e76 100644 --- a/src/decode.rs +++ b/src/decode.rs @@ -1,4 +1,4 @@ -use crate::{options_to_dict, return_ffmpeg_error, StreamInfoChannel}; +use crate::{options_to_dict, return_ffmpeg_error, rstr, StreamInfoChannel}; use std::collections::{HashMap, HashSet}; use std::ffi::CStr; use std::fmt::{Display, Formatter}; @@ -46,9 +46,7 @@ impl Drop for DecoderCodecContext { impl Display for DecoderCodecContext { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { unsafe { - let codec_name = CStr::from_ptr(avcodec_get_name((*self.codec).id)) - .to_str() - .unwrap(); + let codec_name = rstr!(avcodec_get_name((*self.codec).id)); write!( f, "DecoderCodecContext: codec={}, hw={}", @@ -56,9 +54,7 @@ impl Display for DecoderCodecContext { if self.hw_config.is_null() { "no" } else { - CStr::from_ptr(av_hwdevice_get_type_name((*self.hw_config).device_type)) - .to_str() - .unwrap() + rstr!(av_hwdevice_get_type_name((*self.hw_config).device_type)) } ) } @@ -148,7 +144,7 @@ impl Decoder { if codec.is_null() { anyhow::bail!( "Failed to find codec: {}", - CStr::from_ptr(avcodec_get_name((*codec_par).codec_id)).to_str()? + rstr!(avcodec_get_name((*codec_par).codec_id)) ) } let context = avcodec_alloc_context3(codec); @@ -159,7 +155,7 @@ impl Decoder { let mut ret = avcodec_parameters_to_context(context, (*stream).codecpar); return_ffmpeg_error!(ret, "Failed to copy codec parameters to context"); - let codec_name = CStr::from_ptr(avcodec_get_name((*codec).id)).to_str()?; + let codec_name = rstr!(avcodec_get_name((*codec).id)); // try use HW decoder let mut hw_config = ptr::null(); if let Some(ref hw_types) = self.hw_decoder_types { @@ -171,9 +167,7 @@ impl Decoder { if hw_config.is_null() { break; } - let hw_name = - CStr::from_ptr(av_hwdevice_get_type_name((*hw_config).device_type)) - .to_str()?; + let hw_name = rstr!(av_hwdevice_get_type_name((*hw_config).device_type)); if !hw_types.contains(&(*hw_config).device_type) { debug!("skipping hwaccel={}_{}", codec_name, hw_name); continue; @@ -234,7 +228,7 @@ impl Decoder { let mut pkgs = Vec::new(); while ret >= 0 { - let mut frame = av_frame_alloc(); + let frame = av_frame_alloc(); ret = avcodec_receive_frame(ctx.context, frame); if ret < 0 { if ret == AVERROR_EOF || ret == AVERROR(libc::EAGAIN) { diff --git a/src/demux.rs b/src/demux.rs index 4a1c6fb..1199776 100644 --- a/src/demux.rs +++ b/src/demux.rs @@ -1,4 +1,4 @@ -use crate::return_ffmpeg_error; +use crate::{cstr, return_ffmpeg_error}; use crate::{get_ffmpeg_error_msg, DemuxerInfo, StreamChannelType, StreamInfoChannel}; use anyhow::Error; use ffmpeg_sys_the_third::*; @@ -68,7 +68,7 @@ impl Demuxer { match &mut self.input { DemuxerInput::Url(input) => avformat_open_input( &mut self.ctx, - format!("{}\0", input).as_ptr() as *const libc::c_char, + cstr!(input), ptr::null_mut(), ptr::null_mut(), ), @@ -89,7 +89,7 @@ impl Demuxer { avformat_open_input( &mut self.ctx, if let Some(url) = url { - format!("{}\0", url).as_ptr() as *const libc::c_char + cstr!(url) } else { ptr::null_mut() }, diff --git a/src/filter.rs b/src/filter.rs index 65096d2..cc76ad9 100644 --- a/src/filter.rs +++ b/src/filter.rs @@ -1,10 +1,9 @@ -use crate::{cstr, return_ffmpeg_error, set_opts}; +use crate::{cstr, return_ffmpeg_error, rstr, set_opts}; use anyhow::Error; use ffmpeg_sys_the_third::{ av_free, av_strdup, avfilter_get_by_name, avfilter_graph_alloc, avfilter_graph_alloc_filter, avfilter_graph_config, avfilter_graph_create_filter, avfilter_graph_dump, avfilter_graph_parse, - avfilter_graph_parse2, avfilter_graph_parse_ptr, avfilter_inout_alloc, AVFilterContext, - AVFilterGraph, AVFrame, + avfilter_inout_alloc, AVFilterContext, AVFilterGraph, AVFrame, }; use log::debug; use std::collections::HashMap; @@ -26,7 +25,7 @@ impl Filter { /// /// https://ffmpeg.org/ffmpeg-filters.html pub unsafe fn parse(graph: &str) -> Result { - let mut ctx = avfilter_graph_alloc(); + let ctx = avfilter_graph_alloc(); let inputs = avfilter_inout_alloc(); let outputs = avfilter_inout_alloc(); let src = avfilter_get_by_name(cstr!("buffer")); @@ -95,9 +94,8 @@ impl Filter { } pub unsafe fn build(&mut self) -> Result<(), Error> { - let mut d = avfilter_graph_dump(self.graph, ptr::null_mut()); - debug!("{}", CStr::from_ptr(d).to_string_lossy()); - av_free(d as *mut _); + let d = rstr!(avfilter_graph_dump(self.graph, ptr::null_mut())); + debug!("{}", d); let ret = avfilter_graph_config(self.graph, ptr::null_mut()); return_ffmpeg_error!(ret, "Failed to build filter"); diff --git a/src/lib.rs b/src/lib.rs index bf03cdf..e1f9f8c 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -36,24 +36,26 @@ macro_rules! cstr { }; } +#[macro_export] +macro_rules! rstr { + ($str:expr) => { + core::ffi::CStr::from_ptr($str).to_str().unwrap() + }; +} + fn get_ffmpeg_error_msg(ret: libc::c_int) -> String { unsafe { const BUF_SIZE: usize = 512; let mut buf: [libc::c_char; BUF_SIZE] = [0; BUF_SIZE]; av_make_error_string(buf.as_mut_ptr(), BUF_SIZE, ret); - String::from(CStr::from_ptr(buf.as_ptr()).to_str().unwrap()) + rstr!(buf.as_ptr()).to_string() } } unsafe fn options_to_dict(options: HashMap) -> 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, - ); + let ret = av_dict_set(&mut dict, cstr!(key), cstr!(value), 0); return_ffmpeg_error!(ret); } Ok(dict) @@ -93,11 +95,7 @@ fn list_opts(ctx: *mut libc::c_void) -> Result, Error> { break; } - ret.push( - CStr::from_ptr((*opt_ptr).name) - .to_string_lossy() - .to_string(), - ); + ret.push(rstr!((*opt_ptr).name).to_string()); } } Ok(ret) @@ -106,12 +104,7 @@ fn list_opts(ctx: *mut libc::c_void) -> Result, Error> { fn set_opts(ctx: *mut libc::c_void, options: HashMap) -> 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, - ); + let ret = av_opt_set(ctx, cstr!(key), cstr!(value), AV_OPT_SEARCH_CHILDREN); return_ffmpeg_error!(ret); } } diff --git a/src/scale.rs b/src/scale.rs index 4c9f94b..bb7e866 100644 --- a/src/scale.rs +++ b/src/scale.rs @@ -1,12 +1,13 @@ use std::mem::transmute; use std::ptr; -use crate::return_ffmpeg_error; +use crate::{return_ffmpeg_error, rstr}; use anyhow::{bail, Error}; use ffmpeg_sys_the_third::{ - av_frame_alloc, av_frame_copy_props, sws_freeContext, sws_getContext, sws_init_context, + av_frame_alloc, av_frame_copy_props, av_get_pix_fmt_name, sws_freeContext, sws_getContext, sws_scale_frame, AVFrame, AVPixelFormat, SwsContext, SWS_BILINEAR, }; +use log::trace; pub struct Scaler { width: u16, @@ -17,8 +18,6 @@ pub struct Scaler { unsafe impl Send for Scaler {} -unsafe impl Sync for Scaler {} - impl Drop for Scaler { fn drop(&mut self) { unsafe { @@ -28,12 +27,18 @@ impl Drop for Scaler { } } +impl Default for Scaler { + fn default() -> Self { + Self::new() + } +} + impl Scaler { - pub fn new(width: u16, height: u16, format: AVPixelFormat) -> Self { + pub fn new() -> Self { Self { - width, - height, - format, + width: 0, + height: 0, + format: AVPixelFormat::AV_PIX_FMT_YUV420P, ctx: ptr::null_mut(), } } @@ -75,6 +80,16 @@ impl Scaler { bail!("Failed to create scalar context"); } + trace!( + "scale setup: {}x{}@{} => {}x{}@{}", + (*frame).width, + (*frame).height, + rstr!(av_get_pix_fmt_name(transmute((*frame).format))), + width, + height, + rstr!(av_get_pix_fmt_name(format)) + ); + self.width = width; self.height = height; self.format = format; @@ -104,3 +119,54 @@ impl Scaler { Ok(dst_frame) } } + +#[cfg(test)] +mod tests { + use super::*; + use ffmpeg_sys_the_third::{ + av_frame_alloc, av_frame_free, av_frame_get_buffer, AVFrame, AVPixelFormat, + }; + + unsafe fn blank_frame() -> *mut AVFrame { + let frame = av_frame_alloc(); + (*frame).width = 512; + (*frame).height = 512; + (*frame).format = AVPixelFormat::AV_PIX_FMT_RGB24 as libc::c_int; + av_frame_get_buffer(frame, 0); + frame + } + + #[test] + fn scale_rgb24_yuv420() { + unsafe { + let mut frame = blank_frame(); + let mut scaler = Scaler::new(); + + // downscale + let mut out_frame = scaler + .process_frame(frame, 128, 128, AVPixelFormat::AV_PIX_FMT_YUV420P) + .expect("Failed to process frame"); + assert_eq!((*out_frame).width, 128); + assert_eq!((*out_frame).height, 128); + assert_eq!( + (*out_frame).format, + transmute(AVPixelFormat::AV_PIX_FMT_YUV420P) + ); + av_frame_free(&mut out_frame); + + // upscale + let mut out_frame = scaler + .process_frame(frame, 1024, 1024, AVPixelFormat::AV_PIX_FMT_YUV420P) + .expect("Failed to process frame"); + assert_eq!((*out_frame).width, 1024); + assert_eq!((*out_frame).height, 1024); + assert_eq!( + (*out_frame).format, + transmute(AVPixelFormat::AV_PIX_FMT_YUV420P) + ); + av_frame_free(&mut out_frame); + + av_frame_free(&mut frame); + } + } +} diff --git a/src/stream_info.rs b/src/stream_info.rs index 632c5b2..b70b088 100644 --- a/src/stream_info.rs +++ b/src/stream_info.rs @@ -1,4 +1,4 @@ -use crate::format_time; +use crate::{format_time, rstr}; use ffmpeg_sys_the_third::{ av_get_pix_fmt_name, av_get_sample_fmt_name, avcodec_get_name, AVMediaType, AVStream, }; @@ -130,44 +130,36 @@ impl StreamInfoChannel { 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))) }; + let codec_name = unsafe { rstr!(avcodec_get_name(transmute(self.codec as i32))) }; match self.channel_type { StreamChannelType::Video => write!( f, "{} #{}: codec={},size={}x{},fps={:.3},pix_fmt={}", self.channel_type, self.index, - codec_name.to_str().unwrap(), + codec_name, self.width, self.height, self.fps, - unsafe { - CStr::from_ptr(av_get_pix_fmt_name(transmute(self.format as libc::c_int))) - } - .to_str() - .unwrap(), + unsafe { rstr!(av_get_pix_fmt_name(transmute(self.format as libc::c_int))) }, ), StreamChannelType::Audio => write!( f, "{} #{}: codec={},format={},sample_rate={}", self.channel_type, self.index, - codec_name.to_str().unwrap(), + codec_name, unsafe { - CStr::from_ptr(av_get_sample_fmt_name(transmute( + rstr!(av_get_sample_fmt_name(transmute( self.format as libc::c_int, ))) - } - .to_str() - .unwrap(), + }, self.sample_rate, ), StreamChannelType::Subtitle => write!( f, "{} #{}: codec={}", - self.channel_type, - self.index, - codec_name.to_str().unwrap() + self.channel_type, self.index, codec_name ), } }