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

feat: muxer

Browse Source
This commit is contained in:
kieran 2024-11-09 15:47:03 +00:00
parent ddec83a54b
commit 8032966982
No known key found for this signature in database
GPG Key ID: DE71CEB3925BE941
11 changed files with 299 additions and 60 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
.gitignore vendored
View File

@ -1,4 +1,7 @@
/target
/.idea
test.png
# Test output
*.mp4
*.png
*.mkv

2
examples/main.rs
View File

@ -9,7 +9,7 @@ use std::path::PathBuf;
fn main() {
env_logger::init();
let name = args().next().unwrap_or("main".to_string());
let path = if let Some(path) = args().skip(1).next() {
let path = if let Some(path) = args().nth(1) {
PathBuf::from(path)
} else {
error!("Usage: {} <path>", name);

26
src/decode.rs
View File

@ -1,4 +1,4 @@
use crate::{options_to_dict, return_ffmpeg_error, rstr, StreamInfoChannel};
use crate::{bail_ffmpeg, options_to_dict, rstr, StreamInfoChannel};
use std::collections::hash_map::Entry;
use std::collections::{HashMap, HashSet};
use std::fmt::{Display, Formatter};
@ -35,9 +35,9 @@ impl DecoderCodecContext {
/// Get the codec name
pub fn codec_name(&self) -> String {
let codec_name = unsafe { rstr!((*(*self).codec).name) };
let codec_name = unsafe { rstr!((*self.codec).name) };
if self.hw_config.is_null() {
format!("{}", codec_name)
codec_name.to_string()
} else {
let hw = unsafe { rstr!(av_hwdevice_get_type_name((*self.hw_config).device_type)) };
format!("{}_{}", codec_name, hw)
@ -48,9 +48,9 @@ impl DecoderCodecContext {
impl Drop for DecoderCodecContext {
fn drop(&mut self) {
unsafe {
if !self.context.is_null() {
avcodec_free_context(&mut self.context);
self.codec = ptr::null_mut();
self.context = ptr::null_mut();
}
}
}
}
@ -66,8 +66,6 @@ impl Display for DecoderCodecContext {
}
}
unsafe impl Send for DecoderCodecContext {}
pub struct Decoder {
/// Decoder instances by stream index
codecs: HashMap<i32, DecoderCodecContext>,
@ -75,6 +73,12 @@ pub struct Decoder {
hw_decoder_types: Option<HashSet<AVHWDeviceType>>,
}
impl Default for Decoder {
fn default() -> Self {
Self::new()
}
}
impl Decoder {
pub fn new() -> Self {
Self {
@ -177,7 +181,7 @@ impl Decoder {
}
let mut ret = avcodec_parameters_to_context(context, (*stream).codecpar);
return_ffmpeg_error!(ret, "Failed to copy codec parameters to context");
bail_ffmpeg!(ret, "Failed to copy codec parameters to context");
let codec_name = rstr!(avcodec_get_name((*codec).id));
// try use HW decoder
@ -205,7 +209,7 @@ impl Decoder {
ptr::null_mut(),
0,
);
return_ffmpeg_error!(ret, "Failed to create HW ctx");
bail_ffmpeg!(ret, "Failed to create HW ctx");
(*context).hw_device_ctx = av_buffer_ref(hw_buf_ref);
break;
}
@ -218,7 +222,7 @@ impl Decoder {
};
ret = avcodec_open2(context, codec, &mut dict);
return_ffmpeg_error!(ret, "Failed to open codec");
bail_ffmpeg!(ret, "Failed to open codec");
let ctx = DecoderCodecContext {
context,
@ -252,7 +256,7 @@ impl Decoder {
stream: *mut AVStream,
) -> Result<Vec<(*mut AVFrame, *mut AVStream)>, Error> {
let mut ret = avcodec_send_packet(ctx, pkt);
return_ffmpeg_error!(ret, "Failed to decode packet");
bail_ffmpeg!(ret, "Failed to decode packet");
let mut pkgs = Vec::new();
while ret >= 0 {

4
src/demux.rs
View File

@ -1,4 +1,4 @@
use crate::{cstr, return_ffmpeg_error};
use crate::{bail_ffmpeg, cstr};
use crate::{get_ffmpeg_error_msg, DemuxerInfo, StreamChannelType, StreamInfoChannel};
use anyhow::Error;
use ffmpeg_sys_the_third::*;
@ -102,7 +102,7 @@ impl Demuxer {
pub unsafe fn probe_input(&mut self) -> Result<DemuxerInfo, Error> {
let ret = self.open_input();
return_ffmpeg_error!(ret);
bail_ffmpeg!(ret);
if avformat_find_stream_info(self.ctx, ptr::null_mut()) < 0 {
return Err(Error::msg("Could not find stream info"));

56
src/encode.rs
View File

@ -1,13 +1,14 @@
use std::collections::HashMap;
use std::{ptr, slice};
use crate::{get_ffmpeg_error_msg, options_to_dict, return_ffmpeg_error};
use crate::{bail_ffmpeg, get_ffmpeg_error_msg, options_to_dict};
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,
av_channel_layout_default, av_d2q, av_inv_q, av_packet_alloc, av_packet_free,
avcodec_alloc_context3, avcodec_find_encoder, avcodec_free_context,
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;
@ -16,7 +17,15 @@ pub struct Encoder {
codec: *const AVCodec,
}
unsafe impl Send for Encoder {}
impl Drop for Encoder {
fn drop(&mut self) {
unsafe {
if !self.ctx.is_null() {
avcodec_free_context(&mut self.ctx);
}
}
}
}
impl Encoder {
/// Create a new encoder with the specified codec
@ -36,6 +45,16 @@ impl Encoder {
}
}
/// Get the codec
pub fn codec(&self) -> *const AVCodec {
self.codec
}
/// Get the codec context
pub fn codec_context(&self) -> *const AVCodecContext {
self.ctx
}
/// 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();
@ -48,16 +67,10 @@ impl Encoder {
ptr::addr_of_mut!(dst) as _,
&mut num_dst,
);
return_ffmpeg_error!(ret);
bail_ffmpeg!(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;
@ -94,9 +107,11 @@ impl Encoder {
self
}
/// Set the encoder profile (see AV_PROFILE_*)
pub unsafe fn with_framerate(self, rate: AVRational) -> Self {
(*self.ctx).framerate = rate;
/// Set the encoder framerate
pub unsafe fn with_framerate(self, fps: f32) -> Self {
let q = av_d2q(fps as f64, 90_000);
(*self.ctx).framerate = q;
(*self.ctx).time_base = av_inv_q(q);
self
}
@ -137,13 +152,15 @@ impl Encoder {
/// Open the encoder so that you can start encoding frames (see [avcodec_open2])
pub unsafe fn open(self, options: Option<HashMap<String, String>>) -> Result<Self, Error> {
assert!(!self.ctx.is_null());
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);
bail_ffmpeg!(ret);
Ok(self)
}
@ -169,6 +186,7 @@ impl Encoder {
}
bail!(get_ffmpeg_error_msg(ret));
}
(*pkt).time_base = (*self.ctx).time_base;
pkgs.push(pkt);
}
@ -187,8 +205,8 @@ mod tests {
unsafe {
let frame = generate_test_frame();
let mut encoder = Encoder::new(AVCodecID::AV_CODEC_ID_PNG)?
.with_width(512)
.with_height(512);
.with_width((*frame).width)
.with_height((*frame).height);
let pix_fmts: &[AVPixelFormat] =
encoder.list_configs(AVCodecConfig::AV_CODEC_CONFIG_PIX_FORMAT)?;

16
src/filter.rs
View File

@ -1,4 +1,4 @@
use crate::{cstr, return_ffmpeg_error, rstr, set_opts};
use crate::{bail_ffmpeg, cstr, rstr, set_opts};
use anyhow::Error;
use ffmpeg_sys_the_third::{
av_strdup, avfilter_get_by_name, avfilter_graph_alloc, avfilter_graph_alloc_filter,
@ -13,6 +13,12 @@ pub struct Filter {
graph: *mut AVFilterGraph,
}
impl Default for Filter {
fn default() -> Self {
Self::new()
}
}
impl Filter {
pub fn new() -> Self {
Self {
@ -39,7 +45,7 @@ impl Filter {
ptr::null_mut(),
ctx,
);
return_ffmpeg_error!(ret, "Failed to parse graph");
bail_ffmpeg!(ret, "Failed to parse graph");
let ret = avfilter_graph_create_filter(
&mut dst_ctx,
@ -49,7 +55,7 @@ impl Filter {
ptr::null_mut(),
ctx,
);
return_ffmpeg_error!(ret, "Failed to parse graph");
bail_ffmpeg!(ret, "Failed to parse graph");
(*outputs).name = av_strdup((*dst).name);
(*outputs).filter_ctx = dst_ctx;
@ -62,7 +68,7 @@ impl Filter {
(*inputs).next = ptr::null_mut();
let ret = avfilter_graph_parse(ctx, cstr!(graph), inputs, outputs, ptr::null_mut());
return_ffmpeg_error!(ret, "Failed to parse graph");
bail_ffmpeg!(ret, "Failed to parse graph");
let mut ret = Self { graph: ctx };
ret.build()?;
Ok(ret)
@ -97,7 +103,7 @@ impl Filter {
debug!("{}", d);
let ret = avfilter_graph_config(self.graph, ptr::null_mut());
return_ffmpeg_error!(ret, "Failed to build filter");
bail_ffmpeg!(ret, "Failed to build filter");
Ok(())
}

21
src/lib.rs
View File

@ -11,20 +11,22 @@ mod decode;
mod demux;
mod encode;
mod filter;
mod muxer;
mod resample;
mod scale;
mod stream_info;
mod transcode;
#[macro_export]
macro_rules! return_ffmpeg_error {
macro_rules! bail_ffmpeg {
($x:expr) => {
if $x < 0 {
anyhow::bail!(crate::get_ffmpeg_error_msg($x))
anyhow::bail!($crate::get_ffmpeg_error_msg($x))
}
};
($x:expr,$msg:expr) => {
if $x < 0 {
anyhow::bail!(format!("{}: {}", $msg, crate::get_ffmpeg_error_msg($x)))
anyhow::bail!(format!("{}: {}", $msg, $crate::get_ffmpeg_error_msg($x)))
}
};
}
@ -60,7 +62,7 @@ unsafe fn options_to_dict(options: HashMap<String, String>) -> Result<*mut AVDic
let mut dict = ptr::null_mut();
for (key, value) in options {
let ret = av_dict_set(&mut dict, cstr!(key), cstr!(value), 0);
return_ffmpeg_error!(ret);
bail_ffmpeg!(ret);
}
Ok(dict)
}
@ -109,7 +111,7 @@ fn set_opts(ctx: *mut libc::c_void, options: HashMap<String, String>) -> Result<
unsafe {
for (key, value) in options {
let ret = av_opt_set(ctx, cstr!(key), cstr!(value), AV_OPT_SEARCH_CHILDREN);
return_ffmpeg_error!(ret);
bail_ffmpeg!(ret);
}
}
Ok(())
@ -122,7 +124,7 @@ pub unsafe fn get_frame_from_hw(mut frame: *mut AVFrame) -> Result<*mut AVFrame,
} else {
let new_frame = av_frame_alloc();
let ret = av_hwframe_transfer_data(new_frame, frame, 0);
return_ffmpeg_error!(ret);
bail_ffmpeg!(ret);
av_frame_copy_props(new_frame, frame);
av_frame_free(&mut frame);
Ok(new_frame)
@ -135,8 +137,8 @@ pub unsafe fn generate_test_frame() -> *mut AVFrame {
use std::mem::transmute;
let frame = av_frame_alloc();
(*frame).width = 512;
(*frame).height = 512;
(*frame).width = 1024;
(*frame).height = 1024;
(*frame).format = transmute(AVPixelFormat::AV_PIX_FMT_RGB24);
av_frame_get_buffer(frame, 0);
@ -161,8 +163,11 @@ pub unsafe fn generate_test_frame() -> *mut AVFrame {
pub use decode::*;
pub use demux::*;
pub use encode::*;
pub use ffmpeg_sys_the_third;
pub use filter::*;
pub use muxer::*;
pub use resample::*;
pub use scale::*;
pub use stream_info::*;
pub use transcode::*;

185
src/muxer.rs Normal file
View File

@ -0,0 +1,185 @@
use crate::{bail_ffmpeg, cstr, set_opts, Encoder};
use anyhow::{bail, Result};
use ffmpeg_sys_the_third::{
av_dump_format, av_interleaved_write_frame, av_packet_rescale_ts, av_write_trailer,
avcodec_parameters_from_context, avformat_alloc_output_context2, avformat_free_context,
avformat_new_stream, avformat_write_header, avio_flush, avio_open, AVFormatContext, AVPacket,
AVFMT_GLOBALHEADER, AVFMT_NOFILE, AVIO_FLAG_WRITE, AV_CODEC_FLAG_GLOBAL_HEADER,
};
use std::collections::HashMap;
use std::path::PathBuf;
use std::ptr;
pub struct Muxer {
ctx: *mut AVFormatContext,
}
impl Drop for Muxer {
fn drop(&mut self) {
unsafe {
if !self.ctx.is_null() {
avio_flush((*self.ctx).pb);
avformat_free_context(self.ctx);
}
}
}
}
impl Default for Muxer {
fn default() -> Self {
Self::new()
}
}
impl Muxer {
pub fn new() -> Self {
Self {
ctx: ptr::null_mut(),
}
}
/// Open the muxer with a destination path
pub unsafe fn with_output(
mut self,
dst: &PathBuf,
format: Option<&str>,
options: Option<HashMap<String, String>>,
) -> Result<Self> {
if !self.ctx.is_null() {
bail!("context already open");
}
let ret = avformat_alloc_output_context2(
&mut self.ctx,
ptr::null_mut(),
if let Some(ref format) = format {
cstr!(format)
} else {
ptr::null()
},
cstr!(dst.to_str().unwrap()),
);
bail_ffmpeg!(ret);
// Setup global header flag
if (*(*self.ctx).oformat).flags & AVFMT_GLOBALHEADER != 0 {
(*self.ctx).flags |= AV_CODEC_FLAG_GLOBAL_HEADER as libc::c_int;
}
// Set options on ctx
if let Some(opts) = options {
set_opts((*self.ctx).priv_data, opts)?;
}
Ok(self)
}
/// Add a stream to the output using an existing encoder
pub unsafe fn with_stream_encoder(self, encoder: &Encoder) -> Result<Self> {
let stream = avformat_new_stream(self.ctx, encoder.codec());
if stream.is_null() {
bail!("unable to allocate stream");
}
let ret = avcodec_parameters_from_context((*stream).codecpar, encoder.codec_context());
bail_ffmpeg!(ret);
// setup other stream params
let encoder_ctx = encoder.codec_context();
(*stream).time_base = (*encoder_ctx).time_base;
(*stream).avg_frame_rate = (*encoder_ctx).framerate;
(*stream).r_frame_rate = (*encoder_ctx).framerate;
Ok(self)
}
/// Open the output to start sending packets
pub unsafe fn open(self) -> Result<Self> {
if (*(*self.ctx).oformat).flags & AVFMT_NOFILE == 0 {
let ret = avio_open(&mut (*self.ctx).pb, (*self.ctx).url, AVIO_FLAG_WRITE);
bail_ffmpeg!(ret);
}
let ret = avformat_write_header(self.ctx, ptr::null_mut());
bail_ffmpeg!(ret);
av_dump_format(self.ctx, 0, (*self.ctx).url, 1);
Ok(self)
}
/// Write a packet to the output
pub unsafe fn write_packet(&mut self, pkt: *mut AVPacket) -> Result<()> {
assert!((*pkt).stream_index >= 0 && (*pkt).stream_index < (*self.ctx).nb_streams as i32);
assert!((*pkt).time_base.num != 0 && (*pkt).time_base.den != 0);
let stream = *(*self.ctx).streams.add((*pkt).stream_index as usize);
av_packet_rescale_ts(pkt, (*pkt).time_base, (*stream).time_base);
let ret = av_interleaved_write_frame(self.ctx, pkt);
bail_ffmpeg!(ret);
Ok(())
}
/// Close the output and write the trailer
pub unsafe fn close(self) -> Result<()> {
let ret = av_write_trailer(self.ctx);
bail_ffmpeg!(ret);
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::{generate_test_frame, Scaler};
use ffmpeg_sys_the_third::AVCodecID::AV_CODEC_ID_H264;
use ffmpeg_sys_the_third::AVPixelFormat::AV_PIX_FMT_YUV420P;
use ffmpeg_sys_the_third::AV_PROFILE_H264_MAIN;
#[test]
fn encode_mkv() -> Result<()> {
unsafe {
let path = PathBuf::from("test.mp4");
let frame = generate_test_frame();
// convert frame to YUV
let mut scaler = Scaler::new();
let frame = scaler.process_frame(
frame,
(*frame).width as u16,
(*frame).height as u16,
AV_PIX_FMT_YUV420P,
)?;
let mut encoder = Encoder::new(AV_CODEC_ID_H264)?
.with_width((*frame).width)
.with_height((*frame).height)
.with_pix_fmt(AV_PIX_FMT_YUV420P)
.with_bitrate(1_000_000)
.with_framerate(30.0)
.with_profile(AV_PROFILE_H264_MAIN)
.with_level(50)
.open(None)?;
let mut muxer = Muxer::new()
.with_output(&path, None, None)?
.with_stream_encoder(&encoder)?
.open()?;
let mut pts = 0;
for z in 0..100 {
(*frame).pts = pts;
for pkt in encoder.encode_frame(frame)? {
muxer.write_packet(pkt)?;
}
pts += 1;
}
// flush
for f_pk in encoder.encode_frame(ptr::null_mut())? {
muxer.write_packet(f_pk)?;
}
muxer.close()?;
}
Ok(())
}
}

20
src/resample.rs
View File

@ -1,10 +1,10 @@
use crate::bail_ffmpeg;
use crate::get_ffmpeg_error_msg;
use crate::return_ffmpeg_error;
use anyhow::Error;
use ffmpeg_sys_the_third::{
av_channel_layout_default, av_frame_alloc, av_frame_copy_props, av_frame_free,
swr_alloc_set_opts2, swr_convert_frame, swr_init, AVChannelLayout, AVFrame, AVSampleFormat,
SwrContext,
swr_alloc_set_opts2, swr_convert_frame, swr_free, swr_init, AVChannelLayout, AVFrame,
AVSampleFormat, SwrContext,
};
use libc::malloc;
use std::mem::transmute;
@ -17,6 +17,16 @@ pub struct Resample {
ctx: *mut SwrContext,
}
impl Drop for Resample {
fn drop(&mut self) {
unsafe {
if !self.ctx.is_null() {
swr_free(&mut self.ctx);
}
}
}
}
impl Resample {
pub fn new(format: AVSampleFormat, rate: u32, channels: usize) -> Self {
Self {
@ -45,10 +55,10 @@ impl Resample {
0,
ptr::null_mut(),
);
return_ffmpeg_error!(ret);
bail_ffmpeg!(ret);
let ret = swr_init(self.ctx);
return_ffmpeg_error!(ret);
bail_ffmpeg!(ret);
Ok(())
}

11
src/scale.rs
View File

@ -1,7 +1,7 @@
use std::mem::transmute;
use std::ptr;
use crate::{return_ffmpeg_error, rstr};
use crate::{bail_ffmpeg, rstr};
use anyhow::{bail, Error};
use ffmpeg_sys_the_third::{
av_frame_alloc, av_frame_copy_props, av_get_pix_fmt_name, sws_freeContext, sws_getContext,
@ -16,13 +16,12 @@ pub struct Scaler {
ctx: *mut SwsContext,
}
unsafe impl Send for Scaler {}
impl Drop for Scaler {
fn drop(&mut self) {
unsafe {
if !self.ctx.is_null() {
sws_freeContext(self.ctx);
self.ctx = ptr::null_mut();
}
}
}
}
@ -111,10 +110,10 @@ impl Scaler {
let dst_frame = av_frame_alloc();
let ret = av_frame_copy_props(dst_frame, frame);
return_ffmpeg_error!(ret);
bail_ffmpeg!(ret);
let ret = sws_scale_frame(self.ctx, dst_frame, frame);
return_ffmpeg_error!(ret);
bail_ffmpeg!(ret);
Ok(dst_frame)
}

9
src/transcode.rs Normal file
View File

@ -0,0 +1,9 @@
use crate::{Decoder, Demuxer, Encoder, Muxer, Scaler};
pub struct Transcoder {
demuxer: Demuxer,
decoder: Decoder,
scaler: Scaler,
encoder: Encoder,
muxer: Muxer,
}