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
            words = data.get('words', [])
92
93
            if transcript or words:
94
                if data.get('end_of_turn'):
95
                    self.print_final_transcript(transcript)
96
                else:
97
                    self.print_partial_transcript(words)
98
99
    def start_transcription(self):
100
        self.load_audio_channel()
101
102
        self.ws_app = websocket.WebSocketApp(
103
            API_ENDPOINT,
104
            header={"Authorization": YOUR_API_KEY},
105
            on_open=self.on_open,
106
            on_message=self.on_message,
107
        )
108
109
        thread = threading.Thread(target=self.ws_app.run_forever, daemon=True)
110
        thread.start()
111
        return thread
112
113
def play_audio_file():
114
    try:
115
        with wave.open(AUDIO_FILE_PATH, 'rb') as wf:
116
            p = pyaudio.PyAudio()
117
118
            stream = p.open(
119
                format=p.get_format_from_width(wf.getsampwidth()),
120
                channels=wf.getnchannels(),
121
                rate=wf.getframerate(),
122
                output=True
123
            )
124
125
            print(f"Playing audio: {AUDIO_FILE_PATH}")
126
127
            # Play audio in chunks
128
            chunk_size = 1024
129
            data = wf.readframes(chunk_size)
130
131
            while data:
132
                stream.write(data)
133
                data = wf.readframes(chunk_size)
134
135
            stream.stop_stream()
136
            stream.close()
137
            p.terminate()
138
139
            print("Audio playback finished")
140
141
    except Exception as e:
142
        print(f"Error playing audio: {e}")
143
144
145
def transcribe_multichannel():
146
    # Create transcribers for each channel
147
    transcriber_1 = ChannelTranscriber(0, "Speaker 1")
148
    transcriber_2 = ChannelTranscriber(1, "Speaker 2")
149
150
    # Start audio playback
151
    audio_thread = threading.Thread(target=play_audio_file, daemon=True)
152
    audio_thread.start()
153
154
    # Start both transcriptions
155
    thread_1 = transcriber_1.start_transcription()
156
    thread_2 = transcriber_2.start_transcription()
157
158
    # Wait for completion
159
    thread_1.join()
160
    thread_2.join()
161
    audio_thread.join()
162
163
if __name__ == "__main__":
164
    transcribe_multichannel()