Browse Source

Merge branch 'feat-config'

Adds config to isualisation and JSON producers
main
Rob Hallam 2 months ago
parent
commit
65e1a6665e
1 changed files with 12 additions and 6 deletions
  1. +12
    -6
      pipeline/producers.py

+ 12
- 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
@@ -108,10 +107,14 @@ 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


def _fe_colour(self, feature) -> str: def _fe_colour(self, feature) -> str:
"""Return a colour for a feature """Return a colour for a feature
@@ -173,7 +176,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):
@@ -185,12 +188,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