Skip to content

Commit

Permalink
Revise transcode_aac example
Browse files Browse the repository at this point in the history
  • Loading branch information
@ldm0
ldm0 committed Jan 23, 2024
1 parent e7c998c commit f57dfe3
Showing 1 changed file with 78 additions and 55 deletions.
133 changes: 78 additions & 55 deletions tests/transcode_aac.rs
View file
Original file line number Diff line number Diff line change
@@ -1,7 +1,6 @@
// Due to usage of new channel layout api
//! RIIR: https://github.com/FFmpeg/FFmpeg/blob/master/doc/examples/transcode_aac.c
use anyhow::{bail, Context as AnyhowContext, Result};
use cstr::cstr;
use once_cell::sync::Lazy as SyncLazy;
use rsmpeg::{
avcodec::{AVCodec, AVCodecContext},
avformat::{AVFormatContextInput, AVFormatContextOutput},
Expand All @@ -10,18 +9,30 @@ use rsmpeg::{
ffi,
swresample::SwrContext,
};
use std::{ffi::CStr, sync::Mutex};
use std::{
ffi::CStr,
sync::atomic::{AtomicI64, Ordering},
};

/// The output bit rate in bit/s
const OUTPUT_BIT_RATE: i64 = 96000;
/// The number of output channels
const OUTPUT_CHANNELS: i32 = 2;

fn open_input_file(input_file: &CStr) -> Result<(AVFormatContextInput, AVCodecContext, usize)> {
let input_format_context =
AVFormatContextInput::open(input_file, None, &mut None).context("Failed to get encoder")?;
let input_format_context = AVFormatContextInput::open(input_file, None, &mut None)
.context("Could not open input file")?;
let (audio_index, decoder) = input_format_context
.find_best_stream(ffi::AVMediaType_AVMEDIA_TYPE_AUDIO)?
.context("Failed to find audio stream")?;

let stream = &input_format_context.streams()[audio_index];
let mut decode_context = AVCodecContext::new(&decoder);
decode_context.apply_codecpar(&input_format_context.streams()[audio_index].codecpar())?;
decode_context.open(None)?;
decode_context.apply_codecpar(&stream.codecpar())?;
decode_context
.open(None)
.context("Could not open input codec")?;
decode_context.set_pkt_timebase(stream.time_base);
Ok((input_format_context, decode_context, audio_index))
}

Expand All @@ -39,19 +50,13 @@ fn open_output_file(

let mut encode_context = AVCodecContext::new(&encode_codec);

const OUTPUT_CHANNELS: i32 = 2;
const OUTPUT_BIT_RATE: i64 = 96000;
// Set the basic encoder parameters.
// The input file's sample rate is used to avoid a sample rate conversion.
encode_context.set_channels(OUTPUT_CHANNELS);
encode_context.set_ch_layout(AVChannelLayout::from_nb_channels(OUTPUT_CHANNELS).into_inner());
encode_context.set_sample_rate(decode_context.sample_rate);
encode_context.set_sample_fmt(encode_codec.sample_fmts().unwrap()[0]);
encode_context.set_bit_rate(OUTPUT_BIT_RATE);

// Allow the use of the experimental AAC encoder.
encode_context.set_strict_std_compliance(ffi::FF_COMPLIANCE_EXPERIMENTAL);

// Open the encoder for the audio stream to use it later.
encode_context.open(None)?;

Expand All @@ -60,7 +65,7 @@ fn open_output_file(
let mut stream = output_format_context.new_stream();
stream.set_codecpar(encode_context.extract_codecpar());
// Set the sample rate for the container.
stream.set_time_base(ra(decode_context.sample_rate, 1));
stream.set_time_base(ra(1, decode_context.sample_rate));
}

Ok((output_format_context, encode_context))
Expand All @@ -78,38 +83,41 @@ fn init_resampler(
decode_context.sample_fmt,
decode_context.sample_rate,
)
.context("Swrcontext parameters incorrect.")?;
resample_context.init()?;
.context("Could not allocate resample context")?;
resample_context
.init()
.context("Could not open resample context")?;
Ok(resample_context)
}

fn add_samples_to_fifo(
fifo: &mut AVAudioFifo,
samples_buffer: &AVSamples,
num_samples: i32,
frame_size: i32,
) -> Result<()> {
fifo.realloc(fifo.size() + num_samples);
if unsafe { fifo.write(samples_buffer.audio_data.as_ptr(), num_samples) }? < num_samples {
bail!("samples doesn't all written.");
}
fifo.realloc(fifo.size() + frame_size);
unsafe { fifo.write(samples_buffer.audio_data.as_ptr(), frame_size) }
.context("Could not write data to FIFO")?;
Ok(())
}

fn create_output_frame(
fn init_output_frame(
nb_samples: i32,
ch_layout: ffi::AVChannelLayout,
sample_fmt: i32,
sample_rate: i32,
) -> AVFrame {
) -> Result<AVFrame> {
let mut frame = AVFrame::new();
frame.set_nb_samples(nb_samples);
frame.set_ch_layout(ch_layout);
frame.set_format(sample_fmt);
frame.set_sample_rate(sample_rate);

frame.alloc_buffer().unwrap();

frame
.get_buffer(0)
.context("Could not allocate output frame samples")?;

Ok(frame)
}

/// Return boolean: if data is written.
Expand All @@ -118,12 +126,10 @@ fn encode_audio_frame(
output_format_context: &mut AVFormatContextOutput,
encode_context: &mut AVCodecContext,
) -> Result<()> {
static PTS: SyncLazy<Mutex<i64>> = SyncLazy::new(|| Mutex::new(0));
static PTS: AtomicI64 = AtomicI64::new(0);

if let Some(frame) = frame.as_mut() {
let mut pts = PTS.lock().unwrap();
frame.set_pts(*pts);
*pts += frame.nb_samples as i64;
frame.set_pts(PTS.fetch_add(frame.nb_samples as i64, Ordering::Relaxed));
}

encode_context.send_frame(frame.as_ref())?;
Expand All @@ -133,10 +139,12 @@ fn encode_audio_frame(
Err(RsmpegError::EncoderDrainError) | Err(RsmpegError::EncoderFlushedError) => {
break;
}
Err(e) => bail!(e),
Err(e) => Err(e).context("Could not encode frame")?,
};

output_format_context.write_frame(&mut packet)?;
output_format_context
.write_frame(&mut packet)
.context("Could not write frame")?;
}
Ok(())
}
Expand All @@ -146,15 +154,15 @@ fn load_encode_and_write(
output_format_context: &mut AVFormatContextOutput,
encode_context: &mut AVCodecContext,
) -> Result<()> {
let nb_samples = fifo.size().min(encode_context.frame_size);
let mut frame = create_output_frame(
nb_samples,
let frame_size = fifo.size().min(encode_context.frame_size);
let mut frame = init_output_frame(
frame_size,
encode_context.ch_layout().clone().into_inner(),
encode_context.sample_fmt,
encode_context.sample_rate,
);
if unsafe { fifo.read(frame.data_mut().as_mut_ptr(), nb_samples)? } < nb_samples {
bail!("samples doesn't all read");
)?;
if unsafe { fifo.read(frame.data_mut().as_mut_ptr(), frame_size)? } < frame_size {
bail!("Could not read data from FIFO");
}
encode_audio_frame(Some(frame), output_format_context, encode_context)?;
Ok(())
Expand All @@ -170,27 +178,36 @@ fn transcode_aac(input_file: &CStr, output_file: &CStr) -> Result<()> {
open_output_file(output_file, &decode_context)?;

// Initialize the resampler to be able to convert audio sample formats.
let resample_context = init_resampler(&mut decode_context, &mut encode_context)?;
let mut resample_context = init_resampler(&mut decode_context, &mut encode_context)?;

// Initialize the FIFO buffer to store audio samples to be encoded.
let mut fifo = AVAudioFifo::new(encode_context.sample_fmt, encode_context.channels, 1);
let mut fifo = AVAudioFifo::new(
encode_context.sample_fmt,
encode_context.ch_layout.nb_channels,
1,
);

// Write the header of the output file container.
output_format_context.write_header(&mut None)?;

let output_nb_sample = encode_context.frame_size;
output_format_context
.write_header(&mut None)
.context("Could not write output file header")?;

// Loop as long as we have input samples to read or output samples to write;
// abort as soon as we have neither.
loop {
let output_frame_size = encode_context.frame_size;

loop {
// We get enough audio samples.
if fifo.size() >= output_nb_sample {
if fifo.size() >= output_frame_size {
break;
}

// Break when no more input packets.
let packet = match input_format_context.read_packet()? {
let packet = match input_format_context
.read_packet()
.context("Could not read frame")?
{
Some(x) => x,
None => break,
};
Expand All @@ -200,52 +217,58 @@ fn transcode_aac(input_file: &CStr, output_file: &CStr) -> Result<()> {
continue;
}

decode_context.send_packet(Some(&packet))?;
decode_context
.send_packet(Some(&packet))
.context("Could not send packet for decoding")?;

loop {
let frame = match decode_context.receive_frame() {
Ok(frame) => frame,
Err(RsmpegError::DecoderDrainError) | Err(RsmpegError::DecoderFlushedError) => {
break;
}
Err(e) => bail!(e),
Err(e) => Err(e).context("Could not decode frame")?,
};

let mut output_samples = AVSamples::new(
encode_context.channels,
encode_context.ch_layout.nb_channels,
frame.nb_samples,
encode_context.sample_fmt,
0,
)
.context("Create samples buffer failed.")?;

unsafe {
resample_context.convert(
&mut output_samples,
frame.extended_data as *const _,
frame.nb_samples,
)?;
resample_context
.convert(
output_samples.audio_data.as_mut_ptr(),
output_samples.nb_samples,
frame.extended_data as *const _,
frame.nb_samples,
)
.context("Could not convert input samples")?;
}

add_samples_to_fifo(&mut fifo, &output_samples, frame.nb_samples)?;
}
}

// If we still cannot get enough samples, break.
if fifo.size() < output_nb_sample {
if fifo.size() < output_frame_size {
break;
}

// Write frame as much as possible.
while fifo.size() >= output_nb_sample {
while fifo.size() >= output_frame_size {
load_encode_and_write(&mut fifo, &mut output_format_context, &mut encode_context)?;
}
}

// Flush encode context
encode_audio_frame(None, &mut output_format_context, &mut encode_context)?;

output_format_context.write_trailer().unwrap();
output_format_context.write_trailer()?;

Ok(())
}

Expand Down

0 comments on commit f57dfe3

Please sign in to comment.