286 lines
9.9 KiB
Rust
286 lines
9.9 KiB
Rust
//! VAAPI H.264 hardware encoding.
|
|
//!
|
|
//! Two-phase initialization:
|
|
//! - Phase 1 (`VaapiEncoder::new`): store config, no ffmpeg calls.
|
|
//! - Phase 2 (first call to `encode`): build filter graph using `hw_frames_ctx`
|
|
//! from the first DRM_PRIME frame to wire up hwmap, then open the codec.
|
|
//!
|
|
//! The split exists because the filter graph needs the DRM hardware device
|
|
//! context that kmsgrab attaches to each frame — that context isn't available
|
|
//! until the first frame arrives, so the graph can't be validated at
|
|
//! construction time.
|
|
|
|
use anyhow::{Context, Result};
|
|
use tracing::info;
|
|
|
|
/// Encoding configuration.
|
|
pub struct EncoderConfig {
|
|
pub width: u32,
|
|
pub height: u32,
|
|
pub fps: u32,
|
|
/// Quantization parameter (lower = higher quality). Default: 20
|
|
pub qp: u32,
|
|
/// Keyframe interval in frames. Default: 30 (1 keyframe/sec at 30fps)
|
|
pub gop_size: u32,
|
|
}
|
|
|
|
impl Default for EncoderConfig {
|
|
fn default() -> Self {
|
|
Self {
|
|
width: 1920,
|
|
height: 1080,
|
|
fps: 30,
|
|
qp: 20,
|
|
gop_size: 30,
|
|
}
|
|
}
|
|
}
|
|
|
|
/// VAAPI H.264 encoder with lazy filter graph initialization.
|
|
pub struct VaapiEncoder {
|
|
config: EncoderConfig,
|
|
inner: Option<EncoderInner>,
|
|
}
|
|
|
|
struct EncoderInner {
|
|
encoder: ffmpeg::encoder::Video,
|
|
filter_graph: ffmpeg::filter::Graph,
|
|
frame_count: u64,
|
|
time_base: ffmpeg::Rational,
|
|
}
|
|
|
|
impl VaapiEncoder {
|
|
/// Create an encoder from config. No ffmpeg resources are allocated yet.
|
|
pub fn new(config: EncoderConfig) -> Self {
|
|
Self { config, inner: None }
|
|
}
|
|
|
|
/// Encode a frame. On the first call, initializes the filter graph and
|
|
/// codec using the frame's `hw_frames_ctx`. Returns 0 or more packets.
|
|
pub fn encode(&mut self, frame: &ffmpeg::frame::Video) -> Result<Vec<EncodedPacket>> {
|
|
if self.inner.is_none() {
|
|
self.inner = Some(EncoderInner::open(&self.config, frame).context("encoder init")?);
|
|
}
|
|
self.inner.as_mut().unwrap().encode(frame)
|
|
}
|
|
|
|
/// Flush remaining packets out of the encoder.
|
|
pub fn flush(&mut self) -> Result<Vec<EncodedPacket>> {
|
|
match self.inner.as_mut() {
|
|
Some(inner) => inner.flush(),
|
|
None => Ok(vec![]),
|
|
}
|
|
}
|
|
}
|
|
|
|
impl EncoderInner {
|
|
fn open(config: &EncoderConfig, first_frame: &ffmpeg::frame::Video) -> Result<Self> {
|
|
let time_base = ffmpeg::Rational::new(1, config.fps as i32);
|
|
let filter_graph = Self::build_filter_graph(config, first_frame)?;
|
|
|
|
let codec = ffmpeg::encoder::find_by_name("h264_vaapi")
|
|
.context("h264_vaapi encoder not found")?;
|
|
|
|
let mut encoder_ctx = ffmpeg::codec::context::Context::new_with_codec(codec)
|
|
.encoder()
|
|
.video()?;
|
|
|
|
encoder_ctx.set_width(config.width);
|
|
encoder_ctx.set_height(config.height);
|
|
encoder_ctx.set_time_base(time_base);
|
|
encoder_ctx.set_frame_rate(Some(ffmpeg::Rational::new(config.fps as i32, 1)));
|
|
encoder_ctx.set_gop(config.gop_size);
|
|
encoder_ctx.set_max_b_frames(0);
|
|
encoder_ctx.set_format(ffmpeg::format::Pixel::VAAPI);
|
|
|
|
let mut opts = ffmpeg::Dictionary::new();
|
|
opts.set("qp", &config.qp.to_string());
|
|
|
|
let encoder = encoder_ctx
|
|
.open_with(opts)
|
|
.context("open h264_vaapi encoder")?;
|
|
|
|
info!(
|
|
"VAAPI encoder opened: {}x{} @ {}fps, qp={}, gop={}",
|
|
config.width, config.height, config.fps, config.qp, config.gop_size,
|
|
);
|
|
|
|
Ok(Self { encoder, filter_graph, frame_count: 0, time_base })
|
|
}
|
|
|
|
/// Build the filter graph: DRM_PRIME → hwmap(vaapi) → scale_vaapi → buffersink.
|
|
///
|
|
/// Mirrors the ffmpeg CLI approach:
|
|
/// -init_hw_device drm=drm:/dev/dri/card0
|
|
/// -init_hw_device vaapi=va@drm
|
|
/// -filter_hw_device va
|
|
///
|
|
/// We create both devices explicitly and set the VAAPI device on the graph,
|
|
/// then attach the frame's hw_frames_ctx to the buffersrc so hwmap can map
|
|
/// DRM_PRIME frames to VAAPI surfaces.
|
|
fn build_filter_graph(
|
|
config: &EncoderConfig,
|
|
first_frame: &ffmpeg::frame::Video,
|
|
) -> Result<ffmpeg::filter::Graph> {
|
|
let input_width = first_frame.width();
|
|
let input_height = first_frame.height();
|
|
|
|
// Create DRM device, then derive VAAPI from it (like -init_hw_device vaapi=va@drm)
|
|
let (drm_device, vaapi_device) = unsafe {
|
|
let mut drm_ref: *mut ffmpeg::ffi::AVBufferRef = std::ptr::null_mut();
|
|
let dev_path = std::ffi::CString::new("/dev/dri/card0").unwrap();
|
|
let ret = ffmpeg::ffi::av_hwdevice_ctx_create(
|
|
&mut drm_ref,
|
|
ffmpeg::ffi::AVHWDeviceType::AV_HWDEVICE_TYPE_DRM,
|
|
dev_path.as_ptr(),
|
|
std::ptr::null_mut(),
|
|
0,
|
|
);
|
|
anyhow::ensure!(ret >= 0, "av_hwdevice_ctx_create(drm): {ret}");
|
|
info!("DRM device created");
|
|
|
|
let mut vaapi_ref: *mut ffmpeg::ffi::AVBufferRef = std::ptr::null_mut();
|
|
let ret = ffmpeg::ffi::av_hwdevice_ctx_create_derived(
|
|
&mut vaapi_ref,
|
|
ffmpeg::ffi::AVHWDeviceType::AV_HWDEVICE_TYPE_VAAPI,
|
|
drm_ref,
|
|
0,
|
|
);
|
|
anyhow::ensure!(ret >= 0, "av_hwdevice_ctx_create_derived(vaapi): {ret}");
|
|
info!("VAAPI device derived from DRM");
|
|
|
|
(drm_ref, vaapi_ref)
|
|
};
|
|
|
|
let mut graph = ffmpeg::filter::Graph::new();
|
|
|
|
let args = format!(
|
|
"video_size={}x{}:pix_fmt={}:time_base=1/{}:pixel_aspect=1/1",
|
|
input_width,
|
|
input_height,
|
|
ffmpeg::format::Pixel::DRM_PRIME as i32,
|
|
config.fps,
|
|
);
|
|
|
|
graph
|
|
.add(&ffmpeg::filter::find("buffer").unwrap(), "in", &args)
|
|
.context("add buffersrc")?;
|
|
graph
|
|
.add(&ffmpeg::filter::find("buffersink").unwrap(), "out", "")
|
|
.context("add buffersink")?;
|
|
|
|
let filter_spec = format!(
|
|
"hwmap=derive_device=vaapi,scale_vaapi=w={}:h={}:format=nv12",
|
|
config.width, config.height,
|
|
);
|
|
|
|
graph.output("in", 0)?.input("out", 0)?.parse(&filter_spec)?;
|
|
|
|
// Set the VAAPI device on filter contexts that need it
|
|
// (equivalent to -filter_hw_device va in the CLI)
|
|
unsafe {
|
|
let graph_ptr = graph.as_mut_ptr();
|
|
for i in 0..(*graph_ptr).nb_filters {
|
|
let fctx = *(*graph_ptr).filters.add(i as usize);
|
|
(*fctx).hw_device_ctx = ffmpeg::ffi::av_buffer_ref(vaapi_device);
|
|
}
|
|
}
|
|
|
|
// Attach hw_frames_ctx from the first DRM_PRIME frame to the buffersrc
|
|
unsafe {
|
|
let hw_frames_ctx = (*first_frame.as_ptr()).hw_frames_ctx;
|
|
anyhow::ensure!(
|
|
!hw_frames_ctx.is_null(),
|
|
"first KMS frame has no hw_frames_ctx"
|
|
);
|
|
|
|
let mut buffersrc = graph.get("in").unwrap();
|
|
let par = ffmpeg::ffi::av_buffersrc_parameters_alloc();
|
|
anyhow::ensure!(!par.is_null(), "av_buffersrc_parameters_alloc OOM");
|
|
(*par).hw_frames_ctx = ffmpeg::ffi::av_buffer_ref(hw_frames_ctx);
|
|
let ret = ffmpeg::ffi::av_buffersrc_parameters_set(buffersrc.as_mut_ptr(), par);
|
|
ffmpeg::ffi::av_free(par as *mut _);
|
|
anyhow::ensure!(ret >= 0, "av_buffersrc_parameters_set: {ret}");
|
|
}
|
|
|
|
graph.validate().context("validate filter graph")?;
|
|
info!("Filter graph ready: {filter_spec}");
|
|
|
|
// Keep device refs alive — they're ref-counted by the graph now,
|
|
// but we drop our refs here.
|
|
unsafe {
|
|
ffmpeg::ffi::av_buffer_unref(&mut (drm_device as *mut _));
|
|
ffmpeg::ffi::av_buffer_unref(&mut (vaapi_device as *mut _));
|
|
}
|
|
|
|
Ok(graph)
|
|
}
|
|
|
|
fn encode(&mut self, frame: &ffmpeg::frame::Video) -> Result<Vec<EncodedPacket>> {
|
|
self.filter_graph
|
|
.get("in")
|
|
.unwrap()
|
|
.source()
|
|
.add(frame)
|
|
.context("buffersrc add")?;
|
|
|
|
let mut packets = Vec::new();
|
|
let mut filtered = ffmpeg::frame::Video::empty();
|
|
|
|
while self
|
|
.filter_graph
|
|
.get("out")
|
|
.unwrap()
|
|
.sink()
|
|
.frame(&mut filtered)
|
|
.is_ok()
|
|
{
|
|
filtered.set_pts(Some(self.frame_count as i64));
|
|
self.frame_count += 1;
|
|
|
|
self.encoder.send_frame(&filtered)?;
|
|
|
|
let mut encoded = ffmpeg::Packet::empty();
|
|
while self.encoder.receive_packet(&mut encoded).is_ok() {
|
|
packets.push(EncodedPacket {
|
|
data: encoded.data().unwrap_or(&[]).to_vec(),
|
|
pts: encoded.pts().unwrap_or(0),
|
|
dts: encoded.dts().unwrap_or(0),
|
|
keyframe: encoded.is_key(),
|
|
time_base_num: self.time_base.numerator() as u32,
|
|
time_base_den: self.time_base.denominator() as u32,
|
|
});
|
|
}
|
|
}
|
|
|
|
Ok(packets)
|
|
}
|
|
|
|
fn flush(&mut self) -> Result<Vec<EncodedPacket>> {
|
|
self.encoder.send_eof()?;
|
|
let mut packets = Vec::new();
|
|
let mut encoded = ffmpeg::Packet::empty();
|
|
while self.encoder.receive_packet(&mut encoded).is_ok() {
|
|
packets.push(EncodedPacket {
|
|
data: encoded.data().unwrap_or(&[]).to_vec(),
|
|
pts: encoded.pts().unwrap_or(0),
|
|
dts: encoded.dts().unwrap_or(0),
|
|
keyframe: encoded.is_key(),
|
|
time_base_num: self.time_base.numerator() as u32,
|
|
time_base_den: self.time_base.denominator() as u32,
|
|
});
|
|
}
|
|
Ok(packets)
|
|
}
|
|
}
|
|
|
|
/// An encoded video packet ready for transport.
|
|
pub struct EncodedPacket {
|
|
pub data: Vec<u8>,
|
|
pub pts: i64,
|
|
pub dts: i64,
|
|
pub keyframe: bool,
|
|
pub time_base_num: u32,
|
|
pub time_base_den: u32,
|
|
}
|