Parallelize frame processing

.github/workflows/ci-build.yml

@@ -74,8 +74,8 @@ jobs:
         toolchain: stable
         override: true
     - name: Build
-      run: cargo build --all-features --tests --benches
+      run: cargo build --release --all-features --tests --benches
     - name: Run tests
-      run: cargo test --all-features
+      run: cargo test --release --all-features
     - name: Generate docs
-      run: cargo doc --all-features --no-deps
+      run: cargo doc --release --all-features --no-deps

av_metrics/Cargo.toml

@@ -14,14 +14,15 @@ lab = "0.7.2"
 libc = "0.2"
 num-traits = "0.2"
 serde = { version = "1", features = ["derive"], optional = true }
+rayon = { version = "1.2.1", 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 = []
 

av_metrics/src/lib.rs

@@ -12,6 +12,8 @@ extern crate err_derive;
 #[macro_use]
 extern crate itertools;
 
+extern crate rayon;
+
 pub mod video;
 
 #[cfg(cargo_c)]

av_metrics/src/video/ciede/mod.rs

@@ -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;
@@ -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.
@@ -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 {
@@ -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));
@@ -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")]
@@ -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

av_metrics/src/video/decode/mod.rs

@@ -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.

av_metrics/src/video/decode/y4m.rs

@@ -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);

av_metrics/src/video/mod.rs

@@ -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::*;
@@ -185,8 +187,33 @@ 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) => {
+        let frames_pool: Vec<(FrameInfo<$type>, FrameInfo<$type>)>;
+        frames_pool = 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))
+            .collect();
+        $ret = frames_pool
+            .par_iter()
+            .map(|(frame1, frame2)| $self.process_frame(&frame1, &frame2).unwrap())
+            .collect();
+    };
+}
+
 trait VideoMetric {
-    type FrameResult;
+    type FrameResult: Clone + Send;
     type VideoResult;
 
     /// Generic method for internal use that processes multiple frames from a video
@@ -200,55 +227,65 @@ 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::Sync,
+    {
         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;
+        rayon::ThreadPoolBuilder::new()
+            .num_threads(8)
+            .build_global()
+            .unwrap_or(());
+
+        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(

av_metrics/src/video/psnr.rs

@@ -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.
@@ -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]),
@@ -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")]
@@ -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)]