Shiven Mian
11 months ago
parent
a3de4c1286
commit
f749cb878e
@ -0,0 +1,52 @@
|
||||
from fastapi import FastAPI, WebSocket
|
||||
import uvicorn
|
||||
import json
|
||||
from stt import stt
|
||||
import tempfile
|
||||
|
||||
app = FastAPI()
|
||||
|
||||
@app.websocket("/user")
|
||||
async def user(ws: WebSocket):
|
||||
await ws.accept()
|
||||
audio_file = bytearray()
|
||||
mime_type = None
|
||||
|
||||
try:
|
||||
while True:
|
||||
message = await ws.receive()
|
||||
|
||||
if message['type'] == 'websocket.disconnect':
|
||||
break
|
||||
|
||||
if message['type'] == 'websocket.receive':
|
||||
if 'text' in message:
|
||||
control_message = json.loads(message['text'])
|
||||
if control_message.get('action') == 'command' and control_message.get('state') == 'start' and 'mimeType' in control_message:
|
||||
# This indicates the start of a new audio file
|
||||
mime_type = control_message.get('mimeType')
|
||||
elif control_message.get('action') == 'command' and control_message.get('state') == 'end':
|
||||
# This indicates the end of the audio file
|
||||
# Process the complete audio file here
|
||||
transcription = stt(audio_file, mime_type)
|
||||
await ws.send_json({"transcript": transcription})
|
||||
|
||||
print("SENT TRANSCRIPTION!")
|
||||
|
||||
# Reset the bytearray for the next audio file
|
||||
audio_file = bytearray()
|
||||
mime_type = None
|
||||
elif 'bytes' in message:
|
||||
# If it's not a control message, it's part of the audio file
|
||||
audio_file.extend(message['bytes'])
|
||||
|
||||
except Exception as e:
|
||||
print(f"WebSocket connection closed with exception: {e}")
|
||||
finally:
|
||||
await ws.close()
|
||||
print("WebSocket connection closed")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
with tempfile.TemporaryDirectory():
|
||||
uvicorn.run(app, host="0.0.0.0", port=8000)
|
||||
@ -0,0 +1,52 @@
|
||||
from datetime import datetime
|
||||
import os
|
||||
import contextlib
|
||||
import tempfile
|
||||
import ffmpeg
|
||||
import subprocess
|
||||
|
||||
from openai import OpenAI
|
||||
client = OpenAI()
|
||||
|
||||
def convert_mime_type_to_format(mime_type: str) -> str:
|
||||
if mime_type == "audio/x-wav" or mime_type == "audio/wav":
|
||||
return "wav"
|
||||
if mime_type == "audio/webm":
|
||||
return "webm"
|
||||
|
||||
return mime_type
|
||||
|
||||
@contextlib.contextmanager
|
||||
def export_audio_to_wav_ffmpeg(audio: bytearray, mime_type: str) -> str:
|
||||
temp_dir = tempfile.gettempdir()
|
||||
|
||||
# Create a temporary file with the appropriate extension
|
||||
input_ext = convert_mime_type_to_format(mime_type)
|
||||
input_path = os.path.join(temp_dir, f"input_{datetime.now().strftime('%Y%m%d%H%M%S%f')}.{input_ext}")
|
||||
with open(input_path, 'wb') as f:
|
||||
f.write(audio)
|
||||
|
||||
# Export to wav
|
||||
output_path = os.path.join(temp_dir, f"output_{datetime.now().strftime('%Y%m%d%H%M%S%f')}.wav")
|
||||
ffmpeg.input(input_path).output(output_path, acodec='pcm_s16le', ac=1, ar='16k').run()
|
||||
|
||||
print(f"Temporary file path: {output_path}")
|
||||
|
||||
try:
|
||||
yield output_path
|
||||
finally:
|
||||
os.remove(input_path)
|
||||
os.remove(output_path)
|
||||
|
||||
|
||||
def stt(audio_bytes: bytearray, mime_type):
|
||||
with export_audio_to_wav_ffmpeg(audio_bytes, mime_type) as wav_file_path:
|
||||
audio_file = open(wav_file_path, "rb")
|
||||
transcript = client.audio.transcriptions.create(
|
||||
model="whisper-1",
|
||||
file=audio_file,
|
||||
response_format="text"
|
||||
)
|
||||
|
||||
print("Exciting transcription result:", transcript)
|
||||
return transcript
|
||||
@ -0,0 +1,141 @@
|
||||
"""
|
||||
Handles everything the user interacts through.
|
||||
|
||||
Connects to a websocket at /user. Sends shit to it, and displays/plays the shit it sends back.
|
||||
|
||||
For now, just handles a spacebar being pressed— for the duration it's pressed,
|
||||
it should record audio.
|
||||
|
||||
SIMPLEST POSSIBLE: Sends that audio to OpenAI whisper, gets the transcript,
|
||||
sends it to /user in LMC format (role: user, etc)
|
||||
|
||||
MOST FUTUREPROOF: Streams chunks of audio to /user, which will then handle stt in stt.py.
|
||||
"""
|
||||
|
||||
import os
|
||||
import pyaudio
|
||||
import threading
|
||||
import asyncio
|
||||
import websockets
|
||||
import json
|
||||
from pynput import keyboard
|
||||
import wave
|
||||
import tempfile
|
||||
from datetime import datetime
|
||||
|
||||
# Configuration
|
||||
chunk = 1024 # Record in chunks of 1024 samples
|
||||
sample_format = pyaudio.paInt16 # 16 bits per sample
|
||||
channels = 1 # Stereo
|
||||
fs = 48000 # Sample rate
|
||||
|
||||
p = pyaudio.PyAudio() # Create an interface to PortAudio
|
||||
frames = [] # Initialize array to store frames
|
||||
recording = False # Flag to control recording state
|
||||
|
||||
ws_chunk_size = 4096 # Websocket stream chunk size
|
||||
|
||||
async def start_recording():
|
||||
global recording
|
||||
|
||||
if recording:
|
||||
return # Avoid multiple starts
|
||||
recording = True
|
||||
frames.clear() # Clear existing frames
|
||||
|
||||
stream = p.open(format=sample_format,
|
||||
channels=channels,
|
||||
rate=fs,
|
||||
frames_per_buffer=chunk,
|
||||
input=True)
|
||||
|
||||
print("Recording started...")
|
||||
async with websockets.connect("ws://localhost:8000/user") as websocket:
|
||||
# Send the start command with mime type
|
||||
await websocket.send(json.dumps({"action": "command", "state": "start", "mimeType": "audio/wav"}))
|
||||
while recording:
|
||||
data = stream.read(chunk)
|
||||
frames.append(data)
|
||||
|
||||
stream.stop_stream()
|
||||
stream.close()
|
||||
|
||||
try:
|
||||
file_path = save_recording(frames)
|
||||
with open(file_path, 'rb') as audio_file:
|
||||
byte_chunk = audio_file.read(ws_chunk_size)
|
||||
while byte_chunk:
|
||||
await websocket.send(byte_chunk)
|
||||
byte_chunk = audio_file.read(ws_chunk_size)
|
||||
finally:
|
||||
os.remove(file_path)
|
||||
|
||||
# Send the end command
|
||||
await websocket.send(json.dumps({"action": "command", "state": "end"}))
|
||||
|
||||
# Receive a json message and then close the connection
|
||||
message = await websocket.recv()
|
||||
print("Received message:", json.loads(message))
|
||||
|
||||
print("Recording stopped.")
|
||||
|
||||
def save_recording(frames) -> str:
|
||||
# Save the recorded data as a WAV file
|
||||
temp_dir = tempfile.gettempdir()
|
||||
|
||||
# Create a temporary file with the appropriate extension
|
||||
output_path = os.path.join(temp_dir, f"input_{datetime.now().strftime('%Y%m%d%H%M%S%f')}.wav")
|
||||
with wave.open(output_path, 'wb') as wf:
|
||||
wf.setnchannels(channels)
|
||||
wf.setsampwidth(p.get_sample_size(sample_format))
|
||||
wf.setframerate(fs)
|
||||
wf.writeframes(b''.join(frames))
|
||||
|
||||
return output_path
|
||||
|
||||
def start_recording_sync():
|
||||
# Create a new event loop for the thread
|
||||
loop = asyncio.new_event_loop()
|
||||
asyncio.set_event_loop(loop)
|
||||
# Run the asyncio event loop
|
||||
loop.run_until_complete(start_recording())
|
||||
loop.close()
|
||||
|
||||
def stop_recording():
|
||||
global recording
|
||||
recording = False
|
||||
print("Stopped recording")
|
||||
|
||||
def toggle_recording():
|
||||
global recording
|
||||
if recording:
|
||||
stop_recording()
|
||||
else:
|
||||
# Start recording in a new thread to avoid blocking
|
||||
print("Starting recording")
|
||||
threading.Thread(target=start_recording_sync).start()
|
||||
|
||||
is_space_pressed = False # Flag to track the state of the spacebar
|
||||
|
||||
def on_press(key):
|
||||
global is_space_pressed
|
||||
if key == keyboard.Key.space and not is_space_pressed:
|
||||
is_space_pressed = True
|
||||
toggle_recording()
|
||||
|
||||
def on_release(key):
|
||||
global is_space_pressed
|
||||
if key == keyboard.Key.space and is_space_pressed:
|
||||
is_space_pressed = False
|
||||
stop_recording()
|
||||
if key == keyboard.Key.esc:
|
||||
# Stop listener
|
||||
return False
|
||||
|
||||
# Collect events until released
|
||||
with keyboard.Listener(on_press=on_press, on_release=on_release) as listener:
|
||||
with tempfile.TemporaryDirectory():
|
||||
print("Press the spacebar to start/stop recording. Press ESC to exit.")
|
||||
listener.join()
|
||||
|
||||
p.terminate()
|
||||
@ -1,13 +0,0 @@
|
||||
"""
|
||||
Handles everything the user interacts through.
|
||||
|
||||
Connects to a websocket at /user. Sends shit to it, and displays/plays the shit it sends back.
|
||||
|
||||
For now, just handles a spacebar being pressed— for the duration it's pressed,
|
||||
it should record audio.
|
||||
|
||||
SIMPLEST POSSIBLE: Sends that audio to OpenAI whisper, gets the transcript,
|
||||
sends it to /user in LMC format (role: user, etc)
|
||||
|
||||
MOST FUTUREPROOF: Streams chunks of audio to /user, which will then handle stt in stt.py.
|
||||
"""
|
||||
Loading…
Reference in new issue