diff --git a/examples/rc_keyframes.rs b/examples/rc_keyframes.rs new file mode 100644 index 0000000..50c88a6 --- /dev/null +++ b/examples/rc_keyframes.rs @@ -0,0 +1,221 @@ +//! Example: Mid-Stream Key Frames Under Rate Control +//! +//! Encodes AV1 with CBR and a short GOP so key frames land mid-stream, then +//! verifies a full-stream ffmpeg decode plus PSNR against the source. NVIDIA +//! emitted undecodable mid-stream key frames unless the coding state is reset +//! on every key frame; the PSNR floor also catches the decodes-but-corrupt case. + +use pixelforge::{ + Codec, EncodeBitDepth, EncodeConfig, Encoder, InputImage, PixelFormat, RateControlMode, + VideoContextBuilder, +}; +use std::collections::VecDeque; +use std::fs::File; +use std::io::{Read, Write}; +use std::process::Command; + +const WIDTH: u32 = 320; +const HEIGHT: u32 = 240; +const FRAMES: u64 = 30; +/// Short GOP so key frames land mid-stream (frames 0, 10 and 20). +const GOP_SIZE: u32 = 10; +/// PSNR below this means a key frame corrupted the stream without a hard error. +const MIN_PSNR: f64 = 30.0; + +fn main() -> Result<(), Box> { + tracing_subscriber::fmt() + .with_env_filter( + tracing_subscriber::EnvFilter::try_from_default_env() + .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("warn")), + ) + .init(); + + println!("PixelForge mid-stream key frame (rate control) verification\n"); + + let input_path = format!("testdata/test_frames_{WIDTH}x{HEIGHT}_yuv420p.yuv"); + ensure_test_data("yuv420p", &input_path)?; + + let context = VideoContextBuilder::new() + .app_name("RC key frame verification") + .enable_validation(cfg!(debug_assertions)) + .build()?; + + if !context.supports_encode(Codec::AV1) { + println!("AV1: skipped (encode not supported)"); + return Ok(()); + } + + let output_filename = "output_rc_keyframes_AV1.obu"; + let decoded_filename = "decoded_rc_keyframes_AV1.yuv"; + + let config = EncodeConfig::av1(WIDTH, HEIGHT) + .with_rate_control(RateControlMode::Cbr) + .with_pixel_format(PixelFormat::Yuv420) + .with_bit_depth(EncodeBitDepth::Eight) + .with_gop_size(GOP_SIZE) + .with_b_frames(0); + + let mut encoder = Encoder::new(context.clone(), config)?; + let mut input_image = InputImage::new( + context.clone(), + Codec::AV1, + WIDTH, + HEIGHT, + EncodeBitDepth::Eight, + PixelFormat::Yuv420, + )?; + + let mut yuv_data = Vec::new(); + File::open(&input_path)?.read_to_end(&mut yuv_data)?; + let frame_size = (WIDTH * HEIGHT * 3 / 2) as usize; + + let mut output_file = File::create(output_filename)?; + let mut pending: VecDeque = VecDeque::new(); + let mut key_frames = 0u32; + + let drain_one = |pending: &mut VecDeque, + output_file: &mut File, + key_frames: &mut u32| + -> Result<(), Box> { + let packet = pollster::block_on(pending.pop_front().unwrap())?; + if packet.is_key_frame { + *key_frames += 1; + } + output_file.write_all(&packet.data)?; + Ok(()) + }; + + for i in 0..FRAMES { + let start = (i as usize) * frame_size; + let end = start + frame_size; + if end > yuv_data.len() { + break; + } + + let encoder_image = encoder.input_image(); + input_image.upload_yuv420_to(encoder_image, &yuv_data[start..end])?; + pending.push_back(encoder.encode(encoder_image)?); + while pending.len() > 2 { + drain_one(&mut pending, &mut output_file, &mut key_frames)?; + } + } + + encoder.flush()?; + while !pending.is_empty() { + drain_one(&mut pending, &mut output_file, &mut key_frames)?; + } + drop(output_file); + + if key_frames < 2 { + return Err(format!( + "only {key_frames} key frame(s) produced — the mid-stream key frame scenario never engaged" + ) + .into()); + } + + let psnr = decode_and_psnr(output_filename, decoded_filename, &input_path)?; + std::fs::remove_file(output_filename).ok(); + std::fs::remove_file(decoded_filename).ok(); + + if psnr < MIN_PSNR { + return Err(format!( + "full-stream PSNR {psnr:.2} dB below {MIN_PSNR} dB (a mid-stream key frame corrupted the stream)" + ) + .into()); + } + + println!("AV1: PASS — {key_frames} key frames under CBR, full-stream PSNR {psnr:.2} dB"); + Ok(()) +} + +/// Decode the bitstream to raw YUV and return its PSNR against the source. +fn decode_and_psnr( + bitstream: &str, + decoded: &str, + source: &str, +) -> Result> { + let status = Command::new("ffmpeg") + .args([ + "-hide_banner", + "-loglevel", + "error", + "-y", + "-i", + bitstream, + "-pix_fmt", + "yuv420p", + "-f", + "rawvideo", + decoded, + ]) + .output()?; + if !status.status.success() { + return Err(format!( + "ffmpeg decode failed: {}", + String::from_utf8_lossy(&status.stderr) + ) + .into()); + } + + let size = format!("{WIDTH}x{HEIGHT}"); + let output = Command::new("ffmpeg") + .args([ + "-hide_banner", + "-loglevel", + "info", + "-s", + &size, + "-pix_fmt", + "yuv420p", + "-f", + "rawvideo", + "-i", + source, + "-s", + &size, + "-pix_fmt", + "yuv420p", + "-f", + "rawvideo", + "-i", + decoded, + "-lavfi", + "psnr", + "-f", + "null", + "-", + ]) + .output()?; + let stderr = String::from_utf8_lossy(&output.stderr); + let pos = stderr + .find("average:") + .ok_or_else(|| format!("could not parse PSNR: {stderr}"))?; + let rest = &stderr[pos + 8..]; + let end = rest.find(' ').unwrap_or(rest.len()); + Ok(rest[..end].parse()?) +} + +fn ensure_test_data(pix_fmt: &str, path: &str) -> Result<(), Box> { + if std::path::Path::new(path).exists() { + return Ok(()); + } + println!("Generating {path}..."); + let status = Command::new("ffmpeg") + .args([ + "-f", + "lavfi", + "-i", + &format!("testsrc=duration=1:size={WIDTH}x{HEIGHT}:rate=30"), + "-pix_fmt", + pix_fmt, + "-f", + "rawvideo", + "-y", + path, + ]) + .output()?; + if !status.status.success() { + return Err(format!("failed to generate test data: {status:?}").into()); + } + Ok(()) +} diff --git a/src/encoder/av1/record.rs b/src/encoder/av1/record.rs index 622a3d5..12d6cfe 100644 --- a/src/encoder/av1/record.rs +++ b/src/encoder/av1/record.rs @@ -323,7 +323,11 @@ impl Av1 { } all_reference_slots.extend_from_slice(&reference_slots); - let is_first_frame = plan.is_first_frame(); + // Re-issue the coding-state RESET and rate-control setup on every key + // frame, not just the first: NVIDIA otherwise emits undecodable + // mid-stream key frames under active rate control (see rc_keyframes). + let should_reset_coding_state = + plan.is_first_frame() || (is_key_frame && !rc.is_disabled()); // Clamp GOP values to at least 1; a value of 0 is undefined in // Vulkan and causes undefined behavior on some drivers (RADV). let gop_frames = common.config.gop_size.max(1); @@ -336,7 +340,7 @@ impl Av1 { // Reset and write start timestamp reset_start_timestamp(common.device(), command_buffer, timestamp_query_pool); - let begin_coding_info = if is_first_frame { + let begin_coding_info = if should_reset_coding_state { vk::VideoBeginCodingInfoKHR::default() .video_session(common.session) .video_session_parameters(common.session_params) @@ -357,7 +361,7 @@ impl Av1 { .cmd_begin_video_coding(command_buffer, &begin_coding_info); } - if is_first_frame { + if should_reset_coding_state { let mut quality_level_info = vk::VideoEncodeQualityLevelInfoKHR::default().quality_level(0); let control_info = vk::VideoCodingControlInfoKHR::default()