human/tts/Chunk2Mal.py

#encoding = utf8

import logging
import queue
from queue import Queue
from threading import Thread, Event

import numpy as np
import audio


class Chunk2Mal:
    def __init__(self, human):
        self._audio_chunk_queue = Queue()
        self._human = human
        self._thread = None
        self._exit_event = None
        self._chunks = []

    def _on_run(self):
        logging.info('chunk2mal run')
        while not self._exit_event.is_set():
            try:
                chunk, type_ = self.pull_chunk()
                self._chunks.append(chunk)
                self._human.push_audio_frame((chunk, type_))
                print("1")
            except queue.Empty:
                continue

            if len(self._chunks) <= self._human.get_stride_left_size() + self._human.get_stride_right_size():
                continue

            inputs = np.concatenate(self._chunks)  # [N * chunk]
            mel = audio.melspectrogram(inputs)
            left = max(0, self._human.get_stride_left_size() * 80 / 50)
            right = min(len(mel[0]), len(mel[0]) - self._human.get_stride_right_size() * 80 / 50)
            mel_idx_multiplier = 80. * 2 / self._human.get_fps()
            mel_step_size = 16
            i = 0
            mel_chunks = []
            while i < (len(self._chunks) - self._human.get_stride_left_size()
                       - self._human.get_stride_right_size()) / 2:
                print("14")
                start_idx = int(left + i * mel_idx_multiplier)
                # print(start_idx)
                if start_idx + mel_step_size > len(mel[0]):
                    mel_chunks.append(mel[:, len(mel[0]) - mel_step_size:])
                else:
                    mel_chunks.append(mel[:, start_idx: start_idx + mel_step_size])
                i += 1
                print("13")
            self._human.push_feat_queue(mel_chunks)
            print("15")

            # discard the old part to save memory
            self._chunks = self._chunks[-(self._human.get_stride_left_size() + self._human.get_stride_right_size()):]
            print("12")

        logging.info('chunk2mal exit')

    def start(self):
        if self._exit_event is not None:
            return
        self._exit_event = Event()
        self._thread = Thread(target=self._on_run)
        self._thread.start()
        logging.info('chunk2mal start')

    def stop(self):
        if self._exit_event is None:
            return

        self._exit_event.set()
        if self._thread.is_alive():
            self._thread.join()
        logging.info('chunk2mal stop')

    def push_chunk(self, chunk):
        self._audio_chunk_queue.put(chunk)

    def pull_chunk(self):
        try:
            chunk = self._audio_chunk_queue.get(block=True, timeout=1)
            type = 1
        except queue.Empty:
            chunk = np.zeros(self._human.get_chunk(), dtype=np.float32)
            type = 0
        return chunk, type