From 4998d0b4223d50d5df8cc21d0ec6471c7540c70d Mon Sep 17 00:00:00 2001 From: brige Date: Tue, 5 Nov 2024 19:40:03 +0800 Subject: [PATCH] modify code --- human/human_render.py | 11 ++++---- render/base_render.py | 36 +------------------------ render/video_render.py | 57 +-------------------------------------- render/voice_render.py | 46 +------------------------------ ui/pygame_ui.py | 2 +- utils/async_task_queue.py | 27 ++++++++----------- utils/utils.py | 2 +- 7 files changed, 22 insertions(+), 159 deletions(-) diff --git a/human/human_render.py b/human/human_render.py index 673ad66..8189e7b 100644 --- a/human/human_render.py +++ b/human/human_render.py @@ -33,7 +33,7 @@ class HumanRender(AudioHandler): logging.info('human render run') while self._exit_event.is_set(): self._run_step() - time.sleep(0.037) + time.sleep(0.038) logging.info('human render exit') @@ -82,12 +82,13 @@ class HumanRender(AudioHandler): # if self._voice_render.is_full(): # self._context.notify({'msg_id': MessageType.Video_Render_Queue_Full}) - def get_audio_queue_size(self): - return self._voice_render.size() + # def get_audio_queue_size(self): + # return self._voice_render.size() def pause_talk(self): - self._voice_render.pause_talk() - self._video_render.pause_talk() + pass + # self._voice_render.pause_talk() + # self._video_render.pause_talk() def stop(self): logging.info('hunan render stop') diff --git a/render/base_render.py b/render/base_render.py index 59090a0..3bccaf5 100644 --- a/render/base_render.py +++ b/render/base_render.py @@ -11,11 +11,9 @@ logger = logging.getLogger(__name__) class BaseRender(ABC): - def __init__(self, play_clock, context, type_, delay=0.02, thread_name="BaseRenderThread"): + def __init__(self, play_clock, context, type_): self._play_clock = play_clock self._context = context - self._type = type_ - self._delay = delay # self._queue = SyncQueue(context.batch_size, f'{type_}RenderQueue') # self._exit_event = Event() # self._thread = Thread(target=self._on_run, name=thread_name) @@ -25,35 +23,3 @@ class BaseRender(ABC): @abstractmethod def render(self, frame, ps): pass - - def _on_run(self): - logging.info(f'{self._type} render run') - # while self._exit_event.is_set(): - # self._run_step() - # time.sleep(self._delay) - - logging.info(f'{self._type} render exit') - - def put(self, frame, ps): - pass - # self._queue.put((frame, ps)) - - def size(self): - pass - # return self._queue.size() - - def pause_talk(self): - pass - # self._queue.clear() - - def stop(self): - pass - # self._queue.clear() - # self._exit_event.clear() - # self._thread.join() - - @abstractmethod - def _run_step(self): - pass - - diff --git a/render/video_render.py b/render/video_render.py index 57135ea..61ac308 100644 --- a/render/video_render.py +++ b/render/video_render.py @@ -13,7 +13,7 @@ from human.message_type import MessageType class VideoRender(BaseRender): def __init__(self, play_clock, context, human_render): - super().__init__(play_clock, context, 'Video', 0.038, "VideoRenderThread") + super().__init__(play_clock, context, 'Video') self._human_render = human_render self._diff_avg_count = 0 @@ -21,7 +21,6 @@ class VideoRender(BaseRender): res_frame, idx, type_ = frame clock_time = self._play_clock.clock_time() time_difference = clock_time - ps - print("Video frame time", clock_time, ps, time_difference) if abs(time_difference) > self._play_clock.audio_diff_threshold: if self._diff_avg_count < 5: self._diff_avg_count += 1 @@ -56,57 +55,3 @@ class VideoRender(BaseRender): # image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) if self._human_render is not None: self._human_render.put_image(image) - - def _run_step(self): - while self._exit_event.is_set(): - try: - value = self._queue.get() - if value is None: - return - frame, ps = value - res_frame, idx, type_ = frame - except Empty: - return - - clock_time = self._play_clock.clock_time() - time_difference = clock_time - ps - print("Video frame time", clock_time, ps, time_difference) - if abs(time_difference) > self._play_clock.audio_diff_threshold: - if self._diff_avg_count < 5: - self._diff_avg_count += 1 - else: - if time_difference < -self._play_clock.audio_diff_threshold: - sleep_time = abs(time_difference) - print("Video frame waiting to catch up with audio", sleep_time) - if sleep_time <= 1.0: - time.sleep(sleep_time) - - # elif time_difference > self._play_clock.audio_diff_threshold: # 视频比音频快超过10ms - # print("Video frame dropped to catch up with audio") - # continue - - else: - self._diff_avg_count = 0 - - print('video render:', - 'get face', self._queue.size(), - 'audio queue', self._human_render.get_audio_queue_size()) - - if type_ == 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: - print('resize error') - return - combine_frame[y1:y2, x1:x2, :3] = res_frame - - image = combine_frame - # image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) - if self._human_render is not None: - self._human_render.put_image(image) - return diff --git a/render/voice_render.py b/render/voice_render.py index 19b776a..ad2bf5e 100644 --- a/render/voice_render.py +++ b/render/voice_render.py @@ -15,11 +15,7 @@ logger = logging.getLogger(__name__) class VoiceRender(BaseRender): def __init__(self, play_clock, context): self._audio_render = AudioRender() - self._is_empty = True - super().__init__(play_clock, context, 'Voice', 0.03, "VoiceRenderThread") - - def is_full(self): - return self._queue.size() >= self._context.render_batch * 2 + super().__init__(play_clock, context, 'Voice') def render(self, frame, ps): self._play_clock.update_display_time() @@ -36,43 +32,3 @@ class VoiceRender(BaseRender): self._audio_render.write(frame.tobytes(), chunk_len) except Exception as e: logging.error(f'Error writing audio frame: {e}') - - def _run_step(self): - try: - value = self._queue.get() - if value is None: - return - audio_frames, ps = value - # print('voice render queue size', self._queue.size()) - except Empty: - self._context.notify({'msg_id': MessageType.Video_Render_Queue_Empty}) - if not self._is_empty: - print('voice render queue empty') - self._is_empty = True - return - - if self._is_empty: - print('voice render queue not empty') - self._is_empty = False - - status = MessageType.Video_Render_Queue_Not_Empty - if self._queue.size() < self._context.render_batch: - status = MessageType.Video_Render_Queue_Empty - elif self._queue.size() >= self._context.render_batch * 2: - status = MessageType.Video_Render_Queue_Full - self._context.notify({'msg_id': status}) - - self._play_clock.update_display_time() - self._play_clock.current_time = ps - - for audio_frame in audio_frames: - frame, type_ = audio_frame - frame = (frame * 32767).astype(np.int16) - - if self._audio_render is not None: - try: - chunk_len = int(frame.shape[0] * 2) - # print('audio frame:', frame.shape, chunk_len) - self._audio_render.write(frame.tobytes(), chunk_len) - except Exception as e: - logging.error(f'Error writing audio frame: {e}') diff --git a/ui/pygame_ui.py b/ui/pygame_ui.py index 7910329..50e6ab2 100644 --- a/ui/pygame_ui.py +++ b/ui/pygame_ui.py @@ -44,9 +44,9 @@ class PyGameUI: elif event.type == VIDEORESIZE: self.background_display_ = pygame.transform.scale(self._background, event.dict['size']) self.screen_.blit(self.background_display_, (0, 0)) + self._update_human() if self._human_image is not None: self.screen_.blit(self._human_image, (0, 0)) - self._update_human() pygame.display.flip() self.stop() pygame.quit() diff --git a/utils/async_task_queue.py b/utils/async_task_queue.py index 716b9cd..fc03d4e 100644 --- a/utils/async_task_queue.py +++ b/utils/async_task_queue.py @@ -9,7 +9,6 @@ class AsyncTaskQueue: self._queue = asyncio.Queue() self._worker_num = work_num self._current_worker_num = work_num - self._condition = threading.Condition() self._thread = threading.Thread(target=self._run_loop) self._thread.start() self.__loop = None @@ -27,35 +26,31 @@ class AsyncTaskQueue: async def _worker(self): print('_worker') while True: - with self._condition: - self._condition.wait_for(lambda: not self._queue.empty()) - task = await self._queue.get() - func, *args = task # 解包任务 - if func is None: # None 作为结束信号 - break + task = await self._queue.get() + if task is None: # None as a stop signal + break - print(f"Executing task with args: {args}") - await func(*args) # 执行异步函数 - self._queue.task_done() + func, *args = task # Unpack task + print(f"Executing task with args: {args}") + await func(*args) # Execute async function + self._queue.task_done() print('_worker finish') - self._current_worker_num = self._current_worker_num - 1 + self._current_worker_num -= 1 if self._current_worker_num == 0: print('loop stop') self.__loop.stop() def add_task(self, func, *args): - with self._condition: - self._queue.put_nowait((func, *args)) - self._condition.notify() + self.__loop.call_soon_threadsafe(self._queue.put_nowait, (func, *args)) def stop_workers(self): for _ in range(self._worker_num): - self.add_task(None) # 发送结束信号 + self.add_task(None) # Send stop signal def clear(self): while not self._queue.empty(): - self._queue.get() + self._queue.get_nowait() self._queue.task_done() def stop(self): diff --git a/utils/utils.py b/utils/utils.py index 82208f9..78115b5 100644 --- a/utils/utils.py +++ b/utils/utils.py @@ -178,7 +178,7 @@ def load_avatar(path, img_size, device): face_frames = [] coord_frames = [] for face, coord in face_det_results: - resized_crop_frame = cv2.resize(face, (img_size, img_size)) + resized_crop_frame = cv2.resize(face[:, :, :3], (img_size, img_size)) face_frames.append(resized_crop_frame) coord_frames.append(coord)