Discord-Radio-Bot/NoiseGatev2.py

import audioop
import logging
import math
import time

import pyaudio
import discord
import numpy

voice_connection = None

LOGGER = logging.getLogger("Discord_Radio_Bot.NoiseGateV2")


# noinspection PyUnresolvedReferences
class AudioStream:
    def __init__(self, _channels: int = 2, _sample_rate: int = 48000, _frames_per_buffer: int = 1024,
                 _input_device_index: int = None, _output_device_index: int = None, _input: bool = True,
                 _output: bool = True, _init_on_startup: bool = True):
        self.paInstance_kwargs = {
            'format': pyaudio.paInt16,
            'channels': _channels,
            'rate': _sample_rate,
            'input': _input,
            'output': _output,
            'frames_per_buffer': _frames_per_buffer
        }

        if _input_device_index:
            if _input:
                self.paInstance_kwargs['input_device_index'] = _input_device_index
            else:
                LOGGER.warning(f"[AudioStream.__init__]:\tInput was not enabled."
                               f" Reinitialize with '_input=True'")

        if _output_device_index:
            if _output:
                self.paInstance_kwargs['output_device_index'] = _output_device_index
            else:
                LOGGER.warning(f"[AudioStream.__init__]:\tOutput was not enabled."
                               f" Reinitialize with '_output=True'")

        if _init_on_startup:
            # Init PyAudio instance
            LOGGER.info("Creating PyAudio instance")
            self.paInstance = pyaudio.PyAudio()

            # Define and initialize stream object if we have been passed a device ID (pyaudio.open)
            self.stream = None

        if _output_device_index or _input_device_index:
            if _init_on_startup:
                LOGGER.info("Init stream")
                self.init_stream()

    def init_stream(self, _new_output_device_index: int = None, _new_input_device_index: int = None):
        # Check what device was asked to be changed (or set)
        if _new_input_device_index:
            if self.paInstance_kwargs['input']:
                self.paInstance_kwargs['input_device_index'] = _new_input_device_index
            else:
                LOGGER.warning(f"[AudioStream.init_stream]:\tInput was not enabled when initialized."
                               f" Reinitialize with '_input=True'")

        if _new_output_device_index:
            if self.paInstance_kwargs['output']:
                self.paInstance_kwargs['output_device_index'] = _new_output_device_index
            else:
                LOGGER.warning(f"[AudioStream.init_stream]:\tOutput was not enabled when initialized."
                               f" Reinitialize with '_output=True'")

        self.close_if_open()

        # Open the stream
        self.stream = self.paInstance.open(**self.paInstance_kwargs)

    def close_if_open(self):
        # Stop the stream if it is started
        if self.stream:
            if self.stream.is_active():
                self.stream.stop_stream()
                self.stream.close()
                LOGGER.debug(f"[ReopenStream.close_if_open]:\t Stream was open; It was closed.")

    def list_devices(self, _display_input_devices: bool = True, _display_output_devices: bool = True):
        info = self.paInstance.get_host_api_info_by_index(0)
        numdevices = info.get('deviceCount')

        devices = {
            'Input': {},
            'Output': {}
        }
        for i in range(0, numdevices):
            if (self.paInstance.get_device_info_by_host_api_device_index(0, i).get('maxInputChannels')) > 0:
                input_device = self.paInstance.get_device_info_by_host_api_device_index(0, i).get('name')
                devices['Input'][i] = input_device
                if _display_input_devices:
                    LOGGER.debug("Input Device id ", i, " - ", input_device)

            if (self.paInstance.get_device_info_by_host_api_device_index(0, i).get('maxOutputChannels')) > 0:
                output_device = self.paInstance.get_device_info_by_host_api_device_index(0, i).get('name')
                devices['Output'][i] = output_device
                if _display_output_devices:
                    LOGGER.debug("Output Device id ", i, " - ", output_device)

        return devices

    async def stop(self):
        await voice_connection.disconnect()
        self.close_if_open()
        self.stream.close()
        self.paInstance.terminate()


# noinspection PyUnresolvedReferences
class NoiseGate(AudioStream):
    def __init__(self, _voice_connection, _noise_gate_threshold: int, **kwargs):
        super(NoiseGate, self).__init__(_init_on_startup=True, **kwargs)
        global voice_connection
        voice_connection = _voice_connection
        self.THRESHOLD = _noise_gate_threshold
        self.NGStream = NoiseGateStream(self)
        self.Voice_Connection_Thread = None

    def run(self) -> None:
        global voice_connection
        # Start the audio stream
        LOGGER.debug(f"Starting stream")
        self.stream.start_stream()
        # Start the stream to discord
        self.core()

    def core(self, error=None):
        if error:
            LOGGER.warning(error)

        while not voice_connection.is_connected():
            time.sleep(.2)

        if not voice_connection.is_playing():
            LOGGER.debug(f"Playing stream to discord")
            voice_connection.play(self.NGStream, after=self.core)

    async def close(self):
        LOGGER.debug(f"Closing")
        await voice_connection.disconnect()
        if self.stream.is_active:
            self.stream.stop_stream()
            LOGGER.debug(f"Stopping stream")


# noinspection PyUnresolvedReferences
class NoiseGateStream(discord.AudioSource):
    def __init__(self, _stream):
        super(NoiseGateStream, self).__init__()
        self.stream = _stream       # The actual audio stream object
        self.NG_fadeout = 240/20    # Fadeout value used to hold the noisegate after de-triggering
        self.NG_fadeout_count = 0   # A count set when the noisegate is triggered and was de-triggered
        self.process_set_count = 0  # Counts how many processes have been made

    def read(self):
        try:
            while voice_connection.is_connected():
                curr_buffer = bytearray(self.stream.stream.read(960))
                buffer_rms = audioop.rms(curr_buffer, 2)
                if buffer_rms > 0:
                    buffer_decibel = 20 * math.log10(buffer_rms)

                    if self.process_set_count % 10 == 0:
                        if buffer_decibel >= self.stream.THRESHOLD:
                            LOGGER.debug(f"[Noisegate Open] {buffer_decibel} db")
                        else:
                            LOGGER.debug(f"[Noisegate Closed] {buffer_decibel} db")

                    if buffer_decibel >= self.stream.THRESHOLD:
                        self.NG_fadeout_count = self.NG_fadeout
                        self.process_set_count += 1
                        if curr_buffer:
                            return bytes(curr_buffer)

                    else:
                        if self.NG_fadeout_count > 0:
                            self.NG_fadeout_count -= 1
                            LOGGER.debug(f"Frames in fadeout remaining: {self.NG_fadeout_count}")
                            self.process_set_count += 1
                            if curr_buffer:
                                return bytes(curr_buffer)

        except OSError as e:
            LOGGER.warning(e)
            pass

    def audio_datalist_set_volume(self, datalist, volume):
        """ Change value of list of audio chunks """
        sound_level = (volume / 100.)

        for i in range(len(datalist)):
            chunk = numpy.fromstring(datalist[i], numpy.int16)

            chunk = chunk * sound_level

            datalist[i] = chunk.astype(numpy.int16)


if __name__ == '__main__':

    input_index = int(input("Input:\t"))
    output_index = int(input("Output:\t"))

    ng = NoiseGate(_input_device_index=input_index, _output_device_index=output_index)

    ng.list_devices()

    ng.start()