diff --git a/.gitignore b/.gitignore index b983db8..3f5ecf6 100644 --- a/.gitignore +++ b/.gitignore @@ -1,7 +1,7 @@ pretrained/** !pretrained/PUT_MODEL_HERE character -src +!src !tha4/src tha4/distiller-ui-doc tha4/docs diff --git a/args.py b/args.py deleted file mode 100644 index ab6ae1e..0000000 --- a/args.py +++ /dev/null @@ -1,85 +0,0 @@ -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('--eyebrow', 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='512x512') -parser.add_argument('--model', type=str, default='standard_float') -parser.add_argument('--debug_input', action='store_true') -parser.add_argument('--mouse_input', type=str) -parser.add_argument('--perf', type=str) -parser.add_argument('--skip_model', action='store_true') -parser.add_argument('--ifm', type=str) -parser.add_argument('--osf', type=str) -parser.add_argument('--anime4k', action='store_true') -parser.add_argument('--alpha_split', action='store_true') -parser.add_argument('--bongo', 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) - -parser.add_argument('--device_id', type=int, default=0) -parser.add_argument('--use_tensorrt', action='store_true') -parser.add_argument('--use_tha4', action='store_true') -parser.add_argument('--use_tha4_student', action='store_true') -parser.add_argument('--tha4_student_model', type=str, default=None) - -parser.add_argument('--model_seperable', action='store_true') -parser.add_argument('--model_half', action='store_true') -parser.add_argument('--model_cache', action='store_true') -parser.add_argument('--model_vram_cache', action='store_true') -# parser.add_argument('--model_cache_size', type=float, default=1.0) - -parser.add_argument('--use_interpolation', action='store_true') -parser.add_argument('--interpolation_scale', type=int, default=3) -parser.add_argument('--interpolation_half', action='store_true') - -parser.add_argument('--use_cacher', action='store_true') -parser.add_argument('--cacher_quality', type=int, default=85) -# parser.add_argument('--cacher_ram_size', type=float, default=1.0) - -parser.add_argument('--use_sr', action='store_true') -parser.add_argument('--sr_x4', action='store_true') -parser.add_argument('--sr_half', action='store_true') -parser.add_argument('--sr_noise', type=int, default=1) - -parser.add_argument('--frame_rate_limit', type=int, default=30) - -parser.add_argument('--alpha_clean', 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.cache is not None: - args.max_cache_len = convert_to_byte(args.cache) / pow(1024, 3) # In gigabytes -else: - args.max_cache_len = 0 -if args.max_cache_len == 0: - args.use_cacher = False -if args.gpu_cache is not None: - args.max_gpu_cache_len = convert_to_byte(args.gpu_cache) / pow(1024, 3) # In gigabytes -else: - args.max_gpu_cache_len = 0 -if args.output_webcam is None and args.output_dir is None: args.debug = True - -args.model_output_size = 1024 if args.use_sr else 512 - -if args.use_sr and args.use_interpolation: - args.use_interpolation = False diff --git a/ezvtb_rt_interface.py b/ezvtb_rt_interface.py deleted file mode 100644 index 25d2aa8..0000000 --- a/ezvtb_rt_interface.py +++ /dev/null @@ -1,198 +0,0 @@ -import os -import sys -import subprocess -import urllib.request -import shutil - -dir_path = os.path.dirname(os.path.realpath(__file__)) -ezvtb_path = os.path.join(dir_path, 'ezvtuber-rt') -ezvtb_main_path = os.path.join(dir_path, 'ezvtuber-rt-main') - -project_path = '' -if os.path.exists(ezvtb_path): - project_path = ezvtb_path -else: - project_path = ezvtb_main_path - -if project_path not in sys.path: - sys.path.append(project_path) - -# Initialize model data path to point to data/models directory -import ezvtb_rt -models_path = os.path.join(dir_path, 'data', 'models') -ezvtb_rt.init_model_path(models_path) - -def get_core( - #Device setting - device_id:int = 0, - use_tensorrt:bool = True, - #THA3 model setting - model_seperable:bool = True, - model_half:bool = True, #If using directml+half, there is small numerical error on Nvidia, and huge numerical error on AMD - model_cache_size:float = 1.0, #unit of GigaBytes, only works for tensorrt - model_use_eyebrow:bool = True, - #RIFE interpolation setting - use_interpolation:bool = True, - interpolation_scale:int = 2, - interpolation_half:bool = True, #If using directml+half, there is small numerical error on Nvidia, and huge numerical error on AMD - #Cacher setting - cacher_quality:int = 85, - cacher_ram_size:float = 2.0,#unit of GigaBytes - #SR setting - use_sr:bool = False, - sr_x4:bool = True, - sr_half:bool = True, - sr_noise:int = 1, - #THA4 model setting (NEW) - use_tha4: bool = False, # Use THA4 instead of THA3 - use_tha4_student: bool = False, # Use THA4 Student model - tha4_student_model: str = None # Student model name - ): - support_trt = False - if use_tensorrt: - print('Verifying TensorRT') - try: - from ezvtb_rt.trt_utils import cudaSetDevice - cudaSetDevice(device_id) - support_trt = True - except Exception: - support_trt = False - if not support_trt and use_tensorrt: - print('TensorRT option selected but not supported') - use_tensorrt = False - - if use_tensorrt: - os.environ['CUDA_DEVICE'] = str(device_id) - import pycuda.autoinit - print(f'Using device {pycuda.autoinit.device.name()}') - - # Determine model directory based on THA version - if use_tha4_student and tha4_student_model: - # THA4 Student: Use data/models/custom_tha4_models/{model_name} - tha_model_dir = os.path.join( - '.', 'data', 'models', 'custom_tha4_models', - tha4_student_model) - print(f'THA4 Student Mode - Model: {tha4_student_model}') - print(f'THA4 Student Mode - Model Path: {tha_model_dir}') - elif use_tha4: - # THA4: Use data/models/tha4 directory with precision selection - tha_model_dir = os.path.join( - '.', 'data', 'models', 'tha4', - 'fp16' if model_half else 'fp32') - print(f'THA4 Mode - Model Path: {tha_model_dir}') - else: - # THA3: Use data/models/tha3 directory - tha_model_dir = os.path.join( - '.', 'data', 'models', 'tha3', - 'seperable' if model_seperable else 'standard', - 'fp16' if model_half else 'fp32') - print(f'THA3 Mode - Model Path: {tha_model_dir}') - - rife_model_path = '' - if use_interpolation: - rife_model_path = os.path.join( - '.', 'data', 'models', 'rife_512', - f'x{interpolation_scale}', - 'fp16' if interpolation_half else 'fp32') - - sr_model_path = '' - if use_sr: - if sr_x4: - if sr_half: - sr_model_path = os.path.join( - '.', 'data', 'models', 'Real-ESRGAN', - 'exported_256_fp16') - else: - sr_model_path = os.path.join( - '.', 'data', 'models', 'Real-ESRGAN', - 'exported_256') - else: - if sr_half: - sr_model_path = os.path.join( - '.', 'data', 'models', 'waifu2x_upconv', - 'fp16', 'upconv_7', 'art', - f'noise{sr_noise}_scale2x') - else: - sr_model_path = os.path.join( - '.', 'data', 'models', 'waifu2x_upconv', - 'fp32', 'upconv_7', 'art', - f'noise{sr_noise}_scale2x') - - print(f'RIFE Path: {rife_model_path}') - print(f'SR Path: {sr_model_path}') - - core = None - if use_tensorrt: - from ezvtb_rt.core import CoreTRT - if use_tha4_student and tha4_student_model: - core = CoreTRT( - tha_model_version='v4_student', - tha_model_name=tha4_student_model, - rife_model_enable=use_interpolation, - rife_model_scale=interpolation_scale, - rife_model_fp16=interpolation_half, - sr_model_enable=use_sr, - sr_model_scale=4 if sr_x4 else 2, - sr_model_noise=sr_noise, - sr_model_fp16=sr_half, - vram_cache_size=model_cache_size, - cache_max_giga=cacher_ram_size, - use_eyebrow=model_use_eyebrow) - else: - core = CoreTRT( - tha_model_version='v4' if use_tha4 else 'v3', - tha_model_seperable=model_seperable, - tha_model_fp16=model_half, - rife_model_enable=use_interpolation, - rife_model_scale=interpolation_scale, - rife_model_fp16=interpolation_half, - sr_model_enable=use_sr, - sr_model_scale=4 if sr_x4 else 2, - sr_model_noise=sr_noise, - sr_model_fp16=sr_half, - vram_cache_size=model_cache_size, - cache_max_giga=cacher_ram_size, - use_eyebrow=model_use_eyebrow) - else: - from ezvtb_rt.core_ort import CoreORT - if use_tha4_student and tha4_student_model: - core = CoreORT( - tha_model_version='v4_student', - tha_model_name=tha4_student_model, - rife_model_enable=use_interpolation, - rife_model_scale=interpolation_scale, - rife_model_fp16=interpolation_half, - sr_model_enable=use_sr, - sr_model_scale=4 if sr_x4 else 2, - sr_model_noise=sr_noise, - sr_model_fp16=sr_half, - cache_max_giga=cacher_ram_size, - use_eyebrow=model_use_eyebrow) - else: - core = CoreORT( - tha_model_version='v4' if use_tha4 else 'v3', - tha_model_seperable=model_seperable, - tha_model_fp16=model_half, - rife_model_enable=use_interpolation, - rife_model_scale=interpolation_scale, - rife_model_fp16=interpolation_half, - sr_model_enable=use_sr, - sr_model_scale=4 if sr_x4 else 2, - sr_model_noise=sr_noise, - sr_model_fp16=sr_half, - cache_max_giga=cacher_ram_size, - use_eyebrow=model_use_eyebrow) - - return core - -if __name__ == '__main__': - from ezvtb_rt.trt_utils import check_build_all_models, cudaSetDevice - device_id = 0 - import sys - if len(sys.argv) > 1: - device_id = int(sys.argv[1]) - try: - cudaSetDevice(device_id) - check_build_all_models('./data/models') - except: - pass diff --git a/ezvtuber-rt b/ezvtuber-rt index 1bc5fee..2ea780d 160000 --- a/ezvtuber-rt +++ b/ezvtuber-rt @@ -1 +1 @@ -Subproject commit 1bc5feeb09a12915fd86b6c125aee256b48676bf +Subproject commit 2ea780d884afb9740a9c23e6754bc46031eb1e1f diff --git a/launcher2.py b/launcher2.py index 5715feb..24f0a49 100644 --- a/launcher2.py +++ b/launcher2.py @@ -39,14 +39,12 @@ default_arg = { 'is_alpha_clean': False, 'is_eyebrow': False, 'cache_simplify': 'High', - 'cache_compression': 'High', 'ram_cache_size': '2gb', 'vram_cache_size': '2gb', 'model_select': 'seperable_half', 'interpolation': 'x2_half', 'frame_rate_limit': '30', 'sr': 'Off', - 'device_id': '0', 'use_tensorrt': False, 'preset': 'Medium' } @@ -249,10 +247,9 @@ class LauncherPanel(wx.Panel): addOption('osf', title='OpenSeeFace IP:Port', desc='输入OpenSeeFace连接使用的IP:端口号', type=2) addOption('output', title='Output', desc='选择输出目标', - choices=['Spout2', 'OBS VirtualCam', 'UnityCapture', 'Debug Output'], - mapping=[3, 1, 0, 2]) + choices=['Spout2', 'OBS VirtualCam', 'Debug Output'], + mapping=[0, 1, 2]) - addOption('device_id', title='GPU Device', desc='选择需要使用的计算设备(默认0)', choices=['0', '1', '2', '3']) addOption('use_tensorrt', title='TensorRT', desc='运行过预构建脚本后开启TensorRT可进一步提升性能\n（未构建成功时不可用，仅NVIDIA显卡支持）', type=1) @@ -287,9 +284,6 @@ class LauncherPanel(wx.Panel): addOption('cache_simplify', title='Input Simplify', desc='设置输入简化级别\n输入越简化，缓存命中率越高，动作越不平滑', choices=['Off', 'Low', 'Medium', 'High', 'Higher', 'Highest', 'Gaming']) - addOption('cache_compression', title='JPEG Compression', - desc='设置内存缓存压缩等级\n压缩等级越高，缓存命中率越高，输出质量越差', - choices=['Off', 'Low', 'Medium', 'High']) addOption('sr', title='SuperResolution', desc='选择使用的超分模型\n由于性能原因，real-esrgan会进行裁切', choices=['Off', 'anime4k_x2', 'waifu2x_x2_half', 'real-esrgan_x4_half', 'waifu2x_x2_full', @@ -332,13 +326,12 @@ class LauncherPanel(wx.Panel): self.optionDict['ram_cache_size'], self.optionDict['vram_cache_size'], self.optionDict['cache_simplify'], - self.optionDict['cache_compression'], ] presets = { - 'Low': [0, 1, 1, 5, 3], - 'Medium': [1, 1, 1, 4, 2], - 'High': [1, 2, 2, 2, 2], - 'Ultra': [3, 3, 3, 1, 1] + 'Low': [0, 1, 1, 5], + 'Medium': [1, 1, 1, 4], + 'High': [1, 2, 2, 2], + 'Ultra': [3, 3, 3, 1] } if s == 'Custom': @@ -418,18 +411,17 @@ class LauncherPanel(wx.Panel): p = None self.btnLaunch.SetLabelText("Save & Launch") else: - run_args = [sys.executable, 'main.py'] + run_args = [sys.executable, '-m', 'src.main'] if len(args['character']): run_args.append('--character') run_args.append(args['character']) if args['input'] == 0: if len(args['ifm']): - run_args.append('--ifm') + run_args.append('--ifm_input') run_args.append(args['ifm']) elif args['input'] == 1: - run_args.append('--input') - run_args.append('cam') + run_args.append('--cam_input') elif args['input'] == 2: run_args.append('--debug_input') elif args['input'] == 3: @@ -437,38 +429,30 @@ class LauncherPanel(wx.Panel): run_args.append('0,0,' + str(wx.GetDisplaySize().width) + ',' + str(wx.GetDisplaySize().height)) elif args['input'] == 4: if len(args['osf']): - run_args.append('--osf') + run_args.append('--osf_input') run_args.append(args['osf']) if args['output'] == 0: - run_args.append('--output_webcam') - run_args.append('unitycapture') + run_args.append('--output_spout2') elif args['output'] == 1: - run_args.append('--output_webcam') - run_args.append('obs') - elif args['output'] == 3: - run_args.append('--output_webcam') - run_args.append('spout') + run_args.append('--output_virtual_cam') elif args['output'] == 2: - run_args.append('--debug') + run_args.append('--output_debug') + 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['is_bongo']: run_args.append('--bongo') if args['is_alpha_clean']: run_args.append('--alpha_clean') if args['is_eyebrow']: run_args.append('--eyebrow') + if args['cache_simplify'] is not None: run_args.append('--simplify') run_args.append(str(cache_simplify_map[args['cache_simplify']])) - run_args.append('--cacher_quality') - run_args.append(str(cache_simplify_quality_map[args['cache_compression']])) - if args['cache_simplify'] != 'Off': - run_args.append('--use_cacher') if args['ram_cache_size'] is not None: run_args.append('--cache') run_args.append(args['ram_cache_size']) @@ -496,44 +480,37 @@ class LauncherPanel(wx.Panel): # Student model: tha4_student_{model_name} model_name = args['model_select'].replace( 'tha4_student_', '') - run_args.append('--use_tha4_student') - run_args.append('--tha4_student_model') + run_args.append('--model_version') + run_args.append('v4_student') + run_args.append('--model_name') run_args.append(model_name) elif 'tha4' in args['model_select']: - run_args.append('--use_tha4') - if 'half' in args['model_select']: - run_args.append('--model_half') + run_args.append('--model_version') + run_args.append('v4') else: - if 'seperable' in args['model_select']: - run_args.append('--model_seperable') - if 'half' in args['model_select']: - run_args.append('--model_half') + run_args.append('--model_version') + run_args.append('v3') + if 'seperable' in args['model_select']: + run_args.append('--model_seperable') + if 'half' in args['model_select']: + run_args.append('--model_half') if args['frame_rate_limit'] is not None: run_args.append('--frame_rate_limit') run_args.append(args['frame_rate_limit']) if args['sr'] is not None and args['sr'] != 'Off': + run_args.append('--use_sr') if 'anime4k' in args['sr']: - run_args.append('--anime4k') - else: - run_args.append('--use_sr') - if 'x4' in args['sr']: - run_args.append('--sr_x4') - if 'half' in args['sr']: - run_args.append('--sr_half') - - if args['device_id'] is not None: - run_args.append('--device_id') - run_args.append(args['device_id']) + run_args.append('--sr_a4k') + if 'x4' in args['sr']: + run_args.append('--sr_x4') + if 'half' in args['sr']: + run_args.append('--sr_half') if args['use_tensorrt'] is not None and args['use_tensorrt']: run_args.append('--use_tensorrt') - run_args.append('--model_cache') - run_args.append('--model_vram_cache') - run_args.append('--output_size') - run_args.append('512x512') print('Launched: ' + ' '.join(run_args)) p = subprocess.Popen(run_args) self.btnLaunch.SetLabelText('Stop') diff --git a/main.py b/main.py deleted file mode 100644 index f4d0f9b..0000000 --- a/main.py +++ /dev/null @@ -1,1019 +0,0 @@ -import struct - -import cv2 -import torch -import pyvirtualcam -import numpy as np -import mediapipe as mp -from PIL import Image - -import tha2.poser.modes.mode_20_wx -from pose import get_pose -from preprocess import resize_to_512_center, apply_color_curves -from utils import preprocessing_image, postprocessing_image -from ezvtb_rt_interface import get_core - -import errno -import json -import os -import queue -import socket -import time -import math -from pynput.mouse import Button, Controller -import re -from collections import OrderedDict -from multiprocessing import Value, Process, Queue, shared_memory - -from pyanime4k import ac - -from tha2.mocap.ifacialmocap_constants import * - -from args import args -import copy - -from tha3.util import torch_linear_to_srgb, resize_PIL_image, extract_PIL_image_from_filelike, \ - extract_pytorch_image_from_PIL_image - -import collections -from PySpout import SpoutSender -from OpenGL.GL import GL_RGBA - - -def convert_linear_to_srgb(image: torch.Tensor) -> torch.Tensor: - rgb_image = torch_linear_to_srgb(image[0:3, :, :]) - return torch.cat([rgb_image, image[3:4, :, :]], dim=0) - - -class FPS: - def __init__(self, avarageof=50): - self.frametimestamps = collections.deque(maxlen=avarageof) - - def __call__(self): - self.frametimestamps.append(time.time()) - return self.view() - - def view(self): - if len(self.frametimestamps) > 1: - return len(self.frametimestamps) / (self.frametimestamps[-1] - self.frametimestamps[0] + 1e-10) - else: - return 0.0 - - def last(self): - if len(self.frametimestamps): - return self.frametimestamps[-1] - else: - return time.time() - - -device = torch.device('cpu') - - -def create_default_blender_data(): - data = {} - - for blendshape_name in BLENDSHAPE_NAMES: - data[blendshape_name] = 0.0 - - data[HEAD_BONE_X] = 0.0 - data[HEAD_BONE_Y] = 0.0 - data[HEAD_BONE_Z] = 0.0 - data[HEAD_BONE_QUAT] = [0.0, 0.0, 0.0, 1.0] - - data[LEFT_EYE_BONE_X] = 0.0 - data[LEFT_EYE_BONE_Y] = 0.0 - data[LEFT_EYE_BONE_Z] = 0.0 - data[LEFT_EYE_BONE_QUAT] = [0.0, 0.0, 0.0, 1.0] - - data[RIGHT_EYE_BONE_X] = 0.0 - data[RIGHT_EYE_BONE_Y] = 0.0 - data[RIGHT_EYE_BONE_Z] = 0.0 - data[RIGHT_EYE_BONE_QUAT] = [0.0, 0.0, 0.0, 1.0] - - return data - - -class OSFClientProcess(Process): - def __init__(self): - super().__init__() - self.queue = Queue() - self.should_terminate = Value('b', False) - self.address = args.osf.split(':')[0] - self.port = int(args.osf.split(':')[1]) - self.ifm_fps_number = Value('f', 0.0) - self.perf_time = 0 - - def run(self): - self.socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) - self.socket.setblocking(False) - self.socket.bind(("", self.port)) - self.socket.settimeout(0.1) - ifm_fps = FPS() - 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") - osf_raw = (struct.unpack('=di2f2fB1f4f3f3f68f136f210f14f', socket_bytes)) - # print(osf_raw[432:]) - data = {} - OpenSeeDataIndex = [ - 'time', - 'id', - 'cameraResolutionW', - 'cameraResolutionH', - 'rightEyeOpen', - 'leftEyeOpen', - 'got3DPoints', - 'fit3DError', - 'rawQuaternionX', - 'rawQuaternionY', - 'rawQuaternionZ', - 'rawQuaternionW', - 'rawEulerX', - 'rawEulerY', - 'rawEulerZ', - 'translationY', - 'translationX', - 'translationZ', - ] - for i in range(len(OpenSeeDataIndex)): - data[OpenSeeDataIndex[i]] = osf_raw[i] - data['translationY'] *= -1 - data['translationZ'] *= -1 - data['rotationY'] = data['rawEulerY'] - 10 - data['rotationX'] = (-data['rawEulerX'] + 360) % 360 - 180 - data['rotationZ'] = (data['rawEulerZ'] - 90) - OpenSeeFeatureIndex = [ - 'EyeLeft', - 'EyeRight', - 'EyebrowSteepnessLeft', - 'EyebrowUpDownLeft', - 'EyebrowQuirkLeft', - 'EyebrowSteepnessRight', - 'EyebrowUpDownRight', - 'EyebrowQuirkRight', - 'MouthCornerUpDownLeft', - 'MouthCornerInOutLeft', - 'MouthCornerUpDownRight', - 'MouthCornerInOutRight', - 'MouthOpen', - 'MouthWide' - ] - - for i in range(68): - data['confidence' + str(i)] = osf_raw[i + 18] - for i in range(68): - data['pointsX' + str(i)] = osf_raw[i * 2 + 18 + 68] - data['pointsY' + str(i)] = osf_raw[i * 2 + 18 + 68 + 1] - for i in range(70): - data['points3DX' + str(i)] = osf_raw[i * 3 + 18 + 68 + 68 * 2] - data['points3DY' + str(i)] = osf_raw[i * 3 + 18 + 68 + 68 * 2 + 1] - data['points3DZ' + str(i)] = osf_raw[i * 3 + 18 + 68 + 68 * 2 + 2] - - for i in range(len(OpenSeeFeatureIndex)): - data[OpenSeeFeatureIndex[i]] = osf_raw[i + 432] - # print(data['rotationX'],data['rotationY'],data['rotationZ']) - - a = np.array([ - data['points3DX66'] - data['points3DX68'] + data['points3DX67'] - data['points3DX69'], - data['points3DY66'] - data['points3DY68'] + data['points3DY67'] - data['points3DY69'], - data['points3DZ66'] - data['points3DZ68'] + data['points3DZ67'] - data['points3DZ69'] - ]) - a = (a / np.linalg.norm(a)) - data['eyeRotationX'] = a[0] - data['eyeRotationY'] = a[1] - try: - self.queue.put_nowait(data) - except queue.Full: - pass - self.queue.close() - self.socket.close() - - -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): - - 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) - ifm_fps = FPS() - pre_socket_string = '' - 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") - 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) - - 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 - - -class MouseClientProcess(Process): - def __init__(self): - super().__init__() - self.queue = Queue() - - def run(self): - mouse = Controller() - posLimit = [int(x) for x in args.mouse_input.split(',')] - prev = { - 'eye_l_h_temp': 0, - 'eye_r_h_temp': 0, - 'mouth_ratio': 0, - 'eye_y_ratio': 0, - 'eye_x_ratio': 0, - 'x_angle': 0, - 'y_angle': 0, - 'z_angle': 0, - } - while True: - pos = mouse.position - # print(pos) - eye_limit = [0.8, 0.5] - head_eye_reduce = 0.6 - head_slowness = 0.2 - mouse_data = { - 'eye_l_h_temp': 0, - 'eye_r_h_temp': 0, - 'mouth_ratio': 0, - 'eye_y_ratio': np.interp(pos[1], [posLimit[1], posLimit[3]], [1, -1]) * eye_limit[1], - 'eye_x_ratio': np.interp(pos[0], [posLimit[0], posLimit[2]], [1, -1]) * eye_limit[0], - 'x_angle': np.interp(pos[1], [posLimit[1], posLimit[3]], [1, -1]), - 'y_angle': np.interp(pos[0], [posLimit[0], posLimit[2]], [1, -1]), - 'z_angle': 0, - } - mouse_data['x_angle'] = np.interp(head_slowness, [0, 1], [prev['x_angle'], mouse_data['x_angle']]) - mouse_data['y_angle'] = np.interp(head_slowness, [0, 1], [prev['y_angle'], mouse_data['y_angle']]) - mouse_data['eye_y_ratio'] -= mouse_data['x_angle'] * eye_limit[1] * head_eye_reduce - mouse_data['eye_x_ratio'] -= mouse_data['y_angle'] * eye_limit[0] * head_eye_reduce - if args.bongo: - mouse_data['y_angle'] += 0.05 - mouse_data['x_angle'] += 0.05 - prev = mouse_data - self.queue.put_nowait(mouse_data) - time.sleep(1 / 60) - - -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.shms = [shared_memory.SharedMemory(create=True, size=args.model_output_size * args.model_output_size * 4) - for _ in range(args.interpolation_scale if args.use_interpolation else 1)] - 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): - # Use unified ezvtb_rt interface for both THA3 and THA4 - self.model = get_core(device_id=args.device_id, - use_tensorrt=args.use_tensorrt, - use_tha4=args.use_tha4, - use_tha4_student=args.use_tha4_student, - tha4_student_model=args.tha4_student_model, - - model_seperable = args.model_seperable, - model_half=args.model_half, - model_cache_size=args.max_gpu_cache_len, - model_use_eyebrow=args.eyebrow, - - use_interpolation=args.use_interpolation, - interpolation_scale=args.interpolation_scale, - interpolation_half=args.interpolation_half, - - cacher_quality=args.cacher_quality, - cacher_ram_size=args.max_cache_len, - - use_sr=args.use_sr, - sr_half=args.sr_half, - sr_x4=args.sr_x4, - sr_noise=args.sr_noise) - self.model.setImage(self.input_image) - input_pose = np.zeros((1, 45), dtype=np.float32) - - self.shared_nps = [ - np.ndarray((args.model_output_size, args.model_output_size, 4), dtype=np.uint8, buffer=self.shms[i].buf) for - i in range(len(self.shms))] - - model_fps = FPS() - - last_process_time = time.time() - while True: - model_input = None - try: - model_input = self.input_queue.get(block=True, timeout=1.0) # Save cpu from infinite loop - while not self.input_queue.empty(): - model_input = self.input_queue.get_nowait() - except queue.Empty: - continue - if model_input is None: continue - frame_interval = (1 / args.frame_rate_limit) if not args.use_interpolation else ( - 1 / args.frame_rate_limit * args.interpolation_scale) - now = time.time() - if now < frame_interval + last_process_time - 0.005: continue # - last_process_time = now - 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] += model_input[ - ifm_converter.eye_happy_wink_left_index] - model_input[ifm_converter.eye_happy_wink_left_index] = model_input[ - ifm_converter.eye_wink_left_index] / 2 - model_input[ifm_converter.eye_wink_left_index] = model_input[ - ifm_converter.eye_wink_left_index] / 2 - model_input[ifm_converter.eye_wink_right_index] += model_input[ - ifm_converter.eye_happy_wink_right_index] - model_input[ifm_converter.eye_happy_wink_right_index] = model_input[ - ifm_converter.eye_wink_right_index] / 2 - model_input[ifm_converter.eye_wink_right_index] = model_input[ - ifm_converter.eye_wink_right_index] / 2 - - uosum = model_input[ifm_converter.mouth_uuu_index] + \ - model_input[ifm_converter.mouth_ooo_index] - model_input[ifm_converter.mouth_ooo_index] = uosum - model_input[ifm_converter.mouth_uuu_index] = 0 - is_open = (model_input[ifm_converter.mouth_aaa_index] + model_input[ - ifm_converter.mouth_iii_index] + uosum) > 0 - model_input[ifm_converter.mouth_lowered_corner_left_index] = 0 - model_input[ifm_converter.mouth_lowered_corner_right_index] = 0 - model_input[ifm_converter.mouth_raised_corner_left_index] = 0.5 if is_open else 0 - model_input[ifm_converter.mouth_raised_corner_right_index] = 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] = 0 - model_input[ifm_converter.eye_raised_lower_eyelid_right_index] = 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] += model_input[ - ifm_converter.eye_wink_right_index] - model_input[ifm_converter.eye_wink_right_index] = model_input[ - ifm_converter.eye_wink_left_index] / 2 - model_input[ifm_converter.eye_wink_left_index] = model_input[ - ifm_converter.eye_wink_left_index] / 2 - - model_input[ifm_converter.eye_surprised_left_index] += model_input[ - ifm_converter.eye_surprised_right_index] - model_input[ifm_converter.eye_surprised_right_index] = model_input[ - ifm_converter.eye_surprised_left_index] / 2 - model_input[ifm_converter.eye_surprised_left_index] = model_input[ - ifm_converter.eye_surprised_left_index] / 2 - - model_input[ifm_converter.eye_happy_wink_left_index] += model_input[ - ifm_converter.eye_happy_wink_right_index] - model_input[ifm_converter.eye_happy_wink_right_index] = model_input[ - ifm_converter.eye_happy_wink_left_index] / 2 - model_input[ifm_converter.eye_happy_wink_left_index] = model_input[ - ifm_converter.eye_happy_wink_left_index] / 2 - model_input[ifm_converter.mouth_aaa_index] = min( - model_input[ifm_converter.mouth_aaa_index] + - model_input[ifm_converter.mouth_ooo_index] / 2 + - model_input[ifm_converter.mouth_iii_index] / 2 + - model_input[ifm_converter.mouth_uuu_index] / 2, 1 - ) - model_input[ifm_converter.mouth_ooo_index] = 0 - model_input[ifm_converter.mouth_iii_index] = 0 - model_input[ifm_converter.mouth_uuu_index] = 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(0, len(simplify_arr)): - if simplify_arr[i] > 0: - model_input[i] = round(model_input[i] * simplify_arr[i]) / simplify_arr[i] - - if args.perf == 'model': - tic = time.perf_counter() - if args.eyebrow: - for i in range(12): - input_pose[0, i] = model_input[i] - for i in range(27): - input_pose[0, i + 12] = model_input[i + 12] - for i in range(6): - input_pose[0, i + 12 + 27] = model_input[i + 27 + 12] - - output_images = self.model.inference(input_pose) - - if args.perf == 'model': - print("postprocess", (time.perf_counter() - tic) * 1000) - tic = time.perf_counter() - - for i in range(len(self.shared_nps)): - np.copyto(self.shared_nps[i], output_images[i]) - self.output_queue.put_nowait(True) - - if args.debug: - self.model_fps_number.value = model_fps() - if self.model.cacher is not None: - self.cache_hit_ratio.value = self.model.cacher.hits / ( - self.model.cacher.hits + self.model.cacher.miss + 1) - try: - if args.use_tensorrt and args.model_cache and args.model_vram_cache: - self.gpu_cache_hit_ratio.value = self.model.tha.morpher_cacher.hits / ( - self.model.tha.morpher_cacher.hits + self.model.tha.morpher_cacher.miss) - except: - pass - - -@torch.no_grad() -def main(): - # Load character image - img = Image.open(f"data/images/{args.character}.png") - img = img.convert('RGBA') - ow, oh = img.size - for i, px in enumerate(img.getdata()): - if px[3] <= 0: - y = i // ow - x = i % ow - img.putpixel((x, y), (0, 0, 0, 0)) - if ow != 512 or oh != 512: - img = resize_to_512_center(img) - if args.alpha_clean: - curves = { - 'a': [ - (60, 0), - (200, 255) - ] - } - img = apply_color_curves(img, curves) - input_image = np.array(img) - input_image = cv2.cvtColor(input_image, cv2.COLOR_RGBA2BGRA) - - print("Character Image Loaded:", args.character) - cap = None - - output_fps = FPS() - - if not args.debug_input: - - if args.ifm is not None: - client_process = IFMClientProcess() - client_process.daemon = True - client_process.start() - print("iFacialMocap Service Running:", args.ifm) - - elif args.osf is not None: - client_process = OSFClientProcess() - client_process.daemon = True - client_process.start() - print("OpenSeeFace Service Running:", args.osf) - - elif args.mouse_input is not None: - client_process = MouseClientProcess() - client_process.daemon = True - client_process.start() - print("Mouse Input Running") - - else: - - if args.input == 'cam': - cap = cv2.VideoCapture(0) - ret, frame = cap.read() - if ret is None: - raise Exception("Can't find Camera") - else: - cap = cv2.VideoCapture(args.input) - frame_count = 0 - os.makedirs(os.path.join('dst', args.character, args.output_dir), exist_ok=True) - print("Webcam Input Running") - - facemesh = mp.solutions.face_mesh.FaceMesh(refine_landmarks=True) - - if args.output_webcam: - cam_scale = 1 - cam_width_scale = 1 - if args.anime4k: - cam_scale = 2 - if args.alpha_split: - cam_width_scale = 2 - - if args.output_webcam == 'spout': - # Create the sender - sender = SpoutSender("EasyVtuber", args.model_output_size * cam_scale * cam_width_scale, - args.model_output_size * cam_scale, GL_RGBA) - else: - cam = pyvirtualcam.Camera(width=args.model_output_size * cam_scale * cam_width_scale, - height=args.model_output_size * cam_scale, - fps=60, - backend=args.output_webcam, - fmt= - {'unitycapture': pyvirtualcam.PixelFormat.RGBA, - 'obs': pyvirtualcam.PixelFormat.RGB}[ - args.output_webcam]) - print(f'Using virtual camera: {cam.device}') - - a = None - - if args.anime4k: - parameters = ac.Parameters() - # enable HDN for ACNet - parameters.HDN = True - - # a = ac.AC( - # managerList=ac.ManagerList([ac.CUDAManager(dID=0)]), - # type=ac.ProcessorType.Cuda_ACNet, - # ) - - a = ac.AC( - managerList=ac.ManagerList([ac.OpenCLACNetManager(pID=0, dID=0)]), - type=ac.ProcessorType.OpenCL_ACNet, - ) - a.set_arguments(parameters) - print("Anime4K Loaded") - - position_vector = [0, 0, 0, 1] - position_vector_0 = None - pose_vector_0 = None - - pose_queue = [] - blender_data = {} - if (args.ifm): - blender_data = create_default_blender_data() - mouse_data = { - 'eye_l_h_temp': 0, - 'eye_r_h_temp': 0, - 'mouth_ratio': 0, - 'eye_y_ratio': 0, - 'eye_x_ratio': 0, - 'x_angle': 0, - 'y_angle': 0, - 'z_angle': 0, - } - - model_output = None - model_process = ModelClientProcess(input_image) - model_process.daemon = True - model_process.start() - model_output_nps = [np.ndarray((args.model_output_size, args.model_output_size, 4), dtype=np.uint8, - buffer=model_process.shms[i].buf) for i in range(len(model_process.shms))] - model_result_read_ptr = 0 - model_return_fps = FPS(5) - need_a_copy = True - - print("Ready. Close this console to exit.") - - mouth_eye_vector_c_hist = [] - - while True: - # ret, frame = cap.read() - # input_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) - # results = facemesh.process(input_frame) - - if args.perf == 'main': - tic = time.perf_counter() - if args.debug_input: - eyebrow_vector_c = [0.0] * 12 - mouth_eye_vector_c = [0.0] * 27 - pose_vector_c = [0.0] * 6 - - 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_c[14] = 0 - - 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) - - eyebrow_vector_c[6] = math.sin(time.perf_counter() * 1.1) - eyebrow_vector_c[7] = math.sin(time.perf_counter() * 1.1) - - elif args.osf is not None: - try: - new_blender_data = blender_data - while not client_process.should_terminate.value and not client_process.queue.empty(): - new_blender_data = client_process.queue.get_nowait() - blender_data = new_blender_data - except queue.Empty: - pass - eyebrow_vector_c = [0.0] * 12 - mouth_eye_vector_c = [0.0] * 27 - pose_vector_c = [0.0] * 6 - - if len(blender_data) != 0: - mouth_eye_vector_c[2] = 1 - blender_data['leftEyeOpen'] - mouth_eye_vector_c[3] = 1 - blender_data['rightEyeOpen'] - - mouth_eye_vector_c[14] = max(blender_data['MouthOpen'], 0) - # print(mouth_eye_vector_c[14]) - - mouth_eye_vector_c[25] = -blender_data['eyeRotationY'] * 3 - (blender_data['rotationX']) / 57.3 * 1.5 - mouth_eye_vector_c[26] = blender_data['eyeRotationX'] * 3 + (blender_data['rotationY']) / 57.3 - # print(mouth_eye_vector_c[25:27]) - eyebrow_vector_c[6] = blender_data['EyebrowUpDownLeft'] - eyebrow_vector_c[7] = blender_data['EyebrowUpDownRight'] - # print(blender_data['EyebrowUpDownLeft'],blender_data['EyebrowUpDownRight']) - - # if pose_vector_0==None: - # pose_vector_0=[0,0,0] - # pose_vector_0[0] = blender_data['rotationX'] - # pose_vector_0[1] = blender_data['rotationY'] - # pose_vector_0[2] = blender_data['rotationZ'] - # pose_vector_c[0] = (blender_data['rotationX']-pose_vector_0[0])/57.3*3 - # pose_vector_c[1] = -(blender_data['rotationY']-pose_vector_0[1])/57.3*3 - # pose_vector_c[2] = (blender_data['rotationZ']-pose_vector_0[2])/57.3 - pose_vector_c[0] = (blender_data['rotationX']) / 57.3 * 3 - pose_vector_c[1] = -(blender_data['rotationY']) / 57.3 * 3 - pose_vector_c[2] = (blender_data['rotationZ']) / 57.3 * 2 - # print(pose_vector_c) - - if position_vector_0 == None: - position_vector_0 = [0, 0, 0, 1] - position_vector_0[0] = blender_data['translationX'] - position_vector_0[1] = blender_data['translationY'] - position_vector_0[2] = blender_data['translationZ'] - position_vector[0] = -(blender_data['translationX'] - position_vector_0[0]) * 0.1 - position_vector[1] = -(blender_data['translationY'] - position_vector_0[1]) * 0.1 - position_vector[2] = -(blender_data['translationZ'] - position_vector_0[2]) * 0.1 - - # Apply filter - if len(mouth_eye_vector_c_hist) == 0: - for _ in range(3): - mouth_eye_vector_c_hist.append(copy.deepcopy(mouth_eye_vector_c)) - else: - mouth_eye_vector_c_hist.pop(0) - mouth_eye_vector_c_hist.append(mouth_eye_vector_c) - mouth_eye_vector_c[25] = mouth_eye_vector_c_hist[0][25] * 0.2 + mouth_eye_vector_c_hist[1][25] * 0.5 + mouth_eye_vector_c_hist[2][25] * 0.3 - mouth_eye_vector_c[26] = mouth_eye_vector_c_hist[0][26] * 0.2 + mouth_eye_vector_c_hist[1][26] * 0.5 + mouth_eye_vector_c_hist[2][26] * 0.3 - - elif args.ifm is not None: - # get pose from ifm - try: - new_blender_data = blender_data - while not client_process.should_terminate.value and not client_process.queue.empty(): - new_blender_data = client_process.queue.get_nowait() - blender_data = new_blender_data - except queue.Empty: - pass - - ifacialmocap_pose_converted = ifm_converter.convert(blender_data) - - # ifacialmocap_pose = blender_data - # - # eye_l_h_temp = ifacialmocap_pose[EYE_BLINK_LEFT] - # eye_r_h_temp = ifacialmocap_pose[EYE_BLINK_RIGHT] - # mouth_ratio = (ifacialmocap_pose[JAW_OPEN] - 0.10)*1.3 - # x_angle = -ifacialmocap_pose[HEAD_BONE_X] * 1.5 + 1.57 - # y_angle = -ifacialmocap_pose[HEAD_BONE_Y] - # z_angle = ifacialmocap_pose[HEAD_BONE_Z] - 1.57 - # - # eye_x_ratio = (ifacialmocap_pose[EYE_LOOK_IN_LEFT] - - # ifacialmocap_pose[EYE_LOOK_OUT_LEFT] - - # ifacialmocap_pose[EYE_LOOK_IN_RIGHT] + - # ifacialmocap_pose[EYE_LOOK_OUT_RIGHT]) / 2.0 / 0.75 - # - # eye_y_ratio = (ifacialmocap_pose[EYE_LOOK_UP_LEFT] - # + ifacialmocap_pose[EYE_LOOK_UP_RIGHT] - # - ifacialmocap_pose[EYE_LOOK_DOWN_RIGHT] - # + ifacialmocap_pose[EYE_LOOK_DOWN_LEFT]) / 2.0 / 0.75 - - eyebrow_vector_c = [0.0] * 12 - mouth_eye_vector_c = [0.0] * 27 - pose_vector_c = [0.0] * 6 - for i in range(0, 12): - eyebrow_vector_c[i] = ifacialmocap_pose_converted[i] - for i in range(12, 39): - mouth_eye_vector_c[i - 12] = ifacialmocap_pose_converted[i] - for i in range(39, 42): - pose_vector_c[i - 39] = ifacialmocap_pose_converted[i] - - position_vector = blender_data[HEAD_BONE_QUAT] - - elif args.mouse_input is not None: - - try: - new_blender_data = mouse_data - while not client_process.queue.empty(): - new_blender_data = client_process.queue.get_nowait() - mouse_data = new_blender_data - except queue.Empty: - pass - - eye_l_h_temp = mouse_data['eye_l_h_temp'] - eye_r_h_temp = mouse_data['eye_r_h_temp'] - mouth_ratio = mouse_data['mouth_ratio'] - eye_y_ratio = mouse_data['eye_y_ratio'] - eye_x_ratio = mouse_data['eye_x_ratio'] - x_angle = mouse_data['x_angle'] - y_angle = mouse_data['y_angle'] - z_angle = mouse_data['z_angle'] - - eyebrow_vector_c = [0.0] * 12 - mouth_eye_vector_c = [0.0] * 27 - pose_vector_c = [0.0] * 6 - - mouth_eye_vector_c[2] = eye_l_h_temp - mouth_eye_vector_c[3] = eye_r_h_temp - - mouth_eye_vector_c[14] = mouth_ratio * 1.5 - - mouth_eye_vector_c[25] = eye_y_ratio - mouth_eye_vector_c[26] = eye_x_ratio - - pose_vector_c[0] = x_angle - pose_vector_c[1] = y_angle - pose_vector_c[2] = z_angle - - else: - ret, frame = cap.read() - input_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) - results = facemesh.process(input_frame) - - if results.multi_face_landmarks is None: - continue - - facial_landmarks = results.multi_face_landmarks[0].landmark - - if args.debug: - pose, debug_image = get_pose(facial_landmarks, frame) - else: - pose = get_pose(facial_landmarks) - - if len(pose_queue) < 3: - pose_queue.append(pose) - pose_queue.append(pose) - pose_queue.append(pose) - else: - pose_queue.pop(0) - pose_queue.append(pose) - - np_pose = np.average(np.array(pose_queue), axis=0, weights=[0.6, 0.3, 0.1]) - - eye_l_h_temp = np_pose[0] - eye_r_h_temp = np_pose[1] - mouth_ratio = np_pose[2] - eye_y_ratio = np_pose[3] - eye_x_ratio = np_pose[4] - x_angle = np_pose[5] - y_angle = np_pose[6] - z_angle = np_pose[7] - - eyebrow_vector_c = [0.0] * 12 - mouth_eye_vector_c = [0.0] * 27 - pose_vector_c = [0.0] * 6 - - mouth_eye_vector_c[2] = eye_l_h_temp - mouth_eye_vector_c[3] = eye_r_h_temp - - mouth_eye_vector_c[14] = mouth_ratio * 1.5 - - mouth_eye_vector_c[25] = eye_y_ratio - mouth_eye_vector_c[26] = eye_x_ratio - - 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 - - pose_vector_c[3] = pose_vector_c[1] - pose_vector_c[4] = pose_vector_c[2] - - model_input_arr = eyebrow_vector_c - model_input_arr.extend(mouth_eye_vector_c) - model_input_arr.extend(pose_vector_c) - - if args.perf == 'main': - print('===') - print("input", time.perf_counter() - tic) - tic = time.perf_counter() - model_process.input_queue.put_nowait(model_input_arr) - - try: - while not model_process.output_queue.empty(): - model_process.output_queue.get_nowait() - model_result_read_ptr = 0 - model_return_fps() - need_a_copy = True - except queue.Empty: - pass - - if need_a_copy: - model_output = model_output_nps[model_result_read_ptr].copy() - need_a_copy = False - interval = 1 / (model_return_fps.view() + 1e-10) / len(model_output_nps) * 0.8 - if model_result_read_ptr + 1 < len(model_output_nps) and time.time() >= ( - (model_result_read_ptr + 1) * interval + model_return_fps.last()): - model_result_read_ptr += 1 - need_a_copy = True - - postprocessed_image = model_output - if args.perf == 'main': - print("Compute", time.perf_counter() - tic) - tic = time.perf_counter() - - k_scale = 1 - rotate_angle = 0 - dx = 0 - dy = 0 - if args.extend_movement: - k_scale = position_vector[2] * math.sqrt(args.extend_movement) + 1 - rotate_angle = -position_vector[0] * 10 * args.extend_movement - dx = position_vector[0] * 400 * k_scale * args.extend_movement - dy = -position_vector[1] * 600 * k_scale * args.extend_movement - if args.bongo: - rotate_angle -= 5 - - rm = cv2.getRotationMatrix2D((args.model_output_size / 2, args.model_output_size / 2), rotate_angle, k_scale) - rm[0, 2] += dx + args.model_output_size / 2 - args.model_output_size / 2 - rm[1, 2] += dy + args.model_output_size / 2 - args.model_output_size / 2 - - postprocessed_image = cv2.warpAffine( - postprocessed_image, - rm, - (args.model_output_size, args.model_output_size)) - - if args.perf == 'main': - print("postprocess", (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) - - # a.load_image_from_numpy(cv2.cvtColor(postprocessed_image, cv2.COLOR_RGBA2RGB), input_type=ac.AC_INPUT_RGB) - # img = cv2.imread("character/test41.png") - img1 = cv2.cvtColor(postprocessed_image, cv2.COLOR_BGRA2BGR) - # a.load_image_from_numpy(img, input_type=ac.AC_INPUT_BGR) - a.load_image_from_numpy(img1, input_type=ac.AC_INPUT_BGR) - a.process() - postprocessed_image = a.save_image_to_numpy() - postprocessed_image = cv2.merge((postprocessed_image, alpha_channel)) - if args.perf == 'main': - 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_BGR2BGRA) - postprocessed_image = cv2.hconcat([postprocessed_image, alpha_image]) - - if args.debug: - if args.use_sr: - time.sleep( - 0.000001) # Opencv is using different strategy when putting 512x512 and 1024x1024 image to imshow, - # in a way that 1024x1024 is much faster than 512x512 14msvs0.2ms(maybe using zero-copy?) so when using SR, add a delay to prevent overan - # Due to windows system limitation this actually sleep 16ms - output_frame = postprocessed_image - # 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) - cv2.putText(output_frame, str('OUT_FPS:%.1f' % output_fps_number), (0, 16), cv2.FONT_HERSHEY_PLAIN, 1, - (0, 255, 0), 1) - cv2.putText(output_frame, str( - 'GPU_FPS:%.1f' % ( - model_process.model_fps_number.value if not args.use_interpolation else model_process.model_fps_number.value * args.interpolation_scale)), - (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.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) - if args.output_webcam: - # result_image = np.zeros([720, 1280, 3], dtype=np.uint8) - # result_image[720 - 512:, 1280 // 2 - 256:1280 // 2 + 256] = cv2.resize( - # cv2.cvtColor(postprocessing_image(output_image.cpu()), cv2.COLOR_RGBA2RGB), (512, 512)) - result_image = postprocessed_image - if args.output_webcam == 'spout': - sender.send_image(cv2.cvtColor(postprocessed_image, cv2.COLOR_BGRA2RGBA), False) - time.sleep(0.000001) - else: - if args.output_webcam == 'obs': - result_image = cv2.cvtColor(result_image, cv2.COLOR_BGRA2RGB) - cam.send(result_image) - cam.sleep_until_next_frame() - if args.perf == 'main': - print("output", (time.perf_counter() - tic) * 1000) - tic = time.perf_counter() - - -if __name__ == '__main__': - main() diff --git a/pose.py b/pose.py deleted file mode 100644 index 04074da..0000000 --- a/pose.py +++ /dev/null @@ -1,183 +0,0 @@ -import cv2 -import numpy as np -from mediapipe.python.solutions.face_mesh_connections import FACEMESH_CONTOURS -from mediapipe.python.solutions.face_mesh_connections import FACEMESH_FACE_OVAL -from mediapipe.python.solutions.face_mesh_connections import FACEMESH_IRISES -from mediapipe.python.solutions.face_mesh_connections import FACEMESH_LEFT_EYE -from mediapipe.python.solutions.face_mesh_connections import FACEMESH_LEFT_EYEBROW -from mediapipe.python.solutions.face_mesh_connections import FACEMESH_LEFT_IRIS -from mediapipe.python.solutions.face_mesh_connections import FACEMESH_LIPS -from mediapipe.python.solutions.face_mesh_connections import FACEMESH_RIGHT_EYE -from mediapipe.python.solutions.face_mesh_connections import FACEMESH_RIGHT_EYEBROW -from mediapipe.python.solutions.face_mesh_connections import FACEMESH_RIGHT_IRIS -from mediapipe.python.solutions.face_mesh_connections import FACEMESH_TESSELATION - -from facial_points import MOUTH_TOP, MOUTH_BOTTOM, MOUTH_LEFT1, MOUTH_LEFT2, MOUTH_RIGHT -from facial_points import IRIS_L_TOP, IRIS_L_BOTTOM, IRIS_L_LEFT, IRIS_L_RIGHT -from facial_points import IRIS_R_TOP, IRIS_R_BOTTOM, IRIS_R_LEFT, IRIS_R_RIGHT -from utils import get_distance - - -# pose vector index ------------------------------------------ -# wink 0, 1 -# happy wink 2, 3 -# surprised 4, 5 -# relaxed 6, 7 -# unimpressed 8, 9 -# iris Shrinkage L 12 -# iris Shrinkage R MOUTH_TOP -# mouth aaa MOUTH_BOTTOM -# mouth iii 15 -# mouth uuu 16 -# mouth eee 17 -# mouth ooo 18 -# mouth delta 19 -# mouth lowered corner 18,19 -# mouth raised corner 20,21 -# mouth smirk 22 -# Iris rotation 25, 26 - -# x,y,z 27, 28, 29 -# ------------------------------------------------------------ - - -class Landmark: - def __init__(self, x=None, y=None, z=None): - self.x = x - self.y = y - self.z = z - - def __repr__(self): - return f'x: {self.x}\ny: {self.y}\nz: {self.z}' - -# def set_pose(pose, pose_vector) - - -def get_iris_center_point(landmarks, side=None): - """ - calculate right or left side iris center 3d coordinate - Args: - landmarks: - side: - - Returns: - iris center 3d coordinate point (Landmark) - - Raises: - if you didn't specify side - - """ - if side is None: - raise 'please choose side `left` or `right`' - elif side.lower() == 'right' or side.lower() == 'r': - landmark_list = FACEMESH_RIGHT_IRIS - elif side.lower() == 'left' or side.lower() == 'l': - landmark_list = FACEMESH_LEFT_IRIS - - temp = np.zeros(3) - for idx in list(map(lambda x: x[0], landmark_list)): - pos_x, pos_y, pos_z = landmarks[idx].x, landmarks[idx].y, landmarks[idx].z - temp[0] += pos_x - temp[1] += pos_y - temp[2] += pos_z - temp = temp / len(landmark_list) - return Landmark(temp[0], temp[1], temp[2]) - - -def get_pose(facial_landmarks, debug_image=None): - """ - extract pose vector from facial landmarks - - Args: - facial_landmarks (landmarks): extracted facial landmarks by mediapipe - debug_image (numpy array): frame image from video or webcam - - Returns: - tuple: pose vector - numpy array(optional): draw features to frame image - """ - iris_r_center = get_iris_center_point(facial_landmarks, 'r') - iris_l_center = get_iris_center_point(facial_landmarks, 'l') - - mouth_h = facial_landmarks[MOUTH_TOP].y - facial_landmarks[MOUTH_BOTTOM].y - mouth_w = facial_landmarks[MOUTH_RIGHT].x - (facial_landmarks[MOUTH_LEFT1].x + facial_landmarks[MOUTH_LEFT2].x) / 2 - mouth_ratio = mouth_h / mouth_w - - x_angle = np.arctan2(facial_landmarks[197].y - facial_landmarks[9].y, facial_landmarks[197].z - facial_landmarks[9].z) - y_angle = np.arctan2(facial_landmarks[IRIS_L_TOP].z - facial_landmarks[IRIS_R_TOP].z, facial_landmarks[IRIS_L_TOP].x - facial_landmarks[IRIS_R_TOP].x) - z_angle = np.arctan2(facial_landmarks[9].y - facial_landmarks[152].y, facial_landmarks[9].x - facial_landmarks[152].x) - - iris_rotation_l_h = get_distance(facial_landmarks[IRIS_L_TOP], facial_landmarks[IRIS_L_BOTTOM]) - iris_rotation_l_w = get_distance(facial_landmarks[IRIS_L_RIGHT], facial_landmarks[IRIS_L_LEFT]) - iris_rotation_r_h = get_distance(facial_landmarks[IRIS_R_TOP], facial_landmarks[IRIS_R_BOTTOM]) - iris_rotation_r_w = get_distance(facial_landmarks[IRIS_R_RIGHT], facial_landmarks[IRIS_R_LEFT]) - - iris_rotation_l_h_temp = np.sqrt( - (iris_l_center.x - facial_landmarks[IRIS_L_TOP].x) ** 2 + (iris_l_center.y - facial_landmarks[IRIS_L_TOP].y) ** 2) - iris_rotation_l_w_temp = np.sqrt( - (iris_l_center.x - facial_landmarks[IRIS_L_RIGHT].x) ** 2 + (iris_l_center.y - facial_landmarks[IRIS_L_RIGHT].y) ** 2) - iris_rotation_r_h_temp = np.sqrt( - (iris_r_center.x - facial_landmarks[IRIS_R_TOP].x) ** 2 + (iris_r_center.y - facial_landmarks[IRIS_R_TOP].y) ** 2) - iris_rotation_r_w_temp = np.sqrt( - (iris_r_center.x - facial_landmarks[IRIS_R_RIGHT].x) ** 2 + (iris_r_center.y - facial_landmarks[IRIS_R_RIGHT].y) ** 2) - - eye_x_ratio = ((iris_rotation_l_w_temp / iris_rotation_l_w + iris_rotation_r_w_temp / iris_rotation_r_w) - 1) * 3 - eye_y_ratio = ((iris_rotation_l_h_temp / iris_rotation_l_h + iris_rotation_r_h_temp / iris_rotation_r_h) - 1) * 3 - - eye_l_h_temp = 1 - 2 * (facial_landmarks[IRIS_R_BOTTOM].y - facial_landmarks[IRIS_R_TOP].y) / ( - facial_landmarks[IRIS_R_LEFT].x - facial_landmarks[IRIS_R_RIGHT].x) - eye_r_h_temp = 1 - 2 * (facial_landmarks[IRIS_L_BOTTOM].y - facial_landmarks[IRIS_L_TOP].y) / ( - facial_landmarks[IRIS_L_LEFT].x - facial_landmarks[IRIS_L_RIGHT].x) - - if debug_image is not None: - h, w, c = debug_image.shape - - # draw iris - cv2.line(debug_image, [int(iris_l_center.x * w), int(iris_l_center.y * h)], - [int(iris_l_center.x * w), int(iris_l_center.y * h)], (0, 128, 128), 3) - cv2.line(debug_image, [int(iris_r_center.x * w), int(iris_r_center.y * h)], - [int(iris_r_center.x * w), int(iris_r_center.y * h)], (255, 128, 0), 3) - - # draw iris connection - debug_image = cv2.line(debug_image, [int(iris_l_center.x * w), int(iris_l_center.y * h)], - [int(iris_r_center.x * w), int(iris_r_center.y * h)], (0, 128, 255), 1) - - # draw eye edge2iris - debug_image = cv2.line(debug_image, [int(iris_r_center.x * w), int(iris_r_center.y * h)], - [int(facial_landmarks[IRIS_R_RIGHT].x * w), int(facial_landmarks[IRIS_R_RIGHT].y * h)], (0, 255, 128), 1) - debug_image = cv2.line(debug_image, [int(iris_r_center.x * w), int(iris_r_center.y * h)], - [int(facial_landmarks[IRIS_R_BOTTOM].x * w), int(facial_landmarks[IRIS_R_BOTTOM].y * h)], (0, 255, 128), 1) - debug_image = cv2.line(debug_image, [int(iris_r_center.x * w), int(iris_r_center.y * h)], - [int(facial_landmarks[IRIS_R_LEFT].x * w), int(facial_landmarks[IRIS_R_LEFT].y * h)], (0, 255, 128), 1) - debug_image = cv2.line(debug_image, [int(iris_r_center.x * w), int(iris_r_center.y * h)], - [int(facial_landmarks[IRIS_R_TOP].x * w), int(facial_landmarks[IRIS_R_TOP].y * h)], (0, 255, 128), 1) - - debug_image = cv2.line(debug_image, [int(iris_l_center.x * w), int(iris_l_center.y * h)], - [int(facial_landmarks[IRIS_L_RIGHT].x * w), int(facial_landmarks[IRIS_L_RIGHT].y * h)], (255, 128, 0), 1) - debug_image = cv2.line(debug_image, [int(iris_l_center.x * w), int(iris_l_center.y * h)], - [int(facial_landmarks[IRIS_L_TOP].x * w), int(facial_landmarks[IRIS_L_TOP].y * h)], (255, 128, 0), 1) - debug_image = cv2.line(debug_image, [int(iris_l_center.x * w), int(iris_l_center.y * h)], - [int(facial_landmarks[IRIS_L_LEFT].x * w), int(facial_landmarks[IRIS_L_LEFT].y * h)], (255, 128, 0), 1) - debug_image = cv2.line(debug_image, [int(iris_l_center.x * w), int(iris_l_center.y * h)], - [int(facial_landmarks[IRIS_L_BOTTOM].x * w), int(facial_landmarks[IRIS_L_BOTTOM].y * h)], (255, 128, 0), 1) - - # draw mouth - debug_image = cv2.line(debug_image, [int((facial_landmarks[MOUTH_LEFT1].x + facial_landmarks[MOUTH_LEFT2].x) / 2 * w), int((facial_landmarks[MOUTH_LEFT1].y + facial_landmarks[MOUTH_LEFT2].y) / 2 * h)], - [int(facial_landmarks[MOUTH_RIGHT].x * w), int(facial_landmarks[MOUTH_RIGHT].y * h)], (255, 64, 0), 1) - - debug_image = cv2.line(debug_image, [int(facial_landmarks[MOUTH_TOP].x * w), int(facial_landmarks[MOUTH_TOP].y * h)], - [int(facial_landmarks[MOUTH_BOTTOM].x * w), int(facial_landmarks[MOUTH_BOTTOM].y * h)], (255, 64, 0), 1) - - # draw mouth - debug_image = cv2.line(debug_image, [int(facial_landmarks[197].x * w), int(facial_landmarks[197].y * h)], - [int(facial_landmarks[9].x * w), int(facial_landmarks[9].y * h)], (0, 0, 255), 1) - - debug_image = cv2.flip(debug_image, 1) - - cv2.putText(debug_image, f'x angle : {x_angle:.2f}', (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 1) - cv2.putText(debug_image, f'y angle : {y_angle:.2f}', (10, 45), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 1) - cv2.putText(debug_image, f'z angle : {z_angle:.2f}', (10, 60), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 1) - - return (eye_l_h_temp, eye_r_h_temp, mouth_ratio, eye_y_ratio, eye_x_ratio, x_angle , y_angle, z_angle), debug_image - else: - return eye_l_h_temp, eye_r_h_temp, mouth_ratio, eye_y_ratio, eye_x_ratio,x_angle, y_angle, z_angle \ No newline at end of file diff --git a/preprocess.py b/preprocess.py deleted file mode 100644 index ae6bb21..0000000 --- a/preprocess.py +++ /dev/null @@ -1,114 +0,0 @@ -from PIL import Image -import numpy as np -from scipy.interpolate import interp1d - - -def resize_to_512_center(img, fill_color=(0, 0, 0, 0)): - - # 计算缩放比例 - original_width, original_height = img.size - ratio = min(512 / original_width, 512 / original_height) - - # 等比缩放 - new_size = (int(original_width * ratio), int(original_height * ratio)) - resized_img = img.resize(new_size, Image.LANCZOS) # 高质量缩放算法 - - # 创建新画布 - canvas = Image.new(img.mode, (512, 512), fill_color) - - # 计算粘贴位置（居中） - paste_position = ( - (512 - new_size[0]) // 2, - (512 - new_size[1]) // 2 - ) - - # 合并图像（保持透明度） - canvas.paste(resized_img, paste_position) - - - return canvas - -def apply_color_curves(img, curves): - """ - 对图像应用颜色曲线调整 - :param input_path: 输入图像路径 - :param output_path: 输出图像路径 - :param curves: 曲线配置字典，格式为： - { - 'rgb': [(x1, y1), (x2, y2), ...], # 全局通道曲线 - 'r': [(x1, y1), (x2, y2), ...], # 红色通道曲线 - 'g': [(x1, y1), (x2, y2), ...], # 绿色通道曲线 - 'b': [(x1, y1), (x2, y2), ...] # 蓝色通道曲线 - } - """ - if img.mode == 'P': - img = img.convert('RGB') - - # 分离通道 - channels = list(img.split()) - - # 处理每个颜色通道 - for ch_idx, ch_name in enumerate(['r', 'g', 'b','a']): - # 合并全局和通道特定曲线 - combined_points = [] - if 'rgb' in curves and not ch_name=='a': - combined_points.extend(curves['rgb']) - if ch_name in curves: - combined_points.extend(curves[ch_name]) - - if not combined_points: - continue - - # 生成LUT - points = sorted(combined_points, key=lambda x: x[0]) - xs, ys = zip(*points) - - # 自动选择插值方式 - interpolation_kind = 'quadratic' if len(points) >= 3 else 'linear' - interpolator = interp1d( - xs, ys, - kind=interpolation_kind, - bounds_error=False, - fill_value=(ys[0], ys[-1]) - ) - - # 创建查找表 - x_range = np.arange(256) - y_new = interpolator(x_range) - y_new = np.clip(y_new, 0, 255).astype(np.uint8) - lut = y_new.tolist() - - # 应用LUT - channels[ch_idx] = channels[ch_idx].point(lut) - - # 合并通道并保存 - merged_channels = channels - img_out = Image.merge(img.mode, merged_channels) - - return img_out - - - # 使用示例 -if __name__ == "__main__": - # S型曲线增强对比度 - curves = { - 'a': [ - (60, 0), - (127, 127), - (170, 255) - ] - } - - # 红色通道增强中间调 - # curves = { - # 'r': [ - # (0, 0), - # (127, 160), - # (255, 255) - # ] - # } - img = Image.open('input.png') - # img = apply_color_curves(img, curves) - img = resize_to_512_center(img) - - img.save('output.png') \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index 9bd7383..ac93f95 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,12 +1,12 @@ -numpy==1.26.4 +numpy scipy matplotlib protobuf==3.20.1 pynput==1.7.6 -mediapipe +mediapipe==0.10.11 opencv_python==4.5.5.64 -Pillow==9.1.0 -pyanime4k==2.5.2 +Pillow +pyanime4k==3.1.0 pyvirtualcam==0.9.1 pyopengl wxpython @@ -14,4 +14,5 @@ onnx onnxruntime-directml turbojpeg tqdm -torch \ No newline at end of file +torch +OneEuroFilter \ No newline at end of file diff --git a/src/args.py b/src/args.py index 3f32426..7571675 100644 --- a/src/args.py +++ b/src/args.py @@ -29,8 +29,8 @@ parser.add_argument('--alpha_split', action='store_true') parser.add_argument('--bongo', action='store_true') parser.add_argument('--extend_movement', action='store_true') -parser.add_argument('--filter_min_cutoff', type=float, default=0.3) -parser.add_argument('--filter_beta', type=float, default=1.0) +parser.add_argument('--filter_min_cutoff', type=float, default=10.0) +parser.add_argument('--filter_beta', type=float, default=0.3) parser.add_argument('--simplify', type=int, default=1) @@ -85,5 +85,3 @@ if args.output_spout2: args.alpha_split = False # Disable alpha split for spout2 output args.model_output_size = 1024 if args.use_sr else 512 - -print("Arguments:", args) \ No newline at end of file diff --git a/src/debug_input_client.py b/src/debug_input_client.py index 37d8136..09e0323 100644 --- a/src/debug_input_client.py +++ b/src/debug_input_client.py @@ -3,9 +3,7 @@ import numpy as np import time import math from .utils.timer_wait import wait_until -from .utils.channel_shared_mem import SharedMemoryExclusiveChannel -from .utils.filter import OneEuroFilterNumpy -from .args import args +from .utils.shared_mem_guard import SharedMemoryGuard class DebugInputClientProcess(Process): def __init__(self, pose_position_shm: shared_memory.SharedMemory): @@ -16,12 +14,10 @@ class DebugInputClientProcess(Process): def run(self): last_time : float = time.perf_counter() interval : float = 1.0 / 60 # 60 FPS - pose_position_shm_channel = SharedMemoryExclusiveChannel(self.pose_position_shm, ctrl_name="pose_position_shm_ctrl") + pose_position_shm_guard = SharedMemoryGuard(self.pose_position_shm, ctrl_name="pose_position_shm_ctrl") np_pose_shm = np.ndarray((45,), dtype=np.float32, buffer=self.pose_position_shm.buf[:45 * 4]) np_position_shm = np.ndarray((4,), dtype=np.float32, buffer=self.pose_position_shm.buf[45 * 4:45 * 4 + 4 * 4]) - pose_filter = OneEuroFilterNumpy(freq=60, mincutoff=args.filter_min_cutoff, beta=args.filter_beta) - position_filter = OneEuroFilterNumpy(freq=60, mincutoff=args.filter_min_cutoff, beta=args.filter_beta) while True: eyebrow_vector = [0.0] * 12 mouth_eye_vector = [0.0] * 27 @@ -53,8 +49,8 @@ class DebugInputClientProcess(Process): position_vector[2] = math.sin(time.perf_counter() * 0.7) * 0.1 position_vector[3] = 1 - with pose_position_shm_channel.lock(): - np_pose_shm[:] = pose_filter(np.array(model_input_arr, dtype=np.float32)) - np_position_shm[:] = position_filter(np.array(position_vector, dtype=np.float32)) + with pose_position_shm_guard.lock(): + np_pose_shm[:] = np.array(model_input_arr, dtype=np.float32) + np_position_shm[:] = np.array(position_vector, dtype=np.float32) wait_until(last_time + interval) last_time += interval \ No newline at end of file diff --git a/src/face_mesh_client.py b/src/face_mesh_client.py index 5b67cfa..68a09d8 100644 --- a/src/face_mesh_client.py +++ b/src/face_mesh_client.py @@ -3,7 +3,7 @@ from multiprocessing import Process, shared_memory, Value import cv2 import numpy as np from .args import args -from .utils.channel_shared_mem import SharedMemoryExclusiveChannel +from .utils.shared_mem_guard import SharedMemoryGuard from .utils.pose import get_pose from .utils.fps import FPS from .utils.filter import OneEuroFilterNumpy @@ -19,7 +19,7 @@ class FaceMeshClientProcess(Process): def run(self): facemesh = mp.solutions.face_mesh.FaceMesh(refine_landmarks=True) - pose_position_shm_channel = SharedMemoryExclusiveChannel(self.pose_position_shm, ctrl_name="pose_position_shm_ctrl") + pose_position_shm_guard = SharedMemoryGuard(self.pose_position_shm, ctrl_name="pose_position_shm_ctrl") np_pose_shm = np.ndarray((45,), dtype=np.float32, buffer=self.pose_position_shm.buf[:45 * 4]) np_position_shm = np.ndarray((4,), dtype=np.float32, buffer=self.pose_position_shm.buf[45 * 4:45 * 4 + 4 * 4]) @@ -70,7 +70,7 @@ class FaceMeshClientProcess(Process): mouth_eye_vector[2] = min(eye_l_h_temp, eye_r_h_temp) mouth_eye_vector[3] = min(eye_l_h_temp, eye_r_h_temp) - mouth_eye_vector[14] = mouth_ratio * 1.5 + mouth_eye_vector[14] = mouth_ratio * 2.0 mouth_eye_vector[25] = 0.0 # keep iris stable for user demo mouth_eye_vector[26] = 0.0 @@ -86,6 +86,6 @@ class FaceMeshClientProcess(Process): model_input_arr.extend(mouth_eye_vector) model_input_arr.extend(pose_vector) - with pose_position_shm_channel.lock(): + with pose_position_shm_guard.lock(): np_pose_shm[:] = pose_filter(np.array(model_input_arr, dtype=np.float32)) np_position_shm[:] = position_vector \ No newline at end of file diff --git a/src/i_facial_mocap_client.py b/src/i_facial_mocap_client.py index a48d0b7..d877525 100644 --- a/src/i_facial_mocap_client.py +++ b/src/i_facial_mocap_client.py @@ -1,7 +1,7 @@ from multiprocessing import Process, shared_memory, Value from .args import args from .utils.timer_wait import wait_until -from .utils.channel_shared_mem import SharedMemoryExclusiveChannel +from .utils.shared_mem_guard import SharedMemoryGuard import socket import errno import numpy as np @@ -20,7 +20,7 @@ class IFMClientProcess(Process): self.port = int(args.ifm.split(':')[1]) self.fps = Value('f', 0.0) def run(self): - pose_position_shm_channel = SharedMemoryExclusiveChannel(self.pose_position_shm, ctrl_name="pose_position_shm_ctrl") + pose_position_shm_guard = SharedMemoryGuard(self.pose_position_shm, ctrl_name="pose_position_shm_ctrl") np_pose_shm = np.ndarray((45,), dtype=np.float32, buffer=self.pose_position_shm.buf[:45 * 4]) np_position_shm = np.ndarray((4,), dtype=np.float32, buffer=self.pose_position_shm.buf[45 * 4:45 * 4 + 4 * 4]) @@ -49,8 +49,8 @@ class IFMClientProcess(Process): self.socket.settimeout(0.5) # 500 ms - pose_filter = OneEuroFilterNumpy(freq=input_fps.view(), mincutoff=args.filter_min_cutoff, beta=args.filter_beta) - position_filter = OneEuroFilterNumpy(freq=input_fps.view(), mincutoff=args.filter_min_cutoff, beta=args.filter_beta) + # pose_filter = OneEuroFilterNumpy(freq=input_fps.view(), mincutoff=args.filter_min_cutoff, beta=args.filter_beta) + # position_filter = OneEuroFilterNumpy(freq=input_fps.view(), mincutoff=args.filter_min_cutoff, beta=args.filter_beta) print("iFacialMocap Input Running at {:.2f} FPS".format(input_fps.view())) while True: @@ -90,11 +90,11 @@ class IFMClientProcess(Process): model_input_arr.extend(mouth_eye_vector) model_input_arr.extend(pose_vector) - with pose_position_shm_channel.lock(): - np_pose_shm[:] = pose_filter(np.array(model_input_arr, dtype=np.float32)) - np_position_shm[:] = position_filter(np.array(position_vector, dtype=np.float32)) - # np_pose_shm[:] = np.array(model_input_arr, dtype=np.float32) - # np_position_shm[:] = np.array(position_vector, dtype=np.float32) + with pose_position_shm_guard.lock(): + # np_pose_shm[:] = pose_filter(np.array(model_input_arr, dtype=np.float32)) + # np_position_shm[:] = position_filter(np.array(position_vector, dtype=np.float32)) + np_pose_shm[:] = np.array(model_input_arr, dtype=np.float32) + np_position_shm[:] = np.array(position_vector, dtype=np.float32) @staticmethod def convert_from_blender_data(blender_data): diff --git a/src/main.py b/src/main.py index 1375c5a..fa09a59 100644 --- a/src/main.py +++ b/src/main.py @@ -4,7 +4,7 @@ from .model_infer_client import ModelClientProcess from .args import args from .utils.preprocess import resize_to_512_center, apply_color_curves import cv2 -from .utils.channel_shared_mem import SharedMemoryExclusiveChannel +from .utils.shared_mem_guard import SharedMemoryGuard from multiprocessing import shared_memory from .utils.timer_wait import wait_until from PIL import Image @@ -68,7 +68,7 @@ def main(): cam_width_scale = 2 if args.alpha_split else 1 ret_channels = 3 if args.output_virtual_cam or args.output_debug else 4 ret_batch_shm_channels = [ - SharedMemoryExclusiveChannel(infer_process.ret_shared_mem, ctrl_name=f"ret_shm_ctrl_batch_{i}") + SharedMemoryGuard(infer_process.ret_shared_mem, ctrl_name=f"ret_shm_ctrl_batch_{i}") for i in range(args.interpolation_scale) ] np_ret_shms = [ @@ -109,26 +109,20 @@ def main(): elif args.output_spout2: spout_sender.send_image(np_ret_shms[i], False) else: - cv2.imshow("Output Frame Batch {}".format(i), np_ret_shms[i]) + cv2.imshow("EasyVtuber Debug Frame", np_ret_shms[i]) cv2.waitKey(1) wait_until(last_time + interval) last_time += interval ret_batch_shm_channels[i].release() - print("pipeline FPS: {:.2f}, Input FPS: {:.2f}, Model Avg Interval: {:.2f} ms, Cache Hit Ratio: {:.2f}%, GPU Cache Hit Ratio: {:.2f}%, current time {:.5f}".format( + print("Infer Process FPS: {:.2f}, Input FPS: {:.2f}, Model Avg Interval: {:.2f} ms, Cache Hit Ratio: {:.2f}%, GPU Cache Hit Ratio: {:.2f}%, Output Pipeline FPS {:.5f}".format( infer_process.pipeline_fps_number.value, input_fps.value, infer_process.average_model_interval.value * 1000, infer_process.cache_hit_ratio.value * 100, infer_process.gpu_cache_hit_ratio.value * 100, - pipeline_fps() + pipeline_fps() * args.interpolation_scale ), end ='\r', flush=True) - - - - - - if __name__ == "__main__": main() \ No newline at end of file diff --git a/src/model_infer_client.py b/src/model_infer_client.py index 9eb5385..97be6f3 100644 --- a/src/model_infer_client.py +++ b/src/model_infer_client.py @@ -5,7 +5,7 @@ import numpy as np import time from .ezvtb_rt_interface import get_core from .args import args -from .utils.channel_shared_mem import SharedMemoryExclusiveChannel +from .utils.shared_mem_guard import SharedMemoryGuard from .utils.pose_simplify import pose_simplify from .utils.fps import FPS, Interval from typing import List @@ -32,12 +32,12 @@ class ModelClientProcess(Process): self.finish_event = Event() def run(self): - pose_position_shm_channel = SharedMemoryExclusiveChannel(self.pose_position_shm, ctrl_name="pose_position_shm_ctrl") + pose_position_shm_guard = SharedMemoryGuard(self.pose_position_shm, ctrl_name="pose_position_shm_ctrl") np_pose_shm = np.ndarray((45,), dtype=np.float32, buffer=self.pose_position_shm.buf[:45 * 4]) np_position_shm = np.ndarray((4,), dtype=np.float32, buffer=self.pose_position_shm.buf[45 * 4:45 * 4 + 4 * 4]) - ret_batch_shm_channels = [ - SharedMemoryExclusiveChannel(self.ret_shared_mem, ctrl_name=f"ret_shm_ctrl_batch_{i}") + ret_batch_shm_guard = [ + SharedMemoryGuard(self.ret_shared_mem, ctrl_name=f"ret_shm_ctrl_batch_{i}") for i in range(args.interpolation_scale) ] np_ret_shms = [ @@ -81,7 +81,7 @@ class ModelClientProcess(Process): print("Model Inference Ready") while True: - with pose_position_shm_channel.lock(): + with pose_position_shm_guard.lock(): np_pose = np_pose_shm.copy() np_position = np_position_shm.copy() @@ -116,7 +116,7 @@ class ModelClientProcess(Process): self.pipeline_fps_number.value = pipeline_fps() for i in range(args.interpolation_scale): - with ret_batch_shm_channels[i].lock(): # get pressure from main process if ret not consumed + with ret_batch_shm_guard[i].lock(): # get pressure from main process if ret not consumed np_ret_shms[i][:, :, :] = output_images[i] self.finish_event.set() # Back pressure main process loop if infer slow diff --git a/src/mouse_client.py b/src/mouse_client.py index 42e8fb1..a61436b 100644 --- a/src/mouse_client.py +++ b/src/mouse_client.py @@ -3,8 +3,7 @@ from pynput.mouse import Controller import numpy as np import time from .args import args -from .utils.filter import OneEuroFilterNumpy -from .utils.channel_shared_mem import SharedMemoryExclusiveChannel +from .utils.shared_mem_guard import SharedMemoryGuard from .utils.timer_wait import wait_until class MouseClientProcess(Process): @@ -15,7 +14,7 @@ class MouseClientProcess(Process): def run(self): mouse = Controller() - pose_position_shm_channel = SharedMemoryExclusiveChannel(self.pose_position_shm, ctrl_name="pose_position_shm_ctrl") + pose_position_shm_guard = SharedMemoryGuard(self.pose_position_shm, ctrl_name="pose_position_shm_ctrl") np_pose_shm = np.ndarray((45,), dtype=np.float32, buffer=self.pose_position_shm.buf[:45 * 4]) np_position_shm = np.ndarray((4,), dtype=np.float32, buffer=self.pose_position_shm.buf[45 * 4:45 * 4 + 4 * 4]) @@ -35,7 +34,6 @@ class MouseClientProcess(Process): interval : float = 1.0 / 60 # 60 FPS print("Mouse Input Running at 60 FPS") - pose_filter = OneEuroFilterNumpy(freq=60, mincutoff=args.filter_min_cutoff, beta=args.filter_beta) position_vector = np.array([0, 0, 0, 1], dtype=np.float32) while True: pos = mouse.position @@ -93,8 +91,8 @@ class MouseClientProcess(Process): model_input_arr.extend(mouth_eye_vector) model_input_arr.extend(pose_vector) - with pose_position_shm_channel.lock(): - np_pose_shm[:] = pose_filter(np.array(model_input_arr, dtype=np.float32)) + with pose_position_shm_guard.lock(): + np_pose_shm[:] = np.array(model_input_arr, dtype=np.float32) np_position_shm[:] = position_vector wait_until(last_time + interval) diff --git a/src/open_see_face_client.py b/src/open_see_face_client.py index 68e3070..0be9a01 100644 --- a/src/open_see_face_client.py +++ b/src/open_see_face_client.py @@ -3,7 +3,7 @@ from .args import args import socket import struct import numpy as np -from .utils.channel_shared_mem import SharedMemoryExclusiveChannel +from .utils.shared_mem_guard import SharedMemoryGuard from .utils.fps import FPS from .utils.filter import OneEuroFilterNumpy from OneEuroFilter import OneEuroFilter @@ -18,7 +18,7 @@ class OSFClientProcess(Process): self.fps = Value('f', 0.0) def run(self): - pose_position_shm_channel = SharedMemoryExclusiveChannel(self.pose_position_shm, ctrl_name="pose_position_shm_ctrl") + pose_position_shm_guard = SharedMemoryGuard(self.pose_position_shm, ctrl_name="pose_position_shm_ctrl") np_pose_shm = np.ndarray((45,), dtype=np.float32, buffer=self.pose_position_shm.buf[:45 * 4]) np_position_shm = np.ndarray((4,), dtype=np.float32, buffer=self.pose_position_shm.buf[45 * 4:45 * 4 + 4 * 4]) @@ -42,8 +42,8 @@ class OSFClientProcess(Process): pose_filter = OneEuroFilterNumpy(freq=input_fps.view(), mincutoff=args.filter_min_cutoff, beta=args.filter_beta) position_filter = OneEuroFilterNumpy(freq=input_fps.view(), mincutoff=args.filter_min_cutoff, beta=args.filter_beta) - iris_x_filter = OneEuroFilter(freq=input_fps.view(), mincutoff=0.1, beta=0.3) # extra filter for iris movement - iris_y_filter = OneEuroFilter(freq=input_fps.view(), mincutoff=0.1, beta=0.3) # extra filter for iris movement + iris_x_filter = OneEuroFilter(freq=input_fps.view(), mincutoff=0.001, beta=1.0) # extra filter for iris movement + iris_y_filter = OneEuroFilter(freq=input_fps.view(), mincutoff=0.001, beta=1.0) # extra filter for iris movement rotation_offset = None print("OpenSeeFace Input Running at {:.2f} FPS".format(input_fps.view())) while True: @@ -166,7 +166,7 @@ class OSFClientProcess(Process): model_input_arr.extend(mouth_eye_vector) model_input_arr.extend(pose_vector) - with pose_position_shm_channel.lock(): + with pose_position_shm_guard.lock(): np_pose_shm[:] = pose_filter(np.array(model_input_arr, dtype=np.float32)) np_position_shm[:] = position_filter(np.array(position_vector, dtype=np.float32)) # np_pose_shm[:] = np.array(model_input_arr, dtype=np.float32) diff --git a/facial_points.py b/src/utils/facial_points.py similarity index 100% rename from facial_points.py rename to src/utils/facial_points.py diff --git a/src/utils/pose.py b/src/utils/pose.py index 1bc6127..3fa76d6 100644 --- a/src/utils/pose.py +++ b/src/utils/pose.py @@ -12,9 +12,9 @@ from mediapipe.python.solutions.face_mesh_connections import FACEMESH_RIGHT_EYEB from mediapipe.python.solutions.face_mesh_connections import FACEMESH_RIGHT_IRIS from mediapipe.python.solutions.face_mesh_connections import FACEMESH_TESSELATION -from facial_points import MOUTH_TOP, MOUTH_BOTTOM, MOUTH_LEFT1, MOUTH_LEFT2, MOUTH_RIGHT -from facial_points import IRIS_L_TOP, IRIS_L_BOTTOM, IRIS_L_LEFT, IRIS_L_RIGHT -from facial_points import IRIS_R_TOP, IRIS_R_BOTTOM, IRIS_R_LEFT, IRIS_R_RIGHT +from .facial_points import MOUTH_TOP, MOUTH_BOTTOM, MOUTH_LEFT1, MOUTH_LEFT2, MOUTH_RIGHT +from .facial_points import IRIS_L_TOP, IRIS_L_BOTTOM, IRIS_L_LEFT, IRIS_L_RIGHT +from .facial_points import IRIS_R_TOP, IRIS_R_BOTTOM, IRIS_R_LEFT, IRIS_R_RIGHT from .utils import get_distance diff --git a/src/utils/channel_shared_mem.py b/src/utils/shared_mem_guard.py similarity index 99% rename from src/utils/channel_shared_mem.py rename to src/utils/shared_mem_guard.py index dfecece..c349d3e 100644 --- a/src/utils/channel_shared_mem.py +++ b/src/utils/shared_mem_guard.py @@ -46,7 +46,7 @@ MUTEX_ALL_ACCESS = 0x1F0001 ERROR_ALREADY_EXISTS = 183 -class SharedMemoryExclusiveChannel: +class SharedMemoryGuard: """ Two-process exclusive access wrapper using Windows Named Mutex: - payload SharedMemory is provided externally diff --git a/utils.py b/utils.py deleted file mode 100644 index 30bf9fe..0000000 --- a/utils.py +++ /dev/null @@ -1,80 +0,0 @@ -import torch -import numpy as np - - -def linear_rgb2srgb(image): - """ - convert linear rgb image to srgb image - Args: - image (numpy array): linear rgb color space image - - Returns: - image (numpy array): srgb color space image - """ - return np.where(image <= 0.04045, image / 12.92, ((image + 0.055) / 1.055) ** 2.4) - - -def srgb2linear_rgb(image): - """ - convert srgb image to linear rgb image - Args: - image (numpy array): srgb color space image - - Returns: - image (numpy array): linear rgb color space image - """ - return np.where(image <= 0.003130804953560372, image * 12.92, 1.055 * (image ** (1.0 / 2.4)) - 0.055) - - -def preprocessing_image(image): - """ - convert image(numpy array, unit8) to tensor - Args: - image (numpy array): character image (256x256, alpha channel included) - Returns: - tensor - """ - np_image = np.array(image) / 255 - clipped_image = np.clip(np_image, 0, 1) - srgb_image = linear_rgb2srgb(clipped_image) - h, w, c = srgb_image.shape - linear_image = srgb_image.reshape(h * w, c) - for pixel in linear_image: - if pixel[3] == 0.0: - pixel[0:3] = 0.0 - reshaped_image = linear_image.transpose().reshape(c, h, w) - return reshaped_image - - -def postprocessing_image(tensor): - """ - convert tensor to image(numpy array, unit8) - Args: - tensor - Returns: - image (numpy array): character image (256x256, alpha channel included) - """ - tensor = tensor.detach().squeeze(0) - reshaped_tensor = tensor.permute(1, 2, 0) - np_image = reshaped_tensor.numpy() - np_image = (np_image + 1) / 2 - srgb_image = np_image[..., :3] - alpha_image = np_image[..., 3] - clipped_image = np.clip(srgb_image, 0, 1) - rgb_image = srgb2linear_rgb(clipped_image) - rgba_image = np.concatenate([rgb_image, alpha_image[..., np.newaxis]], axis=2) - rgba_image = rgba_image * 255 - return rgba_image.astype(np.uint8) - - -def get_distance(a, b): - """ - calculate euclidean distance a to b - Args: - a (landmark): 3d points - b (landmark): 3d points - - Returns: - L2 distance (float) - """ - return np.sqrt((a.x - b.x) ** 2 + (a.y - b.y) ** 2)