add sentiment editor

src/editors/amplitude/editor.py

@@ -1,6 +1,5 @@
 from ...mediautils.audio import process_audio, resample
 from ...math.average import np_moving_average
-from ...models.moment import Moment
 from ..common import find_moving_average_highlights
 import numpy as np
 import structlog

src/editors/sentiment/editor.py

@@ -1,3 +1,48 @@
+import whisper
+import numpy as np
+import structlog
+
+from flair.models import TextClassifier
+from dataclasses import dataclass
+from flair.data import Sentence
+
+from ...math.average import np_moving_average
+from ..common import find_moving_average_highlights
+
+@dataclass
+class TextGlob:
+    start:float
+    stop:float
+    text:str
+    sentiment:float
 
 class SentimentEditor:
-    pass
+    def __init__(self, video_path, audio_path, params):
+        self.logger = structlog.get_logger("sentiment")
+        self.logger.info("loading whisper model", size=params["model_size"])
+        self.model = whisper.load_model(params["model_size"])
+        self.logger.info("transcribing audio", path=audio_path)
+        self.result = self.model.transcribe(audio_path)
+        self.segments = []
+        for segment in self.result['segments']:
+            self.segments.append(TextGlob(segment['start'], segment['end'], segment['text'], 0))
+        classifier = TextClassifier.load('en-sentiment')
+        self.sentiments = []
+        self.logger.info("calculating sentiment on segments", segments=len(self.segments))
+        for segment in self.segments:
+            sentence = Sentence(segment.text)
+            classifier.predict(sentence)
+            sentsum = sentence.labels[0].score
+            if sentence.labels[0].value == "NEGATIVE":
+                sentsum = sentsum * -1
+            segment.sentiment = sentsum
+            self.sentiments.append(sentsum)
+        self.sentiments = np.array(self.sentiments)
+   
+    def edit(self, large_window, small_window, params):
+        end_time = self.segments[-1].stop
+        window_factor = len(self.sentiments) / end_time
+        long_ma = np_moving_average(self.squared_subsample, large_window * window_factor)
+        short_ma = np_moving_average(self.squared_subsample, small_window * window_factor)
+        highlights = find_moving_average_highlights(short_ma, long_ma, 1.0 / window_factor)
+        return highlights, large_window, small_window