diff --git a/pipeline/feature_extractors.py b/pipeline/feature_extractors.py index 5feb9be..c061a8d 100644 --- a/pipeline/feature_extractors.py +++ b/pipeline/feature_extractors.py @@ -200,12 +200,14 @@ class LoudAudioFeatureExtractor(FeatureExtractor): teardown() is used to clean up temporary files created during setup (if specified by config) """ - def __init__(self, input_files=None, config=None): + _CONFIG_DEFAULT_NUM_FEATURES = 5 # keep the top 5 loudnesses + def __init__(self, input_files=None, config=None, num_features=_CONFIG_DEFAULT_NUM_FEATURES): if not input_files: raise ValueError("No input files provided!") self.input_files = input_files self.config = config self.features = [] + self._num_features = num_features def _audio_file_from_path(self, path: str) -> str: """Return the audio file path given a video file path @@ -249,6 +251,10 @@ class LoudAudioFeatureExtractor(FeatureExtractor): return loudness_features + def _keep_num(self, loudnesses, num=_CONFIG_DEFAULT_NUM_FEATURES) -> list: + """Keep the top n loudnesses (default: 5)""" + return sorted(loudnesses, key=lambda x: x[1], reverse=True)[:num] + def setup(self): """extract audio from video files to be processed by pyloudnorm @@ -266,7 +272,8 @@ class LoudAudioFeatureExtractor(FeatureExtractor): for file in self.input_files: audio_file = self._audio_file_from_path(file.path) loudnesses = self._loudnorm(audio_file) - for time, loudness in loudnesses: + top_loudnesses = self._keep_num(loudnesses, self._num_features) + for time, loudness in top_loudnesses: self.features.append(Feature(interval=Interval(start=time, duration=0.500), source=file, feature_extractor="loudness", score=loudness))