diff --git a/av1an/target_quality/target_quality.py b/av1an/target_quality/target_quality.py index fc9ecf2..ef05193 100644 --- a/av1an/target_quality/target_quality.py +++ b/av1an/target_quality/target_quality.py @@ -1,4 +1,5 @@ import subprocess +import os from math import isnan import numpy as np @@ -37,6 +38,7 @@ class TargetQuality: self.encoder = project.encoder self.ffmpeg_pipe = project.ffmpeg_pipe self.temp = project.temp + self.workers = project.workers def per_frame_target_quality_routine(self, chunk: Chunk): """ @@ -305,7 +307,7 @@ class TargetQuality: :return : path to json file with vmaf scores """ - n_threads = self.n_threads if self.n_threads else 12 + n_threads = self.n_threads if self.n_threads else self.auto_vmaf_threads() cmd = self.probe_cmd( chunk, q, self.ffmpeg_pipe, self.encoder, self.probing_rate, n_threads ) @@ -316,6 +318,19 @@ class TargetQuality: ) return fl + def auto_vmaf_threads(self): + """ + Calculates number of vmaf threads based on CPU cores in system + + :return: Integer value for number of threads + """ + cores = os.cpu_count() + # One thread may not be enough to keep the CPU saturated, so over-provision a bit. + over_provision_factor = 1.25 + minimum_threads = 1 + + return int(max((cores / self.workers) * over_provision_factor, minimum_threads)) + def get_closest(self, q_list, q, positive=True): """ Returns closest value from the list, ascending or descending