Переглянути джерело

Merge branch 'feat-config'

Adds config to isualisation and JSON producers
main
Rob Hallam 2 місяці тому
джерело
коміт
65e1a6665e
1 змінених файлів з 12 додано та 6 видалено
  1. +12
    -6
      pipeline/producers.py

+ 12
- 6
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))

Завантаження…
Відмінити
Зберегти