capture.py

import logging
import collections, queue, os, os.path
import numpy as np
import pyaudio
import wave
import webrtcvad
from halo import Halo
from scipy import signal
import websockets
import asyncio
import traceback
import transcribe
import translate
import punctuate

logging.basicConfig(level=20)

def make_iter():
    loop = asyncio.get_event_loop()
    queue = asyncio.Queue()
    def put(*args):
        loop.call_soon_threadsafe(queue.put_nowait, args)
    return queue, put

class Audio(object):
    """Streams raw audio from microphone. Data is received in a separate thread, and stored in a buffer, to be read from."""

    FORMAT = pyaudio.paInt16
    # Network/VAD rate-space
    RATE_PROCESS = 16000
    CHANNELS = 1
    BLOCKS_PER_SECOND = 50

    def __init__(self, callback=None, device=None, input_rate=RATE_PROCESS, file=None):
        def proxy_callback(in_data, frame_count, time_info, status):
            #pylint: disable=unused-argument
            if self.chunk is not None:
                in_data = self.wf.readframes(self.chunk)
            callback(in_data)
            return (None, pyaudio.paContinue)
        self.buffer_queue, stream_put = make_iter()
        if callback is None: callback = stream_put
        self.device = device
        self.input_rate = input_rate
        self.sample_rate = self.RATE_PROCESS
        self.block_size = int(self.RATE_PROCESS / float(self.BLOCKS_PER_SECOND))
        self.block_size_input = int(self.input_rate / float(self.BLOCKS_PER_SECOND))
        self.pa = pyaudio.PyAudio()

        kwargs = {
            'format': self.FORMAT,
            'channels': self.CHANNELS,
            'rate': self.input_rate,
            'input': True,
            'frames_per_buffer': self.block_size_input,
            'stream_callback': proxy_callback,
        }

        self.chunk = None
        # if not default device
        if self.device:
            kwargs['input_device_index'] = self.device
        elif file is not None:
            self.chunk = 320
            self.wf = wave.open(file, 'rb')

        self.stream = self.pa.open(**kwargs)
        self.stream.start_stream()

    def resample(self, data, input_rate):
        """
        Microphone may not support our native processing sampling rate, so
        resample from input_rate to RATE_PROCESS here for webrtcvad and
        deepspeech

        Args:
            data (binary): Input audio stream
            input_rate (int): Input audio rate to resample from
        """
        data16 = np.fromstring(string=data, dtype=np.int16)
        resample_size = int(len(data16) / self.input_rate * self.RATE_PROCESS)
        resample = signal.resample(data16, resample_size)
        resample16 = np.array(resample, dtype=np.int16)
        return resample16.tostring()

    async def read_resampled(self):
        """Return a block of audio data resampled to 16000hz, blocking if necessary."""
        return self.resample(data=await self.buffer_queue.get(),
                             input_rate=self.input_rate)

    async def read(self):
        # print("Heeee")
        """Return a block of audio data, blocking if necessary."""
        return await self.buffer_queue.get()

    def destroy(self):
        self.stream.stop_stream()
        self.stream.close()
        self.pa.terminate()

    frame_duration_ms = property(lambda self: 1000 * self.block_size // self.sample_rate)

    def write_wav(self, filename, data):
        logging.info("write wav %s", filename)
        wf = wave.open(filename, 'wb')
        wf.setnchannels(self.CHANNELS)
        # wf.setsampwidth(self.pa.get_sample_size(FORMAT))
        assert self.FORMAT == pyaudio.paInt16
        wf.setsampwidth(2)
        wf.setframerate(self.sample_rate)
        wf.writeframes(data)
        wf.close()


class VADAudio(Audio):
    """Filter & segment audio with voice activity detection."""

    def __init__(self, aggressiveness=3, device=None, input_rate=None, file=None):
        super().__init__(device=device, input_rate=input_rate, file=file)
        self.vad = webrtcvad.Vad(aggressiveness)

    async def frame_generator(self):
        """Generator that yields all audio frames from microphone."""
        if self.input_rate == self.RATE_PROCESS:
            while True:
                yield await self.read()
        else:
            while True:
                yield await self.read_resampled()

    async def vad_collector(self, padding_ms=300, ratio=0.75, frames=None):
        """Generator that yields series of consecutive audio frames comprising each utterence, separated by yielding a single None.
            Determines voice activity by ratio of frames in padding_ms. Uses a buffer to include padding_ms prior to being triggered.
            Example: (frame, ..., frame, None, frame, ..., frame, None, ...)
                      |---utterence---|        |---utterence---|
        """
        if frames is None: frames = self.frame_generator()
        num_padding_frames = padding_ms // self.frame_duration_ms
        ring_buffer = collections.deque(maxlen=num_padding_frames)
        triggered = False

        async for frame in frames:
            #I donno why this suddenly became necessary when making the whole pyaudio stuff async but well
            frame = frame[0]
            if len(frame) < 640:
                # print(frame)
                # yield None
                # continue
                return

            is_speech = self.vad.is_speech(frame, self.sample_rate)

            if not triggered:
                ring_buffer.append((frame, is_speech))
                num_voiced = len([f for f, speech in ring_buffer if speech])
                if num_voiced > ratio * ring_buffer.maxlen:
                    triggered = True
                    for f, s in ring_buffer:
                        yield f
                    ring_buffer.clear()

            else:
                yield frame
                ring_buffer.append((frame, is_speech))
                num_unvoiced = len([f for f, speech in ring_buffer if not speech])
                if num_unvoiced > ratio * ring_buffer.maxlen:
                    triggered = False
                    yield None
                    ring_buffer.clear()


def main(vad_aggressiveness, vad_device, vad_rate, source_lang, target_lang):
    # Start audio with VAD
    vad_audio = VADAudio(aggressiveness=vad_aggressiveness,
                         device=vad_device,
                         input_rate=vad_rate)
    print("Listening (ctrl-C to exit)...")

    async def translator():
        frames = vad_audio.vad_collector()
        print(frames)
        # Stream from microphone to DeepSpeech using VAD
        # while True:
        #     yield await frames.__anext__()
        spinner = Halo(spinner='line')
        wav_data = bytearray()
        async for frame in frames:
            if frame is not None:
                spinner.start()
                wav_data.extend(frame)
            else:
                spinner.stop()
                vad_file = os.path.join(os.path.dirname(__file__), "process.wav")
                vad_audio.write_wav(vad_file, wav_data)
                wav_data = bytearray()
                text = transcribe.transcribe(vad_file)
                punctuated = punctuate.punctuate(text.lower())
                print("Recognized: %s" % punctuated)
                yield punctuated
                # translation = translate.translate(punctuated, source_lang, target_lang)
                translation = punctuated
                print("Translation: %s" % translation)
                yield translation
                os.remove(vad_file)
    # #
    async def send_result(websocket, path):
        result = translator()
        async for msg in result:
            try:
                # print(type(msg))
                await websocket.send("_"+msg)
                await websocket.send("hi")
            except Exception:
                traceback.print_exc()

    # async def test():
    #     result = translator()
    #     async for msg in result:
    #         print(msg)

    async def test2():
        frames = vad_audio.vad_collector()
        print(frames)
        async for frame in frames:
            if frame is not None:
                print(frame)
    # asyncio.get_event_loop().run_until_complete(test2())

    ##To connect with Neos websockets
    asyncio.get_event_loop().run_until_complete(websockets.serve(send_result, 'localhost', 8765))

    # asyncio.get_event_loop().run_until_complete(test())
    asyncio.get_event_loop().run_forever()
    # translator()