refactor(avhw): split encoder module
Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
This commit is contained in:
dailz
@@ -0,0 +1,263 @@
|
||||
use std::path::Path;
|
||||
use std::slice;
|
||||
|
||||
use anyhow::{bail, Result};
|
||||
use ffmpeg_next as ff;
|
||||
use ffmpeg_next::ffi;
|
||||
|
||||
use super::filter::build_swenc_filter_graph;
|
||||
use super::{ff_err, AvHwDevCtx, AvHwFrameCtx};
|
||||
use super::{BitrateCommand, CpuNv12Frame, ResolutionChange};
|
||||
|
||||
pub struct SwEncImport {
|
||||
hw_dev: AvHwDevCtx,
|
||||
frames_rgb: AvHwFrameCtx,
|
||||
filter_graph: ff::filter::Graph,
|
||||
width: u32,
|
||||
height: u32,
|
||||
enc_width: u32,
|
||||
enc_height: u32,
|
||||
fps: u32,
|
||||
resolution_rx: Option<crossbeam_channel::Receiver<BitrateCommand>>,
|
||||
encoder_resolution_tx: Option<crossbeam_channel::Sender<ResolutionChange>>,
|
||||
}
|
||||
|
||||
impl SwEncImport {
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn new(
|
||||
drm_device: &Path,
|
||||
width: u32,
|
||||
height: u32,
|
||||
enc_width: u32,
|
||||
enc_height: u32,
|
||||
fps: u32,
|
||||
) -> Result<Self> {
|
||||
let hw_dev = AvHwDevCtx::new_vaapi(drm_device)?;
|
||||
let frames_rgb =
|
||||
AvHwFrameCtx::for_capture(&hw_dev, width, height, ff::format::Pixel::BGRA)?;
|
||||
let filter_graph = build_swenc_filter_graph(
|
||||
&hw_dev,
|
||||
&frames_rgb,
|
||||
width,
|
||||
height,
|
||||
enc_width,
|
||||
enc_height,
|
||||
fps,
|
||||
)?;
|
||||
|
||||
Ok(Self {
|
||||
hw_dev,
|
||||
frames_rgb,
|
||||
filter_graph,
|
||||
width,
|
||||
height,
|
||||
enc_width,
|
||||
enc_height,
|
||||
fps,
|
||||
resolution_rx: None,
|
||||
encoder_resolution_tx: None,
|
||||
})
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn new_with_resolution_control(
|
||||
drm_device: &Path,
|
||||
width: u32,
|
||||
height: u32,
|
||||
enc_width: u32,
|
||||
enc_height: u32,
|
||||
fps: u32,
|
||||
resolution_rx: crossbeam_channel::Receiver<BitrateCommand>,
|
||||
encoder_resolution_tx: crossbeam_channel::Sender<ResolutionChange>,
|
||||
) -> Result<Self> {
|
||||
let mut this = Self::new(drm_device, width, height, enc_width, enc_height, fps)?;
|
||||
this.resolution_rx = Some(resolution_rx);
|
||||
this.encoder_resolution_tx = Some(encoder_resolution_tx);
|
||||
Ok(this)
|
||||
}
|
||||
|
||||
pub fn frames_rgb(&self) -> &AvHwFrameCtx {
|
||||
let _ = self.hw_dev.as_ptr();
|
||||
&self.frames_rgb
|
||||
}
|
||||
|
||||
pub fn import_and_scale(&mut self, hw_frame: &ff::frame::Video) -> Result<CpuNv12Frame> {
|
||||
self.poll_resolution_commands()?;
|
||||
|
||||
let mut filter_src_ctx = self
|
||||
.filter_graph
|
||||
.get("in")
|
||||
.ok_or_else(|| anyhow::anyhow!("filter 'in' not found"))?;
|
||||
let mut filter_src = filter_src_ctx.source();
|
||||
let mut filter_sink_ctx = self
|
||||
.filter_graph
|
||||
.get("out")
|
||||
.ok_or_else(|| anyhow::anyhow!("filter 'out' not found"))?;
|
||||
let mut filter_sink = filter_sink_ctx.sink();
|
||||
|
||||
filter_src
|
||||
.add(hw_frame)
|
||||
.map_err(|e| anyhow::anyhow!("software pipeline filter source add failed: {e}"))?;
|
||||
|
||||
let mut first = None;
|
||||
let mut extra_count = 0usize;
|
||||
loop {
|
||||
let mut filtered = ff::frame::Video::empty();
|
||||
match filter_sink.frame(&mut filtered) {
|
||||
Ok(()) => {
|
||||
if filtered.pts().is_none() {
|
||||
filtered.set_pts(hw_frame.pts());
|
||||
}
|
||||
let cpu_frame = self.transfer_filtered_to_cpu(&filtered)?;
|
||||
if first.is_none() {
|
||||
first = Some(cpu_frame);
|
||||
} else {
|
||||
extra_count += 1;
|
||||
}
|
||||
}
|
||||
Err(ff::Error::Other { errno }) if errno == ffi::EAGAIN => break,
|
||||
Err(e) => bail!("software pipeline filter sink get frame failed: {e}"),
|
||||
}
|
||||
}
|
||||
|
||||
if extra_count > 0 {
|
||||
tracing::warn!(
|
||||
"software import filter produced {extra_count} extra frame(s); dropping extras"
|
||||
);
|
||||
}
|
||||
|
||||
first.ok_or_else(|| anyhow::anyhow!("software pipeline produced no scaled frame"))
|
||||
}
|
||||
|
||||
pub fn flush_import(&mut self) -> Result<Vec<CpuNv12Frame>> {
|
||||
let mut filter_src_ctx = self
|
||||
.filter_graph
|
||||
.get("in")
|
||||
.ok_or_else(|| anyhow::anyhow!("filter 'in' not found"))?;
|
||||
let mut filter_src = filter_src_ctx.source();
|
||||
if let Err(e) = filter_src.flush() {
|
||||
tracing::debug!("filter source flush error: {e}");
|
||||
}
|
||||
|
||||
let mut filter_sink_ctx = self
|
||||
.filter_graph
|
||||
.get("out")
|
||||
.ok_or_else(|| anyhow::anyhow!("filter 'out' not found"))?;
|
||||
let mut filter_sink = filter_sink_ctx.sink();
|
||||
let mut frames = Vec::new();
|
||||
loop {
|
||||
let mut filtered = ff::frame::Video::empty();
|
||||
match filter_sink.frame(&mut filtered) {
|
||||
Ok(()) => frames.push(self.transfer_filtered_to_cpu(&filtered)?),
|
||||
Err(_) => break,
|
||||
}
|
||||
}
|
||||
|
||||
Ok(frames)
|
||||
}
|
||||
|
||||
fn poll_resolution_commands(&mut self) -> Result<()> {
|
||||
let Some(rx) = self.resolution_rx.as_ref().cloned() else {
|
||||
return Ok(());
|
||||
};
|
||||
|
||||
let mut requested = None;
|
||||
while let Ok(cmd) = rx.try_recv() {
|
||||
match cmd {
|
||||
BitrateCommand::UpdateResolution { width, height } => {
|
||||
requested = Some((width & !1, height & !1));
|
||||
}
|
||||
BitrateCommand::UpdateBitrate { .. } => {}
|
||||
BitrateCommand::ForceKeyframe => {}
|
||||
}
|
||||
}
|
||||
|
||||
let Some((width, height)) = requested else {
|
||||
return Ok(());
|
||||
};
|
||||
if width == self.enc_width && height == self.enc_height {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
tracing::info!(
|
||||
from = format_args!("{}x{}", self.enc_width, self.enc_height),
|
||||
to = format_args!("{}x{}", width, height),
|
||||
"rebuilding software import filter graph for resolution change"
|
||||
);
|
||||
let _ = self.flush_import();
|
||||
self.filter_graph = build_swenc_filter_graph(
|
||||
&self.hw_dev,
|
||||
&self.frames_rgb,
|
||||
self.width,
|
||||
self.height,
|
||||
width,
|
||||
height,
|
||||
self.fps,
|
||||
)?;
|
||||
self.enc_width = width;
|
||||
self.enc_height = height;
|
||||
|
||||
if let Some(tx) = &self.encoder_resolution_tx {
|
||||
tx.send(ResolutionChange { width, height })
|
||||
.map_err(|_| anyhow::anyhow!("encoder resolution channel disconnected"))?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn transfer_filtered_to_cpu(&self, filtered: &ff::frame::Video) -> Result<CpuNv12Frame> {
|
||||
// SAFETY: av_frame_alloc returns a newly allocated AVFrame or null,
|
||||
// which is checked below.
|
||||
let mut sw_nv12 = unsafe { ffi::av_frame_alloc() };
|
||||
if sw_nv12.is_null() {
|
||||
bail!("av_frame_alloc failed for NV12 transfer frame");
|
||||
}
|
||||
|
||||
// SAFETY: sw_nv12 is an allocated destination frame; filtered is a valid VAAPI NV12
|
||||
// surface produced by scale_vaapi at encoder dimensions.
|
||||
let transfer_ret = unsafe { ffi::av_hwframe_transfer_data(sw_nv12, filtered.as_ptr(), 0) };
|
||||
if transfer_ret < 0 {
|
||||
// SAFETY: sw_nv12 was allocated above and has not been freed yet.
|
||||
unsafe { ffi::av_frame_free(&mut sw_nv12) };
|
||||
bail!(
|
||||
"av_hwframe_transfer_data failed for GPU-downscaled frame: {}",
|
||||
ff_err(transfer_ret)
|
||||
);
|
||||
}
|
||||
|
||||
// SAFETY: sw_nv12 was filled by av_hwframe_transfer_data. NV12 planes 0 and 1 are
|
||||
// initialized for enc_width x enc_height; linesize values define each row's byte span.
|
||||
let frame = unsafe {
|
||||
let y_ptr = (*sw_nv12).data[0];
|
||||
let uv_ptr = (*sw_nv12).data[1];
|
||||
if y_ptr.is_null() || uv_ptr.is_null() {
|
||||
ffi::av_frame_free(&mut sw_nv12);
|
||||
bail!("NV12 transfer frame missing Y/UV plane data");
|
||||
}
|
||||
let y_stride = (*sw_nv12).linesize[0] as usize;
|
||||
let uv_stride = (*sw_nv12).linesize[1] as usize;
|
||||
if (*sw_nv12).width != self.enc_width as i32
|
||||
|| (*sw_nv12).height != self.enc_height as i32
|
||||
{
|
||||
ffi::av_frame_free(&mut sw_nv12);
|
||||
bail!("NV12 transfer frame has unexpected dimensions");
|
||||
}
|
||||
let y_len = y_stride * self.enc_height as usize;
|
||||
let uv_len = uv_stride * (self.enc_height as usize / 2);
|
||||
let y_data = slice::from_raw_parts(y_ptr, y_len).to_vec();
|
||||
let uv_data = slice::from_raw_parts(uv_ptr, uv_len).to_vec();
|
||||
let pts = filtered.pts().unwrap_or(0);
|
||||
ffi::av_frame_free(&mut sw_nv12);
|
||||
CpuNv12Frame {
|
||||
y_data,
|
||||
uv_data,
|
||||
y_stride,
|
||||
uv_stride,
|
||||
pts,
|
||||
capture_time: std::time::Instant::now(),
|
||||
}
|
||||
};
|
||||
|
||||
Ok(frame)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user