human/human/human_render.py

#encoding = utf8

import logging
import time
from queue import Empty
from threading import Event, Thread

from eventbus import EventBus
from human.message_type import MessageType
from human_handler import AudioHandler
from render import VoiceRender, VideoRender, PlayClock
from utils import SyncQueue

logger = logging.getLogger(__name__)


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

        EventBus().register('stop', self._on_stop)
        EventBus().register('clear_cache', self.on_clear_cache)
        play_clock = PlayClock()
        self._voice_render = VoiceRender(play_clock, context)
        self._video_render = VideoRender(play_clock, context, self)
        self._is_running = True
        self._queue = SyncQueue(context.batch_size, "HumanRender_queue")
        self._exit_event = Event()
        self._thread = Thread(target=self._on_run, name="AudioMalHandlerThread")
        self._exit_event.set()
        self._thread.start()
        self._image_render = None
        self._last_audio_ps = 0
        self._last_video_ps = 0
        self._empty_log = True

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

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

    def on_clear_cache(self, *args, **kwargs):
        self._queue.clear()

    def _on_run(self):
        logging.info('human render run')
        while self._exit_event.is_set() and self._is_running:
            # t = time.time()
            self._run_step()
            # delay = time.time() - t
            delay = 0.03805 # - delay
            # print(delay)
            # if delay <= 0.0:
            #     continue
            time.sleep(delay)

        logging.info('human render exit')

    def _run_step(self):
        try:
            value = self._queue.get(timeout=.005)
            if value is None:
                return
            res_frame, idx, audio_frames = value
            # print('render queue size', self._queue.size())
            if not self._empty_log:
                self._empty_log = True
                logging.info('render render:')
            # print('voice render queue size', self._queue.size())
        except Empty:
            if self._empty_log:
                self._empty_log = False
                logging.info('render queue.Empty:')
            return

        type_ = 1
        if audio_frames[0][1] != 0 and audio_frames[1][1] != 0:
            type_ = 0
        if self._voice_render is not None:
            self._voice_render.render(audio_frames, self._last_audio_ps)
            self._last_audio_ps = self._last_audio_ps + 0.4
        if self._video_render is not None:
            self._video_render.render((res_frame, idx, type_), self._last_video_ps)
            self._last_video_ps = self._last_video_ps + 0.4

    def set_image_render(self, render):
        self._image_render = render

    def put_image(self, image):
        if self._image_render is not None:
            self._image_render.on_render(image)

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

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

        self._queue.put(stream)

    def pause_talk(self):
        logging.info('hunan pause_talk')
        # self._voice_render.pause_talk()
        # self._video_render.pause_talk()

    def stop(self):
        logging.info('hunan render stop')
        self._is_running = False
        if self._exit_event is None:
            return

        self._queue.clear()
        self._exit_event.clear()
        if self._thread.is_alive():
            self._thread.join()
        logging.info('hunan render stop')
        # self._voice_render.stop()
        # self._video_render.stop()
        # self._exit_event.clear()
        # self._thread.join()
'''
        self._exit_event = Event()
        self._thread = Thread(target=self._on_run)
        self._exit_event.set()
        self._thread.start()

    def _on_run(self):
        logging.info('human render run')
        while self._exit_event.is_set():
            self._run_step()
            time.sleep(0.02)

        logging.info('human render exit')

    def _run_step(self):
        try:
            res_frame, idx, audio_frames = self._queue.get(block=True, timeout=.002)
        except queue.Empty:
            # print('render queue.Empty:')
            return None
        if audio_frames[0][1] != 0 and audio_frames[1][1] != 0:
            combine_frame = self._context.frame_list_cycle[idx]
        else:
            bbox = self._context.coord_list_cycle[idx]
            combine_frame = copy.deepcopy(self._context.frame_list_cycle[idx])
            y1, y2, x1, x2 = bbox
            try:
                res_frame = cv2.resize(res_frame.astype(np.uint8), (x2 - x1, y2 - y1))
            except:
                return
            # combine_frame = get_image(ori_frame,res_frame,bbox)
            # t=time.perf_counter()
            combine_frame[y1:y2, x1:x2] = res_frame

        image = combine_frame
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

        if self._image_render is not None:
            self._image_render.on_render(image)

        for audio_frame in audio_frames:
            frame, type_ = audio_frame
            frame = (frame * 32767).astype(np.int16)
            if self._audio_render is not None:
                self._audio_render.write(frame.tobytes(), int(frame.shape[0]*2))
            # new_frame = AudioFrame(format='s16', layout='mono', samples=frame.shape[0])
            # new_frame.planes[0].update(frame.tobytes())
            # new_frame.sample_rate = 16000
'''