feat:launcher ready
This commit is contained in:
zpeng11 committed
1 parent
ebe0073eca
commit
638d6970f7
21 files changed
+84
-1799
No files matched your search
+1
-1
@@ -1,7 +1,7 @@
|
||||
pretrained/**
|
||||
!pretrained/PUT_MODEL_HERE
|
||||
character
|
||||
src
|
||||
!src
|
||||
!tha4/src
|
||||
tha4/distiller-ui-doc
|
||||
tha4/docs
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
+1
-1
Submodule ezvtuber-rt updated: 1bc5feeb09...2ea780d884.
+32
-55
@@ -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')
|
||||
|
||||
@@ -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
|
||||
-114
@@ -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')
|
||||
+6
-5
@@ -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
|
||||
torch
|
||||
OneEuroFilter
|
||||
+2
-4
@@ -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)
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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):
|
||||
|
||||
+5
-11
@@ -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()
|
||||
@@ -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
|
||||
|
||||
+4
-6
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
File renamed without changes.
+3
-3
@@ -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
|
||||
|
||||
|
||||
|
||||
@@ -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
|
||||
@@ -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)
|
||||
Reference in new issue
Block a user