diff --git a/LICENSE b/LICENSE index 1d8c44b..2a663dd 100644 --- a/LICENSE +++ b/LICENSE @@ -1,6 +1,6 @@ MIT License -Copyright (c) 2021 Gunwoo Han +Copyright (c) 2021 Gunwoo Han, yuyuyzl Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal diff --git a/README.md b/README.md index 3ccf518..4ad5bb0 100644 --- a/README.md +++ b/README.md @@ -9,6 +9,8 @@ Fork自 https://github.com/GunwooHan/EasyVtuber 并且省略了ifacialmocap pc端，通过UDP直连的方式使ios面捕刷新率达到30fps，解决了面捕刷新率的瓶颈 最后，将EasyVtuber中使用的OBS虚拟摄像头方案切换为Unity Capture，解锁RGBA输出能力，无需绿背即可直接使用 +[视频介绍和安装说明](https://www.bilibili.com/video/BV1uu411r7DR) + ## Requirements ### 硬件 @@ -87,6 +89,20 @@ https://github.com/pkhungurn/talking-head-anime-2-demo#download-the-model ### 输入输出设备 +#### OBS Virtual Camera + +目前更推荐这个方案，UnityCapture存在未查明的性能瓶颈 +如果你选择自己进行抠像你可以直接输出到obs，如果你需要RGBA支持则需要额外使用一个Shader +下载并安装StreamFX https://github.com/Xaymar/obs-StreamFX +下载Shader（感谢树根的协助） https://github.com/shugen002/shader/blob/master/merge%20alpha2.hlsl +之后，使用`--alpha_split`参数运行 +![alpha split](assets/alphasplit.png) +你会看到这样的输出画面，透明通道单独使用灰度方式发送了 +之后对OBS中的视频采集设备添加滤镜-着色器-选择你下载的`merge alpha2.hlsl`-关闭 +这样透明通道就应用回左边的图像了 +你可能需要手动调整一下裁剪把右侧的无用画面切掉 +(看不到着色器滤镜的话就是StreamFX没装好或者OBS不是最新版) + #### UnityCapture 如果需要使用透明通道输出，参考 https://github.com/schellingb/UnityCapture#installation 安装好UnityCapture diff --git a/args.py b/args.py new file mode 100644 index 0000000..368d1d1 --- /dev/null +++ b/args.py @@ -0,0 +1,41 @@ +import argparse +import re + +def convert_to_byte(size): + result = re.search('(\d+\.?\d*)(b|kb|mb|gb|tb)', size.lower()) + if (result and result.groups()): + unit = result.groups()[1] + amount = float(result.groups()[0]) + index = ['b', 'kb', 'mb', 'gb', 'tb'].index(unit) + return amount * pow(1024, index) + raise ValueError("Invalid size provided, value is " + size) + +parser = argparse.ArgumentParser() +parser.add_argument('--debug', action='store_true') +parser.add_argument('--extend_movement', type=float) +parser.add_argument('--input', type=str, default='cam') +parser.add_argument('--character', type=str, default='y') +parser.add_argument('--output_dir', type=str) +parser.add_argument('--output_webcam', type=str) +parser.add_argument('--output_size', type=str, default='256x256') +parser.add_argument('--debug_input', action='store_true') +parser.add_argument('--perf', action='store_true') +parser.add_argument('--skip_model', action='store_true') +parser.add_argument('--ifm', type=str) +parser.add_argument('--anime4k', action='store_true') +parser.add_argument('--alpha_split', action='store_true') +parser.add_argument('--cache', type=str, default='256mb') +parser.add_argument('--gpu_cache', type=str, default='512mb') +parser.add_argument('--simplify', type=int, default=1) +args = parser.parse_args() +args.output_w = int(args.output_size.split('x')[0]) +args.output_h = int(args.output_size.split('x')[1]) +if args.cache is not None: + args.max_cache_len=int(convert_to_byte(args.cache)/262144) +else: + args.max_cache_len=0 +if args.gpu_cache is not None: + args.max_gpu_cache_len=int(convert_to_byte(args.gpu_cache)/589824) +else: + args.max_gpu_cache_len=0 +if args.output_webcam is None and args.output_dir is None: args.debug = True \ No newline at end of file diff --git a/assets/alphasplit.png b/assets/alphasplit.png new file mode 100644 index 0000000..1d04552 Binary files /dev/null and b/assets/alphasplit.png differ diff --git a/launcher.py b/launcher.py index bbf534c..e94259e 100644 --- a/launcher.py +++ b/launcher.py @@ -6,20 +6,26 @@ import tkinter as tk from tkinter import ttk import json -try: - f = open('launcher.json') - args = json.load(f) - f.close() -except: - args = { +default_arg={ 'character': 'y', 'input': 2, 'output': 2, 'ifm': None, 'is_extend_movement': False, 'is_anime4k': False, + 'is_alpha_split': False, + 'cache_simplify': 1, + 'cache_size': 1, } +try: + f = open('launcher.json') + args = json.load(f) + default_arg.update(args) + f.close() +finally: + args = default_arg + p = None characterList = [] @@ -52,6 +58,25 @@ ttk.Label(launcher, text="iFacialMocap IP:Port").pack(fill='x', expand=True) ifm = tk.StringVar(value=args['ifm']) ttk.Entry(launcher, textvariable=ifm, state=False).pack(fill='x', expand=True) +ttk.Label(launcher, text="Facial Input Simplify").pack(fill='x', expand=True) +cache_simplify = tk.IntVar(value=args['cache_simplify']) +ttk.Radiobutton(launcher, text='Off', value=0, variable=cache_simplify).pack(fill='x', expand=True) +ttk.Radiobutton(launcher, text='Low', value=1, variable=cache_simplify).pack(fill='x', expand=True) +ttk.Radiobutton(launcher, text='Medium', value=2, variable=cache_simplify).pack(fill='x', expand=True) +ttk.Radiobutton(launcher, text='High', value=3, variable=cache_simplify).pack(fill='x', expand=True) +ttk.Radiobutton(launcher, text='Higher', value=4, variable=cache_simplify).pack(fill='x', expand=True) +ttk.Radiobutton(launcher, text='Highest', value=6, variable=cache_simplify).pack(fill='x', expand=True) +ttk.Radiobutton(launcher, text='Gaming', value=8, variable=cache_simplify).pack(fill='x', expand=True) + +ttk.Label(launcher, text="Cache Size (RAM+VRAM)").pack(fill='x', expand=True) +cache_size = tk.IntVar(value=args['cache_size']) +ttk.Radiobutton(launcher, text='Off', value=0, variable=cache_size).pack(fill='x', expand=True) +ttk.Radiobutton(launcher, text='256M+128M', value=1, variable=cache_size).pack(fill='x', expand=True) +ttk.Radiobutton(launcher, text='1GB+512M', value=2, variable=cache_size).pack(fill='x', expand=True) +ttk.Radiobutton(launcher, text='2GB+1GB', value=3, variable=cache_size).pack(fill='x', expand=True) +ttk.Radiobutton(launcher, text='4GB+2GB', value=4, variable=cache_size).pack(fill='x', expand=True) +ttk.Radiobutton(launcher, text='8GB+4GB', value=5, variable=cache_size).pack(fill='x', expand=True) + ttk.Label(launcher, text="Extra Options").pack(fill='x', expand=True) is_extend_movement = tk.BooleanVar(value=args['is_extend_movement']) ttk.Checkbutton(launcher, text='Extend Movement', variable=is_extend_movement).pack(fill='x', expand=True) @@ -59,6 +84,9 @@ ttk.Checkbutton(launcher, text='Extend Movement', variable=is_extend_movement).p is_anime4k = tk.BooleanVar(value=args['is_anime4k']) ttk.Checkbutton(launcher, text='Anime4K', variable=is_anime4k).pack(fill='x', expand=True) +is_alpha_split = tk.BooleanVar(value=args['is_alpha_split']) +ttk.Checkbutton(launcher, text='Alpha Split', variable=is_alpha_split).pack(fill='x', expand=True) + output = tk.IntVar(value=args['output']) ttk.Label(launcher, text="Output").pack(fill='x', expand=True) ttk.Radiobutton(launcher, text='Unity Capture', value=0, variable=output).pack(fill='x', expand=True) @@ -76,6 +104,9 @@ def launch(): 'ifm': ifm.get(), 'is_extend_movement': is_extend_movement.get(), 'is_anime4k': is_anime4k.get(), + 'is_alpha_split': is_alpha_split.get(), + 'cache_simplify': cache_simplify.get(), + 'cache_size': cache_size.get(), } f = open('launcher.json', mode='w') json.dump(args, f) @@ -111,12 +142,22 @@ def launch(): run_args.append('--debug') if args['is_anime4k']: run_args.append('--anime4k') + if args['is_alpha_split']: + run_args.append('--alpha_split') if args['is_extend_movement']: run_args.append('--extend_movement') run_args.append('1') + if args['cache_simplify'] is not None: + run_args.append('--simplify') + run_args.append(str(args['cache_simplify'])) + if args['cache_size'] is not None: + run_args.append('--cache') + run_args.append(['0b','256mb','1gb','2gb','4gb','8gb'][args['cache_size']]) + run_args.append('--gpu_cache') + run_args.append(['0b','128mb','512mb','1gb','2gb','4gb'][args['cache_size']]) run_args.append('--output_size') run_args.append('512x512') - + print(run_args) p = subprocess.Popen(run_args) launch_btn.config(text='Stop') diff --git a/main.py b/main.py index 6608641..36d7f49 100644 --- a/main.py +++ b/main.py @@ -1,5 +1,3 @@ -import argparse - import cv2 import torch import pyvirtualcam @@ -19,30 +17,32 @@ import queue import socket import time import math +import re +from collections import OrderedDict from multiprocessing import Value, Process, Queue -import pyanime4k from pyanime4k import ac from tha2.mocap.ifacialmocap_constants import * -parser = argparse.ArgumentParser() -parser.add_argument('--debug', action='store_true') -parser.add_argument('--extend_movement', type=float) -parser.add_argument('--input', type=str, default='cam') -parser.add_argument('--character', type=str, default='y') -parser.add_argument('--output_dir', type=str) -parser.add_argument('--output_webcam', type=str) -parser.add_argument('--output_size', type=str, default='256x256') -parser.add_argument('--debug_input', action='store_true') -parser.add_argument('--ifm', type=str) -parser.add_argument('--anime4k', action='store_true') -args = parser.parse_args() -args.output_w = int(args.output_size.split('x')[0]) -args.output_h = int(args.output_size.split('x')[1]) -if args.output_webcam is None and args.output_dir is None: args.debug = True +from args import args -device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu') +import collections + + +class FPS: + def __init__(self, avarageof=50): + self.frametimestamps = collections.deque(maxlen=avarageof) + + def __call__(self): + self.frametimestamps.append(time.time()) + if len(self.frametimestamps) > 1: + return len(self.frametimestamps) / (self.frametimestamps[-1] - self.frametimestamps[0]) + else: + return 0.0 + + +device = torch.device('cuda') if torch.cuda.is_available() and not args.skip_model else torch.device('cpu') def create_default_blender_data(): @@ -69,13 +69,17 @@ def create_default_blender_data(): return data -class ClientProcess(Process): +ifm_converter = tha2.poser.modes.mode_20_wx.IFacialMocapPoseConverter20() + + +class IFMClientProcess(Process): def __init__(self): super().__init__() self.queue = Queue() self.should_terminate = Value('b', False) self.address = args.ifm.split(':')[0] self.port = int(args.ifm.split(':')[1]) + self.ifm_fps_number = Value('f', 0.0) self.perf_time = 0 def run(self): @@ -92,6 +96,8 @@ class ClientProcess(Process): self.socket.setblocking(False) self.socket.bind(("", self.port)) self.socket.settimeout(0.1) + ifm_fps = FPS() + pre_socket_string = '' while True: if self.should_terminate.value: break @@ -104,14 +110,13 @@ class ClientProcess(Process): else: raise e socket_string = socket_bytes.decode("utf-8") + if args.debug and pre_socket_string != socket_string: + self.ifm_fps_number.value = ifm_fps() + pre_socket_string = socket_string # print(socket_string) - # blender_data = json.loads(socket_string) data = self.convert_from_blender_data(socket_string) - # cur_time = time.perf_counter() - # fps = 1 / (cur_time - self.perf_time) - # self.perf_time = cur_time - # print(fps) + try: self.queue.put_nowait(data) except queue.Full: @@ -148,12 +153,200 @@ class ClientProcess(Process): return data +class ModelClientProcess(Process): + def __init__(self, input_image): + super().__init__() + self.should_terminate = Value('b', False) + self.updated = Value('b', False) + self.data = None + self.input_image = input_image + self.output_queue = Queue() + self.input_queue = Queue() + self.model_fps_number = Value('f', 0.0) + self.gpu_fps_number = Value('f', 0.0) + self.cache_hit_ratio = Value('f', 0.0) + self.gpu_cache_hit_ratio= Value('f', 0.0) + + def run(self): + model = None + if not args.skip_model: + model = TalkingAnimeLight().to(device) + model = model.eval() + model = model + print("Pretrained Model Loaded") + + mouth_eye_vector = torch.empty(1, 27) + pose_vector = torch.empty(1, 3) + + input_image = self.input_image.to(device) + mouth_eye_vector = mouth_eye_vector.to(device) + pose_vector = pose_vector.to(device) + + model_cache = OrderedDict() + tot = 0 + hit = 0 + model_fps = FPS() + gpu_fps = FPS() + while True: + model_input = None + try: + while not self.input_queue.empty(): + model_input = self.input_queue.get_nowait() + except queue.Empty: + continue + if model_input is None: continue + simplify_arr = [1000] * ifm_converter.pose_size + if args.simplify >= 1: + simplify_arr = [200] * ifm_converter.pose_size + simplify_arr[ifm_converter.eye_wink_left_index] = 50 + simplify_arr[ifm_converter.eye_wink_right_index] = 50 + simplify_arr[ifm_converter.eye_happy_wink_left_index] = 50 + simplify_arr[ifm_converter.eye_happy_wink_right_index] = 50 + simplify_arr[ifm_converter.eye_surprised_left_index] = 30 + simplify_arr[ifm_converter.eye_surprised_right_index] = 30 + simplify_arr[ifm_converter.iris_rotation_x_index] = 25 + simplify_arr[ifm_converter.iris_rotation_y_index] = 25 + simplify_arr[ifm_converter.eye_raised_lower_eyelid_left_index] = 10 + simplify_arr[ifm_converter.eye_raised_lower_eyelid_right_index] = 10 + simplify_arr[ifm_converter.mouth_lowered_corner_left_index] = 5 + simplify_arr[ifm_converter.mouth_lowered_corner_right_index] = 5 + simplify_arr[ifm_converter.mouth_raised_corner_left_index] = 5 + simplify_arr[ifm_converter.mouth_raised_corner_right_index] = 5 + + if args.simplify >= 2: + simplify_arr[ifm_converter.head_x_index] = 100 + simplify_arr[ifm_converter.head_y_index] = 100 + simplify_arr[ifm_converter.eye_surprised_left_index] = 10 + simplify_arr[ifm_converter.eye_surprised_right_index] = 10 + model_input[ifm_converter.eye_wink_left_index - 12] += model_input[ + ifm_converter.eye_happy_wink_left_index - 12] + model_input[ifm_converter.eye_happy_wink_left_index - 12] = model_input[ + ifm_converter.eye_wink_left_index - 12] / 2 + model_input[ifm_converter.eye_wink_left_index - 12] = model_input[ + ifm_converter.eye_wink_left_index - 12] / 2 + model_input[ifm_converter.eye_wink_right_index - 12] += model_input[ + ifm_converter.eye_happy_wink_right_index - 12] + model_input[ifm_converter.eye_happy_wink_right_index - 12] = model_input[ + ifm_converter.eye_wink_right_index - 12] / 2 + model_input[ifm_converter.eye_wink_right_index - 12] = model_input[ + ifm_converter.eye_wink_right_index - 12] / 2 + + uosum = model_input[ifm_converter.mouth_uuu_index - 12] + \ + model_input[ifm_converter.mouth_ooo_index - 12] + model_input[ifm_converter.mouth_ooo_index - 12] = uosum + model_input[ifm_converter.mouth_uuu_index - 12] = 0 + is_open = (model_input[ifm_converter.mouth_aaa_index - 12] + model_input[ + ifm_converter.mouth_iii_index - 12] + uosum) > 0 + model_input[ifm_converter.mouth_lowered_corner_left_index - 12] = 0 + model_input[ifm_converter.mouth_lowered_corner_right_index - 12] = 0 + model_input[ifm_converter.mouth_raised_corner_left_index - 12] = 0.5 if is_open else 0 + model_input[ifm_converter.mouth_raised_corner_right_index - 12] = 0.5 if is_open else 0 + simplify_arr[ifm_converter.mouth_lowered_corner_left_index] = 0 + simplify_arr[ifm_converter.mouth_lowered_corner_right_index] = 0 + simplify_arr[ifm_converter.mouth_raised_corner_left_index] = 0 + simplify_arr[ifm_converter.mouth_raised_corner_right_index] = 0 + if args.simplify >= 3: + simplify_arr[ifm_converter.iris_rotation_x_index] = 20 + simplify_arr[ifm_converter.iris_rotation_y_index] = 20 + simplify_arr[ifm_converter.eye_wink_left_index] = 32 + simplify_arr[ifm_converter.eye_wink_right_index] = 32 + simplify_arr[ifm_converter.eye_happy_wink_left_index] = 32 + simplify_arr[ifm_converter.eye_happy_wink_right_index] = 32 + if args.simplify >= 4: + simplify_arr[ifm_converter.head_x_index] = 50 + simplify_arr[ifm_converter.head_y_index] = 50 + simplify_arr[ifm_converter.neck_z_index] = 100 + model_input[ifm_converter.eye_raised_lower_eyelid_left_index - 12] = 0 + model_input[ifm_converter.eye_raised_lower_eyelid_right_index - 12] = 0 + simplify_arr[ifm_converter.iris_rotation_x_index] = 10 + simplify_arr[ifm_converter.iris_rotation_y_index] = 10 + simplify_arr[ifm_converter.eye_wink_left_index] = 24 + simplify_arr[ifm_converter.eye_wink_right_index] = 24 + simplify_arr[ifm_converter.eye_happy_wink_left_index] = 24 + simplify_arr[ifm_converter.eye_happy_wink_right_index] = 24 + simplify_arr[ifm_converter.eye_surprised_left_index] = 8 + simplify_arr[ifm_converter.eye_surprised_right_index] = 8 + model_input[ifm_converter.eye_wink_left_index - 12] += model_input[ + ifm_converter.eye_wink_right_index - 12] + model_input[ifm_converter.eye_wink_right_index - 12] = model_input[ + ifm_converter.eye_wink_left_index - 12] / 2 + model_input[ifm_converter.eye_wink_left_index - 12] = model_input[ + ifm_converter.eye_wink_left_index - 12] / 2 + + model_input[ifm_converter.eye_surprised_left_index - 12] += model_input[ + ifm_converter.eye_surprised_right_index - 12] + model_input[ifm_converter.eye_surprised_right_index - 12] = model_input[ + ifm_converter.eye_surprised_left_index - 12] / 2 + model_input[ifm_converter.eye_surprised_left_index - 12] = model_input[ + ifm_converter.eye_surprised_left_index - 12] / 2 + + model_input[ifm_converter.eye_happy_wink_left_index - 12] += model_input[ + ifm_converter.eye_happy_wink_right_index - 12] + model_input[ifm_converter.eye_happy_wink_right_index - 12] = model_input[ + ifm_converter.eye_happy_wink_left_index - 12] / 2 + model_input[ifm_converter.eye_happy_wink_left_index - 12] = model_input[ + ifm_converter.eye_happy_wink_left_index - 12] / 2 + model_input[ifm_converter.mouth_aaa_index - 12] = min( + model_input[ifm_converter.mouth_aaa_index - 12] + + model_input[ifm_converter.mouth_ooo_index - 12] / 2 + + model_input[ifm_converter.mouth_iii_index - 12] / 2 + + model_input[ifm_converter.mouth_uuu_index - 12] / 2, 1 + ) + model_input[ifm_converter.mouth_ooo_index - 12] = 0 + model_input[ifm_converter.mouth_iii_index - 12] = 0 + model_input[ifm_converter.mouth_uuu_index - 12] = 0 + for i in range(4, args.simplify): + simplify_arr = [max(math.ceil(x * 0.8), 5) for x in simplify_arr] + for i in range(12, len(simplify_arr)): + if simplify_arr[i] > 0: + model_input[i - 12] = round(model_input[i - 12] * simplify_arr[i]) / simplify_arr[i] + input_hash = hash(tuple(model_input)) + cached = model_cache.get(input_hash) + tot += 1 + mouth_eye_vector_c=[0.0]*27 + if not cached is None: + self.output_queue.put_nowait(cached) + model_cache.move_to_end(input_hash) + hit += 1 + else: + if args.perf: + tic = time.perf_counter() + for i in range(27): + mouth_eye_vector[0, i] = model_input[i] + mouth_eye_vector_c[i]=model_input[i] + for i in range(3): + pose_vector[0, i] = model_input[i + 27] + if model is None: + output_image = input_image + else: + output_image = model(input_image, mouth_eye_vector, pose_vector, mouth_eye_vector_c,self.gpu_cache_hit_ratio) + if args.perf: + torch.cuda.synchronize() + print("model", (time.perf_counter() - tic) * 1000) + tic = time.perf_counter() + postprocessed_image = output_image.cpu() + if args.perf: + print("cpu()", (time.perf_counter() - tic) * 1000) + tic = time.perf_counter() + postprocessed_image = postprocessing_image(postprocessed_image) + if args.perf: + print("postprocess", (time.perf_counter() - tic) * 1000) + tic = time.perf_counter() + + self.output_queue.put_nowait(postprocessed_image) + if args.debug: + self.gpu_fps_number.value = gpu_fps() + if args.max_cache_len > 0: + model_cache[input_hash] = postprocessed_image + if len(model_cache) > args.max_cache_len: + model_cache.popitem(last=False) + if args.debug: + self.model_fps_number.value = model_fps() + self.cache_hit_ratio.value = hit / tot + + @torch.no_grad() def main(): - model = TalkingAnimeLight().to(device) - model = model.eval() - model = model - print("Pretrained Model Loaded") img = Image.open(f"character/{args.character}.png") wRatio = img.size[0] / 256 img = img.resize((256, int(img.size[1] / wRatio))) @@ -163,17 +356,16 @@ def main(): extra_image = np.array(img.crop((0, 256, img.size[0], img.size[1]))) print("Character Image Loaded:", args.character) - - ifm_converter = None cap = None + output_fps = FPS() + if not args.debug_input: if args.ifm is not None: - client_process = ClientProcess() + client_process = IFMClientProcess() client_process.daemon = True client_process.start() - ifm_converter = tha2.poser.modes.mode_20_wx.create_ifacialmocap_pose_converter() print("IFM Service Running:", args.ifm) else: @@ -193,9 +385,13 @@ def main(): if args.output_webcam: cam_scale = 1 + cam_width_scale = 1 if args.anime4k: cam_scale = 2 - cam = pyvirtualcam.Camera(width=args.output_w * cam_scale, height=args.output_h * cam_scale, fps=30, + if args.alpha_split: + cam_width_scale = 2 + cam = pyvirtualcam.Camera(width=args.output_w * cam_scale * cam_width_scale, height=args.output_h * cam_scale, + fps=60, backend=args.output_webcam, fmt= {'unitycapture': pyvirtualcam.PixelFormat.RGBA, 'obs': pyvirtualcam.PixelFormat.RGB}[ @@ -221,21 +417,15 @@ def main(): a.set_arguments(parameters) print("Anime4K Loaded") - mouth_eye_vector = torch.empty(1, 27) - pose_vector = torch.empty(1, 3) - - # input_image = input_image.half() - # mouth_eye_vector = mouth_eye_vector.half() - # pose_vector = pose_vector.half() - - input_image = input_image.to(device) - mouth_eye_vector = mouth_eye_vector.to(device) - pose_vector = pose_vector.to(device) position_vector = [0, 0, 0, 1] pose_queue = [] blender_data = create_default_blender_data() - tic=0 + + model_output = None + model_process = ModelClientProcess(input_image) + model_process.daemon = True + model_process.start() print("Ready. Close this console to exit.") @@ -244,21 +434,24 @@ def main(): # input_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) # results = facemesh.process(input_frame) + if args.perf: + tic = time.perf_counter() if args.debug_input: - mouth_eye_vector[0, :] = 0 - pose_vector[0, :] = 0 + mouth_eye_vector_c = [0.0] * 27 + pose_vector_c = [0.0] * 3 - mouth_eye_vector[0, 2] = math.sin(time.perf_counter()*3) - mouth_eye_vector[0, 3] = math.sin(time.perf_counter()*3) + mouth_eye_vector_c[2] = math.sin(time.perf_counter() * 3) + mouth_eye_vector_c[3] = math.sin(time.perf_counter() * 3) - mouth_eye_vector[0, 14] = 0 + mouth_eye_vector_c[14] = 0 - mouth_eye_vector[0, 25] = math.sin(time.perf_counter()*2.2)*0.2 - mouth_eye_vector[0, 26] = math.sin(time.perf_counter()*3.5)*0.8 + mouth_eye_vector_c[25] = math.sin(time.perf_counter() * 2.2) * 0.2 + mouth_eye_vector_c[26] = math.sin(time.perf_counter() * 3.5) * 0.8 + + pose_vector_c[0] = math.sin(time.perf_counter() * 1.1) + pose_vector_c[1] = math.sin(time.perf_counter() * 1.2) + pose_vector_c[2] = math.sin(time.perf_counter() * 1.5) - pose_vector[0, 0] = math.sin(time.perf_counter()*1.1) - pose_vector[0, 1] = math.sin(time.perf_counter()*1.2) - pose_vector[0, 2] = math.sin(time.perf_counter()*1.5) elif args.ifm is not None: @@ -292,11 +485,15 @@ def main(): # - ifacialmocap_pose[EYE_LOOK_DOWN_RIGHT] # + ifacialmocap_pose[EYE_LOOK_DOWN_LEFT]) / 2.0 / 0.75 - position_vector = blender_data[HEAD_BONE_QUAT] + mouth_eye_vector_c = [0.0] * 27 + pose_vector_c = [0.0] * 3 for i in range(12, 39): - mouth_eye_vector[0, i - 12] = ifacialmocap_pose_converted[i] + mouth_eye_vector_c[i - 12] = ifacialmocap_pose_converted[i] for i in range(39, 42): - pose_vector[0, i - 39] = ifacialmocap_pose_converted[i] + pose_vector_c[i - 39] = ifacialmocap_pose_converted[i] + + position_vector = blender_data[HEAD_BONE_QUAT] + else: ret, frame = cap.read() @@ -332,30 +529,57 @@ def main(): y_angle = np_pose[6] z_angle = np_pose[7] - mouth_eye_vector[0, :] = 0 - pose_vector[0, :] = 0 + mouth_eye_vector_c = [0.0] * 27 + pose_vector_c = [0.0] * 3 - mouth_eye_vector[0, 2] = eye_l_h_temp - mouth_eye_vector[0, 3] = eye_r_h_temp + mouth_eye_vector_c[2] = eye_l_h_temp + mouth_eye_vector_c[3] = eye_r_h_temp - mouth_eye_vector[0, 14] = mouth_ratio * 1.5 + mouth_eye_vector_c[14] = mouth_ratio * 1.5 - mouth_eye_vector[0, 25] = eye_y_ratio - mouth_eye_vector[0, 26] = eye_x_ratio + mouth_eye_vector_c[25] = eye_y_ratio + mouth_eye_vector_c[26] = eye_x_ratio - pose_vector[0, 0] = (x_angle - 1.5) * 1.6 - pose_vector[0, 1] = y_angle * 2.0 # temp weight - pose_vector[0, 2] = (z_angle + 1.5) * 2 # temp weight + pose_vector_c[0] = (x_angle - 1.5) * 1.6 + pose_vector_c[1] = y_angle * 2.0 # temp weight + pose_vector_c[2] = (z_angle + 1.5) * 2 # temp weight + + model_input_arr = mouth_eye_vector_c + model_input_arr.extend(pose_vector_c) + + model_process.input_queue.put_nowait(model_input_arr) + + has_model_output = 0 + try: + new_model_output = model_output + while not model_process.output_queue.empty(): + has_model_output += 1 + new_model_output = model_process.output_queue.get_nowait() + model_output = new_model_output + except queue.Empty: + pass + if model_output is None: + time.sleep(1) + continue + # print(has_model_output) + # should_output=should_output or has_model_output + # if not should_output: + # continue + + postprocessed_image = model_output + + if args.perf: + print('===') + print("input", time.perf_counter() - tic) + tic = time.perf_counter() - output_image = model(input_image, mouth_eye_vector, pose_vector) - postprocessed_image = postprocessing_image(output_image.cpu()) if extra_image is not None: postprocessed_image = cv2.vconcat([postprocessed_image, extra_image]) - k_scale=1 - rotate_angle=0 - dx=0 - dy=0 + k_scale = 1 + rotate_angle = 0 + dx = 0 + dy = 0 if args.extend_movement is not None: k_scale = position_vector[2] * math.sqrt(args.extend_movement) + 1 rotate_angle = -position_vector[0] * 40 * args.extend_movement @@ -370,6 +594,12 @@ def main(): rm, (args.output_w, args.output_h)) + if args.perf: + print("extendmovement", (time.perf_counter() - tic) * 1000) + tic = time.perf_counter() + + output_fps_number = output_fps() + if args.anime4k: alpha_channel = postprocessed_image[:, :, 3] alpha_channel = cv2.resize(alpha_channel, None, fx=2, fy=2) @@ -383,15 +613,40 @@ def main(): postprocessed_image = a.save_image_to_numpy() postprocessed_image = cv2.merge((postprocessed_image, alpha_channel)) postprocessed_image = cv2.cvtColor(postprocessed_image, cv2.COLOR_BGRA2RGBA) + if args.perf: + print("anime4k", (time.perf_counter() - tic) * 1000) + tic = time.perf_counter() + if args.alpha_split: + alpha_image = cv2.merge( + [postprocessed_image[:, :, 3], postprocessed_image[:, :, 3], postprocessed_image[:, :, 3]]) + alpha_image = cv2.cvtColor(alpha_image, cv2.COLOR_RGB2RGBA) + postprocessed_image = cv2.hconcat([postprocessed_image, alpha_image]) if args.debug: output_frame = cv2.cvtColor(postprocessed_image, cv2.COLOR_RGBA2BGRA) # resized_frame = cv2.resize(output_frame, (np.min(debug_image.shape[:2]), np.min(debug_image.shape[:2]))) # output_frame = np.concatenate([debug_image, resized_frame], axis=1) - toc = time.perf_counter() - fps = 1 / (toc - tic) - tic = toc - cv2.putText(output_frame, str('%.1f' % fps),(0,16),cv2.FONT_HERSHEY_PLAIN,1,(0,255,0),1) + cv2.putText(output_frame, str('OUT_FPS:%.1f' % output_fps_number), (0, 16), cv2.FONT_HERSHEY_PLAIN, 1, + (0, 255, 0), 1) + if args.max_cache_len > 0: + cv2.putText(output_frame, str( + 'GPU_FPS:%.1f / %.1f' % (model_process.model_fps_number.value, model_process.gpu_fps_number.value)), + (0, 32), + cv2.FONT_HERSHEY_PLAIN, 1, (0, 255, 0), 1) + else: + cv2.putText(output_frame, str( + 'GPU_FPS:%.1f' % (model_process.model_fps_number.value)), + (0, 32), + cv2.FONT_HERSHEY_PLAIN, 1, (0, 255, 0), 1) + if args.ifm is not None: + cv2.putText(output_frame, str('IFM_FPS:%.1f' % client_process.ifm_fps_number.value), (0, 48), + cv2.FONT_HERSHEY_PLAIN, 1, (0, 255, 0), 1) + if args.max_cache_len > 0: + cv2.putText(output_frame, str('MEMCACHED:%.1f%%' % (model_process.cache_hit_ratio.value * 100)), (0, 64), + cv2.FONT_HERSHEY_PLAIN, 1, (0, 255, 0), 1) + if args.max_gpu_cache_len > 0: + cv2.putText(output_frame, str('GPUCACHED:%.1f%%' % (model_process.gpu_cache_hit_ratio.value * 100)), (0, 80), + cv2.FONT_HERSHEY_PLAIN, 1, (0, 255, 0), 1) cv2.imshow("frame", output_frame) # cv2.imshow("camera", debug_image) cv2.waitKey(1) @@ -404,7 +659,9 @@ def main(): result_image = cv2.cvtColor(result_image, cv2.COLOR_RGBA2RGB) cam.send(result_image) cam.sleep_until_next_frame() - + if args.perf: + print("output", (time.perf_counter() - tic) * 1000) + tic = time.perf_counter() if __name__ == '__main__': diff --git a/models.py b/models.py index bc17477..09f415b 100644 --- a/models.py +++ b/models.py @@ -1,7 +1,12 @@ +import time + import torch.nn as nn from tha2.poser.modes.mode_20 import load_face_morpher, load_face_rotater, load_combiner +from args import args + +from collections import OrderedDict class TalkingAnimeLight(nn.Module): def __init__(self): @@ -9,13 +14,40 @@ class TalkingAnimeLight(nn.Module): self.face_morpher = load_face_morpher('pretrained/face_morpher.pt') self.two_algo_face_rotator = load_face_rotater('pretrained/two_algo_face_rotator.pt') self.combiner = load_combiner('pretrained/combiner.pt') + self.face_cache = OrderedDict() + self.tot = 0 + self.hit = 0 - def forward(self, image, mouth_eye_vector, pose_vector): + def forward(self, image, mouth_eye_vector, pose_vector, mouth_eye_vector_c, ratio=None): x = image.clone() - mouth_eye_morp_image = self.face_morpher(image[:, :, 32:224, 32:224], mouth_eye_vector) + if args.perf: + tic=time.perf_counter() + input_hash = hash(tuple(mouth_eye_vector_c)) + cached = self.face_cache.get(input_hash) + self.tot+=1 + if cached is None: + mouth_eye_morp_image = self.face_morpher(image[:, :, 32:224, 32:224], mouth_eye_vector) + self.face_cache[input_hash]=mouth_eye_morp_image.detach() + if len(self.face_cache) > args.max_gpu_cache_len: + self.face_cache.popitem(last=False) + else: + self.hit+=1 + mouth_eye_morp_image=cached + self.face_cache.move_to_end(input_hash) + if args.debug and ratio is not None: + ratio.value=self.hit / self.tot + if args.perf: + print(" - face_morpher",(time.perf_counter()-tic)*1000) + tic=time.perf_counter() x[:, :, 32:224, 32:224] = mouth_eye_morp_image rotate_image = self.two_algo_face_rotator(x, pose_vector)[:2] + if args.perf: + print(" - rotator",(time.perf_counter()-tic)*1000) + tic=time.perf_counter() output_image = self.combiner(rotate_image[0], rotate_image[1], pose_vector) + if args.perf: + print(" - combiner",(time.perf_counter()-tic)*1000) + tic=time.perf_counter() return output_image diff --git a/test.py b/test.py index 51e7d3f..da42f49 100644 --- a/test.py +++ b/test.py @@ -1,27 +1,65 @@ -from pyanime4k import ac -import pyanime4k - -parameters = ac.Parameters() -# enable HDN for ACNet -parameters.HDN = True - -a = ac.AC( - type=ac.ProcessorType.CPU_ACNet, -) -a.set_arguments(parameters) +# from pyanime4k import ac +# import pyanime4k +# +# parameters = ac.Parameters() +# # enable HDN for ACNet +# parameters.HDN = True +# +# a = ac.AC( +# type=ac.ProcessorType.CPU_ACNet, +# ) +# a.set_arguments(parameters) import cv2 +# +# img = cv2.imread("character/test41.png") +# +# a.load_image_from_numpy(img,input_type=ac.AC_INPUT_BGR) +# +# a.get_processor_info() +# +# # start processing +# a.process() +# +# a.show_image() -img = cv2.imread("character/test41.png") - -a.load_image_from_numpy(img,input_type=ac.AC_INPUT_BGR) - -a.get_processor_info() - -# start processing -a.process() - -a.show_image() - +# import numpy as np +# import zlib +# import bz2 +# import lzma +# import time +# +# postprocessed_image=np.load('out.npy') +# tic=time.perf_counter() +# res=None +# for i in range(100): +# res=zlib.compress(postprocessed_image) +# print("zlib", len(res),(time.perf_counter() - tic) * 1000) +# tic = time.perf_counter() +# for i in range(100): +# zlib.decompress(res) +# print("zlib", len(res),(time.perf_counter() - tic) * 1000) +# tic = time.perf_counter() +# for i in range(100): +# res=bz2.compress(postprocessed_image) +# print("bz2",len(res), (time.perf_counter() - tic) * 1000) +# tic = time.perf_counter() +# for i in range(100): +# bz2.decompress(res) +# print("bz2", len(res),(time.perf_counter() - tic) * 1000) +# tic = time.perf_counter() +# for i in range(100): +# res=lzma.compress(postprocessed_image) +# print("lzma",len(res), (time.perf_counter() - tic) * 1000) +# tic = time.perf_counter() +# for i in range(100): +# lzma.decompress(res) +# print("lzma", len(res),(time.perf_counter() - tic) * 1000) +# tic = time.perf_counter() +# +# +# output_frame = cv2.cvtColor(postprocessed_image, cv2.COLOR_RGBA2BGRA) +# cv2.imshow("frame", output_frame) +# cv2.waitKey(100000) # preview upscaled image # img=a.save_image_to_numpy() # @@ -42,4 +80,117 @@ a.show_image() # [1,2,3,4], # ] # -# print(cv2.merge((np.array(rgb),np.array(a)))) \ No newline at end of file +# print(cv2.merge((np.array(rgb),np.array(a)))) + + +from multiprocessing import Value, Process, Queue +import queue +import socket +import errno +import time +import tha2.poser.modes.mode_20_wx +from tha2.mocap.ifacialmocap_constants import * + +ifm='192.168.31.182:49983' + + + +class ClientProcess(Process): + def __init__(self): + super().__init__() + self.queue = Queue() + self.should_terminate = Value('b', False) + self.address = ifm.split(':')[0] + self.port = int(ifm.split(':')[1]) + self.perf_time = 0 + self.ifm_converter = tha2.poser.modes.mode_20_wx.create_ifacialmocap_pose_converter() + self.a_min=None + self.a_max=None + + def run(self): + + udpClntSock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) + + data = "iFacialMocap_sahuasouryya9218sauhuiayeta91555dy3719" + + data = data.encode('utf-8') + + udpClntSock.sendto(data, (self.address, self.port)) + + self.socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) + self.socket.setblocking(False) + self.socket.bind(("", self.port)) + self.socket.settimeout(0.1) + while True: + if self.should_terminate.value: + break + try: + socket_bytes = self.socket.recv(8192) + except socket.error as e: + err = e.args[0] + if err == errno.EAGAIN or err == errno.EWOULDBLOCK or err == 'timed out': + continue + else: + raise e + socket_string = socket_bytes.decode("utf-8") + # print(socket_string) + + # blender_data = json.loads(socket_string) + data = self.convert_from_blender_data(socket_string) + data=self.ifm_converter.convert(data) + # print(data) + if(self.a_max is None): + self.a_max=[x for x in data] + else: + self.a_max=[max(data[i],self.a_max[i]) for i in range(len(data))] + if(self.a_min is None): + self.a_min=[x for x in data] + else: + self.a_min=[min(data[i],self.a_min[i]) for i in range(len(data))] + print([[self.a_min[i],self.a_max[i]] for i in range(len(data))]) + cur_time = time.perf_counter() + fps = 1 / (cur_time - self.perf_time) + self.perf_time = cur_time + # print(fps) + try: + self.queue.put_nowait(data) + except queue.Full: + pass + self.queue.close() + self.socket.close() + @staticmethod + def convert_from_blender_data(blender_data): + data = {} + + for item in blender_data.split('|'): + if item.find('#') != -1: + k, arr = item.split('#') + arr = [float(n) for n in arr.split(',')] + data[k.replace("_L", "Left").replace("_R", "Right")] = arr + elif item.find('-') != -1: + k, v = item.split("-") + data[k.replace("_L", "Left").replace("_R", "Right")] = float(v) / 100 + + to_rad = 57.3 + data[HEAD_BONE_X] = data["=head"][0] / to_rad + data[HEAD_BONE_Y] = data["=head"][1] / to_rad + data[HEAD_BONE_Z] = data["=head"][2] / to_rad + data[HEAD_BONE_QUAT] = [data["=head"][3], data["=head"][4], data["=head"][5], 1] + # print(data[HEAD_BONE_QUAT][2],min(data[EYE_BLINK_LEFT],data[EYE_BLINK_RIGHT])) + data[RIGHT_EYE_BONE_X] = data["rightEye"][0] / to_rad + data[RIGHT_EYE_BONE_Y] = data["rightEye"][1] / to_rad + data[RIGHT_EYE_BONE_Z] = data["rightEye"][2] / to_rad + data[LEFT_EYE_BONE_X] = data["leftEye"][0] / to_rad + data[LEFT_EYE_BONE_Y] = data["leftEye"][1] / to_rad + data[LEFT_EYE_BONE_Z] = data["leftEye"][2] / to_rad + + return data + + +if __name__ == '__main__': + client_process = ClientProcess() + client_process.daemon = True + client_process.start() + + while True: + pass \ No newline at end of file diff --git a/tha2/poser/modes/mode_20_wx.py b/tha2/poser/modes/mode_20_wx.py index 9d8955c..cbbfdf2 100644 --- a/tha2/poser/modes/mode_20_wx.py +++ b/tha2/poser/modes/mode_20_wx.py @@ -86,7 +86,7 @@ class IFacialMocapPoseConverter20(IFacialMocapPoseConverter): self.eyebrow_serious_left_index = pose_parameters.get_parameter_index("eyebrow_serious_left") self.eyebrow_serious_right_index = pose_parameters.get_parameter_index("eyebrow_serious_right") - self.eye_surprised_left_index = pose_parameters.get_parameter_index("eye_surprised_left") + self.eye_surprised_left_index = pose_parameters.get_parameter_index("eye_surprised_left") # 0 self.eye_surprised_right_index = pose_parameters.get_parameter_index("eye_surprised_right") self.eye_wink_left_index = pose_parameters.get_parameter_index("eye_wink_left") self.eye_wink_right_index = pose_parameters.get_parameter_index("eye_wink_right") @@ -95,13 +95,13 @@ class IFacialMocapPoseConverter20(IFacialMocapPoseConverter): self.eye_relaxed_left_index = pose_parameters.get_parameter_index("eye_relaxed_left") self.eye_relaxed_right_index = pose_parameters.get_parameter_index("eye_relaxed_right") self.eye_raised_lower_eyelid_left_index = pose_parameters.get_parameter_index("eye_raised_lower_eyelid_left") - self.eye_raised_lower_eyelid_right_index = pose_parameters.get_parameter_index("eye_raised_lower_eyelid_right") + self.eye_raised_lower_eyelid_right_index = pose_parameters.get_parameter_index("eye_raised_lower_eyelid_right") # 9 - self.iris_small_left_index = pose_parameters.get_parameter_index("iris_small_left") + self.iris_small_left_index = pose_parameters.get_parameter_index("iris_small_left")#10 self.iris_small_right_index = pose_parameters.get_parameter_index("iris_small_right") self.iris_rotation_x_index = pose_parameters.get_parameter_index("iris_rotation_x") - self.iris_rotation_y_index = pose_parameters.get_parameter_index("iris_rotation_y") + self.iris_rotation_y_index = pose_parameters.get_parameter_index("iris_rotation_y")#13 self.head_x_index = pose_parameters.get_parameter_index("head_x") self.head_y_index = pose_parameters.get_parameter_index("head_y") @@ -270,9 +270,11 @@ class IFacialMocapPoseConverter20(IFacialMocapPoseConverter): pose[self.mouth_iii_index] = restricted_decomp[1] mouth_funnel_denom = self.args.mouth_funnel_max_value - self.args.mouth_funnel_min_value ooo_alpha = clamp((mouth_funnel - self.args.mouth_funnel_min_value) / mouth_funnel_denom, 0.0, 1.0) + # ooo_alpha = 1 uo_value = clamp(restricted_decomp[2] + restricted_decomp[3], 0.0, 1.0) pose[self.mouth_uuu_index] = uo_value * (1.0 - ooo_alpha) pose[self.mouth_ooo_index] = uo_value * ooo_alpha + # print(pose[self.mouth_aaa_index:self.mouth_ooo_index+1]) # pose[28] = restricted_decomp[2] # pose[30] = restricted_decomp[3]