diff --git a/pipeline/video_producers.py b/pipeline/video_producers.py index 014d348..d078cee 100644 --- a/pipeline/video_producers.py +++ b/pipeline/video_producers.py @@ -2,6 +2,7 @@ from abc import ABC import logging +import os import subprocess import tempfile @@ -106,6 +107,8 @@ class VisualisationProducer(Producer): # + position represents time # - save as image plotted_source_videos = [] + bar_labels = [] + fig, ax = plt.subplots() for feature in self.features: # plot source video line if not done already @@ -114,9 +117,25 @@ class VisualisationProducer(Producer): # 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) + bar_labels.append(os.path.basename(feature.source.path)) + # annotate the source video + ax.text(0.25, 0.25, os.path.basename(feature.source.path), ha='left', va='bottom', + fontsize=16) + # 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') + if feature.feature_extractor not in bar_labels: + bar_labels.append(feature.feature_extractor) + # label bar with feature extractor + ax.text(0, 8, feature.feature_extractor, ha='left', va='bottom', + fontsize=16) + + # label the plot's axes + ax.set_xlabel('Time') + # ax.set_yticks([], labels=bar_labels) + ax.set_yticks([]) + # ax.tick_params(axis='y', labelrotation=90, ha='right') # save the plot plt.savefig("/tmp/visualisation.png") plt.close()