From fa1a2dcabbab22281bb1b95c2452324de61e26bf Mon Sep 17 00:00:00 2001
From: Zen <46526140+master-of-zen@users.noreply.github.com>
Date: Tue, 14 Jul 2020 10:41:39 +0300
Subject: [PATCH] added option `n_threads` to set number of threads for vmaf
 calculation

---
 av1an.py           | 3 ++-
 utils/arg_parse.py | 1 +
 utils/vmaf.py      | 9 +++++++--
 3 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/av1an.py b/av1an.py
index 840cf34..f6e13c3 100755
--- a/av1an.py
+++ b/av1an.py
@@ -67,6 +67,7 @@ class Av1an:
         self.min_cq = None
         self.max_cq = None
         self.vmaf_plots = None
+        self.n_treads = None
 
         # get all values from argparse
         self.__dict__.update(arg_parsing())
@@ -129,7 +130,7 @@ class Av1an:
             for count, i in enumerate(cmd):
                 subprocess.run(i[0], shell=True)
 
-                v = call_vmaf(i[1], i[2], model=self.vmaf_path, return_file=True)
+                v = call_vmaf(i[1], i[2], n_threads=self.n_threads, model=self.vmaf_path, return_file=True)
                 # Trying 25 percentile
                 mean = read_vmaf_xml(v, 25)
 
diff --git a/utils/arg_parse.py b/utils/arg_parse.py
index d22eb84..b20483e 100755
--- a/utils/arg_parse.py
+++ b/utils/arg_parse.py
@@ -61,6 +61,7 @@ def arg_parsing():
     parser.add_argument('--min_cq', type=int, default=25, help='Min cq for target vmaf')
     parser.add_argument('--max_cq', type=int, default=50, help='Max cq for target vmaf')
     parser.add_argument('--vmaf_plots', help='Make plots of probes in temp folder', action='store_true')
+    parser.add_argument('--n_threads', type=int, default=None, help='Threads for vmaf calculation')
 
     # Store all vars in dictionary
     return vars(parser.parse_args())
diff --git a/utils/vmaf.py b/utils/vmaf.py
index 26cb589..ca87aa8 100755
--- a/utils/vmaf.py
+++ b/utils/vmaf.py
@@ -30,20 +30,25 @@ def read_vmaf_xml(file, percentile):
         return perc
 
 
-def call_vmaf(source: Path, encoded: Path, model=None, return_file=False):
+def call_vmaf(source: Path, encoded: Path, model, n_threads, return_file=False):
 
     if model:
         mod = f":model_path={model}"
     else:
         mod = ''
 
+    if n_threads:
+        n_threads = f':n_threads={n_threads}'
+    else:
+        n_threads = ''
+
     # For vmaf calculation both source and encoded segment scaled to 1080
     # for proper vmaf calculation
     fl = source.with_name(encoded.stem).with_suffix('.xml').as_posix()
     cmd = f'ffmpeg -loglevel error -hide_banner -r 60 -i {source.as_posix()} -r 60 -i {encoded.as_posix()}  ' \
           f'-filter_complex "[0:v]scale=1920:1080:flags=spline:force_original_aspect_ratio=decrease[scaled1];' \
           f'[1:v]scale=1920:1080:flags=spline:force_original_aspect_ratio=decrease[scaled2];' \
-          f'[scaled2][scaled1]libvmaf=log_path={fl}{mod}" -f null - '
+          f'[scaled2][scaled1]libvmaf=log_path={fl}{mod}{n_threads}" -f null - '
 
     c = subprocess.run(cmd, shell=True, stdout=PIPE, stderr=STDOUT)
     call = c.stdout