Browse Source

refactor!: viz+json producer: add simple config option for output

BREAKING CHANGE: default is now to save to current working directory instead of tmp
main
Rob Hallam 3 months ago
parent
commit
f160a83967
1 changed files with 13 additions and 6 deletions
  1. +13
    -6
      pipeline/producers.py

+ 13
- 6
pipeline/producers.py View File

@@ -24,7 +24,6 @@ class Producer(ABC):
class VideoProducer(Producer): class VideoProducer(Producer):
"""Video producer interface.""" """Video producer interface."""



class FfmpegVideoProducer(VideoProducer): class FfmpegVideoProducer(VideoProducer):
"""Produce videos using ffmpeg""" """Produce videos using ffmpeg"""
# TODO: consider output filename options # TODO: consider output filename options
@@ -105,11 +104,16 @@ class FfmpegVideoProducer(VideoProducer):


class VisualisationProducer(Producer): class VisualisationProducer(Producer):
"""Visualisation producer -- illustrate the features we have extracted""" """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: if not features:
raise ValueError("No features provided") raise ValueError("No features provided")
self.features = features self.features = features


if not output_filepath:
raise ValueError("No output filepath provided")
self.output_filepath = output_filepath # TODO: sanity check this

def produce(self): def produce(self):
"""Produce visualisation""" """Produce visualisation"""
# basic idea: use matplotlib to plot: # basic idea: use matplotlib to plot:
@@ -150,7 +154,7 @@ class VisualisationProducer(Producer):
ax.set_yticks([]) ax.set_yticks([])
# ax.tick_params(axis='y', labelrotation=90, ha='right') # ax.tick_params(axis='y', labelrotation=90, ha='right')
# save the plot # save the plot
plt.savefig("/tmp/visualisation.png")
plt.savefig(self.output_filepath)
plt.close() plt.close()


class PipelineJSONEncoder(json.JSONEncoder): class PipelineJSONEncoder(json.JSONEncoder):
@@ -162,12 +166,15 @@ class PipelineJSONEncoder(json.JSONEncoder):


class JSONProducer(Producer): class JSONProducer(Producer):
"""Produce JSON output""" """Produce JSON output"""
def __init__(self, features):
DEFAULT_OUTPUT_FILEPATH = "features.json"
def __init__(self, features, output_filepath=DEFAULT_OUTPUT_FILEPATH):
if not features: if not features:
raise ValueError("No features provided") raise ValueError("No features provided")
self.features = features self.features = features
if not output_filepath:
raise ValueError("No output filepath provided")
self.output_filepath = output_filepath


def produce(self): 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)) jsonfile.write(json.dumps(self.features, cls=PipelineJSONEncoder, indent=4))

Loading…
Cancel
Save