diff --git a/pipeline/producers.py b/pipeline/producers.py index 37230f3..b656267 100644 --- a/pipeline/producers.py +++ b/pipeline/producers.py @@ -26,15 +26,21 @@ class VideoProducer(Producer): class FfmpegVideoProducer(VideoProducer): """Produce videos using ffmpeg""" - # TODO: consider output filename options + _CONFIG_DEFAULT_OUTPUT_DIR = "/tmp/" + _CONFIG_DEFAULT_OUTPUT_FILENAME = "highlights.mp4" _CONFIG_COMPILE_CLIPS = True - def __init__(self, features, compile_clips=_CONFIG_COMPILE_CLIPS) -> None: + def __init__(self, features, compile_clips=_CONFIG_COMPILE_CLIPS, + output_dir=_CONFIG_DEFAULT_OUTPUT_DIR, + output_filename=_CONFIG_DEFAULT_OUTPUT_FILENAME, + ) -> None: if not features: raise ValueError("No features provided") # TODO: consider if we want to permit empty features (producing no video) self.features = features self._compile_clips = compile_clips + self._output_dir = output_dir + self._output_filename = output_filename def _run_no_output(self, cmd: list, cwd:str=".") -> None: """Run a command and return the output as a string @@ -91,17 +97,17 @@ class FfmpegVideoProducer(VideoProducer): join_file.close() def produce(self): - OUTPUT_DIR = "/tmp/" # TODO: make this a config option + """Produce clips or a video from the features""" clips = [] for num, feature in enumerate(self.features): - output_filepath = f"{OUTPUT_DIR}/highlight_{num}.mp4" + output_filepath = f"{self._output_dir}/highlight_{num}.mp4" self._ffmpeg_feature_to_clip(feature, output_filepath) clips.append(output_filepath) # concatenate the clips if self._compile_clips: - output_filepath = f"{OUTPUT_DIR}/highlights.mp4" + output_filepath = f"{self._output_dir}/{self._output_filename}" self._ffmpeg_concat_clips(clips, output_filepath) logging.info(f"Produced video: {output_filepath}")