From ddec83a54b4b73c76e1dd40cccee735055ceecda Mon Sep 17 00:00:00 2001 From: kieran Date: Fri, 8 Nov 2024 14:28:38 +0000 Subject: [PATCH] feat: encoder --- .gitignore | 4 +- src/encode.rs | 207 ++++++++++++++++++++++++++++++++++++++++++++++++++ src/lib.rs | 31 ++++++++ src/scale.rs | 16 +--- 4 files changed, 244 insertions(+), 14 deletions(-) create mode 100644 src/encode.rs diff --git a/.gitignore b/.gitignore index 40d9aca..adf1101 100644 --- a/.gitignore +++ b/.gitignore @@ -1,2 +1,4 @@ /target -/.idea \ No newline at end of file +/.idea + +test.png \ No newline at end of file diff --git a/src/encode.rs b/src/encode.rs new file mode 100644 index 0000000..c574061 --- /dev/null +++ b/src/encode.rs @@ -0,0 +1,207 @@ +use std::collections::HashMap; +use std::{ptr, slice}; + +use crate::{get_ffmpeg_error_msg, options_to_dict, return_ffmpeg_error}; +use anyhow::{bail, Error}; +use ffmpeg_sys_the_third::{ + av_channel_layout_default, av_packet_alloc, av_packet_free, avcodec_alloc_context3, + avcodec_find_encoder, avcodec_get_supported_config, avcodec_open2, avcodec_receive_packet, + avcodec_send_frame, AVChannelLayout, AVCodec, AVCodecConfig, AVCodecContext, AVCodecID, + AVFrame, AVPacket, AVPixelFormat, AVRational, AVSampleFormat, AVERROR, AVERROR_EOF, +}; +use libc::EAGAIN; + +pub struct Encoder { + ctx: *mut AVCodecContext, + codec: *const AVCodec, +} + +unsafe impl Send for Encoder {} + +impl Encoder { + /// Create a new encoder with the specified codec + pub fn new(codec: AVCodecID) -> Result { + unsafe { + let codec = avcodec_find_encoder(codec); + if codec.is_null() { + bail!("codec not found"); + } + let ctx = avcodec_alloc_context3(codec); + if ctx.is_null() { + bail!("context allocation failed"); + } + // set some defaults + (*ctx).time_base = AVRational { num: 1, den: 1 }; + Ok(Self { ctx, codec }) + } + } + + /// List supported configs (see [avcodec_get_supported_config]) + pub unsafe fn list_configs<'a, T>(&mut self, cfg: AVCodecConfig) -> Result<&'a [T], Error> { + let mut dst = ptr::null_mut(); + let mut num_dst = 0; + let ret = avcodec_get_supported_config( + self.ctx, + self.codec, + cfg, + 0, + ptr::addr_of_mut!(dst) as _, + &mut num_dst, + ); + return_ffmpeg_error!(ret); + Ok(slice::from_raw_parts(dst, num_dst as usize)) + } + + /// Set the encoder timebase + pub unsafe fn with_timebase(self, timebase: AVRational) -> Self { + (*self.ctx).time_base = timebase; + self + } + + /// Set the encoder bitrate + pub unsafe fn with_bitrate(self, bitrate: i64) -> Self { + (*self.ctx).bit_rate = bitrate; + self + } + + /// Set the encoder sample rate (audio) + pub unsafe fn with_sample_rate(self, fmt: i32) -> Self { + (*self.ctx).sample_rate = fmt; + self + } + + /// Set the encoder width in pixels + pub unsafe fn with_width(self, width: i32) -> Self { + (*self.ctx).width = width; + self + } + + /// Set the encoder height in pixels + pub unsafe fn with_height(self, height: i32) -> Self { + (*self.ctx).height = height; + self + } + + /// Set the encoder level (see AV_LEVEL_*) + pub unsafe fn with_level(self, level: i32) -> Self { + (*self.ctx).level = level; + self + } + + /// Set the encoder profile (see AV_PROFILE_*) + pub unsafe fn with_profile(self, profile: i32) -> Self { + (*self.ctx).profile = profile; + self + } + + /// Set the encoder profile (see AV_PROFILE_*) + pub unsafe fn with_framerate(self, rate: AVRational) -> Self { + (*self.ctx).framerate = rate; + self + } + + /// Set the encoder pixel format + pub unsafe fn with_pix_fmt(self, fmt: AVPixelFormat) -> Self { + (*self.ctx).pix_fmt = fmt; + self + } + + /// Set the encoder sample format (audio) + pub unsafe fn with_sample_format(self, fmt: AVSampleFormat) -> Self { + (*self.ctx).sample_fmt = fmt; + self + } + + /// Set the encoder channel layout (audio) + pub unsafe fn with_channel_layout(self, layout: AVChannelLayout) -> Self { + (*self.ctx).ch_layout = layout; + self + } + + /// Set the encoder channel layout using number of channels (audio) + pub unsafe fn with_default_channel_layout(self, channels: i32) -> Self { + let mut layout = AVChannelLayout::empty(); + av_channel_layout_default(&mut layout, channels); + (*self.ctx).ch_layout = layout; + self + } + + /// Apply options to context + pub unsafe fn with_options(self, fx: F) -> Self + where + F: FnOnce(*mut AVCodecContext), + { + fx(self.ctx); + self + } + + /// Open the encoder so that you can start encoding frames (see [avcodec_open2]) + pub unsafe fn open(self, options: Option>) -> Result { + let mut options = if let Some(options) = options { + options_to_dict(options)? + } else { + ptr::null_mut() + }; + let ret = avcodec_open2(self.ctx, self.codec, &mut options); + return_ffmpeg_error!(ret); + Ok(self) + } + + /// Encode a frame, returning a number of [AVPacket] + pub unsafe fn encode_frame( + &mut self, + frame: *mut AVFrame, + ) -> Result, Error> { + let mut pkgs = Vec::new(); + + let mut ret = avcodec_send_frame(self.ctx, frame); + if ret < 0 && ret != AVERROR(EAGAIN) { + bail!(get_ffmpeg_error_msg(ret)); + } + + while ret >= 0 || ret == AVERROR(EAGAIN) { + let mut pkt = av_packet_alloc(); + ret = avcodec_receive_packet(self.ctx, pkt); + if ret != 0 { + av_packet_free(&mut pkt); + if ret == AVERROR(EAGAIN) || ret == AVERROR_EOF { + break; + } + bail!(get_ffmpeg_error_msg(ret)); + } + pkgs.push(pkt); + } + + Ok(pkgs) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::generate_test_frame; + use std::io::Write; + + #[test] + fn test_encode_png() -> Result<(), Error> { + unsafe { + let frame = generate_test_frame(); + let mut encoder = Encoder::new(AVCodecID::AV_CODEC_ID_PNG)? + .with_width(512) + .with_height(512); + + let pix_fmts: &[AVPixelFormat] = + encoder.list_configs(AVCodecConfig::AV_CODEC_CONFIG_PIX_FORMAT)?; + encoder = encoder.with_pix_fmt(pix_fmts[0]).open(None)?; + + let mut test_file = std::fs::File::create("test.png")?; + let mut pkts = encoder.encode_frame(frame)?; + let flush_pkts = encoder.encode_frame(ptr::null_mut())?; + pkts.extend(flush_pkts); + for pkt in pkts { + test_file.write(slice::from_raw_parts((*pkt).data, (*pkt).size as usize))?; + } + } + Ok(()) + } +} diff --git a/src/lib.rs b/src/lib.rs index f87b3be..3e03e84 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -9,6 +9,7 @@ use std::ptr; mod decode; mod demux; +mod encode; mod filter; mod resample; mod scale; @@ -128,6 +129,36 @@ pub unsafe fn get_frame_from_hw(mut frame: *mut AVFrame) -> Result<*mut AVFrame, } } +#[cfg(test)] +pub unsafe fn generate_test_frame() -> *mut AVFrame { + use ffmpeg_sys_the_third::{av_frame_get_buffer, AVPixelFormat}; + use std::mem::transmute; + + let frame = av_frame_alloc(); + (*frame).width = 512; + (*frame).height = 512; + (*frame).format = transmute(AVPixelFormat::AV_PIX_FMT_RGB24); + av_frame_get_buffer(frame, 0); + + let mut lx = 0; + for line in 0..(*frame).height { + let c = lx % 3; + for y in 0..(*frame).width as usize { + let ptr = (*frame).data[0]; + let offset = (line * (*frame).linesize[0]) as usize + (y * 3); + match c { + 0 => *ptr.add(offset) = 0xff, + 1 => *ptr.add(offset + 1) = 0xff, + 2 => *ptr.add(offset + 2) = 0xff, + _ => {} + } + } + lx += 1; + } + + frame +} + pub use decode::*; pub use demux::*; pub use ffmpeg_sys_the_third; diff --git a/src/scale.rs b/src/scale.rs index bb7e866..d45642c 100644 --- a/src/scale.rs +++ b/src/scale.rs @@ -123,23 +123,13 @@ impl Scaler { #[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 - } + use crate::generate_test_frame; + use ffmpeg_sys_the_third::{av_frame_free, AVPixelFormat}; #[test] fn scale_rgb24_yuv420() { unsafe { - let mut frame = blank_frame(); + let mut frame = generate_test_frame(); let mut scaler = Scaler::new(); // downscale