whisper_edit.py

@@ -0,0 +1,219 @@
+#!/usr/bin/env python3
+import whisper
+import sys
+import string
+import math
+from stable_whisper import modify_model
+from stable_whisper import stabilize_timestamps
+from stable_whisper import results_to_word_srt
+from stable_whisper import results_to_word_srt
+from moviepy.editor import AudioClip, VideoFileClip, concatenate_videoclips
+MODEL = 'base'
+# Store phonetic alphabet in a dictionary.
+phonetic_alphabet = {
+  'A': 'Alpha',
+  'B': 'Bravo',
+  'C': 'Charlie',
+  'D': 'Delta',
+  'E': 'Echo',
+  'F': 'Foxtrot',
+  'G': 'Golf',
+  'H': 'Hotel',
+  'I': 'India',
+  'J': 'Juliett',
+  'K': 'Kilo',
+  'L': 'Lima',
+  'M': 'Mike',
+  'N': 'November',
+  'O': 'Oscar',
+  'P': 'Papa',
+  'Q': 'Quebec',
+  'R': 'Romeo',
+  'S': 'Sierra',
+  'T': 'Tango',
+  'U': 'Uniform',
+  'V': 'Victor',
+  'W': 'Whiskey',
+  'X': 'X',
+  'Y': 'Yankee',
+  'Z': 'Zulu',
+}
+
+# Commands
+class Commands:
+  COMMAND_PREFIX = 'victor'
+  keep_segment = [COMMAND_PREFIX, 'kilo']
+  drop_segment = [COMMAND_PREFIX, 'delta']
+  chapter_break = [COMMAND_PREFIX, 'charlie']
+  COMMANDS = {
+    'keep_segment': keep_segment,
+    'drop_segment': drop_segment,
+    'chapter_break': chapter_break,
+  }
+input_file = 'input.mp4'
+output_file = 'output.mp4'
+def speech_to_text():
+  model = whisper.load_model(MODEL)
+  modify_model(model)
+  print('Speech to text...')
+  results = model.transcribe(input_file)
+  print('... done!\n')
+  return results
+def word_level_timestamps(parsed_speech_to_text):
+  print('Getting word-level timestamps...')
+  word_segments = stabilize_timestamps(parsed_speech_to_text, top_focus=True)
+  #print(word_segments)
+  print('... done!\n')
+  return word_segments
+def string_canonical(s):
+  # Remove punctuation.
+  no_punc = s.translate(str.maketrans('', '', string.punctuation))
+  return no_punc.strip().lower()
+
+# Return timestamp of (i + 1)th word. Used to get the ending timestamp
+# of a command.
+def timestamp_of_next_word(seg, i):
+  word_timestamps = seg['whole_word_timestamps']
+  num_words = len(word_timestamps)
+  if (i + 1) >= num_words:
+    # No more words in this segment.
+    return seg['end']
+  else:
+    return seg['whole_word_timestamps'][i + 1]['timestamp']
+
+def find_command_in_segment(seg):
+  word_timestamps = seg['whole_word_timestamps']
+  num_words = len(word_timestamps)
+  i = 0
+  while i < num_words:
+    word = string_canonical(word_timestamps[i]['word'])
+    # print(f'{word} ')
+    if word == Commands.COMMAND_PREFIX:
+      #print('! ')
+      time_command_start = word_timestamps[i]['timestamp']
+      i += 1
+      if i > num_words:
+        break
+      second_word_of_command = string_canonical(word_timestamps[i]['word'])
+      match second_word_of_command:
+        case 'kilo':
+          time_command_end = timestamp_of_next_word(seg, i)
+          return Commands.keep_segment, time_command_start, time_command_end
+        case 'delta':
+          time_command_end = timestamp_of_next_word(seg, i)
+          return Commands.drop_segment, time_command_start, time_command_end
+        case _:
+          continue
+    i += 1
+
+def find_commands(timestamps):
+  commands = []
+  for segment in timestamps:
+    seg_command = find_command_in_segment(segment)
+    if seg_command:
+      commands.append(seg_command)
+  return commands
+
+# Hack. When dropping a segment I see a split-second of the end of
+# the dropped segment. So add a small fuzz factor for that.
+DROP_SEGMENT_DELTA = 0.2
+
+# Returns list of (start, end) tuples of intervals to keep.
+def intervals_to_keep(commands):
+  keep_intervals = []
+  keep_start, keep_end = 0, 0
+  for command_timestamp in commands:
+    cmd, begin_ts, end_ts = command_timestamp
+    match cmd:
+      case Commands.keep_segment:
+        # Keep until the start of the command.
+        keep_end = begin_ts
+        keep_intervals.append([keep_start, keep_end])
+        # Next (possibly) starts at end of command.
+        keep_start = end_ts
+      case Commands.drop_segment:
+        # Next (possibly) starts at end of command.
+        keep_start = end_ts + DROP_SEGMENT_DELTA
+      case _:
+        print(f'Eeek! Unrecognized command: {cmd}')
+  return keep_intervals
+# Iterate over audio to find the non-silent parts. Outputs a list of
+# (speaking_start, speaking_end) intervals.
+# Args:
+#  window_size: (in seconds) hunt for silence in windows of this size
+#  volume_threshold: volume below this threshold is considered to be silence
+#  ease_in: (in seconds) add this much silence around speaking intervals
+def find_speaking_intervals(audio_clip, window_size=0.1, volume_threshold=0.05, ease_in=0.1, audio_fps=44100):
+  # First, iterate over audio to find all silent windows.
+  num_windows = math.floor(audio_clip.end/window_size)
+  window_is_silent = []
+  for i in range(num_windows):
+    s = audio_clip.subclip(i * window_size, (i + 1) * window_size).set_fps(audio_fps)
+    v = s.max_volume()
+    window_is_silent.append(v < volume_threshold)
+
+  # Find speaking intervals.
+  speaking_start = 0
+  speaking_end = 0
+  speaking_intervals = []
+  for i in range(1, len(window_is_silent)):
+    e1 = window_is_silent[i - 1]
+    e2 = window_is_silent[i]
+    # silence -> speaking
+    if e1 and not e2:
+      speaking_start = i * window_size
+    # speaking -> silence, now have a speaking interval
+    if not e1 and e2:
+      speaking_end = i * window_size
+      new_speaking_interval = [max(0, speaking_start - ease_in), speaking_end + ease_in]
+      # With tiny windows, this can sometimes overlap the previous window, so merge.
+      need_to_merge = len(speaking_intervals) > 0 and speaking_intervals[-1][1] > new_speaking_interval[0]
+      if need_to_merge:
+        merged_interval = [speaking_intervals[-1][0], new_speaking_interval[1]]
+        speaking_intervals[-1] = merged_interval
+      else:
+        speaking_intervals.append(new_speaking_interval)
+  return speaking_intervals
+
+def find_speaking(input_clip, input_audio_fps):
+  print("\n\n\n----- Now cutting out dead air... -----")
+  speaking_intervals = find_speaking_intervals(input_clip.audio, audio_fps=input_audio_fps)
+  print("Keeping speaking intervals: " + str(speaking_intervals))
+  speaking_clips = [input_clip.subclip(start, end) for [start, end] in speaking_intervals]
+  final_video = concatenate_videoclips(speaking_clips)
+  return final_video  
+
+def main():
+  sts = speech_to_text()
+  word_ts = word_level_timestamps(sts)
+  commands = find_commands(word_ts)
+  print(commands)
+  keep_intervals = intervals_to_keep(commands)
+  print(keep_intervals)
+
+  vid = VideoFileClip(input_file)
+
+  # Edit with speech-to-text.
+  keep_clips = [vid.subclip(start, end) for [start, end] in keep_intervals]
+  edited_vid = concatenate_videoclips(keep_clips)
+
+  # Cut out dead air.
+  no_dead_air_video = find_speaking(edited_vid, vid.audio.fps)
+
+  print("\n\n\n----- Writing out edited video... -----")
+  no_dead_air_video.write_videofile(output_file,
+        #fps=60,
+        preset='ultrafast',
+        codec='libx264',
+        #codec='h264_videotoolbox',
+        temp_audiofile='temp-audio.m4a',
+        remove_temp=True,
+        audio_codec="aac",
+        #threads=6,
+        ffmpeg_params = ['-threads', '8'],
+    )
+  vid.close()
+
+
+if __name__ == '__main__':
+  main()