#![cfg(feature = "ffmpeg")]
use anyhow::{Result, anyhow};
use bytes::Bytes;
use std::collections::VecDeque;
use std::ffi::CString;
use ffmpeg::codec::{self, encoder};
use ffmpeg::ffi as sys;
use ffmpeg::format::Pixel;
use ffmpeg::util::frame::video::Video as VideoFrameFfmpeg;
use ffmpeg_next as ffmpeg;
use super::{EncodedPacket, Encoder, EncoderConfig};
use crate::frame::{PixelFormat, TransferFn, VideoFrame};
const AV1_ENCODER_PREFERENCE: &[&str] = &[
"av1_nvenc",
"av1_amf",
"av1_qsv",
"av1_vaapi",
"libsvtav1",
"libaom-av1",
"librav1e",
];
fn init_ffmpeg() {
static INIT: std::sync::OnceLock<()> = std::sync::OnceLock::new();
INIT.get_or_init(|| {
let _ = ffmpeg::init();
ffmpeg::util::log::set_level(ffmpeg::util::log::Level::Warning);
});
}
fn transfer_to_cicp(t: TransferFn) -> u8 {
match t {
TransferFn::Bt709 => 1,
TransferFn::St2084 => 16,
TransferFn::AribStdB67 => 18,
_ => 1,
}
}
pub struct FfmpegEncoder {
encoder: encoder::Video,
engaged_name: String,
scratch: VideoFrameFfmpeg,
input_pix_fmt: Pixel,
pending: VecDeque<EncodedPacket>,
pts_counter: u64,
width: u32,
height: u32,
done: bool,
}
impl FfmpegEncoder {
pub fn new(config: EncoderConfig) -> Result<Self> {
init_ffmpeg();
let input_pix_fmt = match config.pixel_format {
PixelFormat::Yuv420p10le => Pixel::YUV420P10LE,
_ => Pixel::YUV420P,
};
let override_name = std::env::var("FFMPEG_AV1_ENCODER").ok();
let preference: Vec<&str> = match override_name.as_deref() {
Some(name) => vec![name],
None => AV1_ENCODER_PREFERENCE.iter().copied().collect(),
};
let mut last_err: Option<String> = None;
for enc_name in preference {
match try_open_encoder(enc_name, &config, input_pix_fmt) {
Ok((enc, scratch)) => {
tracing::info!(
encoder = enc_name,
width = config.width,
height = config.height,
pixel_format = ?config.pixel_format,
"FFmpeg AV1 encoder opened"
);
return Ok(Self {
encoder: enc,
engaged_name: enc_name.to_string(),
scratch,
input_pix_fmt,
pending: VecDeque::new(),
pts_counter: 0,
width: config.width,
height: config.height,
done: false,
});
}
Err(e) => {
tracing::debug!(
encoder = enc_name,
error = %e,
"FFmpeg encoder probe failed; trying next"
);
last_err = Some(format!("{enc_name}: {e}"));
}
}
}
Err(anyhow!(
"FFmpeg: no AV1 encoder from {:?} could open. Last error: {}",
AV1_ENCODER_PREFERENCE,
last_err.unwrap_or_else(|| "(no probes attempted)".to_string())
))
}
pub fn engaged(&self) -> &str {
&self.engaged_name
}
fn drain_packets(&mut self) -> Result<()> {
use ffmpeg::packet::Packet;
let mut pkt = Packet::empty();
loop {
match self.encoder.receive_packet(&mut pkt) {
Ok(()) => {
let is_key = pkt.is_key();
let pts = pkt.pts().unwrap_or(self.pts_counter as i64).max(0) as u64;
let data = pkt
.data()
.map(|b| Bytes::copy_from_slice(b))
.unwrap_or_default();
if !data.is_empty() {
self.pending.push_back(EncodedPacket {
data,
pts,
is_keyframe: is_key,
});
}
unsafe {
sys::av_packet_unref(pkt.as_mut_ptr());
}
}
Err(ffmpeg::Error::Other { errno }) if errno == ffmpeg::util::error::EAGAIN => {
return Ok(());
}
Err(ffmpeg::Error::Eof) => {
self.done = true;
return Ok(());
}
Err(e) => return Err(anyhow!("FFmpeg encoder: receive_packet: {e}")),
}
}
}
}
fn try_open_encoder(
enc_name: &str,
config: &EncoderConfig,
input_pix_fmt: Pixel,
) -> Result<(encoder::Video, VideoFrameFfmpeg)> {
let c_name = CString::new(enc_name).map_err(|e| anyhow!("bad encoder name: {e}"))?;
let ff_codec_ptr = unsafe { sys::avcodec_find_encoder_by_name(c_name.as_ptr()) };
if ff_codec_ptr.is_null() {
return Err(anyhow!(
"encoder '{enc_name}' not present in this libavcodec build"
));
}
let raw_ctx: *mut sys::AVCodecContext = unsafe { sys::avcodec_alloc_context3(ff_codec_ptr) };
if raw_ctx.is_null() {
return Err(anyhow!("avcodec_alloc_context3 returned null"));
}
unsafe {
let ctx = &mut *raw_ctx;
ctx.codec_id = sys::AVCodecID::AV_CODEC_ID_AV1;
ctx.width = config.width as i32;
ctx.height = config.height as i32;
ctx.pix_fmt = encoder_input_fmt(ff_codec_ptr, input_pix_fmt);
let fps = (config.frame_rate.max(1.0)).round() as i32;
ctx.time_base = sys::AVRational {
num: 1,
den: fps.max(1),
};
ctx.framerate = sys::AVRational {
num: fps.max(1),
den: 1,
};
ctx.gop_size = config.keyframe_interval.max(1) as i32;
ctx.bit_rate = (1_500_000i64).max(config.width as i64 * config.height as i64 / 200);
ctx.color_primaries = av_color_primaries(config.color_metadata.colour_primaries);
ctx.color_trc = av_color_trc(config.color_metadata.transfer);
ctx.colorspace = av_color_space(config.color_metadata.matrix_coefficients);
ctx.color_range = if config.color_metadata.full_range {
sys::AVColorRange::AVCOL_RANGE_JPEG
} else {
sys::AVColorRange::AVCOL_RANGE_MPEG
};
let mut opts: *mut sys::AVDictionary = std::ptr::null_mut();
set_quality_opts(&mut opts, enc_name, config)?;
let rc = sys::avcodec_open2(raw_ctx, ff_codec_ptr, &mut opts);
if !opts.is_null() {
sys::av_dict_free(&mut opts);
}
if rc < 0 {
sys::avcodec_free_context(&mut (raw_ctx as *mut _));
return Err(anyhow!("avcodec_open2 on '{enc_name}' returned rc={rc}"));
}
}
let ctx = unsafe { codec::Context::wrap(raw_ctx, None) };
let enc = ctx
.encoder()
.video()
.map_err(|e| anyhow!("encoder().video() on '{enc_name}': {e}"))?;
let scratch = VideoFrameFfmpeg::new(input_pix_fmt, config.width, config.height);
Ok((enc, scratch))
}
unsafe fn encoder_input_fmt(codec: *const sys::AVCodec, fallback: Pixel) -> sys::AVPixelFormat {
let pix_fmts = (*codec).pix_fmts;
if pix_fmts.is_null() {
return fallback.into();
}
let wanted = fallback.into();
let mut i = 0;
loop {
let fmt = *pix_fmts.offset(i);
if fmt == sys::AVPixelFormat::AV_PIX_FMT_NONE {
break;
}
if fmt == wanted {
return wanted;
}
i += 1;
}
*pix_fmts.offset(0)
}
unsafe fn set_quality_opts(
opts: &mut *mut sys::AVDictionary,
enc_name: &str,
config: &EncoderConfig,
) -> Result<()> {
use crate::encode::tuning::{
amf_av1_params, libaom_cq_for_target, nvenc_av1_params, qsv_av1_params, rav1e_params,
};
let set = |key: &str, val: &str| -> Result<()> {
let k = CString::new(key).unwrap();
let v = CString::new(val).unwrap();
let rc = sys::av_dict_set(opts, k.as_ptr(), v.as_ptr(), 0);
if rc < 0 {
return Err(anyhow!("av_dict_set {key}={val} rc={rc}"));
}
Ok(())
};
match enc_name {
"av1_nvenc" => {
let p = nvenc_av1_params(config.target, config.tier, config.width, config.height);
set("cq", &p.cq.to_string())?;
use crate::encode::tuning::NvencRateControl;
let rc_str = match p.rc_mode {
NvencRateControl::ConstQp => "constqp",
NvencRateControl::VbrTargetQuality => "vbr",
};
set("rc", rc_str)?;
let preset = match config.tier {
crate::encode::SpeedTier::Draft => "p5",
crate::encode::SpeedTier::Standard => "p6",
crate::encode::SpeedTier::Archive => "p7",
};
set("preset", preset)?;
set("tune", "hq")?;
if p.lookahead_depth > 0 {
set("rc-lookahead", &p.lookahead_depth.to_string())?;
}
if p.aq_strength > 0 {
set("spatial_aq", "1")?;
set("aq-strength", &p.aq_strength.to_string())?;
}
set("tile-columns", &p.num_tile_columns.to_string())?;
set("tile-rows", &p.num_tile_rows.to_string())?;
}
"av1_amf" => {
let p = amf_av1_params(config.target, config.tier, config.width, config.height);
set("qp_i", &p.q_index_intra.to_string())?;
set("qp_p", &p.q_index_inter.to_string())?;
use crate::encode::tuning::AmfRateControl;
let rc_str = match p.rc_mode {
AmfRateControl::Cqp => "cqp",
AmfRateControl::QualityVbr => "qvbr",
};
set("rc", rc_str)?;
if matches!(p.rc_mode, AmfRateControl::QualityVbr) {
set("qvbr_quality_level", &p.qvbr_quality.to_string())?;
}
use crate::encode::tuning::AmfQualityPreset;
let q_str = match p.quality_preset {
AmfQualityPreset::HighQuality => "quality",
AmfQualityPreset::Quality => "quality",
AmfQualityPreset::Balanced => "balanced",
AmfQualityPreset::Speed => "speed",
};
set("quality", q_str)?;
}
"av1_qsv" => {
let p = qsv_av1_params(config.target, config.tier, config.width, config.height);
use crate::encode::tuning::QsvRateControl;
match p.rc_mode {
QsvRateControl::Icq => {
set("global_quality", &p.icq_quality.to_string())?;
}
QsvRateControl::Cqp => {
set("q", &p.qp_i.to_string())?;
}
}
let preset = match p.target_usage {
1..=2 => "slow",
3..=4 => "medium",
_ => "veryfast",
};
set("preset", preset)?;
}
"av1_vaapi" => {
let q = if config.width >= 1920 {
let p = amf_av1_params(config.target, config.tier, config.width, config.height);
p.q_index_intra
} else {
let p = qsv_av1_params(config.target, config.tier, config.width, config.height);
p.qp_i as u8
};
set("qp", &q.to_string())?;
set("rc_mode", "CQP")?;
}
"libsvtav1" => {
let crf = libaom_cq_for_target(config.target);
set("crf", &crf.to_string())?;
let preset = match config.tier {
crate::encode::SpeedTier::Draft => "10",
crate::encode::SpeedTier::Standard => "7",
crate::encode::SpeedTier::Archive => "4",
};
set("preset", preset)?;
}
"libaom-av1" => {
let crf = libaom_cq_for_target(config.target);
set("crf", &crf.to_string())?;
set("b:v", "0")?; let cpu_used = match config.tier {
crate::encode::SpeedTier::Draft => "8",
crate::encode::SpeedTier::Standard => "4",
crate::encode::SpeedTier::Archive => "2",
};
set("cpu-used", cpu_used)?;
}
"librav1e" => {
let p = rav1e_params(config.target, config.tier, config.width, config.height);
set("qp", &p.quantizer.to_string())?;
set("speed", &p.speed_preset.to_string())?;
set("tile-rows", &p.tile_rows.to_string())?;
set("tile-columns", &p.tile_cols.to_string())?;
}
_ => {}
}
Ok(())
}
fn av_color_primaries(p: u8) -> sys::AVColorPrimaries {
match p {
1 => sys::AVColorPrimaries::AVCOL_PRI_BT709,
5 => sys::AVColorPrimaries::AVCOL_PRI_BT470BG,
6 => sys::AVColorPrimaries::AVCOL_PRI_SMPTE170M,
9 => sys::AVColorPrimaries::AVCOL_PRI_BT2020,
_ => sys::AVColorPrimaries::AVCOL_PRI_BT709,
}
}
fn av_color_trc(t: TransferFn) -> sys::AVColorTransferCharacteristic {
match transfer_to_cicp(t) {
16 => sys::AVColorTransferCharacteristic::AVCOL_TRC_SMPTE2084,
18 => sys::AVColorTransferCharacteristic::AVCOL_TRC_ARIB_STD_B67,
_ => sys::AVColorTransferCharacteristic::AVCOL_TRC_BT709,
}
}
fn av_color_space(m: u8) -> sys::AVColorSpace {
match m {
1 => sys::AVColorSpace::AVCOL_SPC_BT709,
6 => sys::AVColorSpace::AVCOL_SPC_SMPTE170M,
9 => sys::AVColorSpace::AVCOL_SPC_BT2020_NCL,
10 => sys::AVColorSpace::AVCOL_SPC_BT2020_CL,
_ => sys::AVColorSpace::AVCOL_SPC_BT709,
}
}
impl Encoder for FfmpegEncoder {
fn send_frame(&mut self, frame: &VideoFrame) -> Result<()> {
if self.done {
return Err(anyhow!("FFmpeg encoder: send_frame after flush"));
}
let w = self.width as usize;
let h = self.height as usize;
let bytes_per_sample = match self.input_pix_fmt {
Pixel::YUV420P10LE => 2,
_ => 1,
};
let y_len = w * h * bytes_per_sample;
let uv_w = w / 2;
let uv_h = h / 2;
let uv_len = uv_w * uv_h * bytes_per_sample;
let expected = y_len + 2 * uv_len;
if frame.data.len() < expected {
return Err(anyhow!(
"FFmpeg encoder: frame buffer too small ({} < expected {})",
frame.data.len(),
expected
));
}
for plane in 0..3 {
let (pw, ph, src_off) = if plane == 0 {
(w, h, 0)
} else {
(uv_w, uv_h, y_len + if plane == 2 { uv_len } else { 0 })
};
let stride = self.scratch.stride(plane) as usize;
let dst = self.scratch.data_mut(plane);
for row in 0..ph {
let src_start = src_off + row * pw * bytes_per_sample;
let src_end = src_start + pw * bytes_per_sample;
let dst_start = row * stride;
let dst_end = dst_start + pw * bytes_per_sample;
dst[dst_start..dst_end].copy_from_slice(&frame.data[src_start..src_end]);
}
}
unsafe {
let raw = self.scratch.as_mut_ptr();
(*raw).pts = self.pts_counter as i64;
}
self.pts_counter += 1;
self.encoder
.send_frame(&self.scratch)
.map_err(|e| anyhow!("FFmpeg encoder: send_frame: {e}"))?;
self.drain_packets()?;
Ok(())
}
fn flush(&mut self) -> Result<()> {
if self.done {
return Ok(());
}
self.encoder
.send_eof()
.map_err(|e| anyhow!("FFmpeg encoder: send_eof: {e}"))?;
self.drain_packets()?;
self.done = true;
Ok(())
}
fn receive_packet(&mut self) -> Result<Option<EncodedPacket>> {
if self.pending.is_empty() && self.done {
self.drain_packets()?;
}
Ok(self.pending.pop_front())
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::encode::{QualityTarget, SpeedTier};
use crate::frame::ColorMetadata;
fn test_config() -> EncoderConfig {
EncoderConfig {
width: 320,
height: 240,
frame_rate: 24.0,
quality: u8::MAX,
speed_preset: u8::MAX,
keyframe_interval: 48,
target: QualityTarget::Standard,
tier: SpeedTier::Standard,
threads: 0,
pixel_format: PixelFormat::Yuv420p,
color_metadata: ColorMetadata::default(),
}
}
#[test]
fn priority_list_starts_with_hw() {
assert_eq!(AV1_ENCODER_PREFERENCE[0], "av1_nvenc");
assert_eq!(AV1_ENCODER_PREFERENCE[1], "av1_amf");
assert_eq!(AV1_ENCODER_PREFERENCE[2], "av1_qsv");
assert!(AV1_ENCODER_PREFERENCE.contains(&"libsvtav1"));
assert!(AV1_ENCODER_PREFERENCE.contains(&"librav1e"));
}
#[test]
fn transfer_to_cicp_round_trip() {
assert_eq!(transfer_to_cicp(TransferFn::Bt709), 1);
assert_eq!(transfer_to_cicp(TransferFn::St2084), 16);
assert_eq!(transfer_to_cicp(TransferFn::AribStdB67), 18);
}
#[test]
fn set_quality_opts_pulls_cq_from_tuning_adapter_for_nvenc() {
use crate::encode::tuning::nvenc_av1_params;
let config = test_config();
let expected = nvenc_av1_params(config.target, config.tier, config.width, config.height);
assert_eq!(expected.cq, 30);
}
#[test]
fn set_quality_opts_pulls_crf_from_libaom_cq_for_svt_and_libaom() {
use crate::encode::tuning::libaom_cq_for_target;
assert_eq!(libaom_cq_for_target(QualityTarget::VisuallyLossless), 20);
assert_eq!(libaom_cq_for_target(QualityTarget::High), 27);
assert_eq!(libaom_cq_for_target(QualityTarget::Standard), 32);
assert_eq!(libaom_cq_for_target(QualityTarget::Low), 38);
}
#[test]
fn construct_encoder_smoke() {
match FfmpegEncoder::new(test_config()) {
Ok(enc) => {
let name = enc.engaged();
eprintln!("engaged encoder: {name}");
assert!(
AV1_ENCODER_PREFERENCE.contains(&name),
"engaged encoder {name} not in priority list"
);
}
Err(e) => eprintln!("skip: no AV1 encoder available: {e}"),
}
}
#[test]
fn encoder_env_override_forces_specific_backend() {
unsafe {
std::env::set_var("FFMPEG_AV1_ENCODER", "libsvtav1");
}
match FfmpegEncoder::new(test_config()) {
Ok(enc) => {
assert_eq!(enc.engaged(), "libsvtav1");
}
Err(_) => { }
}
unsafe {
std::env::remove_var("FFMPEG_AV1_ENCODER");
}
}
}