Streaming Diarization and Multichannel

1
import websocket
2
import json
3
import threading
4
import numpy as np
5
import wave
6
import time
7
import pyaudio
8
from urllib.parse import urlencode
9
10
# Configuration
11
YOUR_API_KEY = "<YOUR_API_KEY>"
12
AUDIO_FILE_PATH = "<DUAL_CHANNEL_AUDIO_FILE_PATH>"
13
API_BASE_URL = "wss://streaming.assemblyai.com/v3/ws"
14
API_PARAMS = {
15
    "sample_rate": 8000,
16
    "format_turns": "true",
17
    "end_of_turn_confidence_threshold": 0.4,
18
    "min_turn_silence": 160,
19
    "max_turn_silence": 400,
20
}
21
# Build API endpoint with URL encoding
22
API_ENDPOINT = f"{API_BASE_URL}?{urlencode(API_PARAMS)}"
23
24
class ChannelTranscriber:
25
    def __init__(self, channel_id, channel_name):
26
        self.channel_id = channel_id
27
        self.channel_name = channel_name
28
        self.ws_app = None
29
        self.audio_data = []
30
        self.current_turn_line = None
31
        self.line_count = 0
32
33
    def load_audio_channel(self):
34
        """Extract single channel from dual-channel audio file."""
35
        with wave.open(AUDIO_FILE_PATH, 'rb') as wf:
36
            frames = wf.readframes(wf.getnframes())
37
            audio_array = np.frombuffer(frames, dtype=np.int16)
38
39
            if wf.getnchannels() == 2:
40
                audio_array = audio_array.reshape(-1, 2)
41
                channel_audio = audio_array[:, self.channel_id]
42
43
                # Split into chunks for streaming
44
                FRAMES_PER_BUFFER = 400  # 50ms chunks
45
                for i in range(0, len(channel_audio), FRAMES_PER_BUFFER):
46
                    chunk = channel_audio[i:i+FRAMES_PER_BUFFER]
47
                    if len(chunk) < FRAMES_PER_BUFFER:
48
                        chunk = np.pad(chunk, (0, FRAMES_PER_BUFFER - len(chunk)), 'constant')
49
                    self.audio_data.append(chunk.astype(np.int16).tobytes())
50
51
    def on_open(self, ws):
52
        """Stream audio data when connection opens."""
53
        def stream_audio():
54
            for chunk in self.audio_data:
55
                ws.send(chunk, websocket.ABNF.OPCODE_BINARY)
56
                time.sleep(0.05)  # 50ms intervals
57
58
            # Send termination message
59
            terminate_message = {"type": "Terminate"}
60
            ws.send(json.dumps(terminate_message))
61
62
        threading.Thread(target=stream_audio, daemon=True).start()
63
64
    def clear_current_line(self):
65
        if self.current_turn_line is not None:
66
            print("\r" + " " * 100 + "\r", end="", flush=True)
67
68
    def print_partial_transcript(self, words):
69
        self.clear_current_line()
70
        # Build transcript from individual words
71
        word_texts = [word.get('text', '') for word in words]
72
        transcript = ' '.join(word_texts)
73
        partial_text = f"{self.channel_name}: {transcript}"
74
        print(partial_text, end="", flush=True)
75
        self.current_turn_line = len(partial_text)
76
77
    def print_final_transcript(self, transcript):
78
        self.clear_current_line()
79
        final_text = f"{self.channel_name}: {transcript}"
80
        print(final_text, flush=True)
81
        self.current_turn_line = None
82
        self.line_count += 1
83
84
    def on_message(self, ws, message):
85
        """Handle transcription results."""
86
        data = json.loads(message)
87
        msg_type = data.get('type')
88
89
        if msg_type == "Turn":
90
            transcript = data.get('transcript', '').strip()
91
            formatted = data.get('turn_is_formatted', False)
92
            words = data.get('words', [])
93
94
            if transcript or words:
95
                if formatted:
96
                    self.print_final_transcript(transcript)
97
                else:
98
                    self.print_partial_transcript(words)
99
100
    def start_transcription(self):
101
        self.load_audio_channel()
102
103
        self.ws_app = websocket.WebSocketApp(
104
            API_ENDPOINT,
105
            header={"Authorization": YOUR_API_KEY},
106
            on_open=self.on_open,
107
            on_message=self.on_message,
108
        )
109
110
        thread = threading.Thread(target=self.ws_app.run_forever, daemon=True)
111
        thread.start()
112
        return thread
113
114
def play_audio_file():
115
    try:
116
        with wave.open(AUDIO_FILE_PATH, 'rb') as wf:
117
            p = pyaudio.PyAudio()
118
119
            stream = p.open(
120
                format=p.get_format_from_width(wf.getsampwidth()),
121
                channels=wf.getnchannels(),
122
                rate=wf.getframerate(),
123
                output=True
124
            )
125
126
            print(f"Playing audio: {AUDIO_FILE_PATH}")
127
128
            # Play audio in chunks
129
            chunk_size = 1024
130
            data = wf.readframes(chunk_size)
131
132
            while data:
133
                stream.write(data)
134
                data = wf.readframes(chunk_size)
135
136
            stream.stop_stream()
137
            stream.close()
138
            p.terminate()
139
140
            print("Audio playback finished")
141
142
    except Exception as e:
143
        print(f"Error playing audio: {e}")
144
145
146
def transcribe_multichannel():
147
    # Create transcribers for each channel
148
    transcriber_1 = ChannelTranscriber(0, "Speaker 1")
149
    transcriber_2 = ChannelTranscriber(1, "Speaker 2")
150
151
    # Start audio playback
152
    audio_thread = threading.Thread(target=play_audio_file, daemon=True)
153
    audio_thread.start()
154
155
    # Start both transcriptions
156
    thread_1 = transcriber_1.start_transcription()
157
    thread_2 = transcriber_2.start_transcription()
158
159
    # Wait for completion
160
    thread_1.join()
161
    thread_2.join()
162
    audio_thread.join()
163
164
if __name__ == "__main__":
165
    transcribe_multichannel()