More refactoring. Opus now talking to Whisper.

audiosource.py

@@ -0,0 +1,176 @@
+#!/usr/bin/python3
+
+import socket
+import select
+import time
+from queue import Queue
+import json
+import threading
+import speech_recognition
+import wave
+
+from pyogg.opus_decoder import OpusDecoder
+
+
+class AudioSource:
+    def __init__(self):
+        # Thread safe Queue for passing data from the threaded recording
+        # callback.
+        self.data_queue = Queue()
+
+    def is_done(self):
+        return True
+
+
+# -----------------------------------------------------------------------------
+# Microphone
+
+# How real time the recording is in seconds.
+record_timeout = 2
+
+class MicrophoneAudioSource(AudioSource):
+
+    def __init__(self):
+        super().__init__()
+        
+        self._recorder = speech_recognition.Recognizer()
+        self._recorder.energy_threshold = 1200
+
+        # Definitely do this, dynamic energy compensation lowers the energy
+        # threshold dramatically to a point where the SpeechRecognizer
+        # never stops recording.
+        self._recorder.dynamic_energy_threshold = False
+   
+        self._source = speech_recognition.Microphone(sample_rate=16000)
+
+
+        with self._source:
+            self._recorder.adjust_for_ambient_noise(self._source)
+
+        def record_callback(_, audio:speech_recognition.AudioData) -> None:
+            """
+            Threaded callback function to receive audio data when recordings finish.
+            audio: An AudioData containing the recorded bytes.
+            """
+
+            # Grab the raw bytes and push it into the thread safe queue.
+            data = audio.get_raw_data()
+            self.data_queue.put(bytearray(data))
+
+        # Create a background thread that will pass us raw audio bytes.
+        # We could do this manually but SpeechRecognizer provides a nice helper.
+        self._stopper = self._recorder.listen_in_background(
+            self._source, record_callback,
+            phrase_time_limit=record_timeout)
+
+    def stop(self):
+        assert(self._stopper)
+        self._stopper()
+        
+        self._recorder = None
+        self._stopper = None
+        self._source = None
+
+    def is_done(self):
+        return self._recorder == None
+
+# -----------------------------------------------------------------------------
+# Opus stream
+
+# For debugging
+# wave_out = wave.open("wave2.wav", "wb")
+# wave_out.setnchannels(1)
+# wave_out.setframerate(16000)
+# wave_out.setsampwidth(2)
+
+class OpusStreamAudioSource(AudioSource):
+
+    def __init__(self, sock):
+        super().__init__()
+
+        self._socket = sock
+
+        self._opus_decoder = OpusDecoder()
+        self._opus_decoder.set_channels(1)
+        self._opus_decoder.set_sampling_frequency(16000)
+
+        # Fetch user info.
+        user_info_tmp = self._read_packet(self._socket)
+        self._user_info = json.loads(user_info_tmp.decode("utf-8"))
+        print("User connection...")
+        print(json.dumps(self._user_info, indent=4))
+
+        self._is_done = False
+
+        # Start input thread.
+        self._input_thread = threading.Thread(
+            target=self._input_thread_func, daemon=True)
+        self._input_thread.start()
+
+    def _read_packet(self, sock):
+        try:
+            input_buffer = b''
+            #print("Reading packet size...")
+            while len(input_buffer) < 4:
+                input_buffer = input_buffer + sock.recv(1)
+                if not input_buffer:
+                    raise Exception("Failed to read size of packet.")
+
+            packet_size = int.from_bytes(input_buffer, "little")
+            #print("Packet size: ", packet_size)
+
+            input_buffer = b''
+            while len(input_buffer) < packet_size:
+                input_buffer = input_buffer + sock.recv(1)
+                if not input_buffer:
+                    raise Exception("Failed to read packet.")
+
+            return input_buffer
+        
+        except Exception as e: # FIXME: Use socket-specific exception type.
+            return None
+
+    def _input_thread_func(self):
+
+        print("input thread start")
+        try:
+
+            while not self._is_done:
+
+                next_packet = self._read_packet(self._socket)
+
+                if next_packet:
+
+                    # If we don't use bytearray here to copy, we run into a weird
+                    # exception about the memory not being writeable.
+                    decoded_data = self._opus_decoder.decode(bytearray(next_packet))
+
+                    # For debugging.
+                    #wave_out.writeframes(decoded_data)
+                    
+                    # We need to copy decoded_data here or we end up with
+                    # recycled buffers in our queue, which leads to broken
+                    # audio.
+                    self.data_queue.put(bytearray(decoded_data))
+
+                else:
+                    break
+            
+        except Exception as e:
+           # Probably disconnected. We don't care. Just clean up.
+           # FIXME: Limit exception to socket errors.
+           pass
+
+        print("input thread done")
+
+        self._is_done = True
+
+    def stop(self):
+        self._is_done = True
+        
+        # We won't join() the input thread because we don't want to sit around
+        # and wait for a packet. It'll die on its own, so whatever.
+    
+    def is_done(self):
+        return self._is_done
+

difftest.py

@@ -0,0 +1,44 @@
+import difflib
+
+s1 = "1234asdffooMOO"
+s2 = "asdfbarMOOwhatever"
+
+# s1 = "asdffoo"
+# s2 = "asdffooMOO"
+
+def onestepchange(start, dest):
+
+    ret = ""
+
+    for i, s in enumerate(difflib.ndiff(start, dest)):
+        # print(i)
+        # print(s)
+
+        if s[0] == '-':
+            return ret + start[i+1:]
+
+        if s[1] == '+':
+            return ret + s[-1] + start[i:]
+
+        ret = ret + s[-1]
+
+        if len(ret) > len(start):
+            return ret
+
+        if ret[i] != start[i]:
+            return ret + start[i:]
+
+    return ret
+
+
+n = s1
+while n != s2:
+    print(n)
+    n = onestepchange(n, s2)
+
+print(n)
+
+# for i, s in enumerate(difflib.ndiff(s1, s2)):
+#     print(i)
+#     print(s)
+

kiri_reqs.txt

@@ -0,0 +1 @@
+whisper-live tokenizers==0.20.3

requirements.txt

@@ -1,7 +1,8 @@
 setuptools
 pyaudio
 SpeechRecognition
---extra-index-url https://download.pytorch.org/whl/cu116
+--extra-index-url https://download.pytorch.org/whl/rocm6.2.4
 torch
 numpy
 git+https://github.com/openai/whisper.git
+git+https://github.com/TeamPyOgg/PyOgg.git@4118fc40067eb475468726c6bccf1242abfc24fc

requirements2.txt

@@ -0,0 +1,8 @@
+setuptools
+pyaudio
+SpeechRecognition
+--extra-index-url https://download.pytorch.org/whl/rocm6.2.4
+torch
+numpy
+git+https://github.com/openai/whisper.git
+git+https://github.com/TeamPyOgg/PyOgg.git@4118fc40067eb475468726c6bccf1242abfc24fc

transcribe_demo.py

@@ -1,9 +1,16 @@
 #! python3.7
 
+# Recent phrases to include in the text buffer before the current transcription.
+recent_phrase_count = 8
+
+# How real time the recording is in seconds.
+record_timeout = 2
+
+
 import argparse
 import os
 import numpy as np
-import speech_recognition as sr
+import speech_recognition
 import whisper
 import torch
 
@@ -13,6 +20,105 @@ from time import sleep
 from sys import platform
 
 import textwrap
+import difflib
+
+import pygame
+
+import wave
+#from pyogg.opus import OpusEncoder
+
+import socket
+import select
+import time
+import json
+import threading
+
+
+import diffstuff
+import audiosource
+from transcriber import Transcriber
+
+pygame_font_height = 16
+pygame.init()
+pygame_display_surface = pygame.display.set_mode((1280, pygame_font_height * 2))
+pygame.display.set_caption("Transcription")
+pygame_font = pygame.font.Font("/home/kiri/.fonts/Sigmar-Regular.ttf", pygame_font_height)
+
+
+
+
+opus_server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+print("binding")
+opus_server_socket.bind(("127.0.0.1", 9967))
+print("set non blocking")
+#opus_server_socket.setblocking(False)
+print("listening")
+opus_server_socket.listen()
+
+
+wave_out = wave.open("wave.wav", "wb")
+wave_out.setnchannels(1)
+wave_out.setframerate(16000)
+wave_out.setsampwidth(2)
+
+transcriber = Transcriber()
+mic_source = audiosource.MicrophoneAudioSource()
+
+while True:
+    print("looping")
+
+    # Wait for a connection.
+    while True:
+        s = select.select([opus_server_socket], [], [], 0)
+        time.sleep(0.01)
+        if len(s[0]):
+            accepted_socket, addr = opus_server_socket.accept()
+            print(accepted_socket)
+            break
+
+    opusSource = audiosource.OpusStreamAudioSource(accepted_socket)
+
+    transcriber.set_source(opusSource)
+    while not opusSource.is_done():
+        time.sleep(0.1)
+        transcriber.update()
+
+
+exit(0)
+
+
+
+
+def onestepchange(start, dest):
+
+    ret = ""
+
+    for i, s in enumerate(difflib.ndiff(start, dest)):
+        # print(i)
+        # print(s)
+
+        if s[0] == '-':
+            return ret + start[i+1:]
+
+        if s[1] == '+':
+            return ret + s[-1] + start[i:]
+
+        ret = ret + s[-1]
+
+        if len(ret) > len(start):
+            return ret
+
+        if ret[i] != start[i]:
+            return ret + start[i:]
+
+    return ret
+
+def countsteps(start, dest):
+    step_count = 0
+    while start != dest:
+        start = onestepchange(start, dest)
+        step_count += 1
+    return step_count
 
 def main():
     parser = argparse.ArgumentParser()
@@ -22,8 +128,6 @@ def main():
                         help="Don't use the english model.")
     parser.add_argument("--energy_threshold", default=1000,
                         help="Energy level for mic to detect.", type=int)
-    parser.add_argument("--record_timeout", default=2,
-                        help="How real time the recording is in seconds.", type=float)
     parser.add_argument("--phrase_timeout", default=3,
                         help="How much empty space between recordings before we "
                              "consider it a new line in the transcription.", type=float)
@@ -35,30 +139,8 @@ def main():
 
     # The last time a recording was retrieved from the queue.
     phrase_time = None
-    # Thread safe Queue for passing data from the threaded recording callback.
-    data_queue = Queue()
+    #data_queue = Queue()
     # We use SpeechRecognizer to record our audio because it has a nice feature where it can detect when speech ends.
-    recorder = sr.Recognizer()
-    recorder.energy_threshold = args.energy_threshold
-    # Definitely do this, dynamic energy compensation lowers the energy threshold dramatically to a point where the SpeechRecognizer never stops recording.
-    recorder.dynamic_energy_threshold = False
-
-    # Important for linux users.
-    # Prevents permanent application hang and crash by using the wrong Microphone
-    if 'linux' in platform:
-        mic_name = args.default_microphone
-        if not mic_name or mic_name == 'list':
-            print("Available microphone devices are: ")
-            for index, name in enumerate(sr.Microphone.list_microphone_names()):
-                print(f"Microphone with name \"{name}\" found")
-            return
-        else:
-            for index, name in enumerate(sr.Microphone.list_microphone_names()):
-                if mic_name in name:
-                    source = sr.Microphone(sample_rate=16000, device_index=index)
-                    break
-    else:
-        source = sr.Microphone(sample_rate=16000)
 
     # Load / Download model
     model = args.model
@@ -66,107 +148,123 @@ def main():
         model = model + ".en"
     audio_model = whisper.load_model(model)
 
-    record_timeout = args.record_timeout
     phrase_timeout = args.phrase_timeout
 
     transcription = ['']
 
-    with source:
-        recorder.adjust_for_ambient_noise(source)
-
-    def record_callback(_, audio:sr.AudioData) -> None:
-        """
-        Threaded callback function to receive audio data when recordings finish.
-        audio: An AudioData containing the recorded bytes.
-        """
-        # Grab the raw bytes and push it into the thread safe queue.
-        data = audio.get_raw_data()
-        data_queue.put(data)
-
-    # Create a background thread that will pass us raw audio bytes.
-    # We could do this manually but SpeechRecognizer provides a nice helper.
-    recorder.listen_in_background(source, record_callback, phrase_time_limit=record_timeout)
-
     # Cue the user that we're ready to go.
     print("Model loaded.\n")
 
+    # Rolling output text buffer.
+
+    # This is the one that animates. Stored as a single string.
+    rolling_output_text = ""
+    # This is the one that updates in big chunks at lower frequency.
+    # Stored as an array of phrases.
+    output_text = [""]
+
+    mic_audio_source = MicrophoneAudioSource()
+    mic_audio_source.start()
+    data_queue = mic_audio_source.data_queue
+
+    # Rolling audio input buffer.
     audio_data = b''
 
+    diffsize = 0
+
     while True:
         try:
-            now = datetime.utcnow()
-            # Pull raw recorded audio from the queue.
-            if not data_queue.empty():
-                phrase_complete = False
-                # If enough time has passed between recordings, consider the phrase complete.
-                # Clear the current working audio buffer to start over with the new data.
-                if phrase_time and now - phrase_time > timedelta(seconds=phrase_timeout):
-                    phrase_complete = True
-                # This is the last time we received new audio data from the queue.
-                phrase_time = now
 
-                # for d in data_queue:
-                #     if d > 0.5:
-                #         print("Got something: ", d)
+            for event in pygame.event.get():
+                if event.type == pygame.QUIT:
+                    pygame.quit()
+                    exit(0)
 
-                # Combine audio data from queue
-                audio_data += b''.join(data_queue.queue)
-                data_queue.queue.clear()
+            rolling_text_target = " ".join(output_text)[-160:]
+            if rolling_text_target != rolling_output_text:
 
-                # Convert in-ram buffer to something the model can use directly without needing a temp file.
-                # Convert data from 16 bit wide integers to floating point with a width of 32 bits.
-                # Clamp the audio stream frequency to a PCM wavelength compatible default of 32768hz max.
-                audio_np = np.frombuffer(audio_data, dtype=np.int16).astype(np.float32) / 32768.0
+                # Chop off the start all at once. It's not needed for the animation to look good.
+                new_rolling_output_text = onestepchange(rolling_output_text, rolling_text_target)
+                while rolling_output_text.endswith(new_rolling_output_text):
+                    new_rolling_output_text = onestepchange(new_rolling_output_text, rolling_text_target)
+                rolling_output_text = new_rolling_output_text
 
-                # Read the transcription.
-                result = audio_model.transcribe(audio_np, fp16=torch.cuda.is_available())
-                text = result['text'].strip()
+                if countsteps(rolling_output_text, rolling_text_target) > 80:
+                    rolling_output_text = rolling_text_target
 
-                # # If we detected a pause between recordings, add a new item to our transcription.
-                # # Otherwise edit the existing one.
-                # if phrase_complete:
-                #     transcription.append(text)
-                # else:
-                #     transcription[-1] += text
-                print(text)
+                print(rolling_output_text)
 
-                # Update rolling transcription file.
-                f = open("transcription.txt", "w+")
-                output_text = transcription[-4:]
-                output_text.append(text)
-                f.write(" ".join(output_text))
-                f.close()
+                pygame_text_surface = pygame_font.render(rolling_output_text, (0, 0, 0), (255, 255, 255))
+                pygame_text_rect = pygame_text_surface.get_rect()
+                pygame_text_rect.center = (640, pygame_font_height)
+                pygame_text_rect.right = 1280
+                pygame_display_surface.fill((0, 0, 0))
+                pygame_display_surface.blit(pygame_text_surface, pygame_text_rect)
 
-                if phrase_complete:
+                pygame.display.update()
 
-                    # Append to full transcription.
-                    transcription.append(text)
+                diffsize = abs(len(rolling_output_text) - len(rolling_text_target))
 
-                    # text += "\n"
-                    # f = open("transcription.txt", "w+")
-                    # f.write("\n".join(textwrap.wrap(text)))
-                    # f.close()
-
-                    print("* Phrase complete.")
-                    audio_data = b''
-
-
-                # Clear the console to reprint the updated transcription.
-                # os.system('cls' if os.name=='nt' else 'clear')
-                for line in transcription:
-                    print(line)
-                # Flush stdout.
-                print('', end='', flush=True)
             else:
-                # Infinite loops are bad for processors, must sleep.
+
+                now = datetime.utcnow()
+                # Pull raw recorded audio from the queue.
+                if not data_queue.empty():
+
+                    phrase_complete = False
+                    # If enough time has passed between recordings, consider the phrase complete.
+                    # Clear the current working audio buffer to start over with the new data.
+                    #
+                    # FIXME: Shouldn't we cut off the phrase here instead of
+                    # waiting for later?
+                    if phrase_time and now - phrase_time > timedelta(seconds=phrase_timeout):
+                        phrase_complete = True
+
+                    # This is the last time we received new audio data from the queue.
+                    phrase_time = now
+
+                    # Combine audio data from queue
+                    audio_data += b''.join(data_queue.queue)
+                    data_queue.queue.clear()
+
+                    # Convert in-ram buffer to something the model can use directly without needing a temp file.
+                    # Convert data from 16 bit wide integers to floating point with a width of 32 bits.
+                    # Clamp the audio stream frequency to a PCM wavelength compatible default of 32768hz max.
+                    audio_np = np.frombuffer(audio_data, dtype=np.int16).astype(np.float32) / 32768.0
+
+                    # Run the transcription model, and extract the text.
+                    result = audio_model.transcribe(audio_np, fp16=torch.cuda.is_available())
+                    text = result['text'].strip()
+
+                    # Update rolling transcription file.
+                    
+                    # Start with all our recent-but-complete phrases.
+                    output_text = transcription[-recent_phrase_count:]
+
+                    # Append the phrase-in-progress. (TODO: Can we make this a
+                    # different color or something?)
+                    output_text.append(text)
+
+                    # If we're done with the phrase, we can go ahead and stuff
+                    # it into the list and clear out the current audio data
+                    # buffer.
+                    if phrase_complete:
+
+                        # Append to full transcription.
+                        if text != "":
+                            transcription.append(text)
+
+                        # Clear audio buffer.
+                        audio_data = b''
+
+            # Infinite loops are bad for processors, must sleep. Also, limit the anim speed.
+            if diffsize > 30:
                 sleep(0.01)
+            else:
+                sleep(0.05)
+
         except KeyboardInterrupt:
             break
 
-    print("\n\nTranscription:")
-    for line in transcription:
-        print(line)
-
-
 if __name__ == "__main__":
     main()

transcriber.py

@@ -0,0 +1,65 @@
+#!/usr/bin/python3
+
+import numpy as np
+import speech_recognition
+import whisper
+import torch
+import wave
+
+_audio_model = whisper.load_model("medium.en") # "large"
+
+# For debugging...
+# wave_out = wave.open("wave.wav", "wb")
+# wave_out.setnchannels(1)
+# wave_out.setframerate(16000)
+# wave_out.setsampwidth(2)
+
+class Transcriber:
+
+    def __init__(self):
+        self._audio_source = None
+
+        # Audio data for the current phrase.
+        self._current_data = b''
+
+    def set_source(self, source):
+        self._audio_source = source
+
+    def update(self):
+        
+        if self._audio_source:
+
+            if not self._audio_source.data_queue.empty():
+
+                # We got some new data. Let's process it!
+                new_data = []
+                while not self._audio_source.data_queue.empty():
+                    new_packet = self._audio_source.data_queue.get()
+                    new_data.append(new_packet)
+
+                new_data_joined = b''.join(new_data)
+                
+                # For debugging...
+                #wave_out.writeframes(new_data_joined)
+                
+                self._current_data = self._current_data + new_data_joined
+
+                # Convert in-ram buffer to something the model can use
+                # directly without needing a temp file. Convert data from 16
+                # bit wide integers to floating point with a width of 32
+                # bits. Clamp the audio stream frequency to a PCM wavelength
+                # compatible default of 32768hz max.
+                audio_np = np.frombuffer(
+                    self._current_data, dtype=np.int16).astype(np.float32) / 32768.0
+
+                # Run the transcription model, and extract the text.
+                result = _audio_model.transcribe(
+                    audio_np, fp16=torch.cuda.is_available())
+
+                text = result['text'].strip()
+
+                print("text now: ", text)
+
+            # Automatically drop audio sources when we're finished with them.
+            if self._audio_source.is_done():
+                self._audio_source = None