human/human/audio_inference_handler.py

#encoding = utf8
import logging
import os
import queue
import time
from threading import Event, Thread

import cv2
import numpy as np
import torch

from eventbus import EventBus
from human_handler import AudioHandler
from utils import load_model, mirror_index, get_device, SyncQueue
from .huaman_status import HumanStatus

logger = logging.getLogger(__name__)
current_file_path = os.path.dirname(os.path.abspath(__file__))


class AudioInferenceHandler(AudioHandler):
    def __init__(self, context, handler):
        super().__init__(context, handler)

        EventBus().register('stop', self._on_stop)
        self._mal_queue = SyncQueue(1, "AudioInferenceHandler_Mel")
        self._audio_queue = SyncQueue(context.batch_size * 2, "AudioInferenceHandler_Audio")

        self._is_running = True
        self._exit_event = Event()
        self._run_thread = Thread(target=self.__on_run, name="AudioInferenceHandlerThread")
        self._exit_event.set()
        self._run_thread.start()

        logger.info("AudioInferenceHandler init")

    def __del__(self):
        EventBus().unregister('stop', self._on_stop)

    def _on_stop(self, *args, **kwargs):
        self.stop()

    def on_handle(self, stream, type_):
        if not self._is_running:
            return

        if type_ == 1:
            self._mal_queue.put(stream)
        elif type_ == 0:
            self._audio_queue.put(stream)
        # print('AudioInferenceHandler on_handle', type_, self._audio_queue.size())

    def on_message(self, message):
        super().on_message(message)

    def __on_run(self):
        # wav2lip_path = os.path.join(current_file_path, '..', 'checkpoints', 'wav2lip.pth')
        wav2lip_path = os.path.join(current_file_path, '..', 'checkpoints', 'weights', 'wav2lip',
                                    'ema_checkpoint_step000300000.pth')
        logger.info(f'AudioInferenceHandler init, path:{wav2lip_path}')
        model = load_model(wav2lip_path)
        logger.info("Model loaded")

        face_list_cycle = self._context.face_list_cycle
        length = len(face_list_cycle)
        index = 0
        count = 0
        count_time = 0
        logger.info('start inference')
        silence_length = 133
        human_status = HumanStatus(length, silence_length)

        device = get_device()
        logger.info(f'use device:{device}')

        while self._is_running:
            if self._exit_event.is_set():
                start_time = time.perf_counter()
                batch_size = self._context.batch_size
                try:
                    mel_batch = self._mal_queue.get(timeout=0.02)
                    # print('AudioInferenceHandler mel_batch:', len(mel_batch), 'size:', self._mal_queue.size())
                except queue.Empty:
                    continue

                # print('origin mel_batch:', len(mel_batch))
                is_all_silence = True
                audio_frames = []
                current_text = ''
                for _ in range(batch_size * 2):
                    frame, type_ = self._audio_queue.get()
                    # print('AudioInferenceHandler type_', type_)
                    current_text = frame[1]
                    audio_frames.append((frame, type_))
                    if type_ == 0:
                        is_all_silence = False

                if not self._is_running:
                    print('AudioInferenceHandler not running')
                    break

                if is_all_silence:
                    for i in range(batch_size):
                        if not self._is_running:
                            break
                        self.on_next_handle((None, mirror_index(length, index),
                        # self.on_next_handle((None, human_status.get_index(),
                                             audio_frames[i * 2:i * 2 + 2]), 0)
                        index = index + 1
                else:
                    human_status.start_talking()
                    logger.info(f'infer======= {current_text}')
                    # human_status.try_to_talk()
                    t = time.perf_counter()
                    img_batch = []
                    index_list = []
                    # for i in range(batch_size):
                    for i in range(len(mel_batch)):
                        idx = mirror_index(length, index + i)
                        # idx = human_status.get_index()
                        # index_list.append(idx)
                        face = face_list_cycle[idx]
                        img_batch.append(face)

                    # print('orign img_batch:', len(img_batch), 'origin mel_batch:', len(mel_batch))
                    img_batch = np.asarray(img_batch)
                    mel_batch = np.asarray(mel_batch)
                    img_masked = img_batch.copy()
                    img_masked[:, face.shape[0] // 2:] = 0

                    img_batch = np.concatenate((img_masked, img_batch), axis=3) / 255.
                    mel_batch = np.reshape(mel_batch,
                                           [len(mel_batch), mel_batch.shape[1], mel_batch.shape[2], 1])

                    img_batch = torch.FloatTensor(np.transpose(img_batch, (0, 3, 1, 2))).to(device)
                    mel_batch = torch.FloatTensor(np.transpose(mel_batch, (0, 3, 1, 2))).to(device)
                    # print('img_batch:', img_batch.shape, 'mel_batch:', mel_batch.shape)
                    with torch.no_grad():
                        pred = model(mel_batch, img_batch)

                    pred = pred.cpu().numpy().transpose(0, 2, 3, 1) * 255.

                    count_time += (time.perf_counter() - t)
                    count += batch_size

                    if count >= 100:
                        logger.info(f"------actual avg infer fps:{count / count_time:.4f}")
                        count = 0
                        count_time = 0

                    for i, res_frame in enumerate(pred):
                        if not self._is_running:
                            break
                        self.on_next_handle(
                            (res_frame, mirror_index(length, index), audio_frames[i * 2:i * 2 + 2]),
                            # (res_frame, index_list[i], audio_frames[i * 2:i * 2 + 2]),
                            0)
                        index = index + 1

                    logger.info(f'total batch time: {time.perf_counter() - start_time}')
            else:
                time.sleep(1)
                break
        logger.info('AudioInferenceHandler inference processor stop')

    def stop(self):
        logger.info('AudioInferenceHandler stop')
        self._is_running = False
        self._exit_event.clear()
        if self._run_thread.is_alive():
            logger.info('AudioInferenceHandler stop join')
            self._run_thread.join()
        logger.info('AudioInferenceHandler stop exit')

    def pause_talk(self):
        print('AudioInferenceHandler pause_talk', self._audio_queue.size(), self._mal_queue.size())
        self._audio_queue.clear()
        self._mal_queue.clear()