diff --git a/pipeline/producers.py b/pipeline/producers.py index 5f64901..704396c 100644 --- a/pipeline/producers.py +++ b/pipeline/producers.py @@ -113,6 +113,25 @@ class VisualisationProducer(Producer): raise ValueError("No features provided") self.features = features + def _fe_colour(self, feature) -> str: + """Return a colour for a feature + + laughter: red + loudness: blue + video activity: green + words: purple + default: pink + """ + if feature.feature_extractor == "laughter": + return "red" + if feature.feature_extractor == "loudness": + return "blue" + if feature.feature_extractor == "videoactivity": + return "green" + if feature.feature_extractor == "words": + return "purple" + return "black" + def produce(self): """Produce visualisation""" # basic idea: use matplotlib to plot: @@ -140,12 +159,13 @@ class VisualisationProducer(Producer): # 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') + ax.broken_barh([(feature.interval.start, feature.interval.duration)], + (10, 5), facecolors=f'{self._fe_colour(feature)}') 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) + # ax.text(0, 8, feature.feature_extractor, ha='left', va='bottom', + # fontsize=16) # label the plot's axes ax.set_xlabel('Time')