Merge pull request #1 from yuyuyzl/dev

Dev

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

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  

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

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')
 

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__':

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
 
 

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))))
+# 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

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]