use crate::{ error::{Error, FFmpegError}, model::{FFmpegAudioProps, FFmpegCodec, FFmpegProps, FFmpegSubtitleProps, FFmpegVideoProps}, utils::check_error, }; use std::{ ffi::{CStr, CString}, ptr, }; use ffmpeg_sys_next::{ av_bprint_finalize, av_bprint_init, av_channel_layout_describe_bprint, av_chroma_location_name, av_color_primaries_name, av_color_range_name, av_color_space_name, av_color_transfer_name, av_fourcc_make_string, av_get_bits_per_sample, av_get_bytes_per_sample, av_get_media_type_string, av_get_pix_fmt_name, av_get_sample_fmt_name, av_pix_fmt_desc_get, av_reduce, avcodec_alloc_context3, avcodec_flush_buffers, avcodec_free_context, avcodec_get_name, avcodec_open2, avcodec_parameters_to_context, avcodec_profile_name, avcodec_receive_frame, avcodec_send_packet, AVBPrint, AVChromaLocation, AVCodec, AVCodecContext, AVCodecParameters, AVColorPrimaries, AVColorRange, AVColorSpace, AVColorTransferCharacteristic, AVFieldOrder, AVFrame, AVMediaType, AVPacket, AVPixelFormat, AVRational, AVSampleFormat, AVERROR, AVERROR_EOF, AV_FOURCC_MAX_STRING_SIZE, FF_CODEC_PROPERTY_CLOSED_CAPTIONS, FF_CODEC_PROPERTY_FILM_GRAIN, FF_CODEC_PROPERTY_LOSSLESS, }; use libc::EAGAIN; pub struct FFmpegCodecContext(*mut AVCodecContext); impl FFmpegCodecContext { pub(crate) fn new() -> Result { let ptr = unsafe { avcodec_alloc_context3(ptr::null_mut()) }; if ptr.is_null() { Err(FFmpegError::VideoCodecAllocation)?; } Ok(Self(ptr)) } pub(crate) fn as_ref(&self) -> &AVCodecContext { unsafe { self.0.as_ref() }.expect("initialized on struct creation") } pub(crate) fn as_mut(&mut self) -> &mut AVCodecContext { unsafe { self.0.as_mut() }.expect("initialized on struct creation") } pub(crate) fn parameters_to_context( &mut self, codec_params: &AVCodecParameters, ) -> Result<&Self, Error> { check_error( unsafe { avcodec_parameters_to_context(self.as_mut(), codec_params) }, "Fail to fill the codec context with codec parameters", )?; Ok(self) } pub(crate) fn open2(&mut self, video_codec: &AVCodec) -> Result<&Self, Error> { check_error( unsafe { avcodec_open2(self.as_mut(), video_codec, ptr::null_mut()) }, "Failed to open video codec", )?; Ok(self) } pub(crate) fn flush(&mut self) { unsafe { avcodec_flush_buffers(self.as_mut()) }; } pub(crate) fn send_packet(&mut self, packet: *mut AVPacket) -> Result { match unsafe { avcodec_send_packet(self.as_mut(), packet) } { AVERROR_EOF => Ok(false), ret if ret == AVERROR(EAGAIN) => Err(FFmpegError::Again), ret if ret < 0 => Err(FFmpegError::from(ret)), _ => Ok(true), } } pub(crate) fn receive_frame(&mut self, frame: *mut AVFrame) -> Result { match unsafe { avcodec_receive_frame(self.as_mut(), frame) } { AVERROR_EOF => Ok(false), ret if ret == AVERROR(EAGAIN) => Err(FFmpegError::Again), ret if ret < 0 => Err(FFmpegError::from(ret)), _ => Ok(true), } } fn kind(&self) -> (Option, Option) { let kind = unsafe { av_get_media_type_string(self.as_ref().codec_type).as_ref() } .map(|media_type| unsafe { CStr::from_ptr(media_type) }); let sub_kind = unsafe { self.as_ref().codec.as_ref() } .and_then(|codec| unsafe { codec.name.as_ref() }) .map(|name| unsafe { CStr::from_ptr(name) }) .and_then(|sub_kind| { if let Some(kind) = kind { if kind == sub_kind { return None; } } Some(String::from_utf8_lossy(sub_kind.to_bytes()).to_string()) }); ( kind.map(|cstr| String::from_utf8_lossy(cstr.to_bytes()).to_string()), sub_kind, ) } fn name(&self) -> Option { unsafe { avcodec_get_name(self.as_ref().codec_id).as_ref() }.map(|codec_name| { let cstr = unsafe { CStr::from_ptr(codec_name) }; String::from_utf8_lossy(cstr.to_bytes()).to_string() }) } fn profile(&self) -> Option { if self.as_ref().profile == 0 { None } else { unsafe { avcodec_profile_name(self.as_ref().codec_id, self.as_ref().profile).as_ref() } .map(|profile| { let cstr = unsafe { CStr::from_ptr(profile) }; String::from_utf8_lossy(cstr.to_bytes()).to_string() }) } } fn tag(&self) -> Option { if self.as_ref().codec_tag != 0 { CString::new(vec![ 0; usize::try_from(AV_FOURCC_MAX_STRING_SIZE).expect( "AV_FOURCC_MAX_STRING_SIZE is 32, must fit in an usize" ) ]) .ok() .map(|buffer| { let tag = unsafe { CString::from_raw(av_fourcc_make_string( buffer.into_raw(), self.as_ref().codec_tag, )) }; String::from_utf8_lossy(tag.as_bytes()).to_string() }) } else { None } } fn bit_rate(&self) -> i32 { // TODO: use i64 instead of i32 when rspc supports it let ctx = self.as_ref(); match self.as_ref().codec_type { AVMediaType::AVMEDIA_TYPE_VIDEO | AVMediaType::AVMEDIA_TYPE_DATA | AVMediaType::AVMEDIA_TYPE_SUBTITLE | AVMediaType::AVMEDIA_TYPE_ATTACHMENT => ctx.bit_rate.try_into().unwrap_or_default(), AVMediaType::AVMEDIA_TYPE_AUDIO => { let bits_per_sample = unsafe { av_get_bits_per_sample(ctx.codec_id) }; if bits_per_sample != 0 { let bit_rate = ctx.sample_rate * ctx.ch_layout.nb_channels; if bit_rate <= i32::MAX / bits_per_sample { return bit_rate * (bits_per_sample); } } ctx.bit_rate.try_into().unwrap_or_default() } _ => 0, } } fn video_props(&self) -> Option { let ctx = self.as_ref(); if ctx.codec_type != AVMediaType::AVMEDIA_TYPE_VIDEO { return None; } let pixel_format = extract_pixel_format(ctx); let bits_per_channel = extract_bits_per_channel(ctx); let color_range = extract_color_range(ctx); let (color_space, color_primaries, color_transfer) = extract_colors(ctx); // Field Order let field_order = extract_field_order(ctx); // Chroma Sample Location let chroma_location = extract_chroma_location(ctx); let width = ctx.width; let height = ctx.height; let (aspect_ratio_num, aspect_ratio_den) = extract_aspect_ratio(ctx, width, height); let mut properties = vec![]; if ctx.properties & (FF_CODEC_PROPERTY_LOSSLESS.unsigned_abs()) != 0 { properties.push("Closed Captions".to_string()); } if ctx.properties & (FF_CODEC_PROPERTY_CLOSED_CAPTIONS.unsigned_abs()) != 0 { properties.push("Film Grain".to_string()); } if ctx.properties & (FF_CODEC_PROPERTY_FILM_GRAIN.unsigned_abs()) != 0 { properties.push("lossless".to_string()); } Some(FFmpegVideoProps { pixel_format, color_range, bits_per_channel, color_space, color_primaries, color_transfer, field_order, chroma_location, width, height, aspect_ratio_num, aspect_ratio_den, properties, }) } fn audio_props(&self) -> Option { let ctx = self.as_ref(); if ctx.codec_type != AVMediaType::AVMEDIA_TYPE_AUDIO { return None; } let sample_rate = if ctx.sample_rate > 0 { Some(ctx.sample_rate) } else { None }; let mut bprint = AVBPrint { str_: ptr::null_mut(), len: 0, size: 0, size_max: 0, reserved_internal_buffer: [0; 1], reserved_padding: [0; 1000], }; unsafe { av_bprint_init(&mut bprint, 0, u32::MAX /* AV_BPRINT_SIZE_UNLIMITED */); }; let mut channel_layout = ptr::null_mut(); let channel_layout = if unsafe { av_channel_layout_describe_bprint(&ctx.ch_layout, &mut bprint) } < 0 || unsafe { av_bprint_finalize(&mut bprint, &mut channel_layout) } < 0 || channel_layout.is_null() { None } else { let cstr = unsafe { CStr::from_ptr(channel_layout) }; Some(String::from_utf8_lossy(cstr.to_bytes()).to_string()) }; let sample_format = if ctx.sample_fmt == AVSampleFormat::AV_SAMPLE_FMT_NONE { None } else { unsafe { av_get_sample_fmt_name(ctx.sample_fmt).as_ref() }.map(|sample_fmt| { let cstr = unsafe { CStr::from_ptr(sample_fmt) }; String::from_utf8_lossy(cstr.to_bytes()).to_string() }) }; let bit_per_sample = if ctx.bits_per_raw_sample > 0 && ctx.bits_per_raw_sample != unsafe { av_get_bytes_per_sample(ctx.sample_fmt) } * 8 { Some(ctx.bits_per_raw_sample) } else { None }; Some(FFmpegAudioProps { delay: ctx.initial_padding, padding: ctx.trailing_padding, sample_rate, sample_format, bit_per_sample, channel_layout, }) } fn subtitle_props(&self) -> Option { if self.as_ref().codec_type != AVMediaType::AVMEDIA_TYPE_SUBTITLE { return None; } Some(FFmpegSubtitleProps { width: self.as_ref().width, height: self.as_ref().height, }) } fn props(&self) -> Option { match self.as_ref().codec_type { AVMediaType::AVMEDIA_TYPE_VIDEO => self.video_props().map(FFmpegProps::Video), AVMediaType::AVMEDIA_TYPE_AUDIO => self.audio_props().map(FFmpegProps::Audio), AVMediaType::AVMEDIA_TYPE_SUBTITLE => self.subtitle_props().map(FFmpegProps::Subtitle), _ => None, } } } fn extract_aspect_ratio( ctx: &AVCodecContext, width: i32, height: i32, ) -> (Option, Option) { if ctx.sample_aspect_ratio.num == 0 { (None, None) } else { let mut display_aspect_ratio = AVRational { num: 0, den: 0 }; let num = i64::from(width * ctx.sample_aspect_ratio.num); let den = i64::from(height * ctx.sample_aspect_ratio.den); let max = 1024 * 1024; unsafe { av_reduce( &mut display_aspect_ratio.num, &mut display_aspect_ratio.den, num, den, max, ); } ( Some(display_aspect_ratio.num), Some(display_aspect_ratio.den), ) } } fn extract_chroma_location(ctx: &AVCodecContext) -> Option { if ctx.chroma_sample_location == AVChromaLocation::AVCHROMA_LOC_UNSPECIFIED { None } else { unsafe { av_chroma_location_name(ctx.chroma_sample_location).as_ref() }.map( |chroma_location| { let cstr = unsafe { CStr::from_ptr(chroma_location) }; String::from_utf8_lossy(cstr.to_bytes()).to_string() }, ) } } fn extract_field_order(ctx: &AVCodecContext) -> Option { if ctx.field_order == AVFieldOrder::AV_FIELD_UNKNOWN { None } else { Some( (match ctx.field_order { AVFieldOrder::AV_FIELD_TT => "top first", AVFieldOrder::AV_FIELD_BB => "bottom first", AVFieldOrder::AV_FIELD_TB => "top coded first (swapped)", AVFieldOrder::AV_FIELD_BT => "bottom coded first (swapped)", _ => "progressive", }) .to_string(), ) } } fn extract_colors(ctx: &AVCodecContext) -> (Option, Option, Option) { if ctx.colorspace == AVColorSpace::AVCOL_SPC_UNSPECIFIED && ctx.color_primaries == AVColorPrimaries::AVCOL_PRI_UNSPECIFIED && ctx.color_trc == AVColorTransferCharacteristic::AVCOL_TRC_UNSPECIFIED { (None, None, None) } else { let color_space = unsafe { av_color_space_name(ctx.colorspace).as_ref() }.map(|color_space| { let cstr = unsafe { CStr::from_ptr(color_space) }; String::from_utf8_lossy(cstr.to_bytes()).to_string() }); let color_primaries = unsafe { av_color_primaries_name(ctx.color_primaries).as_ref() }.map( |color_primaries| { let cstr = unsafe { CStr::from_ptr(color_primaries) }; String::from_utf8_lossy(cstr.to_bytes()).to_string() }, ); let color_transfer = unsafe { av_color_transfer_name(ctx.color_trc).as_ref() }.map(|color_transfer| { let cstr = unsafe { CStr::from_ptr(color_transfer) }; String::from_utf8_lossy(cstr.to_bytes()).to_string() }); (color_space, color_primaries, color_transfer) } } fn extract_color_range(ctx: &AVCodecContext) -> Option { if ctx.color_range == AVColorRange::AVCOL_RANGE_UNSPECIFIED { None } else { unsafe { av_color_range_name(ctx.color_range).as_ref() }.map(|color_range| { let cstr = unsafe { CStr::from_ptr(color_range) }; String::from_utf8_lossy(cstr.to_bytes()).to_string() }) } } fn extract_bits_per_channel(ctx: &AVCodecContext) -> Option { if ctx.bits_per_raw_sample == 0 || ctx.pix_fmt == AVPixelFormat::AV_PIX_FMT_NONE { None } else { unsafe { av_pix_fmt_desc_get(ctx.pix_fmt).as_ref() }.and_then(|pix_fmt_desc| { let comp = pix_fmt_desc.comp[0]; if ctx.bits_per_raw_sample < comp.depth { Some(ctx.bits_per_raw_sample) } else { None } }) } } fn extract_pixel_format(ctx: &AVCodecContext) -> Option { if ctx.pix_fmt == AVPixelFormat::AV_PIX_FMT_NONE { None } else { unsafe { av_get_pix_fmt_name(ctx.pix_fmt).as_ref() }.map(|pixel_format| { let cstr = unsafe { CStr::from_ptr(pixel_format) }; String::from_utf8_lossy(cstr.to_bytes()).to_string() }) } } impl Drop for FFmpegCodecContext { fn drop(&mut self) { if !self.0.is_null() { unsafe { avcodec_free_context(&mut self.0) }; self.0 = ptr::null_mut(); } } } impl From<&FFmpegCodecContext> for FFmpegCodec { fn from(ctx: &FFmpegCodecContext) -> Self { let (kind, sub_kind) = ctx.kind(); Self { kind, sub_kind, name: ctx.name(), profile: ctx.profile(), tag: ctx.tag(), bit_rate: ctx.bit_rate(), props: ctx.props(), } } }