Commit 50967bdf authored by Josh Holmer's avatar Josh Holmer Committed by Thomas Daede

Output PSNR for each frame and average for video

When enabled via --psnr flat, will output planar and average PSNR for
the video at the end. In verbose mode, will output planar PSNR for each frame.

Closes #672
parent 869fef70
......@@ -15,6 +15,7 @@ use std::collections::BTreeMap;
use std::fmt;
use std::sync::Arc;
use scenechange::SceneChangeDetector;
use metrics::calculate_frame_psnr;
// TODO: use the num crate?
#[derive(Clone, Copy, Debug)]
......@@ -40,6 +41,7 @@ pub struct EncoderConfig {
pub quantizer: usize,
pub tune: Tune,
pub speed_settings: SpeedSettings,
pub show_psnr: bool,
}
impl Default for EncoderConfig {
......@@ -57,7 +59,8 @@ impl EncoderConfig {
low_latency: true,
quantizer: 100,
tune: Tune::Psnr,
speed_settings: SpeedSettings::from_preset(speed)
speed_settings: SpeedSettings::from_preset(speed),
show_psnr: false,
}
}
}
......@@ -244,7 +247,9 @@ pub struct Packet {
pub data: Vec<u8>,
pub rec: Option<Frame>,
pub number: u64,
pub frame_type: FrameType
pub frame_type: FrameType,
/// PSNR for Y, U, and V planes
pub psnr: Option<(f64, f64, f64)>,
}
impl fmt::Display for Packet {
......@@ -563,14 +568,20 @@ impl Context {
// TODO avoid the clone by having rec Arc.
let rec = if self.fi.show_frame { Some(fs.rec.clone()) } else { None };
let mut psnr = None;
if self.fi.config.show_psnr {
if let Some(ref rec) = rec {
psnr = Some(calculate_frame_psnr(&*fs.input, rec, self.seq.bit_depth));
}
}
Ok(Packet { data, rec, number: self.fi.number, frame_type: self.fi.frame_type })
Ok(Packet { data, rec, number: self.fi.number, frame_type: self.fi.frame_type, psnr })
} else {
if let Some(f) = self.frame_q.remove(&self.fi.number) {
self.idx += 1;
if let Some(frame) = f {
let mut fs = FrameState::new_with_frame(&self.fi, frame);
let mut fs = FrameState::new_with_frame(&self.fi, frame.clone());
let data = encode_frame(&mut self.seq, &mut self.fi, &mut fs);
self.packet_data.extend(data);
......@@ -585,7 +596,15 @@ impl Context {
if self.fi.show_frame {
let data = self.packet_data.clone();
self.packet_data = Vec::new();
Ok(Packet { data, rec, number: self.fi.number, frame_type: self.fi.frame_type })
let mut psnr = None;
if self.fi.config.show_psnr {
if let Some(ref rec) = rec {
psnr = Some(calculate_frame_psnr(&*frame, rec, self.seq.bit_depth));
}
}
Ok(Packet { data, rec, number: self.fi.number, frame_type: self.fi.frame_type, psnr })
} else {
Err(EncoderStatus::NeedMoreData)
}
......
......@@ -100,6 +100,10 @@ pub fn parse_cli() -> CliOptions {
.help("verbose logging, output info for every frame")
.long("verbose")
.short("v")
).arg(
Arg::with_name("PSNR")
.help("calculate and display PSNR metrics")
.long("psnr")
).get_matches();
let io = EncoderIO {
......@@ -145,6 +149,7 @@ fn parse_config(matches: &ArgMatches) -> EncoderConfig {
cfg.low_latency = matches.value_of("LOW_LATENCY").unwrap().parse().unwrap();
cfg.tune = matches.value_of("TUNE").unwrap().parse().unwrap();
cfg.quantizer = quantizer;
cfg.show_psnr = matches.is_present("PSNR");
cfg
}
......@@ -153,7 +158,9 @@ pub struct FrameSummary {
/// Frame size in bytes
pub size: usize,
pub number: u64,
pub frame_type: FrameType
pub frame_type: FrameType,
/// PSNR for Y, U, and V planes
pub psnr: Option<(f64, f64, f64)>,
}
impl From<Packet> for FrameSummary {
......@@ -162,6 +169,7 @@ impl From<Packet> for FrameSummary {
size: packet.data.len(),
number: packet.number,
frame_type: packet.frame_type,
psnr: packet.psnr,
}
}
}
......@@ -170,10 +178,13 @@ impl fmt::Display for FrameSummary {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(
f,
"Frame {} - {} - {} bytes",
"Frame {} - {} - {} bytes{}",
self.number,
self.frame_type,
self.size
self.size,
if let Some(psnr) = self.psnr {
format!(" - PSNR: Y: {:.4} Cb: {:.4} Cr: {:.4}", psnr.0, psnr.1, psnr.2)
} else { String::new() }
)
}
}
......@@ -335,16 +346,19 @@ pub struct ProgressInfo {
/// This value will be updated in the CLI very frequently, so we cache the previous value
/// to reduce the overall complexity.
encoded_size: usize,
/// Whether to display PSNR statistics during and at end of encode
show_psnr: bool,
}
impl ProgressInfo {
pub fn new(frame_rate: y4m::Ratio, total_frames: Option<usize>) -> Self {
pub fn new(frame_rate: y4m::Ratio, total_frames: Option<usize>, show_psnr: bool) -> Self {
Self {
frame_rate,
total_frames,
time_started: Instant::now(),
frame_info: Vec::with_capacity(total_frames.unwrap_or_default()),
encoded_size: 0,
show_psnr,
}
}
......@@ -404,11 +418,20 @@ impl ProgressInfo {
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",
Switch: {:>6} avg size: {:>7} B\
{}",
key, key_size / key,
inter, inter_size.checked_div(inter).unwrap_or(0),
ionly, ionly_size / key,
switch, switch_size / key
switch, switch_size / key,
if self.show_psnr {
let psnr_y = self.frame_info.iter().map(|fi| fi.psnr.unwrap().0).sum::<f64>() / self.frame_info.len() as f64;
let psnr_u = self.frame_info.iter().map(|fi| fi.psnr.unwrap().1).sum::<f64>() / self.frame_info.len() as f64;
let psnr_v = self.frame_info.iter().map(|fi| fi.psnr.unwrap().2).sum::<f64>() / self.frame_info.len() as f64;
format!("\nMean PSNR: Y: {:.4} Cb: {:.4} Cr: {:.4} Avg: {:.4}",
psnr_y, psnr_u, psnr_v,
(psnr_y + psnr_u + psnr_v) / 3.0)
} else { String::new() }
)
}
}
......
......@@ -77,7 +77,8 @@ fn main() {
let mut progress = ProgressInfo::new(
framerate,
if cli.limit == 0 { None } else { Some(cli.limit) }
if cli.limit == 0 { None } else { Some(cli.limit) },
cfg.enc.show_psnr
);
ctx.set_frames_to_be_coded(cli.limit as u64);
......
......@@ -36,6 +36,7 @@ pub mod cdef;
pub mod lrf;
pub mod encoder;
pub mod me;
pub mod metrics;
pub mod scan_order;
pub mod scenechange;
......
use encoder::Frame;
use plane::Plane;
/// Calculates the PSNR for a `Frame` by comparing the original (uncompressed) to the compressed
/// version of the frame. Higher PSNR is better--PSNR is capped at 100 in order to avoid skewed
/// statistics from e.g. all black frames, which would otherwise show a PSNR of infinity.
///
/// See https://en.wikipedia.org/wiki/Peak_signal-to-noise_ratio for more details.
pub fn calculate_frame_psnr(original: &Frame, compressed: &Frame, bit_depth: usize) -> (f64, f64, f64) {
(calculate_plane_psnr(&original.planes[0], &compressed.planes[0], bit_depth),
calculate_plane_psnr(&original.planes[1], &compressed.planes[1], bit_depth),
calculate_plane_psnr(&original.planes[2], &compressed.planes[2], bit_depth))
}
/// Calculate the PSNR for a `Plane` by comparing the original (uncompressed) to the compressed
/// version.
fn calculate_plane_psnr(original: &Plane, compressed: &Plane, bit_depth: usize) -> f64 {
let mse = calculate_plane_mse(original, compressed);
if mse <= 0.0000000001 {
return 100.0;
}
let max = ((1 << bit_depth) - 1) as f64;
20.0 * max.log10() - 10.0 * mse.log10()
}
/// Calculate the mean squared error for a `Plane` by comparing the original (uncompressed)
/// to the compressed version.
fn calculate_plane_mse(original: &Plane, compressed: &Plane) -> f64 {
original.iter().zip(compressed.iter())
.map(|(a, b)| (a as i32 - b as i32).abs() as u64)
.map(|err| err * err)
.sum::<u64>() as f64 / (original.cfg.width * original.cfg.height) as f64
}
\ No newline at end of file
......@@ -241,6 +241,54 @@ impl Plane {
}
}
}
/// Iterates over the pixels in the `Plane`, skipping stride data.
pub fn iter(&self) -> PlaneIter {
PlaneIter::new(self)
}
}
#[derive(Debug)]
pub struct PlaneIter<'a> {
plane: &'a Plane,
y: usize,
x: usize,
}
impl<'a> PlaneIter<'a> {
pub fn new(plane: &'a Plane) -> Self {
PlaneIter {
plane,
y: 0,
x: 0,
}
}
fn width(&self) -> usize {
self.plane.cfg.width
}
fn height(&self) -> usize {
self.plane.cfg.height
}
}
impl<'a> Iterator for PlaneIter<'a> {
type Item = u16;
fn next(&mut self) -> Option<<Self as Iterator>::Item> {
if self.y == self.height() - 1 && self.x == self.width() - 1 {
return None;
}
let pixel = self.plane.p(self.x, self.y);
if self.x == self.width() - 1 {
self.x = 0;
self.y += 1;
} else {
self.x += 1;
}
Some(pixel)
}
}
#[derive(Clone, Copy)]
......
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