Working on audio enhancement

This commit is contained in:
Michael Hansen
2023-12-05 16:56:56 -06:00
parent 07bebb8f7a
commit cf60780414
5 changed files with 344 additions and 22 deletions
+244 -21
View File
@@ -2,24 +2,30 @@
import argparse
import asyncio
import logging
import shlex
import sys
import time
import uuid
from enum import Enum
from functools import partial
from pathlib import Path
from typing import Optional
from typing import Callable, Final, Optional
from wyoming.asr import Transcript, Transcribe
from wyoming.asr import Transcribe, Transcript
from wyoming.audio import AudioChunk, AudioStart, AudioStop
from wyoming.client import AsyncClient
from wyoming.event import Event
from wyoming.info import Attribution, Describe, Info, Satellite
from wyoming.mic import MicProcessAsyncClient
from wyoming.pipeline import PipelineStage, RunPipeline
from wyoming.satellite import RunSatellite
from wyoming.server import AsyncEventHandler, AsyncServer, AsyncTcpServer
from wyoming.snd import SndProcessAsyncClient
from wyoming.tts import Synthesize
from wyoming.vad import VoiceStarted, VoiceStopped
from wyoming.wake import Detect, Detection
from wyoming.wake import Detect, Detection, WakeProcessAsyncClient
from .utils import AudioBuffer, chunk_samples, multiply_volume
_LOGGER = logging.getLogger()
_DIR = Path(__file__).parent
@@ -33,11 +39,58 @@ class SatelliteState(str, Enum):
async def main() -> None:
"""Main entry point."""
parser = argparse.ArgumentParser()
# Microphone input
parser.add_argument("--mic-uri", help="URI of Wyoming microphone service")
parser.add_argument("--mic-command", help="Program to run for microphone input")
parser.add_argument(
"--mic-uri", required=True, help="URI of Wyoming microphone service"
"--mic-command-rate",
type=int,
default=16000,
help="Sample rate of mic-command (hertz, default: 16000)",
)
#
parser.add_argument(
"--mic-command-width",
type=int,
default=2,
help="Sample width of mic-command (bytes, default: 2)",
)
parser.add_argument(
"--mic-command-channels",
type=int,
default=1,
help="Sample channels of mic-command (default: 1)",
)
parser.add_argument(
"--mic-command-samples-per-chunk",
type=int,
default=1024,
help="Sample per chunk for mic-command (default: 1024)",
)
# Sound output
parser.add_argument("--snd-uri", help="URI of Wyoming sound service")
parser.add_argument("--snd-command", help="Program to run for sound output")
parser.add_argument(
"--snd-command-rate",
type=int,
default=22050,
help="Sample rate of snd-command (hertz, default: 22050)",
)
parser.add_argument(
"--snd-command-width",
type=int,
default=2,
help="Sample width of snd-command (bytes, default: 2)",
)
parser.add_argument(
"--snd-command-channels",
type=int,
default=1,
help="Sample channels of snd-command (default: 1)",
)
# Local wake word detection
parser.add_argument("--wake-uri", help="URI of Wyoming wake word detection service")
parser.add_argument(
"--wake-word-name",
@@ -45,16 +98,53 @@ async def main() -> None:
default=[],
help="Name of wake word to listen for (requires --wake-uri)",
)
#
parser.add_argument("--wake-command", help="Program to run for wake word detection")
parser.add_argument(
"--wake-command-rate",
type=int,
default=16000,
help="Sample rate of wake-command (hertz, default: 16000)",
)
parser.add_argument(
"--wake-command-width",
type=int,
default=2,
help="Sample width of wake-command (bytes, default: 2)",
)
parser.add_argument(
"--wake-command-channels",
type=int,
default=1,
help="Sample channels of wake-command (default: 1)",
)
# Voice activity detector
parser.add_argument(
"--vad", choices=("disabled", "webrtcvad", "silero"), default="disabled"
)
parser.add_argument("--vad-threshold", type=float, default=0.5)
parser.add_argument("--vad-trigger-level", type=int, default=3)
parser.add_argument("--vad-buffer-chunks", type=int, default=40)
# Audio enhancement (requires webrtc-noise-audio)
parser.add_argument(
"--noise-suppression", type=int, default=0, choices=(0, 1, 2, 3, 4)
)
parser.add_argument("--auto-gain", type=int, default=0, choices=list(range(32)))
parser.add_argument("--volume-multiplier", type=float, default=1.0)
# External event handlers
parser.add_argument(
"--event-uri", help="URI of Wyoming service to forward events to"
)
#
# Satellite details
parser.add_argument("--uri", default="stdio://", help="unix:// or tcp://")
parser.add_argument("--name", required=True, help="Name of the satellite")
parser.add_argument("--area", help="Area name of the satellite")
parser.add_argument("--pipeline", help="Name of server pipeline to run")
#
# Zeroconf
parser.add_argument(
"--no-zeroconf", action="store_true", help="Disable discovery over zeroconf"
)
@@ -66,6 +156,25 @@ async def main() -> None:
parser.add_argument("--debug", action="store_true", help="Log DEBUG messages")
args = parser.parse_args()
# Validate args
if (not args.mic_uri) and (not args.mic_command):
_LOGGER.fatal("Either --mic-uri or --mic-command is required")
sys.exit(1)
if needs_webrtc(args):
try:
import webrtc_noise_gain # noqa: F401
except ImportError:
_LOGGER.fatal("Install extras for webrtc")
sys.exit(1)
if needs_silero(args):
try:
import pysilero_vad # noqa: F401
except ImportError:
_LOGGER.fatal("Install extras for silero")
sys.exit(1)
logging.basicConfig(level=logging.DEBUG if args.debug else logging.INFO)
_LOGGER.debug(args)
@@ -125,11 +234,22 @@ class SatelliteEventHandler(AsyncEventHandler):
self.state: Optional[Satellite] = None
self.satellite_task: Optional[asyncio.Task] = None
self.has_snd = bool(self.cli_args.snd_uri)
self.volume_multiplier = self.cli_args.volume_multiplier
self.has_audio_processing = needs_webrtc(self.cli_args) or (
self.volume_multiplier != 1.0
)
self.has_wake = bool(self.cli_args.wake_uri or self.cli_args.wake_command)
self.has_snd = bool(self.cli_args.snd_uri or self.cli_args.snd_command)
self.snd_client: Optional[AsyncClient] = None
self.events_client: Optional[AsyncClient] = None
self.is_running = True
self.process_webrtc: Optional[Callable[[bytes], bytes]] = None
if needs_webrtc(self.cli_args):
self.process_webrtc = WebRtcAudio(
self.cli_args.auto_gain, self.cli_args.noise_suppression
)
_LOGGER.debug("Client connected: %s", self.client_id)
async def handle_event(self, event: Event) -> bool:
@@ -139,7 +259,7 @@ class SatelliteEventHandler(AsyncEventHandler):
return True
if RunSatellite.is_type(event.type) and (self.satellite_task is None):
if self.cli_args.wake_uri:
if self.has_wake:
# Local wake word detection
await self.write_event(
RunPipeline(
@@ -173,13 +293,27 @@ class SatelliteEventHandler(AsyncEventHandler):
# STT transcript
await self.forward_event(event)
_LOGGER.debug(event)
elif self.has_snd and (
AudioStart.is_type(event.type)
or AudioChunk.is_type(event.type)
or AudioStop.is_type(event.type)
):
if self.snd_client is None:
self.snd_client = AsyncClient.from_uri(self.cli_args.snd_uri)
if self.cli_args.snd_command:
_LOGGER.debug("Running %s", self.cli_args.snd_command)
program, *program_args = shlex.split(self.cli_args.snd_command)
self.snd_client = SndProcessAsyncClient(
rate=self.cli_args.snd_command_rate,
width=self.cli_args.snd_command_width,
channels=self.cli_args.snd_command_channels,
program=program,
program_args=program_args,
)
else:
_LOGGER.debug("Connecting to %s", self.cli_args.snd_uri)
self.snd_client = AsyncClient.from_uri(self.cli_args.snd_uri)
await self.snd_client.connect()
assert self.snd_client is not None
@@ -196,6 +330,7 @@ class SatelliteEventHandler(AsyncEventHandler):
if AudioStart.is_type(event.type) or AudioStop.is_type(event.type):
# TTS start/stop
await self.forward_event(event)
_LOGGER.debug(event)
elif (
Detect.is_type(event.type)
or Detection.is_type(event.type)
@@ -211,15 +346,33 @@ class SatelliteEventHandler(AsyncEventHandler):
# - VoiceStarted/VoiceStopped for when user starts/stops speaking
# - Synthesize for TTS text
await self.forward_event(event)
_LOGGER.debug(event)
else:
_LOGGER.debug("Unexpected event: type=%s, data=%s", event.type, event.data)
return True
def _make_mic_client(self) -> AsyncClient:
if self.cli_args.mic_command:
_LOGGER.debug("Running %s", self.cli_args.mic_command)
program, *program_args = shlex.split(self.cli_args.mic_command)
return MicProcessAsyncClient(
rate=self.cli_args.mic_command_rate,
width=self.cli_args.mic_command_width,
channels=self.cli_args.mic_command_channels,
samples_per_chunk=self.cli_args.mic_command_samples_per_chunk,
program=program,
program_args=program_args,
)
assert self.cli_args.mic_uri
_LOGGER.debug("Connecting to %s", self.cli_args.mic_uri)
return AsyncClient.from_uri(self.cli_args.mic_uri)
async def run_satellite(self) -> None:
try:
_LOGGER.debug("Connecting to %s", self.cli_args.mic_uri)
async with AsyncClient.from_uri(self.cli_args.mic_uri) as mic_client:
mic_client = self._make_mic_client()
async with mic_client:
_LOGGER.info("Streaming audio")
while True:
@@ -229,6 +382,8 @@ class SatelliteEventHandler(AsyncEventHandler):
break
if AudioChunk.is_type(mic_event.type):
mic_event = self._process_audio(mic_event)
# Forward all audio to server
await self.write_event(mic_event)
@@ -237,16 +392,25 @@ class SatelliteEventHandler(AsyncEventHandler):
async def run_satellite_wake(self) -> None:
try:
assert self.cli_args.wake_uri
assert self.cli_args.wake_uri or self.cli_args.wake_command
_LOGGER.debug("Connecting to %s", self.cli_args.mic_uri)
_LOGGER.debug("Connecting to %s", self.cli_args.wake_uri)
mic_client = self._make_mic_client()
async with AsyncClient.from_uri(
self.cli_args.mic_uri
) as mic_client, AsyncClient.from_uri(
self.cli_args.wake_uri
) as wake_client:
if self.cli_args.wake_command:
_LOGGER.debug("Running %s", self.cli_args.wake_command)
program, *program_args = shlex.split(self.cli_args.wake_command)
wake_client = WakeProcessAsyncClient(
rate=self.cli_args.wake_command_rate,
width=self.cli_args.wake_command_width,
channels=self.cli_args.wake_command_channels,
program=program,
program_args=program_args,
)
else:
_LOGGER.debug("Connecting to %s", self.cli_args.wake_uri)
wake_client = AsyncClient.from_uri(self.cli_args.wake_uri)
async with mic_client, wake_client:
_LOGGER.info("Detecting wake word")
if self.cli_args.wake_word_name:
@@ -257,6 +421,7 @@ class SatelliteEventHandler(AsyncEventHandler):
await wake_client.write_event(detect_event)
await self.forward_event(detect_event)
_LOGGER.debug(detect_event)
# Read events in parallel
mic_task = asyncio.create_task(mic_client.read_event())
@@ -275,6 +440,8 @@ class SatelliteEventHandler(AsyncEventHandler):
break
if AudioChunk.is_type(mic_event.type):
mic_event = self._process_audio(mic_event)
# Forward to wake service
await wake_client.write_event(mic_event)
@@ -311,6 +478,28 @@ class SatelliteEventHandler(AsyncEventHandler):
except Exception:
_LOGGER.exception("Unexpected error in run_satellite")
def _process_audio(self, audio_event: Event) -> Event:
"""Perform audio processing."""
if not self.has_audio_processing:
# No processing needed
return audio_event
chunk = AudioChunk.from_event(audio_event)
audio_bytes = chunk.audio
if self.process_webrtc is not None:
audio_bytes = self.process_webrtc(audio_bytes)
if self.volume_multiplier != 1.0:
audio_bytes = multiply_volume(audio_bytes, self.volume_multiplier)
return AudioChunk(
rate=chunk.rate,
width=chunk.width,
channels=chunk.channels,
audio=audio_bytes,
).event()
async def forward_event(self, event: Event) -> None:
"""Forward a Wyoming event to a client."""
if not self.cli_args.event_uri:
@@ -334,12 +523,46 @@ class SatelliteEventHandler(AsyncEventHandler):
def get_mac_address() -> str:
"""Return MAC address formatted as hex with no colons."""
return "".join(
# pylint: disable=consider-using-f-string
["{:02x}".format((uuid.getnode() >> ele) & 0xFF) for ele in range(0, 8 * 6, 8)][
::-1
]
)
def needs_webrtc(args: argparse.Namespace) -> bool:
"""Return True if webrtc must be used."""
return (
(args.noise_suppression > 0) or (args.auto_gain > 0) or (args.vad == "webrtc")
)
def needs_silero(args: argparse.Namespace) -> bool:
"""Return True if silero-vad must be used."""
return args.vad == "silero"
class WebRtcAudio:
_sub_chunk_samples: Final = 160
_sub_chunk_bytes: Final = _sub_chunk_samples * 2 # 16-bit
def __init__(self, auto_gain: int, noise_suppression: int) -> None:
from webrtc_noise_gain import AudioProcessor
self.audio_processor = AudioProcessor(auto_gain, noise_suppression)
self.audio_buffer = AudioBuffer(self._sub_chunk_bytes)
def __call__(self, audio_bytes: bytes) -> bytes:
clean_chunk = bytes()
for sub_chunk in chunk_samples(
audio_bytes, self._sub_chunk_bytes, self.audio_buffer
):
result = self.audio_processor.Process10ms(sub_chunk)
clean_chunk += result.audio
return clean_chunk
# -----------------------------------------------------------------------------
if __name__ == "__main__":