plot_vmaf to rust

av1an-core/Cargo.toml

@@ -19,3 +19,4 @@ clap = "2.33.3"
 num-rational = "0.4.0"
 failure = "0.1.8"
 regex = "1.5.4"
+plotters = "0.3.1"

av1an-core/src/lib.rs

@@ -13,6 +13,7 @@ pub mod file_validation;
 pub mod split;
 pub mod target_quality;
 pub mod vapoursynth;
+pub mod vmaf;
 
 #[allow(non_camel_case_types)]
 #[derive(Clone, Copy)]
@@ -214,7 +215,7 @@ pub fn read_file_to_string(file: PathBuf) -> Result<String, Error> {
   Ok(fs::read_to_string(&file).unwrap_or_else(|_| panic!("Can't open file {:?}", file)))
 }
 
-pub fn read_weighted_vmaf(file: PathBuf, percentile: f64) -> Result<f64, serde_json::Error> {
+pub fn read_vmaf_file(file: PathBuf) -> Result<Vec<f64>, serde_json::Error> {
   let json_str = read_file_to_string(file).unwrap();
   let bazs = serde_json::from_str::<Baz>(&json_str)?;
   let v = bazs
@@ -223,5 +224,11 @@ pub fn read_weighted_vmaf(file: PathBuf, percentile: f64) -> Result<f64, serde_j
     .map(|x| x.metrics.vmaf)
     .collect::<Vec<_>>();
 
-  Ok(get_percentile(v, percentile))
+  Ok(v)
+}
+
+pub fn read_weighted_vmaf(file: PathBuf, percentile: f64) -> Result<f64, serde_json::Error> {
+  let scores = read_vmaf_file(file).unwrap();
+
+  Ok(get_percentile(scores, percentile))
 }

av1an-core/src/vmaf.rs

@@ -0,0 +1,77 @@
+use crate::{read_vmaf_file, read_weighted_vmaf};
+use plotters::prelude::*;
+use std::{path::PathBuf, u32};
+
+pub fn plot_vmaf_score_file(
+  scores_file: PathBuf,
+  plot_path: PathBuf,
+) -> Result<(), Box<dyn std::error::Error>> {
+  let scores = read_vmaf_file(scores_file.clone()).unwrap();
+
+  let plot_width = 2560 + ((scores.len() as f64).log10() as u32 * 200);
+  let plot_heigth = 1440;
+
+  let length = scores.len() as u32;
+  let perc_1 = read_weighted_vmaf(scores_file.clone(), 0.01).unwrap();
+  let perc_25 = read_weighted_vmaf(scores_file.clone(), 0.25).unwrap();
+  let perc_75 = read_weighted_vmaf(scores_file.clone(), 0.75).unwrap();
+  let perc_mean = read_weighted_vmaf(scores_file.clone(), 0.50).unwrap();
+
+  let root =
+    BitMapBackend::new(plot_path.as_os_str(), (plot_width, plot_heigth)).into_drawing_area();
+
+  root.fill(&WHITE)?;
+
+  let mut chart = ChartBuilder::on(&root)
+    .set_label_area_size(LabelAreaPosition::Bottom, (8).percent())
+    .set_label_area_size(LabelAreaPosition::Left, (5).percent())
+    .set_label_area_size(LabelAreaPosition::Left, (5).percent())
+    .set_label_area_size(LabelAreaPosition::Top, (5).percent())
+    .margin((1).percent())
+    .build_cartesian_2d(0u32..length, perc_1.floor()..100.0)?;
+
+  chart.configure_mesh().draw()?;
+
+  // 1%
+  chart
+    .draw_series(LineSeries::new((0..=length).map(|x| (x, perc_1)), &RED))?
+    .label(format!("1%: {}", perc_1))
+    .legend(|(x, y)| PathElement::new(vec![(x, y), (x + 20, y)], &RED));
+
+  // 25%
+  chart
+    .draw_series(LineSeries::new((0..=length).map(|x| (x, perc_25)), &YELLOW))?
+    .label(format!("25%: {}", perc_25))
+    .legend(|(x, y)| PathElement::new(vec![(x, y), (x + 20, y)], &YELLOW));
+
+  // 75%
+  chart
+    .draw_series(LineSeries::new((0..=length).map(|x| (x, perc_75)), &GREEN))?
+    .label(format!("75%: {}", perc_75))
+    .legend(|(x, y)| PathElement::new(vec![(x, y), (x + 20, y)], &GREEN));
+
+  // Mean
+  chart
+    .draw_series(LineSeries::new(
+      (0..=length).map(|x| (x, perc_mean)),
+      &BLACK,
+    ))?
+    .label(format!("Mean: {}", perc_mean))
+    .legend(|(x, y)| PathElement::new(vec![(x, y), (x + 20, y)], &BLACK));
+
+  // Data
+  chart.draw_series(LineSeries::new(
+    (0..).zip(scores.iter()).map(|(x, y)| (x, *y)),
+    &BLUE,
+  ))?;
+
+  chart
+    .configure_series_labels()
+    .background_style(&WHITE.mix(0.8))
+    .border_style(&BLACK)
+    .draw()?;
+
+  root.present().expect("Unable to write result plot to file");
+
+  Ok(())
+}

av1an-pyo3/src/lib.rs

@@ -436,12 +436,18 @@ pub fn finish_progress_bar() {
     .finish();
 }
 
+#[pyfunction]
+pub fn plot_vmaf_score_file(scores_file_string: String, plot_path_string: String) {
+  let scores_file = PathBuf::from(scores_file_string);
+  let plot_path = PathBuf::from(plot_path_string);
+  av1an_core::vmaf::plot_vmaf_score_file(scores_file, plot_path).unwrap()
+}
+
 #[pymodule]
 fn av1an_pyo3(_py: Python, m: &PyModule) -> PyResult<()> {
   m.add_function(wrap_pyfunction!(init_progress_bar, m)?)?;
   m.add_function(wrap_pyfunction!(update_bar, m)?)?;
   m.add_function(wrap_pyfunction!(finish_progress_bar, m)?)?;
-
   m.add_function(wrap_pyfunction!(get_ffmpeg_info, m)?)?;
   m.add_function(wrap_pyfunction!(determine_workers, m)?)?;
   m.add_function(wrap_pyfunction!(create_vs_file, m)?)?;
@@ -482,6 +488,7 @@ fn av1an_pyo3(_py: Python, m: &PyModule) -> PyResult<()> {
   m.add_function(wrap_pyfunction!(probe_cmd, m)?)?;
   m.add_function(wrap_pyfunction!(get_percentile, m)?)?;
   m.add_function(wrap_pyfunction!(read_weighted_vmaf, m)?)?;
+  m.add_function(wrap_pyfunction!(plot_vmaf_score_file, m)?)?;
 
   Ok(())
 }

av1an/vmaf.py

@@ -10,7 +10,7 @@ from pathlib import Path
 from subprocess import PIPE, STDOUT
 
 import numpy as np
-from av1an_pyo3 import get_percentile, log, read_weighted_vmaf
+from av1an_pyo3 import get_percentile, log, read_weighted_vmaf, plot_vmaf_score_file
 from matplotlib import pyplot as plt
 
 from av1an.chunk import Chunk
@@ -195,62 +195,4 @@ class VMAF:
             sys.exit()
 
         file_path = encoded.with_name(f"{encoded.stem}_plot").with_suffix(".png")
-        self.plot_vmaf_score_file(scores, file_path)
-
-    def plot_vmaf_score_file(self, scores: Path, plot_path: Path):
-        if plt is None:
-            log(
-                f"Matplotlib is not installed or could not be loaded, aborting plot_vmaf"
-            )
-            return
-
-        perc_1 = read_weighted_vmaf(str(scores.as_posix()), 0.01)
-        perc_25 = read_weighted_vmaf(str(scores.as_posix()), 0.25)
-        perc_75 = read_weighted_vmaf(str(scores.as_posix()), 0.75)
-        mean = read_weighted_vmaf(str(scores.as_posix()), 0.50)
-
-        with open(scores) as f:
-            file = json.load(f)
-            vmafs = [x["metrics"]["vmaf"] for x in file["frames"]]
-            plot_size = len(vmafs)
-
-        figure_width = 3 + round((4 * log10(plot_size)))
-        plt.figure(figsize=(figure_width, 5))
-
-        plt.plot([1, plot_size], [perc_1, perc_1], "-", color="red")
-        plt.annotate(f"1%: {perc_1}", xy=(0, perc_1), color="red")
-
-        plt.plot([1, plot_size], [perc_25, perc_25], ":", color="orange")
-        plt.annotate(f"25%: {perc_25}", xy=(0, perc_25), color="orange")
-
-        plt.plot([1, plot_size], [perc_75, perc_75], ":", color="green")
-        plt.annotate(f"75%: {perc_75}", xy=(0, perc_75), color="green")
-
-        plt.plot([1, plot_size], [mean, mean], ":", color="black")
-        plt.annotate(f"Mean: {mean}", xy=(0, mean), color="black")
-
-        for i in range(0, 100):
-            plt.axhline(i, color="grey", linewidth=0.4)
-            if i % 5 == 0:
-                plt.axhline(i, color="black", linewidth=0.6)
-
-        plt.plot(
-            range(plot_size),
-            vmafs,
-            label=f"Frames: {plot_size}\nMean:{mean}\n"
-            f"1%: {perc_1} \n25%: {perc_25} \n75%: {perc_75}",
-            linewidth=0.7,
-        )
-        plt.ylabel("VMAF")
-        plt.legend(
-            loc="lower right",
-            markerscale=0,
-            handlelength=0,
-            fancybox=True,
-        )
-        plt.ylim(int(perc_1), 100)
-        plt.tight_layout()
-        plt.margins(0)
-
-        # Save
-        plt.savefig(plot_path, dpi=250)
+        plot_vmaf_score_file(str(scores.as_posix()), str(file_path.as_posix()))