From e9691dbc81006f016da9701383d38c8155e555b9 Mon Sep 17 00:00:00 2001
From: jiegeaiai <jiegeaiai@163.com>
Date: Tue, 3 Dec 2024 00:14:00 +0800
Subject: [PATCH] add main ipc render
---
human/audio_inference_onnx_handler.py | 33 +++++------
human/human_render.py | 2 +-
main.py | 17 ++++++
ui/ipc_render.py | 49 ++++++++++++++++
ui/pygame_ui.py | 83 +--------------------------
utils/__init__.py | 1 +
utils/utils.py | 36 ++++++++++++
7 files changed, 123 insertions(+), 98 deletions(-)
create mode 100644 main.py
create mode 100644 ui/ipc_render.py
diff --git a/human/audio_inference_onnx_handler.py b/human/audio_inference_onnx_handler.py
index f7d28db..8a28627 100644
--- a/human/audio_inference_onnx_handler.py
+++ b/human/audio_inference_onnx_handler.py
@@ -5,7 +5,7 @@ import queue
import time
from threading import Event, Thread
-from gfpgan import GFPGANer
+ # from gfpgan import GFPGANer
from eventbus import EventBus
from human_handler import AudioHandler
from utils import load_model, mirror_index, get_device, SyncQueue
@@ -16,25 +16,26 @@ current_file_path = os.path.dirname(os.path.abspath(__file__))
def load_gfpgan_model(model_path):
logger.info(f'load_gfpgan_model, path:{model_path}')
- model = GFPGANer(
- model_path=model_path,
- upscale=1,
- arch='clean',
- channel_multiplier=2,
- bg_upsampler=None,
- )
- return model
+ # model = GFPGANer(
+ # model_path=model_path,
+ # upscale=1,
+ # arch='clean',
+ # channel_multiplier=2,
+ # bg_upsampler=None,
+ # )
+ return None
+ #model
def load_model(model_path):
- import onnxruntime as ort
-
- sess_opt = ort.SessionOptions()
- sess_opt.intra_op_num_threads = 8
- sess = ort.InferenceSession(model_path, sess_options=sess_opt, providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])
-
- return sess
+ # import onnxruntime as ort
+ # sess_opt = ort.SessionOptions()
+ # sess_opt.intra_op_num_threads = 8
+ # sess = ort.InferenceSession(model_path, sess_options=sess_opt, providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])
+ #
+ # return sess
+ return None
class AudioInferenceOnnxHandler(AudioHandler):
def __init__(self, context, handler):
diff --git a/human/human_render.py b/human/human_render.py
index 6cd816a..186ddbc 100644
--- a/human/human_render.py
+++ b/human/human_render.py
@@ -48,7 +48,7 @@ class HumanRender(AudioHandler):
logging.info('human render run')
while self._exit_event.is_set() and self._is_running:
self._run_step()
- delay = 0.02
+ delay = 0.04
time.sleep(delay)
logging.info('human render exit')
diff --git a/main.py b/main.py
new file mode 100644
index 0000000..ac45dd5
--- /dev/null
+++ b/main.py
@@ -0,0 +1,17 @@
+#encoding = utf8
+
+import logging
+import os
+
+from utils import config_logging
+
+logger = logging.getLogger(__name__)
+current_file_path = os.path.dirname(os.path.abspath(__file__))
+
+if __name__ == '__main__':
+ config_logging('./logs/info.log', logging.INFO, logging.INFO)
+
+ logger.info('------------start------------')
+ ui = PyGameUI()
+ ui.run()
+ logger.info('------------finish------------')
\ No newline at end of file
diff --git a/ui/ipc_render.py b/ui/ipc_render.py
new file mode 100644
index 0000000..bbd1c68
--- /dev/null
+++ b/ui/ipc_render.py
@@ -0,0 +1,49 @@
+#encoding = utf8
+
+import os
+import logging
+from queue import Queue
+
+from human import HumanContext
+from ipc import IPCUtil
+
+logger = logging.getLogger(__name__)
+current_file_path = os.path.dirname(os.path.abspath(__file__))
+
+
+class IpcRender:
+ def __init__(self):
+ self._human_context = None
+ self._queue = None
+ self._ipc = IPCUtil('ipc_sender', 'ipc_sender')
+
+ def run(self):
+ self._queue = Queue()
+ self._human_context = HumanContext()
+ self._human_context.build()
+ render = self._human_context.render_handler
+ render.set_image_render(self)
+
+ def stop(self):
+ pass
+
+ def get_image(self):
+ pass
+
+ def get_audio(self):
+ pass
+
+ def send_audio(self, audio):
+ pass
+
+ def send_image(self, image):
+ pass
+
+ def send_text(self, text):
+ pass
+
+ def send_command(self, command):
+ pass
+
+ def send_binary(self, data, length):
+ pass
diff --git a/ui/pygame_ui.py b/ui/pygame_ui.py
index 29bda7e..02bb31b 100644
--- a/ui/pygame_ui.py
+++ b/ui/pygame_ui.py
@@ -1,18 +1,15 @@
#encoding = utf8
-import copy
import logging
import os
-import time
from queue import Queue
import cv2
-import numpy as np
import pygame
from pygame.locals import *
from human import HumanContext
from ipc import IPCUtil
-from utils import config_logging
+from utils import config_logging, render_image
logger = logging.getLogger(__name__)
current_file_path = os.path.dirname(os.path.abspath(__file__))
@@ -20,9 +17,7 @@ current_file_path = os.path.dirname(os.path.abspath(__file__))
ipc = IPCUtil('ipc_sender', 'ipc_sender')
-def send_image(image):
- identifier = b'\x01'
-
+def send_image(identifier, image):
height, width, channels = image.shape
width_bytes = width.to_bytes(4, byteorder='little')
@@ -40,80 +35,6 @@ def cal_box(inv_m, p):
return x, y
-def img_warp_back_inv_m(img, img_to, inv_m):
- h_up, w_up, c = img_to.shape
- mask = np.ones_like(img).astype(np.float32)
- inv_mask = cv2.warpAffine(mask, inv_m, (w_up, h_up))
- inv_img = cv2.warpAffine(img, inv_m, (w_up, h_up))
- mask_indices = inv_mask == 1
- if 4 == c:
- send_image(img_to)
- img_to[:, :, :3][mask_indices] = inv_img[mask_indices]
- else:
- img_to[inv_mask == 1] = inv_img[inv_mask == 1]
- cv2.imwrite('./full.png', img_to)
- return img_to
- # h_up, w_up, _ = img.shape
- # _, _, c = img_to.shape
- # # cv2.imwrite('./face.png', img)
- # tx = int(inv_m[0][2])
- # ty = int(inv_m[1][2])
- #
- # inv_m[0][2] = 0
- # inv_m[1][2] = 0
- #
- # p0 = cal_box(inv_m, (0, 0))
- # p1 = cal_box(inv_m, (w_up, 0))
- # p2 = cal_box(inv_m, (w_up, h_up))
- # p3 = cal_box(inv_m, (0, h_up))
- # lp = (min(p0[0], p3[0]), min(p0[1], p1[1]))
- # rp = (max(p2[0], p1[0]), min(p2[1], p3[1]))
- #
- # w_up = int(rp[0] - lp[0])
- # h_up = int(rp[1] - lp[1])
- #
- # # print(f'src_x:{w_up}, src_y:{h_up}')
- # inv_m[0][2] = 0
- # inv_m[1][2] = abs(lp[1])
- #
- # mask = np.ones_like(img, dtype=np.float32)
- # inv_mask = cv2.warpAffine(mask, inv_m, (w_up, h_up))
- # inv_img = cv2.warpAffine(img, inv_m, (w_up, h_up))
- #
- # if c == 4:
- # # img_to[30:h, 30:w][:, :, :3] = img
- # img_to[ty:(h_up + ty), tx:(w_up + tx)][:, :, :3][inv_mask == 1] = inv_img[inv_mask == 1]
- # else:
- # img_to[inv_mask == 1] = inv_img[inv_mask == 1]
- #
- # # cv2.imwrite('./full1.png', img_to)
- # return img_to
-
-
-def render_image(context, frame):
- res_frame, idx, type_ = frame
-
- if type_ == 0:
- combine_frame = context.frame_list_cycle[idx]
- else:
- bbox = context.coord_list_cycle[idx]
- combine_frame = copy.deepcopy(context.frame_list_cycle[idx])
- af = context.align_frames[idx]
- inv_m = context.inv_m_frames[idx]
- y1, y2, x1, x2 = bbox
- try:
- t = time.perf_counter()
- res_frame = cv2.resize(res_frame.astype(np.uint8), (x2 - x1, y2 - y1))
- af[y1:y2, x1:x2] = res_frame
- combine_frame = img_warp_back_inv_m(af, combine_frame, inv_m)
- except Exception as e:
- logging.error(f'resize error:{e}')
- return
-
- image = combine_frame
- return image
-
-
class PyGameUI:
def __init__(self):
self._human_context = None
diff --git a/utils/__init__.py b/utils/__init__.py
index 2f3ca11..c3c95a3 100644
--- a/utils/__init__.py
+++ b/utils/__init__.py
@@ -5,4 +5,5 @@ from .sync_queue import SyncQueue
from .utils import mirror_index, load_model, get_device, load_avatar, config_logging
from .utils import read_image, object_stop
from .utils import load_avatar_from_processed, load_avatar_from_256_processed
+from .utils import render_image
from .audio_utils import melspectrogram, save_wav
diff --git a/utils/utils.py b/utils/utils.py
index eb307ba..23b117e 100644
--- a/utils/utils.py
+++ b/utils/utils.py
@@ -1,4 +1,5 @@
#encoding = utf8
+import copy
import glob
import logging
import os
@@ -274,3 +275,38 @@ def config_logging(file_name: str, console_level: int = logging.INFO, file_level
def object_stop(obj):
if obj is not None:
obj.stop()
+
+
+def img_warp_back_inv_m(img, img_to, inv_m):
+ h_up, w_up, c = img_to.shape
+ mask = np.ones_like(img).astype(np.float32)
+ inv_mask = cv2.warpAffine(mask, inv_m, (w_up, h_up))
+ inv_img = cv2.warpAffine(img, inv_m, (w_up, h_up))
+ mask_indices = inv_mask == 1
+ if 4 == c:
+ img_to[:, :, :3][mask_indices] = inv_img[mask_indices]
+ else:
+ img_to[inv_mask == 1] = inv_img[inv_mask == 1]
+ return img_to
+
+
+def render_image(context, frame):
+ res_frame, idx, type_ = frame
+
+ if type_ == 0:
+ combine_frame = context.frame_list_cycle[idx]
+ else:
+ bbox = context.coord_list_cycle[idx]
+ combine_frame = copy.deepcopy(context.frame_list_cycle[idx])
+ af = context.align_frames[idx]
+ inv_m = context.inv_m_frames[idx]
+ y1, y2, x1, x2 = bbox
+ try:
+ res_frame = cv2.resize(res_frame.astype(np.uint8), (x2 - x1, y2 - y1))
+ af[y1:y2, x1:x2] = res_frame
+ combine_frame = img_warp_back_inv_m(af, combine_frame, inv_m)
+ except Exception as e:
+ logging.error(f'resize error:{e}')
+ return None
+
+ return combine_frame