diff --git a/pipeline/video_producers.py b/pipeline/video_producers.py index fb02c0e..014d348 100644 --- a/pipeline/video_producers.py +++ b/pipeline/video_producers.py @@ -5,8 +5,9 @@ import logging import subprocess import tempfile -class VideoProducer(ABC): - """Video producer interface.""" +# for visualisations: +import matplotlib.pyplot as plt + class Producer(ABC): """Generic producer interface.""" def __init__(self, features): @@ -87,3 +88,35 @@ class FfmpegVideoProducer(VideoProducer): output_filepath = f"{OUTPUT_DIR}/highlights.mp4" self._ffmpeg_concat_clips(clips, output_filepath) logging.info(f"Produced video: {output_filepath}") + +class VisualisationProducer(Producer): + """Visualisation producer -- illustrate the features we have extracted""" + def __init__(self, features): + if not features: + raise ValueError("No features provided") + self.features = features + + def produce(self): + """Produce visualisation""" + # basic idea: use matplotlib to plot: + # - a wide line segment representing the source video[s] + # - shorter line segments representing the features extracted where: + # + width represents duration + # + colour represents feature type + # + position represents time + # - save as image + plotted_source_videos = [] + fig, ax = plt.subplots() + for feature in self.features: + # plot source video line if not done already + if feature.source not in plotted_source_videos: + # use video duration as width + # ax.plot([0, feature.source.duration()], [0, 0], color='black', linewidth=10) + ax.broken_barh([(0, feature.source.duration())], (0, 5), facecolors='grey') + plotted_source_videos.append(feature.source) + # plot feature line + # ax.plot([feature.interval.start, feature.interval.end], [1, 1], color='red', linewidth=5) + ax.broken_barh([(feature.interval.start, feature.interval.duration)], (10, 5), facecolors='red') + # save the plot + plt.savefig("/tmp/visualisation.png") + plt.close()