212 files changed
+19818
-2386
No files matched your search
+11
-2
@@ -1,7 +1,11 @@
|
||||
pretrained/**
|
||||
!pretrained/PUT_MODEL_HERE
|
||||
character
|
||||
src
|
||||
!src
|
||||
!tha4/src
|
||||
tha4/distiller-ui-doc
|
||||
tha4/docs
|
||||
tha4/poetry
|
||||
dst
|
||||
.idea
|
||||
bin
|
||||
@@ -142,4 +146,9 @@ data/
|
||||
!data/models/placeholder.txt
|
||||
*.iml
|
||||
.idea/
|
||||
*.pt
|
||||
*.pt
|
||||
|
||||
.serena
|
||||
onnx_model_tha4/fp16
|
||||
onnx_model_tha4/fp32
|
||||
onnx_model_tha4/
|
||||
+2
-2
@@ -1,4 +1,4 @@
|
||||
[submodule "ezvtuber-rt"]
|
||||
path = ezvtuber-rt
|
||||
url = https://github.com/zpeng11/ezvtuber-rt.git
|
||||
branch = main
|
||||
url = git@github.com:yuyuyzl/ezvtuber-rt.git
|
||||
branch = tha4
|
||||
@@ -0,0 +1,7 @@
|
||||
@echo on
|
||||
cd /D "%~dp0"
|
||||
|
||||
SET PATH=%~dp0envs\TensorRT-RTX-1.3.0.35_cu129\bin;%~dp0envs\python_embedded;%~dp0envs\python_embedded\Scripts;%~dp0envs\python_embedded\Library\bin;%PATH%
|
||||
|
||||
start "" pythonw launcher2.py
|
||||
exit
|
||||
@@ -1,41 +0,0 @@
|
||||
@echo on
|
||||
cd /D "%~dp0"
|
||||
|
||||
WHERE conda
|
||||
IF %ERRORLEVEL% NEQ 0 (
|
||||
if not exist "envs" mkdir envs
|
||||
if not exist "envs\miniconda3" mkdir envs\miniconda3
|
||||
|
||||
IF not EXIST %~dp0envs\miniconda3\Scripts (
|
||||
@RD /S /Q %~dp0envs\miniconda3
|
||||
mkdir %~dp0envs\miniconda3
|
||||
echo "Downloading miniconda..."
|
||||
powershell -Command "(New-Object Net.WebClient).DownloadFile('https://repo.anaconda.com/miniconda/Miniconda3-py310_24.9.2-0-Windows-x86_64.exe', '.\envs\miniconda3.exe')"
|
||||
|
||||
echo "Installling minconda..."
|
||||
start /wait "" %~dp0envs\miniconda3.exe /S /AddToPath=0 /RegisterPython=0 /InstallationType=JustMe /D=%~dp0envs\miniconda3
|
||||
echo "Successfully install minconda"
|
||||
del envs\miniconda3.exe
|
||||
)
|
||||
)
|
||||
|
||||
IF EXIST %~dp0envs\miniconda3\Scripts SET PATH=%~dp0envs\miniconda3\Scripts;%PATH%
|
||||
|
||||
call activate
|
||||
call conda env list
|
||||
call conda update -y --all
|
||||
call conda create -y -n ezvtb_rt_venv python=3.10
|
||||
call conda activate ezvtb_rt_venv
|
||||
call conda env list
|
||||
|
||||
call conda install -y nvidia/label/cuda-12.9.0::cuda-nvcc-dev_win-64
|
||||
call conda install -y conda-forge::pycuda
|
||||
|
||||
call python -m pip install --upgrade pip wheel
|
||||
echo yes|python -m pip install nvidia-cudnn-cu12
|
||||
|
||||
echo yes|pip install tensorrt_cu12_libs==10.11.0.33 tensorrt_cu12_bindings==10.11.0.33 tensorrt==10.11.0.33 --extra-index-url https://pypi.nvidia.com
|
||||
|
||||
call python -m pip install -r requirements.txt --no-warn-script-location
|
||||
|
||||
pause
|
||||
@@ -0,0 +1,8 @@
|
||||
@echo on
|
||||
cd /D "%~dp0"
|
||||
|
||||
SET PATH=%~dp0envs\TensorRT-RTX-1.3.0.35_cu129\bin;%~dp0envs\python_embedded;%~dp0envs\python_embedded\Scripts;%~dp0envs\python_embedded\Library\bin;%PATH%
|
||||
|
||||
python launcher2.py
|
||||
|
||||
pause
|
||||
@@ -1,41 +0,0 @@
|
||||
@echo on
|
||||
cd /D "%~dp0"
|
||||
|
||||
WHERE conda
|
||||
IF %ERRORLEVEL% NEQ 0 (
|
||||
if not exist "envs" mkdir envs
|
||||
if not exist "envs\miniconda3" mkdir envs\miniconda3
|
||||
|
||||
IF not EXIST %~dp0envs\miniconda3\Scripts (
|
||||
@RD /S /Q %~dp0envs\miniconda3
|
||||
mkdir %~dp0envs\miniconda3
|
||||
echo "Downloading miniconda..."
|
||||
powershell -Command "(New-Object Net.WebClient).DownloadFile('https://repo.anaconda.com/miniconda/Miniconda3-py310_24.9.2-0-Windows-x86_64.exe', '.\envs\miniconda3.exe')"
|
||||
|
||||
echo "Installling minconda..."
|
||||
start /wait "" %~dp0envs\miniconda3.exe /S /AddToPath=0 /RegisterPython=0 /InstallationType=JustMe /D=%~dp0envs\miniconda3
|
||||
echo "Successfully install minconda"
|
||||
del envs\miniconda3.exe
|
||||
)
|
||||
)
|
||||
|
||||
IF EXIST %~dp0envs\miniconda3\Scripts SET PATH=%~dp0envs\miniconda3\Scripts;%PATH%
|
||||
|
||||
call activate
|
||||
call conda env list
|
||||
call conda update -y --all
|
||||
call conda create -y -n ezvtb_rt_venv python=3.10
|
||||
call conda activate ezvtb_rt_venv
|
||||
call conda env list
|
||||
|
||||
call conda install -y nvidia/label/cuda-12.9.0::cuda-nvcc-dev_win-64
|
||||
call conda install -y pycuda -c https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge/
|
||||
|
||||
call python -m pip install --upgrade pip wheel -i https://mirrors.aliyun.com/pypi/simple/
|
||||
echo yes|python -m pip install nvidia-cudnn-cu12 -i https://mirrors.aliyun.com/pypi/simple/
|
||||
|
||||
echo yes|pip install tensorrt_cu12_libs==10.11.0.33 tensorrt_cu12_bindings==10.11.0.33 tensorrt==10.11.0.33 --extra-index-url https://pypi.nvidia.com
|
||||
|
||||
call python -m pip install -r requirements.txt --no-warn-script-location -i https://mirrors.aliyun.com/pypi/simple/
|
||||
|
||||
pause
|
||||
@@ -1,11 +0,0 @@
|
||||
@echo off
|
||||
cd /D "%~dp0"
|
||||
IF EXIST %~dp0envs\miniconda3\Scripts SET PATH=%~dp0envs\miniconda3\Scripts;%PATH%
|
||||
call activate
|
||||
call conda env list
|
||||
call conda activate ezvtb_rt_venv
|
||||
call conda env list
|
||||
|
||||
python ezvtb_rt_interface.py
|
||||
|
||||
pause
|
||||
-10
@@ -1,10 +0,0 @@
|
||||
@echo off
|
||||
cd /D "%~dp0"
|
||||
IF EXIST %~dp0envs\miniconda3\Scripts SET PATH=%~dp0envs\miniconda3\Scripts;%PATH%
|
||||
call activate
|
||||
call conda env list
|
||||
call conda activate ezvtb_rt_venv
|
||||
call conda env list
|
||||
|
||||
|
||||
pythonw launcher2.py
|
||||
@@ -1,12 +0,0 @@
|
||||
@echo on
|
||||
cd /D "%~dp0"
|
||||
IF EXIST %~dp0envs\miniconda3\Scripts SET PATH=%~dp0envs\miniconda3\Scripts;%PATH%
|
||||
call activate
|
||||
call conda env list
|
||||
call conda activate ezvtb_rt_venv
|
||||
call conda env list
|
||||
|
||||
|
||||
python launcher2.py
|
||||
|
||||
pause
|
||||
+1
-1
@@ -120,4 +120,4 @@ DirectML 允许所有 Windows 显卡设备以及 AIPC 的 NPU 运行此项目，
|
||||
所以，DirectML运行本项目的任意模型出现奇怪结果，或不符合预期的性能倒退，都应该属于意料之中
|
||||
换个模型精度，观察调试输出是否有DirectML黄字警告 通常可以解决大部分问题。
|
||||
具体某个硬件在DirectML的运行异常修复，已远超本项目的目标范围，请自行验证失效kernel，向硬件厂商和微软求助并期待驱动更新带来解决
|
||||
（建议买张N卡吧别折腾自己了）。
|
||||
（建议买张N卡吧别折腾自己了）。
|
||||
@@ -15,21 +15,22 @@ Updates:
|
||||
* 使用 https://github.com/zpeng11/ezvtuber-rt/tree/main 项目转换的ONNX模型，实现TensorRT加速以及非N卡支持，在半精度和全精度下均有加速（具体取决于显卡）
|
||||
* 为A卡和I卡提供DirectML支持，人人都能玩。
|
||||
* 引入[RIFE](https://github.com/hzwer/ECCV2022-RIFE)模型进行插帧，极限帧数提升达到50%-100%，配合锁帧可以有效降低使用占用。
|
||||
* 使用 [TurboJPEG](https://github.com/Dobatymo/turbojpeg-python) 获得成倍缓存命中率提升，在长时间使用上显卡减负效果相当出色。
|
||||
* 使用 `brotli` 压缩器获得成倍缓存命中率提升，在长时间使用上显卡减负效果相当出色。
|
||||
* 使用 [waifu2x](https://github.com/nagadomi/waifu2x) 和 [real-esrgan](https://github.com/xinntao/Real-ESRGAN) 项目带来的输出超分辨率，对比anime4k效果提升显著（也使用更多gpu占用）
|
||||
* 更漂亮的 WxPython 新启动器界面，带中文介绍更加简单易用
|
||||
* 添加Spout2 支持输出原生透明通道给OBS
|
||||
* 添加 [THA V4](https://github.com/pkhungurn/talking-head-anime-4-demo) 支持，拥有更好的图形输出
|
||||
|
||||
## Requirements
|
||||
|
||||
### 硬件
|
||||
|
||||
- 支持FaceID的iPhone（使用ifacialmocap软件，需购买，需要稳定的WIFI连接）或网络摄像头（使用OpenCV）
|
||||
- 支持无摄像头鼠标捕获/USB网络摄像头/iPhone结构光面捕（ifacialmocap软件，需购买）
|
||||
- 任意5年内的游戏级显卡，N卡/A卡/I卡均可使用，某些性能强的核显也可尝试，请参考[性能测试结果](PerformanceTest.md)
|
||||
### 软件
|
||||
|
||||
- 本方案在Windows 10上测试可用
|
||||
- OBS或Unity Capture（虚拟摄像头方案）
|
||||
- 本方案在Windows 10/11上测试可用
|
||||
- OBS 使用经验
|
||||
- Photoshop或其他图片处理软件
|
||||
- 科学上网方案，看懂英文网站和报错的能力
|
||||
|
||||
@@ -40,100 +41,13 @@ Updates:
|
||||
* [谷歌网盘](https://drive.google.com/drive/folders/1cYj18EfVQ2Cl348_rkCu_fgaasHTI_io?usp=drive_link)
|
||||
|
||||
### 安装Spout2 OBS插件(可选)
|
||||
安装整合包附带的 `OBS_Spout2_Plugin_Install_v1.8.exe`， 安装包来自 https://github.com/Off-World-Live/obs-spout2-plugin 项目，可以提供透明通道给OBS。
|
||||
|
||||
### 构建TensorRT模型（N卡专用，可选）
|
||||
双击`02.构建TensorRT模型（N卡使用）.bat`，若是英伟达用户将会开始TensorRT构建, 每台配置仅需构建一次，更新显卡驱动可能需要重新构建
|
||||
额外需要>20min时间(取决于显卡)，结束后可以大幅提升运行速度（可简单理解为玩游戏之前编译着色器，和每个人显卡有关）
|
||||
构建的开头和结束输出如下，请耐心等进度条百分百或完成后自动退出。
|
||||
若不慎关掉窗口也可以双击自动从断点重来，可以先放后台运行构建并尝试下一步运行暂时不勾选TensorRT：
|
||||
```
|
||||
Start testing if TensorRT works on this machine
|
||||
0%| | 0/38 [00:00<?, ?it/s][01/09/2025-22:36:52] [TRT] [I] [MemUsageChange] Init CUDA: CPU +89, GPU +0, now: CPU 6758, GPU 1009 (MiB)
|
||||
[01/09/2025-22:36:57] [TRT] [I] [MemUsageChange] Init builder kernel library: CPU +1358, GPU +194, now: CPU 8423, GPU 1203 (MiB)
|
||||
[01/09/2025-22:36:57] [TRT] [I] Loading ONNX file from path ./data/models\rife_512\x2\fp32.onnx...
|
||||
[01/09/2025-22:36:57] [TRT] [I] Beginning ONNX file parsing
|
||||
[01/09/2025-22:36:57] [TRT] [I] Completed parsing of ONNX file
|
||||
[01/09/2025-22:36:57] [TRT] [I] Input number: 2
|
||||
[01/09/2025-22:36:57] [TRT] [I] Output number: 2
|
||||
[01/09/2025-22:36:57] [TRT] [I] Building an engine from file ./data/models\rife_512\x2\fp32.onnx; this may take a while...
|
||||
....
|
||||
....
|
||||
....
|
||||
[01/10/2025-21:16:53] [TRT] [I] [BlockAssignment] Started assigning block shifts. This will take 66 steps to complete.
|
||||
[01/10/2025-21:16:53] [TRT] [I] [BlockAssignment] Algorithm ShiftNTopDown took 3.4012ms to assign 7 blocks to 66 nodes requiring 218136064 bytes.
|
||||
[01/10/2025-21:16:53] [TRT] [I] Total Activation Memory: 218135552 bytes
|
||||
[01/10/2025-21:16:53] [TRT] [I] Total Weights Memory: 2724242 bytes
|
||||
[01/10/2025-21:16:53] [TRT] [I] Compiler backend is used during engine execution.
|
||||
[01/10/2025-21:16:53] [TRT] [I] Engine generation completed in 160.706 seconds.
|
||||
[01/10/2025-21:16:53] [TRT] [I] [MemUsageStats] Peak memory usage of TRT CPU/GPU memory allocators: CPU 89 MiB, GPU 317 MiB
|
||||
[01/10/2025-21:16:53] [TRT] [I] Completed creating Engine
|
||||
[01/10/2025-21:16:53] [TRT] [I] Saving engine to file ./data/models\waifu2x_upconv\fp16\upconv_7\art\scale2x.trt
|
||||
[01/10/2025-21:16:53] [TRT] [I] Completed saving engine
|
||||
100%|██████████████████████████████████████████████████████████████████████████████████| 38/38 [46:54<00:00, 74.07s/it]
|
||||
```
|
||||
使用Spout2插件可以提供透明通道给OBS，访问 https://github.com/Off-World-Live/obs-spout2-plugin/releases 项目，请选择与你的OBS版本兼容的插件安装exe。
|
||||
|
||||
### 启动项目
|
||||
双击`02A.启动器.bat` 或 `02B.启动器（调试输出）.bat` 启动器，如下图启动。如果上一步构建未完成或者使用A/I卡，需要**手动**将TensorRT去除勾选：
|
||||
双击`01A.启动器.bat` 或 `01B.启动器（调试输出）.bat` 启动器，如下图启动。
|
||||
|
||||
之后请移步后文的输入输出&性能配置板块自行调节。如有问题请先参考本页末的FAQ再提出。
|
||||
|
||||
## Installation(科学上网且使用Git)
|
||||
可使用此安装方法对本项目二次开发
|
||||
|
||||
### 安装Anaconda（可选）
|
||||
前往https://www.anaconda.com/ 安装Anaconda 并保证加入环境变量命令行可以找到，如：
|
||||
```
|
||||
C:\Users\Eleven>conda --version
|
||||
conda 24.11.3
|
||||
```
|
||||
以上方式可以保证conda环境被你选定的Ananconda统一管理。若没有此需求可以忽略直接运行下一步，`01X.构建运行环境（XX源）.bat`脚本会自动安装miniconda。
|
||||
|
||||
### 克隆项目和子项目
|
||||
```
|
||||
git clone https://github.com/zpeng11/EasyVtuber.git
|
||||
cd EasyVtuber
|
||||
git submodule init
|
||||
git submodule update --recursive --remote
|
||||
```
|
||||
|
||||
### 安装环境
|
||||
根据网络环境双击运行`01A.构建运行环境（默认源）.bat`或者`01B.构建运行环境（国内源）.bat`
|
||||
这两个脚本会构建名为`ezvtb_rt_venv`的conda环境
|
||||
或可以用如下命令手动创建：
|
||||
```
|
||||
conda create -y -n ezvtb_rt_venv python=3.10
|
||||
conda activate ezvtb_rt_venv
|
||||
conda install -y nvidia/label/cuda-12.6.3::cuda-nvcc-dev_win-64
|
||||
conda install -y conda-forge::pycuda
|
||||
python -m pip install --upgrade pip wheel
|
||||
python -m pip install nvidia-cudnn-cu12
|
||||
pip install tensorrt_cu12_libs==10.6.0 tensorrt_cu12_bindings==10.6.0 tensorrt==10.6.0 --extra-index-url https://pypi.nvidia.com
|
||||
python -m pip install -r requirements.txt --no-warn-script-location
|
||||
```
|
||||
|
||||
### 下载模型
|
||||
前往模型 release 地址 https://github.com/zpeng11/ezvtuber-rt/releases/tag/0.0.1
|
||||
下载 https://github.com/zpeng11/ezvtuber-rt/releases/download/0.0.1/20241220.zip 并解压到`data/models`文件夹下
|
||||
|
||||
### Pycharm 配置
|
||||
安装完成后，在Pycharm内打开本项目，右下角解释器菜单点开，`Add Interpreter...`->`Conda Environment`->`Existing environment`
|
||||
选好自己电脑上的`conda.exe`和刚才创建好的`ezvtb_rt_venv`环境内的`python.exe`
|
||||
点击OK，依赖全亮即可
|
||||
|
||||
### 检查环境模型，并构建TensorRT
|
||||
在Conda环境中执行以下命令
|
||||
`python ezvtb_rt_interface.py`
|
||||
|
||||
### 运行启动器
|
||||
在Conda环境中执行以下命令
|
||||
`python launcher2.py`
|
||||
|
||||
### 环境错误排查
|
||||
当运行TensorRT构建或启动器出现错误时，请参考如下可能性：
|
||||
1. 环境安装有错误（一般运行时缺少库99%的原因都是这个，pip下载并不稳定，各种网络问题都可能导致安装失败），检查源，手动删除`envs`文件夹再试, 对照 [此log](assets/complete_building_log.txt) 为成功案例排查原因。 实在不懂排查请使用整合包。
|
||||
2. 英伟达显卡但计算架构低于7.5（10系以及之前）不支持TensorRT
|
||||
|
||||
## 输入输出设备
|
||||
#### SPOUT2 OBS插件输出
|
||||
|
||||
@@ -143,7 +57,7 @@ python -m pip install -r requirements.txt --no-warn-script-location
|
||||
|
||||
#### OBS Virtual Camera 无透明通道
|
||||
|
||||
如果你选择自己进行抠像你可以直接仅输出RGB到OBS，使用补帧时可能也需要此方法
|
||||
次选方案，如果你选择自己进行抠像你可以直接仅输出RGB到OBS，使用补帧时可能也需要此方法
|
||||
使用`视频采集设备`源，设置`滤镜`,
|
||||
在`效果滤镜`添加`色值`并自定义颜色为纯黑色，多尝试调整颜色参数直到抠图合理
|
||||
|
||||
@@ -158,29 +72,21 @@ python -m pip install -r requirements.txt --no-warn-script-location
|
||||
之后对OBS中的视频采集设备添加滤镜-着色器-选择你下载的`merge alpha2.hlsl`-关闭
|
||||
这样透明通道就应用回左边的图像了
|
||||
你可能需要手动调整一下裁剪把右侧的无用画面切掉
|
||||
(看不到着色器滤镜的话就是StreamFX没装好或者OBS不是最新版)
|
||||
(看不到着色器滤镜的话就是StreamFX没装好或者OBS不是最新版)d
|
||||
|
||||
#### UnityCapture（Outdated，性能问题）
|
||||
|
||||
存在未查明的性能瓶颈
|
||||
如果需要使用透明通道输出，参考 https://github.com/schellingb/UnityCapture#installation 安装好UnityCapture
|
||||
只需要正常走完Install.bat，在OBS里能看到对应的设备（Unity Video Capture）就行
|
||||
|
||||
在OBS添加完摄像头以后，还需要手动配置一次摄像头属性才能支持ARGB
|
||||
右键属性-取消激活-分辨率类型自定义-分辨率512x512(与`--output_size`参数一致)-视频格式ARGB-激活
|
||||
|
||||
#### iFacialMocap
|
||||
#### iFacialMocap 面捕
|
||||
|
||||
目前最推荐的方案，使用苹果结构光硬件，捕捉效果追踪准确度都最好，demo中使用了此输入。
|
||||
需要输入`ip:port`,其中`ip`会显示在手机上，`port`使用默认`49983`
|
||||
https://www.ifacialmocap.com/download/
|
||||
你大概率需要购买正式版（非广告，只是试用版不太够时长）
|
||||
购买之前确认好自己的设备支持
|
||||
**不需要下载PC软件**，装好iOS端的软件即可，连接信息通过参数传入Python
|
||||
游戏加速器，梯子等网络工具可能会导致ip失效，请自行解决。
|
||||
购买之前确认好自己的设备支持，装好iOS端的软件即可，连接信息通过参数传入Python
|
||||
一般需要iphone与window在同一局域网下
|
||||
且游戏加速器，梯子等网络工具可能会导致ip失效，请自行解决。
|
||||
|
||||
#### OpenSeeFace
|
||||
#### OpenSeeFace 高精度摄像头面捕
|
||||
|
||||
在没有苹果硬件时最推荐使用的捕捉方案，需要更多配置
|
||||
https://github.com/emilianavt/OpenSeeFace/releases
|
||||
直接下载最新版本的Release包并解压，或使用整合包附带的压缩包
|
||||
之后进入解压目录的Binary文件夹
|
||||
@@ -189,28 +95,34 @@ https://github.com/emilianavt/OpenSeeFace/releases
|
||||
之后保存并双击`run.bat`运行，按照提示选择摄像头、分辨率、帧率，捕获正常的话可以看到输出画面
|
||||
最后在启动器中选择OpenSeeFace输入，或添加启动参数`--osf 127.0.0.1:11573`即可接入OpenSeeFace
|
||||
|
||||
#### OpenCV 低精度摄像头面捕
|
||||
|
||||
此输入不需要配置，会尝试开启默认的USB或笔记本电脑摄像头开始面捕，稳定性和精确度略欠佳。
|
||||
|
||||
#### 鼠标音频输入
|
||||
|
||||
隐私考虑不想使用任何摄像头可使用这个输入，定时循环眨眼和呼吸，配合音频音量控制口型，可以在无摄像头时得到不错的动态。
|
||||
|
||||
## 性能配置(调参必读）
|
||||
### 概念须知
|
||||
|
||||
#### 模型&精度
|
||||
本项目完全基于AI算法生成，由3类AI模型提供服务（缝合怪）：
|
||||
* THA模型：指本项目的基础模型，输入静态图像和姿态信息计算出新图像
|
||||
* THA模型：指本项目的基础模型，输入静态图像和姿态信息计算出新图像，有v3_seperable, v3_standard, v4, v4_student 四种版本。
|
||||
* 插帧/补帧模型：指RIFE插帧模型，输入第1帧和第2帧，输出第1.5帧，类似DLSS3或小黄鸭，此模型是整体性能提升的关键
|
||||
* 超分模型：指waifu2x或realesrgan模型，输入一帧图像，解计算出放大且更清晰的同一帧，显著提升画面效果
|
||||
* 超分模型：指Anime4k, waifu2x或realesrgan模型，输入一帧图像，解算出放大且更清晰的同一帧，显著提升画面效果
|
||||
|
||||
每一种模型提供了半精度(标记为half)和全精度(标记为full)选项，半精度通常运行速度比全精度快，但生图效果稍微逊色，需要根据性能进行取舍。
|
||||
THA模型有另外的`seperable`和`standard`两个版本，可以简单理解为`seperable`是性能更强但生图效果稍差的版本，所以加上两种精度THA模型一共有4种组合。
|
||||
每一种模型提供了半精度(标记为half)和全精度(标记为full)选项，半精度通常运行速度比全精度快，但生图效果稍微逊色，且兼容性可能不好。
|
||||
|
||||
#### 缓存
|
||||
|
||||
当我们当前的姿态与一秒前的姿态完全相同时，输入给模型的参数也完全相同，没有必要再做一次完全相同的运算得出完全相同的结果
|
||||
缓存的原理是将过去一段时间内产生过的结果储存起来。
|
||||
倘若发现姿态相同，就使用储存的结果输出而绕过模型运行（又叫缓存命中），从而减轻显卡负担。
|
||||
本项目同时使用显存VRAM和内存RAM用作缓存目的。在内存中还使用了JPEG格式对缓存的图片进行压缩。
|
||||
倘若发现姿态相同，就使用曾经的结果输出而绕过模型运行（又叫缓存命中），从而减轻显卡负担。
|
||||
本项目同时使用显存VRAM和内存RAM用作缓存目的。在内存中还对缓存的图片进行压缩。
|
||||
缓存大小（单位GB）决定了缓存最多能使用的空间大小。更大的缓存能让缓存命中率更高，长时间使用时显卡占用更小。但也会挤占其他应用或游戏的空间。
|
||||
需要自行理解并取舍。
|
||||
|
||||
#### 输入简化
|
||||
#### 输入简化（量化）
|
||||
输入简化与上文的缓存相辅相成，我们将姿态参数分入一些区间，这样相近的参数就可以命中同一个缓存。
|
||||
举个非常简化的例子：缓存中的一个结果对应姿态参数`0.5`, 在不进行任何输入简化时，需要新输入参数完全相同于`0.5`即 `x = 0.5`才能判定命中并绕过GPU运算。
|
||||
而开启了输入简化后，只需要新输入参数 `x` $\approx$ `0.5` 就可以命中。如`x = 0.49`或 `x = 0.51`.
|
||||
@@ -226,9 +138,9 @@ THA模型有另外的`seperable`和`standard`两个版本，可以简单理解
|
||||
#### 超分
|
||||
|
||||
|
||||
超分由waifu2x和realesrgan模型的两种精度一共4种组合，外加一种Anime4K纯cpu实现（效果比较差）。
|
||||
超分由waifu2x和realesrgan模型的两种精度一共4种组合，外加一种Anime4K实现。
|
||||
通常选用半精度，根据使用半身或大头选择`waifu2x_x2_half`和`realesrgan_x4_half`
|
||||
超分启动时会显著增加显卡占用，通常建议在低负载游戏或者杂谈状态下使用。
|
||||
超分启动时会显著增加显卡占用。
|
||||
|
||||
#### 插帧/补帧
|
||||
|
||||
@@ -240,17 +152,70 @@ RIFE插帧的计算消耗比起一次THA计算要少约50%且画面观感差异
|
||||
|
||||
## 脱离启动器运行本体
|
||||
|
||||
`main.py` 的完全体运行命令请参考 `02B.启动器（调试输出）.bat` 的输出以及 `args.py`。 以下为部分基础参数
|
||||
`main.py` 的完全体运行命令请参考 `02B.启动器（调试输出）.bat` 的输出以及 `args.py`。
|
||||
|
||||
参数名 | 值类型 | 说明
|
||||
:---: | :---: | :---:
|
||||
--character|字符串|`character`目录下的输入图像文件名，不需要带扩展名
|
||||
--debug|无|打开OpenCV预览窗口输出渲染结果，如果没有任何输出配置，该参数默认生效
|
||||
--input|字符串|不使用iOS面捕时，传入要使用的摄像头设备名称，默认为设备0，有ifm参数时无效
|
||||
--ifm|字符串|使用iOS面捕时，传入设备的`IP:端口号`，如`192.168.31.182:49983`
|
||||
--output_webcam|字符串|可用值为`obs` `unitycapture`，选择对应的输出种类，不传不输出到摄像头
|
||||
--extend_movement|浮点数|使用iOS面捕返回的头部位置，对模型输出图像进一步进行移动和旋转使得上半身可动<br>传入的数值表示移动倍率（建议值为1）
|
||||
--output_size|字符串|格式为`256x256`，必须是4的倍数。<br>增大它并不会让图像更清晰，但配合extend_movement会增大可动范围
|
||||
## Installation(使用整合包请忽略)
|
||||
可使用此安装方法对本项目二次开发
|
||||
|
||||
### 安装Anaconda（可选）
|
||||
前往https://www.anaconda.com 安装Anaconda 并保证加入环境变量命令行可以找到，如：
|
||||
```
|
||||
C:\Users\Eleven>conda --version
|
||||
conda 24.11.3
|
||||
```
|
||||
以上方式可以保证conda环境被你选定的Ananconda统一管理。若没有此需求可以忽略直接运行下一步。
|
||||
|
||||
### 下载 TensorRT-RTX （Nvidia开发必选）
|
||||
```
|
||||
curl -L -o trt_rtx.zip https://developer.nvidia.com/downloads/trt/rtx_sdk/secure/1.3/TensorRT-RTX-1.3.0.35-win10-amd64-cuda-12.9-Release-external.zip
|
||||
tar -xf trt_rtx.zip "TensorRT-RTX-1.3.0.35" && ren "TensorRT-RTX-1.3.0.35" "TensorRT-RTX-1.3.0.35_cu129" && del trt_rtx.zip
|
||||
# Adding the bin folder to PATH
|
||||
cd TensorRT-RTX-1.3.0.35_cu129\bin && set PATH=%CD%;%PATH%
|
||||
```
|
||||
也可以在Window 环境变量管理器 GUI中将`TensorRT-RTX-1.3.0.35_cu129\bin`在`PATH`环境变量中添加来持久化
|
||||
|
||||
### 准备Conda 环境
|
||||
```
|
||||
conda create -y -n easyvtuber python=3.10
|
||||
conda activate easyvtuber
|
||||
conda install conda-forge::pycuda
|
||||
conda install -c nvidia/label/cuda-12.9.1 cuda-nvcc-dev_win-64 cudnn cuda-runtime
|
||||
```
|
||||
|
||||
### 安装 TensorRT-RTX 的 Python Binding
|
||||
```
|
||||
pip install TensorRT-RTX-1.3.0.35_cu129\python\tensorrt_rtx-1.3.0.35-cp310-none-win_amd64.whl
|
||||
```
|
||||
|
||||
### 克隆项目和子项目
|
||||
```
|
||||
git clone https://github.com/yuyuyzl/EasyVtuber.git
|
||||
cd EasyVtuber
|
||||
git submodule init
|
||||
git submodule update --recursive --remote
|
||||
```
|
||||
|
||||
### 安装环境
|
||||
```
|
||||
pip install -r requirements.txt --no-warn-script-location
|
||||
```
|
||||
|
||||
### 下载模型
|
||||
前往模型 release 地址 https://drive.google.com/file/d/1pWKIpjWeqfpa3Rub185FVvxDr5H09pOi/view?usp=drive_link 并解压到`data/models`文件夹下
|
||||
|
||||
### Pycharm 配置（可选）
|
||||
安装完成后，在Pycharm内打开本项目，右下角解释器菜单点开，`Add Interpreter...`->`Conda Environment`->`Existing environment`
|
||||
选好自己电脑上的`conda.exe`和刚才创建好的`easyvtuber`环境内的`python.exe`
|
||||
点击OK，依赖全亮即可
|
||||
|
||||
### 运行启动器
|
||||
在Conda环境中执行以下命令
|
||||
`python launcher2.py`
|
||||
|
||||
### 环境错误排查
|
||||
当运行TensorRT构建或启动器出现错误时，请参考如下可能性：
|
||||
1. 环境安装有错误（一般运行时缺少库99%的原因都是这个，pip下载并不稳定，各种网络问题都可能导致安装失败），请检查源。
|
||||
2. 英伟达显卡但计算架构低于7.5（10系以及之前）不支持TensorRT
|
||||
|
||||
## FAQ
|
||||
|
||||
@@ -258,9 +223,12 @@ RIFE插帧的计算消耗比起一次THA计算要少约50%且画面观感差异
|
||||
|
||||
A卡I卡使用 DirectML，这个框架与显卡厂商的显卡驱动及 Direct12 实现相关，很多算子实现存在数值错误或未实现并行优化被迫降级到cpu运行。方案：1. 将显卡驱动更新至最新，2. 打开调试启动器运行并检查模型或算子是否被降级到cpu实现，3.阅读上文的模型调参环节，切换不同模型和精度多尝试。 如有兴趣了解更多请参考[性能测试结果](PerformanceTest.md)末尾对DirectML的思考。
|
||||
|
||||
> Q2: 我开启超分辨率后，为什么补帧没效果了？
|
||||
> Q2: 为什么无法找到我自己准备的Vtuber图像/使用自己的图像会报错？
|
||||
|
||||
v0.5的实现中，补帧和超分只能二选一，开启超分后补帧会自动失效。欢迎重新在上文中更新下载v0.7之后的版本。
|
||||
1. 确认图像改为512x512的png，带Alpha透明通道即32比特位宽。
|
||||
2. 确认放置位置为`data/images`，在选项中应当能看到文件名。
|
||||
3. 仍然无法正确启动请尝试删除`launch.json`文件并重新启动软件。
|
||||
4. 99%的错误情况都来自于准备的图像不满足模型需求
|
||||
|
||||
> Q3: 我开启补帧后为什么在 Spout2 输出中边缘会有些抖动？
|
||||
|
||||
@@ -270,16 +238,30 @@ v0.5的实现中，补帧和超分只能二选一，开启超分后补帧会自
|
||||
|
||||
请在GPU页面中点击`Video Decode`或`Video Processing`按钮可以看到下拉，请对这里的选项挨个尝试找到正确的程序负载，一般在 `cuda` 或 `3D` 或 `Compute_X` 或 `Graphics_X`。
|
||||
|
||||
> Q5: 我的配置是 CPU带核显加独立显卡，这个核显性能挺不错，打游戏时还用不上，能否把项目跑到这上面？
|
||||
> Q6: 本机有两张显卡，如何使用第二张副卡使用这个项目？
|
||||
|
||||
完全可行，其实界面中留显卡选项有部分原因就是这方面需求。具体操作如下：去除TensorRT勾选，开启设备管理器观察核显占用，从`1`开始挨个尝试`GPU Device`选项并启动，直到观察到核显开始占用上升即找到了核显对应的编号。（os：DirectML的API就是这么难用没办法）接下来有可能出现Q1中提及的问题，请参照上文解决。
|
||||
因为性能问题很有可能需要使用插帧来提高输出帧数。目前核显相关测试样本较少，欢迎分享运行效果。
|
||||
在运行前配置环境变量`EZVTB_DEVICE_ID`为你想要运行的GPU ID,此变量缺省为`0`
|
||||
|
||||
> Q6: 我主机上插着两张英伟达显卡，如何使用第二张副卡使用这个项目？
|
||||
1. 对第二张卡构建TenosrRT：更改`02.构建TensorRT模型（N卡开启TensorRT使用）.bat` 中的 `python ezvtb_rt_interface.py` 一行，加上空格接参数cuda设备号，例如我的副卡设备号是`1`所以改成 `python ezvtb_rt_interface.py 1` 然后双击运行。
|
||||
2. 在启动器界面中选择相同的设备号。
|
||||
## References
|
||||
```
|
||||
@inproceedings{huang2022rife,
|
||||
title={Real-Time Intermediate Flow Estimation for Video Frame Interpolation},
|
||||
author={Huang, Zhewei and Zhang, Tianyuan and Heng, Wen and Shi, Boxin and Zhou, Shuchang},
|
||||
booktitle={Proceedings of the European Conference on Computer Vision (ECCV)},
|
||||
year={2022}
|
||||
}
|
||||
@misc{Khungurn:2022,
|
||||
author = {Pramook Khungurn},
|
||||
title = {Talking Head(?) Anime from a Single Image 3: Now the Body Too},
|
||||
howpublished = {\url{http://pkhungurn.github.io/talking-head-anime-3/}},
|
||||
year = 2022,
|
||||
note = {Accessed: YYYY-MM-DD},
|
||||
}
|
||||
|
||||
> Q7: 为什么无法使用我自己准备的Vtuber图像/使用自己的图像会报错？
|
||||
1. 确认图像改为512x512的png，带Alpha透明通道即32比特位宽。
|
||||
2. 确认放置位置为`data/images`，在选项中应当能看到文件名。
|
||||
3. 仍然无法正确启动请尝试删除`launch.json`文件并重新启动软件。
|
||||
@InProceedings{wang2021realesrgan,
|
||||
author = {Xintao Wang and Liangbin Xie and Chao Dong and Ying Shan},
|
||||
title = {Real-ESRGAN: Training Real-World Blind Super-Resolution with Pure Synthetic Data},
|
||||
booktitle = {International Conference on Computer Vision Workshops (ICCVW)},
|
||||
date = {2021}
|
||||
}
|
||||
```
|
||||
@@ -1,82 +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('--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,548 +0,0 @@
|
||||
C:\Users\Eleven\Desktop\EasyVtuber-main>cd /D "C:\Users\Eleven\Desktop\EasyVtuber-main\"
|
||||
|
||||
C:\Users\Eleven\Desktop\EasyVtuber-main>WHERE conda
|
||||
INFO: Could not find files for the given pattern(s).
|
||||
|
||||
C:\Users\Eleven\Desktop\EasyVtuber-main>IF 1 NEQ 0 (
|
||||
if not exist "envs" mkdir envs
|
||||
if not exist "envs\miniconda3" mkdir envs\miniconda3
|
||||
IF not EXIST C:\Users\Eleven\Desktop\EasyVtuber-main\envs\miniconda3\Scripts (
|
||||
|
||||
mkdir C:\Users\Eleven\Desktop\EasyVtuber-main\envs\miniconda3
|
||||
echo "Downloading miniconda..."
|
||||
powershell -Command "(New-Object Net.WebClient).DownloadFile('https://repo.anaconda.com/miniconda/Miniconda3-py310_24.9.2-0-Windows-x86_64.exe', '.\envs\miniconda3.exe')"
|
||||
echo "Installling minconda..."
|
||||
start /wait "" C:\Users\Eleven\Desktop\EasyVtuber-main\envs\miniconda3.exe /S /AddToPath=0 /RegisterPython=0 /InstallationType=JustMe /D=C:\Users\Eleven\Desktop\EasyVtuber-main\envs\miniconda3
|
||||
echo "Successfully install minconda"
|
||||
del envs\miniconda3.exe
|
||||
)
|
||||
)
|
||||
"Downloading miniconda..."
|
||||
"Installling minconda..."
|
||||
"Successfully install minconda"
|
||||
|
||||
C:\Users\Eleven\Desktop\EasyVtuber-main>IF EXIST C:\Users\Eleven\Desktop\EasyVtuber-main\envs\miniconda3\Scripts SET PATH=C:\Users\Eleven\Desktop\EasyVtuber-main\envs\miniconda3\Scripts;C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.6\bin;C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.6\libnvvp;;C:\Windows\system32;C:\Windows;C:\Windows\System32\Wbem;C:\Windows\System32\WindowsPowerShell\v1.0\;C:\Windows\System32\OpenSSH\;C:\Program Files (x86)\NVIDIA Corporation\PhysX\Common;C:\Program Files\Git\cmd;C:\Program Files\Tailscale\;C:\Program Files\NVIDIA Corporation\Nsight Compute 2024.3.2\;C:\Program Files\NVIDIA Corporation\NVIDIA NvDLISR;C:\Program Files\dotnet\;C:\Users\Eleven\AppData\Local\Programs\Python\Python312\Scripts\;C:\Users\Eleven\AppData\Local\Programs\Python\Python312\;C:\Users\Eleven\.cargo\bin;C:\Users\Eleven\AppData\Local\Microsoft\WindowsApps;C:\Users\Eleven\AppData\Local\Programs\Microsoft VS Code\bin;C:\Program Files\Git\bin;D:\ffmpeg\bin;
|
||||
|
||||
C:\Users\Eleven\Desktop\EasyVtuber-main>call activate
|
||||
# conda environments:
|
||||
#
|
||||
C:\EasyVtuber\envs\miniconda3
|
||||
C:\EasyVtuber\envs\miniconda3\envs\ezvtb_rt_venv
|
||||
base * C:\Users\Eleven\Desktop\EasyVtuber-main\envs\miniconda3
|
||||
|
||||
Retrieving notices: ...working... done
|
||||
Channels:
|
||||
- defaults
|
||||
Platform: win-64
|
||||
Collecting package metadata (repodata.json): done
|
||||
Solving environment: done
|
||||
|
||||
## Package Plan ##
|
||||
|
||||
environment location: C:\Users\Eleven\Desktop\EasyVtuber-main\envs\miniconda3
|
||||
|
||||
|
||||
The following packages will be downloaded:
|
||||
|
||||
package | build
|
||||
---------------------------|-----------------
|
||||
anaconda-anon-usage-0.5.0 |py310hfc23b7f_100 28 KB
|
||||
brotli-python-1.0.9 | py310h5da7b33_9 346 KB
|
||||
ca-certificates-2024.11.26 | haa95532_0 132 KB
|
||||
certifi-2024.12.14 | py310haa95532_0 162 KB
|
||||
cffi-1.17.1 | py310h827c3e9_1 254 KB
|
||||
conda-24.11.2 | py310haa95532_0 948 KB
|
||||
conda-libmamba-solver-24.11.1| pyhd3eb1b0_0 39 KB
|
||||
conda-package-handling-2.4.0| py310haa95532_0 297 KB
|
||||
conda-package-streaming-0.11.0| py310haa95532_0 31 KB
|
||||
cpp-expected-1.1.0 | h214f63a_0 131 KB
|
||||
cryptography-43.0.3 | py310hbd6ee87_1 1.3 MB
|
||||
libarchive-3.7.7 | h9243413_0 1.8 MB
|
||||
libcurl-8.11.1 | haff574d_0 382 KB
|
||||
libmamba-2.0.5 | hcd6fe79_1 4.7 MB
|
||||
libmambapy-2.0.5 | py310h214f63a_1 459 KB
|
||||
libsolv-0.7.30 | hf2fb9eb_1 473 KB
|
||||
libssh2-1.11.1 | h2addb87_0 312 KB
|
||||
libxml2-2.13.5 | h24da03e_0 2.9 MB
|
||||
menuinst-2.2.0 | py310h5da7b33_0 200 KB
|
||||
nlohmann_json-3.11.2 | h6c2663c_0 126 KB
|
||||
packaging-24.2 | py310haa95532_0 175 KB
|
||||
pluggy-1.5.0 | py310haa95532_0 39 KB
|
||||
requests-2.32.3 | py310haa95532_1 101 KB
|
||||
simdjson-3.10.1 | h214f63a_0 275 KB
|
||||
spdlog-1.11.0 | h59b6b97_0 195 KB
|
||||
vc-14.40 | haa95532_2 10 KB
|
||||
vs2015_runtime-14.42.34433 | h9531ae6_2 1.2 MB
|
||||
zstandard-0.23.0 | py310h4fc1ca9_1 342 KB
|
||||
------------------------------------------------------------
|
||||
Total: 17.3 MB
|
||||
|
||||
The following NEW packages will be INSTALLED:
|
||||
|
||||
cpp-expected pkgs/main/win-64::cpp-expected-1.1.0-h214f63a_0
|
||||
nlohmann_json pkgs/main/win-64::nlohmann_json-3.11.2-h6c2663c_0
|
||||
simdjson pkgs/main/win-64::simdjson-3.10.1-h214f63a_0
|
||||
spdlog pkgs/main/win-64::spdlog-1.11.0-h59b6b97_0
|
||||
|
||||
The following packages will be UPDATED:
|
||||
|
||||
anaconda-anon-usa~ 0.4.4-py310hfc23b7f_100 --> 0.5.0-py310hfc23b7f_100
|
||||
brotli-python 1.0.9-py310hd77b12b_8 --> 1.0.9-py310h5da7b33_9
|
||||
ca-certificates 2024.9.24-haa95532_0 --> 2024.11.26-haa95532_0
|
||||
certifi 2024.8.30-py310haa95532_0 --> 2024.12.14-py310haa95532_0
|
||||
cffi 1.17.1-py310h827c3e9_0 --> 1.17.1-py310h827c3e9_1
|
||||
conda 24.9.2-py310haa95532_0 --> 24.11.2-py310haa95532_0
|
||||
conda-libmamba-so~ 24.9.0-pyhd3eb1b0_0 --> 24.11.1-pyhd3eb1b0_0
|
||||
conda-package-han~ 2.3.0-py310haa95532_0 --> 2.4.0-py310haa95532_0
|
||||
conda-package-str~ 0.10.0-py310haa95532_0 --> 0.11.0-py310haa95532_0
|
||||
cryptography 43.0.0-py310h89fc84f_0 --> 43.0.3-py310hbd6ee87_1
|
||||
libarchive 3.7.4-h9243413_0 --> 3.7.7-h9243413_0
|
||||
libcurl 8.9.1-h0416ee5_0 --> 8.11.1-haff574d_0
|
||||
libmamba 1.5.8-h99b1521_3 --> 2.0.5-hcd6fe79_1
|
||||
libmambapy 1.5.8-py310h77c03ed_3 --> 2.0.5-py310h214f63a_1
|
||||
libsolv 0.7.24-h23ce68f_1 --> 0.7.30-hf2fb9eb_1
|
||||
libssh2 1.11.0-h291bd65_0 --> 1.11.1-h2addb87_0
|
||||
libxml2 2.13.1-h24da03e_2 --> 2.13.5-h24da03e_0
|
||||
menuinst 2.1.2-py310h5da7b33_0 --> 2.2.0-py310h5da7b33_0
|
||||
packaging 24.1-py310haa95532_0 --> 24.2-py310haa95532_0
|
||||
pluggy 1.0.0-py310haa95532_1 --> 1.5.0-py310haa95532_0
|
||||
requests 2.32.3-py310haa95532_0 --> 2.32.3-py310haa95532_1
|
||||
vc 14.40-h2eaa2aa_1 --> 14.40-haa95532_2
|
||||
vs2015_runtime 14.40.33807-h98bb1dd_1 --> 14.42.34433-h9531ae6_2
|
||||
zstandard 0.23.0-py310h4fc1ca9_0 --> 0.23.0-py310h4fc1ca9_1
|
||||
|
||||
|
||||
|
||||
Downloading and Extracting Packages:
|
||||
|
||||
Preparing transaction: done
|
||||
Verifying transaction: done
|
||||
Executing transaction: done
|
||||
Channels:
|
||||
- defaults
|
||||
Platform: win-64
|
||||
Collecting package metadata (repodata.json): done
|
||||
Solving environment: done
|
||||
|
||||
## Package Plan ##
|
||||
|
||||
environment location: C:\Users\Eleven\Desktop\EasyVtuber-main\envs\miniconda3\envs\ezvtb_rt_venv
|
||||
|
||||
added / updated specs:
|
||||
- python=3.10
|
||||
|
||||
|
||||
The following packages will be downloaded:
|
||||
|
||||
package | build
|
||||
---------------------------|-----------------
|
||||
python-3.10.16 | h4607a30_1 16.3 MB
|
||||
------------------------------------------------------------
|
||||
Total: 16.3 MB
|
||||
|
||||
The following NEW packages will be INSTALLED:
|
||||
|
||||
bzip2 pkgs/main/win-64::bzip2-1.0.8-h2bbff1b_6
|
||||
ca-certificates pkgs/main/win-64::ca-certificates-2024.11.26-haa95532_0
|
||||
libffi pkgs/main/win-64::libffi-3.4.4-hd77b12b_1
|
||||
openssl pkgs/main/win-64::openssl-3.0.15-h827c3e9_0
|
||||
pip pkgs/main/win-64::pip-24.2-py310haa95532_0
|
||||
python pkgs/main/win-64::python-3.10.16-h4607a30_1
|
||||
setuptools pkgs/main/win-64::setuptools-75.1.0-py310haa95532_0
|
||||
sqlite pkgs/main/win-64::sqlite-3.45.3-h2bbff1b_0
|
||||
tk pkgs/main/win-64::tk-8.6.14-h0416ee5_0
|
||||
tzdata pkgs/main/noarch::tzdata-2024b-h04d1e81_0
|
||||
vc pkgs/main/win-64::vc-14.40-haa95532_2
|
||||
vs2015_runtime pkgs/main/win-64::vs2015_runtime-14.42.34433-h9531ae6_2
|
||||
wheel pkgs/main/win-64::wheel-0.44.0-py310haa95532_0
|
||||
xz pkgs/main/win-64::xz-5.4.6-h8cc25b3_1
|
||||
zlib pkgs/main/win-64::zlib-1.2.13-h8cc25b3_1
|
||||
|
||||
|
||||
|
||||
Downloading and Extracting Packages:
|
||||
|
||||
Preparing transaction: done
|
||||
Verifying transaction: done
|
||||
Executing transaction: done
|
||||
#
|
||||
# To activate this environment, use
|
||||
#
|
||||
# $ conda activate ezvtb_rt_venv
|
||||
#
|
||||
# To deactivate an active environment, use
|
||||
#
|
||||
# $ conda deactivate
|
||||
|
||||
|
||||
# conda environments:
|
||||
#
|
||||
C:\EasyVtuber\envs\miniconda3
|
||||
C:\EasyVtuber\envs\miniconda3\envs\ezvtb_rt_venv
|
||||
base C:\Users\Eleven\Desktop\EasyVtuber-main\envs\miniconda3
|
||||
ezvtb_rt_venv * C:\Users\Eleven\Desktop\EasyVtuber-main\envs\miniconda3\envs\ezvtb_rt_venv
|
||||
|
||||
Channels:
|
||||
- https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge
|
||||
- defaults
|
||||
Platform: win-64
|
||||
Collecting package metadata (repodata.json): done
|
||||
Solving environment: done
|
||||
|
||||
## Package Plan ##
|
||||
|
||||
environment location: C:\Users\Eleven\Desktop\EasyVtuber-main\envs\miniconda3\envs\ezvtb_rt_venv
|
||||
|
||||
added / updated specs:
|
||||
- pycuda
|
||||
|
||||
|
||||
The following packages will be downloaded:
|
||||
|
||||
package | build
|
||||
---------------------------|-----------------
|
||||
appdirs-1.4.4 | pyhd8ed1ab_1 14 KB https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge
|
||||
boost-1.78.0 | py310h220cb41_4 555 KB https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge
|
||||
boost-cpp-1.78.0 | h5b4e17d_0 17.0 MB https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge
|
||||
ca-certificates-2024.12.14 | h56e8100_0 154 KB https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge
|
||||
cudatoolkit-11.8.0 | h09e9e62_11 638.9 MB https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge
|
||||
importlib-metadata-8.5.0 | pyha770c72_1 28 KB https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge
|
||||
intel-openmp-2024.2.1 | h57928b3_1083 1.8 MB https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge
|
||||
libblas-3.9.0 | 26_win64_mkl 3.6 MB https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge
|
||||
libcblas-3.9.0 | 26_win64_mkl 3.6 MB https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge
|
||||
libhwloc-2.9.1 | h51c2c0f_0 2.4 MB https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge
|
||||
libiconv-1.17 | h8ffe710_0 698 KB https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge
|
||||
liblapack-3.9.0 | 26_win64_mkl 3.6 MB https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge
|
||||
lz4-c-1.9.4 | hcfcfb64_0 131 KB https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge
|
||||
mako-1.3.8 | pyhd8ed1ab_0 66 KB https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge
|
||||
markupsafe-2.1.2 | py310h8d17308_0 25 KB https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge
|
||||
mkl-2024.2.2 | h66d3029_15 98.3 MB https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge
|
||||
numpy-1.24.3 | py310hd02465a_0 5.7 MB https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge
|
||||
openssl-3.1.0 | hcfcfb64_3 7.1 MB https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge
|
||||
platformdirs-4.3.6 | pyhd8ed1ab_1 20 KB https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge
|
||||
pthreads-win32-2.9.1 | hfa6e2cd_3 141 KB https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge
|
||||
pycuda-2022.2.2 | py310ha2c4f5d_0 378 KB https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge
|
||||
python_abi-3.10 | 2_cp310 4 KB https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge
|
||||
pytools-2024.1.6 | pyhd8ed1ab_0 72 KB https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge
|
||||
tbb-2021.9.0 | h91493d7_0 151 KB https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge
|
||||
typing_extensions-4.12.2 | pyha770c72_1 39 KB https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge
|
||||
ucrt-10.0.22621.0 | h57928b3_1 547 KB https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge
|
||||
zipp-3.21.0 | pyhd8ed1ab_1 21 KB https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge
|
||||
------------------------------------------------------------
|
||||
Total: 784.9 MB
|
||||
|
||||
The following NEW packages will be INSTALLED:
|
||||
|
||||
appdirs anaconda/cloud/conda-forge/noarch::appdirs-1.4.4-pyhd8ed1ab_1
|
||||
boost anaconda/cloud/conda-forge/win-64::boost-1.78.0-py310h220cb41_4
|
||||
boost-cpp anaconda/cloud/conda-forge/win-64::boost-cpp-1.78.0-h5b4e17d_0
|
||||
cudatoolkit anaconda/cloud/conda-forge/win-64::cudatoolkit-11.8.0-h09e9e62_11
|
||||
importlib-metadata anaconda/cloud/conda-forge/noarch::importlib-metadata-8.5.0-pyha770c72_1
|
||||
intel-openmp anaconda/cloud/conda-forge/win-64::intel-openmp-2024.2.1-h57928b3_1083
|
||||
libblas anaconda/cloud/conda-forge/win-64::libblas-3.9.0-26_win64_mkl
|
||||
libcblas anaconda/cloud/conda-forge/win-64::libcblas-3.9.0-26_win64_mkl
|
||||
libhwloc anaconda/cloud/conda-forge/win-64::libhwloc-2.9.1-h51c2c0f_0
|
||||
libiconv anaconda/cloud/conda-forge/win-64::libiconv-1.17-h8ffe710_0
|
||||
liblapack anaconda/cloud/conda-forge/win-64::liblapack-3.9.0-26_win64_mkl
|
||||
libxml2 pkgs/main/win-64::libxml2-2.13.5-h24da03e_0
|
||||
lz4-c anaconda/cloud/conda-forge/win-64::lz4-c-1.9.4-hcfcfb64_0
|
||||
mako anaconda/cloud/conda-forge/noarch::mako-1.3.8-pyhd8ed1ab_0
|
||||
markupsafe anaconda/cloud/conda-forge/win-64::markupsafe-2.1.2-py310h8d17308_0
|
||||
mkl anaconda/cloud/conda-forge/win-64::mkl-2024.2.2-h66d3029_15
|
||||
numpy anaconda/cloud/conda-forge/win-64::numpy-1.24.3-py310hd02465a_0
|
||||
platformdirs anaconda/cloud/conda-forge/noarch::platformdirs-4.3.6-pyhd8ed1ab_1
|
||||
pthreads-win32 anaconda/cloud/conda-forge/win-64::pthreads-win32-2.9.1-hfa6e2cd_3
|
||||
pycuda anaconda/cloud/conda-forge/win-64::pycuda-2022.2.2-py310ha2c4f5d_0
|
||||
python_abi anaconda/cloud/conda-forge/win-64::python_abi-3.10-2_cp310
|
||||
pytools anaconda/cloud/conda-forge/noarch::pytools-2024.1.6-pyhd8ed1ab_0
|
||||
tbb anaconda/cloud/conda-forge/win-64::tbb-2021.9.0-h91493d7_0
|
||||
typing_extensions anaconda/cloud/conda-forge/noarch::typing_extensions-4.12.2-pyha770c72_1
|
||||
ucrt anaconda/cloud/conda-forge/win-64::ucrt-10.0.22621.0-h57928b3_1
|
||||
zipp anaconda/cloud/conda-forge/noarch::zipp-3.21.0-pyhd8ed1ab_1
|
||||
zstd pkgs/main/win-64::zstd-1.5.6-h8880b57_0
|
||||
|
||||
The following packages will be UPDATED:
|
||||
|
||||
ca-certificates pkgs/main::ca-certificates-2024.11.26~ --> anaconda/cloud/conda-forge::ca-certificates-2024.12.14-h56e8100_0
|
||||
openssl pkgs/main::openssl-3.0.15-h827c3e9_0 --> anaconda/cloud/conda-forge::openssl-3.1.0-hcfcfb64_3
|
||||
|
||||
|
||||
|
||||
Downloading and Extracting Packages:
|
||||
|
||||
Preparing transaction: done
|
||||
Verifying transaction: done
|
||||
Executing transaction: | "By downloading and using the CUDA Toolkit conda packages, you accept the terms and conditions of the CUDA End User License Agreement (EULA): https://docs.nvidia.com/cuda/eula/index.html"
|
||||
|
||||
done
|
||||
Channels:
|
||||
- https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main
|
||||
- defaults
|
||||
Platform: win-64
|
||||
Collecting package metadata (repodata.json): done
|
||||
Solving environment: done
|
||||
|
||||
## Package Plan ##
|
||||
|
||||
environment location: C:\Users\Eleven\Desktop\EasyVtuber-main\envs\miniconda3\envs\ezvtb_rt_venv
|
||||
|
||||
added / updated specs:
|
||||
- pytorch-cpu
|
||||
|
||||
|
||||
The following packages will be downloaded:
|
||||
|
||||
package | build
|
||||
---------------------------|-----------------
|
||||
blas-1.0 | mkl 6 KB https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main
|
||||
filelock-3.13.1 | py310haa95532_0 22 KB https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main
|
||||
fsspec-2024.6.1 | py310haa95532_0 286 KB https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main
|
||||
gmpy2-2.1.2 | py310h7f96b67_0 160 KB https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main
|
||||
intel-openmp-2023.1.0 | h59b6b97_46320 2.7 MB https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main
|
||||
jinja2-3.1.4 | py310haa95532_1 281 KB https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main
|
||||
libuv-1.48.0 | h827c3e9_0 322 KB https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main
|
||||
mkl-2023.1.0 | h6b88ed4_46358 155.9 MB https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main
|
||||
mkl-service-2.4.0 | py310h2bbff1b_1 44 KB https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main
|
||||
mkl_fft-1.3.11 | py310h827c3e9_0 168 KB https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main
|
||||
mkl_random-1.2.8 | py310hc64d2fc_0 257 KB https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main
|
||||
mpc-1.1.0 | h7edee0f_1 260 KB https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main
|
||||
mpfr-4.0.2 | h62dcd97_1 1.5 MB https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main
|
||||
mpir-3.0.0 | hec2e145_1 1.3 MB https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main
|
||||
mpmath-1.3.0 | py310haa95532_0 834 KB https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main
|
||||
networkx-3.3 | py310haa95532_0 2.5 MB https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main
|
||||
numpy-1.26.4 | py310h055cbcc_0 11 KB https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main
|
||||
numpy-base-1.26.4 | py310h65a83cf_0 8.6 MB https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main
|
||||
pytorch-2.3.1 |cpu_py310h0ce1571_0 111.2 MB https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main
|
||||
pytorch-cpu-2.3.1 | haa95532_0 21 KB https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main
|
||||
sympy-1.13.3 | py310haa95532_0 11.3 MB https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main
|
||||
------------------------------------------------------------
|
||||
Total: 297.6 MB
|
||||
|
||||
The following NEW packages will be INSTALLED:
|
||||
|
||||
blas anaconda/pkgs/main/win-64::blas-1.0-mkl
|
||||
filelock anaconda/pkgs/main/win-64::filelock-3.13.1-py310haa95532_0
|
||||
fsspec anaconda/pkgs/main/win-64::fsspec-2024.6.1-py310haa95532_0
|
||||
gmpy2 anaconda/pkgs/main/win-64::gmpy2-2.1.2-py310h7f96b67_0
|
||||
jinja2 anaconda/pkgs/main/win-64::jinja2-3.1.4-py310haa95532_1
|
||||
libuv anaconda/pkgs/main/win-64::libuv-1.48.0-h827c3e9_0
|
||||
mkl-service anaconda/pkgs/main/win-64::mkl-service-2.4.0-py310h2bbff1b_1
|
||||
mkl_fft anaconda/pkgs/main/win-64::mkl_fft-1.3.11-py310h827c3e9_0
|
||||
mkl_random anaconda/pkgs/main/win-64::mkl_random-1.2.8-py310hc64d2fc_0
|
||||
mpc anaconda/pkgs/main/win-64::mpc-1.1.0-h7edee0f_1
|
||||
mpfr anaconda/pkgs/main/win-64::mpfr-4.0.2-h62dcd97_1
|
||||
mpir anaconda/pkgs/main/win-64::mpir-3.0.0-hec2e145_1
|
||||
mpmath anaconda/pkgs/main/win-64::mpmath-1.3.0-py310haa95532_0
|
||||
networkx anaconda/pkgs/main/win-64::networkx-3.3-py310haa95532_0
|
||||
numpy-base anaconda/pkgs/main/win-64::numpy-base-1.26.4-py310h65a83cf_0
|
||||
pytorch anaconda/pkgs/main/win-64::pytorch-2.3.1-cpu_py310h0ce1571_0
|
||||
pytorch-cpu anaconda/pkgs/main/win-64::pytorch-cpu-2.3.1-haa95532_0
|
||||
sympy anaconda/pkgs/main/win-64::sympy-1.13.3-py310haa95532_0
|
||||
|
||||
The following packages will be REMOVED:
|
||||
|
||||
libblas-3.9.0-26_win64_mkl
|
||||
libcblas-3.9.0-26_win64_mkl
|
||||
liblapack-3.9.0-26_win64_mkl
|
||||
|
||||
The following packages will be UPDATED:
|
||||
|
||||
numpy anaconda/cloud/conda-forge::numpy-1.2~ --> anaconda/pkgs/main::numpy-1.26.4-py310h055cbcc_0
|
||||
|
||||
The following packages will be SUPERSEDED by a higher-priority channel:
|
||||
|
||||
intel-openmp anaconda/cloud/conda-forge::intel-ope~ --> anaconda/pkgs/main::intel-openmp-2023.1.0-h59b6b97_46320
|
||||
mkl anaconda/cloud/conda-forge::mkl-2024.~ --> anaconda/pkgs/main::mkl-2023.1.0-h6b88ed4_46358
|
||||
|
||||
|
||||
|
||||
Downloading and Extracting Packages:
|
||||
|
||||
Preparing transaction: done
|
||||
Verifying transaction: done
|
||||
Executing transaction: done
|
||||
Looking in indexes: https://mirrors.aliyun.com/pypi/simple/
|
||||
Requirement already satisfied: pip in c:\users\eleven\desktop\easyvtuber-main\envs\miniconda3\envs\ezvtb_rt_venv\lib\site-packages (24.2)
|
||||
Collecting pip
|
||||
Downloading https://mirrors.aliyun.com/pypi/packages/ef/7d/500c9ad20238fcfcb4cb9243eede163594d7020ce87bd9610c9e02771876/pip-24.3.1-py3-none-any.whl (1.8 MB)
|
||||
---------------------------------------- 1.8/1.8 MB 1.4 MB/s eta 0:00:00
|
||||
Requirement already satisfied: wheel in c:\users\eleven\desktop\easyvtuber-main\envs\miniconda3\envs\ezvtb_rt_venv\lib\site-packages (0.44.0)
|
||||
Collecting wheel
|
||||
Downloading https://mirrors.aliyun.com/pypi/packages/0b/2c/87f3254fd8ffd29e4c02732eee68a83a1d3c346ae39bc6822dcbcb697f2b/wheel-0.45.1-py3-none-any.whl (72 kB)
|
||||
Installing collected packages: wheel, pip
|
||||
Attempting uninstall: wheel
|
||||
Found existing installation: wheel 0.44.0
|
||||
Uninstalling wheel-0.44.0:
|
||||
Successfully uninstalled wheel-0.44.0
|
||||
Attempting uninstall: pip
|
||||
Found existing installation: pip 24.2
|
||||
Uninstalling pip-24.2:
|
||||
Successfully uninstalled pip-24.2
|
||||
Successfully installed pip-24.3.1 wheel-0.45.1
|
||||
Looking in indexes: https://mirrors.aliyun.com/pypi/simple/
|
||||
Collecting nvidia-cudnn-cu12
|
||||
Downloading https://mirrors.aliyun.com/pypi/packages/07/1d/3c7055e1d189f55ed25183f50e1d82f477024289f925df22ac89c0f4d387/nvidia_cudnn_cu12-9.6.0.74-py3-none-win_amd64.whl (502.8 MB)
|
||||
---------------------------------------- 502.8/502.8 MB 6.3 MB/s eta 0:00:00
|
||||
Collecting nvidia-cublas-cu12 (from nvidia-cudnn-cu12)
|
||||
Downloading https://mirrors.aliyun.com/pypi/packages/84/f7/985e9bdbe3e0ac9298fcc8cfa51a392862a46a0ffaccbbd56939b62a9c83/nvidia_cublas_cu12-12.6.4.1-py3-none-win_amd64.whl (434.5 MB)
|
||||
---------------------------------------- 434.5/434.5 MB 6.2 MB/s eta 0:00:00
|
||||
Installing collected packages: nvidia-cublas-cu12, nvidia-cudnn-cu12
|
||||
Successfully installed nvidia-cublas-cu12-12.6.4.1 nvidia-cudnn-cu12-9.6.0.74
|
||||
Looking in indexes: https://pypi.org/simple, https://pypi.nvidia.com
|
||||
Collecting tensorrt_cu12_libs==10.6.0
|
||||
Downloading https://pypi.nvidia.com/tensorrt-cu12-libs/tensorrt_cu12_libs-10.6.0-py2.py3-none-win_amd64.whl (1075.0 MB)
|
||||
---------------------------------------- 1.1/1.1 GB 8.2 MB/s eta 0:00:00
|
||||
Collecting tensorrt_cu12_bindings==10.6.0
|
||||
Downloading https://pypi.nvidia.com/tensorrt-cu12-bindings/tensorrt_cu12_bindings-10.6.0-cp310-none-win_amd64.whl (814 kB)
|
||||
---------------------------------------- 814.1/814.1 kB 12.0 MB/s eta 0:00:00
|
||||
Collecting tensorrt==10.6.0
|
||||
Downloading https://pypi.nvidia.com/tensorrt/tensorrt-10.6.0.tar.gz (16 kB)
|
||||
Preparing metadata (setup.py) ... done
|
||||
Collecting nvidia-cuda-runtime-cu12 (from tensorrt_cu12_libs==10.6.0)
|
||||
Downloading https://pypi.nvidia.com/nvidia-cuda-runtime-cu12/nvidia_cuda_runtime_cu12-12.6.77-py3-none-win_amd64.whl (891 kB)
|
||||
---------------------------------------- 891.8/891.8 kB 10.2 MB/s eta 0:00:00
|
||||
Collecting tensorrt-cu12==10.6.0 (from tensorrt==10.6.0)
|
||||
Downloading https://pypi.nvidia.com/tensorrt-cu12/tensorrt-cu12-10.6.0.tar.gz (18 kB)
|
||||
Preparing metadata (setup.py) ... done
|
||||
Building wheels for collected packages: tensorrt, tensorrt-cu12
|
||||
Building wheel for tensorrt (setup.py) ... done
|
||||
Created wheel for tensorrt: filename=tensorrt-10.6.0-py2.py3-none-any.whl size=16345 sha256=3fc94cd935389ea9535e279dc81bdf07b97da751150f64a1f45c39ba8b5c2663
|
||||
Stored in directory: c:\users\eleven\appdata\local\pip\cache\wheels\42\bd\86\598d26504705081ad82ae23a08a869322ee5a8bdb410d5d298
|
||||
Building wheel for tensorrt-cu12 (setup.py) ... done
|
||||
Created wheel for tensorrt-cu12: filename=tensorrt_cu12-10.6.0-py2.py3-none-any.whl size=17568 sha256=1856795aec69a0f6767ab62c0962d586ea16c3301c0f7040cde242c5f92f9abc
|
||||
Stored in directory: c:\users\eleven\appdata\local\pip\cache\wheels\13\99\06\326119ed088736cbd6ded866ada42bb29bbb67bed34903e57b
|
||||
Successfully built tensorrt tensorrt-cu12
|
||||
Installing collected packages: tensorrt_cu12_bindings, tensorrt-cu12, nvidia-cuda-runtime-cu12, tensorrt_cu12_libs, tensorrt
|
||||
Successfully installed nvidia-cuda-runtime-cu12-12.6.77 tensorrt-10.6.0 tensorrt-cu12-10.6.0 tensorrt_cu12_bindings-10.6.0 tensorrt_cu12_libs-10.6.0
|
||||
Looking in indexes: https://mirrors.aliyun.com/pypi/simple/
|
||||
Requirement already satisfied: numpy==1.26.4 in c:\users\eleven\desktop\easyvtuber-main\envs\miniconda3\envs\ezvtb_rt_venv\lib\site-packages (from -r requirements.txt (line 1)) (1.26.4)
|
||||
Collecting scipy (from -r requirements.txt (line 2))
|
||||
Downloading https://mirrors.aliyun.com/pypi/packages/51/8c/c2d371111961f737ae08881f654cf54eca796c42ec0429add2a06df97049/scipy-1.15.0-cp310-cp310-win_amd64.whl (43.9 MB)
|
||||
---------------------------------------- 43.9/43.9 MB 6.5 MB/s eta 0:00:00
|
||||
Collecting matplotlib (from -r requirements.txt (line 3))
|
||||
Downloading https://mirrors.aliyun.com/pypi/packages/5a/85/b9a54d64585a6b8737a78a61897450403c30f39e0bd3214270bb0b96f002/matplotlib-3.10.0-cp310-cp310-win_amd64.whl (8.0 MB)
|
||||
---------------------------------------- 8.0/8.0 MB 6.7 MB/s eta 0:00:00
|
||||
Collecting protobuf==3.20.1 (from -r requirements.txt (line 4))
|
||||
Downloading https://mirrors.aliyun.com/pypi/packages/c0/9c/bb88091287418ae1cf8af2bb9ed9710748a562b9abc227e4884d687a8650/protobuf-3.20.1-cp310-cp310-win_amd64.whl (903 kB)
|
||||
---------------------------------------- 903.8/903.8 kB 6.8 MB/s eta 0:00:00
|
||||
Collecting pynput==1.7.6 (from -r requirements.txt (line 5))
|
||||
Downloading https://mirrors.aliyun.com/pypi/packages/02/27/4de87850ff87c8dcecaaf8d27f28cec89ef17eeb6938f250449cb2635e06/pynput-1.7.6-py2.py3-none-any.whl (89 kB)
|
||||
Collecting mediapipe (from -r requirements.txt (line 6))
|
||||
Downloading https://mirrors.aliyun.com/pypi/packages/f6/7a/448885fc960a972d8bf1939602919fdaa8de17f505b6c50782a074e7b70e/mediapipe-0.10.20-cp310-cp310-win_amd64.whl (51.0 MB)
|
||||
---------------------------------------- 51.0/51.0 MB 6.2 MB/s eta 0:00:00
|
||||
Collecting opencv_python==4.5.5.64 (from -r requirements.txt (line 7))
|
||||
Downloading https://mirrors.aliyun.com/pypi/packages/48/c3/798bd7b8f78430f82ec0660b753106717e4e4bb8032ce56f77d8577cb24b/opencv_python-4.5.5.64-cp36-abi3-win_amd64.whl (35.4 MB)
|
||||
---------------------------------------- 35.4/35.4 MB 6.3 MB/s eta 0:00:00
|
||||
Collecting Pillow==9.1.0 (from -r requirements.txt (line 8))
|
||||
Downloading https://mirrors.aliyun.com/pypi/packages/a8/df/1177786a2d1c0bf732ba6d5f05a2fa40f016e81e1c16d62f1101e35d271e/Pillow-9.1.0-cp310-cp310-win_amd64.whl (3.3 MB)
|
||||
---------------------------------------- 3.3/3.3 MB 6.4 MB/s eta 0:00:00
|
||||
Collecting pyanime4k==2.5.2 (from -r requirements.txt (line 9))
|
||||
Downloading https://mirrors.aliyun.com/pypi/packages/81/af/25a4a2e9a3d46820c016f6bb2cc235d953f104546a079c8c67bf09548734/pyanime4k-2.5.2-py3-none-win_amd64.whl (35.4 MB)
|
||||
---------------------------------------- 35.4/35.4 MB 6.3 MB/s eta 0:00:00
|
||||
Collecting pyvirtualcam==0.9.1 (from -r requirements.txt (line 10))
|
||||
Downloading https://mirrors.aliyun.com/pypi/packages/ca/59/e245e2d9832c6cefd212133d5c1ef988e66e7ccd93f4381017e844949334/pyvirtualcam-0.9.1-cp310-cp310-win_amd64.whl (158 kB)
|
||||
Collecting pyopengl (from -r requirements.txt (line 11))
|
||||
Downloading https://mirrors.aliyun.com/pypi/packages/99/48/00e31747821d3fc56faddd00a4725454d1e694a8b67d715cf20f531506a5/PyOpenGL-3.1.7-py3-none-any.whl (2.4 MB)
|
||||
---------------------------------------- 2.4/2.4 MB 6.3 MB/s eta 0:00:00
|
||||
Collecting onnx (from -r requirements.txt (line 12))
|
||||
Downloading https://mirrors.aliyun.com/pypi/packages/0e/d3/d26ebf590a65686dde6b27fef32493026c5be9e42083340d947395f93405/onnx-1.17.0-cp310-cp310-win_amd64.whl (14.5 MB)
|
||||
---------------------------------------- 14.5/14.5 MB 6.0 MB/s eta 0:00:00
|
||||
Collecting onnxruntime-directml (from -r requirements.txt (line 13))
|
||||
Downloading https://mirrors.aliyun.com/pypi/packages/3c/4f/f433239b05304aa9af0217da20508abbbcec1dcd58ee821e3dab8939ecfe/onnxruntime_directml-1.20.1-cp310-cp310-win_amd64.whl (22.8 MB)
|
||||
---------------------------------------- 22.8/22.8 MB 6.3 MB/s eta 0:00:00
|
||||
Collecting turbojpeg (from -r requirements.txt (line 14))
|
||||
Downloading https://mirrors.aliyun.com/pypi/packages/a0/06/1a6ca4e7eb9c8342e5991093a36dcde37033a8f8da58dcf680282bb2b733/turbojpeg-0.0.2-cp310-cp310-win_amd64.whl (439 kB)
|
||||
Collecting tqdm (from -r requirements.txt (line 15))
|
||||
Downloading https://mirrors.aliyun.com/pypi/packages/d0/30/dc54f88dd4a2b5dc8a0279bdd7270e735851848b762aeb1c1184ed1f6b14/tqdm-4.67.1-py3-none-any.whl (78 kB)
|
||||
Collecting six (from pynput==1.7.6->-r requirements.txt (line 5))
|
||||
Downloading https://mirrors.aliyun.com/pypi/packages/b7/ce/149a00dd41f10bc29e5921b496af8b574d8413afcd5e30dfa0ed46c2cc5e/six-1.17.0-py2.py3-none-any.whl (11 kB)
|
||||
Collecting ffmpeg-python>=0.2.0 (from pyanime4k==2.5.2->-r requirements.txt (line 9))
|
||||
Downloading https://mirrors.aliyun.com/pypi/packages/d7/0c/56be52741f75bad4dc6555991fabd2e07b432d333da82c11ad701123888a/ffmpeg_python-0.2.0-py3-none-any.whl (25 kB)
|
||||
Collecting contourpy>=1.0.1 (from matplotlib->-r requirements.txt (line 3))
|
||||
Downloading https://mirrors.aliyun.com/pypi/packages/e0/82/c372be3fc000a3b2005061ca623a0d1ecd2eaafb10d9e883a2fc8566e951/contourpy-1.3.1-cp310-cp310-win_amd64.whl (218 kB)
|
||||
Collecting cycler>=0.10 (from matplotlib->-r requirements.txt (line 3))
|
||||
Downloading https://mirrors.aliyun.com/pypi/packages/e7/05/c19819d5e3d95294a6f5947fb9b9629efb316b96de511b418c53d245aae6/cycler-0.12.1-py3-none-any.whl (8.3 kB)
|
||||
Collecting fonttools>=4.22.0 (from matplotlib->-r requirements.txt (line 3))
|
||||
Downloading https://mirrors.aliyun.com/pypi/packages/19/d1/4dcd865360fb2c499749a913fe80e41c26e8ae18629d87dfffa3de27e831/fonttools-4.55.3-cp310-cp310-win_amd64.whl (2.2 MB)
|
||||
---------------------------------------- 2.2/2.2 MB 6.3 MB/s eta 0:00:00
|
||||
Collecting kiwisolver>=1.3.1 (from matplotlib->-r requirements.txt (line 3))
|
||||
Downloading https://mirrors.aliyun.com/pypi/packages/98/99/0dd05071654aa44fe5d5e350729961e7bb535372935a45ac89a8924316e6/kiwisolver-1.4.8-cp310-cp310-win_amd64.whl (71 kB)
|
||||
Collecting packaging>=20.0 (from matplotlib->-r requirements.txt (line 3))
|
||||
Downloading https://mirrors.aliyun.com/pypi/packages/88/ef/eb23f262cca3c0c4eb7ab1933c3b1f03d021f2c48f54763065b6f0e321be/packaging-24.2-py3-none-any.whl (65 kB)
|
||||
Collecting pyparsing>=2.3.1 (from matplotlib->-r requirements.txt (line 3))
|
||||
Downloading https://mirrors.aliyun.com/pypi/packages/1c/a7/c8a2d361bf89c0d9577c934ebb7421b25dc84bf3a8e3ac0a40aed9acc547/pyparsing-3.2.1-py3-none-any.whl (107 kB)
|
||||
Collecting python-dateutil>=2.7 (from matplotlib->-r requirements.txt (line 3))
|
||||
Downloading https://mirrors.aliyun.com/pypi/packages/ec/57/56b9bcc3c9c6a792fcbaf139543cee77261f3651ca9da0c93f5c1221264b/python_dateutil-2.9.0.post0-py2.py3-none-any.whl (229 kB)
|
||||
Collecting absl-py (from mediapipe->-r requirements.txt (line 6))
|
||||
Downloading https://mirrors.aliyun.com/pypi/packages/a2/ad/e0d3c824784ff121c03cc031f944bc7e139a8f1870ffd2845cc2dd76f6c4/absl_py-2.1.0-py3-none-any.whl (133 kB)
|
||||
Collecting attrs>=19.1.0 (from mediapipe->-r requirements.txt (line 6))
|
||||
Downloading https://mirrors.aliyun.com/pypi/packages/89/aa/ab0f7891a01eeb2d2e338ae8fecbe57fcebea1a24dbb64d45801bfab481d/attrs-24.3.0-py3-none-any.whl (63 kB)
|
||||
Collecting flatbuffers>=2.0 (from mediapipe->-r requirements.txt (line 6))
|
||||
Downloading https://mirrors.aliyun.com/pypi/packages/fb/b4/31c461eef98b96b8ab736d97274548eaf2b2e349bf09e4de3902f7d53084/flatbuffers-24.12.23-py2.py3-none-any.whl (30 kB)
|
||||
Collecting jax (from mediapipe->-r requirements.txt (line 6))
|
||||
Downloading https://mirrors.aliyun.com/pypi/packages/22/49/b4418a7a892c0dd64442bbbeef54e1cdfe722dfc5a7bf0d611d3f5f90e99/jax-0.4.38-py3-none-any.whl (2.2 MB)
|
||||
---------------------------------------- 2.2/2.2 MB 6.3 MB/s eta 0:00:00
|
||||
Collecting jaxlib (from mediapipe->-r requirements.txt (line 6))
|
||||
Downloading https://mirrors.aliyun.com/pypi/packages/0d/04/120c4caac6151f7297fedf9dd776362aa2d417d3f87bda826050b4da45e8/jaxlib-0.4.38-cp310-cp310-win_amd64.whl (64.2 MB)
|
||||
---------------------------------------- 64.2/64.2 MB 6.2 MB/s eta 0:00:00
|
||||
WARNING: Retrying (Retry(total=4, connect=None, read=None, redirect=None, status=None)) after connection broken by 'ReadTimeoutError("HTTPSConnectionPool(host='mirrors.aliyun.com', port=443): Read timed out. (read timeout=15)")': /pypi/simple/opencv-contrib-python/
|
||||
Collecting opencv-contrib-python (from mediapipe->-r requirements.txt (line 6))
|
||||
Downloading https://mirrors.aliyun.com/pypi/packages/a7/9e/7110d2c5d543ab03b9581dbb1f8e2429863e44e0c9b4960b766f230c1279/opencv_contrib_python-4.10.0.84-cp37-abi3-win_amd64.whl (45.5 MB)
|
||||
---------------------------------------- 45.5/45.5 MB 6.3 MB/s eta 0:00:00
|
||||
INFO: pip is looking at multiple versions of mediapipe to determine which version is compatible with other requirements. This could take a while.
|
||||
Collecting mediapipe (from -r requirements.txt (line 6))
|
||||
Downloading https://mirrors.aliyun.com/pypi/packages/d0/ca/4a32e95f1767467f7c461d7d5b4b506faebb1e30e35c1031c1b31cf82846/mediapipe-0.10.18-cp310-cp310-win_amd64.whl (50.9 MB)
|
||||
---------------------------------------- 50.9/50.9 MB 6.3 MB/s eta 0:00:00
|
||||
Downloading https://mirrors.aliyun.com/pypi/packages/ea/1b/79e1ddc041012f5e5a7d60cbfd8cfe6037c5c87f4e77b4fdc0f9dfa2d44e/mediapipe-0.10.14-cp310-cp310-win_amd64.whl (50.8 MB)
|
||||
---------------------------------------- 50.8/50.8 MB 6.3 MB/s eta 0:00:00
|
||||
Downloading https://mirrors.aliyun.com/pypi/packages/bb/b7/ffcdd19de3bf443864d0b0fed2b102275b96cbfbc1ddcf96ad811bbc1af0/mediapipe-0.10.13-cp310-cp310-win_amd64.whl (50.8 MB)
|
||||
---------------------------------------- 50.8/50.8 MB 6.4 MB/s eta 0:00:00
|
||||
Downloading https://mirrors.aliyun.com/pypi/packages/12/50/9c24e158350d3f93be669db291fb452f21a25d874c94c5758374be82fff1/mediapipe-0.10.11-cp310-cp310-win_amd64.whl (50.8 MB)
|
||||
---------------------------------------- 50.8/50.8 MB 6.4 MB/s eta 0:00:00
|
||||
Collecting sounddevice>=0.4.4 (from mediapipe->-r requirements.txt (line 6))
|
||||
Downloading https://mirrors.aliyun.com/pypi/packages/af/9b/15217b04f3b36d30de55fef542389d722de63f1ad81f9c72d8afc98cb6ab/sounddevice-0.5.1-py3-none-win_amd64.whl (363 kB)
|
||||
INFO: pip is looking at multiple versions of onnx to determine which version is compatible with other requirements. This could take a while.
|
||||
Collecting onnx (from -r requirements.txt (line 12))
|
||||
Downloading https://mirrors.aliyun.com/pypi/packages/0c/e4/1bc3ae56e6581587926a50a5c9dce3cfacf510e592e2512c7f5f2a9a4859/onnx-1.16.2-cp310-cp310-win_amd64.whl (14.4 MB)
|
||||
---------------------------------------- 14.4/14.4 MB 6.4 MB/s eta 0:00:00
|
||||
Downloading https://mirrors.aliyun.com/pypi/packages/9c/7c/40fbebcb30f0fb6a773ca36e3a9dd4bdccd4b15455ef6c21a335781db78f/onnx-1.16.1-cp310-cp310-win_amd64.whl (14.4 MB)
|
||||
---------------------------------------- 14.4/14.4 MB 6.3 MB/s eta 0:00:00
|
||||
Downloading https://mirrors.aliyun.com/pypi/packages/3b/e2/471ff83b3862967791d67f630000afce038756afbdf0665a3d767677c851/onnx-1.16.0-cp310-cp310-win_amd64.whl (14.4 MB)
|
||||
---------------------------------------- 14.4/14.4 MB 6.3 MB/s eta 0:00:00
|
||||
Downloading https://mirrors.aliyun.com/pypi/packages/1b/43/6e84bf35a9201fb344b1a98edf7095c4fa1cf1478dfb6638d7b77f5475e6/onnx-1.15.0-cp310-cp310-win_amd64.whl (14.3 MB)
|
||||
---------------------------------------- 14.3/14.3 MB 6.3 MB/s eta 0:00:00
|
||||
Downloading https://mirrors.aliyun.com/pypi/packages/00/36/e7a7e7a85564e7d409e4e8addfa11d41015d2190bfff30064771e7c21ca0/onnx-1.14.1-cp310-cp310-win_amd64.whl (13.3 MB)
|
||||
---------------------------------------- 13.3/13.3 MB 6.2 MB/s eta 0:00:00
|
||||
Downloading https://mirrors.aliyun.com/pypi/packages/22/5c/46298252ea9f92b6b94184e8f001e575f2c346a22011498110fd032fc921/onnx-1.14.0-cp310-cp310-win_amd64.whl (13.3 MB)
|
||||
---------------------------------------- 13.3/13.3 MB 6.2 MB/s eta 0:00:00
|
||||
Downloading https://mirrors.aliyun.com/pypi/packages/bb/cd/e2e63b1f0922dd45f6941a3422ccdb400360955064654a1a3a20e6a613fb/onnx-1.13.1-cp310-cp310-win_amd64.whl (12.2 MB)
|
||||
---------------------------------------- 12.2/12.2 MB 6.2 MB/s eta 0:00:00
|
||||
INFO: pip is still looking at multiple versions of onnx to determine which version is compatible with other requirements. This could take a while.
|
||||
Downloading https://mirrors.aliyun.com/pypi/packages/97/e1/bdb42dde046507250083862520428901af3f5182faac50d3df5f13a4ff3d/onnx-1.13.0-cp310-cp310-win_amd64.whl (12.2 MB)
|
||||
---------------------------------------- 12.2/12.2 MB 6.5 MB/s eta 0:00:00
|
||||
Downloading https://mirrors.aliyun.com/pypi/packages/b0/f0/5e383f64293857e86c300cb70010cce1ca6f038f1563b85501d30404ad2a/onnx-1.12.0-cp310-cp310-win_amd64.whl (11.5 MB)
|
||||
---------------------------------------- 11.5/11.5 MB 6.2 MB/s eta 0:00:00
|
||||
Requirement already satisfied: typing-extensions>=3.6.2.1 in c:\users\eleven\desktop\easyvtuber-main\envs\miniconda3\envs\ezvtb_rt_venv\lib\site-packages (from onnx->-r requirements.txt (line 12)) (4.12.2)
|
||||
Collecting coloredlogs (from onnxruntime-directml->-r requirements.txt (line 13))
|
||||
Downloading https://mirrors.aliyun.com/pypi/packages/a7/06/3d6badcf13db419e25b07041d9c7b4a2c331d3f4e7134445ec5df57714cd/coloredlogs-15.0.1-py2.py3-none-any.whl (46 kB)
|
||||
Requirement already satisfied: sympy in c:\users\eleven\desktop\easyvtuber-main\envs\miniconda3\envs\ezvtb_rt_venv\lib\site-packages (from onnxruntime-directml->-r requirements.txt (line 13)) (1.13.3)
|
||||
Collecting colorama (from tqdm->-r requirements.txt (line 15))
|
||||
Downloading https://mirrors.aliyun.com/pypi/packages/d1/d6/3965ed04c63042e047cb6a3e6ed1a63a35087b6a609aa3a15ed8ac56c221/colorama-0.4.6-py2.py3-none-any.whl (25 kB)
|
||||
Collecting future (from ffmpeg-python>=0.2.0->pyanime4k==2.5.2->-r requirements.txt (line 9))
|
||||
Downloading https://mirrors.aliyun.com/pypi/packages/da/71/ae30dadffc90b9006d77af76b393cb9dfbfc9629f339fc1574a1c52e6806/future-1.0.0-py3-none-any.whl (491 kB)
|
||||
Collecting CFFI>=1.0 (from sounddevice>=0.4.4->mediapipe->-r requirements.txt (line 6))
|
||||
Downloading https://mirrors.aliyun.com/pypi/packages/d1/b6/0b0f5ab93b0df4acc49cae758c81fe4e5ef26c3ae2e10cc69249dfd8b3ab/cffi-1.17.1-cp310-cp310-win_amd64.whl (181 kB)
|
||||
Collecting humanfriendly>=9.1 (from coloredlogs->onnxruntime-directml->-r requirements.txt (line 13))
|
||||
Downloading https://mirrors.aliyun.com/pypi/packages/f0/0f/310fb31e39e2d734ccaa2c0fb981ee41f7bd5056ce9bc29b2248bd569169/humanfriendly-10.0-py2.py3-none-any.whl (86 kB)
|
||||
Collecting ml_dtypes>=0.4.0 (from jax->mediapipe->-r requirements.txt (line 6))
|
||||
Downloading https://mirrors.aliyun.com/pypi/packages/6b/7e/bc54ae885e4d702e60a4bf50aa9066ff35e9c66b5213d11091f6bffb3036/ml_dtypes-0.5.1-cp310-cp310-win_amd64.whl (209 kB)
|
||||
Collecting opt_einsum (from jax->mediapipe->-r requirements.txt (line 6))
|
||||
Downloading https://mirrors.aliyun.com/pypi/packages/23/cd/066e86230ae37ed0be70aae89aabf03ca8d9f39c8aea0dec8029455b5540/opt_einsum-3.4.0-py3-none-any.whl (71 kB)
|
||||
Requirement already satisfied: mpmath<1.4,>=1.1.0 in c:\users\eleven\desktop\easyvtuber-main\envs\miniconda3\envs\ezvtb_rt_venv\lib\site-packages (from sympy->onnxruntime-directml->-r requirements.txt (line 13)) (1.3.0)
|
||||
Collecting pycparser (from CFFI>=1.0->sounddevice>=0.4.4->mediapipe->-r requirements.txt (line 6))
|
||||
Downloading https://mirrors.aliyun.com/pypi/packages/13/a3/a812df4e2dd5696d1f351d58b8fe16a405b234ad2886a0dab9183fb78109/pycparser-2.22-py3-none-any.whl (117 kB)
|
||||
Collecting pyreadline3 (from humanfriendly>=9.1->coloredlogs->onnxruntime-directml->-r requirements.txt (line 13))
|
||||
Downloading https://mirrors.aliyun.com/pypi/packages/5a/dc/491b7661614ab97483abf2056be1deee4dc2490ecbf7bff9ab5cdbac86e1/pyreadline3-3.5.4-py3-none-any.whl (83 kB)
|
||||
Installing collected packages: pyopengl, flatbuffers, turbojpeg, six, scipy, pyvirtualcam, pyreadline3, pyparsing, pycparser, protobuf, Pillow, packaging, opt_einsum, opencv_python, opencv-contrib-python, ml_dtypes, kiwisolver, future, fonttools, cycler, contourpy, colorama, attrs, absl-py, tqdm, python-dateutil, pynput, onnx, jaxlib, humanfriendly, ffmpeg-python, CFFI, sounddevice, pyanime4k, matplotlib, jax, coloredlogs, onnxruntime-directml, mediapipe
|
||||
Successfully installed CFFI-1.17.1 Pillow-9.1.0 absl-py-2.1.0 attrs-24.3.0 colorama-0.4.6 coloredlogs-15.0.1 contourpy-1.3.1 cycler-0.12.1 ffmpeg-python-0.2.0 flatbuffers-24.12.23 fonttools-4.55.3 future-1.0.0 humanfriendly-10.0 jax-0.4.38 jaxlib-0.4.38 kiwisolver-1.4.8 matplotlib-3.10.0 mediapipe-0.10.11 ml_dtypes-0.5.1 onnx-1.12.0 onnxruntime-directml-1.20.1 opencv-contrib-python-4.10.0.84 opencv_python-4.5.5.64 opt_einsum-3.4.0 packaging-24.2 protobuf-3.20.1 pyanime4k-2.5.2 pycparser-2.22 pynput-1.7.6 pyopengl-3.1.7 pyparsing-3.2.1 pyreadline3-3.5.4 python-dateutil-2.9.0.post0 pyvirtualcam-0.9.1 scipy-1.15.0 six-1.17.0 sounddevice-0.5.1 tqdm-4.67.1 turbojpeg-0.0.2
|
||||
Press any key to continue . . .
|
||||
Binary file not shown.
|
Before Width: | Height: | Size: 193 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 108 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 118 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 49 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 71 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 177 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 166 KiB |
@@ -1,127 +0,0 @@
|
||||
import os
|
||||
import sys
|
||||
import subprocess
|
||||
import urllib.request
|
||||
import shutil
|
||||
|
||||
def init_ezvtb_rt(search_dir = '.'):
|
||||
if not os.path.exists(ezvtb_path) and not os.path.exists(ezvtb_main_path): #Check if rt not exist
|
||||
zip_path = 'https://github.com/zpeng11/ezvtuber-rt/archive/refs/heads/main.zip'
|
||||
filehandle, _ = urllib.request.urlretrieve(zip_path)
|
||||
os.rename(filehandle, filehandle + '.zip')
|
||||
shutil.unpack_archive(filehandle + '.zip' , '.')
|
||||
|
||||
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)
|
||||
|
||||
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
|
||||
):
|
||||
support_trt = False
|
||||
if use_tensorrt: #Verify tensorrt
|
||||
print('Verifying TensorRT')
|
||||
try:
|
||||
from ezvtb_rt.trt_utils import cudaSetDevice, check_build_all_models
|
||||
cudaSetDevice(device_id)
|
||||
# check_build_all_models()
|
||||
support_trt = True
|
||||
except:
|
||||
support_trt = False
|
||||
if support_trt == False and use_tensorrt == True:
|
||||
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 devcie {pycuda.autoinit.device.name()}')
|
||||
|
||||
tha_model_dir = os.path.join('.','data', 'models', 'tha3',
|
||||
'seperable' if model_seperable else 'standard',
|
||||
'fp16' if model_half else 'fp32')
|
||||
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: #x2
|
||||
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'THA3 Path:{tha_model_dir}')
|
||||
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
|
||||
core = CoreTRT(tha_model_dir,
|
||||
vram_cache_size = model_cache_size,
|
||||
use_eyebrow = model_use_eyebrow,
|
||||
rife_dir = rife_model_path if len(rife_model_path) > 0 else None,
|
||||
sr_dir = sr_model_path if len(sr_model_path) > 0 else None,
|
||||
cache_max_volume = cacher_ram_size,
|
||||
cache_quality = cacher_quality,
|
||||
)
|
||||
else: #Use directml
|
||||
from ezvtb_rt.core_ort import CoreORT
|
||||
core = CoreORT(tha_model_dir,
|
||||
rife_path = rife_model_path if len(rife_model_path) > 0 else None,
|
||||
sr_path = sr_model_path if len(sr_model_path) > 0 else None,
|
||||
device_id=device_id,
|
||||
cache_max_volume = cacher_ram_size,
|
||||
cache_quality = cacher_quality,
|
||||
use_eyebrow = model_use_eyebrow)
|
||||
|
||||
return core
|
||||
|
||||
if __name__ == '__main__':
|
||||
init_ezvtb_rt()
|
||||
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: d602899f2d...07c302921e.
+424
-80
@@ -1,6 +1,7 @@
|
||||
import ctypes
|
||||
import os
|
||||
import subprocess
|
||||
import threading
|
||||
|
||||
import wx
|
||||
import json
|
||||
@@ -29,26 +30,31 @@ cache_simplify_quality_map = {
|
||||
}
|
||||
default_arg = {
|
||||
'character': 'lambda_00',
|
||||
'input': 2,
|
||||
'input': 3,
|
||||
'output': 2,
|
||||
'ifm': None,
|
||||
'osf': '127.0.0.1:11573',
|
||||
'min_cutoff': 50,
|
||||
'beta': 80,
|
||||
'is_extend_movement': False,
|
||||
'is_alpha_split': False,
|
||||
'is_bongo': False,
|
||||
'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',
|
||||
'interpolation': "Off",
|
||||
'frame_rate_limit': '30',
|
||||
'sr': 'Off',
|
||||
'device_id': '0',
|
||||
'sr': "Off",
|
||||
'use_tensorrt': False,
|
||||
'preset': 'Medium'
|
||||
'preset': 'Low',
|
||||
'mouse_audio_input': False,
|
||||
'audio_sensitivity': '0.02',
|
||||
'audio_threshold': '10.0',
|
||||
'blink_interval': '5.0',
|
||||
'breath_cycle': 'inf'
|
||||
}
|
||||
|
||||
try:
|
||||
@@ -64,15 +70,17 @@ finally:
|
||||
p = None
|
||||
dirPath = 'data/images'
|
||||
characterList = []
|
||||
studentModelList = []
|
||||
studentModelCharacterMap = {}
|
||||
|
||||
hasRTModel = False
|
||||
try:
|
||||
f = open('data/models/tha3/standard/fp16/decomposer.trt')
|
||||
f.close()
|
||||
hasRTModel = True
|
||||
except:
|
||||
pass
|
||||
|
||||
def is_nvidia_gpu():
|
||||
try:
|
||||
# 获取显卡名称列表
|
||||
output = subprocess.check_output("wmic path Win32_VideoController get Name", shell=True).decode('gbk')
|
||||
return "NVIDIA" in output.upper()
|
||||
except Exception:
|
||||
return False
|
||||
hasTRTSupport = is_nvidia_gpu()
|
||||
|
||||
def refreshList():
|
||||
global characterList
|
||||
@@ -82,11 +90,66 @@ def refreshList():
|
||||
characterList.append(item[:-4])
|
||||
|
||||
|
||||
def scanStudentModels():
|
||||
"""Scan custom_tha4_models directory for student models"""
|
||||
global studentModelList, studentModelCharacterMap
|
||||
studentModelList = []
|
||||
studentModelCharacterMap = {}
|
||||
|
||||
custom_models_path = 'data/models/custom_tha4_models'
|
||||
if os.path.exists(custom_models_path):
|
||||
try:
|
||||
for model_name in os.listdir(custom_models_path):
|
||||
model_path = os.path.join(custom_models_path, model_name)
|
||||
if os.path.isdir(model_path):
|
||||
# Check if it's a valid student model
|
||||
face_trt = os.path.join(model_path, 'face_morpher.trt')
|
||||
body_trt = os.path.join(model_path, 'body_morpher.trt')
|
||||
character_png = os.path.join(model_path,
|
||||
'character.png')
|
||||
has_trt = (os.path.exists(face_trt) and
|
||||
os.path.exists(body_trt))
|
||||
has_character = os.path.exists(character_png)
|
||||
|
||||
if has_trt and has_character:
|
||||
studentModelList.append(model_name)
|
||||
studentModelCharacterMap[model_name] = model_name
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Sort alphabetically
|
||||
studentModelList.sort()
|
||||
|
||||
|
||||
refreshList()
|
||||
scanStudentModels()
|
||||
|
||||
def min_cutoff_mapper(value, revert=False):
|
||||
"""
|
||||
非线性映射函数：0-100整数 <-> 0-100浮点数
|
||||
使用平方函数，使得越接近0数字越密集
|
||||
"""
|
||||
if revert:
|
||||
# 浮点 -> 整数: 使用平方根反向映射
|
||||
return int((value / 100.0) ** 0.5 * 100)
|
||||
# 整数 -> 浮点: 使用平方映射
|
||||
return (value / 100.0) ** 2 * 100.0
|
||||
|
||||
|
||||
def beta_mapper(value, revert=False):
|
||||
"""
|
||||
非线性映射函数：0-100整数 <-> 0-1浮点数
|
||||
使用平方函数，使得越接近0数字越密集
|
||||
"""
|
||||
if revert:
|
||||
# 浮点 -> 整数: 使用平方根反向映射
|
||||
return int((value ** 0.5) * 100)
|
||||
# 整数 -> 浮点: 使用平方映射
|
||||
return (value / 100.0) ** 2
|
||||
|
||||
|
||||
class OptionPanel(wx.Panel):
|
||||
def __init__(self, parent, title='', desc='', choices=None, mapping=None, type=0, default=None, disabled=False):
|
||||
def __init__(self, parent, title='', desc='', choices=None, mapping=None, type=0, default=None, disabled=False, mapper=min_cutoff_mapper):
|
||||
wx.Panel.__init__(self, parent)
|
||||
self.type = type
|
||||
if mapping is not None:
|
||||
@@ -140,8 +203,38 @@ class OptionPanel(wx.Panel):
|
||||
self.control.SetValue(default)
|
||||
except:
|
||||
pass
|
||||
elif self.type == 3:
|
||||
# Slider type for float values 0.0 to 1.0
|
||||
sliderSizer = wx.BoxSizer(wx.HORIZONTAL)
|
||||
self.control = wx.Slider(self, wx.ID_ANY, value=50, minValue=0, maxValue=100,
|
||||
style=wx.SL_HORIZONTAL)
|
||||
self.control.SetMinSize(wx.Size(250, -1))
|
||||
|
||||
# Add a label to show the float value
|
||||
self.valueLabel = wx.StaticText(self, wx.ID_ANY, "0.50")
|
||||
self.valueLabel.SetMinSize(wx.Size(50, -1))
|
||||
|
||||
try:
|
||||
if default is not None:
|
||||
self.control.SetValue(default)
|
||||
self.valueLabel.SetLabelText(f"{mapper(default):.4f}")
|
||||
except:
|
||||
pass
|
||||
|
||||
# Update label when slider changes
|
||||
def onSliderChange(event):
|
||||
val = mapper(self.control.GetValue())
|
||||
self.valueLabel.SetLabelText(f"{val:.4f}")
|
||||
self.control.Bind(wx.EVT_SLIDER, onSliderChange)
|
||||
|
||||
sliderSizer.Add(self.control, 1, wx.ALIGN_CENTER_VERTICAL)
|
||||
sliderSizer.Add(self.valueLabel, 0, wx.ALIGN_CENTER_VERTICAL | wx.LEFT, 10)
|
||||
mainSizer.Add(sliderSizer, 0, wx.ALIGN_CENTER_VERTICAL | wx.LEFT, 20)
|
||||
# Skip the normal control addition below
|
||||
self.control._slider_added = True
|
||||
|
||||
mainSizer.Add(self.control, 0, wx.ALIGN_CENTER_VERTICAL | wx.LEFT, 20)
|
||||
if not (self.type == 3 and hasattr(self.control, '_slider_added')):
|
||||
mainSizer.Add(self.control, 0, wx.ALIGN_CENTER_VERTICAL | wx.LEFT, 20)
|
||||
# self.SetBackgroundColour('#000000')
|
||||
|
||||
def GetValue(self):
|
||||
@@ -151,18 +244,103 @@ class OptionPanel(wx.Panel):
|
||||
ret = self.control.GetValue()
|
||||
elif self.type == 2:
|
||||
ret = self.control.GetValue()
|
||||
elif self.type == 3:
|
||||
ret = self.control.GetValue()
|
||||
if self.mapping is not None:
|
||||
return self.mapping[ret]
|
||||
else:
|
||||
return ret
|
||||
|
||||
|
||||
def _important_log_line(line):
|
||||
"""从 main 的一行日志中提取“重要”的简短描述，用于状态栏；无关行返回 None。"""
|
||||
line = line.strip()
|
||||
if not line:
|
||||
return None
|
||||
if line.startswith('Launched:'):
|
||||
return 'Launched'
|
||||
if 'Model Inference Ready' in line:
|
||||
return 'Model Inference Ready'
|
||||
# TRT: Building engine from ONNX: ...\filename.onnx
|
||||
if '[TRT]' in line and 'Building engine from ONNX:' in line:
|
||||
idx = line.find('Building engine from ONNX:')
|
||||
if idx != -1:
|
||||
path = line[idx + len('Building engine from ONNX:'):].strip().rstrip('\r\n')
|
||||
name = os.path.basename(path)
|
||||
if name:
|
||||
return f'Building: {name}'
|
||||
# TRT: Loading ONNX file from path ...\filename.onnx
|
||||
if '[TRT]' in line and 'Loading ONNX file from path' in line:
|
||||
idx = line.find('Loading ONNX file from path')
|
||||
if idx != -1:
|
||||
path = line[idx + len('Loading ONNX file from path'):].strip().strip('.').strip().rstrip('\r\n')
|
||||
name = os.path.basename(path)
|
||||
if name:
|
||||
return f'Loading: {name}'
|
||||
# ORT: Loading ONNX model from path ...\filename.onnx
|
||||
if '[ORT]' in line and 'Loading ONNX model from path' in line:
|
||||
idx = line.find('Loading ONNX model from path')
|
||||
if idx != -1:
|
||||
path = line[idx + len('Loading ONNX model from path'):].strip().strip('.').strip().rstrip('\r\n')
|
||||
name = os.path.basename(path)
|
||||
if name:
|
||||
return f'Loading: {name}'
|
||||
# ORT: Completed loading session: xxx.onnx
|
||||
if '[ORT]' in line and 'Completed loading session:' in line:
|
||||
idx = line.find('Completed loading session:')
|
||||
if idx != -1:
|
||||
name = line[idx + len('Completed loading session:'):].strip().rstrip('\r\n')
|
||||
if name:
|
||||
return f'Loaded: {name}'
|
||||
return None
|
||||
|
||||
|
||||
def _on_main_log_line(panel, line):
|
||||
"""在子线程中调用：若该行是重要日志，则用 wx.CallAfter 更新 panel 的状态框。"""
|
||||
display = _important_log_line(line)
|
||||
if display is not None:
|
||||
wx.CallAfter(panel.statusCtrl.SetValue, display)
|
||||
|
||||
|
||||
def _read_pipe_to_stream(pipe, dest_stream, out_lines=None, on_line_callback=None):
|
||||
"""从 pipe 读行，写回 dest_stream，可选追加到 out_lines，并对每行调用 on_line_callback(line)。"""
|
||||
if pipe is None:
|
||||
return
|
||||
try:
|
||||
for raw in iter(pipe.readline, b''):
|
||||
try:
|
||||
text = raw.decode('utf-8', errors='replace')
|
||||
except Exception:
|
||||
text = raw.decode('gbk', errors='replace')
|
||||
if out_lines is not None:
|
||||
out_lines.append(text)
|
||||
if on_line_callback is not None:
|
||||
on_line_callback(text)
|
||||
# 检查dest_stream是否可用（pythonw环境下可能不可用）
|
||||
if dest_stream is not None:
|
||||
try:
|
||||
dest_stream.write(text)
|
||||
dest_stream.flush()
|
||||
except (AttributeError, OSError, ValueError):
|
||||
# pythonw环境下sys.stdout/sys.stderr可能不可用，忽略错误
|
||||
pass
|
||||
except Exception:
|
||||
pass
|
||||
finally:
|
||||
try:
|
||||
pipe.close()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
class LauncherPanel(wx.Panel):
|
||||
def __init__(self, parent):
|
||||
wx.Panel.__init__(self, parent)
|
||||
self.number_of_buttons = 0
|
||||
self.frame = parent
|
||||
self.optionDict = {}
|
||||
self.main_output_lines = [] # main 的 stdout 副本
|
||||
self.main_stderr_lines = [] # main 的 stderr 副本
|
||||
self.mainSizer = wx.BoxSizer(wx.VERTICAL)
|
||||
controlSizer = wx.BoxSizer(wx.HORIZONTAL)
|
||||
self.widgetSizer = wx.BoxSizer(wx.VERTICAL)
|
||||
@@ -178,7 +356,16 @@ class LauncherPanel(wx.Panel):
|
||||
f = f.MakeLarger()
|
||||
stVtuber.SetFont(f)
|
||||
controlSizer.Add(stEasy, 0, wx.ALL | wx.CENTER, 0)
|
||||
controlSizer.Add(stVtuber, 1, wx.RIGHT | wx.CENTER, 30)
|
||||
controlSizer.Add(stVtuber, 0, wx.RIGHT | wx.CENTER, 30)
|
||||
|
||||
self.statusCtrl = wx.TextCtrl(
|
||||
self, wx.ID_ANY, '',
|
||||
style=wx.TE_READONLY | wx.BORDER_NONE | wx.TE_RIGHT,
|
||||
)
|
||||
self.statusCtrl.SetHint('当前操作')
|
||||
f = self.statusCtrl.GetFont()
|
||||
self.statusCtrl.SetFont(f.Smaller())
|
||||
controlSizer.Add(self.statusCtrl, 1, wx.ALIGN_CENTER_VERTICAL | wx.LEFT | wx.RIGHT, 8)
|
||||
|
||||
self.btnLaunch = wx.Button(self, label="Save && Launch")
|
||||
self.btnLaunch.Bind(wx.EVT_BUTTON, self.OnLaunch)
|
||||
@@ -209,27 +396,54 @@ class LauncherPanel(wx.Panel):
|
||||
addOption('input', title='Input Device', desc='选择希望使用的面捕数据源',
|
||||
choices=['iFacialMocap', 'OpenSeeFace', 'OpenCV(Webcam)', 'Mouse Input', 'Debug Input'],
|
||||
mapping=[0, 4, 1, 3, 2])
|
||||
addOption('ifm', title='iFacialMocap IP:Port', desc='输入iFacialMocap连接使用的IP:端口号', type=2)
|
||||
addOption('is_eyebrow', title='Eyebrow', desc='使用眉毛输入，对性能有一定影响\n仅支持iFacialMocap', type=1,
|
||||
addOption('ifm', title='iFacialMocap IP', desc='输入iFacialMocap连接使用的IP地址，默认连接 49983 端口', type=2)
|
||||
addOption('is_eyebrow', title='Eyebrow', desc='使用眉毛输入，对性能有一定影响', type=1,
|
||||
default=True)
|
||||
addOption('osf', title='OpenSeeFace IP:Port', desc='输入OpenSeeFace连接使用的IP:端口号', type=2)
|
||||
addOption('mouse_audio_input', title='Audio Input', desc='启用WASAPI音频输入控制嘴部动作', type=1)
|
||||
addOption('audio_sensitivity', title='Audio Sensitivity', desc='音频灵敏度，控制音频对嘴部动作的影响程度', type=2)
|
||||
addOption('audio_threshold', title='Audio Threshold', desc='音频阈值，低于此值的音频将被忽略', type=2)
|
||||
addOption('blink_interval', title='Blink Interval', desc='设置眨眼间隔时间',
|
||||
choices=['No Blink', '3 seconds', '5 seconds', '7 seconds'],
|
||||
mapping=['inf', '3.0', '5.0', '7.0'])
|
||||
addOption('min_cutoff', title='Min CutOff', desc='输入滤波频率截断，\n越小越平滑，越大静止时越灵敏',
|
||||
type=3, mapper=min_cutoff_mapper)
|
||||
addOption('beta', title='Beta', desc='输入滤波速度补偿，\n越小越平滑，越大运动时越灵敏',
|
||||
type=3, mapper=beta_mapper)
|
||||
|
||||
addOption('breath_cycle', title='Breath Cycle', desc='设置呼吸循环时间(会增加占用)',
|
||||
choices=['No Breath', '3 seconds', '5 seconds', '7 seconds'],
|
||||
mapping=['inf', '3.0', '5.0', '7.0'])
|
||||
|
||||
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显卡支持）',
|
||||
addOption('use_tensorrt', title='TensorRT加速',
|
||||
desc='需要更长启动和预热时间（仅NVIDIA显卡支持）',
|
||||
type=1)
|
||||
|
||||
addOption('frame_rate_limit', title='FPS Limit', desc='选择帧率限制目标',
|
||||
choices=['10', '15', '20', '30', '60'])
|
||||
addOption('preset', title='Performance Preset', desc='性能预设，注意修改后会覆盖后续配置',
|
||||
choices=['Low', 'Medium', 'High', 'Ultra', 'Custom'])
|
||||
addOption('model_select', title='Model Select', desc='选择使用的模型\nStandard Full精度较高性能较低',
|
||||
choices=['Seperable Half', 'Seperable Full', 'Standard Half', 'Standard Full'],
|
||||
mapping=['seperable_half', 'seperable_full', 'standard_half', 'standard_full'])
|
||||
|
||||
# Build model_select choices
|
||||
model_choices = ['Seperable Half', 'Seperable Full', 'Standard Half',
|
||||
'Standard Full', 'THA4 Half', 'THA4 Full']
|
||||
model_mapping = ['seperable_half', 'seperable_full',
|
||||
'standard_half', 'standard_full', 'tha4_half',
|
||||
'tha4_full']
|
||||
|
||||
# Add student models if available
|
||||
for student_model in studentModelList:
|
||||
model_choices.append(f'THA4 Student ({student_model})')
|
||||
model_mapping.append(f'tha4_student_{student_model}')
|
||||
|
||||
addOption('model_select', title='Model Select',
|
||||
desc='选择使用的模型\nStandard Full精度较高性能较低',
|
||||
choices=model_choices,
|
||||
mapping=model_mapping)
|
||||
addOption('ram_cache_size', title='RAM Cache Size', desc='分配内存缓存大小\n用于存储最终运算结果',
|
||||
choices=['Off', '1GB', '2GB', '4GB', '8GB', '16GB'],
|
||||
mapping=['0b', '1gb', '2gb', '4gb', '8gb', '16gb'])
|
||||
@@ -239,9 +453,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',
|
||||
@@ -270,11 +481,44 @@ class LauncherPanel(wx.Panel):
|
||||
self.optionSizer.Hide(self.optionDict['osf'])
|
||||
else:
|
||||
self.optionSizer.Show(self.optionDict['osf'])
|
||||
self.optionSizer.Show(self.optionDict['is_eyebrow'])
|
||||
if s != 1 and s != 4:
|
||||
self.optionSizer.Hide(self.optionDict['min_cutoff'])
|
||||
self.optionSizer.Hide(self.optionDict['beta'])
|
||||
else:
|
||||
self.optionSizer.Show(self.optionDict['min_cutoff'])
|
||||
self.optionSizer.Show(self.optionDict['beta'])
|
||||
# Show/hide audio input options for Mouse Input (s == 3)
|
||||
if s != 3:
|
||||
self.optionSizer.Hide(self.optionDict['mouse_audio_input'])
|
||||
self.optionSizer.Hide(self.optionDict['audio_sensitivity'])
|
||||
self.optionSizer.Hide(self.optionDict['audio_threshold'])
|
||||
self.optionSizer.Hide(self.optionDict['blink_interval'])
|
||||
else:
|
||||
self.optionSizer.Show(self.optionDict['mouse_audio_input'])
|
||||
self.optionSizer.Show(self.optionDict['blink_interval'])
|
||||
# Update audio fields based on checkbox state
|
||||
audioInputChoice()
|
||||
|
||||
self.frame.fSizer.Layout()
|
||||
self.frame.Fit()
|
||||
|
||||
def audioInputChoice(e=None):
|
||||
"""Handle mouse_audio_input checkbox changes"""
|
||||
enabled = self.optionDict['mouse_audio_input'].GetValue()
|
||||
if enabled:
|
||||
self.optionSizer.Show(self.optionDict['audio_sensitivity'])
|
||||
self.optionSizer.Show(self.optionDict['audio_threshold'])
|
||||
self.optionDict['audio_sensitivity'].control.Enable(True)
|
||||
self.optionDict['audio_threshold'].control.Enable(True)
|
||||
else:
|
||||
self.optionSizer.Hide(self.optionDict['audio_sensitivity'])
|
||||
self.optionSizer.Hide(self.optionDict['audio_threshold'])
|
||||
self.frame.fSizer.Layout()
|
||||
self.frame.Fit()
|
||||
|
||||
self.optionDict['input'].Bind(wx.EVT_CHOICE, inputChoice)
|
||||
self.optionDict['mouse_audio_input'].Bind(wx.EVT_CHECKBOX, audioInputChoice)
|
||||
inputChoice()
|
||||
|
||||
def presetChoice(e=None):
|
||||
@@ -284,13 +528,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':
|
||||
@@ -299,7 +542,7 @@ class LauncherPanel(wx.Panel):
|
||||
for c in presetControls: self.optionSizer.Hide(c)
|
||||
if s in presets:
|
||||
opt = presets[s]
|
||||
for i in range(5): presetControls[i].control.SetSelection(opt[i])
|
||||
for i in range(4): presetControls[i].control.SetSelection(opt[i])
|
||||
|
||||
self.frame.fSizer.Layout()
|
||||
self.frame.Fit()
|
||||
@@ -307,17 +550,54 @@ class LauncherPanel(wx.Panel):
|
||||
self.optionDict['preset'].Bind(wx.EVT_CHOICE, presetChoice)
|
||||
presetChoice()
|
||||
|
||||
def onModelSelect(e=None):
|
||||
"""Handle model selection change"""
|
||||
model_value = self.optionDict['model_select'].GetValue()
|
||||
is_student_model = 'tha4_student_' in model_value
|
||||
|
||||
char_ctrl = self.optionDict['character']
|
||||
|
||||
if is_student_model:
|
||||
# Disable character selection for student models
|
||||
# Student models have their own built-in character
|
||||
char_ctrl.control.Enable(False)
|
||||
char_ctrl.control.SetToolTip(
|
||||
'Locked: Student model includes built-in character')
|
||||
else:
|
||||
# Re-enable character selection for non-student models
|
||||
char_ctrl.control.Enable(True)
|
||||
char_ctrl.control.SetToolTip(
|
||||
'Select a character from data/images')
|
||||
|
||||
self.optionDict['model_select'].Bind(
|
||||
wx.EVT_CHOICE, onModelSelect)
|
||||
|
||||
# Check initial model selection and lock character if needed
|
||||
onModelSelect()
|
||||
|
||||
def onActivate(e):
|
||||
global characterList
|
||||
tName = self.optionDict['character'].GetValue()
|
||||
# 用控件当前选中字符串，避免 GetValue() 用旧 mapping 按下标取导致 IndexError
|
||||
char_ctrl = self.optionDict['character'].control
|
||||
tName = char_ctrl.GetStringSelection() if char_ctrl.GetSelection() >= 0 else ''
|
||||
refreshList()
|
||||
self.optionDict['character'].control.SetItems(characterList)
|
||||
self.optionDict['character'].control.SetSelection(characterList.index(tName))
|
||||
scanStudentModels()
|
||||
self.optionDict['character'].mapping = characterList # 同步 mapping，后续 GetValue() 才正确
|
||||
char_ctrl.SetItems(characterList)
|
||||
try:
|
||||
idx = characterList.index(tName)
|
||||
char_ctrl.SetSelection(idx)
|
||||
except (ValueError, TypeError):
|
||||
if characterList:
|
||||
char_ctrl.SetSelection(0)
|
||||
# 重新应用“当前模型 → 是否锁定 character”的规则
|
||||
onModelSelect()
|
||||
|
||||
if not hasRTModel:
|
||||
if not hasTRTSupport:
|
||||
self.optionDict['use_tensorrt'].control.SetValue(False)
|
||||
self.optionDict['use_tensorrt'].control.Enable(False)
|
||||
self.optionDict['use_tensorrt'].control.SetToolTip('需要先构建TensorRT模型')
|
||||
self.optionDict['use_tensorrt'].control.SetToolTip(
|
||||
'需要NVIDIA显卡支持才能使用TensorRT')
|
||||
|
||||
self.frame.Bind(wx.EVT_ACTIVATE, onActivate)
|
||||
|
||||
@@ -332,62 +612,85 @@ class LauncherPanel(wx.Panel):
|
||||
self.btnLaunch.SetLabelText('Working...')
|
||||
|
||||
if p is not None:
|
||||
subprocess.run(['taskkill', '/F', '/PID', str(p.pid), '/T'], stdout=subprocess.DEVNULL,
|
||||
stderr=subprocess.DEVNULL)
|
||||
creation_flags = 0
|
||||
if sys.platform == 'win32':
|
||||
# CREATE_NO_WINDOW = 0x08000000
|
||||
creation_flags = 0x08000000
|
||||
subprocess.run(['taskkill', '/F', '/PID', str(p.pid), '/T'],
|
||||
stdout=subprocess.DEVNULL,
|
||||
stderr=subprocess.DEVNULL,
|
||||
creationflags=creation_flags)
|
||||
p = None
|
||||
self.statusCtrl.Clear()
|
||||
self.btnLaunch.SetLabelText("Save & Launch")
|
||||
else:
|
||||
run_args = [sys.executable, 'main.py']
|
||||
# 如果启动器是用pythonw启动的，使用python.exe来启动main以便捕获控制台输出
|
||||
python_exe = sys.executable
|
||||
if 'pythonw' in python_exe.lower():
|
||||
python_exe = python_exe.replace('pythonw.exe', 'python.exe').replace('pythonw', 'python')
|
||||
run_args = [python_exe, '-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(args['ifm'])
|
||||
run_args.append('--ifm_input')
|
||||
if ':' in args['ifm']:
|
||||
run_args.append(args['ifm'])
|
||||
else:
|
||||
run_args.append(args['ifm'] + ':49983')
|
||||
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:
|
||||
run_args.append('--mouse_input')
|
||||
run_args.append('0,0,' + str(wx.GetDisplaySize().width) + ',' + str(wx.GetDisplaySize().height))
|
||||
# Add audio input options for mouse input
|
||||
if args['mouse_audio_input']:
|
||||
run_args.append('--mouse_audio_input')
|
||||
if args['audio_sensitivity']:
|
||||
run_args.append('--audio_sensitivity')
|
||||
run_args.append(str(args['audio_sensitivity']))
|
||||
if args['audio_threshold']:
|
||||
run_args.append('--audio_threshold')
|
||||
run_args.append(str(args['audio_threshold']))
|
||||
# Add blink interval for mouse input
|
||||
if args['blink_interval']:
|
||||
run_args.append('--blink_interval')
|
||||
run_args.append(str(args['blink_interval']))
|
||||
elif args['input'] == 4:
|
||||
if len(args['osf']):
|
||||
run_args.append('--osf')
|
||||
run_args.append('--osf_input')
|
||||
run_args.append(args['osf'])
|
||||
|
||||
# Add breath cycle option
|
||||
if args['breath_cycle']:
|
||||
run_args.append('--breath_cycle')
|
||||
run_args.append(str(args['breath_cycle']))
|
||||
|
||||
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'])
|
||||
@@ -411,6 +714,20 @@ class LauncherPanel(wx.Panel):
|
||||
run_args.append('4')
|
||||
|
||||
if args['model_select'] is not None:
|
||||
if 'tha4_student_' in args['model_select']:
|
||||
# Student model: tha4_student_{model_name}
|
||||
model_name = args['model_select'].replace(
|
||||
'tha4_student_', '')
|
||||
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('--model_version')
|
||||
run_args.append('v4')
|
||||
else:
|
||||
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']:
|
||||
@@ -421,28 +738,49 @@ class LauncherPanel(wx.Panel):
|
||||
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')
|
||||
run_args.append('--filter_min_cutoff')
|
||||
run_args.append(str(min_cutoff_mapper(args['min_cutoff'])))
|
||||
|
||||
run_args.append('--filter_beta')
|
||||
run_args.append(str(beta_mapper(args['beta'])))
|
||||
|
||||
print('Launched: ' + ' '.join(run_args))
|
||||
p = subprocess.Popen(run_args)
|
||||
self.main_output_lines.clear()
|
||||
self.main_stderr_lines.clear()
|
||||
self.statusCtrl.SetValue('Launched')
|
||||
on_line = lambda line: _on_main_log_line(self, line)
|
||||
# 使用CREATE_NO_WINDOW标志隐藏控制台窗口，但仍可捕获输出
|
||||
creation_flags = 0
|
||||
if sys.platform == 'win32':
|
||||
# CREATE_NO_WINDOW = 0x08000000
|
||||
creation_flags = 0x08000000
|
||||
p = subprocess.Popen(
|
||||
run_args,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.PIPE,
|
||||
creationflags=creation_flags,
|
||||
)
|
||||
threading.Thread(
|
||||
target=_read_pipe_to_stream,
|
||||
args=(p.stdout, sys.stdout, self.main_output_lines, on_line),
|
||||
daemon=True,
|
||||
).start()
|
||||
threading.Thread(
|
||||
target=_read_pipe_to_stream,
|
||||
args=(p.stderr, sys.stderr, self.main_stderr_lines, on_line),
|
||||
daemon=True,
|
||||
).start()
|
||||
self.btnLaunch.SetLabelText('Stop')
|
||||
|
||||
|
||||
@@ -456,8 +794,14 @@ class MainFrame(wx.Frame):
|
||||
def OnClose(self, e):
|
||||
global p
|
||||
if p is not None:
|
||||
subprocess.run(['taskkill', '/F', '/PID', str(p.pid), '/T'], stdout=subprocess.DEVNULL,
|
||||
stderr=subprocess.DEVNULL)
|
||||
creation_flags = 0
|
||||
if sys.platform == 'win32':
|
||||
# CREATE_NO_WINDOW = 0x08000000
|
||||
creation_flags = 0x08000000
|
||||
subprocess.run(['taskkill', '/F', '/PID', str(p.pid), '/T'],
|
||||
stdout=subprocess.DEVNULL,
|
||||
stderr=subprocess.DEVNULL,
|
||||
creationflags=creation_flags)
|
||||
e.Skip()
|
||||
|
||||
def InitUi(self):
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
import time
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
import tha2.poser.modes.mode_20
|
||||
@@ -13,6 +14,14 @@ from args import args
|
||||
|
||||
from collections import OrderedDict
|
||||
|
||||
# THA4 适配器导入
|
||||
try:
|
||||
from tha4_adapter import THA4Wrapper
|
||||
THA4_AVAILABLE = True
|
||||
except ImportError:
|
||||
THA4_AVAILABLE = False
|
||||
print("Warning: THA4 adapter not available")
|
||||
|
||||
|
||||
class TalkingAnimeLight(nn.Module):
|
||||
def __init__(self):
|
||||
@@ -171,3 +180,58 @@ class TalkingAnime(nn.Module):
|
||||
rotate_image = self.two_algo_face_rotator(x, pose_vector)[:2]
|
||||
output_image = self.combiner(rotate_image[0], rotate_image[1], pose_vector)
|
||||
return output_image
|
||||
|
||||
|
||||
class TalkingAnime4(nn.Module):
|
||||
"""
|
||||
THA4 model class, compatible with THA3 TalkingAnime3 interface
|
||||
Uses THA4's Siren architecture for inference
|
||||
PyTorch only, float32 precision
|
||||
"""
|
||||
def __init__(self, device):
|
||||
super(TalkingAnime4, self).__init__()
|
||||
if not THA4_AVAILABLE:
|
||||
raise RuntimeError(
|
||||
"THA4 adapter not available. Check tha4_adapter.py"
|
||||
)
|
||||
|
||||
# Create THA4 wrapper
|
||||
self.wrapper = THA4Wrapper(device=device)
|
||||
|
||||
print("THA4 model initialized (float32)")
|
||||
|
||||
def forward(self, image, mouth_eye_vector, pose_vector, eyebrow_vector,
|
||||
mouth_eye_vector_c, eyebrow_vector_c, ratio=None):
|
||||
"""
|
||||
Forward inference, compatible with THA3 interface
|
||||
|
||||
Args:
|
||||
image: [batch, 4, 512, 512]
|
||||
mouth_eye_vector: [batch, 27]
|
||||
pose_vector: [batch, 6]
|
||||
eyebrow_vector: [batch, 12]
|
||||
mouth_eye_vector_c: compressed (for caching)
|
||||
eyebrow_vector_c: compressed (for caching)
|
||||
ratio: GPU cache hit ratio
|
||||
|
||||
Returns:
|
||||
output_image: [batch, 4, 512, 512]
|
||||
"""
|
||||
if args.perf == 'model':
|
||||
tic = time.perf_counter()
|
||||
|
||||
# Call THA4 wrapper
|
||||
output_image = self.wrapper.forward(
|
||||
image, mouth_eye_vector, pose_vector, eyebrow_vector,
|
||||
mouth_eye_vector_c, eyebrow_vector_c, ratio
|
||||
)
|
||||
|
||||
if args.perf == 'model':
|
||||
print(" - tha4_inference", (time.perf_counter() - tic) * 1000)
|
||||
|
||||
return output_image
|
||||
|
||||
def to(self, device):
|
||||
"""Move model to specified device"""
|
||||
self.wrapper.to(device)
|
||||
return super().to(device)
|
||||
@@ -0,0 +1,54 @@
|
||||
@echo on
|
||||
REM ========================================================================
|
||||
REM IMPORTANT: Before running this script, please make sure to:
|
||||
REM 1. Checkout to the release_package branch: git checkout release_package
|
||||
REM 2. Rebase with the latest changes: git rebase main
|
||||
REM ========================================================================
|
||||
|
||||
cd /D "%~dp0"
|
||||
|
||||
if not exist "envs" mkdir envs
|
||||
if not exist "envs\miniconda3" mkdir envs\miniconda3
|
||||
|
||||
IF not EXIST %~dp0envs\TensorRT-RTX-1.3.0.35_cu129\bin (
|
||||
@RD /S /Q %~dp0envs\TensorRT-RTX-1.3.0.35_cu129
|
||||
cd /D "%~dp0\envs"
|
||||
curl -L -o trt_rtx.zip https://developer.nvidia.com/downloads/trt/rtx_sdk/secure/1.3/TensorRT-RTX-1.3.0.35-win10-amd64-cuda-12.9-Release-external.zip
|
||||
tar -xf trt_rtx.zip "TensorRT-RTX-1.3.0.35" && ren "TensorRT-RTX-1.3.0.35" "TensorRT-RTX-1.3.0.35_cu129" && del trt_rtx.zip
|
||||
cd /D "%~dp0"
|
||||
)
|
||||
|
||||
IF not EXIST %~dp0envs\miniconda3\Scripts (
|
||||
@RD /S /Q %~dp0envs\miniconda3
|
||||
mkdir %~dp0envs\miniconda3
|
||||
echo "Downloading miniconda..."
|
||||
powershell -Command "Invoke-WebRequest -Uri 'https://repo.anaconda.com/miniconda/Miniconda3-py310_24.9.2-0-Windows-x86_64.exe' -OutFile '.\envs\miniconda3.exe' -UseBasicParsing -Headers @{'User-Agent'='Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/123.0.0.0 Safari/537.36'}"
|
||||
|
||||
echo "Installling minconda..."
|
||||
start /wait "" %~dp0envs\miniconda3.exe /S /AddToPath=0 /RegisterPython=0 /InstallationType=JustMe /D=%~dp0envs\miniconda3
|
||||
echo "Successfully install minconda"
|
||||
del envs\miniconda3.exe
|
||||
)
|
||||
|
||||
SET PATH=%~dp0envs\TensorRT-RTX-1.3.0.35_cu129\bin;%~dp0envs\miniconda3\Scripts;%PATH%
|
||||
|
||||
call activate
|
||||
call conda env list
|
||||
call conda update -y --all
|
||||
call conda create -n ezvtb_rt_venv_release python=3.10 -y
|
||||
call conda activate ezvtb_rt_venv_release
|
||||
call conda env list
|
||||
|
||||
call conda install -y conda-pack
|
||||
call conda install -y conda-forge::pycuda
|
||||
call conda install -y -c nvidia/label/cuda-12.9.1 cuda-nvcc-dev_win-64 cudnn cuda-runtime
|
||||
|
||||
call conda-pack -n ezvtb_rt_venv_release -o %~dp0envs\python_embedded --format no-archive
|
||||
|
||||
|
||||
SET PATH=%~dp0envs\python_embedded;%~dp0envs\python_embedded\Scripts;%~dp0envs\python_embedded\Library\bin;%PATH%
|
||||
call python -m pip install %~dp0envs\TensorRT-RTX-1.3.0.35_cu129\python\tensorrt_rtx-1.3.0.35-cp310-none-win_amd64.whl
|
||||
call python -m pip install -r requirements.txt --no-warn-script-location
|
||||
|
||||
@RD /S /Q %~dp0envs\miniconda3
|
||||
pause
|
||||
+11
-10
@@ -1,17 +1,18 @@
|
||||
numpy==1.26.4
|
||||
numpy
|
||||
scipy
|
||||
matplotlib
|
||||
protobuf==3.20.1
|
||||
pynput==1.7.6
|
||||
mediapipe
|
||||
opencv_python==4.5.5.64
|
||||
Pillow==9.1.0
|
||||
pyanime4k==2.5.2
|
||||
protobuf
|
||||
pynput
|
||||
mediapipe==0.10.11
|
||||
opencv_python
|
||||
Pillow
|
||||
pyanime4k==3.1.0
|
||||
pyvirtualcam==0.9.1
|
||||
pyopengl
|
||||
wxpython
|
||||
onnx
|
||||
onnxruntime-directml
|
||||
turbojpeg
|
||||
tqdm
|
||||
torch
|
||||
torch
|
||||
OneEuroFilter
|
||||
sounddevice
|
||||
brotli
|
||||
Whitespace-only changes.
+97
@@ -0,0 +1,97 @@
|
||||
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('--character', type=str, default='lambda_00')
|
||||
|
||||
parser.add_argument('--debug_input', action='store_true')
|
||||
parser.add_argument('--cam_input', action='store_true')
|
||||
parser.add_argument('--mouse_input', type=str)
|
||||
parser.add_argument('--ifm_input', type=str)
|
||||
parser.add_argument('--osf_input', type=str)
|
||||
|
||||
parser.add_argument('--mouse_audio_input', action='store_true')
|
||||
parser.add_argument('--audio_sensitivity', type=float, default=0.02)
|
||||
parser.add_argument('--audio_threshold', type=float, default=10.0)
|
||||
parser.add_argument('--blink_interval', type=float, default=5.0)
|
||||
parser.add_argument('--breath_cycle', type=float, default=float('inf'))
|
||||
|
||||
parser.add_argument('--output_virtual_cam', action='store_true')
|
||||
parser.add_argument('--output_spout2', action='store_true')
|
||||
parser.add_argument('--output_debug', action='store_true')
|
||||
|
||||
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=10.0)
|
||||
parser.add_argument('--filter_beta', type=float, default=0.3)
|
||||
|
||||
parser.add_argument('--simplify', type=int, default=1)
|
||||
|
||||
parser.add_argument('--use_tensorrt', action='store_true')
|
||||
parser.add_argument('--model_version', type=str, default='v3')
|
||||
parser.add_argument('--model_name', type=str, default='')
|
||||
parser.add_argument('--model_seperable', action='store_true')
|
||||
parser.add_argument('--model_half', action='store_true')
|
||||
parser.add_argument('--gpu_cache', type=str, default='512mb')
|
||||
parser.add_argument('--eyebrow', action='store_true')
|
||||
|
||||
parser.add_argument('--use_interpolation', action='store_true')
|
||||
parser.add_argument('--interpolation_scale', type=int, default=1)
|
||||
parser.add_argument('--interpolation_half', action='store_true')
|
||||
parser.add_argument('--mark_interpolated', action='store_true', help='在插值帧左上角打红点便于确认补帧输出')
|
||||
|
||||
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_a4k', action='store_true')
|
||||
|
||||
parser.add_argument('--cache', type=str, default='256mb')
|
||||
|
||||
parser.add_argument('--frame_rate_limit', type=int, default=30)
|
||||
|
||||
parser.add_argument('--alpha_clean', action='store_true')
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.cache is not None:
|
||||
args.max_ram_cache_len = convert_to_byte(args.cache) / pow(1024, 3) # In gigabytes
|
||||
else:
|
||||
args.max_ram_cache_len = 0
|
||||
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.simplify == 0:
|
||||
args.max_ram_cache_len = 0 # Disable cacher if simplify is Off
|
||||
|
||||
if not args.output_virtual_cam and not args.output_spout2:
|
||||
args.output_debug = True # Default to debug output
|
||||
|
||||
if args.output_spout2:
|
||||
args.alpha_split = False # Disable alpha split for spout2 output
|
||||
|
||||
if not args.cam_input and args.mouse_input is None \
|
||||
and args.ifm_input is None and args.osf_input is None:
|
||||
args.debug_input = True # Default to debug input
|
||||
|
||||
if args.sr_a4k:
|
||||
args.use_sr = True
|
||||
|
||||
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
|
||||
@@ -0,0 +1,65 @@
|
||||
from multiprocessing import Process, shared_memory, Value
|
||||
import numpy as np
|
||||
import time
|
||||
import math
|
||||
from .utils.timer_wait import wait_until
|
||||
from .utils.shared_mem_guard import SharedMemoryGuard
|
||||
from .args import args
|
||||
|
||||
class DebugInputClientProcess(Process):
|
||||
def __init__(self, pose_position_shm: shared_memory.SharedMemory):
|
||||
super().__init__()
|
||||
self.pose_position_shm = pose_position_shm
|
||||
self.fps = Value('f', 60.0)
|
||||
|
||||
def run(self):
|
||||
last_time : float = time.perf_counter()
|
||||
interval : float = 1.0 / 60 # 60 FPS
|
||||
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])
|
||||
# 呼吸循环参数
|
||||
breath_start_time = time.perf_counter()
|
||||
while True:
|
||||
eyebrow_vector = [0.0] * 12
|
||||
mouth_eye_vector = [0.0] * 27
|
||||
pose_vector = [0.0] * 6
|
||||
position_vector = [0, 0, 0, 1]
|
||||
|
||||
# 计算呼吸效果（使用 sin 函数，在 breath_cycle 时间内从 0 到 1 再到 0）
|
||||
breath_elapsed = (time.perf_counter() - breath_start_time) % args.breath_cycle
|
||||
# 使用 sin 函数，让值在一个周期内从 0 -> 1 -> 0
|
||||
# sin 在 0 到 π 之间从 0 到 1 到 0
|
||||
breath_value = np.sin(breath_elapsed / args.breath_cycle * np.pi)
|
||||
|
||||
mouth_eye_vector[2] = math.sin(time.perf_counter() * 3)
|
||||
mouth_eye_vector[3] = math.sin(time.perf_counter() * 3)
|
||||
|
||||
mouth_eye_vector[14] = 0
|
||||
|
||||
mouth_eye_vector[25] = math.sin(time.perf_counter() * 2.2) * 0.2
|
||||
mouth_eye_vector[26] = math.sin(time.perf_counter() * 3.5) * 0.8
|
||||
|
||||
pose_vector[0] = math.sin(time.perf_counter() * 1.1)
|
||||
pose_vector[1] = math.sin(time.perf_counter() * 1.2)
|
||||
pose_vector[2] = math.sin(time.perf_counter() * 1.5)
|
||||
pose_vector[3] = pose_vector[1]
|
||||
pose_vector[4] = pose_vector[2]
|
||||
pose_vector[5] = breath_value
|
||||
eyebrow_vector[6] = math.sin(time.perf_counter() * 1.1)
|
||||
eyebrow_vector[7] = math.sin(time.perf_counter() * 1.1)
|
||||
|
||||
model_input_arr = eyebrow_vector
|
||||
model_input_arr.extend(mouth_eye_vector)
|
||||
model_input_arr.extend(pose_vector)
|
||||
|
||||
position_vector[0] = math.sin(time.perf_counter() * 0.5) * 0.1
|
||||
position_vector[1] = math.sin(time.perf_counter() * 0.6) * 0.1
|
||||
position_vector[2] = math.sin(time.perf_counter() * 0.7) * 0.1
|
||||
position_vector[3] = 1
|
||||
|
||||
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
|
||||
@@ -0,0 +1,72 @@
|
||||
import os
|
||||
import sys
|
||||
from .args import args
|
||||
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(
|
||||
use_tensorrt:bool = True,
|
||||
#THA3 model setting
|
||||
model_version:str = 'v3', #'v3' or 'v4' or 'v4_student'
|
||||
model_name:str = '',
|
||||
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_ram_size:float = 2.0,#unit of GigaBytes
|
||||
#SR setting
|
||||
use_sr:bool = False,
|
||||
sr_x4:bool = True,
|
||||
sr_half:bool = True,
|
||||
sr_a4k:bool = False,
|
||||
):
|
||||
if use_tensorrt:
|
||||
try:
|
||||
from ezvtb_rt.core_trt import CoreTRT as Core
|
||||
except:
|
||||
print("TensorRT is not available, fallback to ONNX Runtime.")
|
||||
args.use_tensorrt = False
|
||||
from ezvtb_rt.core_ort import CoreORT as Core
|
||||
else:
|
||||
from ezvtb_rt.core_ort import CoreORT as Core
|
||||
core = Core(
|
||||
tha_model_version=model_version,
|
||||
tha_model_seperable=model_seperable,
|
||||
tha_model_fp16=model_half,
|
||||
tha_model_name=model_name,
|
||||
|
||||
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_fp16=sr_half,
|
||||
sr_a4k=sr_a4k,
|
||||
|
||||
vram_cache_size=model_cache_size,
|
||||
cache_max_giga=cacher_ram_size,
|
||||
use_eyebrow=model_use_eyebrow,
|
||||
)
|
||||
return core
|
||||
|
||||
@@ -0,0 +1,100 @@
|
||||
import mediapipe as mp
|
||||
from multiprocessing import Process, shared_memory, Value
|
||||
import cv2
|
||||
import numpy as np
|
||||
from .args import args
|
||||
from .utils.shared_mem_guard import SharedMemoryGuard
|
||||
from .utils.pose import get_pose
|
||||
from .utils.fps import FPS
|
||||
from .utils.filter import OneEuroFilterNumpy
|
||||
from OneEuroFilter import OneEuroFilter
|
||||
import time
|
||||
|
||||
|
||||
class FaceMeshClientProcess(Process):
|
||||
def __init__(self, pose_position_shm: shared_memory.SharedMemory):
|
||||
super().__init__()
|
||||
self.pose_position_shm = pose_position_shm
|
||||
self.fps = Value('f', 0.0)
|
||||
|
||||
def run(self):
|
||||
facemesh = mp.solutions.face_mesh.FaceMesh(refine_landmarks=True)
|
||||
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])
|
||||
|
||||
cap = cv2.VideoCapture(0, cv2.CAP_DSHOW)
|
||||
if not cap.isOpened():
|
||||
raise Exception("Can't open webcam")
|
||||
|
||||
print("Warming up webcam...")
|
||||
|
||||
frame_count = 0
|
||||
input_fps = FPS(60)
|
||||
while frame_count < 5:
|
||||
ret, frame = cap.read()
|
||||
if not ret:
|
||||
raise Exception("Can't receive frame (stream end?).")
|
||||
input_fps()
|
||||
frame_count += 1
|
||||
|
||||
# 呼吸循环参数
|
||||
breath_start_time = time.perf_counter()
|
||||
pose_filter = OneEuroFilterNumpy(freq=input_fps.view(), mincutoff=args.filter_min_cutoff, beta=args.filter_beta)
|
||||
position_offset = None
|
||||
print("Webcam Input Running at {:.2f} FPS".format(input_fps.view()))
|
||||
position_vector = np.array([0, 0, 0, 1], dtype=np.float32)
|
||||
while True:
|
||||
ret, frame = cap.read()
|
||||
if not ret:
|
||||
raise Exception("Can't receive frame (stream end?).")
|
||||
self.fps.value = input_fps()
|
||||
rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
|
||||
results = facemesh.process(rgb_frame)
|
||||
if results.multi_face_landmarks is None:
|
||||
continue
|
||||
|
||||
# 计算呼吸效果（使用 sin 函数，在 breath_cycle 时间内从 0 到 1 再到 0）
|
||||
breath_elapsed = (time.perf_counter() - breath_start_time) % args.breath_cycle
|
||||
# 使用 sin 函数，让值在一个周期内从 0 -> 1 -> 0
|
||||
# sin 在 0 到 π 之间从 0 到 1 到 0
|
||||
breath_value = np.sin(breath_elapsed / args.breath_cycle * np.pi)
|
||||
|
||||
facial_landmarks = results.multi_face_landmarks[0].landmark
|
||||
pose = get_pose(facial_landmarks)
|
||||
eye_l_h_temp = pose[0]
|
||||
eye_r_h_temp = pose[1]
|
||||
mouth_ratio = pose[2]
|
||||
eye_y_ratio = pose[3]
|
||||
eye_x_ratio = pose[4]
|
||||
x_angle = pose[5]
|
||||
y_angle = pose[6]
|
||||
z_angle = pose[7]
|
||||
|
||||
eyebrow_vector = [0.0] * 12
|
||||
mouth_eye_vector = [0.0] * 27
|
||||
pose_vector = [0.0] * 6
|
||||
|
||||
mouth_eye_vector[2] = max(eye_l_h_temp, eye_r_h_temp)
|
||||
mouth_eye_vector[3] = max(eye_l_h_temp, eye_r_h_temp)
|
||||
|
||||
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
|
||||
if position_offset is None:
|
||||
position_offset = [(x_angle - 1.5) * 1.6, y_angle * 2.0 , (z_angle + 1.5) * 2]
|
||||
pose_vector[0] = (x_angle - 1.5) * 1.6 - position_offset[0]
|
||||
pose_vector[1] = y_angle * 2.0 - position_offset[1] # temp weight
|
||||
pose_vector[2] = (z_angle + 1.5) * 2 - position_offset[2] # temp weight
|
||||
pose_vector[3] = pose_vector[1]
|
||||
pose_vector[4] = pose_vector[2]
|
||||
pose_vector[5] = breath_value
|
||||
|
||||
model_input_arr = eyebrow_vector
|
||||
model_input_arr.extend(mouth_eye_vector)
|
||||
model_input_arr.extend(pose_vector)
|
||||
|
||||
with pose_position_shm_guard.lock():
|
||||
np_pose_shm[:] = pose_filter(np.array(model_input_arr, dtype=np.float32))
|
||||
np_position_shm[:] = position_vector
|
||||
@@ -0,0 +1,133 @@
|
||||
from multiprocessing import Process, shared_memory, Value
|
||||
from .args import args
|
||||
from .utils.timer_wait import wait_until
|
||||
from .utils.shared_mem_guard import SharedMemoryGuard
|
||||
import socket
|
||||
import errno
|
||||
import numpy as np
|
||||
import time
|
||||
import tha2.poser.modes.mode_20_wx
|
||||
from tha2.mocap.ifacialmocap_constants import *
|
||||
from .utils.fps import FPS
|
||||
from .utils.filter import OneEuroFilterNumpy
|
||||
ifm_converter = tha2.poser.modes.mode_20_wx.IFacialMocapPoseConverter20()
|
||||
|
||||
class IFMClientProcess(Process):
|
||||
def __init__(self, pose_position_shm: shared_memory.SharedMemory):
|
||||
super().__init__()
|
||||
self.pose_position_shm = pose_position_shm
|
||||
self.address = args.ifm_input.split(':')[0]
|
||||
self.port = int(args.ifm_input.split(':')[1])
|
||||
self.fps = Value('f', 0.0)
|
||||
def run(self):
|
||||
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])
|
||||
|
||||
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.bind(("", self.port))
|
||||
self.socket.settimeout(5.0) # give ifm app some time to respond
|
||||
|
||||
print("Warming up iFacialMocap connection...")
|
||||
|
||||
frame_count = 0
|
||||
input_fps = FPS(60)
|
||||
while frame_count < 60:
|
||||
socket_bytes = self.socket.recv(8192)
|
||||
if not socket_bytes:
|
||||
raise Exception("Can't receive iFacialMocap data (stream end?).")
|
||||
input_fps()
|
||||
frame_count += 1
|
||||
|
||||
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)
|
||||
# 呼吸循环参数
|
||||
breath_start_time = time.perf_counter()
|
||||
print("iFacialMocap Input Running at {:.2f} FPS".format(input_fps.view()))
|
||||
while True:
|
||||
socket_bytes = self.socket.recv(8192)
|
||||
|
||||
if not socket_bytes:
|
||||
raise Exception("Can't receive iFacialMocap data (stream end?).")
|
||||
|
||||
self.fps.value = input_fps()
|
||||
|
||||
socket_string = socket_bytes.decode("utf-8")
|
||||
|
||||
try:
|
||||
data = self.convert_from_blender_data(socket_string)
|
||||
except Exception:
|
||||
print("iFacialMocap data parse error:", socket_string)
|
||||
continue
|
||||
|
||||
# 计算呼吸效果（使用 sin 函数，在 breath_cycle 时间内从 0 到 1 再到 0）
|
||||
breath_elapsed = (time.perf_counter() - breath_start_time) % args.breath_cycle
|
||||
# 使用 sin 函数，让值在一个周期内从 0 -> 1 -> 0
|
||||
# sin 在 0 到 π 之间从 0 到 1 到 0
|
||||
breath_value = np.sin(breath_elapsed / args.breath_cycle * np.pi)
|
||||
|
||||
ifacialmocap_pose_converted = ifm_converter.convert(data)
|
||||
position_vector = data[HEAD_BONE_QUAT]
|
||||
|
||||
eyebrow_vector = [0.0] * 12
|
||||
mouth_eye_vector = [0.0] * 27
|
||||
pose_vector = [0.0] * 6
|
||||
for i in range(0, 12):
|
||||
eyebrow_vector[i] = ifacialmocap_pose_converted[i]
|
||||
for i in range(12, 39):
|
||||
mouth_eye_vector[i - 12] = ifacialmocap_pose_converted[i]
|
||||
for i in range(39, 42):
|
||||
pose_vector[i - 39] = ifacialmocap_pose_converted[i]
|
||||
pose_vector[3] = pose_vector[1]
|
||||
pose_vector[4] = pose_vector[2]
|
||||
pose_vector[5] = breath_value
|
||||
|
||||
mouth_eye_vector[14] = mouth_eye_vector[14] * 1.5
|
||||
|
||||
model_input_arr = eyebrow_vector
|
||||
model_input_arr.extend(mouth_eye_vector)
|
||||
model_input_arr.extend(pose_vector)
|
||||
|
||||
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):
|
||||
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
|
||||
+159
@@ -0,0 +1,159 @@
|
||||
import numpy as np
|
||||
import time
|
||||
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.shared_mem_guard import SharedMemoryGuard
|
||||
from multiprocessing import shared_memory
|
||||
from .utils.timer_wait import wait_until
|
||||
from PIL import Image
|
||||
from .utils.fps import FPS
|
||||
import pyvirtualcam
|
||||
from OpenGL.GL import GL_RGBA
|
||||
|
||||
|
||||
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)
|
||||
|
||||
pose_position_shm = shared_memory.SharedMemory(create=True,
|
||||
size=(45 + 4) * 4) # 45 floats for pose, 4 floats for position
|
||||
input_process = None
|
||||
|
||||
if args.cam_input:
|
||||
from .face_mesh_client import FaceMeshClientProcess
|
||||
input_process = FaceMeshClientProcess(pose_position_shm)
|
||||
elif args.ifm_input is not None:
|
||||
from .i_facial_mocap_client import IFMClientProcess
|
||||
input_process = IFMClientProcess(pose_position_shm)
|
||||
elif args.osf_input is not None:
|
||||
from .open_see_face_client import OSFClientProcess
|
||||
input_process = OSFClientProcess(pose_position_shm)
|
||||
elif args.mouse_input is not None:
|
||||
from .mouse_client import MouseClientProcess
|
||||
input_process = MouseClientProcess(pose_position_shm)
|
||||
else:
|
||||
from .debug_input_client import DebugInputClientProcess
|
||||
input_process = DebugInputClientProcess(pose_position_shm)
|
||||
|
||||
input_fps = input_process.fps
|
||||
input_process.daemon = True
|
||||
input_process.start()
|
||||
|
||||
infer_process = ModelClientProcess(input_image, pose_position_shm, input_fps)
|
||||
infer_process.daemon = True
|
||||
infer_process.start()
|
||||
|
||||
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 = [
|
||||
SharedMemoryGuard(infer_process.ret_shared_mem, ctrl_name=f"ret_shm_ctrl_batch_{i}")
|
||||
for i in range(args.interpolation_scale)
|
||||
]
|
||||
np_ret_shms = [
|
||||
np.ndarray((args.model_output_size, cam_width_scale * args.model_output_size, ret_channels), dtype=np.uint8,
|
||||
buffer=infer_process.ret_shared_mem.buf[
|
||||
i * cam_width_scale * args.model_output_size * args.model_output_size * ret_channels:
|
||||
(i + 1) * cam_width_scale * args.model_output_size * args.model_output_size * ret_channels])
|
||||
for i in range(args.interpolation_scale)
|
||||
]
|
||||
|
||||
last_time: float = time.perf_counter()
|
||||
interval: float = 1.0 / args.frame_rate_limit if args.frame_rate_limit > 0 else 0.0
|
||||
|
||||
if args.output_virtual_cam:
|
||||
virtual_cam = pyvirtualcam.Camera(width=cam_width_scale * args.model_output_size,
|
||||
height=args.model_output_size,
|
||||
fps=args.frame_rate_limit,
|
||||
backend='obs',
|
||||
fmt=pyvirtualcam.PixelFormat.RGB)
|
||||
print(f'Using virtual camera: {virtual_cam.device}')
|
||||
elif args.output_spout2:
|
||||
from PySpout import SpoutSender
|
||||
spout_sender = SpoutSender("EasyVtuber", cam_width_scale * args.model_output_size,
|
||||
args.model_output_size, GL_RGBA)
|
||||
else:
|
||||
print("Using OpenCV windows for output display.")
|
||||
|
||||
pipeline_fps = FPS()
|
||||
last_frame_time = None # 上一帧输出时间，用于打印帧时间差
|
||||
last_batch_start_time = None # 上一批就绪时间，用于周期估计
|
||||
n_frames = args.interpolation_scale
|
||||
min_period = n_frames * interval if interval > 0 else n_frames / 60.0 # 60fps 下本批最少占用时间
|
||||
default_period = 1.0 / 15.0 # 约 15fps 推理时的周期，首包无历史时使用
|
||||
|
||||
print("Interval set to {:.3f} seconds".format(interval))
|
||||
while True:
|
||||
infer_process.finish_event.wait()
|
||||
infer_process.finish_event.clear()
|
||||
for i in range(n_frames):
|
||||
ret_batch_shm_channels[i].acquire()
|
||||
|
||||
# 动态周期：本批就绪与上一批就绪的时间间隔，用于本批内均匀排期
|
||||
batch_start_time = time.perf_counter()
|
||||
if last_batch_start_time is not None:
|
||||
observed_period = batch_start_time - last_batch_start_time
|
||||
period = max(min_period, min(observed_period, 1.0))
|
||||
else:
|
||||
period = max(min_period, default_period)
|
||||
last_batch_start_time = batch_start_time
|
||||
|
||||
for i in range(n_frames):
|
||||
# 均匀排期 + frame_rate_limit：取两者中较晚的时间发送
|
||||
target_send_time = batch_start_time + i * (period / n_frames)
|
||||
if interval > 0:
|
||||
target_send_time = max(target_send_time, last_time)
|
||||
wait_until(target_send_time)
|
||||
|
||||
if args.output_virtual_cam:
|
||||
virtual_cam.send(np_ret_shms[i])
|
||||
elif args.output_spout2:
|
||||
spout_sender.send_image(np_ret_shms[i], False)
|
||||
else:
|
||||
cv2.imshow("EasyVtuber Debug Frame", np_ret_shms[i])
|
||||
cv2.waitKey(1)
|
||||
now_send = time.perf_counter()
|
||||
last_frame_time = now_send
|
||||
# 限速：下一帧最早在 last_time + interval，若已落后于当前时间则对齐到 now
|
||||
if interval > 0:
|
||||
last_time += interval
|
||||
if last_time < now_send:
|
||||
last_time = now_send
|
||||
ret_batch_shm_channels[i].release()
|
||||
output_pipeline_fps_val = pipeline_fps() * args.interpolation_scale
|
||||
infer_process.output_pipeline_fps.value = output_pipeline_fps_val
|
||||
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,
|
||||
output_pipeline_fps_val
|
||||
), end='\r', flush=True)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,194 @@
|
||||
import math
|
||||
import os
|
||||
from multiprocessing import Process, Value, shared_memory, Event
|
||||
import cv2
|
||||
import numpy as np
|
||||
import time
|
||||
from .ezvtb_rt_interface import get_core
|
||||
from .args import args
|
||||
from .utils.shared_mem_guard import SharedMemoryGuard
|
||||
from .utils.pose_simplify import pose_simplify
|
||||
from .utils.fps import FPS, Interval
|
||||
from typing import List
|
||||
|
||||
class ModelClientProcess(Process):
|
||||
def __init__(self, input_image, pose_position_shm: shared_memory.SharedMemory , input_fps):
|
||||
super().__init__()
|
||||
self.input_image = input_image
|
||||
self.pose_position_shm = pose_position_shm # 45 floats for pose, 4 floats for position
|
||||
|
||||
self.alpha_width_scale = 2 if args.alpha_split else 1
|
||||
self.ret_channels = 3 if args.output_virtual_cam or args.output_debug else 4
|
||||
self.ret_shape = (args.interpolation_scale, args.model_output_size, self.alpha_width_scale * args.model_output_size, self.ret_channels)
|
||||
self.ret_nbytes = self.alpha_width_scale * args.interpolation_scale * args.model_output_size * args.model_output_size * self.ret_channels # RGBA
|
||||
self.ret_shared_mem = shared_memory.SharedMemory(create=True, size=self.ret_nbytes)
|
||||
|
||||
self.last_model_interval = Value('f', 0.0)
|
||||
self.average_model_interval = Value('f', 0.0)
|
||||
self.cache_hit_ratio = Value('f', 0.0)
|
||||
self.gpu_cache_hit_ratio = Value('f', 0.0)
|
||||
self.pipeline_fps_number = Value('f', 0.0)
|
||||
self.output_pipeline_fps = Value('f', 0.0) # 由 main 进程更新，与 main print 一致
|
||||
self.input_fps = input_fps
|
||||
|
||||
self.finish_event = Event()
|
||||
|
||||
def run(self):
|
||||
# 插值帧红点标记：由 --mark_interpolated 控制，供 ezvtb_rt 读取
|
||||
if getattr(args, 'mark_interpolated', False):
|
||||
os.environ['EZVTB_MARK_INTERPOLATED'] = '1'
|
||||
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_guard = [
|
||||
SharedMemoryGuard(self.ret_shared_mem, ctrl_name=f"ret_shm_ctrl_batch_{i}")
|
||||
for i in range(args.interpolation_scale)
|
||||
]
|
||||
np_ret_shms = [
|
||||
np.ndarray((args.model_output_size, self.alpha_width_scale * args.model_output_size, self.ret_channels), dtype=np.uint8,
|
||||
buffer=self.ret_shared_mem.buf[i * self.alpha_width_scale * args.model_output_size * args.model_output_size * self.ret_channels:
|
||||
(i + 1) * self.alpha_width_scale * args.model_output_size * args.model_output_size * self.ret_channels])
|
||||
for i in range(args.interpolation_scale)
|
||||
]
|
||||
|
||||
model_infer_average_interval: Interval = Interval()
|
||||
pipeline_fps = FPS()
|
||||
|
||||
# Use unified ezvtb_rt interface for both THA3 and THA4
|
||||
model = get_core(use_tensorrt=args.use_tensorrt,
|
||||
model_version=args.model_version,
|
||||
model_name=args.model_name,
|
||||
|
||||
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_ram_size=args.max_ram_cache_len,
|
||||
|
||||
use_sr=args.use_sr,
|
||||
sr_half=args.sr_half,
|
||||
sr_x4=args.sr_x4,
|
||||
sr_a4k=args.sr_a4k,
|
||||
)
|
||||
model.setImage(self.input_image)
|
||||
model_infer_average_interval.start()
|
||||
model.inference([np.zeros((1, 45), dtype=np.float32)]) # Warm up
|
||||
model_infer_average_interval.stop()
|
||||
self.last_model_interval.value = model_infer_average_interval.last()
|
||||
|
||||
last_pose = np.zeros((45,), dtype=np.float32)
|
||||
|
||||
print("Model Inference Ready")
|
||||
while True:
|
||||
with pose_position_shm_guard.lock():
|
||||
np_pose = np_pose_shm.copy()
|
||||
np_position = np_position_shm.copy()
|
||||
|
||||
input_poses = []
|
||||
increment = (np_pose - last_pose) / args.interpolation_scale
|
||||
for i in range(args.interpolation_scale):
|
||||
input_poses.append(pose_simplify(last_pose + increment * (i + 1)))
|
||||
last_pose = np_pose
|
||||
|
||||
model_infer_average_interval.start()
|
||||
output_images = model.inference(input_poses)
|
||||
|
||||
if args.max_ram_cache_len > 0:
|
||||
hits = model.cacher.hits
|
||||
miss = model.cacher.miss
|
||||
if args.use_sr:
|
||||
hits += model.sr_cacher.hits
|
||||
miss += model.sr_cacher.miss
|
||||
total = hits + miss
|
||||
self.cache_hit_ratio.value = (hits / total) if total > 0 else 0.0
|
||||
|
||||
if args.use_tensorrt and args.max_gpu_cache_len > 0:
|
||||
hits = model.tha.cacher.hits
|
||||
miss = model.tha.cacher.miss
|
||||
total = hits + miss
|
||||
self.gpu_cache_hit_ratio.value = (hits / total) if total > 0 else 0.0
|
||||
|
||||
output_images = self.post_process_ret(np_position, output_images)
|
||||
|
||||
self.average_model_interval.value = model_infer_average_interval.stop()
|
||||
self.last_model_interval.value = model_infer_average_interval.last()
|
||||
|
||||
self.pipeline_fps_number.value = pipeline_fps()
|
||||
for i in range(args.interpolation_scale):
|
||||
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
|
||||
|
||||
def post_process_ret(self, np_position: np.ndarray, output_images: np.ndarray) -> List[np.ndarray]:
|
||||
k_scale = 1
|
||||
rotate_angle = 0
|
||||
dx = 0
|
||||
dy = 0
|
||||
if args.extend_movement:
|
||||
k_scale = np_position[2] * math.sqrt(args.extend_movement) + 1
|
||||
rotate_angle = -np_position[0] * 10 * args.extend_movement
|
||||
dx = np_position[0] * 400 * k_scale * args.extend_movement
|
||||
dy = -np_position[1] * 600 * k_scale * args.extend_movement
|
||||
if args.bongo:
|
||||
rotate_angle -= 5
|
||||
|
||||
rm = cv2.getRotationMatrix2D((output_images[0].shape[1] / 2, output_images[0].shape[0] / 2), rotate_angle, k_scale)
|
||||
rm[0, 2] += dx + output_images[0].shape[1] / 2 - output_images[0].shape[1] / 2
|
||||
rm[1, 2] += dy + output_images[0].shape[0] / 2 - output_images[0].shape[0] / 2
|
||||
|
||||
ret = []
|
||||
for i in range(output_images.shape[0]):
|
||||
bgra_image = output_images[i]
|
||||
bgra_image = cv2.warpAffine(
|
||||
bgra_image,
|
||||
rm,
|
||||
(bgra_image.shape[1], bgra_image.shape[0]))
|
||||
|
||||
if args.output_debug:
|
||||
# 与 main.py 输出格式一致
|
||||
y = 16
|
||||
cv2.putText(bgra_image, 'INFER/S: {:.4f}'.format(self.pipeline_fps_number.value), (0, y), cv2.FONT_HERSHEY_PLAIN, 1, (0, 255, 0), 1)
|
||||
y += 16
|
||||
cv2.putText(bgra_image, 'INPUT/S: {:.4f}'.format(self.input_fps.value), (0, y), cv2.FONT_HERSHEY_PLAIN, 1, (0, 255, 0), 1)
|
||||
y += 16
|
||||
cv2.putText(bgra_image, 'OUTPUT/S: {:.4f}'.format(self.output_pipeline_fps.value), (0, y), cv2.FONT_HERSHEY_PLAIN, 1, (0, 255, 0), 1)
|
||||
y += 16
|
||||
cv2.putText(bgra_image, 'CALC: {:.2f} ms'.format(self.average_model_interval.value * 1000), (0, y), cv2.FONT_HERSHEY_PLAIN, 1, (0, 255, 0), 1)
|
||||
y += 16
|
||||
if args.max_ram_cache_len > 0:
|
||||
cv2.putText(bgra_image, 'MEM CACHE: {:.2f}%'.format(self.cache_hit_ratio.value * 100), (0, y), cv2.FONT_HERSHEY_PLAIN, 1, (0, 255, 0), 1)
|
||||
y += 16
|
||||
if args.max_gpu_cache_len > 0:
|
||||
cv2.putText(bgra_image, 'GPU CACHE: {:.2f}%'.format(self.gpu_cache_hit_ratio.value * 100), (0, y), cv2.FONT_HERSHEY_PLAIN, 1, (0, 255, 0), 1)
|
||||
|
||||
if args.alpha_split:
|
||||
rgba_image = cv2.cvtColor(bgra_image, cv2.COLOR_BGRA2RGBA)
|
||||
alpha_channel = rgba_image[:, :, 3]
|
||||
rgb_channels = rgba_image[:,:,:3]
|
||||
alpha_image = cv2.cvtColor(alpha_channel, cv2.COLOR_GRAY2RGB)
|
||||
rgb_channels = cv2.hconcat([rgb_channels, alpha_image])
|
||||
|
||||
if args.output_debug:
|
||||
if args.alpha_split:
|
||||
bgr_channels = cv2.cvtColor(rgb_channels, cv2.COLOR_RGB2BGR)
|
||||
else:
|
||||
bgr_channels = cv2.cvtColor(bgra_image, cv2.COLOR_BGRA2BGR)
|
||||
ret.append(bgr_channels)
|
||||
elif args.output_virtual_cam:
|
||||
if not args.alpha_split:
|
||||
rgb_channels = cv2.cvtColor(bgra_image, cv2.COLOR_BGRA2RGB)
|
||||
ret.append(rgb_channels)
|
||||
else:
|
||||
rgba_image = cv2.cvtColor(bgra_image, cv2.COLOR_BGRA2RGBA)
|
||||
ret.append(rgba_image)
|
||||
return ret
|
||||
|
||||
if __name__ == "__main__":
|
||||
pass
|
||||
@@ -0,0 +1,183 @@
|
||||
from multiprocessing import Process, shared_memory, Value
|
||||
from pynput.mouse import Controller
|
||||
import numpy as np
|
||||
import time
|
||||
import sounddevice as sd
|
||||
from .args import args
|
||||
from .utils.shared_mem_guard import SharedMemoryGuard
|
||||
from .utils.timer_wait import wait_until
|
||||
from OneEuroFilter import OneEuroFilter
|
||||
from .utils.fps import FPS
|
||||
from .utils.timer_wait import wait_until
|
||||
|
||||
class MouseClientProcess(Process):
|
||||
def __init__(self, pose_position_shm: shared_memory.SharedMemory):
|
||||
super().__init__()
|
||||
self.pose_position_shm = pose_position_shm
|
||||
self.fps = Value('f', 60.0)
|
||||
self.audio_volume = Value('f', 0.0)
|
||||
self.audio_running = False
|
||||
self.audio_callback_fps = FPS(60)
|
||||
|
||||
def audio_callback(self, indata, frames, time_info, status):
|
||||
"""音频回调函数，计算音量"""
|
||||
if status:
|
||||
print(f"Audio status: {status}")
|
||||
|
||||
self.audio_callback_fps()
|
||||
# 计算 RMS (均方根) 音量
|
||||
volume_norm = np.linalg.norm(indata) * 10
|
||||
|
||||
# 应用阈值和灵敏度
|
||||
if volume_norm < args.audio_threshold:
|
||||
volume_norm = 0
|
||||
else:
|
||||
volume_norm = (volume_norm - args.audio_threshold) * args.audio_sensitivity
|
||||
# 限制在 0-1 范围
|
||||
volume_norm = np.clip(volume_norm, 0, 1)
|
||||
self.audio_volume.value = volume_norm
|
||||
|
||||
def run(self):
|
||||
mouse = Controller()
|
||||
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])
|
||||
|
||||
# posLimit = [0, 0, 1920, 1080]
|
||||
posLimit = [int(x) for x in args.mouse_input.split(',')]
|
||||
|
||||
if args.mouse_audio_input:
|
||||
# 启动音频流 (使用 WASAPI)
|
||||
try:
|
||||
# 获取默认输入设备（麦克风或系统音频）
|
||||
# 如果要捕获系统音频输出，需要使用 loopback 设备
|
||||
audio_stream = sd.InputStream(
|
||||
callback=self.audio_callback,
|
||||
channels=1,
|
||||
samplerate=16000,
|
||||
blocksize=1024,
|
||||
)
|
||||
audio_stream.start()
|
||||
self.audio_running = True
|
||||
wait_until(time.perf_counter() + 1.0) # 等待一秒以稳定音频输入
|
||||
print("Audio capture started (WASAPI)")
|
||||
except Exception as e:
|
||||
print(f"Failed to start audio capture: {e}")
|
||||
print("Continuing without audio input...")
|
||||
self.audio_running = False
|
||||
audio_filter = OneEuroFilter(freq=self.audio_callback_fps.view(), mincutoff=10.0, beta=0.0)
|
||||
|
||||
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,
|
||||
}
|
||||
last_time : float = time.perf_counter()
|
||||
interval : float = 1.0 / 60 # 60 FPS
|
||||
|
||||
# 眨眼相关参数
|
||||
blink_interval = args.blink_interval # 眨眼间隔时间（秒）
|
||||
blink_duration = 0.25 # 眨眼持续时间（秒）
|
||||
last_blink_time = time.perf_counter()
|
||||
|
||||
# 呼吸循环参数
|
||||
breath_start_time = time.perf_counter()
|
||||
|
||||
print("Mouse Input Running at 60 FPS")
|
||||
position_vector = np.array([0, 0, 0, 1], dtype=np.float32)
|
||||
while True:
|
||||
pos = mouse.position
|
||||
# print(pos)
|
||||
eye_limit = [0.8, 0.5]
|
||||
head_eye_reduce = 0.6
|
||||
head_slowness = 0.2
|
||||
|
||||
# 计算眨眼效果
|
||||
current_time = time.perf_counter()
|
||||
time_since_last_blink = current_time - last_blink_time
|
||||
|
||||
# 每隔 blink_interval 秒眨一次眼
|
||||
if time_since_last_blink >= blink_interval:
|
||||
last_blink_time = current_time
|
||||
time_since_last_blink = 0
|
||||
|
||||
# 计算眨眼动画（使用尖锐的 sin 函数）
|
||||
if time_since_last_blink < blink_duration:
|
||||
# 在眨眼持续时间内，使用 sin 函数从 0 升到接近 1 再返回 0
|
||||
# 使用 sin^2 使曲线更尖锐
|
||||
blink_progress = time_since_last_blink / blink_duration * np.pi
|
||||
blink_value = np.sin(blink_progress) ** 2
|
||||
eye_close = blink_value * 1.0 # 升到1，返回到 0
|
||||
else:
|
||||
eye_close = 0.0
|
||||
|
||||
# 计算呼吸效果（使用 sin 函数，在 breath_cycle 时间内从 0 到 1 再到 0）
|
||||
breath_elapsed = (current_time - breath_start_time) % args.breath_cycle
|
||||
# 使用 sin 函数，让值在一个周期内从 0 -> 1 -> 0
|
||||
# sin 在 0 到 π 之间从 0 到 1 到 0
|
||||
breath_value = np.sin(breath_elapsed / args.breath_cycle * np.pi)
|
||||
|
||||
# 获取音频音量并映射到 mouth_ratio
|
||||
current_mouth_ratio = audio_filter(self.audio_volume.value, time.perf_counter()) if self.audio_running else 0
|
||||
mouse_data = {
|
||||
'eye_l_h_temp': eye_close,
|
||||
'eye_r_h_temp': eye_close,
|
||||
'mouth_ratio': current_mouth_ratio,
|
||||
'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
|
||||
|
||||
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 = [0.0] * 12
|
||||
mouth_eye_vector = [0.0] * 27
|
||||
pose_vector = [0.0] * 6
|
||||
|
||||
mouth_eye_vector[2] = eye_l_h_temp
|
||||
mouth_eye_vector[3] = eye_r_h_temp
|
||||
|
||||
mouth_eye_vector[14] = mouth_ratio * 1.5
|
||||
|
||||
mouth_eye_vector[25] = eye_y_ratio
|
||||
mouth_eye_vector[26] = eye_x_ratio
|
||||
|
||||
pose_vector[0] = x_angle
|
||||
pose_vector[1] = y_angle
|
||||
pose_vector[2] = z_angle
|
||||
pose_vector[3] = pose_vector[1]
|
||||
pose_vector[4] = pose_vector[2]
|
||||
pose_vector[5] = breath_value
|
||||
|
||||
model_input_arr = eyebrow_vector
|
||||
model_input_arr.extend(mouth_eye_vector)
|
||||
model_input_arr.extend(pose_vector)
|
||||
|
||||
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)
|
||||
last_time += interval
|
||||
@@ -0,0 +1,181 @@
|
||||
from multiprocessing import Process, Value, shared_memory
|
||||
from .args import args
|
||||
import socket
|
||||
import struct
|
||||
import numpy as np
|
||||
from .utils.shared_mem_guard import SharedMemoryGuard
|
||||
from .utils.fps import FPS
|
||||
from .utils.filter import OneEuroFilterNumpy
|
||||
from OneEuroFilter import OneEuroFilter
|
||||
import time
|
||||
|
||||
class OSFClientProcess(Process):
|
||||
def __init__(self, pose_position_shm: shared_memory.SharedMemory):
|
||||
super().__init__()
|
||||
self.pose_position_shm = pose_position_shm
|
||||
self.address = args.osf_input.split(':')[0]
|
||||
self.port = int(args.osf_input.split(':')[1])
|
||||
self.fps = Value('f', 0.0)
|
||||
|
||||
def run(self):
|
||||
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])
|
||||
|
||||
self.socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
|
||||
self.socket.bind(("", self.port))
|
||||
self.socket.settimeout(None) # wait indefinitely
|
||||
|
||||
print("Warming up OpenSeeFace connection...")
|
||||
|
||||
frame_count = 0
|
||||
input_fps = FPS(60)
|
||||
while frame_count < 5:
|
||||
socket_bytes = self.socket.recv(8192)
|
||||
if not socket_bytes:
|
||||
raise Exception("Can't receive OpenSeeFace data (stream end?).")
|
||||
input_fps()
|
||||
frame_count += 1
|
||||
|
||||
position_vector_0 = None
|
||||
position_vector = [0, 0, 0, 1]
|
||||
# 呼吸循环参数
|
||||
breath_start_time = time.perf_counter()
|
||||
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.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:
|
||||
socket_bytes = self.socket.recv(8192)
|
||||
|
||||
if not socket_bytes:
|
||||
raise Exception("Can't receive OpenSeeFace data (stream end?).")
|
||||
|
||||
self.fps.value = input_fps()
|
||||
|
||||
try:
|
||||
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'
|
||||
]
|
||||
except Exception:
|
||||
print("OpenSeeFace data parse error:", socket_bytes)
|
||||
continue
|
||||
|
||||
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'])
|
||||
|
||||
# 计算呼吸效果（使用 sin 函数，在 breath_cycle 时间内从 0 到 1 再到 0）
|
||||
breath_elapsed = (time.perf_counter() - breath_start_time) % args.breath_cycle
|
||||
# 使用 sin 函数，让值在一个周期内从 0 -> 1 -> 0
|
||||
# sin 在 0 到 π 之间从 0 到 1 到 0
|
||||
breath_value = np.sin(breath_elapsed / args.breath_cycle * np.pi)
|
||||
|
||||
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]
|
||||
|
||||
eyebrow_vector = [0.0] * 12
|
||||
mouth_eye_vector = [0.0] * 27
|
||||
pose_vector = [0.0] * 6
|
||||
|
||||
mouth_eye_vector[2] = 1 - data['leftEyeOpen']
|
||||
mouth_eye_vector[3] = 1 - data['rightEyeOpen']
|
||||
|
||||
mouth_eye_vector[14] = max(data['MouthOpen'], 0) * 2 # Open larger mouth
|
||||
# print(mouth_eye_vector[14])
|
||||
|
||||
mouth_eye_vector[25] = iris_x_filter(-data['eyeRotationY'] * 3 - (data['rotationX']) / 57.3 * 1.5, timestamp=time.perf_counter())
|
||||
mouth_eye_vector[26] = iris_y_filter(data['eyeRotationX'] * 3 + (data['rotationY']) / 57.3, timestamp=time.perf_counter())
|
||||
# print(mouth_eye_vector[25:27])
|
||||
eyebrow_vector[6] = data['EyebrowUpDownLeft']
|
||||
eyebrow_vector[7] = data['EyebrowUpDownRight']
|
||||
# print(data['EyebrowUpDownLeft'],data['EyebrowUpDownRight'])
|
||||
if rotation_offset is None:
|
||||
rotation_offset = [data['rotationX'], data['rotationY'], data['rotationZ']]
|
||||
pose_vector[0] = (data['rotationX'] - rotation_offset[0]) / 57.3 * 3
|
||||
pose_vector[1] = -(data['rotationY'] - rotation_offset[1]) / 57.3 * 3
|
||||
pose_vector[2] = (data['rotationZ'] - rotation_offset[2]) / 57.3 * 2
|
||||
pose_vector[3] = pose_vector[1]
|
||||
pose_vector[4] = pose_vector[2]
|
||||
pose_vector[5] = breath_value
|
||||
|
||||
if position_vector_0 == None: #Provide an initial reference point
|
||||
position_vector_0 = [0, 0, 0, 1]
|
||||
position_vector_0[0] = data['translationX']
|
||||
position_vector_0[1] = data['translationY']
|
||||
position_vector_0[2] = data['translationZ']
|
||||
#Compute relative translation
|
||||
position_vector[0] = -(data['translationX'] - position_vector_0[0]) * 0.1
|
||||
position_vector[1] = -(data['translationY'] - position_vector_0[1]) * 0.1
|
||||
position_vector[2] = -(data['translationZ'] - position_vector_0[2]) * 0.1
|
||||
|
||||
model_input_arr = eyebrow_vector
|
||||
model_input_arr.extend(mouth_eye_vector)
|
||||
model_input_arr.extend(pose_vector)
|
||||
|
||||
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)
|
||||
Whitespace-only changes.
File renamed without changes.
@@ -0,0 +1,155 @@
|
||||
"""
|
||||
Numpy-compatible wrapper for OneEuroFilter.
|
||||
|
||||
Provides an adaptive One Euro Filter that automatically handles numpy arrays
|
||||
by maintaining separate filter instances for each element.
|
||||
"""
|
||||
|
||||
import numpy as np
|
||||
from OneEuroFilter import OneEuroFilter
|
||||
from typing import Optional
|
||||
import time
|
||||
|
||||
class OneEuroFilterNumpy:
|
||||
"""
|
||||
Auto-initializing One Euro Filter for numpy arrays.
|
||||
|
||||
Automatically determines shape from first input and maintains independent
|
||||
OneEuroFilter instances for each element in the array.
|
||||
|
||||
Example:
|
||||
>>> filter = AdaptiveOneEuroFilterNumpy(freq=30, mincutoff=1.0, beta=0.01)
|
||||
>>> noisy_data = np.array([1.0, 2.0, 3.0])
|
||||
>>> filtered = filter(noisy_data)
|
||||
>>> filtered = filter(noisy_data + np.random.randn(3) * 0.1)
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
freq: float,
|
||||
mincutoff: float = 1.0,
|
||||
beta: float = 0.0,
|
||||
dcutoff: float = 1.0
|
||||
) -> None:
|
||||
"""
|
||||
Initialize an adaptive One Euro Filter.
|
||||
|
||||
Args:
|
||||
freq: An estimate of the frequency in Hz of the signal (> 0)
|
||||
mincutoff: Min cutoff frequency in Hz (> 0). Lower values remove more jitter
|
||||
beta: Parameter to reduce latency (>= 0)
|
||||
dcutoff: Cutoff frequency for derivatives (> 0)
|
||||
"""
|
||||
self.freq = freq
|
||||
self.mincutoff = mincutoff
|
||||
self.beta = beta
|
||||
self.dcutoff = dcutoff
|
||||
self._filters = None
|
||||
self._shape = None
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
x: np.ndarray,
|
||||
timestamp: Optional[float] = None
|
||||
) -> np.ndarray:
|
||||
"""
|
||||
Filter a noisy numpy array.
|
||||
|
||||
On first call, initializes filters based on input shape.
|
||||
|
||||
Args:
|
||||
x: Noisy array to filter
|
||||
timestamp: Timestamp in seconds (optional)
|
||||
|
||||
Returns:
|
||||
Filtered array with same shape as input
|
||||
"""
|
||||
# Initialize filters on first call
|
||||
if self._filters is None:
|
||||
self._shape = x.shape
|
||||
size = np.prod(self._shape).astype(int)
|
||||
self._filters = [
|
||||
OneEuroFilter(
|
||||
freq=self.freq,
|
||||
mincutoff=self.mincutoff,
|
||||
beta=self.beta,
|
||||
dcutoff=self.dcutoff
|
||||
)
|
||||
for _ in range(size)
|
||||
]
|
||||
|
||||
# Validate shape
|
||||
if x.shape != self._shape:
|
||||
raise ValueError(
|
||||
f"Input shape {x.shape} doesn't match initialized shape {self._shape}"
|
||||
)
|
||||
|
||||
if timestamp is None:
|
||||
timestamp = time.perf_counter()
|
||||
|
||||
# Flatten, filter, reshape
|
||||
x_flat = x.flatten()
|
||||
filtered_flat = np.array([
|
||||
self._filters[i](float(val), timestamp)
|
||||
for i, val in enumerate(x_flat)
|
||||
])
|
||||
|
||||
return filtered_flat.reshape(self._shape)
|
||||
|
||||
def filter(
|
||||
self,
|
||||
x: np.ndarray,
|
||||
timestamp: Optional[float] = None
|
||||
) -> np.ndarray:
|
||||
"""Filter a noisy numpy array (alias for __call__)."""
|
||||
return self.__call__(x, timestamp)
|
||||
|
||||
def reset(self) -> None:
|
||||
"""Reset all internal filter states."""
|
||||
if self._filters is not None:
|
||||
for f in self._filters:
|
||||
f.reset()
|
||||
|
||||
def setFrequency(self, freq: float) -> None:
|
||||
"""Set the frequency for all filters."""
|
||||
self.freq = freq
|
||||
if self._filters is not None:
|
||||
for f in self._filters:
|
||||
f.setFrequency(freq)
|
||||
|
||||
def setMinCutoff(self, mincutoff: float) -> None:
|
||||
"""Set the min cutoff frequency for all filters."""
|
||||
self.mincutoff = mincutoff
|
||||
if self._filters is not None:
|
||||
for f in self._filters:
|
||||
f.setMinCutoff(mincutoff)
|
||||
|
||||
def setBeta(self, beta: float) -> None:
|
||||
"""Set the beta parameter for all filters."""
|
||||
self.beta = beta
|
||||
if self._filters is not None:
|
||||
for f in self._filters:
|
||||
f.setBeta(beta)
|
||||
|
||||
def setDerivateCutoff(self, dcutoff: float) -> None:
|
||||
"""Set the derivative cutoff frequency for all filters."""
|
||||
self.dcutoff = dcutoff
|
||||
if self._filters is not None:
|
||||
for f in self._filters:
|
||||
f.setDerivateCutoff(dcutoff)
|
||||
|
||||
def setParameters(
|
||||
self,
|
||||
freq: float,
|
||||
mincutoff: float = 1.0,
|
||||
beta: float = 0.0,
|
||||
dcutoff: float = 1.0
|
||||
) -> None:
|
||||
"""Set all parameters for all filters."""
|
||||
self.freq = freq
|
||||
self.mincutoff = mincutoff
|
||||
self.beta = beta
|
||||
self.dcutoff = dcutoff
|
||||
if self._filters is not None:
|
||||
for f in self._filters:
|
||||
f.setParameters(freq, mincutoff, beta, dcutoff)
|
||||
@@ -0,0 +1,55 @@
|
||||
import time
|
||||
import collections
|
||||
|
||||
class FPS:
|
||||
def __init__(self, avarageof=50):
|
||||
self.frametimestamps = collections.deque(maxlen=avarageof)
|
||||
|
||||
def __call__(self):
|
||||
self.frametimestamps.append(time.perf_counter())
|
||||
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.perf_counter()
|
||||
|
||||
class Interval:
|
||||
def __init__(self, avarageof=50):
|
||||
self.frame_intervals = collections.deque(maxlen=avarageof)
|
||||
self.start_time = None
|
||||
self.interval = 0.0
|
||||
self.moving_sum = 0.0
|
||||
|
||||
def start(self):
|
||||
self.start_time = time.perf_counter()
|
||||
|
||||
def stop(self):
|
||||
if self.start_time is not None:
|
||||
last_interval = time.perf_counter() - self.start_time
|
||||
|
||||
if len(self.frame_intervals) >= self.frame_intervals.maxlen:
|
||||
self.moving_sum -= self.frame_intervals[0]
|
||||
self.moving_sum += last_interval
|
||||
|
||||
self.frame_intervals.append(last_interval)
|
||||
self.start_time = None
|
||||
|
||||
if len(self.frame_intervals) > 0:
|
||||
self.interval = self.moving_sum / len(self.frame_intervals)
|
||||
return self.interval
|
||||
else:
|
||||
raise RuntimeError("FPSInterval.stop() called before start()")
|
||||
|
||||
def view(self):
|
||||
return self.interval
|
||||
|
||||
def last(self):
|
||||
return self.frame_intervals[-1] if len(self.frame_intervals) > 0 else 0.0
|
||||
@@ -12,10 +12,10 @@ 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 utils import get_distance
|
||||
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 ------------------------------------------
|
||||
@@ -0,0 +1,130 @@
|
||||
from ..args import args
|
||||
import numpy as np
|
||||
import math
|
||||
|
||||
# Delay import to avoid loading torch during module initialization
|
||||
ifm_converter = None
|
||||
|
||||
def _get_ifm_converter():
|
||||
global ifm_converter
|
||||
if ifm_converter is None:
|
||||
import tha2.poser.modes.mode_20_wx
|
||||
ifm_converter = tha2.poser.modes.mode_20_wx.IFacialMocapPoseConverter20()
|
||||
return ifm_converter
|
||||
|
||||
def pose_simplify(model_input):
|
||||
ifm_converter = _get_ifm_converter()
|
||||
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]
|
||||
|
||||
input_pose = np.zeros((1, 45), dtype=np.float32)
|
||||
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]
|
||||
return input_pose
|
||||
File renamed without changes.
@@ -0,0 +1,149 @@
|
||||
# windows_only_shm_lock.py
|
||||
from __future__ import annotations
|
||||
import ctypes
|
||||
import ctypes.wintypes as wt
|
||||
import hashlib
|
||||
import time
|
||||
from contextlib import contextmanager
|
||||
from multiprocessing import shared_memory
|
||||
from typing import Optional, Union
|
||||
|
||||
# ---- WinAPI for Named Mutex ----
|
||||
kernel32 = ctypes.WinDLL('kernel32', use_last_error=True)
|
||||
|
||||
# HANDLE CreateMutexW(LPSECURITY_ATTRIBUTES, BOOL, LPCWSTR);
|
||||
_CreateMutexW = kernel32.CreateMutexW
|
||||
_CreateMutexW.argtypes = [wt.LPVOID, wt.BOOL, wt.LPCWSTR]
|
||||
_CreateMutexW.restype = wt.HANDLE
|
||||
|
||||
# HANDLE OpenMutexW(DWORD, BOOL, LPCWSTR);
|
||||
_OpenMutexW = kernel32.OpenMutexW
|
||||
_OpenMutexW.argtypes = [wt.DWORD, wt.BOOL, wt.LPCWSTR]
|
||||
_OpenMutexW.restype = wt.HANDLE
|
||||
|
||||
# DWORD WaitForSingleObject(HANDLE, DWORD);
|
||||
_WaitForSingleObject = kernel32.WaitForSingleObject
|
||||
_WaitForSingleObject.argtypes = [wt.HANDLE, wt.DWORD]
|
||||
_WaitForSingleObject.restype = wt.DWORD
|
||||
|
||||
# BOOL ReleaseMutex(HANDLE);
|
||||
_ReleaseMutex = kernel32.ReleaseMutex
|
||||
_ReleaseMutex.argtypes = [wt.HANDLE]
|
||||
_ReleaseMutex.restype = wt.BOOL
|
||||
|
||||
# BOOL CloseHandle(HANDLE);
|
||||
_CloseHandle = kernel32.CloseHandle
|
||||
_CloseHandle.argtypes = [wt.HANDLE]
|
||||
_CloseHandle.restype = wt.BOOL
|
||||
|
||||
# Constants
|
||||
INFINITE = 0xFFFFFFFF
|
||||
WAIT_OBJECT_0 = 0x00000000
|
||||
WAIT_TIMEOUT = 0x00000102
|
||||
WAIT_ABANDONED = 0x00000080
|
||||
WAIT_FAILED = 0xFFFFFFFF
|
||||
MUTEX_ALL_ACCESS = 0x1F0001
|
||||
ERROR_ALREADY_EXISTS = 183
|
||||
|
||||
|
||||
class SharedMemoryGuard:
|
||||
"""
|
||||
Two-process exclusive access wrapper using Windows Named Mutex:
|
||||
- payload SharedMemory is provided externally
|
||||
- Uses Windows Named Mutex for cross-process synchronization
|
||||
- More direct and simpler than manual InterlockedCompareExchange implementation
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
payload: shared_memory.SharedMemory,
|
||||
ctrl_name: str
|
||||
):
|
||||
# 1) attach payload (external)
|
||||
self.payload_shm = payload
|
||||
|
||||
# 2) create/open named mutex for synchronization
|
||||
self.mutex_name = f"Global\\{ctrl_name}_mutex"
|
||||
self._mutex_handle = self._create_or_open_mutex(self.mutex_name)
|
||||
if self._mutex_handle is None or self._mutex_handle == 0:
|
||||
error = ctypes.get_last_error()
|
||||
raise OSError(f"Failed to create/open mutex {self.mutex_name}, error: {error}")
|
||||
|
||||
@staticmethod
|
||||
def _create_or_open_mutex(name: str) -> wt.HANDLE:
|
||||
"""Create or open a named mutex."""
|
||||
# Try to create the mutex
|
||||
handle = _CreateMutexW(None, False, name)
|
||||
if handle is None or handle == 0:
|
||||
# If creation failed, try to open existing mutex
|
||||
handle = _OpenMutexW(MUTEX_ALL_ACCESS, False, name)
|
||||
return handle
|
||||
|
||||
# ---- payload view: zero-copy for numpy ----
|
||||
def payload_view(self) -> memoryview:
|
||||
return self.payload_shm.buf # user can slice if needed
|
||||
|
||||
# ---- lock API ----
|
||||
def acquire(self, timeout_ms: Optional[int] = None) -> bool:
|
||||
"""
|
||||
Acquire exclusive access using Windows Mutex.
|
||||
- timeout_ms=None => infinite wait
|
||||
- spin parameter is kept for API compatibility but ignored (Mutex handles waiting efficiently)
|
||||
|
||||
Returns True if lock acquired, False if timeout.
|
||||
"""
|
||||
if timeout_ms is None:
|
||||
wait_time = INFINITE
|
||||
else:
|
||||
wait_time = max(0, timeout_ms)
|
||||
|
||||
result = _WaitForSingleObject(self._mutex_handle, wait_time)
|
||||
|
||||
if result == WAIT_OBJECT_0:
|
||||
# Successfully acquired the mutex
|
||||
return True
|
||||
elif result == WAIT_ABANDONED:
|
||||
# Previous owner died without releasing, but we now own it
|
||||
# This is still considered a successful acquisition
|
||||
return True
|
||||
elif result == WAIT_TIMEOUT:
|
||||
# Timeout occurred
|
||||
return False
|
||||
else: # WAIT_FAILED or other error
|
||||
error = ctypes.get_last_error()
|
||||
raise OSError(f"WaitForSingleObject failed with error: {error}")
|
||||
|
||||
def release(self) -> None:
|
||||
"""Release the mutex."""
|
||||
if not _ReleaseMutex(self._mutex_handle):
|
||||
error = ctypes.get_last_error()
|
||||
raise OSError(f"ReleaseMutex failed with error: {error}")
|
||||
|
||||
@contextmanager
|
||||
def lock(self, timeout_ms: Optional[int] = None):
|
||||
"""
|
||||
Context manager for acquiring and releasing the lock.
|
||||
- timeout_ms: timeout in milliseconds (None for infinite)
|
||||
- spin: kept for API compatibility but ignored
|
||||
"""
|
||||
if not self.acquire(timeout_ms=timeout_ms):
|
||||
raise TimeoutError("acquire() timeout")
|
||||
try:
|
||||
yield self
|
||||
finally:
|
||||
self.release()
|
||||
|
||||
# ---- lifecycle ----
|
||||
def close(self) -> None:
|
||||
"""Close handles and shared memory."""
|
||||
if self._mutex_handle and self._mutex_handle != 0:
|
||||
_CloseHandle(self._mutex_handle)
|
||||
self._mutex_handle = None
|
||||
self.payload_shm.close()
|
||||
|
||||
def unlink_control_if_owner(self) -> None:
|
||||
"""
|
||||
Named mutex is automatically cleaned up by Windows when all handles are closed.
|
||||
This method is kept for API compatibility but does nothing.
|
||||
"""
|
||||
pass
|
||||
@@ -0,0 +1,57 @@
|
||||
import ctypes
|
||||
import time
|
||||
|
||||
ntdll = ctypes.WinDLL("ntdll")
|
||||
|
||||
desired = ctypes.c_ulong(5000) # 0.5 ms = 5000 * 100ns
|
||||
current = ctypes.c_ulong()
|
||||
|
||||
ntdll.NtSetTimerResolution(desired, True, ctypes.byref(current)) # Set timer resolution to 0.5ms
|
||||
|
||||
# -------------------------------
|
||||
# 高精度等待函数
|
||||
# -------------------------------
|
||||
def wait_until(
|
||||
target_time: float,
|
||||
spin_threshold: float = 0.0005,
|
||||
sleep_min: float = 0.001,
|
||||
):
|
||||
"""
|
||||
等待直到 perf_counter() >= target_time
|
||||
|
||||
:param target_time: 目标时间（time.perf_counter() 的绝对时间，秒）
|
||||
:param spin_threshold: 最后 busy-wait 窗口（秒，默认 0.5ms）
|
||||
:param sleep_min: 允许 sleep 的最小时间（避免 0 sleep）
|
||||
"""
|
||||
|
||||
# 防止负时间 / NaN
|
||||
if not isinstance(target_time, (int, float)):
|
||||
return
|
||||
|
||||
while True:
|
||||
now = time.perf_counter()
|
||||
remaining = target_time - now
|
||||
|
||||
# 已经过点
|
||||
if remaining <= 0:
|
||||
return
|
||||
|
||||
# 进入 busy-wait 区间
|
||||
if remaining <= spin_threshold:
|
||||
break
|
||||
|
||||
# 可 sleep 区间
|
||||
sleep_time = remaining - spin_threshold
|
||||
|
||||
# 防止 sleep(0) 或极小值
|
||||
if sleep_time >= sleep_min:
|
||||
time.sleep(sleep_time)
|
||||
else:
|
||||
# 太小了，直接进入自旋
|
||||
break
|
||||
|
||||
# -------- busy-wait 精对齐 --------
|
||||
# 使用局部变量减少属性查找开销
|
||||
perf = time.perf_counter
|
||||
while perf() < target_time:
|
||||
pass
|
||||
@@ -1,4 +1,3 @@
|
||||
import torch
|
||||
import numpy as np
|
||||
|
||||
|
||||
@@ -1,258 +0,0 @@
|
||||
# from pyanime4k import ac
|
||||
# import pyanime4k
|
||||
#
|
||||
# parameters = ac.Parameters()
|
||||
# # enable HDN for ACNet
|
||||
# parameters.HDN = True
|
||||
#
|
||||
# a = ac.AC(
|
||||
# type=ac.ProcessorType.CPU_ACNet,
|
||||
# )
|
||||
# a.set_arguments(parameters)
|
||||
import struct
|
||||
|
||||
import cv2
|
||||
#
|
||||
# img = cv2.imread("character/test41.png")
|
||||
#
|
||||
# a.load_image_from_numpy(img,input_type=ac.AC_INPUT_BGR)
|
||||
#
|
||||
# a.get_processor_info()
|
||||
#
|
||||
# # start processing
|
||||
# a.process()
|
||||
#
|
||||
# a.show_image()
|
||||
|
||||
# import numpy as np
|
||||
# import zlib
|
||||
# import bz2
|
||||
# import lzma
|
||||
# import time
|
||||
#
|
||||
# postprocessed_image=np.load('out.npy')
|
||||
# tic=time.perf_counter()
|
||||
# res=None
|
||||
# for i in range(100):
|
||||
# res=zlib.compress(postprocessed_image)
|
||||
# print("zlib", len(res),(time.perf_counter() - tic) * 1000)
|
||||
# tic = time.perf_counter()
|
||||
# for i in range(100):
|
||||
# zlib.decompress(res)
|
||||
# print("zlib", len(res),(time.perf_counter() - tic) * 1000)
|
||||
# tic = time.perf_counter()
|
||||
# for i in range(100):
|
||||
# res=bz2.compress(postprocessed_image)
|
||||
# print("bz2",len(res), (time.perf_counter() - tic) * 1000)
|
||||
# tic = time.perf_counter()
|
||||
# for i in range(100):
|
||||
# bz2.decompress(res)
|
||||
# print("bz2", len(res),(time.perf_counter() - tic) * 1000)
|
||||
# tic = time.perf_counter()
|
||||
# for i in range(100):
|
||||
# res=lzma.compress(postprocessed_image)
|
||||
# print("lzma",len(res), (time.perf_counter() - tic) * 1000)
|
||||
# tic = time.perf_counter()
|
||||
# for i in range(100):
|
||||
# lzma.decompress(res)
|
||||
# print("lzma", len(res),(time.perf_counter() - tic) * 1000)
|
||||
# tic = time.perf_counter()
|
||||
#
|
||||
#
|
||||
# output_frame = cv2.cvtColor(postprocessed_image, cv2.COLOR_RGBA2BGRA)
|
||||
# cv2.imshow("frame", output_frame)
|
||||
# cv2.waitKey(100000)
|
||||
# preview upscaled image
|
||||
# img=a.save_image_to_numpy()
|
||||
#
|
||||
# print(img)
|
||||
# import cv2
|
||||
# import numpy as np
|
||||
#
|
||||
# rgb=[
|
||||
# [[1,2,3],[2,2,3],[3,2,3],[4,2,3]],
|
||||
# [[1,2,3],[2,2,3],[3,2,3],[4,2,3]],
|
||||
# [[1,2,3],[2,2,3],[3,2,3],[4,2,3]],
|
||||
# [[1,2,3],[2,2,3],[3,2,3],[4,2,3]]
|
||||
# ]
|
||||
# a=[
|
||||
# [1,2,3,4],
|
||||
# [1,2,3,4],
|
||||
# [1,2,3,4],
|
||||
# [1,2,3,4],
|
||||
# ]
|
||||
#
|
||||
# print(cv2.merge((np.array(rgb),np.array(a))))
|
||||
|
||||
|
||||
from multiprocessing import Value, Process, Queue
|
||||
import queue
|
||||
import socket
|
||||
import errno
|
||||
import time
|
||||
import tha2.poser.modes.mode_20_wx
|
||||
from tha2.mocap.ifacialmocap_constants import *
|
||||
import numpy as np
|
||||
|
||||
ifm='127.0.0.1:11573'
|
||||
|
||||
|
||||
|
||||
class ClientProcess(Process):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.queue = Queue()
|
||||
self.should_terminate = Value('b', False)
|
||||
self.address = ifm.split(':')[0]
|
||||
self.port = int(ifm.split(':')[1])
|
||||
self.perf_time = 0
|
||||
# self.ifm_converter = tha2.poser.modes.mode_20_wx.create_ifacialmocap_pose_converter()
|
||||
self.a_min=None
|
||||
self.a_max=None
|
||||
|
||||
def run(self):
|
||||
self.socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
|
||||
self.socket.setblocking(False)
|
||||
self.socket.bind(("", self.port))
|
||||
self.socket.settimeout(0.1)
|
||||
while True:
|
||||
if self.should_terminate.value:
|
||||
break
|
||||
try:
|
||||
socket_bytes = self.socket.recv(8192)
|
||||
except socket.error as e:
|
||||
err = e.args[0]
|
||||
if err == errno.EAGAIN or err == errno.EWOULDBLOCK or err == 'timed out':
|
||||
continue
|
||||
else:
|
||||
raise e
|
||||
|
||||
# socket_string = socket_bytes.decode("utf-8")
|
||||
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']
|
||||
data['rotationX']=(-data['rawEulerX']+180)%360
|
||||
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]
|
||||
# blender_data = json.loads(socket_string)
|
||||
# data = self.convert_from_blender_data(socket_string)
|
||||
# data=self.ifm_converter.convert(data)
|
||||
# print(data)
|
||||
# if(self.a_max is None):
|
||||
# self.a_max=[x for x in data]
|
||||
# else:
|
||||
# self.a_max=[max(data[i],self.a_max[i]) for i in range(len(data))]
|
||||
# if(self.a_min is None):
|
||||
# self.a_min=[x for x in data]
|
||||
# else:
|
||||
# self.a_min=[min(data[i],self.a_min[i]) for i in range(len(data))]
|
||||
# print([[self.a_min[i],self.a_max[i]] for i in range(len(data))])
|
||||
cur_time = time.perf_counter()
|
||||
fps = 1 / (cur_time - self.perf_time)
|
||||
self.perf_time = cur_time
|
||||
# print(fps)
|
||||
try:
|
||||
self.queue.put_nowait(data)
|
||||
except queue.Full:
|
||||
pass
|
||||
self.queue.close()
|
||||
self.socket.close()
|
||||
@staticmethod
|
||||
def convert_from_blender_data(blender_data):
|
||||
data = {}
|
||||
|
||||
for item in blender_data.split('|'):
|
||||
if item.find('#') != -1:
|
||||
k, arr = item.split('#')
|
||||
arr = [float(n) for n in arr.split(',')]
|
||||
data[k.replace("_L", "Left").replace("_R", "Right")] = arr
|
||||
elif item.find('-') != -1:
|
||||
k, v = item.split("-")
|
||||
data[k.replace("_L", "Left").replace("_R", "Right")] = float(v) / 100
|
||||
|
||||
to_rad = 57.3
|
||||
data[HEAD_BONE_X] = data["=head"][0] / to_rad
|
||||
data[HEAD_BONE_Y] = data["=head"][1] / to_rad
|
||||
data[HEAD_BONE_Z] = data["=head"][2] / to_rad
|
||||
data[HEAD_BONE_QUAT] = [data["=head"][3], data["=head"][4], data["=head"][5], 1]
|
||||
# print(data[HEAD_BONE_QUAT][2],min(data[EYE_BLINK_LEFT],data[EYE_BLINK_RIGHT]))
|
||||
data[RIGHT_EYE_BONE_X] = data["rightEye"][0] / to_rad
|
||||
data[RIGHT_EYE_BONE_Y] = data["rightEye"][1] / to_rad
|
||||
data[RIGHT_EYE_BONE_Z] = data["rightEye"][2] / to_rad
|
||||
data[LEFT_EYE_BONE_X] = data["leftEye"][0] / to_rad
|
||||
data[LEFT_EYE_BONE_Y] = data["leftEye"][1] / to_rad
|
||||
data[LEFT_EYE_BONE_Z] = data["leftEye"][2] / to_rad
|
||||
|
||||
return data
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
client_process = ClientProcess()
|
||||
client_process.daemon = True
|
||||
client_process.start()
|
||||
|
||||
while True:
|
||||
pass
|
||||
@@ -0,0 +1,232 @@
|
||||
import torch
|
||||
from os.path import join
|
||||
from pathlib import Path
|
||||
import sys
|
||||
import os
|
||||
sys.path.append(os.getcwd())
|
||||
sys.path.append(join(os.getcwd(), 'tha4', 'src'))
|
||||
import numpy as np
|
||||
import PIL.Image
|
||||
import cv2
|
||||
import onnx
|
||||
from onnxsim import simplify
|
||||
import onnxruntime as ort
|
||||
from torch import Tensor
|
||||
from torch.nn import Module
|
||||
from typing import List, Optional
|
||||
from torch.nn.functional import interpolate
|
||||
import matplotlib.pyplot as plt
|
||||
|
||||
# =================================================================================
|
||||
# START: Monkey Patch for torch.nn.functional.affine_grid (CORRECTED, DEVICE-AWARE)
|
||||
# =================================================================================
|
||||
import torch.nn.functional
|
||||
|
||||
def linspace_from_neg_one(num_steps, dtype=torch.float32, device='cpu'):
|
||||
r = torch.linspace(-1, 1, num_steps, dtype=torch.float32, device=device)
|
||||
if num_steps > 1: r = r * (num_steps - 1) / num_steps
|
||||
return r
|
||||
def create_grid(N, C, H, W, device='cpu'):
|
||||
grid = torch.empty((N, H, W, 3), dtype=torch.float32, device=device)
|
||||
grid.select(-1, 0).copy_(linspace_from_neg_one(W, device=device))
|
||||
grid.select(-1, 1).copy_(linspace_from_neg_one(H, device=device).unsqueeze_(-1))
|
||||
grid.select(-1, 2).fill_(1); return grid
|
||||
def custom_affine_grid(theta, size, align_corners=False):
|
||||
N, C, H, W = size; device = theta.device
|
||||
grid = create_grid(N, C, H, W, device=device)
|
||||
grid = grid.view(N, H * W, 3).bmm(theta.transpose(1, 2))
|
||||
grid = grid.view(N, H, W, 2); return grid
|
||||
def apply_affine_grid_patch():
|
||||
torch.nn.functional.affine_grid = custom_affine_grid
|
||||
print("✅ Applied monkey patch to torch.nn.functional.affine_grid (Device-aware version).")
|
||||
# =================================================================================
|
||||
# END: Monkey Patch
|
||||
# =================================================================================
|
||||
|
||||
apply_affine_grid_patch()
|
||||
|
||||
if len(sys.argv) != 2:
|
||||
print("Usage: python onnx_export_tha4.py <model_name>"); print("Available models: float, half")
|
||||
raise ValueError('Do not get model name')
|
||||
|
||||
MODEL_NAME = sys.argv[1]
|
||||
HALF = (MODEL_NAME == 'half')
|
||||
DEVICE_NAME = 'cuda:0'
|
||||
IMAGE_INPUT = os.path.join('.', 'data', 'images', 'character.png')
|
||||
|
||||
TMP_DIR = join('.', 'onnx_model_tha4', 'tmp'); Path(TMP_DIR).mkdir(parents=True, exist_ok=True)
|
||||
MODEL_DIR = join('.', 'onnx_model_tha4', 'fp32' if not HALF else 'fp16'); Path(MODEL_DIR).mkdir(parents=True, exist_ok=True)
|
||||
TMP_FILE_WRITE = join(TMP_DIR, 'tmp.onnx')
|
||||
|
||||
gpu_device = torch.device(DEVICE_NAME); cpu_device = torch.device('cpu')
|
||||
gpu_dtype = torch.float16 if HALF else torch.float32
|
||||
cpu_dtype = torch.float32
|
||||
|
||||
print(f"Using GPU device for validation: {gpu_device}, dtype: {gpu_dtype}")
|
||||
print(f"Using CPU device for ONNX export, dtype: {cpu_dtype}")
|
||||
|
||||
def torch_srgb_to_linear(x: torch.Tensor):
|
||||
x = torch.clip(x, 0.0, 1.0); return torch.where(torch.le(x, 0.04045), x / 12.92, ((x + 0.055) / 1.055) ** 2.4)
|
||||
def torch_linear_to_srgb(x):
|
||||
x = torch.clip(x, 0.0, 1.0); return torch.where(torch.le(x, 0.003130804953560372), x * 12.92, 1.055 * (x ** (1.0 / 2.4)) - 0.055)
|
||||
def numpy_linear_to_srgb(x):
|
||||
x = np.clip(x, 0.0, 1.0); return np.where(x <= 0.003130804953560372, x * 12.92, 1.055 * (x ** (1.0 / 2.4)) - 0.055)
|
||||
def load_tha4_poser(device: torch.device, use_half: bool):
|
||||
from tha4.poser.modes.mode_07 import create_poser
|
||||
m_files = { 'eyebrow_decomposer': 'data/models/tha4/eyebrow_decomposer.pt', 'eyebrow_morphing_combiner': 'data/models/tha4/eyebrow_morphing_combiner.pt', 'face_morpher': 'data/models/tha4/face_morpher.pt', 'body_morpher': 'data/models/tha4/body_morpher.pt', 'upscaler': 'data/models/tha4/upscaler.pt' }
|
||||
print(f"Loading THA4 poser onto {device}..."); poser = create_poser(device=device, module_file_names=m_files)
|
||||
if use_half and device.type == 'cuda': print("Converting poser modules to half precision (FP16)..."); poser.half()
|
||||
return poser
|
||||
|
||||
gpu_poser = load_tha4_poser(gpu_device, HALF)
|
||||
cpu_poser = load_tha4_poser(cpu_device, False)
|
||||
pose_size = gpu_poser.get_num_parameters(); print(f"Pose parameters: {pose_size}")
|
||||
|
||||
def resize_PIL_image(img, size=(512, 512)): return img.resize(size, PIL.Image.Resampling.LANCZOS)
|
||||
def extract_PIL_image_from_filelike(path): return PIL.Image.open(path)
|
||||
def extract_pytorch_image_from_PIL_image(img):
|
||||
arr = np.array(img)
|
||||
if arr.shape[2] != 4: raise ValueError("Input image must be RGBA")
|
||||
t = torch.from_numpy(arr).float() / 255.0
|
||||
t[:, :, :3] = torch_srgb_to_linear(t[:, :, :3])
|
||||
t = t.permute(2, 0, 1) * 2.0 - 1.0; return t
|
||||
|
||||
pil_image = resize_PIL_image(extract_PIL_image_from_filelike(IMAGE_INPUT))
|
||||
pt_img_gpu = extract_pytorch_image_from_PIL_image(pil_image).to(gpu_dtype).unsqueeze(0).to(gpu_device)
|
||||
zero_pose_gpu = torch.zeros(1, pose_size, dtype=gpu_dtype, device=gpu_device)
|
||||
with torch.no_grad(): poser_torch_res = gpu_poser.pose(pt_img_gpu, zero_pose_gpu)
|
||||
print(f"Reference GPU output shape: {poser_torch_res.shape}")
|
||||
|
||||
print("\n" + "="*80); print("Starting ONNX export process on CPU..."); print("="*80)
|
||||
cpu_modules = cpu_poser.get_modules()
|
||||
cv_img_numpy = cv2.imread(IMAGE_INPUT, cv2.IMREAD_UNCHANGED)
|
||||
cv_img_tensor_cpu = torch.from_numpy(cv_img_numpy)
|
||||
|
||||
class EyebrowDecomposerWrapper(Module):
|
||||
def __init__(self, obj): super().__init__(); self.m = obj
|
||||
def forward(self, img: Tensor) -> List[Tensor]:
|
||||
img_f = img.to(cpu_dtype) / 255.0; img_f[:, :, :3] = torch_srgb_to_linear(img_f[:, :, :3])
|
||||
# --- THE FIX: Swap BGR to RGB ---
|
||||
img_f = img_f[:, :, [2, 1, 0, 3]]
|
||||
img_c = img_f.permute(2, 0, 1).unsqueeze(0) * 2.0 - 1.0
|
||||
crop = img_c[:, :, 64:192, 192:320]; res = self.m(crop)
|
||||
return [res[3], res[0], img_c]
|
||||
eyebrow_decomposer_wrapper = EyebrowDecomposerWrapper(cpu_modules['eyebrow_decomposer']).eval()
|
||||
with torch.no_grad(): decomposer_res_cpu = eyebrow_decomposer_wrapper(cv_img_tensor_cpu)
|
||||
torch.onnx.export(eyebrow_decomposer_wrapper, cv_img_tensor_cpu, TMP_FILE_WRITE, opset_version=16, input_names=['input_image'], output_names=["background_layer", "eyebrow_layer", "image_prepared"])
|
||||
onnx_model = onnx.load(TMP_FILE_WRITE); onnx.checker.check_model(onnx_model); onnx_model_sim, check = simplify(onnx_model)
|
||||
onnx.save(onnx_model_sim if check else onnx_model, join(MODEL_DIR, "decomposer.onnx")); print("✓ Decomposer exported")
|
||||
|
||||
eyebrow_pose_cpu = torch.zeros(1, 12, dtype=cpu_dtype)
|
||||
class EyebrowCombinerWrapper(Module):
|
||||
def __init__(self, obj): super().__init__(); self.m = obj
|
||||
def forward(self, img: Tensor, bg: Tensor, layer: Tensor, pose: Tensor) -> Tensor:
|
||||
crop = img[:, :, 32:224, 160:352].clone()
|
||||
morphed = self.m(bg, layer, pose)[2]
|
||||
crop[:, :, 32:160, 32:160] = morphed; return crop
|
||||
eyebrow_combiner_wrapper = EyebrowCombinerWrapper(cpu_modules['eyebrow_morphing_combiner']).eval()
|
||||
combiner_inputs_cpu = (decomposer_res_cpu[2], decomposer_res_cpu[0], decomposer_res_cpu[1], eyebrow_pose_cpu)
|
||||
with torch.no_grad(): combiner_res_cpu = eyebrow_combiner_wrapper(*combiner_inputs_cpu)
|
||||
torch.onnx.export(eyebrow_combiner_wrapper, combiner_inputs_cpu, TMP_FILE_WRITE, opset_version=16, input_names=['image_prepared', 'eyebrow_background_layer', "eyebrow_layer", 'eyebrow_pose'], output_names=['eyebrow_image'])
|
||||
onnx_model = onnx.load(TMP_FILE_WRITE); onnx.checker.check_model(onnx_model); onnx_model_sim, check = simplify(onnx_model)
|
||||
onnx.save(onnx_model_sim if check else onnx_model, join(MODEL_DIR, "combiner.onnx")); print("✓ Combiner exported")
|
||||
|
||||
face_pose_cpu = torch.zeros(1, 27, dtype=cpu_dtype)
|
||||
class FaceMorpherWrapper(Module):
|
||||
def __init__(self, obj): super().__init__(); self.m = obj
|
||||
def forward(self, img: Tensor, crop: Tensor, pose: Tensor) -> List[Tensor]:
|
||||
morphed = self.m(crop, pose)[0]; full = img.clone()
|
||||
full[:, :, 32:224, 160:352] = morphed
|
||||
half = interpolate(full, size=(256, 256), mode='bilinear', align_corners=False); return [full, half]
|
||||
face_morpher_wrapper = FaceMorpherWrapper(cpu_modules['face_morpher']).eval()
|
||||
morpher_inputs_cpu = (decomposer_res_cpu[2], combiner_res_cpu, face_pose_cpu)
|
||||
with torch.no_grad(): morpher_res_cpu = face_morpher_wrapper(*morpher_inputs_cpu)
|
||||
torch.onnx.export(face_morpher_wrapper, morpher_inputs_cpu, TMP_FILE_WRITE, opset_version=16, input_names=['image_prepared', 'im_morpher_crop', 'face_pose'], output_names=['face_morphed_full', 'face_morphed_half'])
|
||||
onnx_model = onnx.load(TMP_FILE_WRITE); onnx.checker.check_model(onnx_model); onnx_model_sim, check = simplify(onnx_model)
|
||||
onnx.save(onnx_model_sim if check else onnx_model, join(MODEL_DIR, "morpher.onnx")); print("✓ Face Morpher exported")
|
||||
|
||||
rotation_pose_cpu = torch.zeros(1, 6, dtype=cpu_dtype)
|
||||
class BodyMorpherWrapper(Module):
|
||||
def __init__(self, obj): super().__init__(); self.m = obj
|
||||
def forward(self, img: Tensor, pose: Tensor) -> List[Tensor]:
|
||||
res = self.m(img, pose); return [res[0], res[3]]
|
||||
body_morpher_wrapper = BodyMorpherWrapper(cpu_modules['body_morpher']).eval()
|
||||
body_morpher_inputs_cpu = (morpher_res_cpu[1], rotation_pose_cpu)
|
||||
with torch.no_grad(): body_morpher_res_cpu = body_morpher_wrapper(*body_morpher_inputs_cpu)
|
||||
torch.onnx.export(body_morpher_wrapper, body_morpher_inputs_cpu, TMP_FILE_WRITE, opset_version=16, input_names=['face_morphed_half', 'rotation_pose'], output_names=['half_res_posed_image', 'half_res_grid_change'])
|
||||
onnx_model = onnx.load(TMP_FILE_WRITE); onnx.checker.check_model(onnx_model); onnx_model_sim, check = simplify(onnx_model)
|
||||
onnx.save(onnx_model_sim if check else onnx_model, join(MODEL_DIR, "body_morpher.onnx")); print("✓ Body Morpher exported")
|
||||
|
||||
class UpscalerWrapper(Module):
|
||||
def __init__(self, obj): super().__init__(); self.m = obj
|
||||
def forward(self, rest_img: Tensor, posed_half: Tensor, grid_half: Tensor, pose: Tensor) -> List[Tensor]:
|
||||
posed_full = interpolate(posed_half, size=(512, 512), mode='bilinear', align_corners=False)
|
||||
grid_full = interpolate(grid_half, size=(512, 512), mode='bilinear', align_corners=False)
|
||||
res = self.m(rest_img, posed_full, grid_full, pose)[0]
|
||||
output = res / 2.0 + 0.5; output[:, :3, :, :] = torch_linear_to_srgb(output[:, :3, :, :])
|
||||
cv_res = (output.squeeze(0).permute(1, 2, 0)[:, :, [2, 1, 0, 3]] * 255).clamp(0, 255).to(torch.uint8)
|
||||
return [output, cv_res]
|
||||
upscaler_wrapper = UpscalerWrapper(cpu_modules['upscaler']).eval()
|
||||
upscaler_inputs_cpu = (morpher_res_cpu[0], body_morpher_res_cpu[0], body_morpher_res_cpu[1], rotation_pose_cpu)
|
||||
with torch.no_grad(): upscaler_res_cpu = upscaler_wrapper(*upscaler_inputs_cpu)
|
||||
torch.onnx.export(upscaler_wrapper, upscaler_inputs_cpu, TMP_FILE_WRITE, opset_version=16, input_names=['rest_image', 'half_res_posed_image', 'half_res_grid_change', 'rotation_pose'], output_names=['result', 'cv_result'])
|
||||
onnx_model = onnx.load(TMP_FILE_WRITE); onnx.checker.check_model(onnx_model); onnx_model_sim, check = simplify(onnx_model)
|
||||
onnx.save(onnx_model_sim if check else onnx_model, join(MODEL_DIR, "upscaler.onnx")); print("✓ Upscaler exported")
|
||||
|
||||
print("\n" + "="*80); print("Validation: Running ONNX inference..."); print("="*80)
|
||||
class THA4ONNXRunner:
|
||||
def __init__(self, model_dir, use_gpu=True):
|
||||
providers = [("CUDAExecutionProvider", {"device_id": 0})] if use_gpu else ['CPUExecutionProvider']
|
||||
self.decomposer_sess = self._create_session(join(model_dir, 'decomposer.onnx'), providers)
|
||||
self.combiner_sess = self._create_session(join(model_dir, "combiner.onnx"), providers)
|
||||
self.morpher_sess = self._create_session(join(model_dir, "morpher.onnx"), providers)
|
||||
self.body_morpher_sess = self._create_session(join(model_dir, "body_morpher.onnx"), providers)
|
||||
self.upscaler_sess = self._create_session(join(model_dir, "upscaler.onnx"), providers)
|
||||
print("✓ All ONNX models loaded")
|
||||
def _create_session(self, path, providers):
|
||||
try: return ort.InferenceSession(path, providers=providers)
|
||||
except Exception: print(f"Failed with {providers}, falling back to CPU for {os.path.basename(path)}"); return ort.InferenceSession(path, providers=['CPUExecutionProvider'])
|
||||
def run(self, img_path, p_eye, p_face, p_rot):
|
||||
img = cv2.imread(img_path, cv2.IMREAD_UNCHANGED)
|
||||
d_res = self.decomposer_sess.run(None, {'input_image': img})
|
||||
c_res = self.combiner_sess.run(None, {'image_prepared': d_res[2], 'eyebrow_background_layer': d_res[0], "eyebrow_layer": d_res[1], 'eyebrow_pose': p_eye})
|
||||
m_res = self.morpher_sess.run(None, {'image_prepared': d_res[2], 'im_morpher_crop': c_res[0], 'face_pose': p_face})
|
||||
b_res = self.body_morpher_sess.run(None, {'face_morphed_half': m_res[1], 'rotation_pose': p_rot})
|
||||
u_res = self.upscaler_sess.run(None, {'rest_image': m_res[0], 'half_res_posed_image': b_res[0], 'half_res_grid_change': b_res[1], 'rotation_pose': p_rot})
|
||||
return u_res[0]
|
||||
|
||||
def save_comparison_images(onnx_img_np, torch_img_np, prefix=""):
|
||||
onnx_hwc = np.transpose(onnx_img_np.squeeze(0), (1, 2, 0))[:, :, :3]
|
||||
torch_hwc = np.transpose(torch_img_np.squeeze(0), (1, 2, 0))[:, :, :3]
|
||||
onnx_hwc = np.clip(onnx_hwc, 0, 1)
|
||||
torch_hwc = np.clip(torch_hwc, 0, 1)
|
||||
onnx_path = f"{prefix}onnx_output.png"; torch_path = f"{prefix}pytorch_output.png"
|
||||
plt.imsave(onnx_path, onnx_hwc); plt.imsave(torch_path, torch_hwc)
|
||||
print(f"✅ Saved comparison images: {onnx_path} and {torch_path}")
|
||||
|
||||
runner = THA4ONNXRunner(MODEL_DIR)
|
||||
p_eye_z, p_face_z, p_rot_z = np.zeros((1, 12), dtype=np.float32), np.zeros((1, 27), dtype=np.float32), np.zeros((1, 6), dtype=np.float32)
|
||||
|
||||
print("\nRunning ONNX inference with zero pose...")
|
||||
onnx_output = runner.run(IMAGE_INPUT, p_eye_z, p_face_z, p_rot_z)
|
||||
ref_np = poser_torch_res.cpu().numpy() / 2.0 + 0.5; ref_np[0, :3, :, :] = numpy_linear_to_srgb(ref_np[0, :3, :, :])
|
||||
save_comparison_images(onnx_output, ref_np, "zero_pose_")
|
||||
mse = ((onnx_output - ref_np) ** 2).mean(); print(f"\n{'='*80}\nMSE (zero pose): {mse:.10f}\n{'='*80}")
|
||||
ACCEPTABLE_MSE_THRESHOLD = 1e-2
|
||||
if mse < ACCEPTABLE_MSE_THRESHOLD: print(f"✅ Validation PASSED! MSE is within the acceptable range of {ACCEPTABLE_MSE_THRESHOLD}.")
|
||||
else: print(f"⚠️ Warning: MSE is higher than expected! But this can be normal due to precision differences.")
|
||||
|
||||
print("\n" + "="*80); print("Testing with random poses..."); print("="*80)
|
||||
p_eye_r = np.random.randn(1, 12).astype(np.float32) * 0.5
|
||||
p_face_r = np.random.randn(1, 27).astype(np.float32) * 0.5
|
||||
p_rot_r = np.random.randn(1, 6).astype(np.float32) * 0.3
|
||||
onnx_output_r = runner.run(IMAGE_INPUT, p_eye_r, p_face_r, p_rot_r)
|
||||
pose_r_gpu = torch.from_numpy(np.concatenate([p_eye_r, p_face_r, p_rot_r], axis=1)).to(gpu_dtype).to(gpu_device)
|
||||
with torch.no_grad(): ref_output_r = gpu_poser.pose(pt_img_gpu, pose_r_gpu)
|
||||
ref_np_r = ref_output_r.cpu().numpy() / 2.0 + 0.5; ref_np_r[0, :3, :, :] = numpy_linear_to_srgb(ref_np_r[0, :3, :, :])
|
||||
save_comparison_images(onnx_output_r, ref_np_r, "random_pose_")
|
||||
mse_r = ((onnx_output_r - ref_np_r) ** 2).mean(); print(f"MSE (random pose): {mse_r:.10f}")
|
||||
if mse_r < ACCEPTABLE_MSE_THRESHOLD: print(f"✅ Random pose test PASSED! (Threshold: {ACCEPTABLE_MSE_THRESHOLD})")
|
||||
else: print(f"⚠️ Warning: High MSE with random pose! But this can be normal.")
|
||||
|
||||
print(f"\n{'='*80}\nExport and validation completed!\nModels saved to: {MODEL_DIR}\n{'='*80}")
|
||||
@@ -0,0 +1,470 @@
|
||||
import torch
|
||||
from os.path import join
|
||||
from pathlib import Path
|
||||
import sys
|
||||
import os
|
||||
sys.path.append(os.getcwd())
|
||||
sys.path.append(join(os.getcwd(), 'tha4', 'src'))
|
||||
import numpy as np
|
||||
import PIL.Image
|
||||
import cv2
|
||||
import onnx
|
||||
from onnxsim import simplify
|
||||
import onnxruntime as ort
|
||||
from torch import Tensor
|
||||
from torch.nn import Module
|
||||
from typing import List
|
||||
import matplotlib.pyplot as plt
|
||||
from onnxconverter_common import float16
|
||||
|
||||
# =================================================================================
|
||||
# START: Monkey Patch for torch.nn.functional.affine_grid (DEVICE-AWARE)
|
||||
# =================================================================================
|
||||
import torch.nn.functional
|
||||
|
||||
def linspace_from_neg_one(num_steps, dtype=torch.float32, device='cpu'):
|
||||
r = torch.linspace(-1, 1, num_steps, dtype=torch.float32, device=device)
|
||||
if num_steps > 1:
|
||||
r = r * (num_steps - 1) / num_steps
|
||||
return r
|
||||
|
||||
def create_grid(N, C, H, W, device='cpu'):
|
||||
grid = torch.empty((N, H, W, 3), dtype=torch.float32, device=device)
|
||||
grid.select(-1, 0).copy_(linspace_from_neg_one(W, device=device))
|
||||
grid.select(-1, 1).copy_(linspace_from_neg_one(H, device=device).unsqueeze_(-1))
|
||||
grid.select(-1, 2).fill_(1)
|
||||
return grid
|
||||
|
||||
def custom_affine_grid(theta, size, align_corners=False):
|
||||
N, C, H, W = size
|
||||
device = theta.device
|
||||
grid = create_grid(N, C, H, W, device=device)
|
||||
grid = grid.view(N, H * W, 3).bmm(theta.transpose(1, 2))
|
||||
grid = grid.view(N, H, W, 2)
|
||||
return grid
|
||||
|
||||
def apply_affine_grid_patch():
|
||||
torch.nn.functional.affine_grid = custom_affine_grid
|
||||
print("✅ Applied monkey patch to torch.nn.functional.affine_grid (Device-aware version).")
|
||||
|
||||
# =================================================================================
|
||||
# END: Monkey Patch
|
||||
# =================================================================================
|
||||
|
||||
apply_affine_grid_patch()
|
||||
|
||||
# Parse command line arguments
|
||||
if len(sys.argv) < 2:
|
||||
print("Usage: python onnx_export_tha4_student.py <model_name> [--fp32 or --fp16] [--cpu]")
|
||||
print("Example: python onnx_export_tha4_student.py kanori_tha4 --fp32")
|
||||
print(" python onnx_export_tha4_student.py kanori_tha4 --fp32 --cpu")
|
||||
raise ValueError('Missing model name')
|
||||
|
||||
MODEL_NAME = sys.argv[1]
|
||||
HALF = False
|
||||
FORCE_CPU = '--cpu' in sys.argv
|
||||
|
||||
# Check CUDA availability
|
||||
if FORCE_CPU:
|
||||
DEVICE_NAME = 'cpu'
|
||||
print("⚠️ Forced CPU mode enabled")
|
||||
elif not torch.cuda.is_available():
|
||||
DEVICE_NAME = 'cpu'
|
||||
print("⚠️ CUDA not available, automatically using CPU mode")
|
||||
else:
|
||||
DEVICE_NAME = 'cuda:0'
|
||||
|
||||
# Use image from model folder
|
||||
MODEL_BASE_PATH = join('.', 'data', 'models', 'custom_tha4_models', MODEL_NAME)
|
||||
IMAGE_INPUT = os.path.join(MODEL_BASE_PATH, 'character.png')
|
||||
|
||||
if not os.path.exists(IMAGE_INPUT):
|
||||
raise FileNotFoundError(f"Image not found: {IMAGE_INPUT}")
|
||||
|
||||
# Create output directories (save models in the model folder)
|
||||
TMP_DIR = join(MODEL_BASE_PATH, 'tmp')
|
||||
Path(TMP_DIR).mkdir(parents=True, exist_ok=True)
|
||||
MODEL_DIR = MODEL_BASE_PATH # Save directly in model folder
|
||||
TMP_FILE_WRITE = join(TMP_DIR, 'tmp.onnx')
|
||||
|
||||
validation_device = torch.device(DEVICE_NAME)
|
||||
export_device = torch.device('cpu')
|
||||
validation_dtype = torch.float16 if (HALF and DEVICE_NAME != 'cpu') else torch.float32
|
||||
export_dtype = torch.float32
|
||||
|
||||
if DEVICE_NAME == 'cpu':
|
||||
print(f"Using CPU device for both validation and ONNX export, dtype: {export_dtype}")
|
||||
else:
|
||||
print(f"Using GPU device for validation: {validation_device}, dtype: {validation_dtype}")
|
||||
print(f"Using CPU device for ONNX export, dtype: {export_dtype}")
|
||||
|
||||
# Color space conversion utilities
|
||||
def torch_srgb_to_linear(x: torch.Tensor):
|
||||
x = torch.clip(x, 0.0, 1.0)
|
||||
return torch.where(torch.le(x, 0.04045), x / 12.92, ((x + 0.055) / 1.055) ** 2.4)
|
||||
|
||||
def torch_linear_to_srgb(x):
|
||||
x = torch.clip(x, 0.0, 1.0)
|
||||
return torch.where(torch.le(x, 0.003130804953560372), x * 12.92, 1.055 * (x ** (1.0 / 2.4)) - 0.055)
|
||||
|
||||
def numpy_linear_to_srgb(x):
|
||||
x = np.clip(x, 0.0, 1.0)
|
||||
return np.where(x <= 0.003130804953560372, x * 12.92, 1.055 * (x ** (1.0 / 2.4)) - 0.055)
|
||||
|
||||
# Load THA4 Student Model (mode_14: 2-step architecture)
|
||||
def load_tha4_student_poser(device: torch.device, use_half: bool, model_base_path: str):
|
||||
from tha4.poser.modes.mode_14 import create_poser
|
||||
|
||||
# Model file paths
|
||||
m_files = {
|
||||
'face_morpher': join(model_base_path, 'face_morpher.pt'),
|
||||
'body_morpher': join(model_base_path, 'body_morpher.pt')
|
||||
}
|
||||
|
||||
print(f"Loading THA4 Student poser from {model_base_path}...")
|
||||
print(f" - Face Morpher: {m_files['face_morpher']}")
|
||||
print(f" - Body Morpher: {m_files['body_morpher']}")
|
||||
|
||||
poser = create_poser(device=device, module_file_names=m_files)
|
||||
|
||||
if use_half and device.type == 'cuda':
|
||||
print("Converting poser modules to half precision (FP16)...")
|
||||
poser.half()
|
||||
|
||||
return poser
|
||||
|
||||
# Load posers for validation and export
|
||||
validation_poser = load_tha4_student_poser(validation_device, HALF and DEVICE_NAME != 'cpu', MODEL_BASE_PATH)
|
||||
if DEVICE_NAME == 'cpu':
|
||||
# If using CPU for validation, reuse the same poser for export
|
||||
export_poser = validation_poser
|
||||
print("Using same poser for validation and export (CPU mode)")
|
||||
else:
|
||||
# If using GPU for validation, load separate CPU poser for export
|
||||
export_poser = load_tha4_student_poser(export_device, False, MODEL_BASE_PATH)
|
||||
|
||||
pose_size = validation_poser.get_num_parameters()
|
||||
print(f"Pose parameters: {pose_size}")
|
||||
|
||||
# Image loading and preprocessing utilities
|
||||
def resize_PIL_image(img, size=(512, 512)):
|
||||
return img.resize(size, PIL.Image.Resampling.LANCZOS)
|
||||
|
||||
def extract_PIL_image_from_filelike(path):
|
||||
return PIL.Image.open(path)
|
||||
|
||||
def extract_pytorch_image_from_PIL_image(img):
|
||||
arr = np.array(img)
|
||||
if arr.shape[2] != 4:
|
||||
raise ValueError("Input image must be RGBA")
|
||||
t = torch.from_numpy(arr).float() / 255.0
|
||||
t[:, :, :3] = torch_srgb_to_linear(t[:, :, :3])
|
||||
t = t.permute(2, 0, 1) * 2.0 - 1.0
|
||||
return t
|
||||
|
||||
# Load test image and run reference inference
|
||||
pil_image = resize_PIL_image(extract_PIL_image_from_filelike(IMAGE_INPUT))
|
||||
pt_img_validation = extract_pytorch_image_from_PIL_image(pil_image).to(validation_dtype).unsqueeze(0).to(validation_device)
|
||||
zero_pose_validation = torch.zeros(1, pose_size, dtype=validation_dtype, device=validation_device)
|
||||
|
||||
with torch.no_grad():
|
||||
poser_torch_res = validation_poser.pose(pt_img_validation, zero_pose_validation)
|
||||
print(f"Reference output shape: {poser_torch_res.shape}")
|
||||
|
||||
print("\n" + "="*80)
|
||||
print("Starting ONNX export process...")
|
||||
print("="*80)
|
||||
|
||||
# Get export modules
|
||||
export_modules = export_poser.get_modules()
|
||||
cv_img_numpy = cv2.imread(IMAGE_INPUT, cv2.IMREAD_UNCHANGED)
|
||||
cv_img_tensor_export = torch.from_numpy(cv_img_numpy)
|
||||
|
||||
# =================================================================================
|
||||
# Export Face Morpher (Student Model Network 1)
|
||||
# =================================================================================
|
||||
print("\n[1/2] Exporting Face Morpher...")
|
||||
|
||||
class FaceMorpherWrapper(Module):
|
||||
"""
|
||||
Face Morpher for Student Model (mode_14)
|
||||
- Input: 45-dimensional pose (first 39 used for face)
|
||||
- Output: 128x128 face region
|
||||
Note: SIREN-based face morpher doesn't need input image, only pose
|
||||
"""
|
||||
def __init__(self, face_morpher_module):
|
||||
super().__init__()
|
||||
self.face_morpher = face_morpher_module
|
||||
|
||||
def forward(self, pose: Tensor) -> Tensor:
|
||||
"""
|
||||
Args:
|
||||
pose: Full pose parameters [1, 39]
|
||||
Returns:
|
||||
face_output: Face morphed image [1, 4, 128, 128] in normalized format
|
||||
"""
|
||||
# Run face morpher (SIREN generates image from pose only)
|
||||
face_output = self.face_morpher.forward(pose)
|
||||
|
||||
return face_output
|
||||
|
||||
face_morpher_wrapper = FaceMorpherWrapper(export_modules['face_morpher']).eval()
|
||||
|
||||
# Test inputs
|
||||
face_pose_export = torch.zeros(1, 39, dtype=export_dtype)
|
||||
|
||||
# Export
|
||||
with torch.no_grad():
|
||||
face_morpher_res_export = face_morpher_wrapper(face_pose_export)
|
||||
|
||||
torch.onnx.export(
|
||||
face_morpher_wrapper,
|
||||
face_pose_export,
|
||||
TMP_FILE_WRITE,
|
||||
opset_version=16,
|
||||
input_names=['pose'],
|
||||
output_names=['face_morphed'],
|
||||
external_data =False,
|
||||
optimize =True,
|
||||
verify =True,
|
||||
dynamo =False,
|
||||
)
|
||||
|
||||
onnx_model = onnx.load(TMP_FILE_WRITE)
|
||||
onnx.checker.check_model(onnx_model)
|
||||
print(f" Original model inputs: {[inp.name for inp in onnx_model.graph.input]}")
|
||||
|
||||
# onnx_model = float16.convert_float_to_float16(onnx_model, keep_io_types=True)
|
||||
# Simplify with input/output names preservation
|
||||
onnx_model_sim, check = simplify(onnx_model, skip_fuse_bn=True)
|
||||
if check:
|
||||
print(f" Simplified model inputs: {[inp.name for inp in onnx_model_sim.graph.input]}")
|
||||
onnx.save(onnx_model_sim, join(MODEL_DIR, "face_morpher.onnx"))
|
||||
else:
|
||||
print(" Simplification failed, using original model")
|
||||
onnx.save(onnx_model, join(MODEL_DIR, "face_morpher.onnx"))
|
||||
|
||||
print("✓ Face Morpher exported")
|
||||
|
||||
# =================================================================================
|
||||
# Export Body Morpher (Student Model Network 2)
|
||||
# =================================================================================
|
||||
print("\n[2/2] Exporting Body Morpher...")
|
||||
|
||||
class BodyMorpherWrapper(Module):
|
||||
"""
|
||||
Body Morpher for Student Model (mode_14)
|
||||
- Input: Full 512x512 image with face morpher output embedded
|
||||
- Output: Final morphed image
|
||||
"""
|
||||
def __init__(self, body_morpher_module):
|
||||
super().__init__()
|
||||
self.body_morpher = body_morpher_module
|
||||
|
||||
def forward(self, img: Tensor, face_morphed: Tensor, pose: Tensor) -> List[Tensor]:
|
||||
"""
|
||||
Args:
|
||||
img: Input image [H, W, 4] in uint8 BGR+A format
|
||||
face_morphed: Face morphed output [1, 4, 128, 128] from face morpher
|
||||
pose: Full pose parameters [1, 45]
|
||||
Returns:
|
||||
result: Final output in [0, 255] uint8 format (OpenCV compatible)
|
||||
"""
|
||||
# Convert image to normalized format
|
||||
img_f = img.to(export_dtype) / 255.0
|
||||
img_f[0, :, :, :3] = torch_srgb_to_linear(img_f[0, :, :, :3])
|
||||
img_f = img_f[:, :, :, [2, 1, 0, 3]] # BGR to RGB
|
||||
img_c = img_f.permute(0, 3, 1, 2) * 2.0 - 1.0
|
||||
|
||||
# Embed face morpher output into the full image
|
||||
center_x = 256
|
||||
center_y = 128 + 16 # 144
|
||||
img_with_face = img_c.clone()
|
||||
img_with_face[:, :, center_y-64:center_y+64, center_x-64:center_x+64] = face_morphed
|
||||
|
||||
# Run body morpher
|
||||
body_output = self.body_morpher.forward(img_with_face, pose)
|
||||
|
||||
# Take the first output (blended image)
|
||||
result = body_output[0]
|
||||
|
||||
# Convert back to [0, 1] range
|
||||
output = result / 2.0 + 0.5
|
||||
output[:, :3, :, :] = torch_linear_to_srgb(output[:, :3, :, :])
|
||||
|
||||
# Convert to OpenCV format (BGR, uint8)
|
||||
cv_result = (output.squeeze(0).permute(1, 2, 0)[:, :, [2, 1, 0, 3]] * 255).clamp(0, 255).to(torch.uint8)
|
||||
|
||||
return [output, cv_result]
|
||||
|
||||
body_morpher_wrapper = BodyMorpherWrapper(export_modules['body_morpher']).eval()
|
||||
|
||||
full_pose_export = torch.zeros(1, 45, dtype=export_dtype)
|
||||
|
||||
cv_img_tensor_export = cv_img_tensor_export.unsqueeze(0)
|
||||
# Export
|
||||
with torch.no_grad():
|
||||
body_morpher_res_export = body_morpher_wrapper(cv_img_tensor_export, face_morpher_res_export, full_pose_export)
|
||||
|
||||
torch.onnx.export(
|
||||
body_morpher_wrapper,
|
||||
(cv_img_tensor_export, face_morpher_res_export, full_pose_export),
|
||||
TMP_FILE_WRITE,
|
||||
opset_version=16,
|
||||
input_names=['input_image', 'face_morphed', 'pose'],
|
||||
output_names=['result', 'cv_result'],
|
||||
external_data =False,
|
||||
optimize =True,
|
||||
verify =True,
|
||||
dynamo =False,
|
||||
)
|
||||
|
||||
onnx_model = onnx.load(TMP_FILE_WRITE)
|
||||
onnx.checker.check_model(onnx_model)
|
||||
print(f" Original model inputs: {[inp.name for inp in onnx_model.graph.input]}")
|
||||
|
||||
# onnx_model = float16.convert_float_to_float16(onnx_model, keep_io_types=True)
|
||||
# for node in onnx_model.graph.node:
|
||||
# if node.name == "/Cast":
|
||||
# for attr in node.attribute:
|
||||
# if attr.name == "to":
|
||||
# attr.i = onnx.TensorProto.FLOAT16 # Set to 10 (FP16)
|
||||
# break
|
||||
# if node.name == "/Squeeze":
|
||||
# node.input[0] = "/ScatterND_14_cast_to_result"
|
||||
|
||||
# Simplify with input/output names preservation
|
||||
onnx_model_sim, check = simplify(onnx_model, skip_fuse_bn=True)
|
||||
if check:
|
||||
print(f" Simplified model inputs: {[inp.name for inp in onnx_model_sim.graph.input]}")
|
||||
onnx.save(onnx_model_sim, join(MODEL_DIR, "body_morpher.onnx"))
|
||||
else:
|
||||
print(" Simplification failed, using original model")
|
||||
onnx.save(onnx_model, join(MODEL_DIR, "body_morpher.onnx"))
|
||||
|
||||
print("✓ Body Morpher exported")
|
||||
|
||||
# =================================================================================
|
||||
# Validation: ONNX Inference
|
||||
# =================================================================================
|
||||
print("\n" + "="*80)
|
||||
print("Validation: Running ONNX inference...")
|
||||
print("="*80)
|
||||
|
||||
class THA4StudentONNXRunner:
|
||||
def __init__(self, model_dir, use_gpu=True):
|
||||
providers = [("CUDAExecutionProvider", {"device_id": 0})] if use_gpu else ['CPUExecutionProvider']
|
||||
|
||||
self.face_morpher_sess = self._create_session(join(model_dir, 'face_morpher.onnx'), providers)
|
||||
self.body_morpher_sess = self._create_session(join(model_dir, 'body_morpher.onnx'), providers)
|
||||
|
||||
# Print actual input names for debugging
|
||||
face_inputs = [inp.name for inp in self.face_morpher_sess.get_inputs()]
|
||||
body_inputs = [inp.name for inp in self.body_morpher_sess.get_inputs()]
|
||||
print(f"✓ All ONNX models loaded")
|
||||
print(f" Face Morpher inputs: {face_inputs}")
|
||||
print(f" Body Morpher inputs: {body_inputs}")
|
||||
|
||||
def _create_session(self, path, providers):
|
||||
try:
|
||||
return ort.InferenceSession(path, providers=providers)
|
||||
except Exception as e:
|
||||
print(f"Failed with {providers}, falling back to CPU for {os.path.basename(path)}")
|
||||
return ort.InferenceSession(path, providers=['CPUExecutionProvider'])
|
||||
|
||||
def run(self, img_path, pose):
|
||||
"""
|
||||
Run inference with student model
|
||||
|
||||
Args:
|
||||
img_path: Path to input image
|
||||
pose: Pose parameters [1, 45]
|
||||
"""
|
||||
img = cv2.imread(img_path, cv2.IMREAD_UNCHANGED)
|
||||
|
||||
# Step 1: Face Morpher (SIREN only needs pose, generates face from scratch)
|
||||
face_res = self.face_morpher_sess.run(None, {
|
||||
'pose': pose[:, :39]
|
||||
})
|
||||
|
||||
# Step 2: Body Morpher (needs original image, face result, and pose)
|
||||
body_res = self.body_morpher_sess.run(None, {
|
||||
'input_image': img,
|
||||
'face_morphed': face_res[0],
|
||||
'pose': pose
|
||||
})
|
||||
|
||||
return body_res[0]
|
||||
|
||||
def save_comparison_images(onnx_img_np, torch_img_np, prefix=""):
|
||||
onnx_hwc = np.transpose(onnx_img_np.squeeze(0), (1, 2, 0))[:, :, :3]
|
||||
torch_hwc = np.transpose(torch_img_np.squeeze(0), (1, 2, 0))[:, :, :3]
|
||||
|
||||
onnx_hwc = np.clip(onnx_hwc, 0, 1)
|
||||
torch_hwc = np.clip(torch_hwc, 0, 1)
|
||||
|
||||
onnx_path = join(MODEL_DIR, f"{prefix}onnx_output.png")
|
||||
torch_path = join(MODEL_DIR, f"{prefix}pytorch_output.png")
|
||||
|
||||
plt.imsave(onnx_path, onnx_hwc)
|
||||
plt.imsave(torch_path, torch_hwc)
|
||||
|
||||
print(f"✅ Saved comparison images: {onnx_path} and {torch_path}")
|
||||
|
||||
# Initialize ONNX runner
|
||||
runner = THA4StudentONNXRunner(MODEL_DIR)
|
||||
pose_zero = np.zeros((1, 45), dtype=np.float32)
|
||||
|
||||
# Test with zero pose
|
||||
print("\nRunning ONNX inference with zero pose...")
|
||||
onnx_output = runner.run(IMAGE_INPUT, pose_zero)
|
||||
|
||||
# Convert reference output to numpy
|
||||
ref_np = poser_torch_res.cpu().numpy() / 2.0 + 0.5
|
||||
ref_np[0, :3, :, :] = numpy_linear_to_srgb(ref_np[0, :3, :, :])
|
||||
|
||||
# Save comparison images
|
||||
save_comparison_images(onnx_output, ref_np, "zero_pose_")
|
||||
|
||||
# Calculate MSE
|
||||
mse = ((onnx_output - ref_np) ** 2).mean()
|
||||
print(f"\n{'='*80}\nMSE (zero pose): {mse:.10f}\n{'='*80}")
|
||||
|
||||
ACCEPTABLE_MSE_THRESHOLD = 1e-2
|
||||
if mse < ACCEPTABLE_MSE_THRESHOLD:
|
||||
print(f"✅ Validation PASSED! MSE is within the acceptable range of {ACCEPTABLE_MSE_THRESHOLD}.")
|
||||
else:
|
||||
print(f"⚠️ Warning: MSE is higher than expected! But this can be normal due to precision differences.")
|
||||
|
||||
# Test with random pose
|
||||
print("\n" + "="*80)
|
||||
print("Testing with random poses...")
|
||||
print("="*80)
|
||||
|
||||
pose_random = np.random.randn(1, 45).astype(np.float32) * 0.3
|
||||
onnx_output_r = runner.run(IMAGE_INPUT, pose_random)
|
||||
|
||||
pose_r_validation = torch.from_numpy(pose_random).to(validation_dtype).to(validation_device)
|
||||
with torch.no_grad():
|
||||
ref_output_r = validation_poser.pose(pt_img_validation, pose_r_validation)
|
||||
|
||||
ref_np_r = ref_output_r.cpu().numpy() / 2.0 + 0.5
|
||||
ref_np_r[0, :3, :, :] = numpy_linear_to_srgb(ref_np_r[0, :3, :, :])
|
||||
|
||||
save_comparison_images(onnx_output_r, ref_np_r, "random_pose_")
|
||||
|
||||
mse_r = ((onnx_output_r - ref_np_r) ** 2).mean()
|
||||
print(f"MSE (random pose): {mse_r:.10f}")
|
||||
|
||||
if mse_r < ACCEPTABLE_MSE_THRESHOLD:
|
||||
print(f"✅ Random pose test PASSED! (Threshold: {ACCEPTABLE_MSE_THRESHOLD})")
|
||||
else:
|
||||
print(f"⚠️ Warning: High MSE with random pose! But this can be normal.")
|
||||
|
||||
print(f"\n{'='*80}")
|
||||
print("Export and validation completed!")
|
||||
print(f"Models saved to: {MODEL_DIR}")
|
||||
print("Generated files:")
|
||||
print(f" - face_morpher.onnx")
|
||||
print(f" - body_morpher.onnx")
|
||||
print(f"{'='*80}")
|
||||
Whitespace-only changes.
Whitespace-only changes.
@@ -0,0 +1,417 @@
|
||||
import os
|
||||
import socket
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
from typing import Optional
|
||||
|
||||
import PIL.Image
|
||||
|
||||
from tha4.shion.base.image_util import torch_linear_to_srgb
|
||||
from tha4.image_util import convert_linear_to_srgb
|
||||
from tha4.mocap.ifacialmocap_pose_converter_25 import create_ifacialmocap_pose_converter
|
||||
from tha4.app.full_manual_poser import resize_PIL_image
|
||||
from tha4.charmodel.character_model import CharacterModel
|
||||
|
||||
sys.path.append(os.getcwd())
|
||||
|
||||
from tha4.mocap.ifacialmocap_pose import create_default_ifacialmocap_pose
|
||||
from tha4.mocap.ifacialmocap_v2 import IFACIALMOCAP_PORT, IFACIALMOCAP_START_STRING, parse_ifacialmocap_v2_pose
|
||||
|
||||
import torch
|
||||
import wx
|
||||
|
||||
from tha4.mocap.ifacialmocap_constants import *
|
||||
from tha4.mocap.ifacialmocap_pose_converter import IFacialMocapPoseConverter
|
||||
|
||||
|
||||
class FpsStatistics:
|
||||
def __init__(self):
|
||||
self.count = 100
|
||||
self.fps = []
|
||||
|
||||
def add_fps(self, fps):
|
||||
self.fps.append(fps)
|
||||
while len(self.fps) > self.count:
|
||||
del self.fps[0]
|
||||
|
||||
def get_average_fps(self):
|
||||
if len(self.fps) == 0:
|
||||
return 0.0
|
||||
else:
|
||||
return sum(self.fps) / len(self.fps)
|
||||
|
||||
|
||||
class MainFrame(wx.Frame):
|
||||
IMAGE_SIZE = 512
|
||||
|
||||
def __init__(self, pose_converter: IFacialMocapPoseConverter, device: torch.device):
|
||||
super().__init__(None, wx.ID_ANY, "iFacialMocap Puppeteer (Fuji)")
|
||||
self.poser = None
|
||||
self.pose_converter = pose_converter
|
||||
self.device = device
|
||||
|
||||
self.ifacialmocap_pose = create_default_ifacialmocap_pose()
|
||||
self.source_image_bitmap = wx.Bitmap(MainFrame.IMAGE_SIZE, MainFrame.IMAGE_SIZE)
|
||||
self.result_image_bitmap = wx.Bitmap(MainFrame.IMAGE_SIZE, MainFrame.IMAGE_SIZE)
|
||||
self.wx_source_image = None
|
||||
self.torch_source_image = None
|
||||
self.last_pose = None
|
||||
self.fps_statistics = FpsStatistics()
|
||||
self.last_update_time = None
|
||||
|
||||
self.create_receiving_socket()
|
||||
self.create_ui()
|
||||
self.create_timers()
|
||||
self.Bind(wx.EVT_CLOSE, self.on_close)
|
||||
|
||||
self.update_source_image_bitmap()
|
||||
self.update_result_image_bitmap()
|
||||
|
||||
def create_receiving_socket(self):
|
||||
self.receiving_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
|
||||
self.receiving_socket.bind(("", IFACIALMOCAP_PORT))
|
||||
self.receiving_socket.setblocking(False)
|
||||
|
||||
def create_timers(self):
|
||||
self.capture_timer = wx.Timer(self, wx.ID_ANY)
|
||||
self.Bind(wx.EVT_TIMER, self.update_capture_panel, id=self.capture_timer.GetId())
|
||||
self.animation_timer = wx.Timer(self, wx.ID_ANY)
|
||||
self.Bind(wx.EVT_TIMER, self.update_result_image_bitmap, id=self.animation_timer.GetId())
|
||||
|
||||
def on_close(self, event: wx.Event):
|
||||
# Stop the timers
|
||||
self.animation_timer.Stop()
|
||||
self.capture_timer.Stop()
|
||||
|
||||
# Close receiving socket
|
||||
self.receiving_socket.close()
|
||||
|
||||
# Destroy the windows
|
||||
self.Destroy()
|
||||
event.Skip()
|
||||
|
||||
def on_start_capture(self, event: wx.Event):
|
||||
capture_device_ip_address = self.capture_device_ip_text_ctrl.GetValue()
|
||||
out_socket = None
|
||||
try:
|
||||
address = (capture_device_ip_address, IFACIALMOCAP_PORT)
|
||||
out_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
|
||||
out_socket.sendto(IFACIALMOCAP_START_STRING, address)
|
||||
except Exception as e:
|
||||
message_dialog = wx.MessageDialog(self, str(e), "Error!", wx.OK)
|
||||
message_dialog.ShowModal()
|
||||
message_dialog.Destroy()
|
||||
finally:
|
||||
if out_socket is not None:
|
||||
out_socket.close()
|
||||
|
||||
def read_ifacialmocap_pose(self):
|
||||
if not self.animation_timer.IsRunning():
|
||||
return self.ifacialmocap_pose
|
||||
socket_bytes = None
|
||||
while True:
|
||||
try:
|
||||
socket_bytes = self.receiving_socket.recv(8192)
|
||||
except socket.error as e:
|
||||
break
|
||||
if socket_bytes is not None:
|
||||
socket_string = socket_bytes.decode("utf-8")
|
||||
self.ifacialmocap_pose = parse_ifacialmocap_v2_pose(socket_string)
|
||||
return self.ifacialmocap_pose
|
||||
|
||||
def on_erase_background(self, event: wx.Event):
|
||||
pass
|
||||
|
||||
def create_animation_panel(self, parent):
|
||||
self.animation_panel = wx.Panel(parent, style=wx.RAISED_BORDER)
|
||||
self.animation_panel_sizer = wx.BoxSizer(wx.HORIZONTAL)
|
||||
self.animation_panel.SetSizer(self.animation_panel_sizer)
|
||||
self.animation_panel.SetAutoLayout(1)
|
||||
|
||||
image_size = MainFrame.IMAGE_SIZE
|
||||
|
||||
if True:
|
||||
self.input_panel = wx.Panel(self.animation_panel, size=(image_size, image_size + 128),
|
||||
style=wx.SIMPLE_BORDER)
|
||||
self.input_panel_sizer = wx.BoxSizer(wx.VERTICAL)
|
||||
self.input_panel.SetSizer(self.input_panel_sizer)
|
||||
self.input_panel.SetAutoLayout(1)
|
||||
self.animation_panel_sizer.Add(self.input_panel, 0, wx.FIXED_MINSIZE)
|
||||
|
||||
self.source_image_panel = wx.Panel(self.input_panel, size=(image_size, image_size), style=wx.SIMPLE_BORDER)
|
||||
self.source_image_panel.Bind(wx.EVT_PAINT, self.paint_source_image_panel)
|
||||
self.source_image_panel.Bind(wx.EVT_ERASE_BACKGROUND, self.on_erase_background)
|
||||
self.input_panel_sizer.Add(self.source_image_panel, 0, wx.FIXED_MINSIZE)
|
||||
|
||||
self.load_model_button = wx.Button(self.input_panel, wx.ID_ANY, "Load Model")
|
||||
self.input_panel_sizer.Add(self.load_model_button, 1, wx.EXPAND)
|
||||
self.load_model_button.Bind(wx.EVT_BUTTON, self.load_model)
|
||||
|
||||
self.input_panel_sizer.Fit(self.input_panel)
|
||||
|
||||
if True:
|
||||
self.pose_converter.init_pose_converter_panel(self.animation_panel)
|
||||
|
||||
if True:
|
||||
self.animation_left_panel = wx.Panel(self.animation_panel, style=wx.SIMPLE_BORDER)
|
||||
self.animation_left_panel_sizer = wx.BoxSizer(wx.VERTICAL)
|
||||
self.animation_left_panel.SetSizer(self.animation_left_panel_sizer)
|
||||
self.animation_left_panel.SetAutoLayout(1)
|
||||
self.animation_panel_sizer.Add(self.animation_left_panel, 0, wx.EXPAND)
|
||||
|
||||
self.result_image_panel = wx.Panel(self.animation_left_panel, size=(image_size, image_size),
|
||||
style=wx.SIMPLE_BORDER)
|
||||
self.result_image_panel.Bind(wx.EVT_PAINT, self.paint_result_image_panel)
|
||||
self.result_image_panel.Bind(wx.EVT_ERASE_BACKGROUND, self.on_erase_background)
|
||||
self.animation_left_panel_sizer.Add(self.result_image_panel, 0, wx.FIXED_MINSIZE)
|
||||
|
||||
separator = wx.StaticLine(self.animation_left_panel, -1, size=(256, 5))
|
||||
self.animation_left_panel_sizer.Add(separator, 0, wx.EXPAND)
|
||||
|
||||
background_text = wx.StaticText(self.animation_left_panel, label="--- Background ---",
|
||||
style=wx.ALIGN_CENTER)
|
||||
self.animation_left_panel_sizer.Add(background_text, 0, wx.EXPAND)
|
||||
|
||||
self.output_background_choice = wx.Choice(
|
||||
self.animation_left_panel,
|
||||
choices=[
|
||||
"TRANSPARENT",
|
||||
"GREEN",
|
||||
"BLUE",
|
||||
"BLACK",
|
||||
"WHITE"
|
||||
])
|
||||
self.output_background_choice.SetSelection(0)
|
||||
self.animation_left_panel_sizer.Add(self.output_background_choice, 0, wx.EXPAND)
|
||||
|
||||
separator = wx.StaticLine(self.animation_left_panel, -1, size=(256, 5))
|
||||
self.animation_left_panel_sizer.Add(separator, 0, wx.EXPAND)
|
||||
|
||||
self.fps_text = wx.StaticText(self.animation_left_panel, label="")
|
||||
self.animation_left_panel_sizer.Add(self.fps_text, wx.SizerFlags().Border())
|
||||
|
||||
self.animation_left_panel_sizer.Fit(self.animation_left_panel)
|
||||
|
||||
self.animation_panel_sizer.Fit(self.animation_panel)
|
||||
|
||||
def create_ui(self):
|
||||
self.main_sizer = wx.BoxSizer(wx.VERTICAL)
|
||||
self.SetSizer(self.main_sizer)
|
||||
self.SetAutoLayout(1)
|
||||
|
||||
self.capture_pose_lock = threading.Lock()
|
||||
|
||||
self.create_connection_panel(self)
|
||||
self.main_sizer.Add(self.connection_panel, wx.SizerFlags(0).Expand().Border(wx.ALL, 5))
|
||||
|
||||
self.create_animation_panel(self)
|
||||
self.main_sizer.Add(self.animation_panel, wx.SizerFlags(0).Expand().Border(wx.ALL, 5))
|
||||
|
||||
self.create_capture_panel(self)
|
||||
self.main_sizer.Add(self.capture_panel, wx.SizerFlags(0).Expand().Border(wx.ALL, 5))
|
||||
|
||||
self.main_sizer.Fit(self)
|
||||
|
||||
def create_connection_panel(self, parent):
|
||||
self.connection_panel = wx.Panel(parent, style=wx.RAISED_BORDER)
|
||||
self.connection_panel_sizer = wx.BoxSizer(wx.HORIZONTAL)
|
||||
self.connection_panel.SetSizer(self.connection_panel_sizer)
|
||||
self.connection_panel.SetAutoLayout(1)
|
||||
|
||||
capture_device_ip_text = wx.StaticText(self.connection_panel, label="Capture Device IP:", style=wx.ALIGN_RIGHT)
|
||||
self.connection_panel_sizer.Add(capture_device_ip_text, wx.SizerFlags(0).FixedMinSize().Border(wx.ALL, 3))
|
||||
|
||||
self.capture_device_ip_text_ctrl = wx.TextCtrl(self.connection_panel, value="192.168.0.1")
|
||||
self.connection_panel_sizer.Add(self.capture_device_ip_text_ctrl, wx.SizerFlags(1).Expand().Border(wx.ALL, 3))
|
||||
|
||||
self.start_capture_button = wx.Button(self.connection_panel, label="START CAPTURE!")
|
||||
self.connection_panel_sizer.Add(self.start_capture_button, wx.SizerFlags(0).FixedMinSize().Border(wx.ALL, 3))
|
||||
self.start_capture_button.Bind(wx.EVT_BUTTON, self.on_start_capture)
|
||||
|
||||
def create_capture_panel(self, parent):
|
||||
self.capture_panel = wx.Panel(parent, style=wx.RAISED_BORDER)
|
||||
self.capture_panel_sizer = wx.FlexGridSizer(cols=5)
|
||||
for i in range(5):
|
||||
self.capture_panel_sizer.AddGrowableCol(i)
|
||||
self.capture_panel.SetSizer(self.capture_panel_sizer)
|
||||
self.capture_panel.SetAutoLayout(1)
|
||||
|
||||
self.rotation_labels = {}
|
||||
self.rotation_value_labels = {}
|
||||
rotation_column_0 = self.create_rotation_column(self.capture_panel, RIGHT_EYE_BONE_ROTATIONS)
|
||||
self.capture_panel_sizer.Add(rotation_column_0, wx.SizerFlags(0).Expand().Border(wx.ALL, 3))
|
||||
rotation_column_1 = self.create_rotation_column(self.capture_panel, LEFT_EYE_BONE_ROTATIONS)
|
||||
self.capture_panel_sizer.Add(rotation_column_1, wx.SizerFlags(0).Expand().Border(wx.ALL, 3))
|
||||
rotation_column_2 = self.create_rotation_column(self.capture_panel, HEAD_BONE_ROTATIONS)
|
||||
self.capture_panel_sizer.Add(rotation_column_2, wx.SizerFlags(0).Expand().Border(wx.ALL, 3))
|
||||
|
||||
def create_rotation_column(self, parent, rotation_names):
|
||||
column_panel = wx.Panel(parent, style=wx.SIMPLE_BORDER)
|
||||
column_panel_sizer = wx.FlexGridSizer(cols=2)
|
||||
column_panel_sizer.AddGrowableCol(1)
|
||||
column_panel.SetSizer(column_panel_sizer)
|
||||
column_panel.SetAutoLayout(1)
|
||||
|
||||
for rotation_name in rotation_names:
|
||||
self.rotation_labels[rotation_name] = wx.StaticText(
|
||||
column_panel, label=rotation_name, style=wx.ALIGN_RIGHT)
|
||||
column_panel_sizer.Add(self.rotation_labels[rotation_name],
|
||||
wx.SizerFlags(1).Expand().Border(wx.ALL, 3))
|
||||
|
||||
self.rotation_value_labels[rotation_name] = wx.TextCtrl(
|
||||
column_panel, style=wx.TE_RIGHT)
|
||||
self.rotation_value_labels[rotation_name].SetValue("0.00")
|
||||
self.rotation_value_labels[rotation_name].Disable()
|
||||
column_panel_sizer.Add(self.rotation_value_labels[rotation_name],
|
||||
wx.SizerFlags(1).Expand().Border(wx.ALL, 3))
|
||||
|
||||
column_panel.GetSizer().Fit(column_panel)
|
||||
return column_panel
|
||||
|
||||
def paint_capture_panel(self, event: wx.Event):
|
||||
self.update_capture_panel(event)
|
||||
|
||||
def update_capture_panel(self, event: wx.Event):
|
||||
data = self.ifacialmocap_pose
|
||||
for rotation_name in ROTATION_NAMES:
|
||||
value = data[rotation_name]
|
||||
self.rotation_value_labels[rotation_name].SetValue("%0.2f" % value)
|
||||
|
||||
@staticmethod
|
||||
def convert_to_100(x):
|
||||
return int(max(0.0, min(1.0, x)) * 100)
|
||||
|
||||
def paint_source_image_panel(self, event: wx.Event):
|
||||
wx.BufferedPaintDC(self.source_image_panel, self.source_image_bitmap)
|
||||
|
||||
def update_source_image_bitmap(self):
|
||||
dc = wx.MemoryDC()
|
||||
dc.SelectObject(self.source_image_bitmap)
|
||||
if self.wx_source_image is None:
|
||||
self.draw_nothing_yet_string(dc)
|
||||
else:
|
||||
dc.Clear()
|
||||
dc.DrawBitmap(self.wx_source_image, 0, 0, True)
|
||||
del dc
|
||||
|
||||
def draw_nothing_yet_string(self, dc):
|
||||
dc.Clear()
|
||||
font = wx.Font(wx.FontInfo(14).Family(wx.FONTFAMILY_SWISS))
|
||||
dc.SetFont(font)
|
||||
w, h = dc.GetTextExtent("Nothing yet!")
|
||||
dc.DrawText("Nothing yet!", (MainFrame.IMAGE_SIZE - w) // 2, (MainFrame.IMAGE_SIZE - h) // 2)
|
||||
|
||||
def paint_result_image_panel(self, event: wx.Event):
|
||||
wx.BufferedPaintDC(self.result_image_panel, self.result_image_bitmap)
|
||||
|
||||
def update_result_image_bitmap(self, event: Optional[wx.Event] = None):
|
||||
ifacialmocap_pose = self.read_ifacialmocap_pose()
|
||||
current_pose = self.pose_converter.convert(ifacialmocap_pose)
|
||||
if self.last_pose is not None and self.last_pose == current_pose:
|
||||
return
|
||||
self.last_pose = current_pose
|
||||
|
||||
if self.torch_source_image is None or self.poser is None:
|
||||
dc = wx.MemoryDC()
|
||||
dc.SelectObject(self.result_image_bitmap)
|
||||
self.draw_nothing_yet_string(dc)
|
||||
del dc
|
||||
return
|
||||
|
||||
pose = torch.tensor(current_pose, device=self.device, dtype=self.poser.get_dtype())
|
||||
|
||||
with torch.no_grad():
|
||||
output_image = self.poser.pose(self.torch_source_image, pose)[0].float()
|
||||
output_image = torch.clip((output_image + 1.0) / 2.0, 0.0, 1.0)
|
||||
output_image = convert_linear_to_srgb(output_image)
|
||||
|
||||
background_choice = self.output_background_choice.GetSelection()
|
||||
if background_choice == 0:
|
||||
pass
|
||||
else:
|
||||
background = torch.zeros(4, output_image.shape[1], output_image.shape[2], device=self.device)
|
||||
background[3, :, :] = 1.0
|
||||
if background_choice == 1:
|
||||
background[1, :, :] = 1.0
|
||||
output_image = self.blend_with_background(output_image, background)
|
||||
elif background_choice == 2:
|
||||
background[2, :, :] = 1.0
|
||||
output_image = self.blend_with_background(output_image, background)
|
||||
elif background_choice == 3:
|
||||
output_image = self.blend_with_background(output_image, background)
|
||||
else:
|
||||
background[0:3, :, :] = 1.0
|
||||
output_image = self.blend_with_background(output_image, background)
|
||||
|
||||
c, h, w = output_image.shape
|
||||
output_image = 255.0 * torch.transpose(output_image.reshape(c, h * w), 0, 1).reshape(h, w, c)
|
||||
output_image = output_image.byte()
|
||||
|
||||
numpy_image = output_image.detach().cpu().numpy()
|
||||
wx_image = wx.ImageFromBuffer(numpy_image.shape[0],
|
||||
numpy_image.shape[1],
|
||||
numpy_image[:, :, 0:3].tobytes(),
|
||||
numpy_image[:, :, 3].tobytes())
|
||||
wx_bitmap = wx_image.ConvertToBitmap()
|
||||
|
||||
dc = wx.MemoryDC()
|
||||
dc.SelectObject(self.result_image_bitmap)
|
||||
dc.Clear()
|
||||
dc.DrawBitmap(wx_bitmap,
|
||||
(MainFrame.IMAGE_SIZE - numpy_image.shape[0]) // 2,
|
||||
(MainFrame.IMAGE_SIZE - numpy_image.shape[1]) // 2, True)
|
||||
del dc
|
||||
|
||||
time_now = time.time_ns()
|
||||
if self.last_update_time is not None:
|
||||
elapsed_time = time_now - self.last_update_time
|
||||
fps = 1.0 / (elapsed_time / 10 ** 9)
|
||||
if self.torch_source_image is not None:
|
||||
self.fps_statistics.add_fps(fps)
|
||||
self.fps_text.SetLabelText("FPS = %0.2f" % self.fps_statistics.get_average_fps())
|
||||
self.last_update_time = time_now
|
||||
|
||||
self.Refresh()
|
||||
|
||||
def blend_with_background(self, numpy_image, background):
|
||||
alpha = numpy_image[3:4, :, :]
|
||||
color = numpy_image[0:3, :, :]
|
||||
new_color = color * alpha + (1.0 - alpha) * background[0:3, :, :]
|
||||
return torch.cat([new_color, background[3:4, :, :]], dim=0)
|
||||
|
||||
def load_model(self, event: wx.Event):
|
||||
dir_name = "data/character_models"
|
||||
file_dialog = wx.FileDialog(self, "Choose a model", dir_name, "", "*.yaml", wx.FD_OPEN)
|
||||
if file_dialog.ShowModal() == wx.ID_OK:
|
||||
character_model_json_file_name = os.path.join(file_dialog.GetDirectory(), file_dialog.GetFilename())
|
||||
try:
|
||||
self.character_model = CharacterModel.load(character_model_json_file_name)
|
||||
self.torch_source_image = self.character_model.get_character_image(self.device)
|
||||
pil_image = resize_PIL_image(
|
||||
PIL.Image.open(self.character_model.character_image_file_name),
|
||||
(MainFrame.IMAGE_SIZE, MainFrame.IMAGE_SIZE))
|
||||
w, h = pil_image.size
|
||||
self.wx_source_image = wx.Bitmap.FromBufferRGBA(w, h, pil_image.convert("RGBA").tobytes())
|
||||
self.update_source_image_bitmap()
|
||||
self.poser = self.character_model.get_poser(self.device)
|
||||
except Exception:
|
||||
message_dialog = wx.MessageDialog(
|
||||
self, "Could not load character model " + character_model_json_file_name, "Poser", wx.OK)
|
||||
message_dialog.ShowModal()
|
||||
message_dialog.Destroy()
|
||||
file_dialog.Destroy()
|
||||
self.Refresh()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
device = torch.device('cuda:0')
|
||||
|
||||
pose_converter = create_ifacialmocap_pose_converter()
|
||||
|
||||
app = wx.App()
|
||||
main_frame = MainFrame(pose_converter, device)
|
||||
main_frame.Show(True)
|
||||
main_frame.capture_timer.Start(10)
|
||||
main_frame.animation_timer.Start(10)
|
||||
app.MainLoop()
|
||||
@@ -0,0 +1,438 @@
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
from typing import List
|
||||
|
||||
from tha4.charmodel.character_model import CharacterModel
|
||||
from tha4.image_util import resize_PIL_image, convert_output_image_from_torch_to_numpy
|
||||
from tha4.poser.modes.mode_14 import get_pose_parameters
|
||||
|
||||
sys.path.append(os.getcwd())
|
||||
|
||||
import PIL.Image
|
||||
import torch
|
||||
import wx
|
||||
|
||||
from tha4.poser.poser import PoseParameterCategory, PoseParameterGroup
|
||||
|
||||
|
||||
class MorphCategoryControlPanel(wx.Panel):
|
||||
def __init__(self,
|
||||
parent,
|
||||
title: str,
|
||||
pose_param_category: PoseParameterCategory,
|
||||
param_groups: List[PoseParameterGroup]):
|
||||
super().__init__(parent, style=wx.SIMPLE_BORDER)
|
||||
self.pose_param_category = pose_param_category
|
||||
self.sizer = wx.BoxSizer(wx.VERTICAL)
|
||||
self.SetSizer(self.sizer)
|
||||
self.SetAutoLayout(1)
|
||||
|
||||
title_text = wx.StaticText(self, label=title, style=wx.ALIGN_CENTER)
|
||||
self.sizer.Add(title_text, 0, wx.EXPAND)
|
||||
|
||||
self.param_groups = [group for group in param_groups if group.get_category() == pose_param_category]
|
||||
self.choice = wx.Choice(self, choices=[group.get_group_name() for group in self.param_groups])
|
||||
if len(self.param_groups) > 0:
|
||||
self.choice.SetSelection(0)
|
||||
self.choice.Bind(wx.EVT_CHOICE, self.on_choice_updated)
|
||||
self.sizer.Add(self.choice, 0, wx.EXPAND)
|
||||
|
||||
self.left_slider = wx.Slider(self, minValue=-1000, maxValue=1000, value=-1000, style=wx.HORIZONTAL)
|
||||
self.sizer.Add(self.left_slider, 0, wx.EXPAND)
|
||||
|
||||
self.right_slider = wx.Slider(self, minValue=-1000, maxValue=1000, value=-1000, style=wx.HORIZONTAL)
|
||||
self.sizer.Add(self.right_slider, 0, wx.EXPAND)
|
||||
|
||||
self.checkbox = wx.CheckBox(self, label="Show")
|
||||
self.checkbox.SetValue(True)
|
||||
self.sizer.Add(self.checkbox, 0, wx.SHAPED | wx.ALIGN_CENTER)
|
||||
|
||||
self.update_ui()
|
||||
|
||||
self.sizer.Fit(self)
|
||||
|
||||
def update_ui(self):
|
||||
param_group = self.param_groups[self.choice.GetSelection()]
|
||||
if param_group.is_discrete():
|
||||
self.left_slider.Enable(False)
|
||||
self.right_slider.Enable(False)
|
||||
self.checkbox.Enable(True)
|
||||
elif param_group.get_arity() == 1:
|
||||
self.left_slider.Enable(True)
|
||||
self.right_slider.Enable(False)
|
||||
self.checkbox.Enable(False)
|
||||
else:
|
||||
self.left_slider.Enable(True)
|
||||
self.right_slider.Enable(True)
|
||||
self.checkbox.Enable(False)
|
||||
|
||||
def on_choice_updated(self, event: wx.Event):
|
||||
param_group = self.param_groups[self.choice.GetSelection()]
|
||||
if param_group.is_discrete():
|
||||
self.checkbox.SetValue(True)
|
||||
self.update_ui()
|
||||
|
||||
def set_param_value(self, pose: List[float]):
|
||||
if len(self.param_groups) == 0:
|
||||
return
|
||||
selected_morph_index = self.choice.GetSelection()
|
||||
param_group = self.param_groups[selected_morph_index]
|
||||
param_index = param_group.get_parameter_index()
|
||||
if param_group.is_discrete():
|
||||
if self.checkbox.GetValue():
|
||||
for i in range(param_group.get_arity()):
|
||||
pose[param_index + i] = 1.0
|
||||
else:
|
||||
param_range = param_group.get_range()
|
||||
alpha = (self.left_slider.GetValue() + 1000) / 2000.0
|
||||
pose[param_index] = param_range[0] + (param_range[1] - param_range[0]) * alpha
|
||||
if param_group.get_arity() == 2:
|
||||
alpha = (self.right_slider.GetValue() + 1000) / 2000.0
|
||||
pose[param_index + 1] = param_range[0] + (param_range[1] - param_range[0]) * alpha
|
||||
|
||||
|
||||
class SimpleParamGroupsControlPanel(wx.Panel):
|
||||
def __init__(self, parent,
|
||||
pose_param_category: PoseParameterCategory,
|
||||
param_groups: List[PoseParameterGroup]):
|
||||
super().__init__(parent, style=wx.SIMPLE_BORDER)
|
||||
self.sizer = wx.BoxSizer(wx.VERTICAL)
|
||||
self.SetSizer(self.sizer)
|
||||
self.SetAutoLayout(1)
|
||||
|
||||
self.param_groups = [group for group in param_groups if group.get_category() == pose_param_category]
|
||||
for param_group in self.param_groups:
|
||||
assert not param_group.is_discrete()
|
||||
assert param_group.get_arity() == 1
|
||||
|
||||
self.sliders = []
|
||||
for param_group in self.param_groups:
|
||||
static_text = wx.StaticText(
|
||||
self,
|
||||
label=" ------------ %s ------------ " % param_group.get_group_name(), style=wx.ALIGN_CENTER)
|
||||
self.sizer.Add(static_text, 0, wx.EXPAND)
|
||||
range = param_group.get_range()
|
||||
min_value = int(range[0] * 1000)
|
||||
max_value = int(range[1] * 1000)
|
||||
slider = wx.Slider(self, minValue=min_value, maxValue=max_value, value=0, style=wx.HORIZONTAL)
|
||||
self.sizer.Add(slider, 0, wx.EXPAND)
|
||||
self.sliders.append(slider)
|
||||
|
||||
self.sizer.Fit(self)
|
||||
|
||||
def set_param_value(self, pose: List[float]):
|
||||
if len(self.param_groups) == 0:
|
||||
return
|
||||
for param_group_index in range(len(self.param_groups)):
|
||||
param_group = self.param_groups[param_group_index]
|
||||
slider = self.sliders[param_group_index]
|
||||
param_range = param_group.get_range()
|
||||
param_index = param_group.get_parameter_index()
|
||||
alpha = (slider.GetValue() - slider.GetMin()) * 1.0 / (slider.GetMax() - slider.GetMin())
|
||||
pose[param_index] = param_range[0] + (param_range[1] - param_range[0]) * alpha
|
||||
|
||||
|
||||
class MainFrame(wx.Frame):
|
||||
IMAGE_SIZE = 512
|
||||
OUTPUT_LENGTH = 6
|
||||
NUM_PARAMETERS = 45
|
||||
|
||||
def __init__(self, device: torch.device):
|
||||
super().__init__(None, wx.ID_ANY, "Poser")
|
||||
self.poser = None
|
||||
self.device = device
|
||||
|
||||
self.wx_source_image = None
|
||||
self.torch_source_image = None
|
||||
|
||||
self.main_sizer = wx.BoxSizer(wx.HORIZONTAL)
|
||||
self.SetSizer(self.main_sizer)
|
||||
self.SetAutoLayout(1)
|
||||
self.init_left_panel()
|
||||
self.init_control_panel()
|
||||
self.init_right_panel()
|
||||
self.main_sizer.Fit(self)
|
||||
|
||||
self.timer = wx.Timer(self, wx.ID_ANY)
|
||||
self.Bind(wx.EVT_TIMER, self.update_images, self.timer)
|
||||
|
||||
save_image_id = wx.NewIdRef()
|
||||
self.Bind(wx.EVT_MENU, self.on_save_image, id=save_image_id)
|
||||
accelerator_table = wx.AcceleratorTable([
|
||||
(wx.ACCEL_CTRL, ord('S'), save_image_id)
|
||||
])
|
||||
self.SetAcceleratorTable(accelerator_table)
|
||||
|
||||
self.last_pose = None
|
||||
self.last_output_index = self.output_index_choice.GetSelection()
|
||||
self.last_output_numpy_image = None
|
||||
|
||||
self.wx_source_image = None
|
||||
self.torch_source_image = None
|
||||
self.source_image_bitmap = wx.Bitmap(MainFrame.IMAGE_SIZE, MainFrame.IMAGE_SIZE)
|
||||
self.result_image_bitmap = wx.Bitmap(MainFrame.IMAGE_SIZE, MainFrame.IMAGE_SIZE)
|
||||
self.source_image_dirty = True
|
||||
|
||||
def init_left_panel(self):
|
||||
self.control_panel = wx.Panel(self, style=wx.SIMPLE_BORDER, size=(MainFrame.IMAGE_SIZE, -1))
|
||||
self.left_panel = wx.Panel(self, style=wx.SIMPLE_BORDER)
|
||||
left_panel_sizer = wx.BoxSizer(wx.VERTICAL)
|
||||
self.left_panel.SetSizer(left_panel_sizer)
|
||||
self.left_panel.SetAutoLayout(1)
|
||||
|
||||
self.source_image_panel = wx.Panel(self.left_panel, size=(MainFrame.IMAGE_SIZE, MainFrame.IMAGE_SIZE),
|
||||
style=wx.SIMPLE_BORDER)
|
||||
self.source_image_panel.Bind(wx.EVT_PAINT, self.paint_source_image_panel)
|
||||
self.source_image_panel.Bind(wx.EVT_ERASE_BACKGROUND, self.on_erase_background)
|
||||
left_panel_sizer.Add(self.source_image_panel, 0, wx.FIXED_MINSIZE)
|
||||
|
||||
self.load_model_button = wx.Button(self.left_panel, wx.ID_ANY, "\nLoad Model\n\n")
|
||||
left_panel_sizer.Add(self.load_model_button, 1, wx.EXPAND)
|
||||
self.load_model_button.Bind(wx.EVT_BUTTON, self.load_model)
|
||||
|
||||
left_panel_sizer.Fit(self.left_panel)
|
||||
self.main_sizer.Add(self.left_panel, 0, wx.FIXED_MINSIZE)
|
||||
|
||||
def on_erase_background(self, event: wx.Event):
|
||||
pass
|
||||
|
||||
def init_control_panel(self):
|
||||
self.control_panel_sizer = wx.BoxSizer(wx.VERTICAL)
|
||||
self.control_panel.SetSizer(self.control_panel_sizer)
|
||||
self.control_panel.SetMinSize(wx.Size(256, 1))
|
||||
|
||||
morph_categories = [
|
||||
PoseParameterCategory.EYEBROW,
|
||||
PoseParameterCategory.EYE,
|
||||
PoseParameterCategory.MOUTH,
|
||||
PoseParameterCategory.IRIS_MORPH
|
||||
]
|
||||
morph_category_titles = {
|
||||
PoseParameterCategory.EYEBROW: " ------------ Eyebrow ------------ ",
|
||||
PoseParameterCategory.EYE: " ------------ Eye ------------ ",
|
||||
PoseParameterCategory.MOUTH: " ------------ Mouth ------------ ",
|
||||
PoseParameterCategory.IRIS_MORPH: " ------------ Iris morphs ------------ ",
|
||||
}
|
||||
self.morph_control_panels = {}
|
||||
param_groups = get_pose_parameters().get_pose_parameter_groups()
|
||||
for category in morph_categories:
|
||||
filtered_param_groups = [group for group in param_groups if group.get_category() == category]
|
||||
if len(filtered_param_groups) == 0:
|
||||
continue
|
||||
control_panel = MorphCategoryControlPanel(
|
||||
self.control_panel,
|
||||
morph_category_titles[category],
|
||||
category,
|
||||
param_groups)
|
||||
self.morph_control_panels[category] = control_panel
|
||||
self.control_panel_sizer.Add(control_panel, 0, wx.EXPAND)
|
||||
|
||||
self.non_morph_control_panels = {}
|
||||
non_morph_categories = [
|
||||
PoseParameterCategory.IRIS_ROTATION,
|
||||
PoseParameterCategory.FACE_ROTATION,
|
||||
PoseParameterCategory.BODY_ROTATION,
|
||||
PoseParameterCategory.BREATHING
|
||||
]
|
||||
for category in non_morph_categories:
|
||||
filtered_param_groups = [group for group in param_groups if group.get_category() == category]
|
||||
if len(filtered_param_groups) == 0:
|
||||
continue
|
||||
control_panel = SimpleParamGroupsControlPanel(
|
||||
self.control_panel,
|
||||
category,
|
||||
param_groups)
|
||||
self.non_morph_control_panels[category] = control_panel
|
||||
self.control_panel_sizer.Add(control_panel, 0, wx.EXPAND)
|
||||
|
||||
self.control_panel_sizer.Fit(self.control_panel)
|
||||
self.main_sizer.Add(self.control_panel, 1, wx.FIXED_MINSIZE)
|
||||
|
||||
def init_right_panel(self):
|
||||
self.right_panel = wx.Panel(self, style=wx.SIMPLE_BORDER)
|
||||
right_panel_sizer = wx.BoxSizer(wx.VERTICAL)
|
||||
self.right_panel.SetSizer(right_panel_sizer)
|
||||
self.right_panel.SetAutoLayout(1)
|
||||
|
||||
self.result_image_panel = wx.Panel(self.right_panel,
|
||||
size=(MainFrame.IMAGE_SIZE, MainFrame.IMAGE_SIZE),
|
||||
style=wx.SIMPLE_BORDER)
|
||||
self.result_image_panel.Bind(wx.EVT_PAINT, self.paint_result_image_panel)
|
||||
self.result_image_panel.Bind(wx.EVT_ERASE_BACKGROUND, self.on_erase_background)
|
||||
self.output_index_choice = wx.Choice(
|
||||
self.right_panel,
|
||||
choices=[str(i) for i in range(MainFrame.OUTPUT_LENGTH)])
|
||||
self.output_index_choice.SetSelection(0)
|
||||
right_panel_sizer.Add(self.result_image_panel, 0, wx.FIXED_MINSIZE)
|
||||
right_panel_sizer.Add(self.output_index_choice, 0, wx.EXPAND)
|
||||
|
||||
self.save_image_button = wx.Button(self.right_panel, wx.ID_ANY, "\nSave Image\n\n")
|
||||
right_panel_sizer.Add(self.save_image_button, 1, wx.EXPAND)
|
||||
self.save_image_button.Bind(wx.EVT_BUTTON, self.on_save_image)
|
||||
|
||||
right_panel_sizer.Fit(self.right_panel)
|
||||
self.main_sizer.Add(self.right_panel, 0, wx.FIXED_MINSIZE)
|
||||
|
||||
def create_param_category_choice(self, param_category: PoseParameterCategory):
|
||||
params = []
|
||||
for param_group in self.poser.get_pose_parameter_groups():
|
||||
if param_group.get_category() == param_category:
|
||||
params.append(param_group.get_group_name())
|
||||
choice = wx.Choice(self.control_panel, choices=params)
|
||||
if len(params) > 0:
|
||||
choice.SetSelection(0)
|
||||
return choice
|
||||
|
||||
def load_model(self, event: wx.Event):
|
||||
dir_name = "data/character_models"
|
||||
file_dialog = wx.FileDialog(self, "Choose a model", dir_name, "", "*.yaml", wx.FD_OPEN)
|
||||
if file_dialog.ShowModal() == wx.ID_OK:
|
||||
character_model_file_name = os.path.join(file_dialog.GetDirectory(), file_dialog.GetFilename())
|
||||
try:
|
||||
self.character_model = CharacterModel.load(character_model_file_name)
|
||||
self.torch_source_image = self.character_model.get_character_image(self.device)
|
||||
pil_image = resize_PIL_image(
|
||||
PIL.Image.open(self.character_model.character_image_file_name),
|
||||
(MainFrame.IMAGE_SIZE, MainFrame.IMAGE_SIZE))
|
||||
w, h = pil_image.size
|
||||
self.wx_source_image = wx.Bitmap.FromBufferRGBA(w, h, pil_image.convert("RGBA").tobytes())
|
||||
self.poser = self.character_model.get_poser(self.device)
|
||||
self.source_image_dirty = True
|
||||
self.Refresh()
|
||||
self.Update()
|
||||
except RuntimeError as e:
|
||||
message_dialog = wx.MessageDialog(
|
||||
self, "Could not load character model " + character_model_file_name, "Poser", wx.OK)
|
||||
message_dialog.ShowModal()
|
||||
message_dialog.Destroy()
|
||||
file_dialog.Destroy()
|
||||
|
||||
def paint_source_image_panel(self, event: wx.Event):
|
||||
wx.BufferedPaintDC(self.source_image_panel, self.source_image_bitmap)
|
||||
|
||||
def paint_result_image_panel(self, event: wx.Event):
|
||||
wx.BufferedPaintDC(self.result_image_panel, self.result_image_bitmap)
|
||||
|
||||
def draw_nothing_yet_string_to_bitmap(self, bitmap):
|
||||
dc = wx.MemoryDC()
|
||||
dc.SelectObject(bitmap)
|
||||
|
||||
dc.Clear()
|
||||
font = wx.Font(wx.FontInfo(14).Family(wx.FONTFAMILY_SWISS))
|
||||
dc.SetFont(font)
|
||||
w, h = dc.GetTextExtent("Nothing yet!")
|
||||
dc.DrawText("Nothing yet!", (MainFrame.IMAGE_SIZE - w) // 2, (MainFrame.IMAGE_SIZE - - h) // 2)
|
||||
|
||||
del dc
|
||||
|
||||
def get_current_pose(self):
|
||||
current_pose = [0.0 for i in range(MainFrame.NUM_PARAMETERS)]
|
||||
for morph_control_panel in self.morph_control_panels.values():
|
||||
morph_control_panel.set_param_value(current_pose)
|
||||
for rotation_control_panel in self.non_morph_control_panels.values():
|
||||
rotation_control_panel.set_param_value(current_pose)
|
||||
return current_pose
|
||||
|
||||
def update_images(self, event: wx.Event):
|
||||
current_pose = self.get_current_pose()
|
||||
if not self.source_image_dirty \
|
||||
and self.last_pose is not None \
|
||||
and self.last_pose == current_pose \
|
||||
and self.last_output_index == self.output_index_choice.GetSelection():
|
||||
return
|
||||
self.last_pose = current_pose
|
||||
self.last_output_index = self.output_index_choice.GetSelection()
|
||||
|
||||
if self.torch_source_image is None or self.poser is None:
|
||||
self.draw_nothing_yet_string_to_bitmap(self.source_image_bitmap)
|
||||
self.draw_nothing_yet_string_to_bitmap(self.result_image_bitmap)
|
||||
self.source_image_dirty = False
|
||||
self.Refresh()
|
||||
self.Update()
|
||||
return
|
||||
|
||||
if self.source_image_dirty:
|
||||
dc = wx.MemoryDC()
|
||||
dc.SelectObject(self.source_image_bitmap)
|
||||
dc.Clear()
|
||||
dc.DrawBitmap(self.wx_source_image, 0, 0)
|
||||
self.source_image_dirty = False
|
||||
|
||||
pose = torch.tensor(current_pose, device=self.device)
|
||||
output_index = self.output_index_choice.GetSelection()
|
||||
with torch.no_grad():
|
||||
start_cuda_event = torch.cuda.Event(enable_timing=True)
|
||||
end_cuda_event = torch.cuda.Event(enable_timing=True)
|
||||
start_cuda_event.record()
|
||||
start_time = time.time()
|
||||
|
||||
output_image = self.poser.pose(self.torch_source_image, pose, output_index)[0].detach().cpu()
|
||||
|
||||
end_time = time.time()
|
||||
end_cuda_event.record()
|
||||
torch.cuda.synchronize()
|
||||
print("cuda time (ms):", start_cuda_event.elapsed_time(end_cuda_event))
|
||||
print("elapsed time (ms):", (end_time - start_time) * 1000.0)
|
||||
|
||||
numpy_image = convert_output_image_from_torch_to_numpy(output_image)
|
||||
self.last_output_numpy_image = numpy_image
|
||||
wx_image = wx.ImageFromBuffer(
|
||||
numpy_image.shape[0],
|
||||
numpy_image.shape[1],
|
||||
numpy_image[:, :, 0:3].tobytes(),
|
||||
numpy_image[:, :, 3].tobytes())
|
||||
wx_bitmap = wx_image.ConvertToBitmap()
|
||||
|
||||
dc = wx.MemoryDC()
|
||||
dc.SelectObject(self.result_image_bitmap)
|
||||
dc.Clear()
|
||||
dc.DrawBitmap(wx_bitmap,
|
||||
(MainFrame.IMAGE_SIZE - numpy_image.shape[0]) // 2,
|
||||
(MainFrame.IMAGE_SIZE - numpy_image.shape[1]) // 2,
|
||||
True)
|
||||
del dc
|
||||
|
||||
self.Refresh()
|
||||
self.Update()
|
||||
|
||||
def on_save_image(self, event: wx.Event):
|
||||
if self.last_output_numpy_image is None:
|
||||
logging.info("There is no output image to save!!!")
|
||||
return
|
||||
|
||||
dir_name = "data/images"
|
||||
file_dialog = wx.FileDialog(self, "Choose an image", dir_name, "", "*.png", wx.FD_SAVE)
|
||||
if file_dialog.ShowModal() == wx.ID_OK:
|
||||
image_file_name = os.path.join(file_dialog.GetDirectory(), file_dialog.GetFilename())
|
||||
try:
|
||||
if os.path.exists(image_file_name):
|
||||
message_dialog = wx.MessageDialog(self, f"Override {image_file_name}", "Manual Poser",
|
||||
wx.YES_NO | wx.ICON_QUESTION)
|
||||
result = message_dialog.ShowModal()
|
||||
if result == wx.ID_YES:
|
||||
self.save_last_numpy_image(image_file_name)
|
||||
message_dialog.Destroy()
|
||||
else:
|
||||
self.save_last_numpy_image(image_file_name)
|
||||
except:
|
||||
message_dialog = wx.MessageDialog(self, f"Could not save {image_file_name}", "Manual Poser", wx.OK)
|
||||
message_dialog.ShowModal()
|
||||
message_dialog.Destroy()
|
||||
file_dialog.Destroy()
|
||||
|
||||
def save_last_numpy_image(self, image_file_name):
|
||||
numpy_image = self.last_output_numpy_image
|
||||
pil_image = PIL.Image.fromarray(numpy_image, mode='RGBA')
|
||||
os.makedirs(os.path.dirname(image_file_name), exist_ok=True)
|
||||
pil_image.save(image_file_name)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
device = torch.device('cuda:0')
|
||||
app = wx.App()
|
||||
main_frame = MainFrame(device)
|
||||
main_frame.Show(True)
|
||||
main_frame.timer.Start(16)
|
||||
app.MainLoop()
|
||||
@@ -0,0 +1,431 @@
|
||||
import os
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
from typing import Optional
|
||||
import PIL.Image
|
||||
|
||||
import cv2
|
||||
import mediapipe
|
||||
from scipy.spatial.transform import Rotation
|
||||
|
||||
from tha4.shion.base.image_util import resize_PIL_image
|
||||
from tha4.charmodel.character_model import CharacterModel
|
||||
from tha4.image_util import convert_linear_to_srgb
|
||||
from tha4.mocap.mediapipe_constants import HEAD_ROTATIONS, HEAD_X, HEAD_Y, HEAD_Z
|
||||
from tha4.mocap.mediapipe_face_pose import MediaPipeFacePose
|
||||
from tha4.mocap.mediapipe_face_pose_converter_00 import MediaPoseFacePoseConverter00
|
||||
|
||||
sys.path.append(os.getcwd())
|
||||
|
||||
import torch
|
||||
import wx
|
||||
|
||||
|
||||
class FpsStatistics:
|
||||
def __init__(self):
|
||||
self.count = 100
|
||||
self.fps = []
|
||||
|
||||
def add_fps(self, fps):
|
||||
self.fps.append(fps)
|
||||
while len(self.fps) > self.count:
|
||||
del self.fps[0]
|
||||
|
||||
def get_average_fps(self):
|
||||
if len(self.fps) == 0:
|
||||
return 0.0
|
||||
else:
|
||||
return sum(self.fps) / len(self.fps)
|
||||
|
||||
|
||||
class MainFrame(wx.Frame):
|
||||
IMAGE_SIZE = 512
|
||||
|
||||
def __init__(self,
|
||||
pose_converter: MediaPoseFacePoseConverter00,
|
||||
video_capture,
|
||||
face_landmarker,
|
||||
device: torch.device):
|
||||
super().__init__(None, wx.ID_ANY, "THA4 Character Model MediaPipe Puppeteer")
|
||||
self.face_landmarker = face_landmarker
|
||||
self.video_capture = video_capture
|
||||
self.pose_converter = pose_converter
|
||||
self.device = device
|
||||
|
||||
self.source_image_bitmap = wx.Bitmap(MainFrame.IMAGE_SIZE, MainFrame.IMAGE_SIZE)
|
||||
self.result_image_bitmap = wx.Bitmap(MainFrame.IMAGE_SIZE, MainFrame.IMAGE_SIZE)
|
||||
self.webcam_capture_bitmap = wx.Bitmap(256, 192)
|
||||
self.wx_source_image = None
|
||||
self.torch_source_image = None
|
||||
self.last_pose = None
|
||||
self.mediapipe_face_pose = None
|
||||
self.fps_statistics = FpsStatistics()
|
||||
self.last_update_time = None
|
||||
self.character_model = None
|
||||
self.poser = None
|
||||
|
||||
self.create_ui()
|
||||
self.create_timers()
|
||||
self.Bind(wx.EVT_CLOSE, self.on_close)
|
||||
|
||||
self.update_source_image_bitmap()
|
||||
self.update_result_image_bitmap()
|
||||
|
||||
def create_timers(self):
|
||||
self.capture_timer = wx.Timer(self, wx.ID_ANY)
|
||||
self.Bind(wx.EVT_TIMER, self.update_capture_panel, id=self.capture_timer.GetId())
|
||||
self.animation_timer = wx.Timer(self, wx.ID_ANY)
|
||||
self.Bind(wx.EVT_TIMER, self.update_result_image_bitmap, id=self.animation_timer.GetId())
|
||||
|
||||
def on_close(self, event: wx.Event):
|
||||
# Stop the timers
|
||||
self.animation_timer.Stop()
|
||||
self.capture_timer.Stop()
|
||||
|
||||
# Destroy the windows
|
||||
self.Destroy()
|
||||
event.Skip()
|
||||
|
||||
def on_erase_background(self, event: wx.Event):
|
||||
pass
|
||||
|
||||
def create_animation_panel(self, parent):
|
||||
self.animation_panel = wx.Panel(parent, style=wx.RAISED_BORDER)
|
||||
self.animation_panel_sizer = wx.BoxSizer(wx.HORIZONTAL)
|
||||
self.animation_panel.SetSizer(self.animation_panel_sizer)
|
||||
self.animation_panel.SetAutoLayout(1)
|
||||
|
||||
image_size = MainFrame.IMAGE_SIZE
|
||||
|
||||
if True:
|
||||
self.input_panel = wx.Panel(self.animation_panel, size=(image_size, image_size + 128),
|
||||
style=wx.SIMPLE_BORDER)
|
||||
self.input_panel_sizer = wx.BoxSizer(wx.VERTICAL)
|
||||
self.input_panel.SetSizer(self.input_panel_sizer)
|
||||
self.input_panel.SetAutoLayout(1)
|
||||
self.animation_panel_sizer.Add(self.input_panel, 0, wx.FIXED_MINSIZE)
|
||||
|
||||
self.source_image_panel = wx.Panel(self.input_panel, size=(image_size, image_size), style=wx.SIMPLE_BORDER)
|
||||
self.source_image_panel.Bind(wx.EVT_PAINT, self.paint_source_image_panel)
|
||||
self.source_image_panel.Bind(wx.EVT_ERASE_BACKGROUND, self.on_erase_background)
|
||||
self.input_panel_sizer.Add(self.source_image_panel, 0, wx.FIXED_MINSIZE)
|
||||
|
||||
self.load_model_button = wx.Button(self.input_panel, wx.ID_ANY, "Load Model")
|
||||
self.input_panel_sizer.Add(self.load_model_button, 1, wx.EXPAND)
|
||||
self.load_model_button.Bind(wx.EVT_BUTTON, self.load_model)
|
||||
|
||||
self.input_panel_sizer.Fit(self.input_panel)
|
||||
|
||||
if True:
|
||||
def current_pose_supplier() -> Optional[MediaPipeFacePose]:
|
||||
return self.mediapipe_face_pose
|
||||
|
||||
self.pose_converter.init_pose_converter_panel(self.animation_panel, current_pose_supplier)
|
||||
|
||||
if True:
|
||||
self.animation_left_panel = wx.Panel(self.animation_panel, style=wx.SIMPLE_BORDER)
|
||||
self.animation_left_panel_sizer = wx.BoxSizer(wx.VERTICAL)
|
||||
self.animation_left_panel.SetSizer(self.animation_left_panel_sizer)
|
||||
self.animation_left_panel.SetAutoLayout(1)
|
||||
self.animation_panel_sizer.Add(self.animation_left_panel, 0, wx.EXPAND)
|
||||
|
||||
self.result_image_panel = wx.Panel(self.animation_left_panel, size=(image_size, image_size),
|
||||
style=wx.SIMPLE_BORDER)
|
||||
self.result_image_panel.Bind(wx.EVT_PAINT, self.paint_result_image_panel)
|
||||
self.result_image_panel.Bind(wx.EVT_ERASE_BACKGROUND, self.on_erase_background)
|
||||
self.animation_left_panel_sizer.Add(self.result_image_panel, 0, wx.FIXED_MINSIZE)
|
||||
|
||||
separator = wx.StaticLine(self.animation_left_panel, -1, size=(256, 5))
|
||||
self.animation_left_panel_sizer.Add(separator, 0, wx.EXPAND)
|
||||
|
||||
background_text = wx.StaticText(self.animation_left_panel, label="--- Background ---",
|
||||
style=wx.ALIGN_CENTER)
|
||||
self.animation_left_panel_sizer.Add(background_text, 0, wx.EXPAND)
|
||||
|
||||
self.output_background_choice = wx.Choice(
|
||||
self.animation_left_panel,
|
||||
choices=[
|
||||
"TRANSPARENT",
|
||||
"GREEN",
|
||||
"BLUE",
|
||||
"BLACK",
|
||||
"WHITE"
|
||||
])
|
||||
self.output_background_choice.SetSelection(0)
|
||||
self.animation_left_panel_sizer.Add(self.output_background_choice, 0, wx.EXPAND)
|
||||
|
||||
separator = wx.StaticLine(self.animation_left_panel, -1, size=(256, 5))
|
||||
self.animation_left_panel_sizer.Add(separator, 0, wx.EXPAND)
|
||||
|
||||
self.fps_text = wx.StaticText(self.animation_left_panel, label="")
|
||||
self.animation_left_panel_sizer.Add(self.fps_text, wx.SizerFlags().Border())
|
||||
|
||||
self.animation_left_panel_sizer.Fit(self.animation_left_panel)
|
||||
|
||||
self.animation_panel_sizer.Fit(self.animation_panel)
|
||||
|
||||
def create_ui(self):
|
||||
self.main_sizer = wx.BoxSizer(wx.VERTICAL)
|
||||
self.SetSizer(self.main_sizer)
|
||||
self.SetAutoLayout(1)
|
||||
|
||||
self.capture_pose_lock = threading.Lock()
|
||||
|
||||
self.create_animation_panel(self)
|
||||
self.main_sizer.Add(self.animation_panel, wx.SizerFlags(0).Expand().Border(wx.ALL, 5))
|
||||
|
||||
self.create_capture_panel(self)
|
||||
self.main_sizer.Add(self.capture_panel, wx.SizerFlags(0).Expand().Border(wx.ALL, 5))
|
||||
|
||||
self.main_sizer.Fit(self)
|
||||
|
||||
def create_capture_panel(self, parent):
|
||||
self.capture_panel = wx.Panel(parent, style=wx.RAISED_BORDER)
|
||||
self.capture_panel_sizer = wx.BoxSizer(wx.HORIZONTAL)
|
||||
self.capture_panel.SetSizer(self.capture_panel_sizer)
|
||||
self.capture_panel.SetAutoLayout(1)
|
||||
|
||||
self.webcam_capture_panel = wx.Panel(self.capture_panel, size=(256, 192), style=wx.SIMPLE_BORDER)
|
||||
self.webcam_capture_panel.Bind(wx.EVT_PAINT, self.paint_webcam_capture_panel)
|
||||
self.webcam_capture_panel.Bind(wx.EVT_ERASE_BACKGROUND, self.on_erase_background)
|
||||
self.capture_panel_sizer.Add(self.webcam_capture_panel, wx.SizerFlags(0).FixedMinSize().Border(wx.ALL, 5))
|
||||
|
||||
self.rotation_labels = {}
|
||||
self.rotation_value_labels = {}
|
||||
rotation_column = self.create_rotation_column(self.capture_panel, HEAD_ROTATIONS)
|
||||
self.capture_panel_sizer.Add(rotation_column, wx.SizerFlags(0).Expand().Border(wx.ALL, 3))
|
||||
|
||||
def paint_webcam_capture_panel(self, event: wx.Event):
|
||||
wx.BufferedPaintDC(self.webcam_capture_panel, self.webcam_capture_bitmap)
|
||||
|
||||
def create_rotation_column(self, parent, rotation_names):
|
||||
column_panel = wx.Panel(parent, style=wx.SIMPLE_BORDER)
|
||||
column_panel_sizer = wx.FlexGridSizer(cols=2)
|
||||
column_panel_sizer.AddGrowableCol(1)
|
||||
column_panel.SetSizer(column_panel_sizer)
|
||||
column_panel.SetAutoLayout(1)
|
||||
|
||||
for rotation_name in rotation_names:
|
||||
self.rotation_labels[rotation_name] = wx.StaticText(
|
||||
column_panel, label=rotation_name, style=wx.ALIGN_RIGHT)
|
||||
column_panel_sizer.Add(self.rotation_labels[rotation_name],
|
||||
wx.SizerFlags(1).Expand().Border(wx.ALL, 3))
|
||||
|
||||
self.rotation_value_labels[rotation_name] = wx.TextCtrl(
|
||||
column_panel, style=wx.TE_RIGHT)
|
||||
self.rotation_value_labels[rotation_name].SetValue("0.00")
|
||||
self.rotation_value_labels[rotation_name].Disable()
|
||||
column_panel_sizer.Add(self.rotation_value_labels[rotation_name],
|
||||
wx.SizerFlags(1).Expand().Border(wx.ALL, 3))
|
||||
|
||||
column_panel.GetSizer().Fit(column_panel)
|
||||
return column_panel
|
||||
|
||||
def update_capture_panel(self, event: wx.Event):
|
||||
there_is_frame, frame = self.video_capture.read()
|
||||
if not there_is_frame:
|
||||
dc = wx.MemoryDC()
|
||||
dc.SelectObject(self.webcam_capture_bitmap)
|
||||
self.draw_nothing_yet_string(dc)
|
||||
del dc
|
||||
return
|
||||
|
||||
rgb_frame = cv2.flip(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB), 1)
|
||||
# Use INTER_AREA for downscaling (better quality than INTER_LINEAR)
|
||||
resized_frame = cv2.resize(
|
||||
rgb_frame, (256, 192),
|
||||
interpolation=cv2.INTER_AREA
|
||||
)
|
||||
wx_image = wx.ImageFromBuffer(256, 192, resized_frame.tobytes())
|
||||
wx_bitmap = wx_image.ConvertToBitmap()
|
||||
|
||||
dc = wx.MemoryDC()
|
||||
dc.SelectObject(self.webcam_capture_bitmap)
|
||||
dc.Clear()
|
||||
dc.DrawBitmap(wx_bitmap, 0, 0, True)
|
||||
del dc
|
||||
|
||||
self.webcam_capture_panel.Refresh()
|
||||
|
||||
time_ms = int(time.time() * 1000)
|
||||
mediapipe_image = mediapipe.Image(image_format=mediapipe.ImageFormat.SRGB, data=rgb_frame)
|
||||
detection_result = self.face_landmarker.detect_for_video(mediapipe_image, time_ms)
|
||||
self.update_mediapipe_face_pose(detection_result)
|
||||
|
||||
def update_mediapipe_face_pose(self, detection_result):
|
||||
if len(detection_result.facial_transformation_matrixes) == 0:
|
||||
return
|
||||
|
||||
xform_matrix = detection_result.facial_transformation_matrixes[0]
|
||||
blendshape_params = {}
|
||||
for item in detection_result.face_blendshapes[0]:
|
||||
blendshape_params[item.category_name] = item.score
|
||||
M = xform_matrix[0:3, 0:3]
|
||||
rot = Rotation.from_matrix(M)
|
||||
euler_angles = rot.as_euler('xyz', degrees=True)
|
||||
|
||||
self.rotation_value_labels[HEAD_X].SetValue("%0.2f" % euler_angles[0])
|
||||
self.rotation_value_labels[HEAD_X].Refresh()
|
||||
self.rotation_value_labels[HEAD_Y].SetValue("%0.2f" % euler_angles[1])
|
||||
self.rotation_value_labels[HEAD_Y].Refresh()
|
||||
self.rotation_value_labels[HEAD_Z].SetValue("%0.2f" % euler_angles[2])
|
||||
self.rotation_value_labels[HEAD_Z].Refresh()
|
||||
|
||||
self.mediapipe_face_pose = MediaPipeFacePose(blendshape_params, xform_matrix)
|
||||
|
||||
@staticmethod
|
||||
def convert_to_100(x):
|
||||
return int(max(0.0, min(1.0, x)) * 100)
|
||||
|
||||
def paint_source_image_panel(self, event: wx.Event):
|
||||
wx.BufferedPaintDC(self.source_image_panel, self.source_image_bitmap)
|
||||
|
||||
def update_source_image_bitmap(self):
|
||||
dc = wx.MemoryDC()
|
||||
dc.SelectObject(self.source_image_bitmap)
|
||||
if self.wx_source_image is None:
|
||||
self.draw_nothing_yet_string(dc)
|
||||
else:
|
||||
dc.Clear()
|
||||
dc.DrawBitmap(self.wx_source_image, 0, 0, True)
|
||||
del dc
|
||||
|
||||
def draw_nothing_yet_string(self, dc):
|
||||
dc.Clear()
|
||||
font = wx.Font(wx.FontInfo(14).Family(wx.FONTFAMILY_SWISS))
|
||||
dc.SetFont(font)
|
||||
w, h = dc.GetTextExtent("Nothing yet!")
|
||||
dc.DrawText("Nothing yet!", (MainFrame.IMAGE_SIZE - w) // 2, (MainFrame.IMAGE_SIZE - h) // 2)
|
||||
|
||||
def paint_result_image_panel(self, event: wx.Event):
|
||||
wx.BufferedPaintDC(self.result_image_panel, self.result_image_bitmap)
|
||||
|
||||
def update_result_image_bitmap(self, event: Optional[wx.Event] = None):
|
||||
if self.mediapipe_face_pose is None or self.poser is None:
|
||||
dc = wx.MemoryDC()
|
||||
dc.SelectObject(self.result_image_bitmap)
|
||||
self.draw_nothing_yet_string(dc)
|
||||
del dc
|
||||
return
|
||||
|
||||
current_pose = self.pose_converter.convert(self.mediapipe_face_pose)
|
||||
if self.last_pose is not None and self.last_pose == current_pose:
|
||||
return
|
||||
self.last_pose = current_pose
|
||||
|
||||
if self.torch_source_image is None:
|
||||
dc = wx.MemoryDC()
|
||||
dc.SelectObject(self.result_image_bitmap)
|
||||
self.draw_nothing_yet_string(dc)
|
||||
del dc
|
||||
return
|
||||
|
||||
pose = torch.tensor(current_pose, device=self.device, dtype=self.poser.get_dtype())
|
||||
|
||||
with torch.no_grad():
|
||||
output_image = self.poser.pose(self.torch_source_image, pose)[0].float()
|
||||
output_image = torch.clip((output_image + 1.0) / 2.0, 0.0, 1.0)
|
||||
output_image = convert_linear_to_srgb(output_image)
|
||||
|
||||
background_choice = self.output_background_choice.GetSelection()
|
||||
if background_choice == 0:
|
||||
pass
|
||||
else:
|
||||
background = torch.zeros(4, output_image.shape[1], output_image.shape[2], device=self.device)
|
||||
background[3, :, :] = 1.0
|
||||
if background_choice == 1:
|
||||
background[1, :, :] = 1.0
|
||||
output_image = self.blend_with_background(output_image, background)
|
||||
elif background_choice == 2:
|
||||
background[2, :, :] = 1.0
|
||||
output_image = self.blend_with_background(output_image, background)
|
||||
elif background_choice == 3:
|
||||
output_image = self.blend_with_background(output_image, background)
|
||||
else:
|
||||
background[0:3, :, :] = 1.0
|
||||
output_image = self.blend_with_background(output_image, background)
|
||||
|
||||
c, h, w = output_image.shape
|
||||
output_image = 255.0 * torch.transpose(output_image.reshape(c, h * w), 0, 1).reshape(h, w, c)
|
||||
output_image = output_image.byte()
|
||||
|
||||
numpy_image = output_image.detach().cpu().numpy()
|
||||
wx_image = wx.ImageFromBuffer(numpy_image.shape[0],
|
||||
numpy_image.shape[1],
|
||||
numpy_image[:, :, 0:3].tobytes(),
|
||||
numpy_image[:, :, 3].tobytes())
|
||||
wx_bitmap = wx_image.ConvertToBitmap()
|
||||
|
||||
dc = wx.MemoryDC()
|
||||
dc.SelectObject(self.result_image_bitmap)
|
||||
dc.Clear()
|
||||
dc.DrawBitmap(wx_bitmap,
|
||||
(MainFrame.IMAGE_SIZE - numpy_image.shape[0]) // 2,
|
||||
(MainFrame.IMAGE_SIZE - numpy_image.shape[1]) // 2, True)
|
||||
del dc
|
||||
|
||||
time_now = time.time_ns()
|
||||
if self.last_update_time is not None:
|
||||
elapsed_time = time_now - self.last_update_time
|
||||
fps = 1.0 / (elapsed_time / 10 ** 9)
|
||||
if self.torch_source_image is not None:
|
||||
self.fps_statistics.add_fps(fps)
|
||||
self.fps_text.SetLabelText("FPS = %0.2f" % self.fps_statistics.get_average_fps())
|
||||
self.last_update_time = time_now
|
||||
|
||||
self.Refresh()
|
||||
|
||||
def blend_with_background(self, numpy_image, background):
|
||||
alpha = numpy_image[3:4, :, :]
|
||||
color = numpy_image[0:3, :, :]
|
||||
new_color = color * alpha + (1.0 - alpha) * background[0:3, :, :]
|
||||
return torch.cat([new_color, background[3:4, :, :]], dim=0)
|
||||
|
||||
def load_model(self, event: wx.Event):
|
||||
dir_name = "data/character_models"
|
||||
file_dialog = wx.FileDialog(self, "Choose a model", dir_name, "", "*.yaml", wx.FD_OPEN)
|
||||
if file_dialog.ShowModal() == wx.ID_OK:
|
||||
character_model_json_file_name = os.path.join(file_dialog.GetDirectory(), file_dialog.GetFilename())
|
||||
try:
|
||||
self.character_model = CharacterModel.load(character_model_json_file_name)
|
||||
self.torch_source_image = self.character_model.get_character_image(self.device)
|
||||
pil_image = resize_PIL_image(
|
||||
PIL.Image.open(self.character_model.character_image_file_name),
|
||||
(MainFrame.IMAGE_SIZE, MainFrame.IMAGE_SIZE))
|
||||
w, h = pil_image.size
|
||||
self.wx_source_image = wx.Bitmap.FromBufferRGBA(w, h, pil_image.convert("RGBA").tobytes())
|
||||
self.update_source_image_bitmap()
|
||||
self.poser = self.character_model.get_poser(self.device)
|
||||
except Exception:
|
||||
message_dialog = wx.MessageDialog(
|
||||
self, "Could not load character model " + character_model_json_file_name, "Poser", wx.OK)
|
||||
message_dialog.ShowModal()
|
||||
message_dialog.Destroy()
|
||||
file_dialog.Destroy()
|
||||
self.Refresh()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
device = torch.device("cuda:0")
|
||||
|
||||
pose_converter = MediaPoseFacePoseConverter00()
|
||||
|
||||
face_landmarker_base_options = mediapipe.tasks.BaseOptions(
|
||||
model_asset_path='data/thirdparty/mediapipe/face_landmarker_v2_with_blendshapes.task')
|
||||
options = mediapipe.tasks.vision.FaceLandmarkerOptions(
|
||||
base_options=face_landmarker_base_options,
|
||||
running_mode=mediapipe.tasks.vision.RunningMode.VIDEO,
|
||||
output_face_blendshapes=True,
|
||||
output_facial_transformation_matrixes=True,
|
||||
num_faces=1)
|
||||
face_landmarker = mediapipe.tasks.vision.FaceLandmarker.create_from_options(options)
|
||||
|
||||
video_capture = cv2.VideoCapture(0)
|
||||
|
||||
app = wx.App()
|
||||
main_frame = MainFrame(pose_converter, video_capture, face_landmarker, device)
|
||||
main_frame.Show(True)
|
||||
main_frame.capture_timer.Start(30)
|
||||
main_frame.animation_timer.Start(30)
|
||||
app.MainLoop()
|
||||
@@ -0,0 +1,25 @@
|
||||
import argparse
|
||||
import logging
|
||||
|
||||
from tha4.distiller.distiller_config import DistillerConfig
|
||||
from tha4.pytasuku.workspace import Workspace
|
||||
|
||||
|
||||
def run_config(config_file_name: str):
|
||||
config = DistillerConfig.load(config_file_name)
|
||||
|
||||
logging.basicConfig(level=logging.INFO, force=True)
|
||||
workspace = Workspace()
|
||||
config.define_tasks(workspace)
|
||||
|
||||
workspace.start_session()
|
||||
workspace.run(f"{config.prefix}/all")
|
||||
workspace.end_session()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(description='Training script.')
|
||||
parser.add_argument("--config_file", type=str, required=True,
|
||||
help="The name of the config file for the distillation process.")
|
||||
args = parser.parse_args()
|
||||
run_config(args.config_file)
|
||||
@@ -0,0 +1,13 @@
|
||||
from tha4.app.distill import run_config
|
||||
from tha4.distiller.ui.distiller_ui_main_frame import DistillerUiMainFrame
|
||||
|
||||
import wx
|
||||
|
||||
if __name__ == "__main__":
|
||||
app = wx.App()
|
||||
main_frame = DistillerUiMainFrame()
|
||||
main_frame.Show(True)
|
||||
app.MainLoop()
|
||||
|
||||
if main_frame.config_file_to_run is not None:
|
||||
run_config(main_frame.config_file_to_run)
|
||||
@@ -0,0 +1,468 @@
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
from typing import List
|
||||
|
||||
from tha4.shion.base.image_util import extract_pytorch_image_from_PIL_image, pytorch_rgba_to_numpy_image, \
|
||||
pytorch_rgb_to_numpy_image
|
||||
from tha4.image_util import grid_change_to_numpy_image, resize_PIL_image
|
||||
|
||||
sys.path.append(os.getcwd())
|
||||
|
||||
import PIL.Image
|
||||
import numpy
|
||||
import torch
|
||||
import wx
|
||||
|
||||
from tha4.poser.poser import Poser, PoseParameterCategory, PoseParameterGroup
|
||||
|
||||
|
||||
class MorphCategoryControlPanel(wx.Panel):
|
||||
def __init__(self,
|
||||
parent,
|
||||
title: str,
|
||||
pose_param_category: PoseParameterCategory,
|
||||
param_groups: List[PoseParameterGroup]):
|
||||
super().__init__(parent, style=wx.SIMPLE_BORDER)
|
||||
self.pose_param_category = pose_param_category
|
||||
self.sizer = wx.BoxSizer(wx.VERTICAL)
|
||||
self.SetSizer(self.sizer)
|
||||
self.SetAutoLayout(1)
|
||||
|
||||
title_text = wx.StaticText(self, label=title, style=wx.ALIGN_CENTER)
|
||||
self.sizer.Add(title_text, 0, wx.EXPAND)
|
||||
|
||||
self.param_groups = [group for group in param_groups if group.get_category() == pose_param_category]
|
||||
self.choice = wx.Choice(self, choices=[group.get_group_name() for group in self.param_groups])
|
||||
if len(self.param_groups) > 0:
|
||||
self.choice.SetSelection(0)
|
||||
self.choice.Bind(wx.EVT_CHOICE, self.on_choice_updated)
|
||||
self.sizer.Add(self.choice, 0, wx.EXPAND)
|
||||
|
||||
self.left_slider = wx.Slider(self, minValue=-1000, maxValue=1000, value=-1000, style=wx.HORIZONTAL)
|
||||
self.sizer.Add(self.left_slider, 0, wx.EXPAND)
|
||||
|
||||
self.right_slider = wx.Slider(self, minValue=-1000, maxValue=1000, value=-1000, style=wx.HORIZONTAL)
|
||||
self.sizer.Add(self.right_slider, 0, wx.EXPAND)
|
||||
|
||||
self.checkbox = wx.CheckBox(self, label="Show")
|
||||
self.checkbox.SetValue(True)
|
||||
self.sizer.Add(self.checkbox, 0, wx.SHAPED | wx.ALIGN_CENTER)
|
||||
|
||||
self.update_ui()
|
||||
|
||||
self.sizer.Fit(self)
|
||||
|
||||
def update_ui(self):
|
||||
param_group = self.param_groups[self.choice.GetSelection()]
|
||||
if param_group.is_discrete():
|
||||
self.left_slider.Enable(False)
|
||||
self.right_slider.Enable(False)
|
||||
self.checkbox.Enable(True)
|
||||
elif param_group.get_arity() == 1:
|
||||
self.left_slider.Enable(True)
|
||||
self.right_slider.Enable(False)
|
||||
self.checkbox.Enable(False)
|
||||
else:
|
||||
self.left_slider.Enable(True)
|
||||
self.right_slider.Enable(True)
|
||||
self.checkbox.Enable(False)
|
||||
|
||||
def on_choice_updated(self, event: wx.Event):
|
||||
param_group = self.param_groups[self.choice.GetSelection()]
|
||||
if param_group.is_discrete():
|
||||
self.checkbox.SetValue(True)
|
||||
self.update_ui()
|
||||
|
||||
def set_param_value(self, pose: List[float]):
|
||||
if len(self.param_groups) == 0:
|
||||
return
|
||||
selected_morph_index = self.choice.GetSelection()
|
||||
param_group = self.param_groups[selected_morph_index]
|
||||
param_index = param_group.get_parameter_index()
|
||||
if param_group.is_discrete():
|
||||
if self.checkbox.GetValue():
|
||||
for i in range(param_group.get_arity()):
|
||||
pose[param_index + i] = 1.0
|
||||
else:
|
||||
param_range = param_group.get_range()
|
||||
alpha = (self.left_slider.GetValue() + 1000) / 2000.0
|
||||
pose[param_index] = param_range[0] + (param_range[1] - param_range[0]) * alpha
|
||||
if param_group.get_arity() == 2:
|
||||
alpha = (self.right_slider.GetValue() + 1000) / 2000.0
|
||||
pose[param_index + 1] = param_range[0] + (param_range[1] - param_range[0]) * alpha
|
||||
|
||||
|
||||
class SimpleParamGroupsControlPanel(wx.Panel):
|
||||
def __init__(self, parent,
|
||||
pose_param_category: PoseParameterCategory,
|
||||
param_groups: List[PoseParameterGroup]):
|
||||
super().__init__(parent, style=wx.SIMPLE_BORDER)
|
||||
self.sizer = wx.BoxSizer(wx.VERTICAL)
|
||||
self.SetSizer(self.sizer)
|
||||
self.SetAutoLayout(1)
|
||||
|
||||
self.param_groups = [group for group in param_groups if group.get_category() == pose_param_category]
|
||||
for param_group in self.param_groups:
|
||||
assert not param_group.is_discrete()
|
||||
assert param_group.get_arity() == 1
|
||||
|
||||
self.sliders = []
|
||||
for param_group in self.param_groups:
|
||||
static_text = wx.StaticText(
|
||||
self,
|
||||
label=" ------------ %s ------------ " % param_group.get_group_name(), style=wx.ALIGN_CENTER)
|
||||
self.sizer.Add(static_text, 0, wx.EXPAND)
|
||||
range = param_group.get_range()
|
||||
min_value = int(range[0] * 1000)
|
||||
max_value = int(range[1] * 1000)
|
||||
slider = wx.Slider(self, minValue=min_value, maxValue=max_value, value=0, style=wx.HORIZONTAL)
|
||||
self.sizer.Add(slider, 0, wx.EXPAND)
|
||||
self.sliders.append(slider)
|
||||
|
||||
self.sizer.Fit(self)
|
||||
|
||||
def set_param_value(self, pose: List[float]):
|
||||
if len(self.param_groups) == 0:
|
||||
return
|
||||
for param_group_index in range(len(self.param_groups)):
|
||||
param_group = self.param_groups[param_group_index]
|
||||
slider = self.sliders[param_group_index]
|
||||
param_range = param_group.get_range()
|
||||
param_index = param_group.get_parameter_index()
|
||||
alpha = (slider.GetValue() - slider.GetMin()) * 1.0 / (slider.GetMax() - slider.GetMin())
|
||||
pose[param_index] = param_range[0] + (param_range[1] - param_range[0]) * alpha
|
||||
|
||||
|
||||
def convert_output_image_from_torch_to_numpy(output_image):
|
||||
if output_image.shape[2] == 2:
|
||||
h, w, c = output_image.shape
|
||||
numpy_image = torch.transpose(output_image.reshape(h * w, c), 0, 1).reshape(c, h, w)
|
||||
elif output_image.shape[0] == 4:
|
||||
numpy_image = pytorch_rgba_to_numpy_image(output_image)
|
||||
elif output_image.shape[0] == 3:
|
||||
numpy_image = pytorch_rgb_to_numpy_image(output_image)
|
||||
elif output_image.shape[0] == 1:
|
||||
c, h, w = output_image.shape
|
||||
alpha_image = torch.cat([output_image.repeat(3, 1, 1) * 2.0 - 1.0, torch.ones(1, h, w)], dim=0)
|
||||
numpy_image = pytorch_rgba_to_numpy_image(alpha_image)
|
||||
elif output_image.shape[0] == 2:
|
||||
numpy_image = grid_change_to_numpy_image(output_image, num_channels=4)
|
||||
else:
|
||||
raise RuntimeError("Unsupported # image channels: %d" % output_image.shape[0])
|
||||
numpy_image = numpy.uint8(numpy.rint(numpy_image * 255.0))
|
||||
return numpy_image
|
||||
|
||||
|
||||
class MainFrame(wx.Frame):
|
||||
def __init__(self, poser: Poser, device: torch.device):
|
||||
super().__init__(None, wx.ID_ANY, "Poser")
|
||||
self.poser = poser
|
||||
self.dtype = self.poser.get_dtype()
|
||||
self.device = device
|
||||
self.image_size = self.poser.get_image_size()
|
||||
|
||||
self.wx_source_image = None
|
||||
self.torch_source_image = None
|
||||
|
||||
self.main_sizer = wx.BoxSizer(wx.HORIZONTAL)
|
||||
self.SetSizer(self.main_sizer)
|
||||
self.SetAutoLayout(1)
|
||||
self.init_left_panel()
|
||||
self.init_control_panel()
|
||||
self.init_right_panel()
|
||||
self.main_sizer.Fit(self)
|
||||
|
||||
self.timer = wx.Timer(self, wx.ID_ANY)
|
||||
self.Bind(wx.EVT_TIMER, self.update_images, self.timer)
|
||||
|
||||
save_image_id = wx.NewIdRef()
|
||||
self.Bind(wx.EVT_MENU, self.on_save_image, id=save_image_id)
|
||||
accelerator_table = wx.AcceleratorTable([
|
||||
(wx.ACCEL_CTRL, ord('S'), save_image_id)
|
||||
])
|
||||
self.SetAcceleratorTable(accelerator_table)
|
||||
|
||||
self.last_pose = None
|
||||
self.last_output_index = self.output_index_choice.GetSelection()
|
||||
self.last_output_numpy_image = None
|
||||
|
||||
self.wx_source_image = None
|
||||
self.torch_source_image = None
|
||||
self.source_image_bitmap = wx.Bitmap(self.image_size, self.image_size)
|
||||
self.result_image_bitmap = wx.Bitmap(self.image_size, self.image_size)
|
||||
self.source_image_dirty = True
|
||||
|
||||
def init_left_panel(self):
|
||||
self.control_panel = wx.Panel(self, style=wx.SIMPLE_BORDER, size=(self.image_size, -1))
|
||||
self.left_panel = wx.Panel(self, style=wx.SIMPLE_BORDER)
|
||||
left_panel_sizer = wx.BoxSizer(wx.VERTICAL)
|
||||
self.left_panel.SetSizer(left_panel_sizer)
|
||||
self.left_panel.SetAutoLayout(1)
|
||||
|
||||
self.source_image_panel = wx.Panel(self.left_panel, size=(self.image_size, self.image_size),
|
||||
style=wx.SIMPLE_BORDER)
|
||||
self.source_image_panel.Bind(wx.EVT_PAINT, self.paint_source_image_panel)
|
||||
self.source_image_panel.Bind(wx.EVT_ERASE_BACKGROUND, self.on_erase_background)
|
||||
left_panel_sizer.Add(self.source_image_panel, 0, wx.FIXED_MINSIZE)
|
||||
|
||||
self.load_image_button = wx.Button(self.left_panel, wx.ID_ANY, "\nLoad Image\n\n")
|
||||
left_panel_sizer.Add(self.load_image_button, 1, wx.EXPAND)
|
||||
self.load_image_button.Bind(wx.EVT_BUTTON, self.load_image)
|
||||
|
||||
left_panel_sizer.Fit(self.left_panel)
|
||||
self.main_sizer.Add(self.left_panel, 0, wx.FIXED_MINSIZE)
|
||||
|
||||
def on_erase_background(self, event: wx.Event):
|
||||
pass
|
||||
|
||||
def init_control_panel(self):
|
||||
self.control_panel_sizer = wx.BoxSizer(wx.VERTICAL)
|
||||
self.control_panel.SetSizer(self.control_panel_sizer)
|
||||
self.control_panel.SetMinSize(wx.Size(256, 1))
|
||||
|
||||
morph_categories = [
|
||||
PoseParameterCategory.EYEBROW,
|
||||
PoseParameterCategory.EYE,
|
||||
PoseParameterCategory.MOUTH,
|
||||
PoseParameterCategory.IRIS_MORPH
|
||||
]
|
||||
morph_category_titles = {
|
||||
PoseParameterCategory.EYEBROW: " ------------ Eyebrow ------------ ",
|
||||
PoseParameterCategory.EYE: " ------------ Eye ------------ ",
|
||||
PoseParameterCategory.MOUTH: " ------------ Mouth ------------ ",
|
||||
PoseParameterCategory.IRIS_MORPH: " ------------ Iris morphs ------------ ",
|
||||
}
|
||||
self.morph_control_panels = {}
|
||||
for category in morph_categories:
|
||||
param_groups = self.poser.get_pose_parameter_groups()
|
||||
filtered_param_groups = [group for group in param_groups if group.get_category() == category]
|
||||
if len(filtered_param_groups) == 0:
|
||||
continue
|
||||
control_panel = MorphCategoryControlPanel(
|
||||
self.control_panel,
|
||||
morph_category_titles[category],
|
||||
category,
|
||||
self.poser.get_pose_parameter_groups())
|
||||
self.morph_control_panels[category] = control_panel
|
||||
self.control_panel_sizer.Add(control_panel, 0, wx.EXPAND)
|
||||
|
||||
self.non_morph_control_panels = {}
|
||||
non_morph_categories = [
|
||||
PoseParameterCategory.IRIS_ROTATION,
|
||||
PoseParameterCategory.FACE_ROTATION,
|
||||
PoseParameterCategory.BODY_ROTATION,
|
||||
PoseParameterCategory.BREATHING
|
||||
]
|
||||
for category in non_morph_categories:
|
||||
param_groups = self.poser.get_pose_parameter_groups()
|
||||
filtered_param_groups = [group for group in param_groups if group.get_category() == category]
|
||||
if len(filtered_param_groups) == 0:
|
||||
continue
|
||||
control_panel = SimpleParamGroupsControlPanel(
|
||||
self.control_panel,
|
||||
category,
|
||||
self.poser.get_pose_parameter_groups())
|
||||
self.non_morph_control_panels[category] = control_panel
|
||||
self.control_panel_sizer.Add(control_panel, 0, wx.EXPAND)
|
||||
|
||||
self.control_panel_sizer.Fit(self.control_panel)
|
||||
self.main_sizer.Add(self.control_panel, 1, wx.FIXED_MINSIZE)
|
||||
|
||||
def init_right_panel(self):
|
||||
self.right_panel = wx.Panel(self, style=wx.SIMPLE_BORDER)
|
||||
right_panel_sizer = wx.BoxSizer(wx.VERTICAL)
|
||||
self.right_panel.SetSizer(right_panel_sizer)
|
||||
self.right_panel.SetAutoLayout(1)
|
||||
|
||||
self.result_image_panel = wx.Panel(self.right_panel,
|
||||
size=(self.image_size, self.image_size),
|
||||
style=wx.SIMPLE_BORDER)
|
||||
self.result_image_panel.Bind(wx.EVT_PAINT, self.paint_result_image_panel)
|
||||
self.result_image_panel.Bind(wx.EVT_ERASE_BACKGROUND, self.on_erase_background)
|
||||
self.output_index_choice = wx.Choice(
|
||||
self.right_panel,
|
||||
choices=[str(i) for i in range(self.poser.get_output_length())])
|
||||
self.output_index_choice.SetSelection(0)
|
||||
right_panel_sizer.Add(self.result_image_panel, 0, wx.FIXED_MINSIZE)
|
||||
right_panel_sizer.Add(self.output_index_choice, 0, wx.EXPAND)
|
||||
|
||||
self.save_image_button = wx.Button(self.right_panel, wx.ID_ANY, "\nSave Image\n\n")
|
||||
right_panel_sizer.Add(self.save_image_button, 1, wx.EXPAND)
|
||||
self.save_image_button.Bind(wx.EVT_BUTTON, self.on_save_image)
|
||||
|
||||
right_panel_sizer.Fit(self.right_panel)
|
||||
self.main_sizer.Add(self.right_panel, 0, wx.FIXED_MINSIZE)
|
||||
|
||||
def create_param_category_choice(self, param_category: PoseParameterCategory):
|
||||
params = []
|
||||
for param_group in self.poser.get_pose_parameter_groups():
|
||||
if param_group.get_category() == param_category:
|
||||
params.append(param_group.get_group_name())
|
||||
choice = wx.Choice(self.control_panel, choices=params)
|
||||
if len(params) > 0:
|
||||
choice.SetSelection(0)
|
||||
return choice
|
||||
|
||||
def load_image(self, event: wx.Event):
|
||||
dir_name = "data/images"
|
||||
file_dialog = wx.FileDialog(self, "Choose an image", dir_name, "", "*.png", wx.FD_OPEN)
|
||||
if file_dialog.ShowModal() == wx.ID_OK:
|
||||
image_file_name = os.path.join(file_dialog.GetDirectory(), file_dialog.GetFilename())
|
||||
try:
|
||||
pil_image = resize_PIL_image(PIL.Image.open(image_file_name),
|
||||
(self.poser.get_image_size(), self.poser.get_image_size()))
|
||||
w, h = pil_image.size
|
||||
if pil_image.mode != 'RGBA':
|
||||
self.source_image_string = "Image must have alpha channel!"
|
||||
self.wx_source_image = None
|
||||
self.torch_source_image = None
|
||||
else:
|
||||
self.wx_source_image = wx.Bitmap.FromBufferRGBA(w, h, pil_image.convert("RGBA").tobytes())
|
||||
self.torch_source_image = extract_pytorch_image_from_PIL_image(pil_image) \
|
||||
.to(self.device).to(self.dtype)
|
||||
self.source_image_dirty = True
|
||||
self.Refresh()
|
||||
self.Update()
|
||||
except:
|
||||
message_dialog = wx.MessageDialog(self, "Could not load image " + image_file_name, "Poser", wx.OK)
|
||||
message_dialog.ShowModal()
|
||||
message_dialog.Destroy()
|
||||
file_dialog.Destroy()
|
||||
|
||||
def paint_source_image_panel(self, event: wx.Event):
|
||||
wx.BufferedPaintDC(self.source_image_panel, self.source_image_bitmap)
|
||||
|
||||
def paint_result_image_panel(self, event: wx.Event):
|
||||
wx.BufferedPaintDC(self.result_image_panel, self.result_image_bitmap)
|
||||
|
||||
def draw_nothing_yet_string_to_bitmap(self, bitmap):
|
||||
dc = wx.MemoryDC()
|
||||
dc.SelectObject(bitmap)
|
||||
|
||||
dc.Clear()
|
||||
font = wx.Font(wx.FontInfo(14).Family(wx.FONTFAMILY_SWISS))
|
||||
dc.SetFont(font)
|
||||
w, h = dc.GetTextExtent("Nothing yet!")
|
||||
dc.DrawText("Nothing yet!", (self.image_size - w) // 2, (self.image_size - - h) // 2)
|
||||
|
||||
del dc
|
||||
|
||||
def get_current_pose(self):
|
||||
current_pose = [0.0 for i in range(self.poser.get_num_parameters())]
|
||||
for morph_control_panel in self.morph_control_panels.values():
|
||||
morph_control_panel.set_param_value(current_pose)
|
||||
for rotation_control_panel in self.non_morph_control_panels.values():
|
||||
rotation_control_panel.set_param_value(current_pose)
|
||||
return current_pose
|
||||
|
||||
def update_images(self, event: wx.Event):
|
||||
current_pose = self.get_current_pose()
|
||||
if not self.source_image_dirty \
|
||||
and self.last_pose is not None \
|
||||
and self.last_pose == current_pose \
|
||||
and self.last_output_index == self.output_index_choice.GetSelection():
|
||||
return
|
||||
self.last_pose = current_pose
|
||||
self.last_output_index = self.output_index_choice.GetSelection()
|
||||
|
||||
if self.torch_source_image is None:
|
||||
self.draw_nothing_yet_string_to_bitmap(self.source_image_bitmap)
|
||||
self.draw_nothing_yet_string_to_bitmap(self.result_image_bitmap)
|
||||
self.source_image_dirty = False
|
||||
self.Refresh()
|
||||
self.Update()
|
||||
return
|
||||
|
||||
if self.source_image_dirty:
|
||||
dc = wx.MemoryDC()
|
||||
dc.SelectObject(self.source_image_bitmap)
|
||||
dc.Clear()
|
||||
dc.DrawBitmap(self.wx_source_image, 0, 0)
|
||||
self.source_image_dirty = False
|
||||
|
||||
pose = torch.tensor(current_pose, device=self.device, dtype=self.dtype)
|
||||
output_index = self.output_index_choice.GetSelection()
|
||||
with torch.no_grad():
|
||||
start_cuda_event = torch.cuda.Event(enable_timing=True)
|
||||
end_cuda_event = torch.cuda.Event(enable_timing=True)
|
||||
start_cuda_event.record()
|
||||
start_time = time.time()
|
||||
|
||||
output_image = self.poser.pose(self.torch_source_image, pose, output_index)[0].detach().cpu()
|
||||
|
||||
end_time = time.time()
|
||||
end_cuda_event.record()
|
||||
torch.cuda.synchronize()
|
||||
print("cuda time (ms):", start_cuda_event.elapsed_time(end_cuda_event))
|
||||
print("elapsed time (ms):", (end_time - start_time) * 1000.0)
|
||||
|
||||
numpy_image = convert_output_image_from_torch_to_numpy(output_image)
|
||||
self.last_output_numpy_image = numpy_image
|
||||
wx_image = wx.ImageFromBuffer(
|
||||
numpy_image.shape[0],
|
||||
numpy_image.shape[1],
|
||||
numpy_image[:, :, 0:3].tobytes(),
|
||||
numpy_image[:, :, 3].tobytes())
|
||||
wx_bitmap = wx_image.ConvertToBitmap()
|
||||
|
||||
dc = wx.MemoryDC()
|
||||
dc.SelectObject(self.result_image_bitmap)
|
||||
dc.Clear()
|
||||
dc.DrawBitmap(wx_bitmap,
|
||||
(self.image_size - numpy_image.shape[0]) // 2,
|
||||
(self.image_size - numpy_image.shape[1]) // 2,
|
||||
True)
|
||||
del dc
|
||||
|
||||
self.Refresh()
|
||||
self.Update()
|
||||
|
||||
def on_save_image(self, event: wx.Event):
|
||||
if self.last_output_numpy_image is None:
|
||||
logging.info("There is no output image to save!!!")
|
||||
return
|
||||
|
||||
dir_name = "data/images"
|
||||
file_dialog = wx.FileDialog(self, "Choose an image", dir_name, "", "*.png", wx.FD_SAVE)
|
||||
if file_dialog.ShowModal() == wx.ID_OK:
|
||||
image_file_name = os.path.join(file_dialog.GetDirectory(), file_dialog.GetFilename())
|
||||
try:
|
||||
if os.path.exists(image_file_name):
|
||||
message_dialog = wx.MessageDialog(self, f"Override {image_file_name}", "Manual Poser",
|
||||
wx.YES_NO | wx.ICON_QUESTION)
|
||||
result = message_dialog.ShowModal()
|
||||
if result == wx.ID_YES:
|
||||
self.save_last_numpy_image(image_file_name)
|
||||
message_dialog.Destroy()
|
||||
else:
|
||||
self.save_last_numpy_image(image_file_name)
|
||||
except:
|
||||
message_dialog = wx.MessageDialog(self, f"Could not save {image_file_name}", "Manual Poser", wx.OK)
|
||||
message_dialog.ShowModal()
|
||||
message_dialog.Destroy()
|
||||
file_dialog.Destroy()
|
||||
|
||||
def save_last_numpy_image(self, image_file_name):
|
||||
numpy_image = self.last_output_numpy_image
|
||||
pil_image = PIL.Image.fromarray(numpy_image, mode='RGBA')
|
||||
os.makedirs(os.path.dirname(image_file_name), exist_ok=True)
|
||||
pil_image.save(image_file_name)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
device = torch.device('cuda:0')
|
||||
try:
|
||||
import tha4.poser.modes.mode_07
|
||||
|
||||
poser = tha4.poser.modes.mode_07.create_poser(device)
|
||||
except RuntimeError as e:
|
||||
print(e)
|
||||
sys.exit()
|
||||
|
||||
app = wx.App()
|
||||
main_frame = MainFrame(poser, device)
|
||||
main_frame.Show(True)
|
||||
main_frame.timer.Start(16)
|
||||
app.MainLoop()
|
||||
Whitespace-only changes.
@@ -0,0 +1,69 @@
|
||||
import json
|
||||
import os.path
|
||||
|
||||
import PIL.Image
|
||||
import torch
|
||||
from omegaconf import OmegaConf
|
||||
|
||||
from tha4.shion.base.image_util import extract_pytorch_image_from_PIL_image
|
||||
from tha4.poser.modes.mode_14 import create_poser, KEY_FACE_MORPHER, KEY_BODY_MORPHER
|
||||
|
||||
|
||||
class CharacterModel:
|
||||
def __init__(self,
|
||||
character_image_file_name: str,
|
||||
face_morpher_file_name: str,
|
||||
body_morpher_file_name: str):
|
||||
self.body_morpher_file_name = body_morpher_file_name
|
||||
self.face_morpher_file_name = face_morpher_file_name
|
||||
self.character_image_file_name = character_image_file_name
|
||||
self.poser = None
|
||||
self.character_image = None
|
||||
|
||||
def get_poser(self, device: torch.device):
|
||||
if self.poser is not None:
|
||||
self.poser.to(device)
|
||||
else:
|
||||
self.poser = create_poser(
|
||||
device,
|
||||
module_file_names={
|
||||
KEY_FACE_MORPHER: self.face_morpher_file_name,
|
||||
KEY_BODY_MORPHER: self.body_morpher_file_name
|
||||
})
|
||||
return self.poser
|
||||
|
||||
def get_character_image(self, device: torch.device):
|
||||
if self.character_image is None:
|
||||
pil_image = PIL.Image.open(self.character_image_file_name)
|
||||
if pil_image.mode != 'RGBA':
|
||||
raise RuntimeError("Character image is not an RGBA image!")
|
||||
self.character_image = extract_pytorch_image_from_PIL_image(pil_image)
|
||||
self.character_image = self.character_image.to(device)
|
||||
return self.character_image
|
||||
|
||||
def save(self, file_name: str):
|
||||
dir = os.path.dirname(file_name)
|
||||
rel_char_image_file_name = os.path.relpath(self.character_image_file_name, dir)
|
||||
rel_face_morpher_file_name = os.path.relpath(self.face_morpher_file_name, dir)
|
||||
rel_body_morpher_file_name = os.path.relpath(self.body_morpher_file_name, dir)
|
||||
data = {
|
||||
"character_image_file_name": rel_char_image_file_name,
|
||||
"face_morpher_file_name": rel_face_morpher_file_name,
|
||||
"body_morpher_file_name": rel_body_morpher_file_name,
|
||||
}
|
||||
conf = OmegaConf.create(data)
|
||||
os.makedirs(dir, exist_ok=True)
|
||||
with open(file_name, "wt") as fout:
|
||||
fout.write(OmegaConf.to_yaml(conf))
|
||||
|
||||
@staticmethod
|
||||
def load(file_name: str):
|
||||
conf = OmegaConf.to_container(OmegaConf.load(file_name))
|
||||
dir = os.path.dirname(file_name)
|
||||
character_image_file_name = os.path.join(dir, conf["character_image_file_name"])
|
||||
face_morpher_file_name = os.path.join(dir, conf["face_morpher_file_name"])
|
||||
body_morpher_file_name = os.path.join(dir, conf["body_morpher_file_name"])
|
||||
return CharacterModel(
|
||||
character_image_file_name,
|
||||
face_morpher_file_name,
|
||||
body_morpher_file_name)
|
||||
Whitespace-only changes.
@@ -0,0 +1,35 @@
|
||||
from typing import List, Callable
|
||||
|
||||
from torch import Tensor
|
||||
from torch.utils.data import Dataset
|
||||
|
||||
|
||||
class ImagePosesAndOtherImagesDataset(Dataset):
|
||||
def __init__(self,
|
||||
main_image_func: Callable[[], Tensor],
|
||||
pose_dataset: Dataset,
|
||||
other_image_funcs: List[Callable[[], Tensor]]):
|
||||
self.main_image_func = main_image_func
|
||||
self.other_image_funcs = other_image_funcs
|
||||
self.pose_dataset = pose_dataset
|
||||
self.main_image = None
|
||||
self.other_images = [None for i in range(len(self.other_image_funcs))]
|
||||
|
||||
def get_main_image(self):
|
||||
if self.main_image is None:
|
||||
self.main_image = self.main_image_func()
|
||||
return self.main_image
|
||||
|
||||
def get_other_image(self, image_index: int):
|
||||
if self.other_images[image_index] is None:
|
||||
self.other_images[image_index] = self.other_image_funcs[image_index]()
|
||||
return self.other_images[image_index]
|
||||
|
||||
def __len__(self):
|
||||
return len(self.pose_dataset)
|
||||
|
||||
def __getitem__(self, index):
|
||||
main_image = self.get_main_image()
|
||||
pose = self.pose_dataset[index][0]
|
||||
other_images = [self.get_other_image(i) for i in range(len(self.other_image_funcs))]
|
||||
return [main_image, pose] + other_images
|
||||
Whitespace-only changes.
@@ -0,0 +1,98 @@
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
from typing import Callable, List, Optional
|
||||
|
||||
from tha4.pytasuku.workspace import Workspace
|
||||
from tha4.shion.core.training.distrib.distributed_trainer import DistributedTrainer
|
||||
from tha4.shion.core.training.distrib.distributed_training_states import DistributedTrainingState
|
||||
|
||||
|
||||
def get_torchrun_executable():
|
||||
return os.path.dirname(sys.executable) + os.path.sep + "torchrun"
|
||||
|
||||
|
||||
class RdzvConfig:
|
||||
def __init__(self, id: int, port: int):
|
||||
self.port = port
|
||||
self.id = id
|
||||
|
||||
|
||||
def run_standalone_config_based_training_script(
|
||||
training_script_file_name: str,
|
||||
config_file_name: str,
|
||||
num_proc_per_node: int,
|
||||
target_checkpoint_examples: Optional[int] = None,
|
||||
rdzv_config: Optional[RdzvConfig] = None):
|
||||
command = f"{get_torchrun_executable()} " \
|
||||
f"--nnodes=1 " \
|
||||
f"--nproc_per_node={num_proc_per_node} "
|
||||
if rdzv_config is not None:
|
||||
command += f"--rdzv_endpoint=localhost:{rdzv_config.port} "
|
||||
command += "--rdzv_backend=c10d "
|
||||
command += f"--rdzv_id={rdzv_config.id} "
|
||||
else:
|
||||
command += "--standalone "
|
||||
command += f"{training_script_file_name} "
|
||||
if target_checkpoint_examples is not None:
|
||||
command += f"--target_checkpoint_examples {target_checkpoint_examples} "
|
||||
command += f"--config_file={config_file_name} "
|
||||
logging.info(f"Executing -- {command}")
|
||||
os.system(command)
|
||||
|
||||
|
||||
def define_standalone_config_based_training_tasks(
|
||||
workspace: Workspace,
|
||||
distributed_trainer_func: Callable[[], DistributedTrainer],
|
||||
training_script_file_name: str,
|
||||
config_file_name: str,
|
||||
num_proc_per_node: int,
|
||||
dependencies: Optional[List[str]] = None,
|
||||
rdzv_config: Optional[RdzvConfig] = None):
|
||||
trainer = distributed_trainer_func()
|
||||
checkpoint_examples = trainer.training_protocol.get_checkpoint_examples()
|
||||
assert len(checkpoint_examples) >= 1
|
||||
assert checkpoint_examples[0] > 0
|
||||
checkpoint_examples = [0] + checkpoint_examples
|
||||
|
||||
if dependencies is None:
|
||||
dependencies = []
|
||||
module_file_dependencies = dependencies[:]
|
||||
for module_name in trainer.pretrained_module_file_names:
|
||||
module_file_dependencies.append(trainer.pretrained_module_file_names[module_name])
|
||||
|
||||
def create_train_func(target_checkpoint_examples: int):
|
||||
return lambda: run_standalone_config_based_training_script(
|
||||
training_script_file_name,
|
||||
config_file_name,
|
||||
num_proc_per_node,
|
||||
target_checkpoint_examples,
|
||||
rdzv_config=rdzv_config)
|
||||
|
||||
train_tasks = []
|
||||
for checkpoint_index in range(0, len(checkpoint_examples)):
|
||||
for module_name in trainer.module_names:
|
||||
module_file_name = DistributedTrainingState.get_module_file_name(
|
||||
trainer.get_checkpoint_prefix(checkpoint_index),
|
||||
module_name)
|
||||
workspace.create_file_task(
|
||||
module_file_name,
|
||||
module_file_dependencies,
|
||||
create_train_func(trainer.checkpoint_examples[checkpoint_index]))
|
||||
for module_name in trainer.accumulators:
|
||||
accumulated_module_file_name = DistributedTrainingState.get_accumulated_module_file_name(
|
||||
trainer.get_checkpoint_prefix(checkpoint_index),
|
||||
module_name)
|
||||
workspace.create_file_task(
|
||||
accumulated_module_file_name,
|
||||
module_file_dependencies,
|
||||
create_train_func(checkpoint_examples[checkpoint_index]))
|
||||
workspace.create_command_task(
|
||||
trainer.get_checkpoint_prefix(checkpoint_index) + "/train_standalone",
|
||||
module_file_dependencies,
|
||||
create_train_func(checkpoint_examples[checkpoint_index]))
|
||||
train_tasks.append(trainer.get_checkpoint_prefix(checkpoint_index) + "/train_standlone")
|
||||
workspace.create_file_task(
|
||||
trainer.prefix + "/train_standalone",
|
||||
module_file_dependencies,
|
||||
create_train_func(checkpoint_examples[-1]))
|
||||
@@ -0,0 +1,16 @@
|
||||
import logging
|
||||
|
||||
from tha4.shion.core.training.distrib.distributed_trainer import DistributedTrainer
|
||||
from tha4.distiller.distiller_config import DistillerConfig
|
||||
|
||||
if __name__ == "__main__":
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
|
||||
parser = DistributedTrainer.get_default_arg_parser()
|
||||
parser.add_argument('--config_file', type=str)
|
||||
args = parser.parse_args()
|
||||
|
||||
config_file_name = args.config_file
|
||||
config = DistillerConfig.load(config_file_name)
|
||||
|
||||
DistributedTrainer.run_with_args(config.get_body_morpher_trainer, args)
|
||||
@@ -0,0 +1,16 @@
|
||||
import logging
|
||||
|
||||
from tha4.shion.core.training.distrib.distributed_trainer import DistributedTrainer
|
||||
from tha4.distiller.distiller_config import DistillerConfig
|
||||
|
||||
if __name__ == "__main__":
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
|
||||
parser = DistributedTrainer.get_default_arg_parser()
|
||||
parser.add_argument('--config_file', type=str)
|
||||
args = parser.parse_args()
|
||||
|
||||
config_file_name = args.config_file
|
||||
config = DistillerConfig.load(config_file_name)
|
||||
|
||||
DistributedTrainer.run_with_args(config.get_face_morpher_trainer, args)
|
||||
@@ -0,0 +1,310 @@
|
||||
import os.path
|
||||
import shutil
|
||||
import PIL.Image
|
||||
from dataclasses import dataclass
|
||||
from typing import Optional
|
||||
|
||||
from omegaconf import OmegaConf
|
||||
from tha4.charmodel.character_model import CharacterModel
|
||||
from tha4.pytasuku.workspace import Workspace, file_task
|
||||
from tha4.distiller.config_based_training_tasks import define_standalone_config_based_training_tasks
|
||||
from tha4.nn.siren.face_morpher.siren_face_morpher_00_trainer import SirenFaceMorpher00TrainerArgs
|
||||
from tha4.nn.siren.morpher.siren_morpher_03_trainer import SirenMorpher03TrainerArgs, TrainingPhases, TrainingPhase, \
|
||||
LossWeights, LossTerm
|
||||
from tha4.shion.base.image_util import pil_image_has_transparency
|
||||
|
||||
POSE_DATASET_FILE_NAME = 'data/pose_dataset.pt'
|
||||
|
||||
|
||||
def copy_file(source_file_name: str, dest_file_name):
|
||||
os.makedirs(os.path.dirname(dest_file_name), exist_ok=True)
|
||||
shutil.copyfile(source_file_name, dest_file_name)
|
||||
|
||||
|
||||
@dataclass
|
||||
class DistillerConfig:
|
||||
prefix: str
|
||||
character_image_file_name: str
|
||||
face_mask_image_file_name: str
|
||||
|
||||
face_morpher_random_seed_0: int = 12771885812175595441
|
||||
face_morpher_random_seed_1: int = 14367217090963479175
|
||||
face_morpher_num_training_examples_per_sample_output: Optional[int] = 10_000
|
||||
face_morpher_batch_size: int = 8
|
||||
|
||||
body_morpher_random_seed_0: int = 2892221210020292507
|
||||
body_morpher_random_seed_1: int = 9998918537095922080
|
||||
body_morpher_num_training_examples_per_sample_output: Optional[int] = 10_000
|
||||
body_morpher_batch_size: int = 8
|
||||
|
||||
num_cpu_workers: int = 1
|
||||
num_gpus: int = 1
|
||||
|
||||
def check(self):
|
||||
DistillerConfig.check_prefix(self.prefix)
|
||||
DistillerConfig.check_character_image_file_name(self.character_image_file_name)
|
||||
DistillerConfig.check_face_mask_image_file_name(self.face_mask_image_file_name)
|
||||
|
||||
DistillerConfig.check_num_cpu_workers(self.num_cpu_workers)
|
||||
DistillerConfig.check_num_gpus(self.num_gpus)
|
||||
|
||||
DistillerConfig.check_random_seed(self.face_morpher_random_seed_0, "face_morpher_random_seed_0")
|
||||
DistillerConfig.check_random_seed(self.face_morpher_random_seed_1, "face_morpher_random_seed_1")
|
||||
DistillerConfig.check_batch_size(self.face_morpher_batch_size, "face_morpher_batch_size")
|
||||
DistillerConfig.check_num_training_examples_per_sample_output(
|
||||
self.face_morpher_num_training_examples_per_sample_output,
|
||||
"face_morpher_num_training_examples_per_sample_output")
|
||||
|
||||
DistillerConfig.check_random_seed(self.body_morpher_random_seed_0, "body_morpher_random_seed_0")
|
||||
DistillerConfig.check_random_seed(self.body_morpher_random_seed_1, "body_morpher_random_seed_1")
|
||||
DistillerConfig.check_batch_size(self.body_morpher_batch_size, "body_morpher_batch_size")
|
||||
DistillerConfig.check_num_training_examples_per_sample_output(
|
||||
self.body_morpher_num_training_examples_per_sample_output,
|
||||
"body_morpher_num_training_examples_per_sample_output")
|
||||
|
||||
@staticmethod
|
||||
def check_prefix(prefix):
|
||||
assert os.path.isdir(prefix), "The 'prefix' must be a directory."
|
||||
assert os.path.exists(prefix), f"The {prefix} directory does not exist."
|
||||
|
||||
@staticmethod
|
||||
def check_character_image_file_name(file_name):
|
||||
_, ext = os.path.splitext(file_name)
|
||||
assert os.path.isfile(file_name), \
|
||||
f"The specified character image file name, {file_name}, does not point to a file."
|
||||
assert ext.lower() == ".png", "The character image file name must have extension '.png'."
|
||||
|
||||
image = PIL.Image.open(file_name)
|
||||
assert pil_image_has_transparency(image), "The character image must have an alpha channel."
|
||||
assert image.width == 512 and image.height == 512, "The character image must be 512x512."
|
||||
image.close()
|
||||
|
||||
@staticmethod
|
||||
def check_face_mask_image_file_name(file_name):
|
||||
_, ext = os.path.splitext(file_name)
|
||||
assert os.path.isfile(file_name), \
|
||||
f"The specified face mask image file name, {file_name}, does not point to a file."
|
||||
assert ext.lower() == ".png", "The face mask image file name must have extension '.png'."
|
||||
|
||||
image = PIL.Image.open(file_name)
|
||||
assert image.width == 512 and image.height == 512, "The face mask image must be 512x512."
|
||||
assert image.mode == "RGB", "The face mask image must be an RGB image."
|
||||
for x in range(512):
|
||||
for y in range(512):
|
||||
r, g, b = image.getpixel((x, y))
|
||||
assert (r == 0) or (r == 255), "The R channel of the face mask image must be 0 or 255"
|
||||
assert (g == 0) or (g == 255), "The G channel of the face mask image must be 0 or 255"
|
||||
assert (b == 0) or (b == 255), "The B channel of the face mask image must be 0 or 255"
|
||||
image.close()
|
||||
|
||||
@staticmethod
|
||||
def check_batch_size(value, field_name: str):
|
||||
assert isinstance(value, int), f"The {field_name} must be an integer."
|
||||
assert value >= 1, f"The {field_name} must be at least 1."
|
||||
assert value <= 8, f"The {field_name} must be at most 8."
|
||||
|
||||
@staticmethod
|
||||
def check_num_cpu_workers(value):
|
||||
assert value >= 1, "The value of 'num_cpu_workers must be at least 1."
|
||||
|
||||
@staticmethod
|
||||
def check_num_gpus(value):
|
||||
assert value >= 1, "The value of 'num_gpus' must be at least 1."
|
||||
|
||||
@staticmethod
|
||||
def check_random_seed(value, field_name: str):
|
||||
assert isinstance(value, int), f"The {field_name} must be an integer."
|
||||
assert value >= 0 and value <= 0x_ffff_ffff_ffff_ffff, "A random seed must be between 0 and 2**64-1."
|
||||
|
||||
@staticmethod
|
||||
def check_num_training_examples_per_sample_output(value, field_name):
|
||||
assert value in [10_000, 100_000, 1_000_000,
|
||||
None], f"The {field_name} must be 10_000, 100_00, 1_000_000_000, or None."
|
||||
|
||||
def save(self, file_name: str):
|
||||
conf = OmegaConf.structured(self)
|
||||
os.makedirs(self.prefix, exist_ok=True)
|
||||
with open(file_name, "wt") as fout:
|
||||
fout.write(OmegaConf.to_yaml(conf))
|
||||
|
||||
def config_yaml_file_name(self):
|
||||
return f"{self.prefix}/config.yaml"
|
||||
|
||||
def create_config_yaml_file(self):
|
||||
if os.path.exists(self.config_yaml_file_name()):
|
||||
return
|
||||
self.save(self.config_yaml_file_name())
|
||||
|
||||
@staticmethod
|
||||
def load(file_name: str) -> 'DistillerConfig':
|
||||
conf = OmegaConf.to_container(OmegaConf.load(file_name))
|
||||
args = DistillerConfig(**conf)
|
||||
args.check()
|
||||
return args
|
||||
|
||||
def face_morpher_prefix(self):
|
||||
return f"{self.prefix}/face_morpher"
|
||||
|
||||
def get_face_morpher_trainer(self, world_size: Optional[int] = None, backend: str = 'gloo'):
|
||||
if world_size is None:
|
||||
world_size = self.num_gpus
|
||||
args = SirenFaceMorpher00TrainerArgs(
|
||||
character_file_name=self.character_image_file_name,
|
||||
face_mask_file_name=self.face_mask_image_file_name,
|
||||
pose_dataset_file_name=POSE_DATASET_FILE_NAME,
|
||||
total_worker=self.num_cpu_workers,
|
||||
num_training_examples_per_sample_output=self.face_morpher_num_training_examples_per_sample_output,
|
||||
total_batch_size=self.face_morpher_batch_size,
|
||||
training_random_seed=self.face_morpher_random_seed_0,
|
||||
sample_output_random_seed=self.face_morpher_random_seed_1)
|
||||
return args.create_trainer(self.face_morpher_prefix(), world_size, backend)
|
||||
|
||||
def body_morpher_prefix(self):
|
||||
return f"{self.prefix}/body_morpher"
|
||||
|
||||
def get_body_morpher_trainer(self, world_size: Optional[int] = None, backend: str = 'gloo'):
|
||||
if world_size is None:
|
||||
world_size = self.num_gpus
|
||||
args = SirenMorpher03TrainerArgs(
|
||||
character_file_name=self.character_image_file_name,
|
||||
pose_dataset_file_name=POSE_DATASET_FILE_NAME,
|
||||
total_worker=self.num_cpu_workers,
|
||||
num_training_examples_per_sample_output=self.body_morpher_num_training_examples_per_sample_output,
|
||||
training_random_seed=self.body_morpher_random_seed_0,
|
||||
sample_output_random_seed=self.body_morpher_random_seed_1,
|
||||
total_batch_size=self.body_morpher_batch_size,
|
||||
sample_output_batch_size=1,
|
||||
training_phases=TrainingPhases([
|
||||
TrainingPhase(
|
||||
num_examples_upper_bound=200_000,
|
||||
learning_rate=1e-4,
|
||||
loss_weights=LossWeights(weights={
|
||||
LossTerm.full_blended: 0.25,
|
||||
LossTerm.full_warped: 0.25,
|
||||
LossTerm.full_grid_change: 0.5,
|
||||
LossTerm.full_color_change: 2.0,
|
||||
})),
|
||||
TrainingPhase(
|
||||
num_examples_upper_bound=400_000,
|
||||
learning_rate=3e-5,
|
||||
loss_weights=LossWeights(weights={
|
||||
LossTerm.full_blended: 0.25,
|
||||
LossTerm.full_warped: 0.25,
|
||||
LossTerm.full_grid_change: 0.5,
|
||||
LossTerm.full_color_change: 2.0,
|
||||
})),
|
||||
TrainingPhase(
|
||||
num_examples_upper_bound=600_000,
|
||||
learning_rate=3e-5,
|
||||
loss_weights=LossWeights(weights={
|
||||
LossTerm.full_blended: 1.0,
|
||||
LossTerm.full_warped: 2.5,
|
||||
LossTerm.full_grid_change: 5.0,
|
||||
LossTerm.full_color_change: 1.0,
|
||||
})),
|
||||
TrainingPhase(
|
||||
num_examples_upper_bound=800_000,
|
||||
learning_rate=1e-5,
|
||||
loss_weights=LossWeights(weights={
|
||||
LossTerm.full_blended: 1.0,
|
||||
LossTerm.full_warped: 2.5,
|
||||
LossTerm.full_grid_change: 5.0,
|
||||
LossTerm.full_color_change: 1.0,
|
||||
})),
|
||||
TrainingPhase(
|
||||
num_examples_upper_bound=1_300_000,
|
||||
learning_rate=1e-5,
|
||||
loss_weights=LossWeights(weights={
|
||||
LossTerm.full_blended: 10.0,
|
||||
LossTerm.full_warped: 1.0,
|
||||
LossTerm.full_grid_change: 1.0,
|
||||
LossTerm.full_color_change: 1.0,
|
||||
})),
|
||||
TrainingPhase(
|
||||
num_examples_upper_bound=1_500_000,
|
||||
learning_rate=3e-6,
|
||||
loss_weights=LossWeights(weights={
|
||||
LossTerm.full_blended: 10.0,
|
||||
LossTerm.full_warped: 1.0,
|
||||
LossTerm.full_grid_change: 1.0,
|
||||
LossTerm.full_color_change: 1.0,
|
||||
})),
|
||||
]))
|
||||
return args.create_trainer(self.body_morpher_prefix(), world_size, backend)
|
||||
|
||||
def character_model_prefix(self):
|
||||
return f"{self.prefix}/character_model"
|
||||
|
||||
def character_model_face_morpher_file_name(self):
|
||||
return f"{self.character_model_prefix()}/face_morpher.pt"
|
||||
|
||||
def character_model_body_morpher_file_name(self):
|
||||
return f"{self.character_model_prefix()}/body_morpher.pt"
|
||||
|
||||
def character_model_character_png_file_name(self):
|
||||
return f"{self.character_model_prefix()}/character.png"
|
||||
|
||||
def character_model_yaml_file_name(self):
|
||||
return f"{self.character_model_prefix()}/character_model.yaml"
|
||||
|
||||
def define_tasks(self, workspace: Workspace):
|
||||
workspace.create_file_task(self.config_yaml_file_name(), [], self.create_config_yaml_file)
|
||||
|
||||
define_standalone_config_based_training_tasks(
|
||||
workspace,
|
||||
self.get_face_morpher_trainer,
|
||||
"src/tha4/distiller/distill_face_morpher.py",
|
||||
self.config_yaml_file_name(),
|
||||
num_proc_per_node=self.num_gpus,
|
||||
dependencies=[
|
||||
self.config_yaml_file_name(),
|
||||
])
|
||||
|
||||
define_standalone_config_based_training_tasks(
|
||||
workspace,
|
||||
self.get_body_morpher_trainer,
|
||||
"src/tha4/distiller/distill_body_morpher.py",
|
||||
self.config_yaml_file_name(),
|
||||
num_proc_per_node=self.num_gpus,
|
||||
dependencies=[
|
||||
self.config_yaml_file_name(),
|
||||
])
|
||||
|
||||
@file_task(workspace, self.character_model_character_png_file_name(), [self.character_image_file_name])
|
||||
def copy_character_image_file_name():
|
||||
copy_file(self.character_image_file_name, self.character_model_character_png_file_name())
|
||||
|
||||
@file_task(workspace, self.character_model_face_morpher_file_name(), [
|
||||
f"{self.face_morpher_prefix()}/checkpoint/0010/module_module.pt",
|
||||
])
|
||||
def copy_face_morpher():
|
||||
copy_file(
|
||||
f"{self.face_morpher_prefix()}/checkpoint/0010/module_module.pt",
|
||||
self.character_model_face_morpher_file_name())
|
||||
|
||||
@file_task(workspace, self.character_model_body_morpher_file_name(), [
|
||||
f"{self.body_morpher_prefix()}/checkpoint/0015/module_module.pt",
|
||||
])
|
||||
def copy_face_morpher():
|
||||
copy_file(
|
||||
f"{self.body_morpher_prefix()}/checkpoint/0015/module_module.pt",
|
||||
self.character_model_body_morpher_file_name())
|
||||
|
||||
@file_task(workspace, self.character_model_yaml_file_name(), [])
|
||||
def create_character_model_yaml_file():
|
||||
character_model = CharacterModel(
|
||||
self.character_model_character_png_file_name(),
|
||||
self.character_model_face_morpher_file_name(),
|
||||
self.character_model_body_morpher_file_name())
|
||||
character_model.save(self.character_model_yaml_file_name())
|
||||
|
||||
workspace.create_command_task(
|
||||
f"{self.prefix}/all",
|
||||
[
|
||||
f"{self.face_morpher_prefix()}/train_standalone",
|
||||
f"{self.body_morpher_prefix()}/train_standalone",
|
||||
self.character_model_character_png_file_name(),
|
||||
self.character_model_face_morpher_file_name(),
|
||||
self.character_model_body_morpher_file_name(),
|
||||
self.character_model_yaml_file_name(),
|
||||
])
|
||||
Whitespace-only changes.
@@ -0,0 +1,142 @@
|
||||
import os.path
|
||||
from contextlib import contextmanager
|
||||
from pathlib import PurePath, Path
|
||||
from typing import Callable, Any, Optional
|
||||
|
||||
from tha4.distiller.distiller_config import DistillerConfig
|
||||
|
||||
|
||||
class DistillerConfigState:
|
||||
def __init__(self):
|
||||
self.config = DistillerConfig(prefix="", character_image_file_name="", face_mask_image_file_name="")
|
||||
self.last_saved_timestamp = None
|
||||
self.dirty = False
|
||||
|
||||
def load(self, file_name):
|
||||
self.config = DistillerConfig.load(file_name)
|
||||
if os.path.exists(self.config.config_yaml_file_name()):
|
||||
self.last_saved_timestamp = os.path.getmtime(self.config.config_yaml_file_name())
|
||||
else:
|
||||
self.last_saved_timestamp = None
|
||||
self.dirty = False
|
||||
|
||||
def need_to_check_overwrite(self):
|
||||
if self.last_saved_timestamp is None:
|
||||
return True
|
||||
if not os.path.exists(self.config.config_yaml_file_name()):
|
||||
return False
|
||||
if self.last_saved_timestamp < os.path.getmtime(self.config.config_yaml_file_name()):
|
||||
return True
|
||||
return False
|
||||
|
||||
def save(self):
|
||||
self.config.save(self.config.config_yaml_file_name())
|
||||
self.dirty = False
|
||||
self.last_saved_timestamp = os.path.getmtime(self.config.config_yaml_file_name())
|
||||
|
||||
@contextmanager
|
||||
def updating_value(self, value_func: Callable[[], Any]):
|
||||
old_value = value_func()
|
||||
yield
|
||||
new_value = value_func()
|
||||
if new_value != old_value:
|
||||
self.dirty = True
|
||||
|
||||
def set_prefix(self, new_value):
|
||||
with self.updating_value(lambda: self.config.prefix):
|
||||
new_relative_path = self.get_relative_path_to_cwd(
|
||||
new_value,
|
||||
"The prefix directory must be a subdirectory of the talking-head-anime-4-demo's source code directory.")
|
||||
DistillerConfig.check_prefix(new_relative_path)
|
||||
self.config.prefix = new_relative_path
|
||||
|
||||
def set_character_image_file_name(self, new_value):
|
||||
with self.updating_value(lambda: self.config.character_image_file_name):
|
||||
new_relative_path = self.get_relative_path_to_cwd(
|
||||
new_value,
|
||||
"The character image file must be under talking-head-anime-4-demo's source code directory.")
|
||||
DistillerConfig.check_character_image_file_name(new_relative_path)
|
||||
self.config.character_image_file_name = new_relative_path
|
||||
|
||||
def set_face_mask_image_file_name(self, new_value):
|
||||
with self.updating_value(lambda: self.config.face_mask_image_file_name):
|
||||
new_relative_path = self.get_relative_path_to_cwd(
|
||||
new_value,
|
||||
"The face mask image file must be under talking-head-anime-4-demo's source code directory.")
|
||||
DistillerConfig.check_face_mask_image_file_name(new_relative_path)
|
||||
self.config.face_mask_image_file_name = new_relative_path
|
||||
|
||||
def set_num_cpu_workers(self, new_value: int):
|
||||
with self.updating_value(lambda: self.config.num_cpu_workers):
|
||||
DistillerConfig.check_num_cpu_workers(new_value)
|
||||
self.config.num_cpu_workers = new_value
|
||||
|
||||
def set_num_gpus(self, new_value: int):
|
||||
with self.updating_value(lambda: self.config.num_gpus):
|
||||
DistillerConfig.check_num_cpu_workers(new_value)
|
||||
self.config.num_gpus = new_value
|
||||
|
||||
def set_face_morpher_random_seed_0(self, new_value: int):
|
||||
with self.updating_value(lambda: self.config.face_morpher_random_seed_0):
|
||||
DistillerConfig.check_random_seed(new_value, "face_morpher_random_seed_0")
|
||||
self.config.face_morpher_random_seed_0 = new_value
|
||||
|
||||
def set_face_morpher_random_seed_1(self, new_value: int):
|
||||
with self.updating_value(lambda: self.config.face_morpher_random_seed_1):
|
||||
DistillerConfig.check_random_seed(new_value, "face_morpher_random_seed_1")
|
||||
self.config.face_morpher_random_seed_1 = new_value
|
||||
|
||||
def set_face_morpher_num_training_examples_per_sample_output(self, new_value: Optional[int]):
|
||||
with self.updating_value(lambda: self.config.face_morpher_num_training_examples_per_sample_output):
|
||||
DistillerConfig.check_num_training_examples_per_sample_output(
|
||||
new_value, "face_morpher_num_training_examples_per_sample_output")
|
||||
self.config.face_morpher_num_training_examples_per_sample_output = new_value
|
||||
|
||||
def set_face_morpher_batch_size(self, new_value: int):
|
||||
with self.updating_value(lambda: self.config.face_morpher_batch_size):
|
||||
DistillerConfig.check_batch_size(new_value, "face_morpher_batch_size")
|
||||
self.config.face_morpher_batch_size = new_value
|
||||
|
||||
def set_body_morpher_random_seed_0(self, new_value: int):
|
||||
with self.updating_value(lambda: self.config.body_morpher_random_seed_0):
|
||||
DistillerConfig.check_random_seed(new_value, "body_morpher_random_seed_0")
|
||||
self.config.body_morpher_random_seed_0 = new_value
|
||||
|
||||
def set_body_morpher_random_seed_1(self, new_value: int):
|
||||
with self.updating_value(lambda: self.config.body_morpher_random_seed_1):
|
||||
DistillerConfig.check_random_seed(new_value, "body_morpher_random_seed_1")
|
||||
self.config.body_morpher_random_seed_1 = new_value
|
||||
|
||||
def set_body_morpher_num_training_examples_per_sample_output(self, new_value: Optional[int]):
|
||||
with self.updating_value(lambda: self.config.body_morpher_num_training_examples_per_sample_output):
|
||||
DistillerConfig.check_num_training_examples_per_sample_output(
|
||||
new_value, "body_morpher_num_training_examples_per_sample_output")
|
||||
self.config.body_morpher_num_training_examples_per_sample_output = new_value
|
||||
|
||||
def set_body_morpher_batch_size(self, new_value: int):
|
||||
with self.updating_value(lambda: self.config.body_morpher_batch_size):
|
||||
DistillerConfig.check_batch_size(new_value, "body_morpher_batch_size")
|
||||
self.config.body_morpher_batch_size = new_value
|
||||
|
||||
def get_relative_path_to_cwd(self, file_name: str, message: str):
|
||||
cwd = os.getcwd()
|
||||
assert os.path.commonprefix([cwd, file_name]) == cwd, message
|
||||
cwd_path = Path(cwd).as_posix()
|
||||
new_path = Path(file_name).as_posix()
|
||||
new_relative_path = os.path.relpath(str(new_path), cwd_path)
|
||||
new_relative_path = str(Path(new_relative_path).as_posix())
|
||||
return new_relative_path
|
||||
|
||||
def can_show_character_image(self):
|
||||
return os.path.isfile(self.config.character_image_file_name)
|
||||
|
||||
def can_show_face_mask_image(self):
|
||||
return os.path.isfile(self.config.face_mask_image_file_name)
|
||||
|
||||
def can_show_mask_on_face_image(self):
|
||||
return self.can_show_character_image() and self.can_show_face_mask_image()
|
||||
|
||||
def can_save(self):
|
||||
return os.path.isdir(self.config.prefix) \
|
||||
and os.path.isfile(self.config.character_image_file_name) \
|
||||
and os.path.isfile(self.config.face_mask_image_file_name)
|
||||
@@ -0,0 +1,761 @@
|
||||
import multiprocessing
|
||||
import random
|
||||
from contextlib import contextmanager
|
||||
from typing import Callable
|
||||
import PIL.Image
|
||||
|
||||
import torch
|
||||
import wx
|
||||
import wx.html
|
||||
import wx.lib.intctrl
|
||||
from tha4.distiller.ui.distiller_config_state import DistillerConfigState
|
||||
from tha4.image_util import convert_output_image_from_torch_to_numpy
|
||||
from tha4.shion.base.image_util import extract_pytorch_image_from_PIL_image
|
||||
|
||||
|
||||
def wx_bind_event(widget, evt):
|
||||
def f(handler):
|
||||
widget.Bind(evt, handler)
|
||||
return handler
|
||||
|
||||
return f
|
||||
|
||||
|
||||
class DistillerUiMainFrame(wx.Frame):
|
||||
PARAM_NAME_STATIC_TEXT_MIN_WIDTH = 400
|
||||
NUM_TRAINING_EXAMPLES_PER_SAMPLE_OUTPUT_CHOICES = [
|
||||
"10_000", "100_000", "1_000_000", "Do not generate sample outputs"]
|
||||
|
||||
def __init__(self):
|
||||
super().__init__(None, wx.ID_ANY, "Distiller UI")
|
||||
|
||||
self.init_ui()
|
||||
self.init_menus()
|
||||
self.init_bitmaps()
|
||||
self.Bind(wx.EVT_CLOSE, self.on_close)
|
||||
|
||||
self.state = DistillerConfigState()
|
||||
self.update_ui()
|
||||
|
||||
self.config_file_to_run = None
|
||||
|
||||
def init_ui(self):
|
||||
main_sizer = wx.BoxSizer(wx.HORIZONTAL)
|
||||
|
||||
self.SetSizer(main_sizer)
|
||||
self.SetAutoLayout(1)
|
||||
|
||||
left_panel = self.init_left_panel(self)
|
||||
main_sizer.Add(left_panel, 0, wx.FIXED_MINSIZE)
|
||||
|
||||
middle_panel = self.init_middle_panel(self)
|
||||
main_sizer.Add(middle_panel, 0, wx.EXPAND)
|
||||
|
||||
right_panel = self.init_right_panel(self)
|
||||
main_sizer.Add(right_panel, 1, wx.EXPAND)
|
||||
|
||||
main_sizer.Fit(self)
|
||||
|
||||
def init_menus(self):
|
||||
self.file_menu = wx.Menu()
|
||||
|
||||
self.new_menu_id = wx.Window.NewControlId()
|
||||
self.file_menu.Append(
|
||||
self.new_menu_id, item="&New\tCTRL+N", helpString="Create a new distiller configuration.")
|
||||
self.Bind(wx.EVT_MENU, self.on_new, id=self.new_menu_id)
|
||||
|
||||
self.open_menu_id = wx.Window.NewControlId()
|
||||
self.file_menu.Append(
|
||||
self.open_menu_id, item="&Open\tCTRL+O", helpString="Open a distiller confuguration.")
|
||||
self.Bind(wx.EVT_MENU, self.on_open, id=self.open_menu_id)
|
||||
|
||||
self.save_menu_id = wx.Window.NewControlId()
|
||||
self.save_menu_item = wx.MenuItem(
|
||||
self.file_menu, id=self.save_menu_id, text="&Save\tCTRL+S",
|
||||
helpString="Save the current distiller configuration. Error message will be shown it it is not well formed.")
|
||||
self.Bind(wx.EVT_MENU, self.on_save, id=self.save_menu_id)
|
||||
self.file_menu.Append(self.save_menu_item)
|
||||
|
||||
self.file_menu.AppendSeparator()
|
||||
|
||||
self.exit_menu_id = wx.ID_EXIT
|
||||
self.file_menu.Append(
|
||||
self.exit_menu_id, item="E&xit\tCTRL+Q", helpString="Exit the application.")
|
||||
self.Bind(wx.EVT_MENU, self.on_close, id=self.exit_menu_id)
|
||||
|
||||
self.menu_bar = wx.MenuBar()
|
||||
self.menu_bar.Append(self.file_menu, "&File")
|
||||
|
||||
self.SetMenuBar(self.menu_bar)
|
||||
|
||||
def init_bitmaps(self):
|
||||
self.face_image_bitmap = wx.Bitmap(128, 128)
|
||||
self.face_image_pytorch = None
|
||||
self.face_mask_image_bitmap = wx.Bitmap(128, 128)
|
||||
self.face_mask_image_pytorch = None
|
||||
self.mask_on_face_image_bitmap = wx.Bitmap(128, 128)
|
||||
self.draw_nothing_yet_string_to_bitmap(self.face_image_bitmap, 128, 128)
|
||||
self.draw_nothing_yet_string_to_bitmap(self.face_mask_image_bitmap, 128, 128)
|
||||
self.draw_nothing_yet_string_to_bitmap(self.mask_on_face_image_bitmap, 128, 128)
|
||||
|
||||
@contextmanager
|
||||
def create_panel(self, parent, sizer, *args, **kwargs):
|
||||
panel = wx.Panel(parent, *args, **kwargs)
|
||||
panel.SetSizer(sizer)
|
||||
panel.SetAutoLayout(1)
|
||||
|
||||
try:
|
||||
yield panel, sizer
|
||||
finally:
|
||||
sizer.Fit(panel)
|
||||
|
||||
def init_left_panel(self, parent):
|
||||
with self.create_panel(parent, wx.BoxSizer(wx.VERTICAL), style=wx.BORDER_SIMPLE) as (panel, sizer):
|
||||
self.face_image_panel = wx.Panel(panel, size=(128, 128), style=wx.SIMPLE_BORDER)
|
||||
self.face_image_panel.Bind(wx.EVT_PAINT, self.on_face_image_panel_paint)
|
||||
sizer.Add(self.face_image_panel, 0, wx.EXPAND)
|
||||
|
||||
static_text = wx.StaticText(panel, label="Face", style=wx.ALIGN_CENTER)
|
||||
sizer.Add(static_text, 0, wx.EXPAND)
|
||||
|
||||
self.face_mask_image_panel = wx.Panel(panel, size=(128, 128), style=wx.SIMPLE_BORDER)
|
||||
self.face_mask_image_panel.Bind(wx.EVT_PAINT, self.on_face_mask_image_panel_paint)
|
||||
sizer.Add(self.face_mask_image_panel, 0, wx.EXPAND)
|
||||
|
||||
static_text = wx.StaticText(panel, label="Face mask", style=wx.ALIGN_CENTER)
|
||||
sizer.Add(static_text, 0, wx.EXPAND)
|
||||
|
||||
self.mask_on_face_image_panel = wx.Panel(panel, size=(128, 128), style=wx.SIMPLE_BORDER)
|
||||
self.mask_on_face_image_panel.Bind(wx.EVT_PAINT, self.on_mask_on_face_image_panel_paint)
|
||||
sizer.Add(self.mask_on_face_image_panel, 0, wx.EXPAND)
|
||||
|
||||
static_text = wx.StaticText(panel, label="Mask upon face", style=wx.ALIGN_CENTER)
|
||||
sizer.Add(static_text, 0, wx.EXPAND)
|
||||
|
||||
return panel
|
||||
|
||||
def on_erase_background(self, event):
|
||||
pass
|
||||
|
||||
def on_face_image_panel_paint(self, event):
|
||||
wx.BufferedPaintDC(self.face_image_panel, self.face_image_bitmap)
|
||||
|
||||
def on_face_mask_image_panel_paint(self, event):
|
||||
wx.BufferedPaintDC(self.face_mask_image_panel, self.face_mask_image_bitmap)
|
||||
|
||||
def on_mask_on_face_image_panel_paint(self, event):
|
||||
wx.BufferedPaintDC(self.mask_on_face_image_panel, self.mask_on_face_image_bitmap)
|
||||
|
||||
def init_middle_panel(self, parent):
|
||||
with self.create_panel(parent, wx.BoxSizer(wx.VERTICAL), style=wx.BORDER_SIMPLE) as (panel, sizer):
|
||||
sizer.Add(self.init_prefix_panel(panel), 0, wx.EXPAND)
|
||||
sizer.Add(self.init_character_image_file_name_panel(panel), 0, wx.EXPAND)
|
||||
sizer.Add(self.init_face_mask_image_file_name_panel(panel), 0, wx.EXPAND)
|
||||
sizer.Add(self.init_num_cpu_workers_panel(panel), 0, wx.EXPAND)
|
||||
sizer.Add(self.init_num_gpus_panel(panel), 0, wx.EXPAND)
|
||||
sizer.Add(self.init_face_morpher_random_seed_0_panel(panel), 0, wx.EXPAND)
|
||||
sizer.Add(self.init_face_morpher_random_seed_1_panel(panel), 0, wx.EXPAND)
|
||||
sizer.Add(self.init_face_morpher_batch_size_panel(panel), 0, wx.EXPAND)
|
||||
sizer.Add(self.init_body_morpher_random_seed_0_panel(panel), 0, wx.EXPAND)
|
||||
sizer.Add(self.init_body_morpher_random_seed_1_panel(panel), 0, wx.EXPAND)
|
||||
sizer.Add(self.init_body_morpher_batch_size_panel(panel), 0, wx.EXPAND)
|
||||
sizer.Add(self.init_num_training_examples_per_sample_output_panel(panel), 0, wx.EXPAND)
|
||||
|
||||
self.run_button = wx.Button(panel, label="RUN")
|
||||
self.run_button.SetMinSize((-1, 64))
|
||||
self.run_button.Bind(wx.EVT_BUTTON, self.on_run)
|
||||
sizer.Add(self.run_button, 1, wx.EXPAND)
|
||||
|
||||
return panel
|
||||
|
||||
def init_prefix_panel(self, parent):
|
||||
with self.create_panel(parent, wx.BoxSizer(wx.VERTICAL), style=wx.BORDER_SIMPLE) as (panel, panel_sizer):
|
||||
prefix_param_name_panel = self.create_param_name_panel_with_help_button(
|
||||
panel,
|
||||
"prefix (i.e. project directory)",
|
||||
self.create_help_button_func("distiller-ui-doc/params/prefix.html"))
|
||||
panel_sizer.Add(prefix_param_name_panel, 1, wx.EXPAND)
|
||||
|
||||
with self.create_panel(panel, wx.BoxSizer(wx.HORIZONTAL), style=wx.BORDER_NONE) \
|
||||
as (prefix_panel, prefix_sizer):
|
||||
self.prefix_text_ctrl = wx.TextCtrl(prefix_panel, value="")
|
||||
self.prefix_text_ctrl.SetEditable(False)
|
||||
prefix_sizer.Add(self.prefix_text_ctrl, 1, wx.EXPAND)
|
||||
|
||||
self.prefix_change_button = wx.Button(prefix_panel, label="Change...")
|
||||
self.prefix_change_button.Bind(wx.EVT_BUTTON, self.on_prefix_change_button)
|
||||
prefix_sizer.Add(self.prefix_change_button, 0, wx.EXPAND)
|
||||
panel_sizer.Add(prefix_panel, 1, wx.EXPAND)
|
||||
|
||||
return panel
|
||||
|
||||
def on_prefix_change_button(self, event):
|
||||
dir_dialog = wx.DirDialog(self, "Choose a directory.", style=wx.DD_DEFAULT_STYLE | wx.DD_NEW_DIR_BUTTON)
|
||||
if dir_dialog.ShowModal() != wx.ID_OK:
|
||||
return
|
||||
prefix_value = dir_dialog.GetPath()
|
||||
try:
|
||||
self.state.set_prefix(prefix_value)
|
||||
self.update_ui()
|
||||
except Exception as e:
|
||||
message_dialog = wx.MessageDialog(self, str(e), "Error", wx.OK | wx.ICON_ERROR)
|
||||
message_dialog.ShowModal()
|
||||
|
||||
def init_character_image_file_name_panel(self, parent):
|
||||
with self.create_panel(parent, wx.BoxSizer(wx.VERTICAL), style=wx.BORDER_SIMPLE) as (panel, panel_sizer):
|
||||
prefix_param_name_panel = self.create_param_name_panel_with_help_button(
|
||||
panel,
|
||||
"character_image_file_name",
|
||||
self.create_help_button_func("distiller-ui-doc/params/character_image_file_name.html"))
|
||||
panel_sizer.Add(prefix_param_name_panel, 1, wx.EXPAND)
|
||||
|
||||
with self.create_panel(panel, wx.BoxSizer(wx.HORIZONTAL), style=wx.BORDER_NONE) as (sub_panel, sub_sizer):
|
||||
self.character_image_file_name_text_ctrl = wx.TextCtrl(sub_panel, value="")
|
||||
self.character_image_file_name_text_ctrl.SetEditable(False)
|
||||
sub_sizer.Add(self.character_image_file_name_text_ctrl, 1, wx.EXPAND)
|
||||
|
||||
self.character_image_change_button = wx.Button(sub_panel, label="Change...")
|
||||
self.character_image_change_button.Bind(wx.EVT_BUTTON, self.on_character_image_change_button)
|
||||
sub_sizer.Add(self.character_image_change_button, 0, wx.EXPAND)
|
||||
panel_sizer.Add(sub_panel, 1, wx.EXPAND)
|
||||
|
||||
return panel
|
||||
|
||||
def on_character_image_change_button(self, event):
|
||||
file_dialog = wx.FileDialog(self, "Choose a PNG file", wildcard="*.png", style=wx.FD_OPEN)
|
||||
if file_dialog.ShowModal() != wx.ID_OK:
|
||||
return
|
||||
file_name = file_dialog.GetPath()
|
||||
try:
|
||||
self.state.set_character_image_file_name(file_name)
|
||||
self.update_face_image_bitmap(file_name)
|
||||
self.update_ui()
|
||||
except Exception as e:
|
||||
message_dialog = wx.MessageDialog(self, str(e), "Error", wx.OK | wx.ICON_ERROR)
|
||||
message_dialog.ShowModal()
|
||||
|
||||
def update_face_image_bitmap(self, new_file_name: str):
|
||||
pil_image = PIL.Image.open(new_file_name)
|
||||
subimage = pil_image.crop((256 - 64, 80, 256 + 64, 208))
|
||||
self.face_image_bitmap = wx.Bitmap.FromBufferRGBA(128, 128, subimage.convert("RGBA").tobytes())
|
||||
self.face_image_pytorch = extract_pytorch_image_from_PIL_image(subimage).to(torch.float)
|
||||
self.update_mask_on_face_image_bitmap()
|
||||
|
||||
def init_face_mask_image_file_name_panel(self, parent):
|
||||
with self.create_panel(parent, wx.BoxSizer(wx.VERTICAL), style=wx.BORDER_SIMPLE) as (panel, panel_sizer):
|
||||
prefix_param_name_panel = self.create_param_name_panel_with_help_button(
|
||||
panel,
|
||||
"face_mask_image_file_name",
|
||||
self.create_help_button_func("distiller-ui-doc/params/face_mask_image_file_name.html"))
|
||||
panel_sizer.Add(prefix_param_name_panel, 1, wx.EXPAND)
|
||||
|
||||
with self.create_panel(panel, wx.BoxSizer(wx.HORIZONTAL), style=wx.BORDER_NONE) as (sub_panel, sub_sizer):
|
||||
self.face_mask_image_file_name_text_ctrl = wx.TextCtrl(sub_panel, value="")
|
||||
self.face_mask_image_file_name_text_ctrl.SetEditable(False)
|
||||
sub_sizer.Add(self.face_mask_image_file_name_text_ctrl, 1, wx.EXPAND)
|
||||
|
||||
self.face_mask_image_file_name_change_button = wx.Button(sub_panel, label="Change...")
|
||||
self.face_mask_image_file_name_change_button.Bind(wx.EVT_BUTTON, self.on_face_mask_image_change_button)
|
||||
sub_sizer.Add(self.face_mask_image_file_name_change_button, 0, wx.EXPAND)
|
||||
panel_sizer.Add(sub_panel, 1, wx.EXPAND)
|
||||
|
||||
return panel
|
||||
|
||||
def on_face_mask_image_change_button(self, event):
|
||||
file_dialog = wx.FileDialog(self, "Choose a PNG file", wildcard="*.png", style=wx.FD_OPEN)
|
||||
if file_dialog.ShowModal() != wx.ID_OK:
|
||||
return
|
||||
file_name = file_dialog.GetPath()
|
||||
try:
|
||||
self.state.set_face_mask_image_file_name(file_name)
|
||||
self.update_face_mask_image_bitmap(file_name)
|
||||
self.update_ui()
|
||||
except Exception as e:
|
||||
message_dialog = wx.MessageDialog(self, str(e), "Error", wx.OK | wx.ICON_ERROR)
|
||||
message_dialog.ShowModal()
|
||||
|
||||
def update_face_mask_image_bitmap(self, new_file_name):
|
||||
pil_image = PIL.Image.open(new_file_name)
|
||||
subimage = pil_image.crop((256 - 64, 80, 256 + 64, 208))
|
||||
self.face_mask_image_bitmap = wx.Bitmap.FromBufferRGBA(128, 128, subimage.convert("RGBA").tobytes())
|
||||
self.face_mask_image_pytorch = extract_pytorch_image_from_PIL_image(subimage).to(torch.float)
|
||||
self.face_mask_image_pytorch = self.face_mask_image_pytorch[0:1, :, :]
|
||||
self.update_mask_on_face_image_bitmap()
|
||||
|
||||
def update_mask_on_face_image_bitmap(self):
|
||||
if self.face_image_pytorch is None:
|
||||
return
|
||||
if self.face_mask_image_pytorch is None:
|
||||
return
|
||||
|
||||
mask_on_face_image = (0.5 * self.face_image_pytorch) + (0.5 * self.face_mask_image_pytorch)
|
||||
numpy_image = convert_output_image_from_torch_to_numpy(mask_on_face_image)
|
||||
wx_image = wx.ImageFromBuffer(
|
||||
numpy_image.shape[0],
|
||||
numpy_image.shape[1],
|
||||
numpy_image[:, :, 0:3].tobytes(),
|
||||
numpy_image[:, :, 3].tobytes())
|
||||
self.mask_on_face_image_bitmap = wx_image.ConvertToBitmap()
|
||||
|
||||
def init_num_cpu_workers_panel(self, parent):
|
||||
with self.create_panel(parent, wx.BoxSizer(wx.VERTICAL), style=wx.BORDER_SIMPLE) as (panel, panel_sizer):
|
||||
prefix_param_name_panel = self.create_param_name_panel_with_help_button(
|
||||
panel,
|
||||
"num_cpu_workers",
|
||||
self.create_help_button_func("distiller-ui-doc/params/num_cpu_workers.html"))
|
||||
panel_sizer.Add(prefix_param_name_panel, 1, wx.EXPAND)
|
||||
|
||||
num_cpus = multiprocessing.cpu_count()
|
||||
self.num_cpu_workers_spin_ctrl = wx.SpinCtrl(panel, initial=1, min=1, max=num_cpus)
|
||||
|
||||
@wx_bind_event(self.num_cpu_workers_spin_ctrl, wx.EVT_SPINCTRL)
|
||||
def on_num_cpu_workers_spin_ctrl(event):
|
||||
self.state.set_num_cpu_workers(self.num_cpu_workers_spin_ctrl.GetValue())
|
||||
self.Refresh()
|
||||
|
||||
panel_sizer.Add(self.num_cpu_workers_spin_ctrl, 1, wx.EXPAND)
|
||||
|
||||
return panel
|
||||
|
||||
def init_num_gpus_panel(self, parent):
|
||||
with self.create_panel(parent, wx.BoxSizer(wx.VERTICAL), style=wx.BORDER_SIMPLE) as (panel, panel_sizer):
|
||||
prefix_param_name_panel = self.create_param_name_panel_with_help_button(
|
||||
panel,
|
||||
"num_gpus",
|
||||
self.create_help_button_func("distiller-ui-doc/params/num_gpus.html"))
|
||||
panel_sizer.Add(prefix_param_name_panel, 1, wx.EXPAND)
|
||||
|
||||
num_gpus = torch.cuda.device_count()
|
||||
self.num_gpus_spin_ctrl = wx.SpinCtrl(panel, initial=1, min=1, max=max(1, num_gpus))
|
||||
|
||||
@wx_bind_event(self.num_gpus_spin_ctrl, wx.EVT_SPINCTRL)
|
||||
def on_num_cpu_workers_spin_ctrl(event):
|
||||
self.state.set_num_gpus(self.num_gpus_spin_ctrl.GetValue())
|
||||
self.Refresh()
|
||||
|
||||
panel_sizer.Add(self.num_gpus_spin_ctrl, 1, wx.EXPAND)
|
||||
|
||||
return panel
|
||||
|
||||
def init_face_morpher_random_seed_0_panel(self, parent):
|
||||
with self.create_panel(parent, wx.BoxSizer(wx.VERTICAL), style=wx.BORDER_SIMPLE) as (panel, panel_sizer):
|
||||
prefix_param_name_panel = self.create_param_name_panel_with_help_button(
|
||||
panel,
|
||||
"face_morpher_random_seed_0",
|
||||
self.create_help_button_func("distiller-ui-doc/params/face_morpher_random_seed_0.html"))
|
||||
panel_sizer.Add(prefix_param_name_panel, 1, wx.EXPAND)
|
||||
|
||||
with self.create_panel(panel, wx.BoxSizer(wx.HORIZONTAL), style=wx.BORDER_NONE) as (sub_panel, sub_sizer):
|
||||
initial_value = random.randint(0, 2 ** 64 - 1)
|
||||
self.face_morpher_random_seed_0_int_ctrl = wx.lib.intctrl.IntCtrl(
|
||||
sub_panel, value=initial_value, min=0, max=0x_ffff_ffff_ffff_ffff)
|
||||
|
||||
@wx_bind_event(self.face_morpher_random_seed_0_int_ctrl, wx.EVT_TEXT)
|
||||
def on_face_morpher_random_seed_0_int_ctrl_text(event):
|
||||
self.state.set_face_morpher_random_seed_0(self.face_morpher_random_seed_0_int_ctrl.GetValue())
|
||||
|
||||
sub_sizer.Add(self.face_morpher_random_seed_0_int_ctrl, 1, wx.EXPAND)
|
||||
|
||||
self.face_morpher_random_seed_0_randomize_button = wx.Button(sub_panel, label="Randomize")
|
||||
|
||||
@wx_bind_event(self.face_morpher_random_seed_0_randomize_button, wx.EVT_BUTTON)
|
||||
def on_face_morpher_random_seed_0_randomize_button(event):
|
||||
new_value = random.randint(0, 0x_ffff_ffff_ffff_ffff)
|
||||
self.face_morpher_random_seed_0_int_ctrl.SetValue(new_value)
|
||||
self.state.set_face_morpher_random_seed_0(new_value)
|
||||
|
||||
sub_sizer.Add(self.face_morpher_random_seed_0_randomize_button, 0, wx.EXPAND)
|
||||
panel_sizer.Add(sub_panel, 1, wx.EXPAND)
|
||||
|
||||
return panel
|
||||
|
||||
def init_face_morpher_random_seed_1_panel(self, parent):
|
||||
with self.create_panel(parent, wx.BoxSizer(wx.VERTICAL), style=wx.BORDER_SIMPLE) as (panel, panel_sizer):
|
||||
prefix_param_name_panel = self.create_param_name_panel_with_help_button(
|
||||
panel,
|
||||
"face_morpher_random_seed_1",
|
||||
self.create_help_button_func("distiller-ui-doc/params/face_morpher_random_seed_1.html"))
|
||||
panel_sizer.Add(prefix_param_name_panel, 1, wx.EXPAND)
|
||||
|
||||
with self.create_panel(panel, wx.BoxSizer(wx.HORIZONTAL), style=wx.BORDER_NONE) as (sub_panel, sub_sizer):
|
||||
initial_value = random.randint(0, 2 ** 64 - 1)
|
||||
self.face_morpher_random_seed_1_int_ctrl = wx.lib.intctrl.IntCtrl(
|
||||
sub_panel, value=initial_value, min=0, max=0x_ffff_ffff_ffff_ffff)
|
||||
|
||||
@wx_bind_event(self.face_morpher_random_seed_1_int_ctrl, wx.EVT_TEXT)
|
||||
def on_face_morpher_random_seed_1_int_ctrl_text(event):
|
||||
self.state.set_face_morpher_random_seed_1(self.face_morpher_random_seed_1_int_ctrl.GetValue())
|
||||
|
||||
sub_sizer.Add(self.face_morpher_random_seed_1_int_ctrl, 1, wx.EXPAND)
|
||||
|
||||
self.face_morpher_random_seed_1_randomize_button = wx.Button(sub_panel, label="Randomize")
|
||||
|
||||
@wx_bind_event(self.face_morpher_random_seed_1_randomize_button, wx.EVT_BUTTON)
|
||||
def on_face_morpher_random_seed_1_randomize_button(event):
|
||||
new_value = random.randint(0, 0x_ffff_ffff_ffff_ffff)
|
||||
self.face_morpher_random_seed_1_int_ctrl.SetValue(new_value)
|
||||
self.state.set_face_morpher_random_seed_1(new_value)
|
||||
|
||||
sub_sizer.Add(self.face_morpher_random_seed_1_randomize_button, 0, wx.EXPAND)
|
||||
panel_sizer.Add(sub_panel, 1, wx.EXPAND)
|
||||
|
||||
return panel
|
||||
|
||||
def init_face_morpher_batch_size_panel(self, parent):
|
||||
with self.create_panel(parent, wx.BoxSizer(wx.VERTICAL), style=wx.BORDER_SIMPLE) as (panel, panel_sizer):
|
||||
prefix_param_name_panel = self.create_param_name_panel_with_help_button(
|
||||
panel,
|
||||
"face_morpher_batch_size",
|
||||
self.create_help_button_func("distiller-ui-doc/params/face_morpher_batch_size.html"))
|
||||
panel_sizer.Add(prefix_param_name_panel, 1, wx.EXPAND)
|
||||
|
||||
self.face_morpher_batch_size_spin_ctrl = wx.SpinCtrl(panel, initial=8, min=1, max=8)
|
||||
|
||||
@wx_bind_event(self.face_morpher_batch_size_spin_ctrl, wx.EVT_SPINCTRL)
|
||||
def on_face_morpher_batch_size_spin_ctrl(event):
|
||||
self.state.set_face_morpher_batch_size(self.face_morpher_batch_size_spin_ctrl.GetValue())
|
||||
|
||||
panel_sizer.Add(self.face_morpher_batch_size_spin_ctrl, 1, wx.EXPAND)
|
||||
|
||||
return panel
|
||||
|
||||
def init_body_morpher_random_seed_0_panel(self, parent):
|
||||
with self.create_panel(parent, wx.BoxSizer(wx.VERTICAL), style=wx.BORDER_SIMPLE) as (panel, panel_sizer):
|
||||
prefix_param_name_panel = self.create_param_name_panel_with_help_button(
|
||||
panel,
|
||||
"body_morpher_random_seed_0",
|
||||
self.create_help_button_func("distiller-ui-doc/params/body_morpher_random_seed_0.html"))
|
||||
panel_sizer.Add(prefix_param_name_panel, 1, wx.EXPAND)
|
||||
|
||||
with self.create_panel(panel, wx.BoxSizer(wx.HORIZONTAL), style=wx.BORDER_NONE) as (sub_panel, sub_sizer):
|
||||
initial_value = random.randint(0, 2 ** 64 - 1)
|
||||
self.body_morpher_random_seed_0_int_ctrl = wx.lib.intctrl.IntCtrl(
|
||||
sub_panel, value=initial_value, min=0, max=0x_ffff_ffff_ffff_ffff)
|
||||
|
||||
@wx_bind_event(self.body_morpher_random_seed_0_int_ctrl, wx.EVT_TEXT)
|
||||
def on_body_morpher_random_seed_0_int_ctrl_text(event):
|
||||
self.state.set_body_morpher_random_seed_0(self.body_morpher_random_seed_0_int_ctrl.GetValue())
|
||||
|
||||
sub_sizer.Add(self.body_morpher_random_seed_0_int_ctrl, 1, wx.EXPAND)
|
||||
|
||||
self.body_morpher_random_seed_0_randomize_button = wx.Button(sub_panel, label="Randomize")
|
||||
|
||||
@wx_bind_event(self.body_morpher_random_seed_0_randomize_button, wx.EVT_BUTTON)
|
||||
def on_body_morpher_random_seed_0_randomize_button(event):
|
||||
new_value = random.randint(0, 0x_ffff_ffff_ffff_ffff)
|
||||
self.body_morpher_random_seed_0_int_ctrl.SetValue(new_value)
|
||||
self.state.set_body_morpher_random_seed_0(new_value)
|
||||
|
||||
sub_sizer.Add(self.body_morpher_random_seed_0_randomize_button, 0, wx.EXPAND)
|
||||
panel_sizer.Add(sub_panel, 1, wx.EXPAND)
|
||||
|
||||
return panel
|
||||
|
||||
def init_body_morpher_random_seed_1_panel(self, parent):
|
||||
with self.create_panel(parent, wx.BoxSizer(wx.VERTICAL), style=wx.BORDER_SIMPLE) as (panel, panel_sizer):
|
||||
prefix_param_name_panel = self.create_param_name_panel_with_help_button(
|
||||
panel,
|
||||
"body_morpher_random_seed_1",
|
||||
self.create_help_button_func("distiller-ui-doc/params/body_morpher_random_seed_1.html"))
|
||||
panel_sizer.Add(prefix_param_name_panel, 1, wx.EXPAND)
|
||||
|
||||
with self.create_panel(panel, wx.BoxSizer(wx.HORIZONTAL), style=wx.BORDER_NONE) as (sub_panel, sub_sizer):
|
||||
initial_value = random.randint(0, 2 ** 64 - 1)
|
||||
self.body_morpher_random_seed_1_int_ctrl = wx.lib.intctrl.IntCtrl(
|
||||
sub_panel, value=initial_value, min=0, max=0x_ffff_ffff_ffff_ffff)
|
||||
|
||||
@wx_bind_event(self.body_morpher_random_seed_1_int_ctrl, wx.EVT_TEXT)
|
||||
def on_body_morpher_random_seed_1_int_ctrl_text(event):
|
||||
self.state.set_body_morpher_random_seed_1(self.body_morpher_random_seed_1_int_ctrl.GetValue())
|
||||
|
||||
sub_sizer.Add(self.body_morpher_random_seed_1_int_ctrl, 1, wx.EXPAND)
|
||||
|
||||
self.body_morpher_random_seed_1_randomize_button = wx.Button(sub_panel, label="Randomize")
|
||||
|
||||
@wx_bind_event(self.body_morpher_random_seed_1_randomize_button, wx.EVT_BUTTON)
|
||||
def on_body_morpher_random_seed_1_randomize_button(event):
|
||||
new_value = random.randint(0, 0x_ffff_ffff_ffff_ffff)
|
||||
self.body_morpher_random_seed_1_int_ctrl.SetValue(new_value)
|
||||
self.state.set_body_morpher_random_seed_1(new_value)
|
||||
|
||||
sub_sizer.Add(self.body_morpher_random_seed_1_randomize_button, 0, wx.EXPAND)
|
||||
panel_sizer.Add(sub_panel, 1, wx.EXPAND)
|
||||
|
||||
return panel
|
||||
|
||||
def init_body_morpher_batch_size_panel(self, parent):
|
||||
with self.create_panel(parent, wx.BoxSizer(wx.VERTICAL), style=wx.BORDER_SIMPLE) as (panel, panel_sizer):
|
||||
prefix_param_name_panel = self.create_param_name_panel_with_help_button(
|
||||
panel,
|
||||
"body_morpher_batch_size",
|
||||
self.create_help_button_func("distiller-ui-doc/params/body_morpher_batch_size.html"))
|
||||
panel_sizer.Add(prefix_param_name_panel, 1, wx.EXPAND)
|
||||
|
||||
self.body_morpher_batch_size_spin_ctrl = wx.SpinCtrl(panel, initial=8, min=1, max=8)
|
||||
|
||||
@wx_bind_event(self.body_morpher_batch_size_spin_ctrl, wx.EVT_SPINCTRL)
|
||||
def on_body_morpher_batch_size_spin_ctrl(event):
|
||||
self.state.set_body_morpher_batch_size(self.body_morpher_batch_size_spin_ctrl.GetValue())
|
||||
|
||||
panel_sizer.Add(self.body_morpher_batch_size_spin_ctrl, 1, wx.EXPAND)
|
||||
|
||||
return panel
|
||||
|
||||
def init_num_training_examples_per_sample_output_panel(self, parent):
|
||||
with self.create_panel(parent, wx.BoxSizer(wx.VERTICAL), style=wx.BORDER_SIMPLE) as (panel, panel_sizer):
|
||||
prefix_param_name_panel = self.create_param_name_panel_with_help_button(
|
||||
panel,
|
||||
"num_training_examples_per_sample_output",
|
||||
self.create_help_button_func("distiller-ui-doc/params/num_training_examples_per_sample_output.html"))
|
||||
panel_sizer.Add(prefix_param_name_panel, 1, wx.EXPAND)
|
||||
|
||||
self.num_training_examples_per_sample_output_combobox = \
|
||||
wx.ComboBox(panel,
|
||||
value="10_000",
|
||||
choices=DistillerUiMainFrame.NUM_TRAINING_EXAMPLES_PER_SAMPLE_OUTPUT_CHOICES)
|
||||
|
||||
@wx_bind_event(self.num_training_examples_per_sample_output_combobox, wx.EVT_COMBOBOX)
|
||||
def on_num_training_examples_per_sample_output_combobox(event):
|
||||
index = self.num_training_examples_per_sample_output_combobox.GetSelection()
|
||||
if index == 3:
|
||||
self.state.set_face_morpher_num_training_examples_per_sample_output(None)
|
||||
self.state.set_body_morpher_num_training_examples_per_sample_output(None)
|
||||
else:
|
||||
selected = DistillerUiMainFrame.NUM_TRAINING_EXAMPLES_PER_SAMPLE_OUTPUT_CHOICES[index]
|
||||
new_value = int(selected)
|
||||
self.state.set_face_morpher_num_training_examples_per_sample_output(new_value)
|
||||
self.state.set_body_morpher_num_training_examples_per_sample_output(new_value)
|
||||
|
||||
panel_sizer.Add(self.num_training_examples_per_sample_output_combobox, 1, wx.EXPAND)
|
||||
|
||||
return panel
|
||||
|
||||
def on_close(self, event):
|
||||
if self.state.dirty:
|
||||
confirmation_dialog = wx.MessageDialog(
|
||||
parent=self,
|
||||
message=f"You have not saved your work. Do you want to exit anyway?",
|
||||
caption="Confirmation",
|
||||
style=wx.YES_NO | wx.ICON_QUESTION)
|
||||
result = confirmation_dialog.ShowModal()
|
||||
if result == wx.ID_NO:
|
||||
return
|
||||
|
||||
self.Destroy()
|
||||
|
||||
def create_help_button_func(self, html_file_name: str):
|
||||
def init_help_button_func(parent):
|
||||
button = wx.Button(parent, label="Help")
|
||||
|
||||
@wx_bind_event(button, wx.EVT_BUTTON)
|
||||
def on_prefix_button(event):
|
||||
self.html_window.LoadPage(html_file_name)
|
||||
self.Refresh()
|
||||
|
||||
return button
|
||||
|
||||
return init_help_button_func
|
||||
|
||||
def create_param_name_panel_with_help_button(
|
||||
self, parent, param_name: str, help_button_func: Callable[[wx.Window], wx.Button]):
|
||||
with self.create_panel(parent, wx.BoxSizer(wx.HORIZONTAL), style=wx.NO_BORDER) \
|
||||
as (panel, sizer):
|
||||
title_text_panel = self.create_vertically_centered_text_panel(
|
||||
panel, param_name, DistillerUiMainFrame.PARAM_NAME_STATIC_TEXT_MIN_WIDTH)
|
||||
sizer.Add(title_text_panel, 1, wx.EXPAND)
|
||||
|
||||
help_button = help_button_func(panel)
|
||||
sizer.Add(help_button, 0, wx.EXPAND)
|
||||
return panel
|
||||
|
||||
def create_vertically_centered_text_panel(self, parent, text: str, min_width: int):
|
||||
with self.create_panel(parent, wx.BoxSizer(wx.VERTICAL), style=wx.NO_BORDER) as (panel, sizer):
|
||||
sizer.AddStretchSpacer(1)
|
||||
text = wx.StaticText(
|
||||
panel,
|
||||
label=text,
|
||||
style=wx.ALIGN_CENTER)
|
||||
text.SetMinSize((min_width, -1))
|
||||
sizer.Add(text, 0, wx.EXPAND)
|
||||
sizer.AddStretchSpacer(1)
|
||||
return panel
|
||||
|
||||
def init_right_panel(self, parent):
|
||||
with self.create_panel(parent, wx.BoxSizer(wx.VERTICAL), style=wx.BORDER_SIMPLE) as (panel, sizer):
|
||||
self.html_window = wx.html.HtmlWindow(panel)
|
||||
self.html_window.SetMinSize((600, 600))
|
||||
self.html_window.SetFonts("Times New Roman", "Courier New", sizes=[10, 12, 14, 16, 18, 20, 24])
|
||||
self.html_window.LoadPage("distiller-ui-doc/index.html")
|
||||
sizer.Add(self.html_window, 1, wx.EXPAND)
|
||||
|
||||
go_to_main_documentation_button = wx.Button(panel, label="Go to Main Documentation")
|
||||
sizer.Add(go_to_main_documentation_button, 0, wx.EXPAND)
|
||||
|
||||
@wx_bind_event(go_to_main_documentation_button, wx.EVT_BUTTON)
|
||||
def on_go_to_main_documentation_button(event):
|
||||
self.html_window.LoadPage("distiller-ui-doc/index.html")
|
||||
self.Refresh()
|
||||
|
||||
return panel
|
||||
|
||||
def populate_distiller_config(self):
|
||||
self.state.config.prefix = self.prefix_text_ctrl.GetValue()
|
||||
self.state.config.character_image_file_name = self.character_image_file_name_text_ctrl.GetValue()
|
||||
self.state.config.face_mask_image_file_name = self.face_mask_image_file_name_text_ctrl.GetValue()
|
||||
|
||||
self.state.config.num_cpu_workers = self.num_cpu_workers_spin_ctrl.GetValue()
|
||||
self.state.config.num_gpus = self.num_gpus_spin_ctrl.GetValue()
|
||||
|
||||
self.state.config.face_morpher_random_seed_0 = self.face_morpher_random_seed_0_int_ctrl.GetValue()
|
||||
self.state.config.face_morpher_random_seed_1 = self.face_morpher_random_seed_1_int_ctrl.GetValue()
|
||||
self.state.config.face_morpher_batch_size = self.face_morpher_batch_size_spin_ctrl.GetValue()
|
||||
|
||||
self.state.config.body_morpher_random_seed_0 = self.body_morpher_random_seed_0_int_ctrl.GetValue()
|
||||
self.state.config.body_morpher_random_seed_1 = self.body_morpher_random_seed_1_int_ctrl.GetValue()
|
||||
self.state.config.body_morpher_batch_size = self.body_morpher_batch_size_spin_ctrl.GetValue()
|
||||
|
||||
if self.num_training_examples_per_sample_output_combobox.GetValue() == \
|
||||
DistillerUiMainFrame.NUM_TRAINING_EXAMPLES_PER_SAMPLE_OUTPUT_CHOICES[-1]:
|
||||
self.state.config.face_morpher_num_training_examples_per_sample_output = None
|
||||
self.state.config.body_morpher_num_training_examples_per_sample_output = None
|
||||
else:
|
||||
value = int(self.num_training_examples_per_sample_output_combobox.GetValue())
|
||||
self.state.config.face_morpher_num_training_examples_per_sample_output = value
|
||||
self.state.config.body_morpher_num_training_examples_per_sample_output = value
|
||||
|
||||
def update_ui(self):
|
||||
self.prefix_text_ctrl.SetValue(self.state.config.prefix)
|
||||
self.character_image_file_name_text_ctrl.SetValue(self.state.config.character_image_file_name)
|
||||
self.face_mask_image_file_name_text_ctrl.SetValue(self.state.config.face_mask_image_file_name)
|
||||
|
||||
if not self.state.can_show_character_image():
|
||||
self.draw_nothing_yet_string_to_bitmap(self.face_image_bitmap, 128, 128)
|
||||
if not self.state.can_show_face_mask_image():
|
||||
self.draw_nothing_yet_string_to_bitmap(self.face_mask_image_bitmap, 128, 128)
|
||||
if not self.state.can_show_mask_on_face_image():
|
||||
self.draw_nothing_yet_string_to_bitmap(self.mask_on_face_image_bitmap, 128, 128)
|
||||
|
||||
self.num_cpu_workers_spin_ctrl.SetValue(self.state.config.num_cpu_workers)
|
||||
self.num_gpus_spin_ctrl.SetValue(self.state.config.num_gpus)
|
||||
|
||||
self.face_morpher_random_seed_0_int_ctrl.SetValue(self.state.config.face_morpher_random_seed_0)
|
||||
self.face_morpher_random_seed_1_int_ctrl.SetValue(self.state.config.face_morpher_random_seed_1)
|
||||
self.face_morpher_batch_size_spin_ctrl.SetValue(self.state.config.face_morpher_batch_size)
|
||||
|
||||
self.body_morpher_random_seed_0_int_ctrl.SetValue(self.state.config.body_morpher_random_seed_0)
|
||||
self.body_morpher_random_seed_1_int_ctrl.SetValue(self.state.config.body_morpher_random_seed_1)
|
||||
self.body_morpher_batch_size_spin_ctrl.SetValue(self.state.config.body_morpher_batch_size)
|
||||
|
||||
if self.state.config.body_morpher_num_training_examples_per_sample_output is None:
|
||||
self.num_training_examples_per_sample_output_combobox.SetSelection(3)
|
||||
else:
|
||||
choices = [int(x) for x in DistillerUiMainFrame.NUM_TRAINING_EXAMPLES_PER_SAMPLE_OUTPUT_CHOICES[:-1]]
|
||||
self.num_training_examples_per_sample_output_combobox.SetSelection(
|
||||
choices.index(self.state.config.body_morpher_num_training_examples_per_sample_output))
|
||||
|
||||
self.save_menu_item.Enable(self.state.can_save())
|
||||
|
||||
self.Refresh()
|
||||
|
||||
def draw_nothing_yet_string_to_bitmap(self, bitmap, width: int, height: int):
|
||||
dc = wx.MemoryDC()
|
||||
dc.SelectObject(bitmap)
|
||||
|
||||
dc.Clear()
|
||||
font = wx.Font(wx.FontInfo(14).Family(wx.FONTFAMILY_SWISS))
|
||||
dc.SetFont(font)
|
||||
w, h = dc.GetTextExtent("Nothing yet!")
|
||||
dc.DrawText("Nothing yet!", (width - w) // 2, (height - h) // 2)
|
||||
|
||||
del dc
|
||||
|
||||
def try_saving(self):
|
||||
if not self.state.can_save():
|
||||
message_dialog = wx.MessageDialog(
|
||||
self,
|
||||
"Cannot save yet! Please make sure you set the prefix, character_image_file_name, "
|
||||
"and face_mask_image_file_name first.",
|
||||
"Error",
|
||||
wx.OK | wx.ICON_ERROR)
|
||||
message_dialog.ShowModal()
|
||||
return False
|
||||
else:
|
||||
if self.state.need_to_check_overwrite():
|
||||
confirmation_dialog = wx.MessageDialog(
|
||||
parent=self,
|
||||
message=f"Overwriting {self.state.config.config_yaml_file_name()}?",
|
||||
caption="Confirmation",
|
||||
style=wx.YES_NO | wx.CANCEL | wx.ICON_QUESTION)
|
||||
result = confirmation_dialog.ShowModal()
|
||||
if result == wx.ID_YES:
|
||||
self.state.save()
|
||||
return True
|
||||
elif result == wx.ID_NO:
|
||||
return False
|
||||
else:
|
||||
return False
|
||||
else:
|
||||
self.state.save()
|
||||
return True
|
||||
|
||||
def on_save(self, event):
|
||||
return self.try_saving()
|
||||
|
||||
def on_new(self, event):
|
||||
if self.state.dirty:
|
||||
confirmation_dialog = wx.MessageDialog(
|
||||
parent=self,
|
||||
message=f"You have not saved the current config. Do you want to proceed?",
|
||||
caption="Confirmation",
|
||||
style=wx.YES_NO | wx.ICON_QUESTION)
|
||||
result = confirmation_dialog.ShowModal()
|
||||
if result == wx.ID_NO:
|
||||
return
|
||||
self.state = DistillerConfigState()
|
||||
self.update_ui()
|
||||
|
||||
def on_open(self, event):
|
||||
if self.state.dirty:
|
||||
confirmation_dialog = wx.MessageDialog(
|
||||
parent=self,
|
||||
message=f"You have not saved the current config. Do you want to proceed?",
|
||||
caption="Confirmation",
|
||||
style=wx.YES_NO | wx.ICON_QUESTION)
|
||||
result = confirmation_dialog.ShowModal()
|
||||
if result == wx.ID_NO:
|
||||
return
|
||||
|
||||
file_dialog = wx.FileDialog(self, "Choose a YAML file", wildcard="*.yaml", style=wx.FD_OPEN)
|
||||
if file_dialog.ShowModal() != wx.ID_OK:
|
||||
return
|
||||
file_name = file_dialog.GetPath()
|
||||
try:
|
||||
self.state.load(file_name)
|
||||
self.face_image_pytorch = None
|
||||
self.face_mask_image_pytorch = None
|
||||
self.update_face_image_bitmap(self.state.config.character_image_file_name)
|
||||
self.update_face_mask_image_bitmap(self.state.config.face_mask_image_file_name)
|
||||
self.update_ui()
|
||||
except Exception as e:
|
||||
message_dialog = wx.MessageDialog(self, str(e), "Error", wx.OK | wx.ICON_ERROR)
|
||||
message_dialog.ShowModal()
|
||||
|
||||
def on_run(self, event):
|
||||
try:
|
||||
self.state.config.check()
|
||||
except Exception as e:
|
||||
message_dialog = wx.MessageDialog(self, str(e), "Error", wx.OK | wx.ICON_ERROR)
|
||||
message_dialog.ShowModal()
|
||||
return
|
||||
|
||||
if self.state.dirty:
|
||||
message_dialog = wx.MessageDialog(
|
||||
self,
|
||||
"Please save the configuration first.",
|
||||
"Error",
|
||||
wx.OK | wx.ICON_ERROR)
|
||||
message_dialog.ShowModal()
|
||||
return
|
||||
|
||||
self.config_file_to_run = self.state.config.config_yaml_file_name()
|
||||
self.Destroy()
|
||||
@@ -0,0 +1,58 @@
|
||||
import math
|
||||
|
||||
import PIL.Image
|
||||
import numpy
|
||||
import torch
|
||||
from matplotlib import cm
|
||||
from tha4.shion.base.image_util import numpy_linear_to_srgb, pytorch_rgba_to_numpy_image, pytorch_rgb_to_numpy_image, \
|
||||
torch_linear_to_srgb
|
||||
|
||||
|
||||
def grid_change_to_numpy_image(torch_image, num_channels=3):
|
||||
height = torch_image.shape[1]
|
||||
width = torch_image.shape[2]
|
||||
size_image = (torch_image[0, :, :] ** 2 + torch_image[1, :, :] ** 2).sqrt().view(height, width, 1).numpy()
|
||||
hsv = cm.get_cmap('hsv')
|
||||
angle_image = hsv(((torch.atan2(
|
||||
torch_image[0, :, :].view(height * width),
|
||||
torch_image[1, :, :].view(height * width)).view(height, width) + math.pi) / (2 * math.pi)).numpy()) * 3
|
||||
numpy_image = size_image * angle_image[:, :, 0:3]
|
||||
rgb_image = numpy_linear_to_srgb(numpy_image)
|
||||
if num_channels == 3:
|
||||
return rgb_image
|
||||
elif num_channels == 4:
|
||||
return numpy.concatenate([rgb_image, numpy.ones_like(size_image)], axis=2)
|
||||
else:
|
||||
raise RuntimeError("Unsupported num_channels: " + str(num_channels))
|
||||
|
||||
|
||||
def resize_PIL_image(pil_image, size=(256, 256)):
|
||||
w, h = pil_image.size
|
||||
d = min(w, h)
|
||||
r = ((w - d) // 2, (h - d) // 2, (w + d) // 2, (h + d) // 2)
|
||||
return pil_image.resize(size, resample=PIL.Image.LANCZOS, box=r)
|
||||
|
||||
|
||||
def convert_output_image_from_torch_to_numpy(output_image):
|
||||
if output_image.shape[2] == 2:
|
||||
h, w, c = output_image.shape
|
||||
numpy_image = torch.transpose(output_image.reshape(h * w, c), 0, 1).reshape(c, h, w)
|
||||
elif output_image.shape[0] == 4:
|
||||
numpy_image = pytorch_rgba_to_numpy_image(output_image)
|
||||
elif output_image.shape[0] == 3:
|
||||
numpy_image = pytorch_rgb_to_numpy_image(output_image)
|
||||
elif output_image.shape[0] == 1:
|
||||
c, h, w = output_image.shape
|
||||
alpha_image = torch.cat([output_image.repeat(3, 1, 1) * 2.0 - 1.0, torch.ones(1, h, w)], dim=0)
|
||||
numpy_image = pytorch_rgba_to_numpy_image(alpha_image)
|
||||
elif output_image.shape[0] == 2:
|
||||
numpy_image = grid_change_to_numpy_image(output_image, num_channels=4)
|
||||
else:
|
||||
raise RuntimeError("Unsupported # image channels: %d" % output_image.shape[0])
|
||||
numpy_image = numpy.uint8(numpy.rint(numpy_image * 255.0))
|
||||
return numpy_image
|
||||
|
||||
|
||||
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)
|
||||
Whitespace-only changes.
@@ -0,0 +1,239 @@
|
||||
EYE_LOOK_IN_LEFT = "eyeLookInLeft"
|
||||
EYE_LOOK_OUT_LEFT = "eyeLookOutLeft"
|
||||
EYE_LOOK_DOWN_LEFT = "eyeLookDownLeft"
|
||||
EYE_LOOK_UP_LEFT = "eyeLookUpLeft"
|
||||
EYE_BLINK_LEFT = "eyeBlinkLeft"
|
||||
EYE_SQUINT_LEFT = "eyeSquintLeft"
|
||||
EYE_WIDE_LEFT = "eyeWideLeft"
|
||||
EYE_LOOK_IN_RIGHT = "eyeLookInRight"
|
||||
EYE_LOOK_OUT_RIGHT = "eyeLookOutRight"
|
||||
EYE_LOOK_DOWN_RIGHT = "eyeLookDownRight"
|
||||
EYE_LOOK_UP_RIGHT = "eyeLookUpRight"
|
||||
EYE_BLINK_RIGHT = "eyeBlinkRight"
|
||||
EYE_SQUINT_RIGHT = "eyeSquintRight"
|
||||
EYE_WIDE_RIGHT = "eyeWideRight"
|
||||
BROW_DOWN_LEFT = "browDownLeft"
|
||||
BROW_OUTER_UP_LEFT = "browOuterUpLeft"
|
||||
BROW_DOWN_RIGHT = "browDownRight"
|
||||
BROW_OUTER_UP_RIGHT = "browOuterUpRight"
|
||||
BROW_INNER_UP = "browInnerUp"
|
||||
NOSE_SNEER_LEFT = "noseSneerLeft"
|
||||
NOSE_SNEER_RIGHT = "noseSneerRight"
|
||||
CHEEK_SQUINT_LEFT = "cheekSquintLeft"
|
||||
CHEEK_SQUINT_RIGHT = "cheekSquintRight"
|
||||
CHEEK_PUFF = "cheekPuff"
|
||||
MOUTH_LEFT = "mouthLeft"
|
||||
MOUTH_DIMPLE_LEFT = "mouthDimpleLeft"
|
||||
MOUTH_FROWN_LEFT = "mouthFrownLeft"
|
||||
MOUTH_LOWER_DOWN_LEFT = "mouthLowerDownLeft"
|
||||
MOUTH_PRESS_LEFT = "mouthPressLeft"
|
||||
MOUTH_SMILE_LEFT = "mouthSmileLeft"
|
||||
MOUTH_STRETCH_LEFT = "mouthStretchLeft"
|
||||
MOUTH_UPPER_UP_LEFT = "mouthUpperUpLeft"
|
||||
MOUTH_RIGHT = "mouthRight"
|
||||
MOUTH_DIMPLE_RIGHT = "mouthDimpleRight"
|
||||
MOUTH_FROWN_RIGHT = "mouthFrownRight"
|
||||
MOUTH_LOWER_DOWN_RIGHT = "mouthLowerDownRight"
|
||||
MOUTH_PRESS_RIGHT = "mouthPressRight"
|
||||
MOUTH_SMILE_RIGHT = "mouthSmileRight"
|
||||
MOUTH_STRETCH_RIGHT = "mouthStretchRight"
|
||||
MOUTH_UPPER_UP_RIGHT = "mouthUpperUpRight"
|
||||
MOUTH_CLOSE = "mouthClose"
|
||||
MOUTH_FUNNEL = "mouthFunnel"
|
||||
MOUTH_PUCKER = "mouthPucker"
|
||||
MOUTH_ROLL_LOWER = "mouthRollLower"
|
||||
MOUTH_ROLL_UPPER = "mouthRollUpper"
|
||||
MOUTH_SHRUG_LOWER = "mouthShrugLower"
|
||||
MOUTH_SHRUG_UPPER = "mouthShrugUpper"
|
||||
JAW_LEFT = "jawLeft"
|
||||
JAW_RIGHT = "jawRight"
|
||||
JAW_FORWARD = "jawForward"
|
||||
JAW_OPEN = "jawOpen"
|
||||
TONGUE_OUT = "tongueOut"
|
||||
|
||||
BLENDSHAPE_NAMES = [
|
||||
EYE_LOOK_IN_LEFT, # 0
|
||||
EYE_LOOK_OUT_LEFT, # 1
|
||||
EYE_LOOK_DOWN_LEFT, # 2
|
||||
EYE_LOOK_UP_LEFT, # 3
|
||||
EYE_BLINK_LEFT, # 4
|
||||
EYE_SQUINT_LEFT, # 5
|
||||
EYE_WIDE_LEFT, # 6
|
||||
EYE_LOOK_IN_RIGHT, # 7
|
||||
EYE_LOOK_OUT_RIGHT, # 8
|
||||
EYE_LOOK_DOWN_RIGHT, # 9
|
||||
EYE_LOOK_UP_RIGHT, # 10
|
||||
EYE_BLINK_RIGHT, # 11
|
||||
EYE_SQUINT_RIGHT, # 12
|
||||
EYE_WIDE_RIGHT, # 13
|
||||
BROW_DOWN_LEFT, # 14
|
||||
BROW_OUTER_UP_LEFT, # 15
|
||||
BROW_DOWN_RIGHT, # 16
|
||||
BROW_OUTER_UP_RIGHT, # 17
|
||||
BROW_INNER_UP, # 18
|
||||
NOSE_SNEER_LEFT, # 19
|
||||
NOSE_SNEER_RIGHT, # 20
|
||||
CHEEK_SQUINT_LEFT, # 21
|
||||
CHEEK_SQUINT_RIGHT, # 22
|
||||
CHEEK_PUFF, # 23
|
||||
MOUTH_LEFT, # 24
|
||||
MOUTH_DIMPLE_LEFT, # 25
|
||||
MOUTH_FROWN_LEFT, # 26
|
||||
MOUTH_LOWER_DOWN_LEFT, # 27
|
||||
MOUTH_PRESS_LEFT, # 28
|
||||
MOUTH_SMILE_LEFT, # 29
|
||||
MOUTH_STRETCH_LEFT, # 30
|
||||
MOUTH_UPPER_UP_LEFT, # 31
|
||||
MOUTH_RIGHT, # 32
|
||||
MOUTH_DIMPLE_RIGHT, # 33
|
||||
MOUTH_FROWN_RIGHT, # 34
|
||||
MOUTH_LOWER_DOWN_RIGHT, # 35
|
||||
MOUTH_PRESS_RIGHT, # 36
|
||||
MOUTH_SMILE_RIGHT, # 37
|
||||
MOUTH_STRETCH_RIGHT, # 38
|
||||
MOUTH_UPPER_UP_RIGHT, # 39
|
||||
MOUTH_CLOSE, # 40
|
||||
MOUTH_FUNNEL, # 41
|
||||
MOUTH_PUCKER, # 42
|
||||
MOUTH_ROLL_LOWER, # 43
|
||||
MOUTH_ROLL_UPPER, # 44
|
||||
MOUTH_SHRUG_LOWER, # 45
|
||||
MOUTH_SHRUG_UPPER, # 46
|
||||
JAW_LEFT, # 47
|
||||
JAW_RIGHT, # 48
|
||||
JAW_FORWARD, # 49
|
||||
JAW_OPEN, # 50
|
||||
TONGUE_OUT, # 51
|
||||
]
|
||||
|
||||
EYE_LEFT_BLENDSHAPES = [
|
||||
EYE_LOOK_IN_LEFT, # 0
|
||||
EYE_LOOK_OUT_LEFT, # 1
|
||||
EYE_LOOK_DOWN_LEFT, # 2
|
||||
EYE_LOOK_UP_LEFT, # 3
|
||||
EYE_BLINK_LEFT, # 4
|
||||
EYE_SQUINT_LEFT, # 5
|
||||
EYE_WIDE_LEFT, # 6
|
||||
]
|
||||
|
||||
EYE_RIGHT_BLENDSHAPES = [
|
||||
EYE_LOOK_IN_RIGHT, # 7
|
||||
EYE_LOOK_OUT_RIGHT, # 8
|
||||
EYE_LOOK_DOWN_RIGHT, # 9
|
||||
EYE_LOOK_UP_RIGHT, # 10
|
||||
EYE_BLINK_RIGHT, # 11
|
||||
EYE_SQUINT_RIGHT, # 12
|
||||
EYE_WIDE_RIGHT, # 13
|
||||
]
|
||||
|
||||
BROW_LEFT_BLENDSHAPES = [
|
||||
BROW_DOWN_LEFT, # 14
|
||||
BROW_OUTER_UP_LEFT, # 15
|
||||
|
||||
]
|
||||
|
||||
BROW_RIGHT_BLENDSHAPES = [
|
||||
BROW_DOWN_RIGHT, # 16
|
||||
BROW_OUTER_UP_RIGHT, # 17
|
||||
|
||||
]
|
||||
|
||||
BROW_BOTH_BLENDSHAPES = [
|
||||
BROW_INNER_UP, # 18
|
||||
]
|
||||
|
||||
NOSE_BLENDSHAPES = [
|
||||
NOSE_SNEER_LEFT, # 19
|
||||
NOSE_SNEER_RIGHT, # 20
|
||||
]
|
||||
|
||||
CHECK_BLENDSHAPES = [
|
||||
CHEEK_SQUINT_LEFT, # 21
|
||||
CHEEK_SQUINT_RIGHT, # 22
|
||||
CHEEK_PUFF, # 23
|
||||
]
|
||||
|
||||
MOUTH_LEFT_BLENDSHAPES = [
|
||||
MOUTH_LEFT, # 24
|
||||
MOUTH_DIMPLE_LEFT, # 25
|
||||
MOUTH_FROWN_LEFT, # 26
|
||||
MOUTH_LOWER_DOWN_LEFT, # 27
|
||||
MOUTH_PRESS_LEFT, # 28
|
||||
MOUTH_SMILE_LEFT, # 29
|
||||
MOUTH_STRETCH_LEFT, # 30
|
||||
MOUTH_UPPER_UP_LEFT, # 31
|
||||
]
|
||||
|
||||
MOUTH_RIGHT_BLENDSHAPES = [
|
||||
MOUTH_RIGHT, # 32
|
||||
MOUTH_DIMPLE_RIGHT, # 33
|
||||
MOUTH_FROWN_RIGHT, # 34
|
||||
MOUTH_LOWER_DOWN_RIGHT, # 35
|
||||
MOUTH_PRESS_RIGHT, # 36
|
||||
MOUTH_SMILE_RIGHT, # 37
|
||||
MOUTH_STRETCH_RIGHT, # 38
|
||||
MOUTH_UPPER_UP_RIGHT, # 39
|
||||
]
|
||||
|
||||
MOUTH_BOTH_BLENDSHAPES = [
|
||||
MOUTH_CLOSE, # 40
|
||||
MOUTH_FUNNEL, # 41
|
||||
MOUTH_PUCKER, # 42
|
||||
MOUTH_ROLL_LOWER, # 43
|
||||
MOUTH_ROLL_UPPER, # 44
|
||||
MOUTH_SHRUG_LOWER, # 45
|
||||
MOUTH_SHRUG_UPPER, # 46
|
||||
]
|
||||
|
||||
JAW_BLENDSHAPES = [
|
||||
JAW_LEFT, # 47
|
||||
JAW_RIGHT, # 48
|
||||
JAW_FORWARD, # 49
|
||||
JAW_OPEN, # 50
|
||||
]
|
||||
|
||||
TONGUE_BLENDSHAPES = [
|
||||
TONGUE_OUT, # 51
|
||||
]
|
||||
|
||||
COLUMN_0_BLENDSHAPES = EYE_RIGHT_BLENDSHAPES + BROW_RIGHT_BLENDSHAPES + [NOSE_SNEER_RIGHT, CHEEK_SQUINT_RIGHT]
|
||||
COLUMN_1_BLENDSHAPES = EYE_LEFT_BLENDSHAPES + BROW_LEFT_BLENDSHAPES + [NOSE_SNEER_LEFT, CHEEK_SQUINT_LEFT]
|
||||
COLUMN_2_BLENDSHAPES = MOUTH_RIGHT_BLENDSHAPES + [JAW_RIGHT]
|
||||
COLUMN_3_BLENDSHAPES = MOUTH_LEFT_BLENDSHAPES + [JAW_LEFT]
|
||||
COLUMN_4_BLENDSHAPES = [BROW_INNER_UP, CHEEK_PUFF] + MOUTH_BOTH_BLENDSHAPES + [JAW_FORWARD, JAW_OPEN, TONGUE_OUT]
|
||||
|
||||
BLENDSHAPE_COLUMNS = [
|
||||
COLUMN_0_BLENDSHAPES,
|
||||
COLUMN_1_BLENDSHAPES,
|
||||
COLUMN_2_BLENDSHAPES,
|
||||
COLUMN_3_BLENDSHAPES,
|
||||
COLUMN_4_BLENDSHAPES,
|
||||
]
|
||||
|
||||
RIGHT_EYE_BONE_X = "rightEyeBoneX"
|
||||
RIGHT_EYE_BONE_Y = "rightEyeBoneY"
|
||||
RIGHT_EYE_BONE_Z = "rightEyeBoneZ"
|
||||
RIGHT_EYE_BONE_ROTATIONS = [RIGHT_EYE_BONE_X, RIGHT_EYE_BONE_Y, RIGHT_EYE_BONE_Z]
|
||||
|
||||
LEFT_EYE_BONE_X = "leftEyeBoneX"
|
||||
LEFT_EYE_BONE_Y = "leftEyeBoneY"
|
||||
LEFT_EYE_BONE_Z = "leftEyeBoneZ"
|
||||
LEFT_EYE_BONE_ROTATIONS = [LEFT_EYE_BONE_X, LEFT_EYE_BONE_Y, LEFT_EYE_BONE_Z]
|
||||
|
||||
HEAD_BONE_X = "headBoneX"
|
||||
HEAD_BONE_Y = "headBoneY"
|
||||
HEAD_BONE_Z = "headBoneZ"
|
||||
HEAD_BONE_ROTATIONS = [HEAD_BONE_X, HEAD_BONE_Y, HEAD_BONE_Z]
|
||||
|
||||
ROTATION_NAMES = RIGHT_EYE_BONE_ROTATIONS + LEFT_EYE_BONE_ROTATIONS + HEAD_BONE_ROTATIONS
|
||||
|
||||
RIGHT_EYE_BONE_QUAT = "rightEyeBoneQuat"
|
||||
LEFT_EYE_BONE_QUAT = "leftEyeBoneQuat"
|
||||
HEAD_BONE_QUAT = "headBoneQuat"
|
||||
QUATERNION_NAMES = [
|
||||
RIGHT_EYE_BONE_QUAT,
|
||||
LEFT_EYE_BONE_QUAT,
|
||||
HEAD_BONE_QUAT
|
||||
]
|
||||
|
||||
IFACIALMOCAP_DATETIME_FORMAT = "%Y/%m/%d-%H:%M:%S.%f"
|
||||
@@ -0,0 +1,27 @@
|
||||
from tha4.mocap.ifacialmocap_constants import BLENDSHAPE_NAMES, HEAD_BONE_X, HEAD_BONE_Y, HEAD_BONE_Z, \
|
||||
HEAD_BONE_QUAT, LEFT_EYE_BONE_X, LEFT_EYE_BONE_Y, LEFT_EYE_BONE_Z, LEFT_EYE_BONE_QUAT, RIGHT_EYE_BONE_X, \
|
||||
RIGHT_EYE_BONE_Y, RIGHT_EYE_BONE_Z, RIGHT_EYE_BONE_QUAT
|
||||
|
||||
|
||||
def create_default_ifacialmocap_pose():
|
||||
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
|
||||
@@ -0,0 +1,12 @@
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Dict, List
|
||||
|
||||
|
||||
class IFacialMocapPoseConverter(ABC):
|
||||
@abstractmethod
|
||||
def convert(self, ifacialmocap_pose: Dict[str, float]) -> List[float]:
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def init_pose_converter_panel(self, parent):
|
||||
pass
|
||||
@@ -0,0 +1,614 @@
|
||||
import math
|
||||
import time
|
||||
from enum import Enum
|
||||
from typing import Optional, Dict, List, Callable
|
||||
|
||||
import numpy
|
||||
import scipy.optimize
|
||||
import wx
|
||||
|
||||
from tha4.mocap.ifacialmocap_constants import MOUTH_SMILE_LEFT, MOUTH_SHRUG_UPPER, MOUTH_SMILE_RIGHT, \
|
||||
BROW_INNER_UP, BROW_OUTER_UP_RIGHT, BROW_OUTER_UP_LEFT, BROW_DOWN_LEFT, BROW_DOWN_RIGHT, EYE_WIDE_LEFT, \
|
||||
EYE_WIDE_RIGHT, EYE_BLINK_LEFT, EYE_BLINK_RIGHT, CHEEK_SQUINT_LEFT, CHEEK_SQUINT_RIGHT, EYE_LOOK_IN_LEFT, \
|
||||
EYE_LOOK_OUT_LEFT, EYE_LOOK_IN_RIGHT, EYE_LOOK_OUT_RIGHT, EYE_LOOK_UP_LEFT, EYE_LOOK_UP_RIGHT, EYE_LOOK_DOWN_RIGHT, \
|
||||
EYE_LOOK_DOWN_LEFT, HEAD_BONE_X, HEAD_BONE_Y, HEAD_BONE_Z, JAW_OPEN, MOUTH_FROWN_LEFT, MOUTH_FROWN_RIGHT, \
|
||||
MOUTH_LOWER_DOWN_LEFT, MOUTH_LOWER_DOWN_RIGHT, MOUTH_FUNNEL, MOUTH_PUCKER
|
||||
from tha4.mocap.ifacialmocap_pose_converter import IFacialMocapPoseConverter
|
||||
from tha4.poser.modes.pose_parameters import get_pose_parameters
|
||||
|
||||
|
||||
class EyebrowDownMode(Enum):
|
||||
TROUBLED = 1
|
||||
ANGRY = 2
|
||||
LOWERED = 3
|
||||
SERIOUS = 4
|
||||
|
||||
|
||||
class WinkMode(Enum):
|
||||
NORMAL = 1
|
||||
RELAXED = 2
|
||||
|
||||
|
||||
def rad_to_deg(rad):
|
||||
return rad * 180.0 / math.pi
|
||||
|
||||
|
||||
def deg_to_rad(deg):
|
||||
return deg * math.pi / 180.0
|
||||
|
||||
|
||||
def clamp(x, min_value, max_value):
|
||||
return max(min_value, min(max_value, x))
|
||||
|
||||
|
||||
class IFacialMocapPoseConverter25Args:
|
||||
def __init__(self,
|
||||
smile_threshold_min: float = 0.4,
|
||||
smile_threshold_max: float = 0.6,
|
||||
eyebrow_down_mode: EyebrowDownMode = EyebrowDownMode.ANGRY,
|
||||
wink_mode: WinkMode = WinkMode.NORMAL,
|
||||
eye_surprised_max: float = 0.5,
|
||||
eye_blink_max: float = 0.8,
|
||||
eyebrow_down_max: float = 0.4,
|
||||
cheek_squint_min: float = 0.1,
|
||||
cheek_squint_max: float = 0.7,
|
||||
eye_rotation_factor: float = 1.0 / 0.75,
|
||||
jaw_open_min: float = 0.1,
|
||||
jaw_open_max: float = 0.4,
|
||||
mouth_frown_max: float = 0.6,
|
||||
mouth_funnel_min: float = 0.25,
|
||||
mouth_funnel_max: float = 0.5,
|
||||
iris_small_left=0.0,
|
||||
iris_small_right=0.0):
|
||||
self.iris_small_right = iris_small_left
|
||||
self.iris_small_left = iris_small_right
|
||||
|
||||
self.wink_mode = wink_mode
|
||||
|
||||
self.mouth_funnel_max = mouth_funnel_max
|
||||
self.mouth_funnel_min = mouth_funnel_min
|
||||
self.mouth_frown_max = mouth_frown_max
|
||||
|
||||
self.jaw_open_max = jaw_open_max
|
||||
self.jaw_open_min = jaw_open_min
|
||||
|
||||
self.eye_rotation_factor = eye_rotation_factor
|
||||
|
||||
self.cheek_squint_max = cheek_squint_max
|
||||
self.cheek_squint_min = cheek_squint_min
|
||||
|
||||
self.eyebrow_down_max = eyebrow_down_max
|
||||
|
||||
self.eye_blink_max = eye_blink_max
|
||||
self.eye_surprised_max = eye_surprised_max
|
||||
self.eyebrow_down_mode = eyebrow_down_mode
|
||||
|
||||
self.smile_threshold_min = smile_threshold_min
|
||||
self.smile_threshold_max = smile_threshold_max
|
||||
|
||||
def set_smile_threshold_min(self, new_value: float):
|
||||
self.smile_threshold_min = new_value
|
||||
|
||||
def set_smile_threshold_max(self, new_value: float):
|
||||
self.smile_threshold_max = new_value
|
||||
|
||||
def set_eye_surprised_max(self, new_value: float):
|
||||
self.eye_surprised_max = new_value
|
||||
|
||||
def set_eye_blink_max(self, new_value: float):
|
||||
self.eye_blink_max = new_value
|
||||
|
||||
def set_eyebrow_down_max(self, new_value: float):
|
||||
self.eyebrow_down_max = new_value
|
||||
|
||||
def set_cheek_squint_min(self, new_value: float):
|
||||
self.cheek_squint_min = new_value
|
||||
|
||||
def set_cheek_squint_max(self, new_value: float):
|
||||
self.cheek_squint_max = new_value
|
||||
|
||||
def set_jaw_open_min(self, new_value: float):
|
||||
self.jaw_open_min = new_value
|
||||
|
||||
def set_jaw_open_max(self, new_value: float):
|
||||
self.jaw_open_max = new_value
|
||||
|
||||
def set_mouth_frown_max(self, new_value: float):
|
||||
self.mouth_frown_max = new_value
|
||||
|
||||
def set_mouth_funnel_min(self, new_value: float):
|
||||
self.mouth_funnel_min = new_value
|
||||
|
||||
def set_mouth_funnel_max(self, new_value: float):
|
||||
self.mouth_funnel_min = new_value
|
||||
|
||||
|
||||
class IFacialMocapPoseConverter25(IFacialMocapPoseConverter):
|
||||
def __init__(self, args: Optional[IFacialMocapPoseConverter25Args] = None):
|
||||
super().__init__()
|
||||
if args is None:
|
||||
args = IFacialMocapPoseConverter25Args()
|
||||
self.args = args
|
||||
pose_parameters = get_pose_parameters()
|
||||
self.pose_size = 45
|
||||
|
||||
self.eyebrow_troubled_left_index = pose_parameters.get_parameter_index("eyebrow_troubled_left")
|
||||
self.eyebrow_troubled_right_index = pose_parameters.get_parameter_index("eyebrow_troubled_right")
|
||||
self.eyebrow_angry_left_index = pose_parameters.get_parameter_index("eyebrow_angry_left")
|
||||
self.eyebrow_angry_right_index = pose_parameters.get_parameter_index("eyebrow_angry_right")
|
||||
self.eyebrow_happy_left_index = pose_parameters.get_parameter_index("eyebrow_happy_left")
|
||||
self.eyebrow_happy_right_index = pose_parameters.get_parameter_index("eyebrow_happy_right")
|
||||
self.eyebrow_raised_left_index = pose_parameters.get_parameter_index("eyebrow_raised_left")
|
||||
self.eyebrow_raised_right_index = pose_parameters.get_parameter_index("eyebrow_raised_right")
|
||||
self.eyebrow_lowered_left_index = pose_parameters.get_parameter_index("eyebrow_lowered_left")
|
||||
self.eyebrow_lowered_right_index = pose_parameters.get_parameter_index("eyebrow_lowered_right")
|
||||
self.eyebrow_serious_left_index = pose_parameters.get_parameter_index("eyebrow_serious_left")
|
||||
self.eyebrow_serious_right_index = pose_parameters.get_parameter_index("eyebrow_serious_right")
|
||||
|
||||
self.eye_surprised_left_index = pose_parameters.get_parameter_index("eye_surprised_left")
|
||||
self.eye_surprised_right_index = pose_parameters.get_parameter_index("eye_surprised_right")
|
||||
self.eye_wink_left_index = pose_parameters.get_parameter_index("eye_wink_left")
|
||||
self.eye_wink_right_index = pose_parameters.get_parameter_index("eye_wink_right")
|
||||
self.eye_happy_wink_left_index = pose_parameters.get_parameter_index("eye_happy_wink_left")
|
||||
self.eye_happy_wink_right_index = pose_parameters.get_parameter_index("eye_happy_wink_right")
|
||||
self.eye_relaxed_left_index = pose_parameters.get_parameter_index("eye_relaxed_left")
|
||||
self.eye_relaxed_right_index = pose_parameters.get_parameter_index("eye_relaxed_right")
|
||||
self.eye_raised_lower_eyelid_left_index = pose_parameters.get_parameter_index("eye_raised_lower_eyelid_left")
|
||||
self.eye_raised_lower_eyelid_right_index = pose_parameters.get_parameter_index("eye_raised_lower_eyelid_right")
|
||||
|
||||
self.iris_small_left_index = pose_parameters.get_parameter_index("iris_small_left")
|
||||
self.iris_small_right_index = pose_parameters.get_parameter_index("iris_small_right")
|
||||
|
||||
self.iris_rotation_x_index = pose_parameters.get_parameter_index("iris_rotation_x")
|
||||
self.iris_rotation_y_index = pose_parameters.get_parameter_index("iris_rotation_y")
|
||||
|
||||
self.head_x_index = pose_parameters.get_parameter_index("head_x")
|
||||
self.head_y_index = pose_parameters.get_parameter_index("head_y")
|
||||
self.neck_z_index = pose_parameters.get_parameter_index("neck_z")
|
||||
|
||||
self.mouth_aaa_index = pose_parameters.get_parameter_index("mouth_aaa")
|
||||
self.mouth_iii_index = pose_parameters.get_parameter_index("mouth_iii")
|
||||
self.mouth_uuu_index = pose_parameters.get_parameter_index("mouth_uuu")
|
||||
self.mouth_eee_index = pose_parameters.get_parameter_index("mouth_eee")
|
||||
self.mouth_ooo_index = pose_parameters.get_parameter_index("mouth_ooo")
|
||||
|
||||
self.mouth_lowered_corner_left_index = pose_parameters.get_parameter_index("mouth_lowered_corner_left")
|
||||
self.mouth_lowered_corner_right_index = pose_parameters.get_parameter_index("mouth_lowered_corner_right")
|
||||
self.mouth_raised_corner_left_index = pose_parameters.get_parameter_index("mouth_raised_corner_left")
|
||||
self.mouth_raised_corner_right_index = pose_parameters.get_parameter_index("mouth_raised_corner_right")
|
||||
|
||||
self.body_y_index = pose_parameters.get_parameter_index("body_y")
|
||||
self.body_z_index = pose_parameters.get_parameter_index("body_z")
|
||||
self.breathing_index = pose_parameters.get_parameter_index("breathing")
|
||||
|
||||
self.breathing_start_time = time.time()
|
||||
|
||||
self.panel = None
|
||||
|
||||
def init_pose_converter_panel(self, parent):
|
||||
self.panel = wx.Panel(parent, style=wx.SIMPLE_BORDER)
|
||||
self.panel_sizer = wx.BoxSizer(wx.VERTICAL)
|
||||
self.panel.SetSizer(self.panel_sizer)
|
||||
self.panel.SetAutoLayout(1)
|
||||
parent.GetSizer().Add(self.panel, 0, wx.EXPAND)
|
||||
|
||||
if True:
|
||||
eyebrow_down_mode_text = wx.StaticText(self.panel, label=" --- Eyebrow Down Mode --- ",
|
||||
style=wx.ALIGN_CENTER)
|
||||
self.panel_sizer.Add(eyebrow_down_mode_text, 0, wx.EXPAND)
|
||||
|
||||
self.eyebrow_down_mode_choice = wx.Choice(
|
||||
self.panel,
|
||||
choices=[
|
||||
"ANGRY",
|
||||
"TROUBLED",
|
||||
"SERIOUS",
|
||||
"LOWERED",
|
||||
])
|
||||
self.eyebrow_down_mode_choice.SetSelection(0)
|
||||
self.panel_sizer.Add(self.eyebrow_down_mode_choice, 0, wx.EXPAND)
|
||||
self.eyebrow_down_mode_choice.Bind(wx.EVT_CHOICE, self.change_eyebrow_down_mode)
|
||||
|
||||
separator = wx.StaticLine(self.panel, -1, size=(256, 5))
|
||||
self.panel_sizer.Add(separator, 0, wx.EXPAND)
|
||||
|
||||
if True:
|
||||
wink_mode_text = wx.StaticText(self.panel, label=" --- Wink Mode --- ", style=wx.ALIGN_CENTER)
|
||||
self.panel_sizer.Add(wink_mode_text, 0, wx.EXPAND)
|
||||
|
||||
self.wink_mode_choice = wx.Choice(
|
||||
self.panel,
|
||||
choices=[
|
||||
"NORMAL",
|
||||
"RELAXED",
|
||||
])
|
||||
self.wink_mode_choice.SetSelection(0)
|
||||
self.panel_sizer.Add(self.wink_mode_choice, 0, wx.EXPAND)
|
||||
self.wink_mode_choice.Bind(wx.EVT_CHOICE, self.change_wink_mode)
|
||||
|
||||
separator = wx.StaticLine(self.panel, -1, size=(256, 5))
|
||||
self.panel_sizer.Add(separator, 0, wx.EXPAND)
|
||||
|
||||
if True:
|
||||
iris_size_text = wx.StaticText(self.panel, label=" --- Iris Size --- ", style=wx.ALIGN_CENTER)
|
||||
self.panel_sizer.Add(iris_size_text, 0, wx.EXPAND)
|
||||
|
||||
self.iris_left_slider = wx.Slider(self.panel, minValue=0, maxValue=1000, value=0, style=wx.HORIZONTAL)
|
||||
self.panel_sizer.Add(self.iris_left_slider, 0, wx.EXPAND)
|
||||
self.iris_left_slider.Bind(wx.EVT_SLIDER, self.change_iris_size)
|
||||
|
||||
self.iris_right_slider = wx.Slider(self.panel, minValue=0, maxValue=1000, value=0, style=wx.HORIZONTAL)
|
||||
self.panel_sizer.Add(self.iris_right_slider, 0, wx.EXPAND)
|
||||
self.iris_right_slider.Bind(wx.EVT_SLIDER, self.change_iris_size)
|
||||
self.iris_right_slider.Enable(False)
|
||||
|
||||
self.link_left_right_irises = wx.CheckBox(
|
||||
self.panel, label="Use same value for both sides")
|
||||
self.link_left_right_irises.SetValue(True)
|
||||
self.panel_sizer.Add(self.link_left_right_irises, wx.SizerFlags().CenterHorizontal().Border())
|
||||
self.link_left_right_irises.Bind(wx.EVT_CHECKBOX, self.link_left_right_irises_clicked)
|
||||
|
||||
separator = wx.StaticLine(self.panel, -1, size=(256, 5))
|
||||
self.panel_sizer.Add(separator, 0, wx.EXPAND)
|
||||
|
||||
if True:
|
||||
breathing_frequency_text = wx.StaticText(
|
||||
self.panel, label=" --- Breathing --- ", style=wx.ALIGN_CENTER)
|
||||
self.panel_sizer.Add(breathing_frequency_text, 0, wx.EXPAND)
|
||||
|
||||
self.restart_breathing_cycle_button = wx.Button(self.panel, label="Restart Breathing Cycle")
|
||||
self.restart_breathing_cycle_button.Bind(wx.EVT_BUTTON, self.restart_breathing_cycle_clicked)
|
||||
self.panel_sizer.Add(self.restart_breathing_cycle_button, 0, wx.EXPAND)
|
||||
|
||||
self.breathing_frequency_slider = wx.Slider(
|
||||
self.panel, minValue=0, maxValue=60, value=20, style=wx.HORIZONTAL)
|
||||
self.panel_sizer.Add(self.breathing_frequency_slider, 0, wx.EXPAND)
|
||||
|
||||
self.breathing_gauge = wx.Gauge(self.panel, style=wx.GA_HORIZONTAL, range=1000)
|
||||
self.panel_sizer.Add(self.breathing_gauge, 0, wx.EXPAND)
|
||||
|
||||
if True:
|
||||
separator = wx.StaticLine(self.panel, -1, size=(256, 5))
|
||||
self.panel_sizer.Add(separator, 0, wx.EXPAND)
|
||||
|
||||
convertion_parameters_text = wx.StaticText(
|
||||
self.panel, label="--- Conversion Parameters ---", style=wx.ALIGN_CENTER)
|
||||
self.panel_sizer.Add(convertion_parameters_text, 0, wx.EXPAND)
|
||||
|
||||
conversion_param_panel = wx.Panel(self.panel)
|
||||
self.panel_sizer.Add(conversion_param_panel, 0, wx.EXPAND)
|
||||
conversion_panel_sizer = wx.FlexGridSizer(cols=2)
|
||||
conversion_panel_sizer.AddGrowableCol(1)
|
||||
conversion_param_panel.SetSizer(conversion_panel_sizer)
|
||||
conversion_param_panel.SetAutoLayout(1)
|
||||
|
||||
self.smile_thresold_min_spin = self.create_spin_control(
|
||||
conversion_param_panel,
|
||||
"Smile Threshold Min:", self.args.smile_threshold_min, self.args.set_smile_threshold_min)
|
||||
self.smile_thresold_max_spin = self.create_spin_control(
|
||||
conversion_param_panel,
|
||||
"Smile Threshold Max:", self.args.smile_threshold_max, self.args.set_smile_threshold_max)
|
||||
self.eye_surprised_max_spin = self.create_spin_control(
|
||||
conversion_param_panel,
|
||||
"Eye Surprised Max:", self.args.eye_surprised_max, self.args.set_eye_surprised_max)
|
||||
self.eye_blink_max_spin = self.create_spin_control(
|
||||
conversion_param_panel,
|
||||
"Eye Blink Max:", self.args.eye_blink_max, self.args.set_eye_blink_max)
|
||||
self.eyebrow_down_max_spin = self.create_spin_control(
|
||||
conversion_param_panel,
|
||||
"Eyebrow Down Max:", self.args.eyebrow_down_max, self.args.set_eyebrow_down_max)
|
||||
self.cheek_squint_min_spin = self.create_spin_control(
|
||||
conversion_param_panel,
|
||||
"Cheek Squint Min:", self.args.cheek_squint_min, self.args.set_cheek_squint_min)
|
||||
self.cheek_squint_max_spin = self.create_spin_control(
|
||||
conversion_param_panel,
|
||||
"Cheek Squint Max:", self.args.cheek_squint_max, self.args.set_cheek_squint_max)
|
||||
self.jaw_open_min_spin = self.create_spin_control(
|
||||
conversion_param_panel,
|
||||
"Jaw Open Min:", self.args.jaw_open_min, self.args.set_jaw_open_min)
|
||||
self.jaw_open_max_spin = self.create_spin_control(
|
||||
conversion_param_panel,
|
||||
"Jaw Open Max:", self.args.jaw_open_max, self.args.set_jaw_open_max)
|
||||
self.mouth_frown_max_spin = self.create_spin_control(
|
||||
conversion_param_panel,
|
||||
"Mouth Frown Max:", self.args.mouth_frown_max, self.args.set_mouth_frown_max)
|
||||
self.mouth_funnel_min_spin = self.create_spin_control(
|
||||
conversion_param_panel,
|
||||
"Mouth Funnel Min:", self.args.mouth_funnel_min, self.args.set_mouth_funnel_min)
|
||||
self.mouth_funnel_max_spin = self.create_spin_control(
|
||||
conversion_param_panel,
|
||||
"Mouth Funnel Max:", self.args.mouth_funnel_max, self.args.set_mouth_funnel_max)
|
||||
|
||||
self.panel_sizer.Fit(self.panel)
|
||||
|
||||
def create_spin_control(self, parent, label: str, initial_value: float, set_func: Callable[[float], None]):
|
||||
sizer = parent.GetSizer()
|
||||
|
||||
text = wx.StaticText(parent, label=label)
|
||||
sizer.Add(text, wx.SizerFlags().Right().Border(wx.ALL, 2))
|
||||
|
||||
spin_ctrl = wx.SpinCtrlDouble(
|
||||
parent,
|
||||
wx.ID_ANY,
|
||||
min=0.0,
|
||||
max=1.0,
|
||||
initial=initial_value,
|
||||
inc=0.01)
|
||||
sizer.Add(spin_ctrl, wx.SizerFlags().Border(wx.ALL, 2).Expand())
|
||||
|
||||
def handler(event: wx.Event):
|
||||
new_value = spin_ctrl.GetValue()
|
||||
set_func(new_value)
|
||||
|
||||
spin_ctrl.Bind(wx.EVT_SPINCTRLDOUBLE, handler)
|
||||
|
||||
return spin_ctrl
|
||||
|
||||
def restart_breathing_cycle_clicked(self, event: wx.Event):
|
||||
self.breathing_start_time = time.time()
|
||||
|
||||
def change_eyebrow_down_mode(self, event: wx.Event):
|
||||
selected_index = self.eyebrow_down_mode_choice.GetSelection()
|
||||
if selected_index == 0:
|
||||
self.args.eyebrow_down_mode = EyebrowDownMode.ANGRY
|
||||
elif selected_index == 1:
|
||||
self.args.eyebrow_down_mode = EyebrowDownMode.TROUBLED
|
||||
elif selected_index == 2:
|
||||
self.args.eyebrow_down_mode = EyebrowDownMode.SERIOUS
|
||||
else:
|
||||
self.args.eyebrow_down_mode = EyebrowDownMode.LOWERED
|
||||
|
||||
def change_wink_mode(self, event: wx.Event):
|
||||
selected_index = self.wink_mode_choice.GetSelection()
|
||||
if selected_index == 0:
|
||||
self.args.wink_mode = WinkMode.NORMAL
|
||||
else:
|
||||
self.args.wink_mode = WinkMode.RELAXED
|
||||
|
||||
def change_iris_size(self, event: wx.Event):
|
||||
if self.link_left_right_irises.GetValue():
|
||||
left_value = self.iris_left_slider.GetValue()
|
||||
right_value = self.iris_right_slider.GetValue()
|
||||
if left_value != right_value:
|
||||
self.iris_right_slider.SetValue(left_value)
|
||||
self.args.iris_small_left = left_value / 1000.0
|
||||
self.args.iris_small_right = left_value / 1000.0
|
||||
else:
|
||||
self.args.iris_small_left = self.iris_left_slider.GetValue() / 1000.0
|
||||
self.args.iris_small_right = self.iris_right_slider.GetValue() / 1000.0
|
||||
|
||||
def link_left_right_irises_clicked(self, event: wx.Event):
|
||||
if self.link_left_right_irises.GetValue():
|
||||
self.iris_right_slider.Enable(False)
|
||||
else:
|
||||
self.iris_right_slider.Enable(True)
|
||||
self.change_iris_size(event)
|
||||
|
||||
def decompose_head_body_param(self, param, threshold=2.0 / 3):
|
||||
if abs(param) < threshold:
|
||||
return (param, 0.0)
|
||||
else:
|
||||
if param < 0:
|
||||
sign = -1.0
|
||||
else:
|
||||
sign = 1.0
|
||||
return (threshold * sign, (abs(param) - threshold) * sign)
|
||||
|
||||
def convert(self, ifacialmocap_pose: Dict[str, float]) -> List[float]:
|
||||
pose = [0.0 for i in range(self.pose_size)]
|
||||
|
||||
smile_value = \
|
||||
(ifacialmocap_pose[MOUTH_SMILE_LEFT] + ifacialmocap_pose[MOUTH_SMILE_RIGHT]) / 2.0 \
|
||||
+ ifacialmocap_pose[MOUTH_SHRUG_UPPER]
|
||||
if self.args.smile_threshold_min >= self.args.smile_threshold_max:
|
||||
smile_degree = 0.0
|
||||
else:
|
||||
if smile_value < self.args.smile_threshold_min:
|
||||
smile_degree = 0.0
|
||||
elif smile_value > self.args.smile_threshold_max:
|
||||
smile_degree = 1.0
|
||||
else:
|
||||
smile_degree = (smile_value - self.args.smile_threshold_min) / (
|
||||
self.args.smile_threshold_max - self.args.smile_threshold_min)
|
||||
|
||||
# Eyebrow
|
||||
if True:
|
||||
brow_inner_up = ifacialmocap_pose[BROW_INNER_UP]
|
||||
brow_outer_up_right = ifacialmocap_pose[BROW_OUTER_UP_RIGHT]
|
||||
brow_outer_up_left = ifacialmocap_pose[BROW_OUTER_UP_LEFT]
|
||||
|
||||
brow_up_left = clamp(brow_inner_up + brow_outer_up_left, 0.0, 1.0)
|
||||
brow_up_right = clamp(brow_inner_up + brow_outer_up_right, 0.0, 1.0)
|
||||
pose[self.eyebrow_raised_left_index] = brow_up_left
|
||||
pose[self.eyebrow_raised_right_index] = brow_up_right
|
||||
|
||||
if self.args.eyebrow_down_max <= 0.0:
|
||||
brow_down_left = 0.0
|
||||
brow_down_right = 0.0
|
||||
else:
|
||||
brow_down_left = (1.0 - smile_degree) \
|
||||
* clamp(ifacialmocap_pose[BROW_DOWN_LEFT] / self.args.eyebrow_down_max, 0.0, 1.0)
|
||||
brow_down_right = (1.0 - smile_degree) \
|
||||
* clamp(ifacialmocap_pose[BROW_DOWN_RIGHT] / self.args.eyebrow_down_max, 0.0, 1.0)
|
||||
|
||||
if self.args.eyebrow_down_mode == EyebrowDownMode.TROUBLED:
|
||||
pose[self.eyebrow_troubled_left_index] = brow_down_left
|
||||
pose[self.eyebrow_troubled_right_index] = brow_down_right
|
||||
elif self.args.eyebrow_down_mode == EyebrowDownMode.ANGRY:
|
||||
pose[self.eyebrow_angry_left_index] = brow_down_left
|
||||
pose[self.eyebrow_angry_right_index] = brow_down_right
|
||||
elif self.args.eyebrow_down_mode == EyebrowDownMode.LOWERED:
|
||||
pose[self.eyebrow_lowered_left_index] = brow_down_left
|
||||
pose[self.eyebrow_lowered_right_index] = brow_down_right
|
||||
elif self.args.eyebrow_down_mode == EyebrowDownMode.SERIOUS:
|
||||
pose[self.eyebrow_serious_left_index] = brow_down_left
|
||||
pose[self.eyebrow_serious_right_index] = brow_down_right
|
||||
|
||||
brow_happy_value = clamp(smile_value, 0.0, 1.0) * smile_degree
|
||||
pose[self.eyebrow_happy_left_index] = brow_happy_value
|
||||
pose[self.eyebrow_happy_right_index] = brow_happy_value
|
||||
|
||||
# Eye
|
||||
if True:
|
||||
# Surprised
|
||||
if self.args.eye_surprised_max <= 0.0:
|
||||
pose[self.eye_surprised_left_index] = 0.0
|
||||
pose[self.eye_surprised_right_index] = 0.0
|
||||
else:
|
||||
pose[self.eye_surprised_left_index] = clamp(
|
||||
ifacialmocap_pose[EYE_WIDE_LEFT] / self.args.eye_surprised_max, 0.0, 1.0)
|
||||
pose[self.eye_surprised_right_index] = clamp(
|
||||
ifacialmocap_pose[EYE_WIDE_RIGHT] / self.args.eye_surprised_max, 0.0, 1.0)
|
||||
|
||||
# Wink
|
||||
if self.args.wink_mode == WinkMode.NORMAL:
|
||||
wink_left_index = self.eye_wink_left_index
|
||||
wink_right_index = self.eye_wink_right_index
|
||||
else:
|
||||
wink_left_index = self.eye_relaxed_left_index
|
||||
wink_right_index = self.eye_relaxed_right_index
|
||||
if self.args.eye_blink_max <= 0:
|
||||
pose[wink_left_index] = 0.0
|
||||
pose[wink_right_index] = 0.0
|
||||
pose[self.eye_happy_wink_left_index] = 0.0
|
||||
pose[self.eye_happy_wink_right_index] = 0.0
|
||||
else:
|
||||
pose[wink_left_index] = (1.0 - smile_degree) * clamp(
|
||||
ifacialmocap_pose[EYE_BLINK_LEFT] / self.args.eye_blink_max, 0.0, 1.0)
|
||||
pose[wink_right_index] = (1.0 - smile_degree) * clamp(
|
||||
ifacialmocap_pose[EYE_BLINK_RIGHT] / self.args.eye_blink_max, 0.0, 1.0)
|
||||
pose[self.eye_happy_wink_left_index] = smile_degree * clamp(
|
||||
ifacialmocap_pose[EYE_BLINK_LEFT] / self.args.eye_blink_max, 0.0, 1.0)
|
||||
pose[self.eye_happy_wink_right_index] = smile_degree * clamp(
|
||||
ifacialmocap_pose[EYE_BLINK_RIGHT] / self.args.eye_blink_max, 0.0, 1.0)
|
||||
|
||||
# Lower eyelid
|
||||
cheek_squint_denom = self.args.cheek_squint_max - self.args.cheek_squint_min
|
||||
if cheek_squint_denom <= 0.0:
|
||||
pose[self.eye_raised_lower_eyelid_left_index] = 0.0
|
||||
pose[self.eye_raised_lower_eyelid_right_index] = 0.0
|
||||
else:
|
||||
pose[self.eye_raised_lower_eyelid_left_index] = \
|
||||
clamp(
|
||||
(ifacialmocap_pose[CHEEK_SQUINT_LEFT] - self.args.cheek_squint_min) / cheek_squint_denom,
|
||||
0.0, 1.0)
|
||||
pose[self.eye_raised_lower_eyelid_right_index] = \
|
||||
clamp(
|
||||
(ifacialmocap_pose[CHEEK_SQUINT_RIGHT] - self.args.cheek_squint_min) / cheek_squint_denom,
|
||||
0.0, 1.0)
|
||||
|
||||
# Iris rotation
|
||||
if True:
|
||||
eye_rotation_y = (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 * self.args.eye_rotation_factor
|
||||
pose[self.iris_rotation_y_index] = clamp(eye_rotation_y, -1.0, 1.0)
|
||||
|
||||
eye_rotation_x = (ifacialmocap_pose[EYE_LOOK_UP_LEFT]
|
||||
+ ifacialmocap_pose[EYE_LOOK_UP_RIGHT]
|
||||
- ifacialmocap_pose[EYE_LOOK_DOWN_LEFT]
|
||||
- ifacialmocap_pose[EYE_LOOK_DOWN_RIGHT]) / 2.0 * self.args.eye_rotation_factor
|
||||
pose[self.iris_rotation_x_index] = clamp(eye_rotation_x, -1.0, 1.0)
|
||||
|
||||
# Iris size
|
||||
if True:
|
||||
pose[self.iris_small_left_index] = self.args.iris_small_left
|
||||
pose[self.iris_small_right_index] = self.args.iris_small_right
|
||||
|
||||
# Head rotation
|
||||
if True:
|
||||
x_param = clamp(-ifacialmocap_pose[HEAD_BONE_X] * 180.0 / math.pi, -15.0, 15.0) / 15.0
|
||||
pose[self.head_x_index] = x_param
|
||||
|
||||
y_param = clamp(-ifacialmocap_pose[HEAD_BONE_Y] * 180.0 / math.pi, -10.0, 10.0) / 10.0
|
||||
pose[self.head_y_index] = y_param
|
||||
pose[self.body_y_index] = y_param
|
||||
|
||||
z_param = clamp(ifacialmocap_pose[HEAD_BONE_Z] * 180.0 / math.pi, -15.0, 15.0) / 15.0
|
||||
pose[self.neck_z_index] = z_param
|
||||
pose[self.body_z_index] = z_param
|
||||
|
||||
# Mouth
|
||||
if True:
|
||||
jaw_open_denom = self.args.jaw_open_max - self.args.jaw_open_min
|
||||
if jaw_open_denom <= 0:
|
||||
mouth_open = 0.0
|
||||
else:
|
||||
mouth_open = clamp((ifacialmocap_pose[JAW_OPEN] - self.args.jaw_open_min) / jaw_open_denom, 0.0, 1.0)
|
||||
pose[self.mouth_aaa_index] = mouth_open
|
||||
pose[self.mouth_raised_corner_left_index] = clamp(smile_value, 0.0, 1.0)
|
||||
pose[self.mouth_raised_corner_right_index] = clamp(smile_value, 0.0, 1.0)
|
||||
|
||||
is_mouth_open = mouth_open > 0.0
|
||||
if not is_mouth_open:
|
||||
if self.args.mouth_frown_max > 0:
|
||||
mouth_frown_value = 0.0
|
||||
else:
|
||||
mouth_frown_value = clamp(
|
||||
(ifacialmocap_pose[MOUTH_FROWN_LEFT] + ifacialmocap_pose[
|
||||
MOUTH_FROWN_RIGHT]) / self.args.mouth_frown_max, 0.0, 1.0)
|
||||
pose[self.mouth_lowered_corner_left_index] = mouth_frown_value
|
||||
pose[self.mouth_lowered_corner_right_index] = mouth_frown_value
|
||||
else:
|
||||
mouth_lower_down = clamp(
|
||||
ifacialmocap_pose[MOUTH_LOWER_DOWN_LEFT] + ifacialmocap_pose[MOUTH_LOWER_DOWN_RIGHT], 0.0, 1.0)
|
||||
mouth_funnel = ifacialmocap_pose[MOUTH_FUNNEL]
|
||||
mouth_pucker = ifacialmocap_pose[MOUTH_PUCKER]
|
||||
|
||||
mouth_point = [mouth_open, mouth_lower_down, mouth_funnel, mouth_pucker]
|
||||
|
||||
aaa_point = [1.0, 1.0, 0.0, 0.0]
|
||||
iii_point = [0.0, 1.0, 0.0, 0.0]
|
||||
uuu_point = [0.5, 0.3, 0.25, 0.75]
|
||||
ooo_point = [1.0, 0.5, 0.5, 0.4]
|
||||
|
||||
decomp = numpy.array([0, 0, 0, 0])
|
||||
M = numpy.array([
|
||||
aaa_point,
|
||||
iii_point,
|
||||
uuu_point,
|
||||
ooo_point
|
||||
])
|
||||
|
||||
def loss(decomp):
|
||||
return numpy.linalg.norm(numpy.matmul(decomp, M) - mouth_point) \
|
||||
+ 0.01 * numpy.linalg.norm(decomp, ord=1)
|
||||
|
||||
opt_result = scipy.optimize.minimize(
|
||||
loss, decomp, bounds=[(0.0, 1.0), (0.0, 1.0), (0.0, 1.0), (0.0, 1.0)])
|
||||
decomp = opt_result["x"]
|
||||
restricted_decomp = [decomp.item(0), decomp.item(1), decomp.item(2), decomp.item(3)]
|
||||
pose[self.mouth_aaa_index] = restricted_decomp[0]
|
||||
pose[self.mouth_iii_index] = restricted_decomp[1]
|
||||
mouth_funnel_denom = self.args.mouth_funnel_max - self.args.mouth_funnel_min
|
||||
if mouth_funnel_denom <= 0:
|
||||
ooo_alpha = 0.0
|
||||
uo_value = 0.0
|
||||
else:
|
||||
ooo_alpha = clamp((mouth_funnel - self.args.mouth_funnel_min) / mouth_funnel_denom, 0.0, 1.0)
|
||||
uo_value = clamp(restricted_decomp[2] + restricted_decomp[3], 0.0, 1.0)
|
||||
pose[self.mouth_uuu_index] = uo_value * (1.0 - ooo_alpha)
|
||||
pose[self.mouth_ooo_index] = uo_value * ooo_alpha
|
||||
|
||||
if self.panel is not None:
|
||||
frequency = self.breathing_frequency_slider.GetValue()
|
||||
if frequency == 0:
|
||||
value = 0.0
|
||||
pose[self.breathing_index] = value
|
||||
self.breathing_start_time = time.time()
|
||||
else:
|
||||
period = 60.0 / frequency
|
||||
now = time.time()
|
||||
diff = now - self.breathing_start_time
|
||||
frac = (diff % period) / period
|
||||
value = (-math.cos(2 * math.pi * frac) + 1.0) / 2.0
|
||||
pose[self.breathing_index] = value
|
||||
self.breathing_gauge.SetValue(int(1000 * value))
|
||||
|
||||
return pose
|
||||
|
||||
|
||||
def create_ifacialmocap_pose_converter(
|
||||
args: Optional[IFacialMocapPoseConverter25Args] = None) -> IFacialMocapPoseConverter:
|
||||
return IFacialMocapPoseConverter25(args)
|
||||
@@ -0,0 +1,89 @@
|
||||
import math
|
||||
|
||||
from tha4.mocap.ifacialmocap_constants import BLENDSHAPE_NAMES, HEAD_BONE_X, HEAD_BONE_Y, HEAD_BONE_Z, \
|
||||
RIGHT_EYE_BONE_X, RIGHT_EYE_BONE_Y, RIGHT_EYE_BONE_Z, LEFT_EYE_BONE_X, LEFT_EYE_BONE_Y, LEFT_EYE_BONE_Z, \
|
||||
HEAD_BONE_QUAT, LEFT_EYE_BONE_QUAT, RIGHT_EYE_BONE_QUAT
|
||||
|
||||
IFACIALMOCAP_PORT = 49983
|
||||
IFACIALMOCAP_START_STRING = "iFacialMocap_sahuasouryya9218sauhuiayeta91555dy3719|sendDataVersion=v2".encode('utf-8')
|
||||
|
||||
|
||||
def parse_ifacialmocap_v2_pose(ifacialmocap_output):
|
||||
output = {}
|
||||
parts = ifacialmocap_output.split("|")
|
||||
for part in parts:
|
||||
part = part.strip()
|
||||
if len(part) == 0:
|
||||
continue
|
||||
if "&" in part:
|
||||
components = part.split("&")
|
||||
assert len(components) == 2
|
||||
key = components[0]
|
||||
value = float(components[1]) / 100.0
|
||||
if key.endswith("_L"):
|
||||
key = key[:-2] + "Left"
|
||||
elif key.endswith("_R"):
|
||||
key = key[:-2] + "Right"
|
||||
if key in BLENDSHAPE_NAMES:
|
||||
output[key] = value
|
||||
elif part.startswith("=head#"):
|
||||
components = part[len("=head#"):].split(",")
|
||||
assert len(components) == 6
|
||||
output[HEAD_BONE_X] = float(components[0]) * math.pi / 180
|
||||
output[HEAD_BONE_Y] = float(components[1]) * math.pi / 180
|
||||
output[HEAD_BONE_Z] = float(components[2]) * math.pi / 180
|
||||
elif part.startswith("rightEye#"):
|
||||
components = part[len("rightEye#"):].split(",")
|
||||
output[RIGHT_EYE_BONE_X] = float(components[0]) * math.pi / 180
|
||||
output[RIGHT_EYE_BONE_Y] = float(components[1]) * math.pi / 180
|
||||
output[RIGHT_EYE_BONE_Z] = float(components[2]) * math.pi / 180
|
||||
elif part.startswith("leftEye#"):
|
||||
components = part[len("leftEye#"):].split(",")
|
||||
output[LEFT_EYE_BONE_X] = float(components[0]) * math.pi / 180
|
||||
output[LEFT_EYE_BONE_Y] = float(components[1]) * math.pi / 180
|
||||
output[LEFT_EYE_BONE_Z] = float(components[2]) * math.pi / 180
|
||||
output[HEAD_BONE_QUAT] = [0.0, 0.0, 0.0, 1.0]
|
||||
output[LEFT_EYE_BONE_QUAT] = [0.0, 0.0, 0.0, 1.0]
|
||||
output[RIGHT_EYE_BONE_QUAT] = [0.0, 0.0, 0.0, 1.0]
|
||||
return output
|
||||
|
||||
|
||||
def parse_ifacialmocap_v1_pose(ifacialmocap_output):
|
||||
output = {}
|
||||
parts = ifacialmocap_output.split("|")
|
||||
for part in parts:
|
||||
part = part.strip()
|
||||
if len(part) == 0:
|
||||
continue
|
||||
if part.startswith("=head#"):
|
||||
components = part[len("=head#"):].split(",")
|
||||
assert len(components) == 6
|
||||
output[HEAD_BONE_X] = float(components[0]) * math.pi / 180
|
||||
output[HEAD_BONE_Y] = float(components[1]) * math.pi / 180
|
||||
output[HEAD_BONE_Z] = float(components[2]) * math.pi / 180
|
||||
elif part.startswith("rightEye#"):
|
||||
components = part[len("rightEye#"):].split(",")
|
||||
output[RIGHT_EYE_BONE_X] = float(components[0]) * math.pi / 180
|
||||
output[RIGHT_EYE_BONE_Y] = float(components[1]) * math.pi / 180
|
||||
output[RIGHT_EYE_BONE_Z] = float(components[2]) * math.pi / 180
|
||||
elif part.startswith("leftEye#"):
|
||||
components = part[len("leftEye#"):].split(",")
|
||||
output[LEFT_EYE_BONE_X] = float(components[0]) * math.pi / 180
|
||||
output[LEFT_EYE_BONE_Y] = float(components[1]) * math.pi / 180
|
||||
output[LEFT_EYE_BONE_Z] = float(components[2]) * math.pi / 180
|
||||
else:
|
||||
components = part.split("-")
|
||||
assert len(components) == 2
|
||||
key = components[0]
|
||||
value = float(components[1]) / 100.0
|
||||
if key.endswith("_L"):
|
||||
key = key[:-2] + "Left"
|
||||
elif key.endswith("_R"):
|
||||
key = key[:-2] + "Right"
|
||||
if key in BLENDSHAPE_NAMES:
|
||||
output[key] = value
|
||||
output[HEAD_BONE_QUAT] = [0.0, 0.0, 0.0, 1.0]
|
||||
output[LEFT_EYE_BONE_QUAT] = [0.0, 0.0, 0.0, 1.0]
|
||||
output[RIGHT_EYE_BONE_QUAT] = [0.0, 0.0, 0.0, 1.0]
|
||||
return output
|
||||
|
||||
@@ -0,0 +1,216 @@
|
||||
EYE_LOOK_IN_LEFT = "eyeLookInLeft"
|
||||
EYE_LOOK_OUT_LEFT = "eyeLookOutLeft"
|
||||
EYE_LOOK_DOWN_LEFT = "eyeLookDownLeft"
|
||||
EYE_LOOK_UP_LEFT = "eyeLookUpLeft"
|
||||
EYE_BLINK_LEFT = "eyeBlinkLeft"
|
||||
EYE_SQUINT_LEFT = "eyeSquintLeft"
|
||||
EYE_WIDE_LEFT = "eyeWideLeft"
|
||||
EYE_LOOK_IN_RIGHT = "eyeLookInRight"
|
||||
EYE_LOOK_OUT_RIGHT = "eyeLookOutRight"
|
||||
EYE_LOOK_DOWN_RIGHT = "eyeLookDownRight"
|
||||
EYE_LOOK_UP_RIGHT = "eyeLookUpRight"
|
||||
EYE_BLINK_RIGHT = "eyeBlinkRight"
|
||||
EYE_SQUINT_RIGHT = "eyeSquintRight"
|
||||
EYE_WIDE_RIGHT = "eyeWideRight"
|
||||
BROW_DOWN_LEFT = "browDownLeft"
|
||||
BROW_OUTER_UP_LEFT = "browOuterUpLeft"
|
||||
BROW_DOWN_RIGHT = "browDownRight"
|
||||
BROW_OUTER_UP_RIGHT = "browOuterUpRight"
|
||||
BROW_INNER_UP = "browInnerUp"
|
||||
NOSE_SNEER_LEFT = "noseSneerLeft"
|
||||
NOSE_SNEER_RIGHT = "noseSneerRight"
|
||||
CHEEK_SQUINT_LEFT = "cheekSquintLeft"
|
||||
CHEEK_SQUINT_RIGHT = "cheekSquintRight"
|
||||
CHEEK_PUFF = "cheekPuff"
|
||||
MOUTH_LEFT = "mouthLeft"
|
||||
MOUTH_DIMPLE_LEFT = "mouthDimpleLeft"
|
||||
MOUTH_FROWN_LEFT = "mouthFrownLeft"
|
||||
MOUTH_LOWER_DOWN_LEFT = "mouthLowerDownLeft"
|
||||
MOUTH_PRESS_LEFT = "mouthPressLeft"
|
||||
MOUTH_SMILE_LEFT = "mouthSmileLeft"
|
||||
MOUTH_STRETCH_LEFT = "mouthStretchLeft"
|
||||
MOUTH_UPPER_UP_LEFT = "mouthUpperUpLeft"
|
||||
MOUTH_RIGHT = "mouthRight"
|
||||
MOUTH_DIMPLE_RIGHT = "mouthDimpleRight"
|
||||
MOUTH_FROWN_RIGHT = "mouthFrownRight"
|
||||
MOUTH_LOWER_DOWN_RIGHT = "mouthLowerDownRight"
|
||||
MOUTH_PRESS_RIGHT = "mouthPressRight"
|
||||
MOUTH_SMILE_RIGHT = "mouthSmileRight"
|
||||
MOUTH_STRETCH_RIGHT = "mouthStretchRight"
|
||||
MOUTH_UPPER_UP_RIGHT = "mouthUpperUpRight"
|
||||
MOUTH_CLOSE = "mouthClose"
|
||||
MOUTH_FUNNEL = "mouthFunnel"
|
||||
MOUTH_PUCKER = "mouthPucker"
|
||||
MOUTH_ROLL_LOWER = "mouthRollLower"
|
||||
MOUTH_ROLL_UPPER = "mouthRollUpper"
|
||||
MOUTH_SHRUG_LOWER = "mouthShrugLower"
|
||||
MOUTH_SHRUG_UPPER = "mouthShrugUpper"
|
||||
JAW_LEFT = "jawLeft"
|
||||
JAW_RIGHT = "jawRight"
|
||||
JAW_FORWARD = "jawForward"
|
||||
JAW_OPEN = "jawOpen"
|
||||
NEUTRAL = "_neutral"
|
||||
|
||||
BLENDSHAPE_NAMES = [
|
||||
EYE_LOOK_IN_LEFT, # 0
|
||||
EYE_LOOK_OUT_LEFT, # 1
|
||||
EYE_LOOK_DOWN_LEFT, # 2
|
||||
EYE_LOOK_UP_LEFT, # 3
|
||||
EYE_BLINK_LEFT, # 4
|
||||
EYE_SQUINT_LEFT, # 5
|
||||
EYE_WIDE_LEFT, # 6
|
||||
EYE_LOOK_IN_RIGHT, # 7
|
||||
EYE_LOOK_OUT_RIGHT, # 8
|
||||
EYE_LOOK_DOWN_RIGHT, # 9
|
||||
EYE_LOOK_UP_RIGHT, # 10
|
||||
EYE_BLINK_RIGHT, # 11
|
||||
EYE_SQUINT_RIGHT, # 12
|
||||
EYE_WIDE_RIGHT, # 13
|
||||
BROW_DOWN_LEFT, # 14
|
||||
BROW_OUTER_UP_LEFT, # 15
|
||||
BROW_DOWN_RIGHT, # 16
|
||||
BROW_OUTER_UP_RIGHT, # 17
|
||||
BROW_INNER_UP, # 18
|
||||
NOSE_SNEER_LEFT, # 19
|
||||
NOSE_SNEER_RIGHT, # 20
|
||||
CHEEK_SQUINT_LEFT, # 21
|
||||
CHEEK_SQUINT_RIGHT, # 22
|
||||
CHEEK_PUFF, # 23
|
||||
MOUTH_LEFT, # 24
|
||||
MOUTH_DIMPLE_LEFT, # 25
|
||||
MOUTH_FROWN_LEFT, # 26
|
||||
MOUTH_LOWER_DOWN_LEFT, # 27
|
||||
MOUTH_PRESS_LEFT, # 28
|
||||
MOUTH_SMILE_LEFT, # 29
|
||||
MOUTH_STRETCH_LEFT, # 30
|
||||
MOUTH_UPPER_UP_LEFT, # 31
|
||||
MOUTH_RIGHT, # 32
|
||||
MOUTH_DIMPLE_RIGHT, # 33
|
||||
MOUTH_FROWN_RIGHT, # 34
|
||||
MOUTH_LOWER_DOWN_RIGHT, # 35
|
||||
MOUTH_PRESS_RIGHT, # 36
|
||||
MOUTH_SMILE_RIGHT, # 37
|
||||
MOUTH_STRETCH_RIGHT, # 38
|
||||
MOUTH_UPPER_UP_RIGHT, # 39
|
||||
MOUTH_CLOSE, # 40
|
||||
MOUTH_FUNNEL, # 41
|
||||
MOUTH_PUCKER, # 42
|
||||
MOUTH_ROLL_LOWER, # 43
|
||||
MOUTH_ROLL_UPPER, # 44
|
||||
MOUTH_SHRUG_LOWER, # 45
|
||||
MOUTH_SHRUG_UPPER, # 46
|
||||
JAW_LEFT, # 47
|
||||
JAW_RIGHT, # 48
|
||||
JAW_FORWARD, # 49
|
||||
JAW_OPEN, # 50
|
||||
NEUTRAL, # 51
|
||||
]
|
||||
|
||||
EYE_LEFT_BLENDSHAPES = [
|
||||
EYE_LOOK_IN_LEFT, # 0
|
||||
EYE_LOOK_OUT_LEFT, # 1
|
||||
EYE_LOOK_DOWN_LEFT, # 2
|
||||
EYE_LOOK_UP_LEFT, # 3
|
||||
EYE_BLINK_LEFT, # 4
|
||||
EYE_SQUINT_LEFT, # 5
|
||||
EYE_WIDE_LEFT, # 6
|
||||
]
|
||||
|
||||
EYE_RIGHT_BLENDSHAPES = [
|
||||
EYE_LOOK_IN_RIGHT, # 7
|
||||
EYE_LOOK_OUT_RIGHT, # 8
|
||||
EYE_LOOK_DOWN_RIGHT, # 9
|
||||
EYE_LOOK_UP_RIGHT, # 10
|
||||
EYE_BLINK_RIGHT, # 11
|
||||
EYE_SQUINT_RIGHT, # 12
|
||||
EYE_WIDE_RIGHT, # 13
|
||||
]
|
||||
|
||||
BROW_LEFT_BLENDSHAPES = [
|
||||
BROW_DOWN_LEFT, # 14
|
||||
BROW_OUTER_UP_LEFT, # 15
|
||||
|
||||
]
|
||||
|
||||
BROW_RIGHT_BLENDSHAPES = [
|
||||
BROW_DOWN_RIGHT, # 16
|
||||
BROW_OUTER_UP_RIGHT, # 17
|
||||
|
||||
]
|
||||
|
||||
BROW_BOTH_BLENDSHAPES = [
|
||||
BROW_INNER_UP, # 18
|
||||
]
|
||||
|
||||
NOSE_BLENDSHAPES = [
|
||||
NOSE_SNEER_LEFT, # 19
|
||||
NOSE_SNEER_RIGHT, # 20
|
||||
]
|
||||
|
||||
CHECK_BLENDSHAPES = [
|
||||
CHEEK_SQUINT_LEFT, # 21
|
||||
CHEEK_SQUINT_RIGHT, # 22
|
||||
CHEEK_PUFF, # 23
|
||||
]
|
||||
|
||||
MOUTH_LEFT_BLENDSHAPES = [
|
||||
MOUTH_LEFT, # 24
|
||||
MOUTH_DIMPLE_LEFT, # 25
|
||||
MOUTH_FROWN_LEFT, # 26
|
||||
MOUTH_LOWER_DOWN_LEFT, # 27
|
||||
MOUTH_PRESS_LEFT, # 28
|
||||
MOUTH_SMILE_LEFT, # 29
|
||||
MOUTH_STRETCH_LEFT, # 30
|
||||
MOUTH_UPPER_UP_LEFT, # 31
|
||||
]
|
||||
|
||||
MOUTH_RIGHT_BLENDSHAPES = [
|
||||
MOUTH_RIGHT, # 32
|
||||
MOUTH_DIMPLE_RIGHT, # 33
|
||||
MOUTH_FROWN_RIGHT, # 34
|
||||
MOUTH_LOWER_DOWN_RIGHT, # 35
|
||||
MOUTH_PRESS_RIGHT, # 36
|
||||
MOUTH_SMILE_RIGHT, # 37
|
||||
MOUTH_STRETCH_RIGHT, # 38
|
||||
MOUTH_UPPER_UP_RIGHT, # 39
|
||||
]
|
||||
|
||||
MOUTH_BOTH_BLENDSHAPES = [
|
||||
MOUTH_CLOSE, # 40
|
||||
MOUTH_FUNNEL, # 41
|
||||
MOUTH_PUCKER, # 42
|
||||
MOUTH_ROLL_LOWER, # 43
|
||||
MOUTH_ROLL_UPPER, # 44
|
||||
MOUTH_SHRUG_LOWER, # 45
|
||||
MOUTH_SHRUG_UPPER, # 46
|
||||
]
|
||||
|
||||
JAW_BLENDSHAPES = [
|
||||
JAW_LEFT, # 47
|
||||
JAW_RIGHT, # 48
|
||||
JAW_FORWARD, # 49
|
||||
JAW_OPEN, # 50
|
||||
]
|
||||
|
||||
NEUTRAL_BLENDSHAPES = [
|
||||
NEUTRAL, # 51
|
||||
]
|
||||
|
||||
COLUMN_0_BLENDSHAPES = EYE_RIGHT_BLENDSHAPES + BROW_RIGHT_BLENDSHAPES + [NOSE_SNEER_RIGHT, CHEEK_SQUINT_RIGHT]
|
||||
COLUMN_1_BLENDSHAPES = EYE_LEFT_BLENDSHAPES + BROW_LEFT_BLENDSHAPES + [NOSE_SNEER_LEFT, CHEEK_SQUINT_LEFT]
|
||||
COLUMN_2_BLENDSHAPES = MOUTH_RIGHT_BLENDSHAPES + [JAW_RIGHT]
|
||||
COLUMN_3_BLENDSHAPES = MOUTH_LEFT_BLENDSHAPES + [JAW_LEFT]
|
||||
COLUMN_4_BLENDSHAPES = [BROW_INNER_UP, CHEEK_PUFF] + MOUTH_BOTH_BLENDSHAPES + [JAW_FORWARD, JAW_OPEN, NEUTRAL]
|
||||
|
||||
BLENDSHAPE_COLUMNS = [
|
||||
COLUMN_0_BLENDSHAPES,
|
||||
COLUMN_1_BLENDSHAPES,
|
||||
COLUMN_2_BLENDSHAPES,
|
||||
COLUMN_3_BLENDSHAPES,
|
||||
COLUMN_4_BLENDSHAPES,
|
||||
]
|
||||
|
||||
HEAD_X = "headX"
|
||||
HEAD_Y = "headY"
|
||||
HEAD_Z = "headZ"
|
||||
HEAD_ROTATIONS = [HEAD_X, HEAD_Y, HEAD_Z]
|
||||
@@ -0,0 +1,41 @@
|
||||
import json
|
||||
import os
|
||||
from typing import Optional, Dict
|
||||
|
||||
import numpy
|
||||
|
||||
|
||||
class MediaPipeFacePose:
|
||||
KEY_BLENDSHAPE_PARAMS = "blendshape_params"
|
||||
KEY_XFORM_MATRIX = "xform_matrix"
|
||||
|
||||
def __init__(self, blendshape_params: Optional[Dict[str, float]], xform_matrix: Optional[numpy.ndarray]):
|
||||
if blendshape_params is None:
|
||||
blendshape_params = {}
|
||||
if xform_matrix is None:
|
||||
self.xform_matrix = numpy.zeros(4, 4)
|
||||
for i in range(4):
|
||||
self.xform_matrix[i, i] = 1.0
|
||||
|
||||
self.blendshape_params = blendshape_params
|
||||
self.xform_matrix = xform_matrix
|
||||
|
||||
def get_json(self):
|
||||
return {
|
||||
MediaPipeFacePose.KEY_BLENDSHAPE_PARAMS: self.blendshape_params.copy(),
|
||||
MediaPipeFacePose.KEY_XFORM_MATRIX: self.xform_matrix.tolist()
|
||||
}
|
||||
|
||||
def save(self, file_name: str):
|
||||
os.makedirs(os.path.dirname(file_name), exist_ok=True)
|
||||
with open(file_name, "wt") as fout:
|
||||
fout.write(json.dumps(self.get_json()))
|
||||
|
||||
@staticmethod
|
||||
def load(file_name: str):
|
||||
with open(file_name, "rt") as fin:
|
||||
s = fin.read()
|
||||
json_data = json.loads(s)
|
||||
return MediaPipeFacePose(
|
||||
json_data[MediaPipeFacePose.KEY_BLENDSHAPE_PARAMS],
|
||||
xform_matrix = numpy.array(json_data[MediaPipeFacePose.KEY_XFORM_MATRIX]))
|
||||
@@ -0,0 +1,17 @@
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import List, Callable, Optional
|
||||
|
||||
from tha4.mocap.mediapipe_face_pose import MediaPipeFacePose
|
||||
|
||||
|
||||
class MediaPipeFacePoseConverter(ABC):
|
||||
@abstractmethod
|
||||
def convert(self, mediapipe_face_pose: MediaPipeFacePose) -> List[float]:
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def init_pose_converter_panel(
|
||||
self,
|
||||
parent,
|
||||
current_pose_supplier: Callable[[], Optional[MediaPipeFacePose]]):
|
||||
pass
|
||||
@@ -0,0 +1,662 @@
|
||||
import math
|
||||
import time
|
||||
from enum import Enum
|
||||
from typing import Optional, List, Callable
|
||||
|
||||
import numpy
|
||||
import scipy.optimize
|
||||
import wx
|
||||
from scipy.spatial.transform import Rotation
|
||||
|
||||
from tha4.poser.modes.pose_parameters import get_pose_parameters
|
||||
from tha4.mocap.mediapipe_constants import MOUTH_SMILE_LEFT, MOUTH_SHRUG_UPPER, MOUTH_SMILE_RIGHT, \
|
||||
BROW_INNER_UP, BROW_OUTER_UP_RIGHT, BROW_OUTER_UP_LEFT, BROW_DOWN_LEFT, BROW_DOWN_RIGHT, EYE_WIDE_LEFT, \
|
||||
EYE_WIDE_RIGHT, EYE_BLINK_LEFT, EYE_BLINK_RIGHT, CHEEK_SQUINT_LEFT, CHEEK_SQUINT_RIGHT, EYE_LOOK_IN_LEFT, \
|
||||
EYE_LOOK_OUT_LEFT, EYE_LOOK_IN_RIGHT, EYE_LOOK_OUT_RIGHT, EYE_LOOK_UP_LEFT, EYE_LOOK_UP_RIGHT, EYE_LOOK_DOWN_RIGHT, \
|
||||
EYE_LOOK_DOWN_LEFT, JAW_OPEN, MOUTH_FROWN_LEFT, MOUTH_FROWN_RIGHT, \
|
||||
MOUTH_LOWER_DOWN_LEFT, MOUTH_LOWER_DOWN_RIGHT, MOUTH_FUNNEL, MOUTH_PUCKER
|
||||
from tha4.mocap.mediapipe_face_pose import MediaPipeFacePose
|
||||
from tha4.mocap.mediapipe_face_pose_converter import MediaPipeFacePoseConverter
|
||||
|
||||
|
||||
class EyebrowDownMode(Enum):
|
||||
TROUBLED = 1
|
||||
ANGRY = 2
|
||||
LOWERED = 3
|
||||
SERIOUS = 4
|
||||
|
||||
|
||||
class WinkMode(Enum):
|
||||
NORMAL = 1
|
||||
RELAXED = 2
|
||||
|
||||
|
||||
def rad_to_deg(rad):
|
||||
return rad * 180.0 / math.pi
|
||||
|
||||
|
||||
def deg_to_rad(deg):
|
||||
return deg * math.pi / 180.0
|
||||
|
||||
|
||||
def clamp(x, min_value, max_value):
|
||||
return max(min_value, min(max_value, x))
|
||||
|
||||
|
||||
class MediaPipeFacePoseConverter00Args:
|
||||
def __init__(self,
|
||||
smile_threshold_min: float = 0.4,
|
||||
smile_threshold_max: float = 0.6,
|
||||
eyebrow_down_mode: EyebrowDownMode = EyebrowDownMode.ANGRY,
|
||||
wink_mode: WinkMode = WinkMode.NORMAL,
|
||||
eye_surprised_max: float = 0.5,
|
||||
eye_blink_max: float = 0.8,
|
||||
eyebrow_down_max: float = 0.4,
|
||||
cheek_squint_min: float = 0.1,
|
||||
cheek_squint_max: float = 0.7,
|
||||
eye_rotation_factor: float = 1.0 / 0.75,
|
||||
jaw_open_min: float = 0.1,
|
||||
jaw_open_max: float = 0.4,
|
||||
mouth_frown_max: float = 0.6,
|
||||
mouth_funnel_min: float = 0.25,
|
||||
mouth_funnel_max: float = 0.5,
|
||||
iris_small_left=0.0,
|
||||
iris_small_right=0.0,
|
||||
head_x_offset=0.0,
|
||||
head_y_offset=0.0,
|
||||
head_z_offset=0.0):
|
||||
self.iris_small_right = iris_small_left
|
||||
self.iris_small_left = iris_small_right
|
||||
|
||||
self.wink_mode = wink_mode
|
||||
|
||||
self.mouth_funnel_max = mouth_funnel_max
|
||||
self.mouth_funnel_min = mouth_funnel_min
|
||||
self.mouth_frown_max = mouth_frown_max
|
||||
|
||||
self.jaw_open_max = jaw_open_max
|
||||
self.jaw_open_min = jaw_open_min
|
||||
|
||||
self.eye_rotation_factor = eye_rotation_factor
|
||||
|
||||
self.cheek_squint_max = cheek_squint_max
|
||||
self.cheek_squint_min = cheek_squint_min
|
||||
|
||||
self.eyebrow_down_max = eyebrow_down_max
|
||||
|
||||
self.eye_blink_max = eye_blink_max
|
||||
self.eye_surprised_max = eye_surprised_max
|
||||
|
||||
self.smile_threshold_min = smile_threshold_min
|
||||
self.smile_threshold_max = smile_threshold_max
|
||||
|
||||
self.head_z_offset = head_z_offset
|
||||
self.head_y_offset = head_y_offset
|
||||
self.head_x_offset = head_x_offset
|
||||
|
||||
self.eyebrow_down_mode = eyebrow_down_mode
|
||||
|
||||
def set_smile_threshold_min(self, new_value: float):
|
||||
self.smile_threshold_min = new_value
|
||||
|
||||
def set_smile_threshold_max(self, new_value: float):
|
||||
self.smile_threshold_max = new_value
|
||||
|
||||
def set_eye_surprised_max(self, new_value: float):
|
||||
self.eye_surprised_max = new_value
|
||||
|
||||
def set_eye_blink_max(self, new_value: float):
|
||||
self.eye_blink_max = new_value
|
||||
|
||||
def set_eyebrow_down_max(self, new_value: float):
|
||||
self.eyebrow_down_max = new_value
|
||||
|
||||
def set_cheek_squint_min(self, new_value: float):
|
||||
self.cheek_squint_min = new_value
|
||||
|
||||
def set_cheek_squint_max(self, new_value: float):
|
||||
self.cheek_squint_max = new_value
|
||||
|
||||
def set_jaw_open_min(self, new_value: float):
|
||||
self.jaw_open_min = new_value
|
||||
|
||||
def set_jaw_open_max(self, new_value: float):
|
||||
self.jaw_open_max = new_value
|
||||
|
||||
def set_mouth_frown_max(self, new_value: float):
|
||||
self.mouth_frown_max = new_value
|
||||
|
||||
def set_mouth_funnel_min(self, new_value: float):
|
||||
self.mouth_funnel_min = new_value
|
||||
|
||||
def set_mouth_funnel_max(self, new_value: float):
|
||||
self.mouth_funnel_min = new_value
|
||||
|
||||
|
||||
class MediaPoseFacePoseConverter00(MediaPipeFacePoseConverter):
|
||||
def __init__(self, args: Optional[MediaPipeFacePoseConverter00Args] = None):
|
||||
super().__init__()
|
||||
if args is None:
|
||||
args = MediaPipeFacePoseConverter00Args()
|
||||
self.args = args
|
||||
pose_parameters = get_pose_parameters()
|
||||
self.pose_size = 45
|
||||
|
||||
self.eyebrow_troubled_left_index = pose_parameters.get_parameter_index("eyebrow_troubled_left")
|
||||
self.eyebrow_troubled_right_index = pose_parameters.get_parameter_index("eyebrow_troubled_right")
|
||||
self.eyebrow_angry_left_index = pose_parameters.get_parameter_index("eyebrow_angry_left")
|
||||
self.eyebrow_angry_right_index = pose_parameters.get_parameter_index("eyebrow_angry_right")
|
||||
self.eyebrow_happy_left_index = pose_parameters.get_parameter_index("eyebrow_happy_left")
|
||||
self.eyebrow_happy_right_index = pose_parameters.get_parameter_index("eyebrow_happy_right")
|
||||
self.eyebrow_raised_left_index = pose_parameters.get_parameter_index("eyebrow_raised_left")
|
||||
self.eyebrow_raised_right_index = pose_parameters.get_parameter_index("eyebrow_raised_right")
|
||||
self.eyebrow_lowered_left_index = pose_parameters.get_parameter_index("eyebrow_lowered_left")
|
||||
self.eyebrow_lowered_right_index = pose_parameters.get_parameter_index("eyebrow_lowered_right")
|
||||
self.eyebrow_serious_left_index = pose_parameters.get_parameter_index("eyebrow_serious_left")
|
||||
self.eyebrow_serious_right_index = pose_parameters.get_parameter_index("eyebrow_serious_right")
|
||||
|
||||
self.eye_surprised_left_index = pose_parameters.get_parameter_index("eye_surprised_left")
|
||||
self.eye_surprised_right_index = pose_parameters.get_parameter_index("eye_surprised_right")
|
||||
self.eye_wink_left_index = pose_parameters.get_parameter_index("eye_wink_left")
|
||||
self.eye_wink_right_index = pose_parameters.get_parameter_index("eye_wink_right")
|
||||
self.eye_happy_wink_left_index = pose_parameters.get_parameter_index("eye_happy_wink_left")
|
||||
self.eye_happy_wink_right_index = pose_parameters.get_parameter_index("eye_happy_wink_right")
|
||||
self.eye_relaxed_left_index = pose_parameters.get_parameter_index("eye_relaxed_left")
|
||||
self.eye_relaxed_right_index = pose_parameters.get_parameter_index("eye_relaxed_right")
|
||||
self.eye_raised_lower_eyelid_left_index = pose_parameters.get_parameter_index("eye_raised_lower_eyelid_left")
|
||||
self.eye_raised_lower_eyelid_right_index = pose_parameters.get_parameter_index("eye_raised_lower_eyelid_right")
|
||||
|
||||
self.iris_small_left_index = pose_parameters.get_parameter_index("iris_small_left")
|
||||
self.iris_small_right_index = pose_parameters.get_parameter_index("iris_small_right")
|
||||
|
||||
self.iris_rotation_x_index = pose_parameters.get_parameter_index("iris_rotation_x")
|
||||
self.iris_rotation_y_index = pose_parameters.get_parameter_index("iris_rotation_y")
|
||||
|
||||
self.head_x_index = pose_parameters.get_parameter_index("head_x")
|
||||
self.head_y_index = pose_parameters.get_parameter_index("head_y")
|
||||
self.neck_z_index = pose_parameters.get_parameter_index("neck_z")
|
||||
|
||||
self.mouth_aaa_index = pose_parameters.get_parameter_index("mouth_aaa")
|
||||
self.mouth_iii_index = pose_parameters.get_parameter_index("mouth_iii")
|
||||
self.mouth_uuu_index = pose_parameters.get_parameter_index("mouth_uuu")
|
||||
self.mouth_eee_index = pose_parameters.get_parameter_index("mouth_eee")
|
||||
self.mouth_ooo_index = pose_parameters.get_parameter_index("mouth_ooo")
|
||||
|
||||
self.mouth_lowered_corner_left_index = pose_parameters.get_parameter_index("mouth_lowered_corner_left")
|
||||
self.mouth_lowered_corner_right_index = pose_parameters.get_parameter_index("mouth_lowered_corner_right")
|
||||
self.mouth_raised_corner_left_index = pose_parameters.get_parameter_index("mouth_raised_corner_left")
|
||||
self.mouth_raised_corner_right_index = pose_parameters.get_parameter_index("mouth_raised_corner_right")
|
||||
|
||||
self.body_y_index = pose_parameters.get_parameter_index("body_y")
|
||||
self.body_z_index = pose_parameters.get_parameter_index("body_z")
|
||||
self.breathing_index = pose_parameters.get_parameter_index("breathing")
|
||||
|
||||
self.breathing_start_time = time.time()
|
||||
|
||||
self.panel = None
|
||||
self.current_pose_supplier = None
|
||||
|
||||
def init_pose_converter_panel(
|
||||
self,
|
||||
parent,
|
||||
current_pose_supplier: Callable[[], Optional[MediaPipeFacePose]]):
|
||||
self.panel = wx.Panel(parent, style=wx.SIMPLE_BORDER)
|
||||
self.panel_sizer = wx.BoxSizer(wx.VERTICAL)
|
||||
self.panel.SetSizer(self.panel_sizer)
|
||||
self.panel.SetAutoLayout(1)
|
||||
parent.GetSizer().Add(self.panel, 0, wx.EXPAND)
|
||||
|
||||
self.current_pose_supplier = current_pose_supplier
|
||||
|
||||
if True:
|
||||
eyebrow_down_mode_text = wx.StaticText(self.panel, label=" --- Eyebrow Down Mode --- ",
|
||||
style=wx.ALIGN_CENTER)
|
||||
self.panel_sizer.Add(eyebrow_down_mode_text, 0, wx.EXPAND)
|
||||
|
||||
self.eyebrow_down_mode_choice = wx.Choice(
|
||||
self.panel,
|
||||
choices=[
|
||||
"ANGRY",
|
||||
"TROUBLED",
|
||||
"SERIOUS",
|
||||
"LOWERED",
|
||||
])
|
||||
self.eyebrow_down_mode_choice.SetSelection(0)
|
||||
self.panel_sizer.Add(self.eyebrow_down_mode_choice, 0, wx.EXPAND)
|
||||
self.eyebrow_down_mode_choice.Bind(wx.EVT_CHOICE, self.change_eyebrow_down_mode)
|
||||
|
||||
if True:
|
||||
separator = wx.StaticLine(self.panel, -1, size=(256, 5))
|
||||
self.panel_sizer.Add(separator, 0, wx.EXPAND)
|
||||
|
||||
wink_mode_text = wx.StaticText(self.panel, label=" --- Wink Mode --- ", style=wx.ALIGN_CENTER)
|
||||
self.panel_sizer.Add(wink_mode_text, 0, wx.EXPAND)
|
||||
|
||||
self.wink_mode_choice = wx.Choice(
|
||||
self.panel,
|
||||
choices=[
|
||||
"NORMAL",
|
||||
"RELAXED",
|
||||
])
|
||||
self.wink_mode_choice.SetSelection(0)
|
||||
self.panel_sizer.Add(self.wink_mode_choice, 0, wx.EXPAND)
|
||||
self.wink_mode_choice.Bind(wx.EVT_CHOICE, self.change_wink_mode)
|
||||
|
||||
if True:
|
||||
separator = wx.StaticLine(self.panel, -1, size=(256, 5))
|
||||
self.panel_sizer.Add(separator, 0, wx.EXPAND)
|
||||
|
||||
iris_size_text = wx.StaticText(self.panel, label=" --- Iris Size --- ", style=wx.ALIGN_CENTER)
|
||||
self.panel_sizer.Add(iris_size_text, 0, wx.EXPAND)
|
||||
|
||||
self.iris_left_slider = wx.Slider(self.panel, minValue=0, maxValue=1000, value=0, style=wx.HORIZONTAL)
|
||||
self.panel_sizer.Add(self.iris_left_slider, 0, wx.EXPAND)
|
||||
self.iris_left_slider.Bind(wx.EVT_SLIDER, self.change_iris_size)
|
||||
|
||||
self.iris_right_slider = wx.Slider(self.panel, minValue=0, maxValue=1000, value=0, style=wx.HORIZONTAL)
|
||||
self.panel_sizer.Add(self.iris_right_slider, 0, wx.EXPAND)
|
||||
self.iris_right_slider.Bind(wx.EVT_SLIDER, self.change_iris_size)
|
||||
self.iris_right_slider.Enable(False)
|
||||
|
||||
self.link_left_right_irises = wx.CheckBox(
|
||||
self.panel, label="Use same value for both sides")
|
||||
self.link_left_right_irises.SetValue(True)
|
||||
self.panel_sizer.Add(self.link_left_right_irises, wx.SizerFlags().CenterHorizontal().Border())
|
||||
self.link_left_right_irises.Bind(wx.EVT_CHECKBOX, self.link_left_right_irises_clicked)
|
||||
|
||||
if True:
|
||||
separator = wx.StaticLine(self.panel, -1, size=(256, 5))
|
||||
self.panel_sizer.Add(separator, 0, wx.EXPAND)
|
||||
|
||||
breathing_frequency_text = wx.StaticText(
|
||||
self.panel, label=" --- Breathing --- ", style=wx.ALIGN_CENTER)
|
||||
self.panel_sizer.Add(breathing_frequency_text, 0, wx.EXPAND)
|
||||
|
||||
self.restart_breathing_cycle_button = wx.Button(self.panel, label="Restart Breathing Cycle")
|
||||
self.restart_breathing_cycle_button.Bind(wx.EVT_BUTTON, self.restart_breathing_cycle_clicked)
|
||||
self.panel_sizer.Add(self.restart_breathing_cycle_button, 0, wx.EXPAND)
|
||||
|
||||
self.breathing_frequency_slider = wx.Slider(
|
||||
self.panel, minValue=0, maxValue=60, value=20, style=wx.HORIZONTAL)
|
||||
self.panel_sizer.Add(self.breathing_frequency_slider, 0, wx.EXPAND)
|
||||
|
||||
self.breathing_gauge = wx.Gauge(self.panel, style=wx.GA_HORIZONTAL, range=1000)
|
||||
self.panel_sizer.Add(self.breathing_gauge, 0, wx.EXPAND)
|
||||
|
||||
if True:
|
||||
separator = wx.StaticLine(self.panel, -1, size=(256, 5))
|
||||
self.panel_sizer.Add(separator, 0, wx.EXPAND)
|
||||
|
||||
face_orientation_text = wx.StaticText(
|
||||
self.panel, label="--- Face Orientation ---", style=wx.ALIGN_CENTER)
|
||||
self.panel_sizer.Add(face_orientation_text, 0, wx.EXPAND)
|
||||
|
||||
self.calibrate_face_orientation_button = wx.Button(self.panel, label="Calibrate (I'm looking forward)")
|
||||
self.calibrate_face_orientation_button.Bind(wx.EVT_BUTTON, self.calibrate_face_orientation_clicked)
|
||||
self.panel_sizer.Add(self.calibrate_face_orientation_button, 0, wx.EXPAND)
|
||||
|
||||
if True:
|
||||
separator = wx.StaticLine(self.panel, -1, size=(256, 5))
|
||||
self.panel_sizer.Add(separator, 0, wx.EXPAND)
|
||||
|
||||
convertion_parameters_text = wx.StaticText(
|
||||
self.panel, label="--- Conversion Parameters ---", style=wx.ALIGN_CENTER)
|
||||
self.panel_sizer.Add(convertion_parameters_text, 0, wx.EXPAND)
|
||||
|
||||
conversion_param_panel = wx.Panel(self.panel)
|
||||
self.panel_sizer.Add(conversion_param_panel, 0, wx.EXPAND)
|
||||
conversion_panel_sizer = wx.FlexGridSizer(cols=2)
|
||||
conversion_panel_sizer.AddGrowableCol(1)
|
||||
conversion_param_panel.SetSizer(conversion_panel_sizer)
|
||||
conversion_param_panel.SetAutoLayout(1)
|
||||
|
||||
self.smile_thresold_min_spin = self.create_spin_control(
|
||||
conversion_param_panel,
|
||||
"Smile Threshold Min:", self.args.smile_threshold_min, self.args.set_smile_threshold_min)
|
||||
self.smile_thresold_max_spin = self.create_spin_control(
|
||||
conversion_param_panel,
|
||||
"Smile Threshold Max:", self.args.smile_threshold_max, self.args.set_smile_threshold_max)
|
||||
self.eye_surprised_max_spin = self.create_spin_control(
|
||||
conversion_param_panel,
|
||||
"Eye Surprised Max:", self.args.eye_surprised_max, self.args.set_eye_surprised_max)
|
||||
self.eye_blink_max_spin = self.create_spin_control(
|
||||
conversion_param_panel,
|
||||
"Eye Blink Max:", self.args.eye_blink_max, self.args.set_eye_blink_max)
|
||||
self.eyebrow_down_max_spin = self.create_spin_control(
|
||||
conversion_param_panel,
|
||||
"Eyebrow Down Max:", self.args.eyebrow_down_max, self.args.set_eyebrow_down_max)
|
||||
self.cheek_squint_min_spin = self.create_spin_control(
|
||||
conversion_param_panel,
|
||||
"Cheek Squint Min:", self.args.cheek_squint_min, self.args.set_cheek_squint_min)
|
||||
self.cheek_squint_max_spin = self.create_spin_control(
|
||||
conversion_param_panel,
|
||||
"Cheek Squint Max:", self.args.cheek_squint_max, self.args.set_cheek_squint_max)
|
||||
self.jaw_open_min_spin = self.create_spin_control(
|
||||
conversion_param_panel,
|
||||
"Jaw Open Min:", self.args.jaw_open_min, self.args.set_jaw_open_min)
|
||||
self.jaw_open_max_spin = self.create_spin_control(
|
||||
conversion_param_panel,
|
||||
"Jaw Open Max:", self.args.jaw_open_max, self.args.set_jaw_open_max)
|
||||
self.mouth_frown_max_spin = self.create_spin_control(
|
||||
conversion_param_panel,
|
||||
"Mouth Frown Max:", self.args.mouth_frown_max, self.args.set_mouth_frown_max)
|
||||
self.mouth_funnel_min_spin = self.create_spin_control(
|
||||
conversion_param_panel,
|
||||
"Mouth Funnel Min:", self.args.mouth_funnel_min, self.args.set_mouth_funnel_min)
|
||||
self.mouth_funnel_max_spin = self.create_spin_control(
|
||||
conversion_param_panel,
|
||||
"Mouth Funnel Max:", self.args.mouth_funnel_max, self.args.set_mouth_funnel_max)
|
||||
|
||||
self.panel_sizer.Fit(self.panel)
|
||||
|
||||
def create_spin_control(self, parent, label: str, initial_value: float, set_func: Callable[[float], None]):
|
||||
sizer = parent.GetSizer()
|
||||
|
||||
text = wx.StaticText(parent, label=label)
|
||||
sizer.Add(text, wx.SizerFlags().Right().Border(wx.ALL, 2))
|
||||
|
||||
spin_ctrl = wx.SpinCtrlDouble(
|
||||
parent,
|
||||
wx.ID_ANY,
|
||||
min=0.0,
|
||||
max=1.0,
|
||||
initial=initial_value,
|
||||
inc=0.01)
|
||||
sizer.Add(spin_ctrl, wx.SizerFlags().Border(wx.ALL, 2).Expand())
|
||||
|
||||
def handler(event: wx.Event):
|
||||
new_value = spin_ctrl.GetValue()
|
||||
set_func(new_value)
|
||||
|
||||
spin_ctrl.Bind(wx.EVT_SPINCTRLDOUBLE, handler)
|
||||
|
||||
return spin_ctrl
|
||||
|
||||
def extract_euler_angles(self, mediapipe_face_pose: MediaPipeFacePose):
|
||||
M = mediapipe_face_pose.xform_matrix[0:3, 0:3]
|
||||
rot = Rotation.from_matrix(M)
|
||||
return rot.as_euler('xyz', degrees=False)
|
||||
|
||||
def calibrate_face_orientation_clicked(self, event: wx.Event):
|
||||
if self.current_pose_supplier is None:
|
||||
return
|
||||
|
||||
mediapipe_face_pose = self.current_pose_supplier()
|
||||
if mediapipe_face_pose is None:
|
||||
return
|
||||
|
||||
euler_angles = self.extract_euler_angles(mediapipe_face_pose)
|
||||
self.args.head_x_offset = euler_angles[0]
|
||||
self.args.head_y_offset = euler_angles[1]
|
||||
self.args.head_z_offset = euler_angles[2]
|
||||
|
||||
def restart_breathing_cycle_clicked(self, event: wx.Event):
|
||||
self.breathing_start_time = time.time()
|
||||
|
||||
def change_eyebrow_down_mode(self, event: wx.Event):
|
||||
selected_index = self.eyebrow_down_mode_choice.GetSelection()
|
||||
if selected_index == 0:
|
||||
self.args.eyebrow_down_mode = EyebrowDownMode.ANGRY
|
||||
elif selected_index == 1:
|
||||
self.args.eyebrow_down_mode = EyebrowDownMode.TROUBLED
|
||||
elif selected_index == 2:
|
||||
self.args.eyebrow_down_mode = EyebrowDownMode.SERIOUS
|
||||
else:
|
||||
self.args.eyebrow_down_mode = EyebrowDownMode.LOWERED
|
||||
|
||||
def change_wink_mode(self, event: wx.Event):
|
||||
selected_index = self.wink_mode_choice.GetSelection()
|
||||
if selected_index == 0:
|
||||
self.args.wink_mode = WinkMode.NORMAL
|
||||
else:
|
||||
self.args.wink_mode = WinkMode.RELAXED
|
||||
|
||||
def change_iris_size(self, event: wx.Event):
|
||||
if self.link_left_right_irises.GetValue():
|
||||
left_value = self.iris_left_slider.GetValue()
|
||||
right_value = self.iris_right_slider.GetValue()
|
||||
if left_value != right_value:
|
||||
self.iris_right_slider.SetValue(left_value)
|
||||
self.args.iris_small_left = left_value / 1000.0
|
||||
self.args.iris_small_right = left_value / 1000.0
|
||||
else:
|
||||
self.args.iris_small_left = self.iris_left_slider.GetValue() / 1000.0
|
||||
self.args.iris_small_right = self.iris_right_slider.GetValue() / 1000.0
|
||||
|
||||
def link_left_right_irises_clicked(self, event: wx.Event):
|
||||
if self.link_left_right_irises.GetValue():
|
||||
self.iris_right_slider.Enable(False)
|
||||
else:
|
||||
self.iris_right_slider.Enable(True)
|
||||
self.change_iris_size(event)
|
||||
|
||||
def decompose_head_body_param(self, param, threshold=2.0 / 3):
|
||||
if abs(param) < threshold:
|
||||
return (param, 0.0)
|
||||
else:
|
||||
if param < 0:
|
||||
sign = -1.0
|
||||
else:
|
||||
sign = 1.0
|
||||
return (threshold * sign, (abs(param) - threshold) * sign)
|
||||
|
||||
def convert(self, mediapipe_face_pose: MediaPipeFacePose) -> List[float]:
|
||||
pose = [0.0 for i in range(self.pose_size)]
|
||||
|
||||
blendshape_params = mediapipe_face_pose.blendshape_params
|
||||
|
||||
smile_value = \
|
||||
(blendshape_params[MOUTH_SMILE_LEFT] + blendshape_params[MOUTH_SMILE_RIGHT]) / 2.0 \
|
||||
+ blendshape_params[MOUTH_SHRUG_UPPER]
|
||||
if self.args.smile_threshold_min >= self.args.smile_threshold_max:
|
||||
smile_degree = 0.0
|
||||
else:
|
||||
if smile_value < self.args.smile_threshold_min:
|
||||
smile_degree = 0.0
|
||||
elif smile_value > self.args.smile_threshold_max:
|
||||
smile_degree = 1.0
|
||||
else:
|
||||
smile_degree = (smile_value - self.args.smile_threshold_min) / (
|
||||
self.args.smile_threshold_max - self.args.smile_threshold_min)
|
||||
|
||||
# Eyebrow
|
||||
if True:
|
||||
brow_inner_up = blendshape_params[BROW_INNER_UP]
|
||||
brow_outer_up_right = blendshape_params[BROW_OUTER_UP_RIGHT]
|
||||
brow_outer_up_left = blendshape_params[BROW_OUTER_UP_LEFT]
|
||||
|
||||
brow_up_left = clamp(brow_inner_up + brow_outer_up_left, 0.0, 1.0)
|
||||
brow_up_right = clamp(brow_inner_up + brow_outer_up_right, 0.0, 1.0)
|
||||
pose[self.eyebrow_raised_left_index] = brow_up_left
|
||||
pose[self.eyebrow_raised_right_index] = brow_up_right
|
||||
|
||||
if self.args.eyebrow_down_max <= 0.0:
|
||||
brow_down_left = 0.0
|
||||
brow_down_right = 0.0
|
||||
else:
|
||||
brow_down_left = (1.0 - smile_degree) \
|
||||
* clamp(blendshape_params[BROW_DOWN_LEFT] / self.args.eyebrow_down_max, 0.0, 1.0)
|
||||
brow_down_right = (1.0 - smile_degree) \
|
||||
* clamp(blendshape_params[BROW_DOWN_RIGHT] / self.args.eyebrow_down_max, 0.0, 1.0)
|
||||
|
||||
if self.args.eyebrow_down_mode == EyebrowDownMode.TROUBLED:
|
||||
pose[self.eyebrow_troubled_left_index] = brow_down_left
|
||||
pose[self.eyebrow_troubled_right_index] = brow_down_right
|
||||
elif self.args.eyebrow_down_mode == EyebrowDownMode.ANGRY:
|
||||
pose[self.eyebrow_angry_left_index] = brow_down_left
|
||||
pose[self.eyebrow_angry_right_index] = brow_down_right
|
||||
elif self.args.eyebrow_down_mode == EyebrowDownMode.LOWERED:
|
||||
pose[self.eyebrow_lowered_left_index] = brow_down_left
|
||||
pose[self.eyebrow_lowered_right_index] = brow_down_right
|
||||
elif self.args.eyebrow_down_mode == EyebrowDownMode.SERIOUS:
|
||||
pose[self.eyebrow_serious_left_index] = brow_down_left
|
||||
pose[self.eyebrow_serious_right_index] = brow_down_right
|
||||
|
||||
brow_happy_value = clamp(smile_value, 0.0, 1.0) * smile_degree
|
||||
pose[self.eyebrow_happy_left_index] = brow_happy_value
|
||||
pose[self.eyebrow_happy_right_index] = brow_happy_value
|
||||
|
||||
# Eye
|
||||
if True:
|
||||
# Surprised
|
||||
if self.args.eye_surprised_max <= 0.0:
|
||||
pose[self.eye_surprised_left_index] = 0.0
|
||||
pose[self.eye_surprised_right_index] = 0.0
|
||||
else:
|
||||
pose[self.eye_surprised_left_index] = clamp(
|
||||
blendshape_params[EYE_WIDE_LEFT] / self.args.eye_surprised_max, 0.0, 1.0)
|
||||
pose[self.eye_surprised_right_index] = clamp(
|
||||
blendshape_params[EYE_WIDE_RIGHT] / self.args.eye_surprised_max, 0.0, 1.0)
|
||||
|
||||
# Wink
|
||||
if self.args.wink_mode == WinkMode.NORMAL:
|
||||
wink_left_index = self.eye_wink_left_index
|
||||
wink_right_index = self.eye_wink_right_index
|
||||
else:
|
||||
wink_left_index = self.eye_relaxed_left_index
|
||||
wink_right_index = self.eye_relaxed_right_index
|
||||
if self.args.eye_blink_max <= 0:
|
||||
pose[wink_left_index] = 0.0
|
||||
pose[wink_right_index] = 0.0
|
||||
pose[self.eye_happy_wink_left_index] = 0.0
|
||||
pose[self.eye_happy_wink_right_index] = 0.0
|
||||
else:
|
||||
pose[wink_left_index] = (1.0 - smile_degree) * clamp(
|
||||
blendshape_params[EYE_BLINK_LEFT] / self.args.eye_blink_max, 0.0, 1.0)
|
||||
pose[wink_right_index] = (1.0 - smile_degree) * clamp(
|
||||
blendshape_params[EYE_BLINK_RIGHT] / self.args.eye_blink_max, 0.0, 1.0)
|
||||
pose[self.eye_happy_wink_left_index] = smile_degree * clamp(
|
||||
blendshape_params[EYE_BLINK_LEFT] / self.args.eye_blink_max, 0.0, 1.0)
|
||||
pose[self.eye_happy_wink_right_index] = smile_degree * clamp(
|
||||
blendshape_params[EYE_BLINK_RIGHT] / self.args.eye_blink_max, 0.0, 1.0)
|
||||
|
||||
# Lower eyelid
|
||||
cheek_squint_denom = self.args.cheek_squint_max - self.args.cheek_squint_min
|
||||
if cheek_squint_denom <= 0.0:
|
||||
pose[self.eye_raised_lower_eyelid_left_index] = 0.0
|
||||
pose[self.eye_raised_lower_eyelid_right_index] = 0.0
|
||||
else:
|
||||
pose[self.eye_raised_lower_eyelid_left_index] = \
|
||||
clamp(
|
||||
(blendshape_params[CHEEK_SQUINT_LEFT] - self.args.cheek_squint_min) / cheek_squint_denom,
|
||||
0.0, 1.0)
|
||||
pose[self.eye_raised_lower_eyelid_right_index] = \
|
||||
clamp(
|
||||
(blendshape_params[CHEEK_SQUINT_RIGHT] - self.args.cheek_squint_min) / cheek_squint_denom,
|
||||
0.0, 1.0)
|
||||
|
||||
# Iris rotation
|
||||
if True:
|
||||
eye_rotation_y = (blendshape_params[EYE_LOOK_IN_LEFT]
|
||||
- blendshape_params[EYE_LOOK_OUT_LEFT]
|
||||
- blendshape_params[EYE_LOOK_IN_RIGHT]
|
||||
+ blendshape_params[EYE_LOOK_OUT_RIGHT]) / 2.0 * self.args.eye_rotation_factor
|
||||
pose[self.iris_rotation_y_index] = clamp(eye_rotation_y, -1.0, 1.0)
|
||||
|
||||
eye_rotation_x = (blendshape_params[EYE_LOOK_UP_LEFT]
|
||||
+ blendshape_params[EYE_LOOK_UP_RIGHT]
|
||||
- blendshape_params[EYE_LOOK_DOWN_LEFT]
|
||||
- blendshape_params[EYE_LOOK_DOWN_RIGHT]) / 2.0 * self.args.eye_rotation_factor
|
||||
pose[self.iris_rotation_x_index] = clamp(eye_rotation_x, -1.0, 1.0)
|
||||
|
||||
# Iris size
|
||||
if True:
|
||||
pose[self.iris_small_left_index] = self.args.iris_small_left
|
||||
pose[self.iris_small_right_index] = self.args.iris_small_right
|
||||
|
||||
# Head rotation
|
||||
if True:
|
||||
euler_angles = self.extract_euler_angles(mediapipe_face_pose)
|
||||
euler_angles[0] -= self.args.head_x_offset
|
||||
euler_angles[1] -= self.args.head_y_offset
|
||||
euler_angles[2] -= self.args.head_z_offset
|
||||
|
||||
x_param = clamp(-euler_angles[0] * 180.0 / math.pi, -15.0, 15.0) / 15.0
|
||||
pose[self.head_x_index] = x_param
|
||||
|
||||
y_param = clamp(-euler_angles[1] * 180.0 / math.pi, -10.0, 10.0) / 10.0
|
||||
pose[self.head_y_index] = y_param
|
||||
pose[self.body_y_index] = y_param
|
||||
|
||||
z_param = clamp(euler_angles[2] * 180.0 / math.pi, -15.0, 15.0) / 15.0
|
||||
pose[self.neck_z_index] = z_param
|
||||
pose[self.body_z_index] = z_param
|
||||
|
||||
# Mouth
|
||||
if True:
|
||||
jaw_open_denom = self.args.jaw_open_max - self.args.jaw_open_min
|
||||
if jaw_open_denom <= 0:
|
||||
mouth_open = 0.0
|
||||
else:
|
||||
mouth_open = clamp((blendshape_params[JAW_OPEN] - self.args.jaw_open_min) / jaw_open_denom, 0.0, 1.0)
|
||||
pose[self.mouth_aaa_index] = mouth_open
|
||||
pose[self.mouth_raised_corner_left_index] = clamp(smile_value, 0.0, 1.0)
|
||||
pose[self.mouth_raised_corner_right_index] = clamp(smile_value, 0.0, 1.0)
|
||||
|
||||
is_mouth_open = mouth_open > 0.0
|
||||
if not is_mouth_open:
|
||||
if self.args.mouth_frown_max <= 0:
|
||||
mouth_frown_value = 0.0
|
||||
else:
|
||||
mouth_frown_value = clamp(
|
||||
(blendshape_params[MOUTH_FROWN_LEFT] + blendshape_params[
|
||||
MOUTH_FROWN_RIGHT]) / self.args.mouth_frown_max, 0.0, 1.0)
|
||||
pose[self.mouth_lowered_corner_left_index] = mouth_frown_value
|
||||
pose[self.mouth_lowered_corner_right_index] = mouth_frown_value
|
||||
else:
|
||||
mouth_lower_down = clamp(
|
||||
blendshape_params[MOUTH_LOWER_DOWN_LEFT] + blendshape_params[MOUTH_LOWER_DOWN_RIGHT], 0.0, 1.0)
|
||||
mouth_funnel = blendshape_params[MOUTH_FUNNEL]
|
||||
mouth_pucker = blendshape_params[MOUTH_PUCKER]
|
||||
|
||||
mouth_point = [mouth_open, mouth_lower_down, mouth_funnel, mouth_pucker]
|
||||
|
||||
aaa_point = [1.0, 1.0, 0.0, 0.0]
|
||||
iii_point = [0.0, 1.0, 0.0, 0.0]
|
||||
uuu_point = [0.5, 0.3, 0.25, 0.75]
|
||||
ooo_point = [1.0, 0.5, 0.5, 0.4]
|
||||
|
||||
decomp = numpy.array([0, 0, 0, 0])
|
||||
M = numpy.array([
|
||||
aaa_point,
|
||||
iii_point,
|
||||
uuu_point,
|
||||
ooo_point
|
||||
])
|
||||
|
||||
def loss(decomp):
|
||||
return numpy.linalg.norm(numpy.matmul(decomp, M) - mouth_point) \
|
||||
+ 0.01 * numpy.linalg.norm(decomp, ord=1)
|
||||
|
||||
opt_result = scipy.optimize.minimize(
|
||||
loss, decomp, bounds=[(0.0, 1.0), (0.0, 1.0), (0.0, 1.0), (0.0, 1.0)])
|
||||
decomp = opt_result["x"]
|
||||
restricted_decomp = [decomp.item(0), decomp.item(1), decomp.item(2), decomp.item(3)]
|
||||
pose[self.mouth_aaa_index] = restricted_decomp[0]
|
||||
pose[self.mouth_iii_index] = restricted_decomp[1]
|
||||
mouth_funnel_denom = self.args.mouth_funnel_max - self.args.mouth_funnel_min
|
||||
if mouth_funnel_denom <= 0:
|
||||
ooo_alpha = 0.0
|
||||
uo_value = 0.0
|
||||
else:
|
||||
ooo_alpha = clamp((mouth_funnel - self.args.mouth_funnel_min) / mouth_funnel_denom, 0.0, 1.0)
|
||||
uo_value = clamp(restricted_decomp[2] + restricted_decomp[3], 0.0, 1.0)
|
||||
pose[self.mouth_uuu_index] = uo_value * (1.0 - ooo_alpha)
|
||||
pose[self.mouth_ooo_index] = uo_value * ooo_alpha
|
||||
|
||||
if self.panel is not None:
|
||||
frequency = self.breathing_frequency_slider.GetValue()
|
||||
if frequency == 0:
|
||||
value = 0.0
|
||||
pose[self.breathing_index] = value
|
||||
self.breathing_start_time = time.time()
|
||||
else:
|
||||
period = 60.0 / frequency
|
||||
now = time.time()
|
||||
diff = now - self.breathing_start_time
|
||||
frac = (diff % period) / period
|
||||
value = (-math.cos(2 * math.pi * frac) + 1.0) / 2.0
|
||||
pose[self.breathing_index] = value
|
||||
self.breathing_gauge.SetValue(int(1000 * value))
|
||||
|
||||
return pose
|
||||
Whitespace-only changes.
Whitespace-only changes.
@@ -0,0 +1,55 @@
|
||||
from typing import Optional
|
||||
|
||||
from tha4.nn.conv import create_conv7_block_from_block_args, create_conv3_block_from_block_args, \
|
||||
create_downsample_block_from_block_args, create_conv3
|
||||
from tha4.nn.resnet_block import ResnetBlock
|
||||
from tha4.nn.resnet_block_seperable import ResnetBlockSeparable
|
||||
from tha4.nn.separable_conv import create_separable_conv7_block, create_separable_conv3_block, \
|
||||
create_separable_downsample_block, create_separable_conv3
|
||||
from tha4.nn.util import BlockArgs
|
||||
|
||||
|
||||
class ConvBlockFactory:
|
||||
def __init__(self,
|
||||
block_args: BlockArgs,
|
||||
use_separable_convolution: bool = False):
|
||||
self.use_separable_convolution = use_separable_convolution
|
||||
self.block_args = block_args
|
||||
|
||||
def create_conv3(self,
|
||||
in_channels: int,
|
||||
out_channels: int,
|
||||
bias: bool,
|
||||
initialization_method: Optional[str] = None):
|
||||
if initialization_method is None:
|
||||
initialization_method = self.block_args.initialization_method
|
||||
if self.use_separable_convolution:
|
||||
return create_separable_conv3(
|
||||
in_channels, out_channels, bias, initialization_method, self.block_args.use_spectral_norm)
|
||||
else:
|
||||
return create_conv3(
|
||||
in_channels, out_channels, bias, initialization_method, self.block_args.use_spectral_norm)
|
||||
|
||||
def create_conv7_block(self, in_channels: int, out_channels: int):
|
||||
if self.use_separable_convolution:
|
||||
return create_separable_conv7_block(in_channels, out_channels, self.block_args)
|
||||
else:
|
||||
return create_conv7_block_from_block_args(in_channels, out_channels, self.block_args)
|
||||
|
||||
def create_conv3_block(self, in_channels: int, out_channels: int):
|
||||
if self.use_separable_convolution:
|
||||
return create_separable_conv3_block(in_channels, out_channels, self.block_args)
|
||||
else:
|
||||
return create_conv3_block_from_block_args(in_channels, out_channels, self.block_args)
|
||||
|
||||
def create_downsample_block(self, in_channels: int, out_channels: int, is_output_1x1: bool):
|
||||
if self.use_separable_convolution:
|
||||
return create_separable_downsample_block(in_channels, out_channels, is_output_1x1, self.block_args)
|
||||
else:
|
||||
return create_downsample_block_from_block_args(in_channels, out_channels, is_output_1x1)
|
||||
|
||||
def create_resnet_block(self, num_channels: int, is_1x1: bool):
|
||||
if self.use_separable_convolution:
|
||||
return ResnetBlockSeparable.create(num_channels, is_1x1, block_args=self.block_args)
|
||||
else:
|
||||
return ResnetBlock.create(num_channels, is_1x1, block_args=self.block_args)
|
||||
@@ -0,0 +1,68 @@
|
||||
from typing import Optional
|
||||
|
||||
from torch.nn import Sigmoid, Sequential, Tanh
|
||||
|
||||
from tha4.nn.conv import create_conv3, create_conv3_from_block_args
|
||||
from tha4.nn.nonlinearity_factory import ReLUFactory
|
||||
from tha4.nn.normalization import InstanceNorm2dFactory
|
||||
from tha4.nn.util import BlockArgs
|
||||
|
||||
|
||||
class PoserArgs00:
|
||||
def __init__(self,
|
||||
image_size: int,
|
||||
input_image_channels: int,
|
||||
output_image_channels: int,
|
||||
start_channels: int,
|
||||
num_pose_params: int,
|
||||
block_args: Optional[BlockArgs] = None):
|
||||
self.num_pose_params = num_pose_params
|
||||
self.start_channels = start_channels
|
||||
self.output_image_channels = output_image_channels
|
||||
self.input_image_channels = input_image_channels
|
||||
self.image_size = image_size
|
||||
if block_args is None:
|
||||
self.block_args = BlockArgs(
|
||||
normalization_layer_factory=InstanceNorm2dFactory(),
|
||||
nonlinearity_factory=ReLUFactory(inplace=True))
|
||||
else:
|
||||
self.block_args = block_args
|
||||
|
||||
def create_alpha_block(self):
|
||||
from torch.nn import Sequential
|
||||
return Sequential(
|
||||
create_conv3(
|
||||
in_channels=self.start_channels,
|
||||
out_channels=1,
|
||||
bias=True,
|
||||
initialization_method=self.block_args.initialization_method,
|
||||
use_spectral_norm=False),
|
||||
Sigmoid())
|
||||
|
||||
def create_all_channel_alpha_block(self):
|
||||
from torch.nn import Sequential
|
||||
return Sequential(
|
||||
create_conv3(
|
||||
in_channels=self.start_channels,
|
||||
out_channels=self.output_image_channels,
|
||||
bias=True,
|
||||
initialization_method=self.block_args.initialization_method,
|
||||
use_spectral_norm=False),
|
||||
Sigmoid())
|
||||
|
||||
def create_color_change_block(self):
|
||||
return Sequential(
|
||||
create_conv3_from_block_args(
|
||||
in_channels=self.start_channels,
|
||||
out_channels=self.output_image_channels,
|
||||
bias=True,
|
||||
block_args=self.block_args),
|
||||
Tanh())
|
||||
|
||||
def create_grid_change_block(self):
|
||||
return create_conv3(
|
||||
in_channels=self.start_channels,
|
||||
out_channels=2,
|
||||
bias=False,
|
||||
initialization_method='zero',
|
||||
use_spectral_norm=False)
|
||||
@@ -0,0 +1,121 @@
|
||||
import math
|
||||
from typing import Optional, List
|
||||
|
||||
import torch
|
||||
from torch import Tensor
|
||||
from torch.nn import ModuleList, Module
|
||||
|
||||
from tha4.nn.common.poser_args import PoserArgs00
|
||||
from tha4.nn.conv import create_conv3_block_from_block_args, create_downsample_block_from_block_args, \
|
||||
create_upsample_block_from_block_args
|
||||
from tha4.nn.nonlinearity_factory import ReLUFactory
|
||||
from tha4.nn.normalization import InstanceNorm2dFactory
|
||||
from tha4.nn.resnet_block import ResnetBlock
|
||||
from tha4.nn.util import BlockArgs
|
||||
|
||||
|
||||
class PoserEncoderDecoder00Args(PoserArgs00):
|
||||
def __init__(self,
|
||||
image_size: int,
|
||||
input_image_channels: int,
|
||||
output_image_channels: int,
|
||||
num_pose_params: int ,
|
||||
start_channels: int,
|
||||
bottleneck_image_size,
|
||||
num_bottleneck_blocks,
|
||||
max_channels: int,
|
||||
block_args: Optional[BlockArgs] = None):
|
||||
super().__init__(
|
||||
image_size, input_image_channels, output_image_channels, start_channels, num_pose_params, block_args)
|
||||
self.max_channels = max_channels
|
||||
self.num_bottleneck_blocks = num_bottleneck_blocks
|
||||
self.bottleneck_image_size = bottleneck_image_size
|
||||
assert bottleneck_image_size > 1
|
||||
|
||||
if block_args is None:
|
||||
self.block_args = BlockArgs(
|
||||
normalization_layer_factory=InstanceNorm2dFactory(),
|
||||
nonlinearity_factory=ReLUFactory(inplace=True))
|
||||
else:
|
||||
self.block_args = block_args
|
||||
|
||||
|
||||
class PoserEncoderDecoder00(Module):
|
||||
def __init__(self, args: PoserEncoderDecoder00Args):
|
||||
super().__init__()
|
||||
self.args = args
|
||||
|
||||
self.num_levels = int(math.log2(args.image_size // args.bottleneck_image_size)) + 1
|
||||
|
||||
self.downsample_blocks = ModuleList()
|
||||
self.downsample_blocks.append(
|
||||
create_conv3_block_from_block_args(
|
||||
args.input_image_channels,
|
||||
args.start_channels,
|
||||
args.block_args))
|
||||
current_image_size = args.image_size
|
||||
current_num_channels = args.start_channels
|
||||
while current_image_size > args.bottleneck_image_size:
|
||||
next_image_size = current_image_size // 2
|
||||
next_num_channels = self.get_num_output_channels_from_image_size(next_image_size)
|
||||
self.downsample_blocks.append(create_downsample_block_from_block_args(
|
||||
in_channels=current_num_channels,
|
||||
out_channels=next_num_channels,
|
||||
is_output_1x1=False,
|
||||
block_args=args.block_args))
|
||||
current_image_size = next_image_size
|
||||
current_num_channels = next_num_channels
|
||||
assert len(self.downsample_blocks) == self.num_levels
|
||||
|
||||
self.bottleneck_blocks = ModuleList()
|
||||
self.bottleneck_blocks.append(create_conv3_block_from_block_args(
|
||||
in_channels=current_num_channels + args.num_pose_params,
|
||||
out_channels=current_num_channels,
|
||||
block_args=args.block_args))
|
||||
for i in range(1, args.num_bottleneck_blocks):
|
||||
self.bottleneck_blocks.append(
|
||||
ResnetBlock.create(
|
||||
num_channels=current_num_channels,
|
||||
is1x1=False,
|
||||
block_args=args.block_args))
|
||||
|
||||
self.upsample_blocks = ModuleList()
|
||||
while current_image_size < args.image_size:
|
||||
next_image_size = current_image_size * 2
|
||||
next_num_channels = self.get_num_output_channels_from_image_size(next_image_size)
|
||||
self.upsample_blocks.append(create_upsample_block_from_block_args(
|
||||
in_channels=current_num_channels,
|
||||
out_channels=next_num_channels,
|
||||
block_args=args.block_args))
|
||||
current_image_size = next_image_size
|
||||
current_num_channels = next_num_channels
|
||||
|
||||
def get_num_output_channels_from_level(self, level: int):
|
||||
return self.get_num_output_channels_from_image_size(self.args.image_size // (2 ** level))
|
||||
|
||||
def get_num_output_channels_from_image_size(self, image_size: int):
|
||||
return min(self.args.start_channels * (self.args.image_size // image_size), self.args.max_channels)
|
||||
|
||||
def forward(self, image: Tensor, pose: Optional[Tensor] = None) -> List[Tensor]:
|
||||
if self.args.num_pose_params != 0:
|
||||
assert pose is not None
|
||||
else:
|
||||
assert pose is None
|
||||
outputs = []
|
||||
feature = image
|
||||
outputs.append(feature)
|
||||
for block in self.downsample_blocks:
|
||||
feature = block(feature)
|
||||
outputs.append(feature)
|
||||
if pose is not None:
|
||||
n, c = pose.shape
|
||||
pose = pose.view(n, c, 1, 1).repeat(1, 1, self.args.bottleneck_image_size, self.args.bottleneck_image_size)
|
||||
feature = torch.cat([feature, pose], dim=1)
|
||||
for block in self.bottleneck_blocks:
|
||||
feature = block(feature)
|
||||
outputs.append(feature)
|
||||
for block in self.upsample_blocks:
|
||||
feature = block(feature)
|
||||
outputs.append(feature)
|
||||
outputs.reverse()
|
||||
return outputs
|
||||
@@ -0,0 +1,92 @@
|
||||
import math
|
||||
from typing import Optional, List
|
||||
|
||||
import torch
|
||||
from torch import Tensor
|
||||
from torch.nn import ModuleList, Module
|
||||
|
||||
from tha4.nn.common.poser_encoder_decoder_00 import PoserEncoderDecoder00Args
|
||||
from tha4.nn.resnet_block_seperable import ResnetBlockSeparable
|
||||
from tha4.nn.separable_conv import create_separable_conv3_block, create_separable_downsample_block, \
|
||||
create_separable_upsample_block
|
||||
|
||||
|
||||
class PoserEncoderDecoder00Separable(Module):
|
||||
def __init__(self, args: PoserEncoderDecoder00Args):
|
||||
super().__init__()
|
||||
self.args = args
|
||||
|
||||
self.num_levels = int(math.log2(args.image_size // args.bottleneck_image_size)) + 1
|
||||
|
||||
self.downsample_blocks = ModuleList()
|
||||
self.downsample_blocks.append(
|
||||
create_separable_conv3_block(
|
||||
args.input_image_channels,
|
||||
args.start_channels,
|
||||
args.block_args))
|
||||
current_image_size = args.image_size
|
||||
current_num_channels = args.start_channels
|
||||
while current_image_size > args.bottleneck_image_size:
|
||||
next_image_size = current_image_size // 2
|
||||
next_num_channels = self.get_num_output_channels_from_image_size(next_image_size)
|
||||
self.downsample_blocks.append(create_separable_downsample_block(
|
||||
in_channels=current_num_channels,
|
||||
out_channels=next_num_channels,
|
||||
is_output_1x1=False,
|
||||
block_args=args.block_args))
|
||||
current_image_size = next_image_size
|
||||
current_num_channels = next_num_channels
|
||||
assert len(self.downsample_blocks) == self.num_levels
|
||||
|
||||
self.bottleneck_blocks = ModuleList()
|
||||
self.bottleneck_blocks.append(create_separable_conv3_block(
|
||||
in_channels=current_num_channels + args.num_pose_params,
|
||||
out_channels=current_num_channels,
|
||||
block_args=args.block_args))
|
||||
for i in range(1, args.num_bottleneck_blocks):
|
||||
self.bottleneck_blocks.append(
|
||||
ResnetBlockSeparable.create(
|
||||
num_channels=current_num_channels,
|
||||
is1x1=False,
|
||||
block_args=args.block_args))
|
||||
|
||||
self.upsample_blocks = ModuleList()
|
||||
while current_image_size < args.image_size:
|
||||
next_image_size = current_image_size * 2
|
||||
next_num_channels = self.get_num_output_channels_from_image_size(next_image_size)
|
||||
self.upsample_blocks.append(create_separable_upsample_block(
|
||||
in_channels=current_num_channels,
|
||||
out_channels=next_num_channels,
|
||||
block_args=args.block_args))
|
||||
current_image_size = next_image_size
|
||||
current_num_channels = next_num_channels
|
||||
|
||||
def get_num_output_channels_from_level(self, level: int):
|
||||
return self.get_num_output_channels_from_image_size(self.args.image_size // (2 ** level))
|
||||
|
||||
def get_num_output_channels_from_image_size(self, image_size: int):
|
||||
return min(self.args.start_channels * (self.args.image_size // image_size), self.args.max_channels)
|
||||
|
||||
def forward(self, image: Tensor, pose: Optional[Tensor] = None) -> List[Tensor]:
|
||||
if self.args.num_pose_params != 0:
|
||||
assert pose is not None
|
||||
else:
|
||||
assert pose is None
|
||||
outputs = []
|
||||
feature = image
|
||||
outputs.append(feature)
|
||||
for block in self.downsample_blocks:
|
||||
feature = block(feature)
|
||||
outputs.append(feature)
|
||||
if pose is not None:
|
||||
n, c = pose.shape
|
||||
pose = pose.view(n, c, 1, 1).repeat(1, 1, self.args.bottleneck_image_size, self.args.bottleneck_image_size)
|
||||
feature = torch.cat([feature, pose], dim=1)
|
||||
for block in self.bottleneck_blocks:
|
||||
feature = block(feature)
|
||||
outputs.append(feature)
|
||||
for block in self.upsample_blocks:
|
||||
feature = block(feature)
|
||||
outputs.append(feature)
|
||||
outputs.reverse()
|
||||
return outputs
|
||||
@@ -0,0 +1,100 @@
|
||||
import math
|
||||
from typing import Optional, List
|
||||
|
||||
import torch
|
||||
from torch import Tensor
|
||||
from torch.nn import Module, ModuleList, Sequential, Upsample
|
||||
|
||||
from tha4.nn.common.conv_block_factory import ConvBlockFactory
|
||||
from tha4.nn.nonlinearity_factory import LeakyReLUFactory
|
||||
from tha4.nn.normalization import InstanceNorm2dFactory
|
||||
from tha4.nn.util import BlockArgs
|
||||
|
||||
|
||||
class ResizeConvEncoderDecoderArgs:
|
||||
def __init__(self,
|
||||
image_size: int,
|
||||
input_channels: int,
|
||||
start_channels: int,
|
||||
bottleneck_image_size,
|
||||
num_bottleneck_blocks,
|
||||
max_channels: int,
|
||||
block_args: Optional[BlockArgs] = None,
|
||||
upsample_mode: str = 'bilinear',
|
||||
use_separable_convolution=False):
|
||||
self.use_separable_convolution = use_separable_convolution
|
||||
self.upsample_mode = upsample_mode
|
||||
self.block_args = block_args
|
||||
self.max_channels = max_channels
|
||||
self.num_bottleneck_blocks = num_bottleneck_blocks
|
||||
self.bottleneck_image_size = bottleneck_image_size
|
||||
self.start_channels = start_channels
|
||||
self.image_size = image_size
|
||||
self.input_channels = input_channels
|
||||
|
||||
|
||||
class ResizeConvEncoderDecoder(Module):
|
||||
def __init__(self, args: ResizeConvEncoderDecoderArgs):
|
||||
super().__init__()
|
||||
self.args = args
|
||||
|
||||
self.num_levels = int(math.log2(args.image_size // args.bottleneck_image_size)) + 1
|
||||
|
||||
conv_block_factory = ConvBlockFactory(args.block_args, args.use_separable_convolution)
|
||||
|
||||
self.downsample_blocks = ModuleList()
|
||||
self.downsample_blocks.append(conv_block_factory.create_conv7_block(args.input_channels, args.start_channels))
|
||||
current_image_size = args.image_size
|
||||
current_num_channels = args.start_channels
|
||||
while current_image_size > args.bottleneck_image_size:
|
||||
next_image_size = current_image_size // 2
|
||||
next_num_channels = self.get_num_output_channels_from_image_size(next_image_size)
|
||||
self.downsample_blocks.append(conv_block_factory.create_downsample_block(
|
||||
in_channels=current_num_channels,
|
||||
out_channels=next_num_channels,
|
||||
is_output_1x1=False))
|
||||
current_image_size = next_image_size
|
||||
current_num_channels = next_num_channels
|
||||
assert len(self.downsample_blocks) == self.num_levels
|
||||
|
||||
self.bottleneck_blocks = ModuleList()
|
||||
for i in range(args.num_bottleneck_blocks):
|
||||
self.bottleneck_blocks.append(conv_block_factory.create_resnet_block(current_num_channels, is_1x1=False))
|
||||
|
||||
self.output_image_sizes = [current_image_size]
|
||||
self.output_num_channels = [current_num_channels]
|
||||
self.upsample_blocks = ModuleList()
|
||||
if args.upsample_mode == 'nearest':
|
||||
align_corners = None
|
||||
else:
|
||||
align_corners = False
|
||||
while current_image_size < args.image_size:
|
||||
next_image_size = current_image_size * 2
|
||||
next_num_channels = self.get_num_output_channels_from_image_size(next_image_size)
|
||||
self.upsample_blocks.append(
|
||||
Sequential(
|
||||
Upsample(scale_factor=2, mode=args.upsample_mode, align_corners=align_corners),
|
||||
conv_block_factory.create_conv3_block(
|
||||
in_channels=current_num_channels, out_channels=next_num_channels)))
|
||||
current_image_size = next_image_size
|
||||
current_num_channels = next_num_channels
|
||||
self.output_image_sizes.append(current_image_size)
|
||||
self.output_num_channels.append(current_num_channels)
|
||||
|
||||
def get_num_output_channels_from_level(self, level: int):
|
||||
return self.get_num_output_channels_from_image_size(self.args.image_size // (2 ** level))
|
||||
|
||||
def get_num_output_channels_from_image_size(self, image_size: int):
|
||||
return min(self.args.start_channels * (self.args.image_size // image_size), self.args.max_channels)
|
||||
|
||||
def forward(self, feature: Tensor) -> List[Tensor]:
|
||||
outputs = []
|
||||
for block in self.downsample_blocks:
|
||||
feature = block(feature)
|
||||
for block in self.bottleneck_blocks:
|
||||
feature = block(feature)
|
||||
outputs.append(feature)
|
||||
for block in self.upsample_blocks:
|
||||
feature = block(feature)
|
||||
outputs.append(feature)
|
||||
return outputs
|
||||
@@ -0,0 +1,107 @@
|
||||
from typing import Optional, List
|
||||
|
||||
import torch
|
||||
from torch import Tensor
|
||||
from torch.nn import ModuleList, Module, Upsample
|
||||
|
||||
from tha4.nn.common.conv_block_factory import ConvBlockFactory
|
||||
from tha4.nn.nonlinearity_factory import ReLUFactory
|
||||
from tha4.nn.normalization import InstanceNorm2dFactory
|
||||
from tha4.nn.util import BlockArgs
|
||||
|
||||
|
||||
class ResizeConvUNetArgs:
|
||||
def __init__(self,
|
||||
image_size: int,
|
||||
input_channels: int,
|
||||
start_channels: int,
|
||||
bottleneck_image_size: int,
|
||||
num_bottleneck_blocks: int,
|
||||
max_channels: int,
|
||||
upsample_mode: str = 'bilinear',
|
||||
block_args: Optional[BlockArgs] = None,
|
||||
use_separable_convolution: bool = False):
|
||||
if block_args is None:
|
||||
block_args = BlockArgs(
|
||||
normalization_layer_factory=InstanceNorm2dFactory(),
|
||||
nonlinearity_factory=ReLUFactory(inplace=False))
|
||||
|
||||
self.use_separable_convolution = use_separable_convolution
|
||||
self.block_args = block_args
|
||||
self.upsample_mode = upsample_mode
|
||||
self.max_channels = max_channels
|
||||
self.num_bottleneck_blocks = num_bottleneck_blocks
|
||||
self.bottleneck_image_size = bottleneck_image_size
|
||||
self.input_channels = input_channels
|
||||
self.start_channels = start_channels
|
||||
self.image_size = image_size
|
||||
|
||||
|
||||
class ResizeConvUNet(Module):
|
||||
def __init__(self, args: ResizeConvUNetArgs):
|
||||
super().__init__()
|
||||
self.args = args
|
||||
conv_block_factory = ConvBlockFactory(args.block_args, args.use_separable_convolution)
|
||||
|
||||
self.downsample_blocks = ModuleList()
|
||||
self.downsample_blocks.append(conv_block_factory.create_conv3_block(
|
||||
self.args.input_channels,
|
||||
self.args.start_channels))
|
||||
current_channels = self.args.start_channels
|
||||
current_size = self.args.image_size
|
||||
|
||||
size_to_channel = {
|
||||
current_size: current_channels
|
||||
}
|
||||
while current_size > self.args.bottleneck_image_size:
|
||||
next_size = current_size // 2
|
||||
next_channels = min(self.args.max_channels, current_channels * 2)
|
||||
self.downsample_blocks.append(conv_block_factory.create_downsample_block(
|
||||
current_channels,
|
||||
next_channels,
|
||||
is_output_1x1=False))
|
||||
current_size = next_size
|
||||
current_channels = next_channels
|
||||
size_to_channel[current_size] = current_channels
|
||||
|
||||
self.bottleneck_blocks = ModuleList()
|
||||
for i in range(self.args.num_bottleneck_blocks):
|
||||
self.bottleneck_blocks.append(conv_block_factory.create_resnet_block(current_channels, is_1x1=False))
|
||||
|
||||
self.output_image_sizes = [current_size]
|
||||
self.output_num_channels = [current_channels]
|
||||
self.upsample_blocks = ModuleList()
|
||||
while current_size < self.args.image_size:
|
||||
next_size = current_size * 2
|
||||
next_channels = size_to_channel[next_size]
|
||||
self.upsample_blocks.append(conv_block_factory.create_conv3_block(
|
||||
current_channels + next_channels,
|
||||
next_channels))
|
||||
current_size = next_size
|
||||
current_channels = next_channels
|
||||
self.output_image_sizes.append(current_size)
|
||||
self.output_num_channels.append(current_channels)
|
||||
|
||||
if args.upsample_mode == 'nearest':
|
||||
align_corners = None
|
||||
else:
|
||||
align_corners = False
|
||||
self.double_resolution = Upsample(scale_factor=2, mode=args.upsample_mode, align_corners=align_corners)
|
||||
|
||||
def forward(self, feature: Tensor) -> List[Tensor]:
|
||||
downsampled_features = []
|
||||
for block in self.downsample_blocks:
|
||||
feature = block(feature)
|
||||
downsampled_features.append(feature)
|
||||
|
||||
for block in self.bottleneck_blocks:
|
||||
feature = block(feature)
|
||||
|
||||
outputs = [feature]
|
||||
for i in range(0, len(self.upsample_blocks)):
|
||||
feature = self.double_resolution(feature)
|
||||
feature = torch.cat([feature, downsampled_features[-i - 2]], dim=1)
|
||||
feature = self.upsample_blocks[i](feature)
|
||||
outputs.append(feature)
|
||||
|
||||
return outputs
|
||||
@@ -0,0 +1,658 @@
|
||||
import math
|
||||
from enum import Enum
|
||||
from typing import Optional, List
|
||||
|
||||
import torch
|
||||
from torch import zero_, Tensor
|
||||
from torch.nn import Module, GroupNorm, Sequential, SiLU, Conv2d, AvgPool2d, Linear, Dropout, ModuleList
|
||||
from torch.nn.functional import interpolate
|
||||
|
||||
from tha4.shion.core.module_factory import ModuleFactory
|
||||
|
||||
|
||||
class Identity(Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
|
||||
def forward(self, x):
|
||||
return x
|
||||
|
||||
|
||||
class IdentityFactory(ModuleFactory):
|
||||
def create(self) -> Module:
|
||||
return Identity()
|
||||
|
||||
|
||||
def init_to_zero(module: Module):
|
||||
with torch.no_grad():
|
||||
zero_(module.weight)
|
||||
zero_(module.bias)
|
||||
return module
|
||||
|
||||
|
||||
class Upsample(Module):
|
||||
def __init__(self, in_channels: int, out_channels: Optional[int] = None, use_conv: bool = False):
|
||||
super().__init__()
|
||||
if out_channels is None:
|
||||
out_channels = in_channels
|
||||
self.in_channels = in_channels
|
||||
if use_conv or in_channels != out_channels:
|
||||
self.postprocess = Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1)
|
||||
else:
|
||||
self.postprocess = Identity()
|
||||
|
||||
def forward(self, x):
|
||||
assert x.shape[1] == self.in_channels
|
||||
return self.postprocess(interpolate(x, scale_factor=2, mode="nearest"))
|
||||
|
||||
|
||||
class Downsample(Module):
|
||||
def __init__(self, in_channels: int, out_channels: Optional[int] = None, use_conv: bool = False):
|
||||
super().__init__()
|
||||
if out_channels is None:
|
||||
out_channels = in_channels
|
||||
self.in_channels = in_channels
|
||||
if use_conv or in_channels != out_channels:
|
||||
self.op = Conv2d(in_channels, out_channels, kernel_size=3, stride=2, padding=1)
|
||||
else:
|
||||
self.op = AvgPool2d(kernel_size=2, stride=2)
|
||||
|
||||
def forward(self, x):
|
||||
assert x.shape[1] == self.in_channels
|
||||
return self.op(x)
|
||||
|
||||
|
||||
def GroupNorm32(channels):
|
||||
return GroupNorm(min(32, channels), channels)
|
||||
|
||||
|
||||
class SamplingMode(Enum):
|
||||
SAME_RESOLUTION = 0
|
||||
UPSAMPLING = 1
|
||||
DOWNSAMPING = 2
|
||||
|
||||
|
||||
class ResBlockArgs:
|
||||
def __init__(self,
|
||||
dropout_prob: float,
|
||||
use_cond0: bool = True,
|
||||
use_cond1: bool = False,
|
||||
init_conditioned_residual_to_zero: bool = False,
|
||||
use_conv_on_skip_connection: bool = False):
|
||||
assert not use_cond1 or use_cond0
|
||||
self.use_conv_on_skip_connection = use_conv_on_skip_connection
|
||||
self.use_cond1 = use_cond1
|
||||
self.use_cond0 = use_cond0
|
||||
self.init_conditioned_residual_to_zero = init_conditioned_residual_to_zero
|
||||
self.dropout_prob = dropout_prob
|
||||
|
||||
|
||||
def apply_scaleshift(x: Tensor, scaleshift: Tensor, condition_bias: float = 1.0) -> Tensor:
|
||||
assert len(scaleshift.shape) == 2
|
||||
assert len(x.shape) == 4
|
||||
assert x.shape[0] == scaleshift.shape[0]
|
||||
assert 2 * x.shape[1] == scaleshift.shape[1]
|
||||
scaleshift = scaleshift.reshape(scaleshift.shape[0], scaleshift.shape[1], 1, 1)
|
||||
scale, shift = torch.chunk(scaleshift, 2, dim=1)
|
||||
return x * (condition_bias + scale) + shift
|
||||
|
||||
|
||||
class ResBlock(Module):
|
||||
def __init__(self,
|
||||
in_channels: int,
|
||||
out_channels: int,
|
||||
cond0_channels: Optional[int] = None,
|
||||
cond1_channels: Optional[int] = None,
|
||||
sampling_mode: SamplingMode = SamplingMode.SAME_RESOLUTION,
|
||||
dropout_prob: float = 0.1,
|
||||
condition_bias: float = 1.0):
|
||||
super().__init__()
|
||||
assert cond0_channels is not None or cond1_channels is None
|
||||
|
||||
self.in_channels = in_channels
|
||||
self.out_channels = out_channels
|
||||
self.sampling_mode = sampling_mode
|
||||
self.cond0_channels = cond0_channels
|
||||
self.cond1_channels = cond1_channels
|
||||
self.condition_bias = condition_bias
|
||||
|
||||
if sampling_mode == SamplingMode.UPSAMPLING:
|
||||
self.x_resample = Upsample(in_channels)
|
||||
self.h_resample = Upsample(in_channels)
|
||||
elif sampling_mode == SamplingMode.DOWNSAMPING:
|
||||
self.x_resample = Downsample(in_channels)
|
||||
self.h_resample = Downsample(in_channels)
|
||||
else:
|
||||
self.x_resample = Identity()
|
||||
self.h_resample = Identity()
|
||||
|
||||
self.nonlinear = SiLU()
|
||||
|
||||
# Layers before conditioning
|
||||
self.norm0 = GroupNorm32(in_channels)
|
||||
self.conv0 = Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1)
|
||||
|
||||
# Conditioning layers
|
||||
if cond0_channels is not None:
|
||||
self.cond0_layers = Sequential(
|
||||
SiLU(),
|
||||
Linear(cond0_channels, 2 * out_channels))
|
||||
self.norm1 = GroupNorm32(out_channels)
|
||||
self.dropout = Dropout(dropout_prob)
|
||||
self.conv1 = init_to_zero(Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1))
|
||||
if cond1_channels is not None:
|
||||
self.cond1_layers = Sequential(
|
||||
SiLU(),
|
||||
Linear(cond0_channels, 2 * out_channels))
|
||||
|
||||
# Skip layer
|
||||
if in_channels == out_channels:
|
||||
self.skip = Identity()
|
||||
else:
|
||||
self.skip = Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0)
|
||||
|
||||
def forward(self, x: Tensor, cond0: Optional[Tensor] = None, cond1: Optional[Tensor] = None) -> Tensor:
|
||||
assert self.cond0_channels is None or cond0 is not None
|
||||
assert self.cond1_channels is None or cond1 is not None
|
||||
|
||||
h = self.conv0(self.h_resample(self.nonlinear(self.norm0(x))))
|
||||
if self.cond0_channels is not None:
|
||||
h = self.norm1(h)
|
||||
h = apply_scaleshift(h, self.cond0_layers(cond0), self.condition_bias)
|
||||
if self.cond1_channels is not None:
|
||||
h = apply_scaleshift(h, self.cond1_layers(cond1), self.condition_bias)
|
||||
h = self.conv1(self.dropout(self.nonlinear(h)))
|
||||
return self.skip(self.x_resample(x)) + h
|
||||
|
||||
|
||||
class AttentionBlockArgs:
|
||||
def __init__(self,
|
||||
num_heads: Optional[int] = 1,
|
||||
num_head_channels: Optional[int] = None,
|
||||
use_new_attention_order: bool = False):
|
||||
self.use_new_attention_order = use_new_attention_order
|
||||
self.num_head_channels = num_head_channels
|
||||
self.num_heads = num_heads
|
||||
|
||||
|
||||
def qkv_attention_legacy(qkv: torch.Tensor, num_heads: int):
|
||||
assert len(qkv.shape) == 3
|
||||
B, W, L = qkv.shape
|
||||
H = num_heads
|
||||
assert W % (3 * H) == 0
|
||||
C = W // (3 * H)
|
||||
q, k, v = qkv.reshape(B * H, C * 3, L).split(C, dim=1)
|
||||
scale = 1.0 / math.sqrt(math.sqrt(C))
|
||||
weight = torch.einsum('bct,bcs->bts', q * scale, k * scale)
|
||||
weight = torch.softmax(weight, dim=-1)
|
||||
output = torch.einsum("bts,bcs->bct", weight, v)
|
||||
return output.reshape(B, H * C, L)
|
||||
|
||||
|
||||
def qkv_attention(qkv: torch.Tensor, num_heads: int):
|
||||
B, W, L = qkv.shape
|
||||
H = num_heads
|
||||
assert W % (3 * H) == 0
|
||||
C = W // (3 * H)
|
||||
q, k, v = qkv.chunk(3, dim=1)
|
||||
scale = 1.0 / math.sqrt(math.sqrt(C))
|
||||
weight = torch.einsum("bct,bcs->bts", (q * scale).view(B * H, C, L), (k * scale).view(B * H, C, L))
|
||||
weight = torch.softmax(weight, dim=-1)
|
||||
output = torch.einsum("bts,bcs->bct", weight, v.reshape(B * H, C, L))
|
||||
return output.reshape(B, H * C, L)
|
||||
|
||||
|
||||
class AttentionBlock(Module):
|
||||
def __init__(self,
|
||||
num_channels: int,
|
||||
args: AttentionBlockArgs):
|
||||
super().__init__()
|
||||
self.use_new_attention_order = args.use_new_attention_order
|
||||
|
||||
if args.num_head_channels is None:
|
||||
assert args.num_heads is not None
|
||||
assert num_channels % args.num_heads == 0
|
||||
self.num_heads = args.num_heads
|
||||
self.num_head_channels = num_channels // self.num_heads
|
||||
elif args.num_heads is None:
|
||||
assert args.num_head_channels is not None
|
||||
assert num_channels % args.num_head_channels == 0
|
||||
self.num_heads = num_channels // args.num_head_channels
|
||||
self.num_head_channels = args.num_head_channels
|
||||
|
||||
self.norm = GroupNorm32(num_channels)
|
||||
self.qkv = Conv2d(num_channels, 3 * num_channels, kernel_size=1, stride=1, padding=0)
|
||||
self.conv = Conv2d(num_channels, num_channels, kernel_size=1, stride=1, padding=0)
|
||||
with torch.no_grad():
|
||||
zero_(self.conv.weight)
|
||||
zero_(self.conv.bias)
|
||||
|
||||
def forward(self, x: torch.Tensor):
|
||||
assert len(x.shape) == 4
|
||||
B, C, H, W = x.shape
|
||||
qkv = self.qkv(self.norm(x)).reshape(B, 3 * C, H * W)
|
||||
if self.use_new_attention_order:
|
||||
h = qkv_attention(qkv, self.num_heads)
|
||||
else:
|
||||
h = qkv_attention_legacy(qkv, self.num_heads)
|
||||
h = self.conv(h.reshape(B, C, H, W))
|
||||
return x + h
|
||||
|
||||
|
||||
class Arity3To1(Module):
|
||||
def __init__(self, module: Module):
|
||||
super().__init__()
|
||||
self.module = module
|
||||
|
||||
def forward(self, x: Tensor, y: Optional[Tensor] = None, z: Optional[Tensor] = None):
|
||||
return self.module(x)
|
||||
|
||||
|
||||
class DownsamplingBlock(Module):
|
||||
def __init__(self,
|
||||
in_channels: int,
|
||||
out_channels: int,
|
||||
cond0_channels: Optional[int],
|
||||
cond1_channels: Optional[int],
|
||||
num_res_blocks: int,
|
||||
dropout_prob: float,
|
||||
use_attention: bool,
|
||||
perform_downsampling: bool,
|
||||
resample_with_res_block: bool,
|
||||
use_conv_to_resample: bool,
|
||||
attention_block_args: AttentionBlockArgs,
|
||||
condition_bias: float = 1.0):
|
||||
super().__init__()
|
||||
self.use_attention = use_attention
|
||||
self.res_blocks = ModuleList()
|
||||
self.attention_blocks = ModuleList()
|
||||
self.perform_downsampling = perform_downsampling
|
||||
self.output_channels = []
|
||||
for j in range(num_res_blocks):
|
||||
self.res_blocks.append(ResBlock(
|
||||
in_channels=in_channels if j == 0 else out_channels,
|
||||
out_channels=out_channels,
|
||||
cond0_channels=cond0_channels,
|
||||
cond1_channels=cond1_channels,
|
||||
dropout_prob=dropout_prob,
|
||||
condition_bias=condition_bias))
|
||||
if use_attention:
|
||||
self.attention_blocks.append(AttentionBlock(out_channels, attention_block_args))
|
||||
self.output_channels.append(out_channels)
|
||||
if perform_downsampling:
|
||||
if resample_with_res_block:
|
||||
self.downsample = ResBlock(
|
||||
in_channels=out_channels,
|
||||
out_channels=out_channels,
|
||||
cond0_channels=cond0_channels,
|
||||
cond1_channels=cond1_channels,
|
||||
dropout_prob=dropout_prob,
|
||||
sampling_mode=SamplingMode.DOWNSAMPING,
|
||||
condition_bias=condition_bias)
|
||||
else:
|
||||
self.downsample = Arity3To1(Downsample(out_channels, use_conv_to_resample))
|
||||
self.output_channels.append(out_channels)
|
||||
|
||||
def forward(self, h: Tensor, cond0: Optional[Tensor] = None, cond1: Optional[Tensor] = None) -> List[Tensor]:
|
||||
hs = []
|
||||
for i in range(len(self.res_blocks)):
|
||||
h = self.res_blocks[i].forward(h, cond0, cond1)
|
||||
if self.use_attention:
|
||||
h = self.attention_blocks[i].forward(h)
|
||||
hs.append(h)
|
||||
if self.perform_downsampling:
|
||||
hs.append(self.downsample(h, cond0, cond1))
|
||||
return hs
|
||||
|
||||
|
||||
class UpsamplingBlock(Module):
|
||||
def __init__(self,
|
||||
in_channels: int,
|
||||
out_channels: int,
|
||||
cond0_channels: Optional[int],
|
||||
cond1_channels: Optional[int],
|
||||
num_resnet_blocks: int,
|
||||
skip_channels: List[int],
|
||||
dropout_prob: float,
|
||||
use_attention: bool,
|
||||
perform_upsampling: bool,
|
||||
resample_with_res_block: bool,
|
||||
use_conv_to_resample: bool,
|
||||
attention_block_args: AttentionBlockArgs,
|
||||
condition_bias: float = 1.0):
|
||||
super().__init__()
|
||||
self.use_attention = use_attention
|
||||
self.resnet_blocks = ModuleList()
|
||||
self.attention_blocks = ModuleList()
|
||||
self.perform_upsampling = perform_upsampling
|
||||
for i in range(num_resnet_blocks):
|
||||
self.resnet_blocks.append(ResBlock(
|
||||
in_channels=(in_channels if i == 0 else out_channels) + skip_channels[i],
|
||||
out_channels=out_channels,
|
||||
cond0_channels=cond0_channels,
|
||||
cond1_channels=cond1_channels,
|
||||
dropout_prob=dropout_prob,
|
||||
condition_bias=condition_bias))
|
||||
if use_attention:
|
||||
self.attention_blocks.append(AttentionBlock(out_channels, attention_block_args))
|
||||
if perform_upsampling:
|
||||
if resample_with_res_block:
|
||||
self.upsample = ResBlock(
|
||||
in_channels=out_channels,
|
||||
out_channels=out_channels,
|
||||
cond0_channels=cond0_channels,
|
||||
cond1_channels=cond1_channels,
|
||||
sampling_mode=SamplingMode.UPSAMPLING,
|
||||
dropout_prob=dropout_prob,
|
||||
condition_bias=condition_bias)
|
||||
else:
|
||||
self.upsample = Arity3To1(Upsample(out_channels, use_conv_to_resample))
|
||||
|
||||
def forward(self,
|
||||
h: Tensor,
|
||||
skips: List[Tensor],
|
||||
cond0: Optional[Tensor] = None,
|
||||
cond1: Optional[Tensor] = None) -> Tensor:
|
||||
for i in range(len(self.resnet_blocks)):
|
||||
h = self.resnet_blocks[i].forward(torch.concat([h, skips[i]], dim=1), cond0, cond1)
|
||||
if self.use_attention:
|
||||
h = self.attention_blocks[i].forward(h)
|
||||
if self.perform_upsampling:
|
||||
h = self.upsample.forward(h, cond0, cond1)
|
||||
return h
|
||||
|
||||
|
||||
def compute_timestep_embedding(t: Tensor, out_channels: int):
|
||||
assert len(t.shape) == 2
|
||||
b, c = t.shape
|
||||
assert c == 1
|
||||
half_channels = out_channels // 2
|
||||
scale = -math.log(10000.0) / (half_channels - 1)
|
||||
log_times = scale * torch.arange(0, half_channels, device=t.device)
|
||||
times = torch.exp(log_times).reshape(1, half_channels) * t
|
||||
t_emb = torch.cat([torch.cos(times), torch.sin(times)], dim=1)
|
||||
if out_channels % 2 == 1:
|
||||
t_emb = torch.nn.functional.pad(t_emb, (1, 1), mode='constant')
|
||||
return t_emb
|
||||
|
||||
|
||||
class TimeEmbedding(Module):
|
||||
def __init__(self, out_channels: int):
|
||||
super().__init__()
|
||||
self.out_channels = out_channels
|
||||
|
||||
def forward(self, t: Tensor):
|
||||
return compute_timestep_embedding(t, self.out_channels)
|
||||
|
||||
|
||||
class UnetArgs:
|
||||
def __init__(self,
|
||||
in_channels: int = 3,
|
||||
out_channels: int = 3,
|
||||
model_channels: int = 64,
|
||||
level_channel_multipliers: Optional[List[int]] = None,
|
||||
level_use_attention: Optional[List[bool]] = None,
|
||||
num_res_blocks_per_level: int = 2,
|
||||
num_middle_res_blocks: int = 2,
|
||||
time_embedding_channels: Optional[int] = None,
|
||||
cond_input_channels: int = 4,
|
||||
cond_internal_channels: int = 512,
|
||||
attention_block_args: Optional[AttentionBlockArgs] = None,
|
||||
dropout_prob: float = 0.1,
|
||||
resample_with_res_block: bool = True,
|
||||
use_conv_to_resample=False,
|
||||
condition_bias: float = 1.0):
|
||||
assert len(level_channel_multipliers) == len(level_use_attention)
|
||||
assert not use_conv_to_resample or not resample_with_res_block
|
||||
|
||||
if time_embedding_channels is None:
|
||||
time_embedding_channels = model_channels
|
||||
if level_channel_multipliers is None:
|
||||
level_channel_multipliers = [1, 2, 4, 8]
|
||||
if level_use_attention is None:
|
||||
level_use_attention = [False for _ in level_channel_multipliers]
|
||||
if attention_block_args is None:
|
||||
attention_block_args = AttentionBlockArgs(
|
||||
num_heads=1,
|
||||
num_head_channels=None,
|
||||
use_new_attention_order=False)
|
||||
|
||||
self.condition_bias = condition_bias
|
||||
self.use_conv_to_resample = use_conv_to_resample
|
||||
self.resample_with_res_block = resample_with_res_block
|
||||
self.cond_internal_channels = cond_internal_channels
|
||||
self.dropout_prob = dropout_prob
|
||||
self.attention_block_args = attention_block_args
|
||||
self.time_embedding_channels = time_embedding_channels
|
||||
self.num_res_blocks_per_level = num_res_blocks_per_level
|
||||
self.level_use_attention = level_use_attention
|
||||
self.level_channel_multipliers = level_channel_multipliers
|
||||
self.model_channels = model_channels
|
||||
self.out_channels = out_channels
|
||||
self.in_channels = in_channels
|
||||
self.num_levels = len(level_channel_multipliers)
|
||||
self.num_middle_res_blocks = num_middle_res_blocks
|
||||
self.cond_input_channels = cond_input_channels
|
||||
|
||||
|
||||
class Unet(Module):
|
||||
def __init__(self, args: UnetArgs):
|
||||
super().__init__()
|
||||
self.args = args
|
||||
|
||||
self.time_embed = Sequential(
|
||||
TimeEmbedding(self.args.time_embedding_channels),
|
||||
Linear(self.args.time_embedding_channels, self.args.cond_internal_channels),
|
||||
SiLU(),
|
||||
Linear(self.args.cond_internal_channels, self.args.cond_internal_channels))
|
||||
|
||||
self.cond_embed = Sequential(
|
||||
Linear(self.args.cond_input_channels, self.args.cond_internal_channels),
|
||||
SiLU(),
|
||||
Linear(self.args.cond_internal_channels, self.args.cond_internal_channels))
|
||||
|
||||
self.first_conv = Conv2d(args.in_channels, args.model_channels, kernel_size=3, stride=1, padding=1)
|
||||
current_channels = args.model_channels
|
||||
channels = [current_channels]
|
||||
|
||||
# Downsampling blocks
|
||||
self.down_blocks = ModuleList()
|
||||
for i in range(args.num_levels):
|
||||
out_channels = args.model_channels * args.level_channel_multipliers[i]
|
||||
perform_downsampling = i < args.num_levels - 1
|
||||
down_block = DownsamplingBlock(
|
||||
in_channels=current_channels,
|
||||
out_channels=out_channels,
|
||||
cond0_channels=args.cond_internal_channels,
|
||||
cond1_channels=args.cond_internal_channels,
|
||||
num_res_blocks=args.num_res_blocks_per_level,
|
||||
dropout_prob=args.dropout_prob,
|
||||
use_attention=args.level_use_attention[i],
|
||||
perform_downsampling=perform_downsampling,
|
||||
attention_block_args=args.attention_block_args,
|
||||
resample_with_res_block=args.resample_with_res_block,
|
||||
use_conv_to_resample=args.use_conv_to_resample,
|
||||
condition_bias=args.condition_bias)
|
||||
self.down_blocks.append(down_block)
|
||||
current_channels = out_channels
|
||||
channels += down_block.output_channels
|
||||
|
||||
# Middle blocks
|
||||
self.middle_blocks = ModuleList()
|
||||
for i in range(self.args.num_middle_res_blocks - 1):
|
||||
self.middle_blocks.append(ResBlock(
|
||||
in_channels=current_channels,
|
||||
out_channels=current_channels,
|
||||
cond0_channels=args.cond_internal_channels,
|
||||
cond1_channels=args.cond_internal_channels,
|
||||
dropout_prob=args.dropout_prob,
|
||||
condition_bias=args.condition_bias))
|
||||
self.middle_blocks.append(
|
||||
Arity3To1(AttentionBlock(num_channels=current_channels, args=args.attention_block_args)))
|
||||
self.middle_blocks.append(ResBlock(
|
||||
in_channels=current_channels,
|
||||
out_channels=current_channels,
|
||||
cond0_channels=args.cond_internal_channels,
|
||||
cond1_channels=args.cond_internal_channels,
|
||||
dropout_prob=args.dropout_prob,
|
||||
condition_bias=args.condition_bias))
|
||||
|
||||
# Upsampling blocks
|
||||
self.up_blocks = ModuleList()
|
||||
for i in reversed(range(args.num_levels)):
|
||||
skip_channels = []
|
||||
for j in range(args.num_res_blocks_per_level + 1):
|
||||
skip_channels.append(channels.pop())
|
||||
perform_upsampling = i > 0
|
||||
out_channels = args.model_channels * args.level_channel_multipliers[i]
|
||||
up_block = UpsamplingBlock(
|
||||
in_channels=current_channels,
|
||||
out_channels=out_channels,
|
||||
cond0_channels=args.cond_internal_channels,
|
||||
cond1_channels=args.cond_internal_channels,
|
||||
num_resnet_blocks=args.num_res_blocks_per_level + 1,
|
||||
skip_channels=skip_channels,
|
||||
dropout_prob=args.dropout_prob,
|
||||
use_attention=args.level_use_attention[i],
|
||||
perform_upsampling=perform_upsampling,
|
||||
attention_block_args=args.attention_block_args,
|
||||
resample_with_res_block=args.resample_with_res_block,
|
||||
use_conv_to_resample=args.use_conv_to_resample,
|
||||
condition_bias=args.condition_bias)
|
||||
self.up_blocks.append(up_block)
|
||||
current_channels = out_channels
|
||||
assert len(channels) == 0
|
||||
|
||||
self.last = Sequential(
|
||||
GroupNorm32(current_channels),
|
||||
SiLU(),
|
||||
init_to_zero(Conv2d(current_channels, args.out_channels, kernel_size=3, stride=1, padding=1)))
|
||||
|
||||
def forward(self, x: Tensor, t: Tensor, cond: Tensor):
|
||||
t_emb = self.time_embed(t)
|
||||
cond_emb = self.cond_embed(cond)
|
||||
hs = [self.first_conv(x)]
|
||||
for block in self.down_blocks:
|
||||
hs += block.forward(hs[-1], t_emb, cond_emb)
|
||||
h = hs[-1]
|
||||
for block in self.middle_blocks:
|
||||
h = block(h, t_emb, cond_emb)
|
||||
for block in self.up_blocks:
|
||||
skips = []
|
||||
for i in range(self.args.num_res_blocks_per_level + 1):
|
||||
skips.append(hs.pop())
|
||||
h = block.forward(h, skips, t_emb, cond_emb)
|
||||
assert len(hs) == 0
|
||||
return self.last(h)
|
||||
|
||||
|
||||
class UnetWithFirstConvAddition(Module):
|
||||
def __init__(self, args: UnetArgs):
|
||||
super().__init__()
|
||||
self.args = args
|
||||
|
||||
self.time_embed = Sequential(
|
||||
TimeEmbedding(self.args.time_embedding_channels),
|
||||
Linear(self.args.time_embedding_channels, self.args.cond_internal_channels),
|
||||
SiLU(),
|
||||
Linear(self.args.cond_internal_channels, self.args.cond_internal_channels))
|
||||
|
||||
self.cond_embed = Sequential(
|
||||
Linear(self.args.cond_input_channels, self.args.cond_internal_channels),
|
||||
SiLU(),
|
||||
Linear(self.args.cond_internal_channels, self.args.cond_internal_channels))
|
||||
|
||||
self.first_conv = Conv2d(args.in_channels, args.model_channels, kernel_size=3, stride=1, padding=1)
|
||||
current_channels = args.model_channels
|
||||
channels = [current_channels]
|
||||
|
||||
# Downsampling blocks
|
||||
self.down_blocks = ModuleList()
|
||||
for i in range(args.num_levels):
|
||||
out_channels = args.model_channels * args.level_channel_multipliers[i]
|
||||
perform_downsampling = i < args.num_levels - 1
|
||||
down_block = DownsamplingBlock(
|
||||
in_channels=current_channels,
|
||||
out_channels=out_channels,
|
||||
cond0_channels=args.cond_internal_channels,
|
||||
cond1_channels=args.cond_internal_channels,
|
||||
num_res_blocks=args.num_res_blocks_per_level,
|
||||
dropout_prob=args.dropout_prob,
|
||||
use_attention=args.level_use_attention[i],
|
||||
perform_downsampling=perform_downsampling,
|
||||
attention_block_args=args.attention_block_args,
|
||||
resample_with_res_block=args.resample_with_res_block,
|
||||
use_conv_to_resample=args.use_conv_to_resample,
|
||||
condition_bias=args.condition_bias)
|
||||
self.down_blocks.append(down_block)
|
||||
current_channels = out_channels
|
||||
channels += down_block.output_channels
|
||||
|
||||
# Middle blocks
|
||||
self.middle_blocks = ModuleList()
|
||||
for i in range(self.args.num_middle_res_blocks - 1):
|
||||
self.middle_blocks.append(ResBlock(
|
||||
in_channels=current_channels,
|
||||
out_channels=current_channels,
|
||||
cond0_channels=args.cond_internal_channels,
|
||||
cond1_channels=args.cond_internal_channels,
|
||||
dropout_prob=args.dropout_prob,
|
||||
condition_bias=args.condition_bias))
|
||||
self.middle_blocks.append(
|
||||
Arity3To1(AttentionBlock(num_channels=current_channels, args=args.attention_block_args)))
|
||||
self.middle_blocks.append(ResBlock(
|
||||
in_channels=current_channels,
|
||||
out_channels=current_channels,
|
||||
cond0_channels=args.cond_internal_channels,
|
||||
cond1_channels=args.cond_internal_channels,
|
||||
dropout_prob=args.dropout_prob,
|
||||
condition_bias=args.condition_bias))
|
||||
|
||||
# Upsampling blocks
|
||||
self.up_blocks = ModuleList()
|
||||
for i in reversed(range(args.num_levels)):
|
||||
skip_channels = []
|
||||
for j in range(args.num_res_blocks_per_level + 1):
|
||||
skip_channels.append(channels.pop())
|
||||
perform_upsampling = i > 0
|
||||
out_channels = args.model_channels * args.level_channel_multipliers[i]
|
||||
up_block = UpsamplingBlock(
|
||||
in_channels=current_channels,
|
||||
out_channels=out_channels,
|
||||
cond0_channels=args.cond_internal_channels,
|
||||
cond1_channels=args.cond_internal_channels,
|
||||
num_resnet_blocks=args.num_res_blocks_per_level + 1,
|
||||
skip_channels=skip_channels,
|
||||
dropout_prob=args.dropout_prob,
|
||||
use_attention=args.level_use_attention[i],
|
||||
perform_upsampling=perform_upsampling,
|
||||
attention_block_args=args.attention_block_args,
|
||||
resample_with_res_block=args.resample_with_res_block,
|
||||
use_conv_to_resample=args.use_conv_to_resample,
|
||||
condition_bias=args.condition_bias)
|
||||
self.up_blocks.append(up_block)
|
||||
current_channels = out_channels
|
||||
assert len(channels) == 0
|
||||
|
||||
self.last = Sequential(
|
||||
GroupNorm32(current_channels),
|
||||
SiLU(),
|
||||
init_to_zero(Conv2d(current_channels, args.out_channels, kernel_size=3, stride=1, padding=1)))
|
||||
|
||||
def forward(self, x: Tensor, t: Tensor, cond: Tensor, first_conv_addition: Tensor):
|
||||
t_emb = self.time_embed(t)
|
||||
cond_emb = self.cond_embed(cond)
|
||||
first_conv = self.first_conv(x)
|
||||
hs = [first_conv + first_conv_addition]
|
||||
for block in self.down_blocks:
|
||||
hs += block.forward(hs[-1], t_emb, cond_emb)
|
||||
h = hs[-1]
|
||||
for block in self.middle_blocks:
|
||||
h = block(h, t_emb, cond_emb)
|
||||
for block in self.up_blocks:
|
||||
skips = []
|
||||
for i in range(self.args.num_res_blocks_per_level + 1):
|
||||
skips.append(hs.pop())
|
||||
h = block.forward(h, skips, t_emb, cond_emb)
|
||||
assert len(hs) == 0
|
||||
return self.last(h)
|
||||
@@ -0,0 +1,189 @@
|
||||
from typing import Optional, Union, Callable
|
||||
|
||||
from torch.nn import Conv2d, Module, Sequential, ConvTranspose2d
|
||||
|
||||
from tha4.shion.core.module_factory import ModuleFactory
|
||||
from tha4.nn.nonlinearity_factory import resolve_nonlinearity_factory
|
||||
from tha4.nn.normalization import NormalizationLayerFactory
|
||||
from tha4.nn.util import wrap_conv_or_linear_module, BlockArgs
|
||||
|
||||
|
||||
def create_conv7(in_channels: int, out_channels: int,
|
||||
bias: bool = False,
|
||||
initialization_method: Union[str, Callable[[Module], Module]] = 'he',
|
||||
use_spectral_norm: bool = False) -> Module:
|
||||
return wrap_conv_or_linear_module(
|
||||
Conv2d(in_channels, out_channels, kernel_size=7, stride=1, padding=3, bias=bias),
|
||||
initialization_method,
|
||||
use_spectral_norm)
|
||||
|
||||
|
||||
def create_conv7_from_block_args(in_channels: int,
|
||||
out_channels: int,
|
||||
bias: bool = False,
|
||||
block_args: Optional[BlockArgs] = None) -> Module:
|
||||
if block_args is None:
|
||||
block_args = BlockArgs()
|
||||
return create_conv7(
|
||||
in_channels, out_channels, bias,
|
||||
block_args.initialization_method,
|
||||
block_args.use_spectral_norm)
|
||||
|
||||
|
||||
def create_conv3(in_channels: int,
|
||||
out_channels: int,
|
||||
bias: bool = False,
|
||||
initialization_method: Union[str, Callable[[Module], Module]] = 'he',
|
||||
use_spectral_norm: bool = False) -> Module:
|
||||
return wrap_conv_or_linear_module(
|
||||
Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=bias),
|
||||
initialization_method,
|
||||
use_spectral_norm)
|
||||
|
||||
|
||||
def create_conv3_from_block_args(in_channels: int, out_channels: int,
|
||||
bias: bool = False,
|
||||
block_args: Optional[BlockArgs] = None):
|
||||
if block_args is None:
|
||||
block_args = BlockArgs()
|
||||
return create_conv3(in_channels, out_channels, bias,
|
||||
block_args.initialization_method,
|
||||
block_args.use_spectral_norm)
|
||||
|
||||
|
||||
def create_conv1(in_channels: int, out_channels: int,
|
||||
initialization_method: Union[str, Callable[[Module], Module]] = 'he',
|
||||
bias: bool = False,
|
||||
use_spectral_norm: bool = False) -> Module:
|
||||
return wrap_conv_or_linear_module(
|
||||
Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0, bias=bias),
|
||||
initialization_method,
|
||||
use_spectral_norm)
|
||||
|
||||
|
||||
def create_conv1_from_block_args(in_channels: int,
|
||||
out_channels: int,
|
||||
bias: bool = False,
|
||||
block_args: Optional[BlockArgs] = None) -> Module:
|
||||
if block_args is None:
|
||||
block_args = BlockArgs()
|
||||
return create_conv1(
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
initialization_method=block_args.initialization_method,
|
||||
bias=bias,
|
||||
use_spectral_norm=block_args.use_spectral_norm)
|
||||
|
||||
|
||||
def create_conv7_block(in_channels: int, out_channels: int,
|
||||
initialization_method: Union[str, Callable[[Module], Module]] = 'he',
|
||||
nonlinearity_factory: Optional[ModuleFactory] = None,
|
||||
normalization_layer_factory: Optional[NormalizationLayerFactory] = None,
|
||||
use_spectral_norm: bool = False) -> Module:
|
||||
nonlinearity_factory = resolve_nonlinearity_factory(nonlinearity_factory)
|
||||
return Sequential(
|
||||
create_conv7(in_channels, out_channels,
|
||||
bias=False, initialization_method=initialization_method, use_spectral_norm=use_spectral_norm),
|
||||
NormalizationLayerFactory.resolve_2d(normalization_layer_factory).create(out_channels, affine=True),
|
||||
resolve_nonlinearity_factory(nonlinearity_factory).create())
|
||||
|
||||
|
||||
def create_conv7_block_from_block_args(
|
||||
in_channels: int, out_channels: int,
|
||||
block_args: Optional[BlockArgs] = None) -> Module:
|
||||
if block_args is None:
|
||||
block_args = BlockArgs()
|
||||
return create_conv7_block(in_channels, out_channels,
|
||||
block_args.initialization_method,
|
||||
block_args.nonlinearity_factory,
|
||||
block_args.normalization_layer_factory,
|
||||
block_args.use_spectral_norm)
|
||||
|
||||
|
||||
def create_conv3_block(in_channels: int, out_channels: int,
|
||||
initialization_method: Union[str, Callable[[Module], Module]] = 'he',
|
||||
nonlinearity_factory: Optional[ModuleFactory] = None,
|
||||
normalization_layer_factory: Optional[NormalizationLayerFactory] = None,
|
||||
use_spectral_norm: bool = False) -> Module:
|
||||
nonlinearity_factory = resolve_nonlinearity_factory(nonlinearity_factory)
|
||||
return Sequential(
|
||||
create_conv3(in_channels, out_channels,
|
||||
bias=False, initialization_method=initialization_method, use_spectral_norm=use_spectral_norm),
|
||||
NormalizationLayerFactory.resolve_2d(normalization_layer_factory).create(out_channels, affine=True),
|
||||
resolve_nonlinearity_factory(nonlinearity_factory).create())
|
||||
|
||||
|
||||
def create_conv3_block_from_block_args(
|
||||
in_channels: int, out_channels: int, block_args: Optional[BlockArgs] = None):
|
||||
if block_args is None:
|
||||
block_args = BlockArgs()
|
||||
return create_conv3_block(in_channels, out_channels,
|
||||
block_args.initialization_method,
|
||||
block_args.nonlinearity_factory,
|
||||
block_args.normalization_layer_factory,
|
||||
block_args.use_spectral_norm)
|
||||
|
||||
|
||||
def create_downsample_block(in_channels: int, out_channels: int,
|
||||
is_output_1x1: bool = False,
|
||||
initialization_method: Union[str, Callable[[Module], Module]] = 'he',
|
||||
nonlinearity_factory: Optional[ModuleFactory] = None,
|
||||
normalization_layer_factory: Optional[NormalizationLayerFactory] = None,
|
||||
use_spectral_norm: bool = False) -> Module:
|
||||
if is_output_1x1:
|
||||
return Sequential(
|
||||
wrap_conv_or_linear_module(
|
||||
Conv2d(in_channels, out_channels, kernel_size=4, stride=2, padding=1, bias=False),
|
||||
initialization_method,
|
||||
use_spectral_norm),
|
||||
resolve_nonlinearity_factory(nonlinearity_factory).create())
|
||||
else:
|
||||
return Sequential(
|
||||
wrap_conv_or_linear_module(
|
||||
Conv2d(in_channels, out_channels, kernel_size=4, stride=2, padding=1, bias=False),
|
||||
initialization_method,
|
||||
use_spectral_norm),
|
||||
NormalizationLayerFactory.resolve_2d(normalization_layer_factory).create(out_channels, affine=True),
|
||||
resolve_nonlinearity_factory(nonlinearity_factory).create())
|
||||
|
||||
|
||||
def create_downsample_block_from_block_args(in_channels: int, out_channels: int,
|
||||
is_output_1x1: bool = False,
|
||||
block_args: Optional[BlockArgs] = None):
|
||||
if block_args is None:
|
||||
block_args = BlockArgs()
|
||||
return create_downsample_block(
|
||||
in_channels, out_channels,
|
||||
is_output_1x1,
|
||||
block_args.initialization_method,
|
||||
block_args.nonlinearity_factory,
|
||||
block_args.normalization_layer_factory,
|
||||
block_args.use_spectral_norm)
|
||||
|
||||
|
||||
def create_upsample_block(in_channels: int,
|
||||
out_channels: int,
|
||||
initialization_method: Union[str, Callable[[Module], Module]] = 'he',
|
||||
nonlinearity_factory: Optional[ModuleFactory] = None,
|
||||
normalization_layer_factory: Optional[NormalizationLayerFactory] = None,
|
||||
use_spectral_norm: bool = False) -> Module:
|
||||
nonlinearity_factory = resolve_nonlinearity_factory(nonlinearity_factory)
|
||||
return Sequential(
|
||||
wrap_conv_or_linear_module(
|
||||
ConvTranspose2d(in_channels, out_channels, kernel_size=4, stride=2, padding=1, bias=False),
|
||||
initialization_method,
|
||||
use_spectral_norm),
|
||||
NormalizationLayerFactory.resolve_2d(normalization_layer_factory).create(out_channels, affine=True),
|
||||
resolve_nonlinearity_factory(nonlinearity_factory).create())
|
||||
|
||||
|
||||
def create_upsample_block_from_block_args(in_channels: int,
|
||||
out_channels: int,
|
||||
block_args: Optional[BlockArgs] = None) -> Module:
|
||||
if block_args is None:
|
||||
block_args = BlockArgs()
|
||||
return create_upsample_block(in_channels, out_channels,
|
||||
block_args.initialization_method,
|
||||
block_args.nonlinearity_factory,
|
||||
block_args.normalization_layer_factory,
|
||||
block_args.use_spectral_norm)
|
||||
Whitespace-only changes.
@@ -0,0 +1,81 @@
|
||||
from typing import List, Optional
|
||||
|
||||
import torch
|
||||
from torch import Tensor
|
||||
from torch.nn import Module
|
||||
|
||||
from tha4.nn.common.poser_encoder_decoder_00 import PoserEncoderDecoder00Args, PoserEncoderDecoder00
|
||||
from tha4.nn.image_processing_util import apply_color_change
|
||||
from tha4.shion.core.module_factory import ModuleFactory
|
||||
from tha4.nn.nonlinearity_factory import ReLUFactory
|
||||
from tha4.nn.normalization import InstanceNorm2dFactory
|
||||
from tha4.nn.util import BlockArgs
|
||||
|
||||
|
||||
class EyebrowDecomposer00Args(PoserEncoderDecoder00Args):
|
||||
def __init__(self,
|
||||
image_size: int = 128,
|
||||
image_channels: int = 4,
|
||||
start_channels: int = 64,
|
||||
bottleneck_image_size=16,
|
||||
num_bottleneck_blocks=6,
|
||||
max_channels: int = 512,
|
||||
block_args: Optional[BlockArgs] = None):
|
||||
super().__init__(
|
||||
image_size,
|
||||
image_channels,
|
||||
image_channels,
|
||||
0,
|
||||
start_channels,
|
||||
bottleneck_image_size,
|
||||
num_bottleneck_blocks,
|
||||
max_channels,
|
||||
block_args)
|
||||
|
||||
|
||||
class EyebrowDecomposer00(Module):
|
||||
def __init__(self, args: EyebrowDecomposer00Args):
|
||||
super().__init__()
|
||||
self.args = args
|
||||
self.body = PoserEncoderDecoder00(args)
|
||||
self.background_layer_alpha = self.args.create_alpha_block()
|
||||
self.background_layer_color_change = self.args.create_color_change_block()
|
||||
self.eyebrow_layer_alpha = self.args.create_alpha_block()
|
||||
self.eyebrow_layer_color_change = self.args.create_color_change_block()
|
||||
|
||||
def forward(self, image: Tensor, *args) -> List[Tensor]:
|
||||
feature = self.body(image)[0]
|
||||
|
||||
background_layer_alpha = self.background_layer_alpha(feature)
|
||||
background_layer_color_change = self.background_layer_color_change(feature)
|
||||
background_layer_1 = apply_color_change(background_layer_alpha, background_layer_color_change, image)
|
||||
|
||||
eyebrow_layer_alpha = self.eyebrow_layer_alpha(feature)
|
||||
eyebrow_layer_color_change = self.eyebrow_layer_color_change(feature)
|
||||
eyebrow_layer = apply_color_change(eyebrow_layer_alpha, image, eyebrow_layer_color_change)
|
||||
|
||||
return [
|
||||
eyebrow_layer, # 0
|
||||
eyebrow_layer_alpha, # 1
|
||||
eyebrow_layer_color_change, # 2
|
||||
background_layer_1, # 3
|
||||
background_layer_alpha, # 4
|
||||
background_layer_color_change, # 5
|
||||
]
|
||||
|
||||
EYEBROW_LAYER_INDEX = 0
|
||||
EYEBROW_LAYER_ALPHA_INDEX = 1
|
||||
EYEBROW_LAYER_COLOR_CHANGE_INDEX = 2
|
||||
BACKGROUND_LAYER_INDEX = 3
|
||||
BACKGROUND_LAYER_ALPHA_INDEX = 4
|
||||
BACKGROUND_LAYER_COLOR_CHANGE_INDEX = 5
|
||||
OUTPUT_LENGTH = 6
|
||||
|
||||
|
||||
class EyebrowDecomposer00Factory(ModuleFactory):
|
||||
def __init__(self, args: EyebrowDecomposer00Args):
|
||||
super().__init__()
|
||||
self.args = args
|
||||
|
||||
def create(self) -> Module:
|
||||
return EyebrowDecomposer00(self.args)
|
||||
Whitespace-only changes.
@@ -0,0 +1,91 @@
|
||||
from typing import List, Optional
|
||||
|
||||
import torch
|
||||
from torch import Tensor
|
||||
from torch.nn import Module
|
||||
|
||||
from tha4.nn.common.poser_encoder_decoder_00 import PoserEncoderDecoder00Args, PoserEncoderDecoder00
|
||||
from tha4.nn.image_processing_util import apply_color_change, apply_grid_change, apply_rgb_change
|
||||
from tha4.shion.core.module_factory import ModuleFactory
|
||||
from tha4.nn.nonlinearity_factory import ReLUFactory
|
||||
from tha4.nn.normalization import InstanceNorm2dFactory
|
||||
from tha4.nn.util import BlockArgs
|
||||
|
||||
|
||||
class EyebrowMorphingCombiner00Args(PoserEncoderDecoder00Args):
|
||||
def __init__(self,
|
||||
image_size: int = 128,
|
||||
image_channels: int = 4,
|
||||
num_pose_params: int = 12,
|
||||
start_channels: int = 64,
|
||||
bottleneck_image_size=16,
|
||||
num_bottleneck_blocks=6,
|
||||
max_channels: int = 512,
|
||||
block_args: Optional[BlockArgs] = None):
|
||||
super().__init__(
|
||||
image_size,
|
||||
2 * image_channels,
|
||||
image_channels,
|
||||
num_pose_params,
|
||||
start_channels,
|
||||
bottleneck_image_size,
|
||||
num_bottleneck_blocks,
|
||||
max_channels,
|
||||
block_args)
|
||||
|
||||
|
||||
class EyebrowMorphingCombiner00(Module):
|
||||
def __init__(self, args: EyebrowMorphingCombiner00Args):
|
||||
super().__init__()
|
||||
self.args = args
|
||||
self.body = PoserEncoderDecoder00(args)
|
||||
self.morphed_eyebrow_layer_grid_change = self.args.create_grid_change_block()
|
||||
self.morphed_eyebrow_layer_alpha = self.args.create_alpha_block()
|
||||
self.morphed_eyebrow_layer_color_change = self.args.create_color_change_block()
|
||||
self.combine_alpha = self.args.create_alpha_block()
|
||||
|
||||
def forward(self, background_layer: Tensor, eyebrow_layer: Tensor, pose: Tensor, *args) -> List[Tensor]:
|
||||
combined_image = torch.cat([background_layer, eyebrow_layer], dim=1)
|
||||
feature = self.body(combined_image, pose)[0]
|
||||
|
||||
morphed_eyebrow_layer_grid_change = self.morphed_eyebrow_layer_grid_change(feature)
|
||||
morphed_eyebrow_layer_alpha = self.morphed_eyebrow_layer_alpha(feature)
|
||||
morphed_eyebrow_layer_color_change = self.morphed_eyebrow_layer_color_change(feature)
|
||||
warped_eyebrow_layer = apply_grid_change(morphed_eyebrow_layer_grid_change, eyebrow_layer)
|
||||
morphed_eyebrow_layer = apply_color_change(
|
||||
morphed_eyebrow_layer_alpha, morphed_eyebrow_layer_color_change, warped_eyebrow_layer)
|
||||
|
||||
combine_alpha = self.combine_alpha(feature)
|
||||
eyebrow_image = apply_rgb_change(combine_alpha, morphed_eyebrow_layer, background_layer)
|
||||
eyebrow_image_no_combine_alpha = apply_rgb_change(
|
||||
(morphed_eyebrow_layer[:, 3:4, :, :] + 1.0) / 2.0, morphed_eyebrow_layer, background_layer)
|
||||
|
||||
return [
|
||||
eyebrow_image, # 0
|
||||
combine_alpha, # 1
|
||||
eyebrow_image_no_combine_alpha, # 2
|
||||
morphed_eyebrow_layer, # 3
|
||||
morphed_eyebrow_layer_alpha, # 4
|
||||
morphed_eyebrow_layer_color_change, # 5
|
||||
warped_eyebrow_layer, # 6
|
||||
morphed_eyebrow_layer_grid_change, # 7
|
||||
]
|
||||
|
||||
EYEBROW_IMAGE_INDEX = 0
|
||||
COMBINE_ALPHA_INDEX = 1
|
||||
EYEBROW_IMAGE_NO_COMBINE_ALPHA_INDEX = 2
|
||||
MORPHED_EYEBROW_LAYER_INDEX = 3
|
||||
MORPHED_EYEBROW_LAYER_ALPHA_INDEX = 4
|
||||
MORPHED_EYEBROW_LAYER_COLOR_CHANGE_INDEX = 5
|
||||
WARPED_EYEBROW_LAYER_INDEX = 6
|
||||
MORPHED_EYEBROW_LAYER_GRID_CHANGE_INDEX = 7
|
||||
OUTPUT_LENGTH = 8
|
||||
|
||||
|
||||
class EyebrowMorphingCombiner00Factory(ModuleFactory):
|
||||
def __init__(self, args: EyebrowMorphingCombiner00Args):
|
||||
super().__init__()
|
||||
self.args = args
|
||||
|
||||
def create(self) -> Module:
|
||||
return EyebrowMorphingCombiner00(self.args)
|
||||
Whitespace-only changes.
@@ -0,0 +1,211 @@
|
||||
import math
|
||||
from typing import List, Optional
|
||||
|
||||
import torch
|
||||
from torch import Tensor
|
||||
from torch.nn import ModuleList, Sequential, Sigmoid, Tanh, Module
|
||||
from torch.nn.functional import affine_grid, grid_sample
|
||||
|
||||
from tha4.shion.core.module_factory import ModuleFactory
|
||||
from tha4.nn.conv import create_conv3_block_from_block_args, \
|
||||
create_downsample_block_from_block_args, create_upsample_block_from_block_args, create_conv3_from_block_args, \
|
||||
create_conv3
|
||||
from tha4.nn.nonlinearity_factory import LeakyReLUFactory
|
||||
from tha4.nn.normalization import InstanceNorm2dFactory
|
||||
from tha4.nn.resnet_block import ResnetBlock
|
||||
from tha4.nn.util import BlockArgs
|
||||
|
||||
|
||||
class FaceMorpher08Args:
|
||||
def __init__(self,
|
||||
image_size: int = 256,
|
||||
image_channels: int = 4,
|
||||
num_expression_params: int = 67,
|
||||
start_channels: int = 16,
|
||||
bottleneck_image_size=4,
|
||||
num_bottleneck_blocks=3,
|
||||
max_channels: int = 512,
|
||||
block_args: Optional[BlockArgs] = None,
|
||||
output_iris_mouth_grid_change: bool = False):
|
||||
self.max_channels = max_channels
|
||||
self.num_bottleneck_blocks = num_bottleneck_blocks
|
||||
assert bottleneck_image_size > 1
|
||||
self.bottleneck_image_size = bottleneck_image_size
|
||||
self.start_channels = start_channels
|
||||
self.image_channels = image_channels
|
||||
self.num_expression_params = num_expression_params
|
||||
self.image_size = image_size
|
||||
self.output_iris_mouth_grid_change = output_iris_mouth_grid_change
|
||||
|
||||
if block_args is None:
|
||||
self.block_args = BlockArgs(
|
||||
normalization_layer_factory=InstanceNorm2dFactory(),
|
||||
nonlinearity_factory=LeakyReLUFactory(negative_slope=0.2, inplace=True))
|
||||
else:
|
||||
self.block_args = block_args
|
||||
|
||||
|
||||
class FaceMorpher08(Module):
|
||||
def __init__(self, args: FaceMorpher08Args):
|
||||
super().__init__()
|
||||
self.args = args
|
||||
self.num_levels = int(math.log2(args.image_size // args.bottleneck_image_size)) + 1
|
||||
|
||||
self.downsample_blocks = ModuleList()
|
||||
self.downsample_blocks.append(
|
||||
create_conv3_block_from_block_args(
|
||||
args.image_channels,
|
||||
args.start_channels,
|
||||
args.block_args))
|
||||
current_image_size = args.image_size
|
||||
current_num_channels = args.start_channels
|
||||
while current_image_size > args.bottleneck_image_size:
|
||||
next_image_size = current_image_size // 2
|
||||
next_num_channels = self.get_num_output_channels_from_image_size(next_image_size)
|
||||
self.downsample_blocks.append(create_downsample_block_from_block_args(
|
||||
in_channels=current_num_channels,
|
||||
out_channels=next_num_channels,
|
||||
is_output_1x1=False,
|
||||
block_args=args.block_args))
|
||||
current_image_size = next_image_size
|
||||
current_num_channels = next_num_channels
|
||||
assert len(self.downsample_blocks) == self.num_levels
|
||||
|
||||
self.bottleneck_blocks = ModuleList()
|
||||
self.bottleneck_blocks.append(create_conv3_block_from_block_args(
|
||||
in_channels=current_num_channels + args.num_expression_params,
|
||||
out_channels=current_num_channels,
|
||||
block_args=args.block_args))
|
||||
for i in range(1, args.num_bottleneck_blocks):
|
||||
self.bottleneck_blocks.append(
|
||||
ResnetBlock.create(
|
||||
num_channels=current_num_channels,
|
||||
is1x1=False,
|
||||
block_args=args.block_args))
|
||||
|
||||
self.upsample_blocks = ModuleList()
|
||||
while current_image_size < args.image_size:
|
||||
next_image_size = current_image_size * 2
|
||||
next_num_channels = self.get_num_output_channels_from_image_size(next_image_size)
|
||||
self.upsample_blocks.append(create_upsample_block_from_block_args(
|
||||
in_channels=current_num_channels,
|
||||
out_channels=next_num_channels,
|
||||
block_args=args.block_args))
|
||||
current_image_size = next_image_size
|
||||
current_num_channels = next_num_channels
|
||||
|
||||
self.iris_mouth_grid_change = self.create_grid_change_block()
|
||||
self.iris_mouth_color_change = self.create_color_change_block()
|
||||
self.iris_mouth_alpha = self.create_alpha_block()
|
||||
|
||||
self.eye_color_change = self.create_color_change_block()
|
||||
self.eye_alpha = self.create_alpha_block()
|
||||
|
||||
def create_alpha_block(self):
|
||||
return Sequential(
|
||||
create_conv3(
|
||||
in_channels=self.args.start_channels,
|
||||
out_channels=1,
|
||||
bias=True,
|
||||
initialization_method=self.args.block_args.initialization_method,
|
||||
use_spectral_norm=False),
|
||||
Sigmoid())
|
||||
|
||||
def create_color_change_block(self):
|
||||
return Sequential(
|
||||
create_conv3_from_block_args(
|
||||
in_channels=self.args.start_channels,
|
||||
out_channels=self.args.image_channels,
|
||||
bias=True,
|
||||
block_args=self.args.block_args),
|
||||
Tanh())
|
||||
|
||||
def create_grid_change_block(self):
|
||||
return create_conv3(
|
||||
in_channels=self.args.start_channels,
|
||||
out_channels=2,
|
||||
bias=False,
|
||||
initialization_method='zero',
|
||||
use_spectral_norm=False)
|
||||
|
||||
def get_num_output_channels_from_level(self, level: int):
|
||||
return self.get_num_output_channels_from_image_size(self.args.image_size // (2 ** level))
|
||||
|
||||
def get_num_output_channels_from_image_size(self, image_size: int):
|
||||
return min(self.args.start_channels * (self.args.image_size // image_size), self.args.max_channels)
|
||||
|
||||
def merge_down(self, top_layer: Tensor, bottom_layer: Tensor):
|
||||
top_layer_rgb = top_layer[:, 0:3, :, :]
|
||||
top_layer_a = top_layer[:, 3:4, :, :]
|
||||
return bottom_layer * (1 - top_layer_a) + torch.cat([top_layer_rgb * top_layer_a, top_layer_a], dim=1)
|
||||
|
||||
def apply_grid_change(self, grid_change, image: Tensor) -> Tensor:
|
||||
n, c, h, w = image.shape
|
||||
device = grid_change.device
|
||||
grid_change = torch.transpose(grid_change.view(n, 2, h * w), 1, 2).view(n, h, w, 2)
|
||||
identity = torch.tensor(
|
||||
[[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]],
|
||||
device=device,
|
||||
dtype=grid_change.dtype).unsqueeze(0).repeat(n, 1, 1)
|
||||
base_grid = affine_grid(identity, [n, c, h, w], align_corners=False)
|
||||
grid = base_grid + grid_change
|
||||
resampled_image = grid_sample(image, grid, mode='bilinear', padding_mode='border', align_corners=False)
|
||||
return resampled_image
|
||||
|
||||
def apply_color_change(self, alpha, color_change, image: Tensor) -> Tensor:
|
||||
return color_change * alpha + image * (1 - alpha)
|
||||
|
||||
def forward(self, image: Tensor, pose: Tensor, *args) -> List[Tensor]:
|
||||
feature = image
|
||||
for block in self.downsample_blocks:
|
||||
feature = block(feature)
|
||||
n, c = pose.shape
|
||||
pose = pose.view(n, c, 1, 1).repeat(1, 1, self.args.bottleneck_image_size, self.args.bottleneck_image_size)
|
||||
feature = torch.cat([feature, pose], dim=1)
|
||||
for block in self.bottleneck_blocks:
|
||||
feature = block(feature)
|
||||
for block in self.upsample_blocks:
|
||||
feature = block(feature)
|
||||
|
||||
iris_mouth_grid_change = self.iris_mouth_grid_change(feature)
|
||||
iris_mouth_image_0 = self.apply_grid_change(iris_mouth_grid_change, image)
|
||||
iris_mouth_color_change = self.iris_mouth_color_change(feature)
|
||||
iris_mouth_alpha = self.iris_mouth_alpha(feature)
|
||||
iris_mouth_image_1 = self.apply_color_change(iris_mouth_alpha, iris_mouth_color_change, iris_mouth_image_0)
|
||||
|
||||
eye_color_change = self.eye_color_change(feature)
|
||||
eye_alpha = self.eye_alpha(feature)
|
||||
output_image = self.apply_color_change(eye_alpha, eye_color_change, iris_mouth_image_1.detach())
|
||||
|
||||
outputs = [
|
||||
output_image, # 0
|
||||
eye_alpha, # 1
|
||||
eye_color_change, # 2
|
||||
iris_mouth_image_1, # 3
|
||||
iris_mouth_alpha, # 4
|
||||
iris_mouth_color_change, # 5
|
||||
iris_mouth_image_0, # 6
|
||||
]
|
||||
|
||||
if self.args.output_iris_mouth_grid_change:
|
||||
outputs.append(iris_mouth_grid_change)
|
||||
|
||||
return outputs
|
||||
|
||||
OUTPUT_IMAGE_INDEX = 0
|
||||
EYE_ALPHA_INDEX = 1
|
||||
EYE_COLOR_CHANGE_INDEX = 2
|
||||
IRIS_MOUTH_IMAGE_1_INDEX = 3
|
||||
IRIS_MOUTH_ALPHA_INDEX = 4
|
||||
IRIS_MOUTH_COLOR_CHANGE_INDEX = 5
|
||||
IRIS_MOUTH_IMAGE_0_INDEX = 6
|
||||
IRIS_MOUTH_GRID_CHANGE_INDEX = 7
|
||||
|
||||
|
||||
class FaceMorpher08Factory(ModuleFactory):
|
||||
def __init__(self, args: FaceMorpher08Args):
|
||||
super().__init__()
|
||||
self.args = args
|
||||
|
||||
def create(self) -> Module:
|
||||
return FaceMorpher08(self.args)
|
||||
@@ -0,0 +1,58 @@
|
||||
import torch
|
||||
from torch import Tensor
|
||||
from torch.nn.functional import affine_grid, grid_sample
|
||||
|
||||
|
||||
def apply_rgb_change(alpha: Tensor, color_change: Tensor, image: Tensor):
|
||||
image_rgb = image[:, 0:3, :, :]
|
||||
color_change_rgb = color_change[:, 0:3, :, :]
|
||||
output_rgb = color_change_rgb * alpha + image_rgb * (1 - alpha)
|
||||
return torch.cat([output_rgb, image[:, 3:4, :, :]], dim=1)
|
||||
|
||||
|
||||
def apply_grid_change(grid_change, image: Tensor) -> Tensor:
|
||||
n, c, h, w = image.shape
|
||||
device = grid_change.device
|
||||
grid_change = torch.transpose(grid_change.view(n, 2, h * w), 1, 2).view(n, h, w, 2)
|
||||
identity = torch.tensor(
|
||||
[[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]],
|
||||
dtype=grid_change.dtype,
|
||||
device=device).unsqueeze(0).repeat(n, 1, 1)
|
||||
base_grid = affine_grid(identity, [n, c, h, w], align_corners=False)
|
||||
grid = base_grid + grid_change
|
||||
resampled_image = grid_sample(image, grid, mode='bilinear', padding_mode='border', align_corners=False)
|
||||
return resampled_image
|
||||
|
||||
|
||||
class GridChangeApplier:
|
||||
def __init__(self):
|
||||
self.last_n = None
|
||||
self.last_device = None
|
||||
self.last_identity = None
|
||||
|
||||
def apply(self, grid_change: Tensor, image: Tensor, align_corners: bool = False) -> Tensor:
|
||||
n, c, h, w = image.shape
|
||||
device = grid_change.device
|
||||
grid_change = torch.transpose(grid_change.view(n, 2, h * w), 1, 2).view(n, h, w, 2)
|
||||
|
||||
if n == self.last_n and device == self.last_device:
|
||||
identity = self.last_identity
|
||||
else:
|
||||
identity = torch.tensor(
|
||||
[[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]],
|
||||
dtype=grid_change.dtype,
|
||||
device=device,
|
||||
requires_grad=False) \
|
||||
.unsqueeze(0).repeat(n, 1, 1)
|
||||
self.last_identity = identity
|
||||
self.last_n = n
|
||||
self.last_device = device
|
||||
base_grid = affine_grid(identity, [n, c, h, w], align_corners=align_corners)
|
||||
|
||||
grid = base_grid + grid_change
|
||||
resampled_image = grid_sample(image, grid, mode='bilinear', padding_mode='border', align_corners=align_corners)
|
||||
return resampled_image
|
||||
|
||||
|
||||
def apply_color_change(alpha, color_change, image: Tensor) -> Tensor:
|
||||
return color_change * alpha + image * (1 - alpha)
|
||||
@@ -0,0 +1,76 @@
|
||||
from typing import Callable
|
||||
|
||||
import torch
|
||||
from torch import zero_
|
||||
from torch.nn import Module
|
||||
from torch.nn.init import kaiming_normal_, xavier_normal_, normal_
|
||||
|
||||
|
||||
def create_init_function(method: str = 'none') -> Callable[[Module], Module]:
|
||||
def init(module: Module):
|
||||
if method == 'none':
|
||||
return module
|
||||
elif method == 'he':
|
||||
kaiming_normal_(module.weight)
|
||||
return module
|
||||
elif method == 'xavier':
|
||||
xavier_normal_(module.weight)
|
||||
return module
|
||||
elif method == 'dcgan':
|
||||
normal_(module.weight, 0.0, 0.02)
|
||||
return module
|
||||
elif method == 'dcgan_001':
|
||||
normal_(module.weight, 0.0, 0.01)
|
||||
return module
|
||||
elif method == "zero":
|
||||
with torch.no_grad():
|
||||
zero_(module.weight)
|
||||
return module
|
||||
else:
|
||||
raise ("Invalid initialization method %s" % method)
|
||||
|
||||
return init
|
||||
|
||||
|
||||
class HeInitialization:
|
||||
def __init__(self, a: int = 0, mode: str = 'fan_in', nonlinearity: str = 'leaky_relu'):
|
||||
self.nonlinearity = nonlinearity
|
||||
self.mode = mode
|
||||
self.a = a
|
||||
|
||||
def __call__(self, module: Module) -> Module:
|
||||
with torch.no_grad():
|
||||
kaiming_normal_(module.weight, a=self.a, mode=self.mode, nonlinearity=self.nonlinearity)
|
||||
return module
|
||||
|
||||
|
||||
class NormalInitialization:
|
||||
def __init__(self, mean: float = 0.0, std: float = 1.0):
|
||||
self.std = std
|
||||
self.mean = mean
|
||||
|
||||
def __call__(self, module: Module) -> Module:
|
||||
with torch.no_grad():
|
||||
normal_(module.weight, self.mean, self.std)
|
||||
return module
|
||||
|
||||
|
||||
class XavierInitialization:
|
||||
def __init__(self, gain: float = 1.0):
|
||||
self.gain = gain
|
||||
|
||||
def __call__(self, module: Module) -> Module:
|
||||
with torch.no_grad():
|
||||
xavier_normal_(module.weight, self.gain)
|
||||
return module
|
||||
|
||||
|
||||
class ZeroInitialization:
|
||||
def __call__(self, module: Module) -> Module:
|
||||
with torch.no_grad:
|
||||
zero_(module.weight)
|
||||
return module
|
||||
|
||||
class NoInitialization:
|
||||
def __call__(self, module: Module) -> Module:
|
||||
return module
|
||||
Whitespace-only changes.
Loaded 100 of 212 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user