444 lines
17 KiB
Python
444 lines
17 KiB
Python
import sys
|
||
import json
|
||
import wave
|
||
import time
|
||
import pyttsx3
|
||
import torch
|
||
import requests
|
||
import soundfile
|
||
import yaml
|
||
import pygame
|
||
import pygame.locals
|
||
import numpy as np
|
||
import pyaudio
|
||
import whisper
|
||
import logging
|
||
import threading
|
||
import queue
|
||
import asyncio
|
||
import edge_tts
|
||
import os
|
||
|
||
# Configure logging
|
||
logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s')
|
||
|
||
BACK_COLOR = (0,0,0)
|
||
REC_COLOR = (255,0,0)
|
||
TEXT_COLOR = (255,255,255)
|
||
REC_SIZE = 80
|
||
FONT_SIZE = 24
|
||
WIDTH = 320
|
||
HEIGHT = 240
|
||
KWIDTH = 20
|
||
KHEIGHT = 6
|
||
MAX_TEXT_LEN_DISPLAY = 32
|
||
|
||
INPUT_DEFAULT_DURATION_SECONDS = 5
|
||
INPUT_FORMAT = pyaudio.paInt16
|
||
INPUT_CHANNELS = 1
|
||
INPUT_RATE = 16000
|
||
INPUT_CHUNK = 1024
|
||
OLLAMA_REST_HEADERS = {'Content-Type': 'application/json'}
|
||
INPUT_CONFIG_PATH ="assistant.yaml"
|
||
|
||
class Assistant:
|
||
def __init__(self):
|
||
logging.info("Initializing Assistant")
|
||
self.config = self.init_config()
|
||
|
||
# 预初始化 pygame mixer，避免每次语音播放时的初始化开销
|
||
pygame.mixer.init()
|
||
|
||
programIcon = pygame.image.load('assistant.png')
|
||
|
||
self.clock = pygame.time.Clock()
|
||
pygame.display.set_icon(programIcon)
|
||
pygame.display.set_caption("Assistant")
|
||
|
||
self.windowSurface = pygame.display.set_mode((WIDTH, HEIGHT), 0, 32)
|
||
self.font = pygame.font.SysFont(None, FONT_SIZE)
|
||
|
||
self.audio = pyaudio.PyAudio()
|
||
|
||
# Initialize TTS engines
|
||
self.tts_engine = pyttsx3.init()
|
||
self.tts_engine.setProperty('rate', self.tts_engine.getProperty('rate') - 50)
|
||
|
||
# Set default voice for edge-tts
|
||
self.edge_voice = self.config.tts.edge_voice
|
||
|
||
try:
|
||
self.audio.open(format=INPUT_FORMAT,
|
||
channels=INPUT_CHANNELS,
|
||
rate=INPUT_RATE,
|
||
input=True,
|
||
frames_per_buffer=INPUT_CHUNK).close()
|
||
except Exception as e:
|
||
logging.error(f"Error opening audio stream: {str(e)}")
|
||
self.wait_exit()
|
||
|
||
self.display_message(self.config.messages.loadingModel)
|
||
self.model = whisper.load_model(self.config.whisperRecognition.modelPath)
|
||
self.context = []
|
||
|
||
self.text_to_speech(self.config.conversation.greeting)
|
||
time.sleep(0.5)
|
||
self.display_message(self.config.messages.pressSpace)
|
||
|
||
def wait_exit(self):
|
||
while True:
|
||
self.display_message(self.config.messages.noAudioInput)
|
||
self.clock.tick(60)
|
||
for event in pygame.event.get():
|
||
if event.type == pygame.locals.QUIT:
|
||
self.shutdown()
|
||
|
||
def shutdown(self):
|
||
logging.info("Shutting down Assistant")
|
||
self.audio.terminate()
|
||
pygame.quit()
|
||
sys.exit()
|
||
|
||
def init_config(self):
|
||
logging.info("Initializing configuration")
|
||
class Inst:
|
||
pass
|
||
|
||
with open('assistant.yaml', encoding='utf-8') as data:
|
||
configYaml = yaml.safe_load(data)
|
||
|
||
config = Inst()
|
||
config.messages = Inst()
|
||
config.messages.loadingModel = configYaml["messages"]["loadingModel"]
|
||
config.messages.pressSpace = configYaml["messages"]["pressSpace"]
|
||
config.messages.noAudioInput = configYaml["messages"]["noAudioInput"]
|
||
|
||
config.conversation = Inst()
|
||
config.conversation.greeting = configYaml["conversation"]["greeting"]
|
||
|
||
config.ollama = Inst()
|
||
config.ollama.url = configYaml["ollama"]["url"]
|
||
config.ollama.model = configYaml["ollama"]["model"]
|
||
|
||
config.whisperRecognition = Inst()
|
||
config.whisperRecognition.modelPath = configYaml["whisperRecognition"]["modelPath"]
|
||
config.whisperRecognition.lang = configYaml["whisperRecognition"]["lang"]
|
||
|
||
config.tts = Inst()
|
||
config.tts.engine = configYaml["tts"]["engine"] # 'edge-tts' or 'pyttsx3'
|
||
config.tts.edge_voice = configYaml["tts"]["edge_voice"]
|
||
|
||
return config
|
||
|
||
def display_rec_start(self):
|
||
logging.info("Displaying recording start")
|
||
self.windowSurface.fill(BACK_COLOR)
|
||
pygame.draw.circle(self.windowSurface, REC_COLOR, (WIDTH/2, HEIGHT/2), REC_SIZE)
|
||
pygame.display.flip()
|
||
|
||
def display_sound_energy(self, energy):
|
||
logging.info(f"Displaying sound energy: {energy}")
|
||
COL_COUNT = 5
|
||
RED_CENTER = 100
|
||
FACTOR = 10
|
||
MAX_AMPLITUDE = 100
|
||
|
||
self.windowSurface.fill(BACK_COLOR)
|
||
amplitude = int(MAX_AMPLITUDE*energy)
|
||
hspace, vspace = 2*KWIDTH, int(KHEIGHT/2)
|
||
def rect_coords(x, y):
|
||
return (int(x-KWIDTH/2), int(y-KHEIGHT/2),
|
||
KWIDTH, KHEIGHT)
|
||
for i in range(-int(np.floor(COL_COUNT/2)), int(np.ceil(COL_COUNT/2))):
|
||
x, y, count = WIDTH/2+(i*hspace), HEIGHT/2, amplitude-2*abs(i)
|
||
|
||
mid = int(np.ceil(count/2))
|
||
for i in range(0, mid):
|
||
offset = i*(KHEIGHT+vspace)
|
||
pygame.draw.rect(self.windowSurface, RED_CENTER,
|
||
rect_coords(x, y+offset))
|
||
#mirror:
|
||
pygame.draw.rect(self.windowSurface, RED_CENTER,
|
||
rect_coords(x, y-offset))
|
||
pygame.display.flip()
|
||
|
||
def display_message(self, text):
|
||
logging.info(f"Displaying message: {text}")
|
||
self.windowSurface.fill(BACK_COLOR)
|
||
|
||
label = self.font.render(text
|
||
if (len(text)<MAX_TEXT_LEN_DISPLAY)
|
||
else (text[0:MAX_TEXT_LEN_DISPLAY]+"..."),
|
||
1,
|
||
TEXT_COLOR)
|
||
|
||
size = label.get_rect()[2:4]
|
||
self.windowSurface.blit(label, (WIDTH/2 - size[0]/2, HEIGHT/2 - size[1]/2))
|
||
|
||
pygame.display.flip()
|
||
|
||
def waveform_from_mic(self, key = pygame.K_SPACE) -> np.ndarray:
|
||
logging.info("Capturing waveform from microphone")
|
||
self.display_rec_start()
|
||
|
||
stream = self.audio.open(format=INPUT_FORMAT,
|
||
channels=INPUT_CHANNELS,
|
||
rate=INPUT_RATE,
|
||
input=True,
|
||
frames_per_buffer=INPUT_CHUNK)
|
||
frames = []
|
||
|
||
while True:
|
||
pygame.event.pump() # process event queue
|
||
pressed = pygame.key.get_pressed()
|
||
if pressed[key]:
|
||
data = stream.read(INPUT_CHUNK)
|
||
frames.append(data)
|
||
else:
|
||
break
|
||
|
||
stream.stop_stream()
|
||
stream.close()
|
||
|
||
return np.frombuffer(b''.join(frames), np.int16).astype(np.float32) * (1 / 32768.0)
|
||
|
||
def speech_to_text(self, waveform):
|
||
logging.info("Converting speech to text")
|
||
result_queue = queue.Queue()
|
||
|
||
def transcribe_speech():
|
||
try:
|
||
logging.info("Starting transcription")
|
||
transcript = self.model.transcribe(waveform,
|
||
language=self.config.whisperRecognition.lang,
|
||
fp16=torch.cuda.is_available())
|
||
logging.info("Transcription completed")
|
||
text = transcript["text"]
|
||
print('\nMe:\n', text.strip())
|
||
result_queue.put(text)
|
||
except Exception as e:
|
||
logging.error(f"An error occurred during transcription: {str(e)}")
|
||
result_queue.put("")
|
||
|
||
transcription_thread = threading.Thread(target=transcribe_speech)
|
||
transcription_thread.start()
|
||
transcription_thread.join()
|
||
|
||
return result_queue.get()
|
||
|
||
def ask_ollama(self, prompt, responseCallback):
|
||
logging.info(f"Asking OLLaMa with prompt: {prompt}")
|
||
full_prompt = prompt if hasattr(self, "contextSent") else (prompt)
|
||
self.contextSent = True
|
||
jsonParam = {
|
||
"model": self.config.ollama.model,
|
||
"stream": True,
|
||
"context": self.context,
|
||
"prompt": full_prompt
|
||
}
|
||
|
||
try:
|
||
response = requests.post(self.config.ollama.url,
|
||
json=jsonParam,
|
||
headers=OLLAMA_REST_HEADERS,
|
||
stream=True,
|
||
timeout=30)
|
||
response.raise_for_status()
|
||
|
||
full_response = ""
|
||
for line in response.iter_lines():
|
||
if not line:
|
||
continue
|
||
|
||
try:
|
||
body = json.loads(line)
|
||
token = body.get('response', '')
|
||
full_response += token
|
||
|
||
if 'error' in body:
|
||
logging.error(f"Error from OLLaMa: {body['error']}")
|
||
responseCallback("Error: " + body['error'])
|
||
return
|
||
|
||
if body.get('done', False) and 'context' in body:
|
||
self.context = body['context']
|
||
break
|
||
except json.JSONDecodeError as e:
|
||
logging.error(f"Failed to decode JSON response: {str(e)}")
|
||
continue
|
||
|
||
if full_response.strip():
|
||
try:
|
||
responseCallback(full_response.strip())
|
||
except Exception as e:
|
||
logging.error(f"Error in response callback: {str(e)}")
|
||
self.display_message("Error processing response")
|
||
else:
|
||
logging.warning("Received empty response from OLLaMa")
|
||
self.display_message("Received empty response")
|
||
|
||
except requests.exceptions.ReadTimeout as e:
|
||
logging.error(f"ReadTimeout occurred while asking OLLaMa: {str(e)}")
|
||
self.display_message("Request timed out. Please try again.")
|
||
except requests.exceptions.RequestException as e:
|
||
logging.error(f"An error occurred while asking OLLaMa: {str(e)}")
|
||
self.display_message("Connection error. Please try again.")
|
||
except Exception as e:
|
||
logging.error(f"Unexpected error in ask_ollama: {str(e)}")
|
||
self.display_message("An unexpected error occurred")
|
||
|
||
async def edge_tts_speak(self, text):
|
||
try:
|
||
logging.info(f"Using edge-tts with voice: {self.edge_voice}")
|
||
communicate = edge_tts.Communicate(text, self.edge_voice)
|
||
|
||
# 使用异步方式并行处理音频生成
|
||
audio_task = asyncio.create_task(communicate.save("temp_speech.mp3"))
|
||
|
||
# 在音频生成的同时执行其他初始化
|
||
if not pygame.mixer.get_init():
|
||
pygame.mixer.init()
|
||
|
||
# 等待音频生成完成
|
||
await audio_task
|
||
|
||
if not os.path.exists("temp_speech.mp3") or os.path.getsize("temp_speech.mp3") == 0:
|
||
raise Exception("Generated audio file is empty or does not exist")
|
||
|
||
pygame.mixer.music.load("temp_speech.mp3")
|
||
pygame.mixer.music.play()
|
||
|
||
while pygame.mixer.music.get_busy():
|
||
await asyncio.sleep(0.1) # 使用异步等待替代 pygame.time.wait
|
||
|
||
# 不要每次都退出 mixer，只清理文件
|
||
if os.path.exists("temp_speech.mp3"):
|
||
os.remove("temp_speech.mp3")
|
||
|
||
except Exception as e:
|
||
logging.error(f"An error occurred during edge-tts speech playback: {str(e)}")
|
||
logging.error(f"Voice being used: {self.edge_voice}")
|
||
logging.info("Falling back to pyttsx3...")
|
||
self.tts_engine.say(text)
|
||
self.tts_engine.runAndWait()
|
||
|
||
def text_to_speech(self, text):
|
||
logging.info(f"Converting text to speech: {text}")
|
||
print('\nAI:\n', text.strip())
|
||
|
||
def play_speech():
|
||
try:
|
||
if self.config.tts.engine == "edge-tts":
|
||
# Create temp file for visualization
|
||
tempPath = './temp.wav'
|
||
|
||
async def process_speech():
|
||
communicate = edge_tts.Communicate(text, self.edge_voice)
|
||
await communicate.save("temp_speech.mp3")
|
||
|
||
# Convert mp3 to wav for visualization
|
||
data, samplerate = soundfile.read("temp_speech.mp3")
|
||
soundfile.write(tempPath, data, samplerate)
|
||
|
||
# Play audio with visualization
|
||
wf = wave.open(tempPath, 'rb')
|
||
stream = self.audio.open(format=self.audio.get_format_from_width(wf.getsampwidth()),
|
||
channels=wf.getnchannels(),
|
||
rate=wf.getframerate(),
|
||
output=True)
|
||
|
||
chunkSize = 1024
|
||
chunk = wf.readframes(chunkSize)
|
||
while chunk:
|
||
stream.write(chunk)
|
||
tmp = np.array(np.frombuffer(chunk, np.int16), np.float32) * (1 / 32768.0)
|
||
energy_of_chunk = np.sqrt(np.mean(tmp**2))
|
||
self.display_sound_energy(energy_of_chunk)
|
||
chunk = wf.readframes(chunkSize)
|
||
|
||
wf.close()
|
||
stream.stop_stream()
|
||
stream.close()
|
||
|
||
# Cleanup temp files
|
||
if os.path.exists("temp_speech.mp3"):
|
||
os.remove("temp_speech.mp3")
|
||
if os.path.exists(tempPath):
|
||
os.remove(tempPath)
|
||
|
||
asyncio.run(process_speech())
|
||
|
||
else: # pyttsx3
|
||
tempPath = './temp.wav'
|
||
self.tts_engine.save_to_file(text, tempPath)
|
||
self.tts_engine.runAndWait()
|
||
|
||
# Play audio with visualization
|
||
data, samplerate = soundfile.read(tempPath)
|
||
soundfile.write(tempPath, data, samplerate)
|
||
|
||
wf = wave.open(tempPath, 'rb')
|
||
stream = self.audio.open(format=self.audio.get_format_from_width(wf.getsampwidth()),
|
||
channels=wf.getnchannels(),
|
||
rate=wf.getframerate(),
|
||
output=True)
|
||
|
||
chunkSize = 1024
|
||
chunk = wf.readframes(chunkSize)
|
||
while chunk:
|
||
stream.write(chunk)
|
||
tmp = np.array(np.frombuffer(chunk, np.int16), np.float32) * (1 / 32768.0)
|
||
energy_of_chunk = np.sqrt(np.mean(tmp**2))
|
||
self.display_sound_energy(energy_of_chunk)
|
||
chunk = wf.readframes(chunkSize)
|
||
|
||
wf.close()
|
||
stream.stop_stream()
|
||
stream.close()
|
||
|
||
if os.path.exists(tempPath):
|
||
os.remove(tempPath)
|
||
|
||
logging.info("Speech playback completed")
|
||
# self.display_message(text)
|
||
|
||
except Exception as e:
|
||
logging.error(f"An error occurred during speech playback: {str(e)}")
|
||
|
||
# Use daemon thread so main program can exit
|
||
speech_thread = threading.Thread(target=play_speech, daemon=True)
|
||
speech_thread.start()
|
||
|
||
|
||
def main():
|
||
logging.info("Starting Assistant")
|
||
pygame.init()
|
||
|
||
ass = Assistant()
|
||
|
||
push_to_talk_key = pygame.K_SPACE
|
||
quit_key = pygame.K_ESCAPE
|
||
|
||
while True:
|
||
ass.clock.tick(60)
|
||
for event in pygame.event.get():
|
||
if event.type == pygame.KEYDOWN:
|
||
if event.key == push_to_talk_key:
|
||
logging.info("Push-to-talk key pressed")
|
||
speech = ass.waveform_from_mic(push_to_talk_key)
|
||
|
||
transcription = ass.speech_to_text(waveform=speech)
|
||
|
||
ass.ask_ollama(transcription, ass.text_to_speech)
|
||
|
||
time.sleep(1)
|
||
ass.display_message(ass.config.messages.pressSpace)
|
||
|
||
elif event.key == quit_key:
|
||
logging.info("Quit key pressed")
|
||
ass.shutdown()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|