Stream Audio for Deepfake Detection

1
#!/usr/bin/env python3
2
3
from __future__ import annotations
4
5
import argparse
6
import asyncio
7
import json
8
import os
9
import shutil
10
import struct
11
from pathlib import Path
12
from typing import Any
13
14
import aiohttp
15
16
17
STREAM_URL = "wss://stream.resemble.ai/api/v1/detect/audio"
18
SAMPLE_RATE = 16_000
19
FRAME_MILLISECONDS = 100
20
21
22
def streaming_wav_header() -> bytes:
23
    """Create a mono PCM WAV header for a stream of unknown final length."""
24
    channels = 1
25
    bits_per_sample = 16
26
    bytes_per_sample = bits_per_sample // 8
27
    byte_rate = SAMPLE_RATE * channels * bytes_per_sample
28
    block_align = channels * bytes_per_sample
29
    return struct.pack(
30
        "<4sL4s4sLHHLLHH4sL",
31
        b"RIFF",
32
        0xFFFFFFFF,
33
        b"WAVE",
34
        b"fmt ",
35
        16,
36
        1,
37
        channels,
38
        SAMPLE_RATE,
39
        byte_rate,
40
        block_align,
41
        bits_per_sample,
42
        b"data",
43
        0xFFFFFFFF,
44
    )
45
46
47
def decode_message(message: aiohttp.WSMessage) -> dict[str, Any]:
48
    if message.type != aiohttp.WSMsgType.TEXT:
49
        raise RuntimeError(f"Expected a text message, received {message.type.name}")
50
    payload = json.loads(message.data.strip())
51
    if not isinstance(payload, dict):
52
        raise RuntimeError("The server returned an unexpected JSON value")
53
    return payload
54
55
56
def print_result(payload: dict[str, Any]) -> None:
57
    message_type = payload.get("type")
58
    if message_type == "ready":
59
        print(f"[ready] stream_id={payload.get('stream_id')}")
60
    elif message_type == "chunk":
61
        chunk = payload.get("chunk_info") or {}
62
        label = chunk.get("chunk_label")
63
        score = chunk.get("chunk_aggregated_score")
64
        score_text = "" if score is None or label == "skipped" else f" score={float(score):.4f}"
65
        print(
66
            f"[chunk {chunk.get('chunk_id')}] "
67
            f"{float(chunk.get('begin_timestamp_s', 0)):.1f}-"
68
            f"{float(chunk.get('end_timestamp_s', 0)):.1f}s "
69
            f"label={label}{score_text}"
70
        )
71
    elif message_type == "final":
72
        score = payload.get("aggregated_score")
73
        score_text = "N/A" if score is None else f"{float(score):.4f}"
74
        print(f"[final] label={payload.get('label')} aggregated_score={score_text}")
75
    elif message_type == "error":
76
        code = payload.get("error_code") or "stream_error"
77
        message = payload.get("error") or payload.get("error_message") or "Unknown error"
78
        raise RuntimeError(f"{code}: {message}")
79
80
81
async def send_audio(websocket: aiohttp.ClientWebSocketResponse, audio_path: Path) -> float:
82
    """Decode the input with ffmpeg and send paced binary PCM frames."""
83
    ffmpeg = shutil.which("ffmpeg")
84
    if not ffmpeg:
85
        raise RuntimeError("ffmpeg was not found on PATH")
86
87
    bytes_per_second = SAMPLE_RATE * 2  # mono signed 16-bit PCM
88
    frame_bytes = bytes_per_second * FRAME_MILLISECONDS // 1000
89
    process = await asyncio.create_subprocess_exec(
90
        ffmpeg,
91
        "-v",
92
        "error",
93
        "-i",
94
        str(audio_path),
95
        "-f",
96
        "s16le",
97
        "-acodec",
98
        "pcm_s16le",
99
        "-ar",
100
        str(SAMPLE_RATE),
101
        "-ac",
102
        "1",
103
        "pipe:1",
104
        stdout=asyncio.subprocess.PIPE,
105
        stderr=asyncio.subprocess.PIPE,
106
    )
107
    assert process.stdout is not None
108
    assert process.stderr is not None
109
110
    sent_bytes = 0
111
    started_at = asyncio.get_running_loop().time()
112
    await websocket.send_bytes(streaming_wav_header())
113
114
    try:
115
        while True:
116
            frame = await process.stdout.read(frame_bytes)
117
            if not frame:
118
                break
119
120
            await websocket.send_bytes(frame)
121
            sent_bytes += len(frame)
122
123
            target_elapsed = sent_bytes / bytes_per_second
124
            actual_elapsed = asyncio.get_running_loop().time() - started_at
125
            if target_elapsed > actual_elapsed:
126
                await asyncio.sleep(target_elapsed - actual_elapsed)
127
    except BaseException:
128
        if process.returncode is None:
129
            process.kill()
130
        await process.wait()
131
        raise
132
133
    try:
134
        await asyncio.wait_for(process.wait(), timeout=5)
135
    except asyncio.TimeoutError:
136
        process.kill()
137
        await process.wait()
138
139
    stderr = (await process.stderr.read()).decode(errors="replace").strip()
140
    if process.returncode != 0:
141
        raise RuntimeError(f"ffmpeg failed: {stderr or f'exit code {process.returncode}'}")
142
    if sent_bytes == 0:
143
        raise RuntimeError("ffmpeg decoded no audio")
144
145
    await websocket.send_str(json.dumps({"type": "end"}))
146
    duration = sent_bytes / bytes_per_second
147
    print(f"[sent] {duration:.2f}s of audio; end marker sent")
148
    return duration
149
150
151
async def receive_results(
152
    websocket: aiohttp.ClientWebSocketResponse,
153
) -> dict[str, Any] | None:
154
    final_result: dict[str, Any] | None = None
155
    async for message in websocket:
156
        if message.type == aiohttp.WSMsgType.ERROR:
157
            raise RuntimeError(f"WebSocket error: {websocket.exception()}")
158
        if message.type != aiohttp.WSMsgType.TEXT:
159
            continue
160
161
        payload = decode_message(message)
162
        print_result(payload)
163
        if payload.get("type") == "final":
164
            final_result = payload
165
    return final_result
166
167
168
async def stream_audio(audio_path: Path, api_key: str) -> None:
169
    headers = {"Authorization": f"Bearer {api_key}"}
170
    params = {"filename": audio_path.name}
171
    timeout = aiohttp.ClientTimeout(total=None, sock_connect=30)
172
    final_result: dict[str, Any] | None
173
174
    async with aiohttp.ClientSession(timeout=timeout) as session:
175
        try:
176
            async with session.ws_connect(
177
                STREAM_URL,
178
                headers=headers,
179
                params=params,
180
                heartbeat=30,
181
            ) as websocket:
182
                first_message = await asyncio.wait_for(websocket.receive(), timeout=30)
183
                first_payload = decode_message(first_message)
184
                print_result(first_payload)
185
                if first_payload.get("type") != "ready":
186
                    raise RuntimeError("The first server message was not ready")
187
188
                sender = asyncio.create_task(send_audio(websocket, audio_path))
189
                receiver = asyncio.create_task(receive_results(websocket))
190
                try:
191
                    await asyncio.gather(sender, receiver)
192
                finally:
193
                    for task in (sender, receiver):
194
                        if not task.done():
195
                            task.cancel()
196
                    await asyncio.gather(sender, receiver, return_exceptions=True)
197
198
                final_result = receiver.result()
199
200
                print(f"[closed] code={websocket.close_code}")
201
                if websocket.close_code != 1000:
202
                    raise RuntimeError(
203
                        f"The WebSocket closed unexpectedly with code {websocket.close_code}"
204
                    )
205
        except aiohttp.WSServerHandshakeError as exc:
206
            raise RuntimeError(
207
                f"WebSocket handshake rejected with HTTP {exc.status}: {exc.message}"
208
            ) from exc
209
210
    if final_result is None:
211
        raise RuntimeError("The server closed without returning a final result")
212
213
214
def main() -> None:
215
    parser = argparse.ArgumentParser(description="Stream audio for deepfake detection")
216
    parser.add_argument("audio_path", type=Path)
217
    args = parser.parse_args()
218
219
    audio_path = args.audio_path.expanduser().resolve()
220
    if not audio_path.is_file():
221
        parser.error(f"Audio file not found: {audio_path}")
222
223
    api_key = os.environ.get("RESEMBLE_API_KEY", "").strip()
224
    if not api_key:
225
        parser.error("Set RESEMBLE_API_KEY before running the client")
226
227
    asyncio.run(stream_audio(audio_path, api_key))
228
229
230
if __name__ == "__main__":
231
    main()