Bläddra i källkod

feat: add VisualisationProducer

Creates a plot showing where in source video(s) the pipeline features
came from
main
Rob Hallam 3 månader sedan
förälder
incheckning
214f8c974a
1 ändrade filer med 35 tillägg och 2 borttagningar
  1. +35
    -2
      pipeline/video_producers.py

+ 35
- 2
pipeline/video_producers.py Visa fil

@@ -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()

Laddar…
Avbryt
Spara