Refactor pad and merge timestamps into one function

This also fixes a bunch of issues regarding when the timestamps
should be merged.

src/segments.py

@@ -0,0 +1,47 @@
+from typing import Any, Dict, List
+
+import copy
+
+def merge_timestamps(timestamps: List[Dict[str, Any]], merge_window: float = 5, max_merge_size: float = 30, padding_left: float = 1, padding_right: float = 1):
+    result = []
+
+    if len(timestamps) == 0:
+        return result
+
+    processed_time = 0
+    current_segment = None
+
+    for i in range(len(timestamps)):
+        next_segment = timestamps[i]
+
+        delta = next_segment['start'] - processed_time
+
+        # Note that segments can still be longer than the max merge size, they just won't be merged in that case
+        if current_segment is None or delta > merge_window or next_segment['end'] - current_segment['start'] > max_merge_size:
+            # Finish the current segment
+            if current_segment is not None:
+                # Add right padding
+                finish_padding = min(padding_right, delta / 2) if delta < padding_left + padding_right else padding_right
+                current_segment['end'] += finish_padding
+                delta -= finish_padding
+
+                result.append(current_segment)
+
+            # Start a new segment
+            current_segment = copy.deepcopy(next_segment)
+
+            # Pad the segment
+            current_segment['start'] = current_segment['start'] - min(padding_left, delta)
+            processed_time = current_segment['end']
+
+        else:
+            # Merge the segment
+            current_segment['end'] = next_segment['end']
+            processed_time = current_segment['end']
+        
+    # Add the last segment
+    if current_segment is not None:
+        current_segment['end'] += padding_right
+        result.append(current_segment)
+    
+    return result

src/vad.py

@@ -5,6 +5,8 @@ from typing import Any, Deque, Iterator, List, Dict
 
 from pprint import pprint
 
+from src.segments import merge_timestamps
+
 # Workaround for https://github.com/tensorflow/tensorflow/issues/48797
 try:
     import tensorflow as tf
@@ -110,8 +112,10 @@ class AbstractTranscription(ABC):
         # get speech timestamps from full audio file
         seconds_timestamps = self.get_transcribe_timestamps(audio)
 
-        padded = self.pad_timestamps(seconds_timestamps, self.segment_padding_left, self.segment_padding_right)
-        merged = self.merge_timestamps(padded, self.max_silent_period, self.max_merge_size, self.min_force_merge_gap, self.max_force_merge_size)
+        #for seconds_timestamp in seconds_timestamps:
+        #    print("VAD timestamp ", format_timestamp(seconds_timestamp['start']), " to ", format_timestamp(seconds_timestamp['end']))
+
+        merged = merge_timestamps(seconds_timestamps, self.max_silent_period, self.max_merge_size, self.segment_padding_left, self.segment_padding_right)
 
         # A deque of transcribed segments that is passed to the next segment as a prompt
         prompt_window = deque()
@@ -346,70 +350,6 @@ class AbstractTranscription(ABC):
             result.append(new_segment)
         return result
 
-    def pad_timestamps(self, timestamps: List[Dict[str, Any]], padding_left: float, padding_right: float):
-        if (padding_left == 0 and padding_right == 0):
-            return timestamps
-        
-        result = []
-        prev_entry = None
-
-        for i in range(len(timestamps)):
-            curr_entry = timestamps[i]
-            next_entry = timestamps[i + 1] if i < len(timestamps) - 1 else None
-
-            segment_start = curr_entry['start']
-            segment_end = curr_entry['end']
-
-            if padding_left is not None:
-                segment_start = max(prev_entry['end'] if prev_entry else 0, segment_start - padding_left)
-            if padding_right is not None:
-                segment_end = segment_end + padding_right
-
-                # Do not pad past the next segment
-                if (next_entry is not None):
-                    segment_end = min(next_entry['start'], segment_end)
-
-            new_entry = { 'start': segment_start, 'end': segment_end }
-            prev_entry = new_entry
-            result.append(new_entry)
-
-        return result
-
-    def merge_timestamps(self, timestamps: List[Dict[str, Any]], max_merge_gap: float, max_merge_size: float, 
-                                min_force_merge_gap: float, max_force_merge_size: float):
-        if max_merge_gap is None:
-            return timestamps
-
-        result = []
-        current_entry = None
-
-        for entry in timestamps:
-            if current_entry is None:
-                current_entry = entry
-                continue
-
-            # Get distance to the previous entry
-            distance = entry['start'] - current_entry['end']
-            current_entry_size = current_entry['end'] - current_entry['start']
-
-            if distance <= max_merge_gap and (max_merge_size is None or current_entry_size <= max_merge_size):
-                # Regular merge
-                current_entry['end'] = entry['end']
-            elif min_force_merge_gap is not None and distance <= min_force_merge_gap and \
-                 (max_force_merge_size is None or current_entry_size <= max_force_merge_size):
-                # Force merge if the distance is small (up to a certain maximum size)
-                current_entry['end'] = entry['end']
-            else:
-                # Output current entry
-                result.append(current_entry)
-                current_entry = entry
-        
-        # Add final entry
-        if current_entry is not None:
-            result.append(current_entry)
-
-        return result
-
     def multiply_timestamps(self, timestamps: List[Dict[str, Any]], factor: float):
         result = []
 

tests/segments_test.py

@@ -0,0 +1,48 @@
+import sys
+import unittest
+
+sys.path.append('../whisper-webui')
+
+from src.segments import merge_timestamps
+
+class TestSegments(unittest.TestCase):
+    def __init__(self, *args, **kwargs):
+        super(TestSegments, self).__init__(*args, **kwargs)
+        
+    def test_merge_segments(self):
+        segments = [
+            {'start': 10.0, 'end': 20.0},
+            {'start': 22.0, 'end': 27.0},
+            {'start': 31.0, 'end': 35.0},
+            {'start': 45.0, 'end': 60.0},
+            {'start': 61.0, 'end': 65.0},
+            {'start': 68.0, 'end': 98.0},
+            {'start': 100.0, 'end': 102.0},
+            {'start': 110.0, 'end': 112.0}
+        ]
+
+        result = merge_timestamps(segments, merge_window=5, max_merge_size=30, padding_left=1, padding_right=1)
+
+        self.assertListEqual(result, [
+            {'start': 9.0, 'end': 36.0},
+            {'start': 44.0, 'end': 66.0},
+            {'start': 67.0, 'end': 99.0},
+            {'start': 99.0, 'end': 103.0},
+            {'start': 109.0, 'end': 113.0}
+        ])
+
+    def test_overlap_next(self):
+        segments = [
+            {'start': 5.0, 'end': 39.182},
+            {'start': 39.986, 'end': 40.814}
+        ]
+
+        result = merge_timestamps(segments, merge_window=5, max_merge_size=30, padding_left=1, padding_right=1)
+
+        self.assertListEqual(result, [
+            {'start': 4.0, 'end': 39.584},
+            {'start': 39.584, 'end': 41.814}
+        ])
+
+if __name__ == '__main__':
+    unittest.main()