Commit 4322cfb7 authored by Timothy B. Terriberry's avatar Timothy B. Terriberry Committed by Derek Buitenhuis
Browse files

Initial implementations of 2-pass rate control.

PUBLIC API CHANGE:
-Addition of EncoderStatus::NotReady
-Addition of Context::twopass_out()
-Addition of Context::twopass_bytes_needed()
-Addition of Context::twopass_in()
parent cc071643
...@@ -661,7 +661,7 @@ pub(crate) struct ContextInner<T: Pixel> { ...@@ -661,7 +661,7 @@ pub(crate) struct ContextInner<T: Pixel> {
/// A storage space for reordered frames. /// A storage space for reordered frames.
packet_data: Vec<u8>, packet_data: Vec<u8>,
segment_output_frameno_start: u64, segment_output_frameno_start: u64,
segment_input_frameno_start: u64, pub(crate) segment_input_frameno_start: u64,
keyframe_detector: SceneChangeDetector<T>, keyframe_detector: SceneChangeDetector<T>,
pub(crate) config: EncoderConfig, pub(crate) config: EncoderConfig,
rc_state: RCState, rc_state: RCState,
...@@ -691,6 +691,10 @@ pub enum EncoderStatus { ...@@ -691,6 +691,10 @@ pub enum EncoderStatus {
Encoded, Encoded,
/// Generic fatal error /// Generic fatal error
Failure, Failure,
/// A frame was encoded in the first pass of a 2-pass encode, but its stats
/// data was not retrieved with twopass_out(), or not enough stats data was
/// provided in the second pass of a 2-pass encode to encode the next frame.
NotReady
} }
pub struct Packet<T: Pixel> { pub struct Packet<T: Pixel> {
...@@ -741,6 +745,44 @@ impl<T: Pixel> Context<T> { ...@@ -741,6 +745,44 @@ impl<T: Pixel> Context<T> {
self.inner.send_frame(frame) self.inner.send_frame(frame)
} }
/// Retrieve the first-pass data of a two-pass encode for the frame that was
/// just encoded. This should be called BEFORE every call to receive_packet()
/// (including the very first one), even if no packet was produced by the
/// last call to receive_packet, if any (i.e., *EncoderStatus::Encoded* was
/// returned). It needs to be called once more after
/// *EncoderStatus::LimitReached* is returned, to retrieve the header that
/// should be written to the front of the stats file (overwriting the
/// placeholder header that was emitted at the start of encoding).
///
/// It is still safe to call this function when receive_packet() returns any
/// other error. It will return *None* instead of returning a duplicate copy
/// of the previous frame's data.
pub fn twopass_out(&mut self) -> Option<&[u8]> {
let params = self.inner.rc_state.get_twopass_out_params(&self.inner);
self.inner.rc_state.twopass_out(params)
}
/// Ask how many bytes of the stats file are needed before the next frame
/// of the second pass in a two-pass encode can be encoded. This is a lower
/// bound (more might be required), but if 0 is returned, then encoding can
/// proceed. This is just a hint to the application, and does not need to
/// be called for encoding the second pass to work, so long as the
/// application continues to provide more data to twopass_in() in a loop
/// until twopass_in() returns 0.
pub fn twopass_bytes_needed(&mut self) -> usize {
self.inner.rc_state.twopass_in(None).unwrap_or(0)
}
/// Provide stats data produced in the first pass of a two-pass encode to the
/// second pass. On success this returns the number of bytes of that data
/// which were consumed. When encoding the second pass of a two-pass encode,
/// this should be called repeatedly in a loop before every call to
/// receive_packet() (including the very first one) until no bytes are
/// consumed, or until twopass_bytes_needed() returns 0.
pub fn twopass_in(&mut self, buf: &[u8]) -> Result<usize, EncoderStatus> {
self.inner.rc_state.twopass_in(Some(buf)).or(Err(EncoderStatus::Failure))
}
pub fn receive_packet(&mut self) -> Result<Packet<T>, EncoderStatus> { pub fn receive_packet(&mut self) -> Result<Packet<T>, EncoderStatus> {
let inner = &mut self.inner; let inner = &mut self.inner;
let pool = &mut self.pool; let pool = &mut self.pool;
...@@ -970,8 +1012,12 @@ impl<T: Pixel> ContextInner<T> { ...@@ -970,8 +1012,12 @@ impl<T: Pixel> ContextInner<T> {
Ok((fi, true)) Ok((fi, true))
} }
pub(crate) fn done_processing(&self) -> bool {
self.limit != 0 && self.frames_processed == self.limit
}
pub fn receive_packet(&mut self) -> Result<Packet<T>, EncoderStatus> { pub fn receive_packet(&mut self) -> Result<Packet<T>, EncoderStatus> {
if self.limit != 0 && self.frames_processed == self.limit { if self.done_processing() {
return Err(EncoderStatus::LimitReached); return Err(EncoderStatus::LimitReached);
} }
...@@ -993,6 +1039,9 @@ impl<T: Pixel> ContextInner<T> { ...@@ -993,6 +1039,9 @@ impl<T: Pixel> ContextInner<T> {
let ret = { let ret = {
let fi = self.frame_invariants.get_mut(&cur_output_frameno).unwrap(); let fi = self.frame_invariants.get_mut(&cur_output_frameno).unwrap();
if fi.show_existing_frame { if fi.show_existing_frame {
if !self.rc_state.ready() {
return Err(EncoderStatus::NotReady);
}
let mut fs = FrameState::new(fi); let mut fs = FrameState::new(fi);
let sef_data = encode_show_existing_frame(fi, &mut fs); let sef_data = encode_show_existing_frame(fi, &mut fs);
...@@ -1011,6 +1060,9 @@ impl<T: Pixel> ContextInner<T> { ...@@ -1011,6 +1060,9 @@ impl<T: Pixel> ContextInner<T> {
self.output_frameno += 1; self.output_frameno += 1;
self.finalize_packet(rec, &fi) self.finalize_packet(rec, &fi)
} else if let Some(f) = self.frame_q.get(&fi.input_frameno) { } else if let Some(f) = self.frame_q.get(&fi.input_frameno) {
if !self.rc_state.ready() {
return Err(EncoderStatus::NotReady);
}
if let Some(frame) = f.clone() { if let Some(frame) = f.clone() {
let fti = fi.get_frame_subtype(); let fti = fi.get_frame_subtype();
let qps = let qps =
......
...@@ -34,6 +34,8 @@ pub struct CliOptions { ...@@ -34,6 +34,8 @@ pub struct CliOptions {
pub skip: usize, pub skip: usize,
pub verbose: bool, pub verbose: bool,
pub threads: usize, pub threads: usize,
pub pass1file_name: Option<String>,
pub pass2file_name: Option<String>
} }
pub fn parse_cli() -> CliOptions { pub fn parse_cli() -> CliOptions {
...@@ -82,12 +84,16 @@ pub fn parse_cli() -> CliOptions { ...@@ -82,12 +84,16 @@ pub fn parse_cli() -> CliOptions {
) )
// ENCODING SETTINGS // ENCODING SETTINGS
.arg( .arg(
Arg::with_name("PASS") Arg::with_name("FIRST_PASS")
.help("Specify first-pass or second-pass to run as a two-pass encode; If not provided, will run a one-pass encode") .help("Perform the first pass of a two-pass encode, saving the pass data to the specified file for future passes")
.short("p") .long("first-pass")
.long("pass") .takes_value(true)
)
.arg(
Arg::with_name("SECOND_PASS")
.help("Perform the second pass of a two-pass encode, reading the pass data saved from a previous pass from the specified file")
.long("second-pass")
.takes_value(true) .takes_value(true)
.possible_values(&["1", "2"])
) )
.arg( .arg(
Arg::with_name("LIMIT") Arg::with_name("LIMIT")
...@@ -341,6 +347,8 @@ pub fn parse_cli() -> CliOptions { ...@@ -341,6 +347,8 @@ pub fn parse_cli() -> CliOptions {
skip: matches.value_of("SKIP").unwrap().parse().unwrap(), skip: matches.value_of("SKIP").unwrap().parse().unwrap(),
verbose: matches.is_present("VERBOSE"), verbose: matches.is_present("VERBOSE"),
threads, threads,
pass1file_name: matches.value_of("FIRST_PASS").map(|s| s.to_owned()),
pass2file_name: matches.value_of("SECOND_PASS").map(|s| s.to_owned())
} }
} }
...@@ -467,7 +475,7 @@ fn parse_config(matches: &ArgMatches<'_>) -> EncoderConfig { ...@@ -467,7 +475,7 @@ fn parse_config(matches: &ArgMatches<'_>) -> EncoderConfig {
cfg.bitrate = bitrate.checked_mul(1000).expect("Bitrate too high"); cfg.bitrate = bitrate.checked_mul(1000).expect("Bitrate too high");
cfg.reservoir_frame_delay = matches.value_of("RESERVOIR_FRAME_DELAY").map(|reservior_frame_delay| reservior_frame_delay.parse().unwrap()); cfg.reservoir_frame_delay = matches.value_of("RESERVOIR_FRAME_DELAY").map(|reservior_frame_delay| reservior_frame_delay.parse().unwrap());
cfg.show_psnr = matches.is_present("PSNR"); cfg.show_psnr = matches.is_present("PSNR");
cfg.pass = matches.value_of("PASS").map(|pass| pass.parse().unwrap()); cfg.pass = None;
cfg.stats_file = if cfg.pass.is_some() { cfg.stats_file = if cfg.pass.is_some() {
Some(PathBuf::from(matches.value_of("STATS_FILE").unwrap())) Some(PathBuf::from(matches.value_of("STATS_FILE").unwrap()))
} else { } else {
......
...@@ -18,6 +18,7 @@ use rav1e::prelude::*; ...@@ -18,6 +18,7 @@ use rav1e::prelude::*;
use std::io; use std::io;
use std::io::Write; use std::io::Write;
use std::io::Read; use std::io::Read;
use std::io::Seek;
use std::path::Path; use std::path::Path;
use std::sync::Arc; use std::sync::Arc;
use crate::decoder::Decoder; use crate::decoder::Decoder;
...@@ -69,10 +70,51 @@ impl<D: Decoder> Source<D> { ...@@ -69,10 +70,51 @@ impl<D: Decoder> Source<D> {
fn process_frame<T: Pixel, D: Decoder>( fn process_frame<T: Pixel, D: Decoder>(
ctx: &mut Context<T>, output_file: &mut dyn Muxer, ctx: &mut Context<T>, output_file: &mut dyn Muxer,
source: &mut Source<D>, source: &mut Source<D>,
pass1file: Option<&mut File>,
pass2file: Option<&mut File>,
buffer: &mut [u8],
buf_pos: &mut usize,
mut y4m_enc: Option<&mut y4m::Encoder<'_, Box<dyn Write>>> mut y4m_enc: Option<&mut y4m::Encoder<'_, Box<dyn Write>>>
) -> Option<Vec<FrameSummary>> { ) -> Option<Vec<FrameSummary>> {
let y4m_details = source.input.get_video_details(); let y4m_details = source.input.get_video_details();
let mut frame_summaries = Vec::new(); let mut frame_summaries = Vec::new();
let mut pass1file = pass1file;
let mut pass2file = pass2file;
// Submit first pass data to pass 2.
if let Some(passfile) = pass2file.as_mut() {
loop {
let mut bytes = ctx.twopass_bytes_needed();
if bytes == 0 {
break;
}
// Read in some more bytes, if necessary.
bytes = bytes.min(buffer.len() - *buf_pos);
if bytes > 0 {
passfile.read_exact(&mut buffer[*buf_pos..*buf_pos + bytes])
.expect("Could not read frame data from two-pass data file!");
}
// And pass them off.
let consumed = ctx.twopass_in(&buffer[*buf_pos..*buf_pos + bytes])
.expect("Error submitting pass data in second pass.");
// If the encoder consumed the whole buffer, reset it.
if consumed >= bytes {
*buf_pos = 0;
}
else {
*buf_pos += consumed;
}
}
}
// Extract first pass data from pass 1.
// We call this before encoding any frames to ensure we are in 2-pass mode
// and to get the placeholder header data.
if let Some(passfile) = pass1file.as_mut() {
if let Some(outbuf) = ctx.twopass_out() {
passfile.write_all(outbuf)
.expect("Unable to write to two-pass data file.");
}
}
let pkt_wrapped = ctx.receive_packet(); let pkt_wrapped = ctx.receive_packet();
match pkt_wrapped { match pkt_wrapped {
Ok(pkt) => { Ok(pkt) => {
...@@ -89,11 +131,25 @@ fn process_frame<T: Pixel, D: Decoder>( ...@@ -89,11 +131,25 @@ fn process_frame<T: Pixel, D: Decoder>(
unreachable!(); unreachable!();
} }
Err(EncoderStatus::LimitReached) => { Err(EncoderStatus::LimitReached) => {
if let Some(passfile) = pass1file.as_mut() {
if let Some(outbuf) = ctx.twopass_out() {
// The last block of data we get is the summary data that needs to go
// at the start of the pass file.
// Seek to the start so we can write it there.
passfile.seek(std::io::SeekFrom::Start(0))
.expect("Unable to seek in two-pass data file.");
passfile.write_all(outbuf)
.expect("Unable to write to two-pass data file.");
}
}
return None; return None;
} }
Err(EncoderStatus::Failure) => { Err(EncoderStatus::Failure) => {
panic!("Failed to encode video"); panic!("Failed to encode video");
} }
Err(EncoderStatus::NotReady) => {
panic!("Mis-managed handling of two-pass stats data");
}
Err(EncoderStatus::Encoded) => {} Err(EncoderStatus::Encoded) => {}
} }
Some(frame_summaries) Some(frame_summaries)
...@@ -110,13 +166,27 @@ fn do_encode<T: Pixel, D: Decoder>( ...@@ -110,13 +166,27 @@ fn do_encode<T: Pixel, D: Decoder>(
cfg: Config, verbose: bool, mut progress: ProgressInfo, cfg: Config, verbose: bool, mut progress: ProgressInfo,
output: &mut dyn Muxer, output: &mut dyn Muxer,
source: &mut Source<D>, source: &mut Source<D>,
pass1file_name: Option<&String>,
pass2file_name: Option<&String>,
mut y4m_enc: Option<y4m::Encoder<'_, Box<dyn Write>>> mut y4m_enc: Option<y4m::Encoder<'_, Box<dyn Write>>>
) { ) {
let mut ctx: Context<T> = cfg.new_context(); let mut ctx: Context<T> = cfg.new_context();
let mut pass2file = pass2file_name.map(|f| {
File::open(f)
.unwrap_or_else(|_| panic!("Unable to open \"{}\" for reading two-pass data.", f))
});
let mut pass1file = pass1file_name.map(|f| {
File::create(f)
.unwrap_or_else(|_| panic!("Unable to open \"{}\" for writing two-pass data.", f))
});
let mut buffer: [u8; 80] = [0; 80];
let mut buf_pos = 0;
while let Some(frame_info) = while let Some(frame_info) =
process_frame(&mut ctx, &mut *output, source, y4m_enc.as_mut()) process_frame(&mut ctx, &mut *output, source, pass1file.as_mut(),
pass2file.as_mut(), &mut buffer, &mut buf_pos, y4m_enc.as_mut())
{ {
for frame in frame_info { for frame in frame_info {
progress.add_frame(frame); progress.add_frame(frame);
...@@ -235,11 +305,13 @@ fn main() { ...@@ -235,11 +305,13 @@ fn main() {
if video_info.bit_depth == 8 { if video_info.bit_depth == 8 {
do_encode::<u8, y4m::Decoder<'_, Box<dyn Read>>>( do_encode::<u8, y4m::Decoder<'_, Box<dyn Read>>>(
cfg, cli.verbose, progress, &mut *cli.io.output, &mut source, y4m_enc cfg, cli.verbose, progress, &mut *cli.io.output, &mut source,
cli.pass1file_name.as_ref(), cli.pass2file_name.as_ref(), y4m_enc
) )
} else { } else {
do_encode::<u16, y4m::Decoder<'_, Box<dyn Read>>>( do_encode::<u16, y4m::Decoder<'_, Box<dyn Read>>>(
cfg, cli.verbose, progress, &mut *cli.io.output, &mut source, y4m_enc cfg, cli.verbose, progress, &mut *cli.io.output, &mut source,
cli.pass1file_name.as_ref(), cli.pass2file_name.as_ref(), y4m_enc
) )
} }
} }
This diff is collapsed.
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment