98 lines
3.3 KiB
Python
98 lines
3.3 KiB
Python
import audioop
|
|
import math
|
|
import pyaudio
|
|
import asyncio
|
|
from internal.logger import create_logger
|
|
|
|
# You need to import the base AudioSource class from your specific library.
|
|
# This is a common path, but yours might be different.
|
|
from discord import AudioSource
|
|
|
|
LOGGER = create_logger(__name__)
|
|
|
|
# Constants for audio processing
|
|
SAMPLES_PER_FRAME = 960
|
|
CHANNELS = 2
|
|
SAMPLE_RATE = 48000
|
|
FRAME_SIZE = SAMPLES_PER_FRAME * CHANNELS * 2 # 16-bit PCM
|
|
SILENT_FRAME = b'\x00' * FRAME_SIZE
|
|
|
|
|
|
class NoiseGateSource(AudioSource):
|
|
def __init__(self, audio_stream, threshold: int):
|
|
self.audio_stream = audio_stream
|
|
self.threshold = threshold
|
|
self.ng_fadeout_count = 0
|
|
self.NG_FADEOUT_FRAMES = 12 # 240ms fadeout time
|
|
|
|
def read(self) -> bytes:
|
|
"""
|
|
Reads data from the audio stream, applies the noise gate,
|
|
and returns a 20ms audio frame.
|
|
"""
|
|
try:
|
|
# Read a frame's worth of data from the input stream.
|
|
pcm_data = self.audio_stream.read(SAMPLES_PER_FRAME, exception_on_overflow=False)
|
|
|
|
# Ensure we have a full frame of data.
|
|
if len(pcm_data) != FRAME_SIZE:
|
|
return SILENT_FRAME
|
|
|
|
# Calculate volume to check against the threshold.
|
|
rms = audioop.rms(pcm_data, 2)
|
|
if rms == 0:
|
|
# If there's no volume, check if we're in the fadeout period.
|
|
if self.ng_fadeout_count > 0:
|
|
self.ng_fadeout_count -= 1
|
|
return pcm_data # Return the (silent) data to complete the fade
|
|
return SILENT_FRAME
|
|
|
|
db = 20 * math.log10(rms)
|
|
|
|
# If volume is above the threshold, send the audio and reset fadeout.
|
|
if db >= self.threshold:
|
|
self.ng_fadeout_count = self.NG_FADEOUT_FRAMES
|
|
return pcm_data
|
|
|
|
# If below threshold but still in the fadeout period, send the audio.
|
|
if self.ng_fadeout_count > 0:
|
|
self.ng_fadeout_count -= 1
|
|
return pcm_data
|
|
|
|
# Otherwise, the gate is closed. Send silence.
|
|
return SILENT_FRAME
|
|
|
|
except Exception as e:
|
|
LOGGER.error(f"Error in NoiseGateSource.read: {e}", exc_info=True)
|
|
return SILENT_FRAME
|
|
|
|
def cleanup(self) -> None:
|
|
"""Called when the player stops."""
|
|
# The AudioStreamManager now handles cleanup.
|
|
LOGGER.info("Audio source cleanup called.")
|
|
pass
|
|
|
|
class AudioStreamManager:
|
|
"""Manages the PyAudio instance and input stream."""
|
|
def __init__(self, input_device_index: int):
|
|
self.pa = pyaudio.PyAudio()
|
|
self.stream = self.pa.open(
|
|
format=pyaudio.paInt16,
|
|
channels=CHANNELS,
|
|
rate=SAMPLE_RATE,
|
|
input=True,
|
|
frames_per_buffer=SAMPLES_PER_FRAME,
|
|
input_device_index=input_device_index
|
|
)
|
|
self.stream.start_stream()
|
|
LOGGER.info(f"Audio stream started on device {input_device_index}")
|
|
|
|
def get_stream(self):
|
|
return self.stream
|
|
|
|
def terminate(self):
|
|
if self.stream and self.stream.is_active():
|
|
self.stream.stop_stream()
|
|
self.stream.close()
|
|
self.pa.terminate()
|
|
LOGGER.info("PyAudio instance terminated.") |