Skip to content

Commit

Permalink
Parallelize frame processing
Browse files Browse the repository at this point in the history
  • Loading branch information
Luni-4 committed Nov 22, 2019
1 parent 9d06a5f commit 65fd63c
Show file tree
Hide file tree
Showing 10 changed files with 192 additions and 78 deletions.
1 change: 1 addition & 0 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

3 changes: 2 additions & 1 deletion av_metrics/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -13,14 +13,15 @@ itertools = "0.8.1"
lab = "0.7.2"
num-traits = "0.2"
serde = { version = "1", features = ["derive"], optional = true }
rayon = { version = "1.2", optional = true }
y4m = { version = "0.4", optional = true }

[dev-dependencies]
criterion = "0.3"

[features]
default = ["y4m-decode"]
decode = []
decode = ["rayon"]
y4m-decode = ["y4m", "decode"]
bench = []

Expand Down
2 changes: 2 additions & 0 deletions av_metrics/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,8 @@ extern crate err_derive;
#[macro_use]
extern crate itertools;

extern crate rayon;

pub mod video;

/// Possible errors that may occur during processing of a metric.
Expand Down
24 changes: 19 additions & 5 deletions av_metrics/src/video/ciede/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,9 @@
#[cfg(feature = "decode")]
use crate::video::decode::Decoder;
use crate::video::pixel::{CastFromPrimitive, Pixel};
use crate::video::ParallelMethod;
use crate::video::{FrameInfo, VideoMetric};
use crate::MetricsError;
use std::f64;

mod rgbtolab;
Expand Down Expand Up @@ -54,7 +56,10 @@ pub fn calculate_frame_ciede<T: Pixel>(
frame1: &FrameInfo<T>,
frame2: &FrameInfo<T>,
) -> Result<f64, Box<dyn Error>> {
Ciede2000::default().process_frame(frame1, frame2)
match Ciede2000::default().process_frame(frame1, frame2) {
Ok((val, _)) => Ok(val),
Err(val) => Err(val.into()),
}
}

/// Calculate the CIEDE2000 metric between two video frames. Higher is better.
Expand All @@ -67,7 +72,10 @@ pub fn calculate_frame_ciede_nosimd<T: Pixel>(
frame1: &FrameInfo<T>,
frame2: &FrameInfo<T>,
) -> Result<f64, Box<dyn Error>> {
(Ciede2000 { use_simd: false }).process_frame(frame1, frame2)
match (Ciede2000 { use_simd: false }).process_frame(frame1, frame2) {
Ok((val, _)) => Ok(val),
Err(val) => Err(val.into()),
}
}

struct Ciede2000 {
Expand All @@ -85,10 +93,10 @@ impl VideoMetric for Ciede2000 {
type VideoResult = f64;

fn process_frame<T: Pixel>(
&mut self,
&self,
frame1: &FrameInfo<T>,
frame2: &FrameInfo<T>,
) -> Result<Self::FrameResult, Box<dyn Error>> {
) -> Result<(Self::FrameResult, Option<f64>), MetricsError> {
frame1.can_compare(&frame2)?;

let dec = frame1.chroma_sampling.get_decimation().unwrap_or((1, 1));
Expand Down Expand Up @@ -123,7 +131,7 @@ impl VideoMetric for Ciede2000 {
* (delta_e_vec.iter().map(|x| *x as f64).sum::<f64>()
/ ((y_width * y_height) as f64))
.log10();
Ok(score.min(100.))
Ok((score.min(100.), None))
}

#[cfg(feature = "decode")]
Expand All @@ -133,6 +141,12 @@ impl VideoMetric for Ciede2000 {
) -> Result<Self::VideoResult, Box<dyn Error>> {
Ok(metrics.iter().copied().sum::<f64>() / metrics.len() as f64)
}

fn which_method(&self) -> ParallelMethod {
ParallelMethod::Ciede
}

fn set_cweight(&mut self, _cweight: f64) {}
}

// Arguments for delta e
Expand Down
2 changes: 1 addition & 1 deletion av_metrics/src/video/decode/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@ pub use self::y4m::*;
/// Currently, y4m decoding support using the `y4m` crate is built-in
/// to this crate. This trait is extensible so users may implement
/// their own decoders.
pub trait Decoder {
pub trait Decoder: Send {
/// Read the next frame from the input video.
///
/// Expected to return `Err` if the end of the video is reached.
Expand Down
2 changes: 1 addition & 1 deletion av_metrics/src/video/decode/y4m.rs
Original file line number Diff line number Diff line change
Expand Up @@ -38,7 +38,7 @@ fn copy_from_raw_u8<T: Pixel>(source: &[u8]) -> Vec<T> {
}
}

impl<R: Read> Decoder for y4m::Decoder<'_, R> {
impl<R: Read + Send> Decoder for y4m::Decoder<'_, R> {
fn read_video_frame<T: Pixel>(&mut self) -> Result<FrameInfo<T>, ()> {
let bit_depth = self.get_bit_depth();
let (chroma_sampling, chroma_sample_pos) = get_chroma_sampling(self);
Expand Down
108 changes: 78 additions & 30 deletions av_metrics/src/video/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,9 @@ pub mod psnr_hvs;
pub mod ssim;

use crate::MetricsError;
use rayon::prelude::*;
use std::error::Error;
use std::iter;

#[cfg(feature = "decode")]
pub use decode::*;
Expand Down Expand Up @@ -185,6 +187,44 @@ pub struct PlanarMetrics {
pub avg: f64,
}

/// Establishes a different parallel method according to the metric
#[derive(Copy, Clone, Debug, PartialEq)]
pub enum ParallelMethod {
/// Method used for Psnr
Psnr,
/// Method used for Ciede
Ciede,
/// Method used for all the other metrics
Others,
}

macro_rules! process_video {
($self:ident, $decoder1:ident, $decoder2:ident, $frame_limit:ident, $type:ident, $ret:ident) => {
$ret = iter::from_fn(|| $decoder1.read_video_frame::<$type>().ok())
.zip(iter::from_fn(|| $decoder2.read_video_frame::<$type>().ok()))
.take($frame_limit.unwrap_or(std::usize::MAX))
.par_bridge()
.try_fold(
|| vec![],
|mut acc, (frame1, frame2)| {
$self
.process_frame(&frame1, &frame2)
.map(|(metrics, cweight)| {
acc.push((metrics, cweight));
acc
})
},
)
.try_reduce(
|| vec![],
|mut a, b| {
a.extend_from_slice(&b);
Ok(a)
},
)?;
};
}

trait VideoMetric {
type FrameResult;
type VideoResult;
Expand All @@ -200,55 +240,63 @@ trait VideoMetric {
decoder1: &mut D,
decoder2: &mut D,
frame_limit: Option<usize>,
) -> Result<Self::VideoResult, Box<dyn Error>> {
) -> Result<Self::VideoResult, Box<dyn Error>>
where
Self: std::marker::Send,
Self: std::marker::Sync,
Self::FrameResult: std::marker::Send,
Self::FrameResult: std::clone::Clone,
{
if decoder1.get_bit_depth() != decoder2.get_bit_depth() {
return Err(Box::new(MetricsError::InputMismatch {
reason: "Bit depths do not match",
}));
}

let mut metrics = Vec::with_capacity(frame_limit.unwrap_or(0));
let mut frame_no = 0;
while frame_limit.map(|limit| limit > frame_no).unwrap_or(true) {
if decoder1.get_bit_depth() > 8 {
let frame1 = decoder1.read_video_frame::<u16>();
let frame2 = decoder2.read_video_frame::<u16>();
if let Ok(frame1) = frame1 {
if let Ok(frame2) = frame2 {
metrics.push(self.process_frame(&frame1, &frame2)?);
frame_no += 1;
continue;
}
}
} else {
let frame1 = decoder1.read_video_frame::<u8>();
let frame2 = decoder2.read_video_frame::<u8>();
if let Ok(frame1) = frame1 {
if let Ok(frame2) = frame2 {
metrics.push(self.process_frame(&frame1, &frame2)?);
frame_no += 1;
continue;
}
}
}
// At end of video
break;
let metrics_values: Vec<(Self::FrameResult, Option<f64>)>;
if decoder1.get_bit_depth() > 8 {
process_video!(self, decoder1, decoder2, frame_limit, u16, metrics_values);
} else {
process_video!(self, decoder1, decoder2, frame_limit, u8, metrics_values);
}
if frame_no == 0 {

if metrics_values.is_empty() {
return Err(MetricsError::UnsupportedInput {
reason: "No readable frames found in one or more input files",
}
.into());
}

let metrics = metrics_values.iter().fold(vec![], |mut acc, v| {
let (metric, _) = v;
acc.push(metric.clone());
acc
});

if let ParallelMethod::Others = self.which_method() {
let metric_option = metrics_values
.iter()
.find(|v| {
let (_, cweight) = v;
cweight.is_some()
})
.unwrap();
let (_, cweight) = metric_option;
self.set_cweight(cweight.unwrap());
}

self.aggregate_frame_results(&metrics)
}

fn which_method(&self) -> ParallelMethod;

fn set_cweight(&mut self, cweight: f64);

fn process_frame<T: Pixel>(
&mut self,
&self,
frame1: &FrameInfo<T>,
frame2: &FrameInfo<T>,
) -> Result<Self::FrameResult, Box<dyn Error>>;
) -> Result<(Self::FrameResult, Option<f64>), MetricsError>;

#[cfg(feature = "decode")]
fn aggregate_frame_results(
Expand Down
16 changes: 12 additions & 4 deletions av_metrics/src/video/psnr.rs
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,9 @@
use crate::video::decode::Decoder;
use crate::video::pixel::CastFromPrimitive;
use crate::video::pixel::Pixel;
use crate::video::ParallelMethod;
use crate::video::{FrameInfo, PlanarMetrics, PlaneData, VideoMetric};
use crate::MetricsError;
use std::error::Error;

/// Calculates the PSNR for two videos. Higher is better.
Expand Down Expand Up @@ -53,7 +55,7 @@ pub fn calculate_frame_psnr<T: Pixel>(
frame1: &FrameInfo<T>,
frame2: &FrameInfo<T>,
) -> Result<PlanarMetrics, Box<dyn Error>> {
let metrics = Psnr.process_frame(frame1, frame2)?;
let (metrics, _) = Psnr.process_frame(frame1, frame2)?;
Ok(PlanarMetrics {
y: calculate_psnr(metrics[0]),
u: calculate_psnr(metrics[1]),
Expand All @@ -75,17 +77,17 @@ impl VideoMetric for Psnr {
type VideoResult = PsnrResults;

fn process_frame<T: Pixel>(
&mut self,
&self,
frame1: &FrameInfo<T>,
frame2: &FrameInfo<T>,
) -> Result<Self::FrameResult, Box<dyn Error>> {
) -> Result<(Self::FrameResult, Option<f64>), MetricsError> {
frame1.can_compare(&frame2)?;

let bit_depth = frame1.bit_depth;
let y = calculate_plane_psnr_metrics(&frame1.planes[0], &frame2.planes[0], bit_depth);
let u = calculate_plane_psnr_metrics(&frame1.planes[1], &frame2.planes[1], bit_depth);
let v = calculate_plane_psnr_metrics(&frame1.planes[2], &frame2.planes[2], bit_depth);
Ok([y, u, v])
Ok(([y, u, v], None))
}

#[cfg(feature = "decode")]
Expand All @@ -111,6 +113,12 @@ impl VideoMetric for Psnr {
};
Ok(PsnrResults { psnr, apsnr })
}

fn which_method(&self) -> ParallelMethod {
ParallelMethod::Psnr
}

fn set_cweight(&mut self, _cweight: f64) {}
}

#[derive(Debug, Clone, Copy, Default)]
Expand Down
38 changes: 26 additions & 12 deletions av_metrics/src/video/psnr_hvs.rs
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,9 @@
use crate::video::decode::Decoder;
use crate::video::pixel::CastFromPrimitive;
use crate::video::pixel::Pixel;
use crate::video::ParallelMethod;
use crate::video::{FrameInfo, PlanarMetrics, PlaneData, VideoMetric};
use crate::MetricsError;
use std::error::Error;

/// Calculates the PSNR-HVS score between two videos. Higher is better.
Expand All @@ -30,8 +32,8 @@ pub fn calculate_frame_psnr_hvs<T: Pixel>(
frame1: &FrameInfo<T>,
frame2: &FrameInfo<T>,
) -> Result<PlanarMetrics, Box<dyn Error>> {
let mut processor = PsnrHvs::default();
let result = processor.process_frame(frame1, frame2)?;
let processor = PsnrHvs::default();
let (result, _) = processor.process_frame(frame1, frame2)?;
let cweight = processor.cweight.unwrap();
Ok(PlanarMetrics {
y: log10_convert(result.y, 1.0),
Expand All @@ -56,26 +58,30 @@ impl VideoMetric for PsnrHvs {
/// Returns the *unweighted* scores. Depending on whether we output per-frame
/// or per-video, these will be weighted at different points.
fn process_frame<T: Pixel>(
&mut self,
&self,
frame1: &FrameInfo<T>,
frame2: &FrameInfo<T>,
) -> Result<Self::FrameResult, Box<dyn Error>> {
) -> Result<(Self::FrameResult, Option<f64>), MetricsError> {
frame1.can_compare(&frame2)?;
let mut cweight = None;
if self.cweight.is_none() {
self.cweight = Some(frame1.chroma_sampling.get_chroma_weight());
cweight = Some(frame1.chroma_sampling.get_chroma_weight());
}

let bit_depth = frame1.bit_depth;
let y = calculate_plane_psnr_hvs(&frame1.planes[0], &frame2.planes[0], 0, bit_depth);
let u = calculate_plane_psnr_hvs(&frame1.planes[1], &frame2.planes[1], 1, bit_depth);
let v = calculate_plane_psnr_hvs(&frame1.planes[2], &frame2.planes[2], 2, bit_depth);
Ok(PlanarMetrics {
y,
u,
v,
// field not used here
avg: 0.,
})
Ok((
PlanarMetrics {
y,
u,
v,
// field not used here
avg: 0.,
},
cweight,
))
}

#[cfg(feature = "decode")]
Expand All @@ -97,6 +103,14 @@ impl VideoMetric for PsnrHvs {
),
})
}

fn which_method(&self) -> ParallelMethod {
ParallelMethod::Others
}

fn set_cweight(&mut self, cweight: f64) {
self.cweight = Some(cweight);
}
}

// Normalized inverse quantization matrix for 8x8 DCT at the point of transparency.
Expand Down
Loading

0 comments on commit 65fd63c

Please sign in to comment.