whisper_api.py 2.2 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061
  1. from fastapi import FastAPI, WebSocket
  2. import whisper
  3. import numpy as np
  4. import scipy.io.wavfile as wavfile
  5. import base64
  6. import logging
  7. import os
  8. app = FastAPI()
  9. model = whisper.load_model("large") # 정확도 향상을 위해 medium 사용
  10. logging.basicConfig(level=logging.INFO)
  11. @app.websocket("/audio-stream")
  12. async def transcribe_audio(websocket: WebSocket):
  13. await websocket.accept()
  14. buffer = bytearray() # 전체 오디오 데이터 저장
  15. sample_rate = 16000
  16. try:
  17. while True:
  18. data = await websocket.receive_text()
  19. if data == "STOP": # 종료 신호 수신
  20. break
  21. audio_data = base64.b64decode(data)
  22. buffer.extend(audio_data)
  23. logging.info(f"Received audio chunk of length: {len(audio_data)}")
  24. # 녹음 종료 후 파일 저장 및 처리
  25. if len(buffer) > 0:
  26. audio_np = np.frombuffer(buffer, dtype=np.int16)
  27. output_file = "recorded_audio.wav"
  28. wavfile.write(output_file, sample_rate, audio_np)
  29. logging.info(f"Saved audio file: {output_file}, size: {os.path.getsize(output_file)} bytes")
  30. # STT: 한국어 음성 인식
  31. stt_result = model.transcribe(output_file, language="ko")
  32. transcription = stt_result["text"]
  33. language = stt_result["language"]
  34. if transcription.strip():
  35. logging.info(f"Transcription: {transcription} [lang: {language}]")
  36. await websocket.send_json({
  37. "text": transcription,
  38. "language": language,
  39. # "translated_text": translated_text,
  40. "audio_url": "/audio"
  41. })
  42. else:
  43. await websocket.send_json({"text": "No transcription available", "language": "unknown"})
  44. else:
  45. await websocket.send_json({"text": "No audio data received", "language": "unknown"})
  46. except Exception as e:
  47. logging.error(f"Error: {e}")
  48. await websocket.send_json({"text": f"Error: {str(e)}", "language": "unknown"})
  49. finally:
  50. await websocket.close()
  51. @app.get("/health")
  52. async def health_check():
  53. return {"status": "healthy"}