From 0bd39a141573ed7fac66f0ac63165f5d6ffd405a Mon Sep 17 00:00:00 2001
From: Rob Hallam <0504004h@student.gla.ac.uk>
Date: Sun, 15 Sep 2024 23:18:24 +0100
Subject: [PATCH] feat: [VizP] different colours for different FEs

---
 pipeline/producers.py | 26 +++++++++++++++++++++++---
 1 file changed, 23 insertions(+), 3 deletions(-)

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')