diff --git a/pipeline/producers.py b/pipeline/producers.py index 704396c..37230f3 100644 --- a/pipeline/producers.py +++ b/pipeline/producers.py @@ -24,7 +24,6 @@ class Producer(ABC): class VideoProducer(Producer): """Video producer interface.""" - class FfmpegVideoProducer(VideoProducer): """Produce videos using ffmpeg""" # TODO: consider output filename options @@ -108,10 +107,14 @@ class FfmpegVideoProducer(VideoProducer): class VisualisationProducer(Producer): """Visualisation producer -- illustrate the features we have extracted""" - def __init__(self, features): + DEFAULT_OUTPUT_FILEPATH = "visualisation.png" + def __init__(self, features, output_filepath=DEFAULT_OUTPUT_FILEPATH): if not features: raise ValueError("No features provided") self.features = features + if not output_filepath: + raise ValueError("No output filepath provided") + self.output_filepath = output_filepath def _fe_colour(self, feature) -> str: """Return a colour for a feature @@ -173,7 +176,7 @@ class VisualisationProducer(Producer): ax.set_yticks([]) # ax.tick_params(axis='y', labelrotation=90, ha='right') # save the plot - plt.savefig("/tmp/visualisation.png") + plt.savefig(self.output_filepath) plt.close() class PipelineJSONEncoder(json.JSONEncoder): @@ -185,12 +188,15 @@ class PipelineJSONEncoder(json.JSONEncoder): class JSONProducer(Producer): """Produce JSON output""" - def __init__(self, features): + DEFAULT_OUTPUT_FILEPATH = "features.json" + def __init__(self, features, output_filepath=DEFAULT_OUTPUT_FILEPATH): if not features: raise ValueError("No features provided") self.features = features + if not output_filepath: + raise ValueError("No output filepath provided") + self.output_filepath = output_filepath def produce(self): - # FIXME: config option for output path - with open("/tmp/features.json", "w") as jsonfile: + with open(self.output_filepath, "w") as jsonfile: jsonfile.write(json.dumps(self.features, cls=PipelineJSONEncoder, indent=4))