added option n_threads to set number of threads for vmaf calculation

This commit is contained in:
Zen 2020-07-14 10:41:39 +03:00
parent f7a9eb4fde
commit fa1a2dcabb
3 changed files with 10 additions and 3 deletions

View file

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

View file

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

View file

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