Commit 714c5ca7 authored by Josh Holmer's avatar Josh Holmer Committed by Luca Barbato
Browse files

Improve CLI progress indicator and print stats after encode

Adds a new `--verbose` flag to display frame-by-frame info similar to
the old behavior.
parent 7862710a
use clap::{App, Arg};
use rav1e::*;
use std::fmt;
use std::fs::File;
use std::io;
use std::io::prelude::*;
use std::ops::Add;
use std::slice;
use std::sync::Arc;
use std::time::Instant;
use y4m;
pub struct EncoderIO {
......@@ -13,7 +16,14 @@ pub struct EncoderIO {
pub rec: Option<Box<dyn Write>>
}
pub fn parse_cli() -> (EncoderIO, EncoderConfig, usize) {
pub struct CliOptions {
pub io: EncoderIO,
pub enc: EncoderConfig,
pub limit: usize,
pub verbose: bool,
}
pub fn parse_cli() -> CliOptions {
let matches = App::new("rav1e")
.version("0.1.0")
.about("AV1 video encoder")
......@@ -70,6 +80,11 @@ pub fn parse_cli() -> (EncoderIO, EncoderConfig, usize) {
.possible_values(&Tune::variants())
.default_value("psnr")
.case_insensitive(true)
).arg(
Arg::with_name("VERBOSE")
.help("verbose logging, output info for every frame")
.long("verbose")
.short("v")
).get_matches();
let io = EncoderIO {
......@@ -101,18 +116,64 @@ pub fn parse_cli() -> (EncoderIO, EncoderConfig, usize) {
panic!("argument out of range");
}
let limit = matches.value_of("LIMIT").unwrap().parse().unwrap();
CliOptions {
io,
enc: config,
limit: matches.value_of("LIMIT").unwrap().parse().unwrap(),
verbose: matches.is_present("VERBOSE"),
}
}
#[derive(Debug, Clone, Copy)]
pub struct FrameSummary {
/// Frame size in bytes
pub size: usize,
pub number: u64,
pub frame_type: FrameType
}
impl From<Packet> for FrameSummary {
fn from(packet: Packet) -> Self {
Self {
size: packet.data.len(),
number: packet.number,
frame_type: packet.frame_type,
}
}
}
impl Add<Packet> for FrameSummary {
type Output = FrameSummary;
(io, config, limit)
fn add(self, rhs: Packet) -> <Self as Add<Packet>>::Output {
assert!(rhs.number == self.number);
Self {
size: self.size + rhs.data.len(),
number: self.number,
frame_type: self.frame_type,
}
}
}
impl fmt::Display for FrameSummary {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(
f,
"Frame {} - {} - {} bytes",
self.number,
self.frame_type,
self.size
)
}
}
/// Encode and write a frame.
/// returns wheter done with sequence
/// Returns frame information in a `Result`.
pub fn process_frame(
ctx: &mut Context, output_file: &mut dyn Write,
y4m_dec: &mut y4m::Decoder<'_, Box<dyn Read>>,
mut y4m_enc: Option<&mut y4m::Encoder<'_, Box<dyn Write>>>
) -> bool {
) -> Result<Option<FrameSummary>, ()> {
let width = y4m_dec.get_width();
let height = y4m_dec.get_height();
let y4m_bits = y4m_dec.get_bit_depth();
......@@ -156,91 +217,217 @@ pub fn process_frame(
};
let mut has_data = true;
let mut frame_summary = None;
while has_data {
let pkt_wrapped = ctx.receive_packet();
match pkt_wrapped {
Ok(pkt) => {
eprintln!("{}", pkt);
write_ivf_frame(output_file, pkt.number as u64, pkt.data.as_ref());
if let Some(y4m_enc_uw) = y4m_enc.as_mut() {
if let Some(rec) = pkt.rec {
let pitch_y = if bit_depth > 8 { width * 2 } else { width };
let pitch_uv = pitch_y / 2;
let (mut rec_y, mut rec_u, mut rec_v) = (
vec![128u8; pitch_y * height],
vec![128u8; pitch_uv * (height / 2)],
vec![128u8; pitch_uv * (height / 2)]
);
Ok(pkt) => {
write_ivf_frame(output_file, pkt.number as u64, pkt.data.as_ref());
if let Some(y4m_enc_uw) = y4m_enc.as_mut() {
if let Some(ref rec) = pkt.rec {
let pitch_y = if bit_depth > 8 { width * 2 } else { width };
let pitch_uv = pitch_y / 2;
let (stride_y, stride_u, stride_v) = (
rec.planes[0].cfg.stride,
rec.planes[1].cfg.stride,
rec.planes[2].cfg.stride
);
let (mut rec_y, mut rec_u, mut rec_v) = (
vec![128u8; pitch_y * height],
vec![128u8; pitch_uv * (height / 2)],
vec![128u8; pitch_uv * (height / 2)]
);
for (line, line_out) in rec.planes[0]
.data_origin()
.chunks(stride_y)
.zip(rec_y.chunks_mut(pitch_y))
{
if bit_depth > 8 {
unsafe {
line_out.copy_from_slice(slice::from_raw_parts::<u8>(
line.as_ptr() as (*const u8),
pitch_y
));
}
} else {
line_out.copy_from_slice(
&line.iter().map(|&v| v as u8).collect::<Vec<u8>>()[..pitch_y]
let (stride_y, stride_u, stride_v) = (
rec.planes[0].cfg.stride,
rec.planes[1].cfg.stride,
rec.planes[2].cfg.stride
);
}
}
for (line, line_out) in rec.planes[1]
.data_origin()
.chunks(stride_u)
.zip(rec_u.chunks_mut(pitch_uv))
{
if bit_depth > 8 {
unsafe {
line_out.copy_from_slice(slice::from_raw_parts::<u8>(
line.as_ptr() as (*const u8),
pitch_uv
));
for (line, line_out) in rec.planes[0]
.data_origin()
.chunks(stride_y)
.zip(rec_y.chunks_mut(pitch_y))
{
if bit_depth > 8 {
unsafe {
line_out.copy_from_slice(slice::from_raw_parts::<u8>(
line.as_ptr() as (*const u8),
pitch_y
));
}
} else {
line_out.copy_from_slice(
&line.iter().map(|&v| v as u8).collect::<Vec<u8>>()[..pitch_y]
);
}
}
} else {
line_out.copy_from_slice(
&line.iter().map(|&v| v as u8).collect::<Vec<u8>>()[..pitch_uv]
);
}
}
for (line, line_out) in rec.planes[2]
.data_origin()
.chunks(stride_v)
.zip(rec_v.chunks_mut(pitch_uv))
{
if bit_depth > 8 {
unsafe {
line_out.copy_from_slice(slice::from_raw_parts::<u8>(
line.as_ptr() as (*const u8),
pitch_uv
));
for (line, line_out) in rec.planes[1]
.data_origin()
.chunks(stride_u)
.zip(rec_u.chunks_mut(pitch_uv))
{
if bit_depth > 8 {
unsafe {
line_out.copy_from_slice(slice::from_raw_parts::<u8>(
line.as_ptr() as (*const u8),
pitch_uv
));
}
} else {
line_out.copy_from_slice(
&line.iter().map(|&v| v as u8).collect::<Vec<u8>>()[..pitch_uv]
);
}
}
} else {
line_out.copy_from_slice(
&line.iter().map(|&v| v as u8).collect::<Vec<u8>>()[..pitch_uv]
);
for (line, line_out) in rec.planes[2]
.data_origin()
.chunks(stride_v)
.zip(rec_v.chunks_mut(pitch_uv))
{
if bit_depth > 8 {
unsafe {
line_out.copy_from_slice(slice::from_raw_parts::<u8>(
line.as_ptr() as (*const u8),
pitch_uv
));
}
} else {
line_out.copy_from_slice(
&line.iter().map(|&v| v as u8).collect::<Vec<u8>>()[..pitch_uv]
);
}
}
let rec_frame = y4m::Frame::new([&rec_y, &rec_u, &rec_v], None);
y4m_enc_uw.write_frame(&rec_frame).unwrap();
}
}
if frame_summary.is_none() {
frame_summary = Some(pkt.into());
} else {
frame_summary = frame_summary.map(|f| f + pkt);
}
},
_ => { has_data = false; }
}
}
if read_frame {
Ok(frame_summary)
} else {
Err(())
}
}
let rec_frame = y4m::Frame::new([&rec_y, &rec_u, &rec_v], None);
y4m_enc_uw.write_frame(&rec_frame).unwrap();
}
}
},
_ => { has_data = false; }
#[derive(Debug, Clone)]
pub struct ProgressInfo {
/// Frame rate of the video
frame_rate: y4m::Ratio,
/// The length of the whole video, in frames, if known
total_frames: Option<usize>,
/// The time the encode was started
time_started: Instant,
/// List of frames encoded so far
frame_info: Vec<FrameSummary>,
/// Video size so far in bytes.
///
/// This value will be updated in the CLI very frequently, so we cache the previous value
/// to reduce the overall complexity.
encoded_size: usize,
}
impl ProgressInfo {
pub fn new(frame_rate: y4m::Ratio, total_frames: Option<usize>) -> Self {
Self {
frame_rate,
total_frames,
time_started: Instant::now(),
frame_info: Vec::with_capacity(total_frames.unwrap_or_default()),
encoded_size: 0,
}
}
pub fn add_frame(&mut self, frame: FrameSummary) {
self.encoded_size += frame.size;
self.frame_info.push(frame);
}
pub fn frames_encoded(&self) -> usize {
self.frame_info.len()
}
pub fn encoding_fps(&self) -> f64 {
let seconds = Instant::now().duration_since(self.time_started).as_secs();
self.frame_info.len() as f64 / seconds as f64
}
pub fn video_fps(&self) -> f64 {
self.frame_rate.num as f64 / self.frame_rate.den as f64
}
/// Returns the bitrate of the frames so far, in bits/second
pub fn bitrate(&self) -> usize {
let bits = self.encoded_size * 8;
let seconds = self.frame_info.len() as f64 / self.video_fps();
(bits as f64 / seconds) as usize
}
/// Estimates the final filesize in bytes, if the number of frames is known
pub fn estimated_size(&self) -> usize {
self.total_frames
.map(|frames| self.encoded_size * frames / self.frames_encoded())
.unwrap_or_default()
}
/// Number of frames of given type which appear in the video
pub fn get_frame_type_count(&self, frame_type: FrameType) -> usize {
self.frame_info.iter()
.filter(|frame| frame.frame_type == frame_type)
.count()
}
/// Size in bytes of all frames of given frame type
pub fn get_frame_type_size(&self, frame_type: FrameType) -> usize {
self.frame_info.iter()
.filter(|frame| frame.frame_type == frame_type)
.map(|frame| frame.size)
.sum()
}
pub fn print_stats(&self) -> String {
let (key, key_size) = (self.get_frame_type_count(FrameType::KEY), self.get_frame_type_size(FrameType::KEY));
let (inter, inter_size) = (self.get_frame_type_count(FrameType::INTER), self.get_frame_type_size(FrameType::INTER));
let (ionly, ionly_size) = (self.get_frame_type_count(FrameType::INTRA_ONLY), self.get_frame_type_size(FrameType::INTRA_ONLY));
let (switch, switch_size) = (self.get_frame_type_count(FrameType::SWITCH), self.get_frame_type_size(FrameType::SWITCH));
format!("\
Key Frames: {:>6} avg size: {:>7} B\n\
Inter: {:>6} avg size: {:>7} B\n\
Intra Only: {:>6} avg size: {:>7} B\n\
Switch: {:>6} avg size: {:>7} B",
key, key_size / key,
inter, inter_size / inter,
ionly, ionly_size / key,
switch, switch_size / key
)
}
}
impl fmt::Display for ProgressInfo {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
if let Some(total_frames) = self.total_frames {
write!(
f,
"encoded {}/{} frames, {:.2} fps, {:.2} Kb/s, est. size: {:.2} MB",
self.frames_encoded(),
total_frames,
self.encoding_fps(),
self.bitrate() as f64 / 1024f64,
self.estimated_size() as f64 / (1024 * 1024) as f64
)
} else {
write!(
f,
"encoded {} frames, {:.2} fps, {:.2} Kb/s",
self.frames_encoded(),
self.encoding_fps(),
self.bitrate() as f64 / 1024f64
)
}
}
read_frame
}
......@@ -14,18 +14,19 @@ extern crate y4m;
mod common;
use common::*;
use std::io;
use std::io::Write;
use rav1e::*;
fn main() {
let (mut io, enc, limit) = parse_cli();
let mut y4m_dec = y4m::decode(&mut io.input).unwrap();
let mut cli = parse_cli();
let mut y4m_dec = y4m::decode(&mut cli.io.input).unwrap();
let width = y4m_dec.get_width();
let height = y4m_dec.get_height();
let framerate = y4m_dec.get_framerate();
let color_space = y4m_dec.get_colorspace();
let mut count = 0;
let mut y4m_enc = match io.rec.as_mut() {
let mut y4m_enc = match cli.io.rec.as_mut() {
Some(rec) => Some(
y4m::encode(width, height, framerate)
.with_colorspace(color_space)
......@@ -56,31 +57,49 @@ fn main() {
let cfg = Config {
frame_info: FrameInfo { width, height, bit_depth, chroma_sampling },
timebase: Rational::new(framerate.den as u64, framerate.num as u64),
enc
enc: cli.enc
};
let mut ctx = cfg.new_context();
let stderr = io::stderr();
let mut err = stderr.lock();
let _ = writeln!(err, "{}x{} @ {}/{} fps", width, height, framerate.num, framerate.den);
write_ivf_header(
&mut io.output,
&mut cli.io.output,
width,
height,
framerate.num,
framerate.den
);
loop {
if !process_frame(&mut ctx, &mut io.output, &mut y4m_dec, y4m_enc.as_mut())
{
break;
}
let mut progress = ProgressInfo::new(
framerate,
if cli.limit == 0 { None } else { Some(cli.limit) }
);
count += 1;
loop {
match process_frame(&mut ctx, &mut cli.io.output, &mut y4m_dec, y4m_enc.as_mut()) {
Ok(Some(frame_info)) => {
progress.add_frame(frame_info);
let _ = if cli.verbose {
writeln!(err, "{} - {}", frame_info, progress)
} else {
write!(err, "\r{} ", progress)
};
},
Ok(_) => (),
Err(_) => break,
};
if limit != 0 && count >= limit {
if cli.limit != 0 && progress.frames_encoded() >= cli.limit {
break;
}
io.output.flush().unwrap();
cli.io.output.flush().unwrap();
}
let _ = write!(err, "\n{}\n", progress.print_stats());
}
......@@ -76,14 +76,11 @@ fn main() {
rl.add_history_entry(&line);
match line.split_whitespace().next() {
Some("process_frame") => {
if !process_frame(
&mut ctx,
&mut io.output,
&mut y4m_dec,
y4m_enc.as_mut()
) {
break;
}
match process_frame(&mut ctx, &mut io.output, &mut y4m_dec, y4m_enc.as_mut()) {
Ok(Some(frame_info)) => eprintln!("{}", frame_info),
Ok(_) => (),
Err(_) => break,
};
io.output.flush().unwrap();
}
......
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