mirror of
https://github.com/storytold/storyteller-ml.git
synced 2026-10-09 00:09:55 +00:00
Add SadTalker @ ae0b26359d44c9379450a8e02a38b45eeec15549
MIT license Commit: https://github.com/OpenTalker/SadTalker commit ae0b26359d44c9379450a8e02a38b45eeec15549 (HEAD -> main, origin/main, origin/HEAD) Author: shadowcun <shadowcun@tencent.com> Date: Mon Jul 3 13:48:35 2023 +0800 fixed bug: hacking facexlib of FAN
This commit is contained in:
Brandon Thomas
@@ -0,0 +1,174 @@
|
||||
# Byte-compiled / optimized / DLL files
|
||||
__pycache__/
|
||||
*.py[cod]
|
||||
*$py.class
|
||||
|
||||
# C extensions
|
||||
*.so
|
||||
|
||||
# Distribution / packaging
|
||||
.Python
|
||||
build/
|
||||
develop-eggs/
|
||||
dist/
|
||||
downloads/
|
||||
eggs/
|
||||
.eggs/
|
||||
lib/
|
||||
lib64/
|
||||
parts/
|
||||
sdist/
|
||||
var/
|
||||
wheels/
|
||||
share/python-wheels/
|
||||
*.egg-info/
|
||||
.installed.cfg
|
||||
*.egg
|
||||
MANIFEST
|
||||
|
||||
# PyInstaller
|
||||
# Usually these files are written by a python script from a template
|
||||
# before PyInstaller builds the exe, so as to inject date/other infos into it.
|
||||
*.manifest
|
||||
*.spec
|
||||
|
||||
# Installer logs
|
||||
pip-log.txt
|
||||
pip-delete-this-directory.txt
|
||||
|
||||
# Unit test / coverage reports
|
||||
htmlcov/
|
||||
.tox/
|
||||
.nox/
|
||||
.coverage
|
||||
.coverage.*
|
||||
.cache
|
||||
nosetests.xml
|
||||
coverage.xml
|
||||
*.cover
|
||||
*.py,cover
|
||||
.hypothesis/
|
||||
.pytest_cache/
|
||||
cover/
|
||||
|
||||
# Translations
|
||||
*.mo
|
||||
*.pot
|
||||
|
||||
# Django stuff:
|
||||
*.log
|
||||
local_settings.py
|
||||
db.sqlite3
|
||||
db.sqlite3-journal
|
||||
|
||||
# Flask stuff:
|
||||
instance/
|
||||
.webassets-cache
|
||||
|
||||
# Scrapy stuff:
|
||||
.scrapy
|
||||
|
||||
# Sphinx documentation
|
||||
docs/_build/
|
||||
|
||||
# PyBuilder
|
||||
.pybuilder/
|
||||
target/
|
||||
|
||||
# Jupyter Notebook
|
||||
.ipynb_checkpoints
|
||||
|
||||
# IPython
|
||||
profile_default/
|
||||
ipython_config.py
|
||||
|
||||
# pyenv
|
||||
# For a library or package, you might want to ignore these files since the code is
|
||||
# intended to run in multiple environments; otherwise, check them in:
|
||||
# .python-version
|
||||
|
||||
# pipenv
|
||||
# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
|
||||
# However, in case of collaboration, if having platform-specific dependencies or dependencies
|
||||
# having no cross-platform support, pipenv may install dependencies that don't work, or not
|
||||
# install all needed dependencies.
|
||||
#Pipfile.lock
|
||||
|
||||
# poetry
|
||||
# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
|
||||
# This is especially recommended for binary packages to ensure reproducibility, and is more
|
||||
# commonly ignored for libraries.
|
||||
# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
|
||||
#poetry.lock
|
||||
|
||||
# pdm
|
||||
# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
|
||||
#pdm.lock
|
||||
# pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
|
||||
# in version control.
|
||||
# https://pdm.fming.dev/#use-with-ide
|
||||
.pdm.toml
|
||||
|
||||
# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
|
||||
__pypackages__/
|
||||
|
||||
# Celery stuff
|
||||
celerybeat-schedule
|
||||
celerybeat.pid
|
||||
|
||||
# SageMath parsed files
|
||||
*.sage.py
|
||||
|
||||
# Environments
|
||||
.env
|
||||
.venv
|
||||
env/
|
||||
venv/
|
||||
ENV/
|
||||
env.bak/
|
||||
venv.bak/
|
||||
|
||||
# Spyder project settings
|
||||
.spyderproject
|
||||
.spyproject
|
||||
|
||||
# Rope project settings
|
||||
.ropeproject
|
||||
|
||||
# mkdocs documentation
|
||||
/site
|
||||
|
||||
# mypy
|
||||
.mypy_cache/
|
||||
.dmypy.json
|
||||
dmypy.json
|
||||
|
||||
# Pyre type checker
|
||||
.pyre/
|
||||
|
||||
# pytype static type analyzer
|
||||
.pytype/
|
||||
|
||||
# Cython debug symbols
|
||||
cython_debug/
|
||||
|
||||
# PyCharm
|
||||
# JetBrains specific template is maintained in a separate JetBrains.gitignore that can
|
||||
# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
|
||||
# and can be added to the global gitignore or merged into this file. For a more nuclear
|
||||
# option (not recommended) you can uncomment the following to ignore the entire idea folder.
|
||||
.idea/
|
||||
|
||||
examples/results/*
|
||||
gfpgan/*
|
||||
checkpoints/*
|
||||
assets/*
|
||||
results/*
|
||||
Dockerfile
|
||||
start_docker.sh
|
||||
start.sh
|
||||
|
||||
checkpoints
|
||||
|
||||
# Mac
|
||||
.DS_Store
|
||||
@@ -0,0 +1,21 @@
|
||||
MIT License
|
||||
|
||||
Copyright (c) 2023 Tencent AI Lab
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
of this software and associated documentation files (the "Software"), to deal
|
||||
in the Software without restriction, including without limitation the rights
|
||||
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||
copies of the Software, and to permit persons to whom the Software is
|
||||
furnished to do so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in all
|
||||
copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||
SOFTWARE.
|
||||
@@ -0,0 +1,268 @@
|
||||
<div align="center">
|
||||
|
||||
<img src='https://user-images.githubusercontent.com/4397546/229094115-862c747e-7397-4b54-ba4a-bd368bfe2e0f.png' width='500px'/>
|
||||
|
||||
|
||||
<!--<h2> 😭 SadTalker： <span style="font-size:12px">Learning Realistic 3D Motion Coefficients for Stylized Audio-Driven Single Image Talking Face Animation </span> </h2> -->
|
||||
|
||||
<a href='https://arxiv.org/abs/2211.12194'><img src='https://img.shields.io/badge/ArXiv-PDF-red'></a> <a href='https://sadtalker.github.io'><img src='https://img.shields.io/badge/Project-Page-Green'></a> [](https://colab.research.google.com/github/Winfredy/SadTalker/blob/main/quick_demo.ipynb) [](https://huggingface.co/spaces/vinthony/SadTalker) [](https://colab.research.google.com/github/camenduru/stable-diffusion-webui-colab/blob/main/video/stable/stable_diffusion_1_5_video_webui_colab.ipynb) [](https://replicate.com/cjwbw/sadtalker)
|
||||
|
||||
<div>
|
||||
<a target='_blank'>Wenxuan Zhang <sup>*,1,2</sup> </a> 
|
||||
<a href='https://vinthony.github.io/' target='_blank'>Xiaodong Cun <sup>*,2</a> 
|
||||
<a href='https://xuanwangvc.github.io/' target='_blank'>Xuan Wang <sup>3</sup></a> 
|
||||
<a href='https://yzhang2016.github.io/' target='_blank'>Yong Zhang <sup>2</sup></a> 
|
||||
<a href='https://xishen0220.github.io/' target='_blank'>Xi Shen <sup>2</sup></a>  </br>
|
||||
<a href='https://yuguo-xjtu.github.io/' target='_blank'>Yu Guo<sup>1</sup> </a> 
|
||||
<a href='https://scholar.google.com/citations?hl=zh-CN&user=4oXBp9UAAAAJ' target='_blank'>Ying Shan <sup>2</sup> </a> 
|
||||
<a target='_blank'>Fei Wang <sup>1</sup> </a> 
|
||||
</div>
|
||||
<br>
|
||||
<div>
|
||||
<sup>1</sup> Xi'an Jiaotong University   <sup>2</sup> Tencent AI Lab   <sup>3</sup> Ant Group  
|
||||
</div>
|
||||
<br>
|
||||
<i><strong><a href='https://arxiv.org/abs/2211.12194' target='_blank'>CVPR 2023</a></strong></i>
|
||||
<br>
|
||||
<br>
|
||||
|
||||
|
||||
|
||||
|
||||
<b>TL;DR: single portrait image 🙎♂️ + audio 🎤 = talking head video 🎞.</b>
|
||||
|
||||
<br>
|
||||
|
||||
</div>
|
||||
|
||||
|
||||
|
||||
## 🔥 Highlight
|
||||
|
||||
- 🔥 The extension of the [stable-diffusion-webui](https://github.com/AUTOMATIC1111/stable-diffusion-webui) is online. Checkout more details [here](docs/webui_extension.md).
|
||||
|
||||
https://user-images.githubusercontent.com/4397546/231495639-5d4bb925-ea64-4a36-a519-6389917dac29.mp4
|
||||
|
||||
- 🔥 `full image mode` is online! checkout [here](https://github.com/Winfredy/SadTalker#full-bodyimage-generation) for more details.
|
||||
|
||||
| still+enhancer in v0.0.1 | still + enhancer in v0.0.2 | [input image @bagbag1815](https://twitter.com/bagbag1815/status/1642754319094108161) |
|
||||
|:--------------------: |:--------------------: | :----: |
|
||||
| <video src="https://user-images.githubusercontent.com/48216707/229484996-5d7be64f-2553-4c9e-a452-c5cf0b8ebafe.mp4" type="video/mp4"> </video> | <video src="https://user-images.githubusercontent.com/4397546/230717873-355b7bf3-d3de-49f9-a439-9220e623fce7.mp4" type="video/mp4"> </video> | <img src='./examples/source_image/full_body_2.png' width='380'>
|
||||
|
||||
- 🔥 Several new mode, eg, `still mode`, `reference mode`, `resize mode` are online for better and custom applications.
|
||||
|
||||
- 🔥 Happy to see more community demos at [bilibili](https://search.bilibili.com/all?keyword=sadtalker&from_source=webtop_search&spm_id_from=333.1007&search_source=3
|
||||
), [Youtube](https://www.youtube.com/results?search_query=sadtalker&sp=CAM%253D) and [twitter #sadtalker](https://twitter.com/search?q=%23sadtalker&src=typed_query).
|
||||
|
||||
## 📋 Changelog (Previous changelog can be founded [here](docs/changlelog.md))
|
||||
|
||||
- __[2023.06.12]__: add more new features in WEBUI extension, see the discussion [here](https://github.com/OpenTalker/SadTalker/discussions/386).
|
||||
|
||||
- __[2023.06.05]__: release a new 512 beta face model. Fixed some bugs and improve the performance.
|
||||
|
||||
- __[2023.04.15]__: Adding automatic1111 colab by @camenduru, thanks for this awesome colab: [](https://colab.research.google.com/github/camenduru/stable-diffusion-webui-colab/blob/main/video/stable/stable_diffusion_1_5_video_webui_colab.ipynb).
|
||||
|
||||
- __[2023.04.12]__: adding a more detailed sd-webui installation document, fixed reinstallation problem.
|
||||
|
||||
- __[2023.04.12]__: Fixed the sd-webui safe issues becasue of the 3rd packages, optimize the output path in `sd-webui-extension`.
|
||||
|
||||
- __[2023.04.08]__: ❗️❗️❗️ In v0.0.2, we add a logo watermark to the generated video to prevent abusing since it is very realistic.
|
||||
|
||||
- __[2023.04.08]__: v0.0.2, full image animation, adding baidu driver for download checkpoints. Optimizing the logic about enhancer.
|
||||
|
||||
|
||||
## 🚧 TODO: See the Discussion https://github.com/OpenTalker/SadTalker/issues/280
|
||||
|
||||
## If you have any problem, please view our [FAQ](docs/FAQ.md) before opening an issue.
|
||||
|
||||
|
||||
|
||||
## ⚙️ 1. Installation.
|
||||
|
||||
Tutorials from communities: [中文windows教程](https://www.bilibili.com/video/BV1Dc411W7V6/) | [日本語コース](https://br-d.fanbox.cc/posts/5685086?utm_campaign=manage_post_page&utm_medium=share&utm_source=twitter)
|
||||
|
||||
### Linux:
|
||||
|
||||
1. Installing [anaconda](https://www.anaconda.com/), python and git.
|
||||
|
||||
2. Creating the env and install the requirements.
|
||||
```bash
|
||||
git clone https://github.com/Winfredy/SadTalker.git
|
||||
|
||||
cd SadTalker
|
||||
|
||||
conda create -n sadtalker python=3.8
|
||||
|
||||
conda activate sadtalker
|
||||
|
||||
pip install torch==1.12.1+cu113 torchvision==0.13.1+cu113 torchaudio==0.12.1 --extra-index-url https://download.pytorch.org/whl/cu113
|
||||
|
||||
conda install ffmpeg
|
||||
|
||||
pip install -r requirements.txt
|
||||
|
||||
### tts is optional for gradio demo.
|
||||
### pip install TTS
|
||||
|
||||
```
|
||||
### Windows ([中文windows教程](https://www.bilibili.com/video/BV1Dc411W7V6/)):
|
||||
|
||||
1. Install [Python 3.10.6](https://www.python.org/downloads/windows/), checking "Add Python to PATH".
|
||||
2. Install [git](https://git-scm.com/download/win) manually (OR `scoop install git` via [scoop](https://scoop.sh/)).
|
||||
3. Install `ffmpeg`, following [this instruction](https://www.wikihow.com/Install-FFmpeg-on-Windows) (OR using `scoop install ffmpeg` via [scoop](https://scoop.sh/)).
|
||||
4. Download our SadTalker repository, for example by running `git clone https://github.com/Winfredy/SadTalker.git`.
|
||||
5. Download the `checkpoint` and `gfpgan` [below↓](https://github.com/Winfredy/SadTalker#-2-download-trained-models).
|
||||
5. Run `start.bat` from Windows Explorer as normal, non-administrator, user, a gradio WebUI demo will be started.
|
||||
|
||||
### Macbook:
|
||||
|
||||
More tips about installnation on Macbook and the Docker file can be founded [here](docs/install.md)
|
||||
|
||||
## 📥 2. Download Trained Models.
|
||||
|
||||
You can run the following script to put all the models in the right place.
|
||||
|
||||
```bash
|
||||
bash scripts/download_models.sh
|
||||
```
|
||||
|
||||
Other alternatives:
|
||||
> we also provide an offline patch (`gfpgan/`), thus, no model will be downloaded when generating.
|
||||
|
||||
**Google Driver**: download our pre-trained model from [ this link (main checkpoints)](https://drive.google.com/file/d/1gwWh45pF7aelNP_P78uDJL8Sycep-K7j/view?usp=sharing) and [ gfpgan (offline patch)](https://drive.google.com/file/d/19AIBsmfcHW6BRJmeqSFlG5fL445Xmsyi?usp=sharing)
|
||||
|
||||
**Github Release Page**: download all the files from the [lastest github release page](https://github.com/Winfredy/SadTalker/releases), and then, put it in ./checkpoints.
|
||||
|
||||
**百度云盘**: we provided the downloaded model in [checkpoints, 提取码: sadt.](https://pan.baidu.com/s/1P4fRgk9gaSutZnn8YW034Q?pwd=sadt) And [gfpgan, 提取码: sadt.](https://pan.baidu.com/s/1kb1BCPaLOWX1JJb9Czbn6w?pwd=sadt)
|
||||
|
||||
|
||||
|
||||
<details><summary>Model Details</summary>
|
||||
|
||||
|
||||
Model explains:
|
||||
|
||||
##### New version
|
||||
| Model | Description
|
||||
| :--- | :----------
|
||||
|checkpoints/mapping_00229-model.pth.tar | Pre-trained MappingNet in Sadtalker.
|
||||
|checkpoints/mapping_00109-model.pth.tar | Pre-trained MappingNet in Sadtalker.
|
||||
|checkpoints/SadTalker_V0.0.2_256.safetensors | packaged sadtalker checkpoints of old version, 256 face render).
|
||||
|checkpoints/SadTalker_V0.0.2_512.safetensors | packaged sadtalker checkpoints of old version, 512 face render).
|
||||
|gfpgan/weights | Face detection and enhanced models used in `facexlib` and `gfpgan`.
|
||||
|
||||
|
||||
##### Old version
|
||||
| Model | Description
|
||||
| :--- | :----------
|
||||
|checkpoints/auido2exp_00300-model.pth | Pre-trained ExpNet in Sadtalker.
|
||||
|checkpoints/auido2pose_00140-model.pth | Pre-trained PoseVAE in Sadtalker.
|
||||
|checkpoints/mapping_00229-model.pth.tar | Pre-trained MappingNet in Sadtalker.
|
||||
|checkpoints/mapping_00109-model.pth.tar | Pre-trained MappingNet in Sadtalker.
|
||||
|checkpoints/facevid2vid_00189-model.pth.tar | Pre-trained face-vid2vid model from [the reappearance of face-vid2vid](https://github.com/zhanglonghao1992/One-Shot_Free-View_Neural_Talking_Head_Synthesis).
|
||||
|checkpoints/epoch_20.pth | Pre-trained 3DMM extractor in [Deep3DFaceReconstruction](https://github.com/microsoft/Deep3DFaceReconstruction).
|
||||
|checkpoints/wav2lip.pth | Highly accurate lip-sync model in [Wav2lip](https://github.com/Rudrabha/Wav2Lip).
|
||||
|checkpoints/shape_predictor_68_face_landmarks.dat | Face landmark model used in [dilb](http://dlib.net/).
|
||||
|checkpoints/BFM | 3DMM library file.
|
||||
|checkpoints/hub | Face detection models used in [face alignment](https://github.com/1adrianb/face-alignment).
|
||||
|gfpgan/weights | Face detection and enhanced models used in `facexlib` and `gfpgan`.
|
||||
|
||||
The final folder will be shown as:
|
||||
|
||||
<img width="331" alt="image" src="https://user-images.githubusercontent.com/4397546/232511411-4ca75cbf-a434-48c5-9ae0-9009e8316484.png">
|
||||
|
||||
|
||||
</details>
|
||||
|
||||
## 🔮 3. Quick Start ([Best Practice](docs/best_practice.md)).
|
||||
|
||||
### WebUI Demos:
|
||||
|
||||
**Online**: [Huggingface](https://huggingface.co/spaces/vinthony/SadTalker) | [SDWebUI-Colab](https://colab.research.google.com/github/camenduru/stable-diffusion-webui-colab/blob/main/video/stable/stable_diffusion_1_5_video_webui_colab.ipynb) | [Colab](https://colab.research.google.com/github/Winfredy/SadTalker/blob/main/quick_demo.ipynb)
|
||||
|
||||
**Local Autiomatic1111 stable-diffusion webui extension**: please refer to [Autiomatic1111 stable-diffusion webui docs](docs/webui_extension.md).
|
||||
|
||||
**Local gradio demo(highly recommanded!)**: Similar to our [hugging-face demo](https://huggingface.co/spaces/vinthony/SadTalker) can be run by:
|
||||
|
||||
```bash
|
||||
## you need manually install TTS(https://github.com/coqui-ai/TTS) via `pip install tts` in advanced.
|
||||
python app.py
|
||||
```
|
||||
|
||||
**Local gradio demo(highly recommanded!)**:
|
||||
|
||||
- windows: just double click `webui.bat`, the requirements will be installed automatically.
|
||||
- Linux/Mac OS: run `bash webui.sh` to start the webui.
|
||||
|
||||
|
||||
### Manually usages:
|
||||
|
||||
##### Animating a portrait image from default config:
|
||||
```bash
|
||||
python inference.py --driven_audio <audio.wav> \
|
||||
--source_image <video.mp4 or picture.png> \
|
||||
--enhancer gfpgan
|
||||
```
|
||||
The results will be saved in `results/$SOME_TIMESTAMP/*.mp4`.
|
||||
|
||||
##### Full body/image Generation:
|
||||
|
||||
Using `--still` to generate a natural full body video. You can add `enhancer` to improve the quality of the generated video.
|
||||
|
||||
```bash
|
||||
python inference.py --driven_audio <audio.wav> \
|
||||
--source_image <video.mp4 or picture.png> \
|
||||
--result_dir <a file to store results> \
|
||||
--still \
|
||||
--preprocess full \
|
||||
--enhancer gfpgan
|
||||
```
|
||||
|
||||
More examples and configuration and tips can be founded in the [ >>> best practice documents <<<](docs/best_practice.md).
|
||||
|
||||
## 🛎 Citation
|
||||
|
||||
If you find our work useful in your research, please consider citing:
|
||||
|
||||
```bibtex
|
||||
@article{zhang2022sadtalker,
|
||||
title={SadTalker: Learning Realistic 3D Motion Coefficients for Stylized Audio-Driven Single Image Talking Face Animation},
|
||||
author={Zhang, Wenxuan and Cun, Xiaodong and Wang, Xuan and Zhang, Yong and Shen, Xi and Guo, Yu and Shan, Ying and Wang, Fei},
|
||||
journal={arXiv preprint arXiv:2211.12194},
|
||||
year={2022}
|
||||
}
|
||||
```
|
||||
|
||||
|
||||
|
||||
## 💗 Acknowledgements
|
||||
|
||||
Facerender code borrows heavily from [zhanglonghao's reproduction of face-vid2vid](https://github.com/zhanglonghao1992/One-Shot_Free-View_Neural_Talking_Head_Synthesis) and [PIRender](https://github.com/RenYurui/PIRender). We thank the authors for sharing their wonderful code. In training process, We also use the model from [Deep3DFaceReconstruction](https://github.com/microsoft/Deep3DFaceReconstruction) and [Wav2lip](https://github.com/Rudrabha/Wav2Lip). We thank for their wonderful work.
|
||||
|
||||
See also these wonderful 3rd libraries we use:
|
||||
|
||||
- **Face Utils**: https://github.com/xinntao/facexlib
|
||||
- **Face Enhancement**: https://github.com/TencentARC/GFPGAN
|
||||
- **Image/Video Enhancement**:https://github.com/xinntao/Real-ESRGAN
|
||||
|
||||
## 🥂 Extensions:
|
||||
|
||||
- [SadTalker-Video-Lip-Sync](https://github.com/Zz-ww/SadTalker-Video-Lip-Sync) from [@Zz-ww](https://github.com/Zz-ww): SadTalker for Video Lip Editing
|
||||
|
||||
## 🥂 Related Works
|
||||
- [StyleHEAT: One-Shot High-Resolution Editable Talking Face Generation via Pre-trained StyleGAN (ECCV 2022)](https://github.com/FeiiYin/StyleHEAT)
|
||||
- [CodeTalker: Speech-Driven 3D Facial Animation with Discrete Motion Prior (CVPR 2023)](https://github.com/Doubiiu/CodeTalker)
|
||||
- [VideoReTalking: Audio-based Lip Synchronization for Talking Head Video Editing In the Wild (SIGGRAPH Asia 2022)](https://github.com/vinthony/video-retalking)
|
||||
- [DPE: Disentanglement of Pose and Expression for General Video Portrait Editing (CVPR 2023)](https://github.com/Carlyx/DPE)
|
||||
- [3D GAN Inversion with Facial Symmetry Prior (CVPR 2023)](https://github.com/FeiiYin/SPI/)
|
||||
- [T2M-GPT: Generating Human Motion from Textual Descriptions with Discrete Representations (CVPR 2023)](https://github.com/Mael-zys/T2M-GPT)
|
||||
|
||||
## 📢 Disclaimer
|
||||
|
||||
This is not an official product of Tencent. This repository can only be used for personal/research/non-commercial purposes.
|
||||
|
||||
LOGO: color and font suggestion: [ChatGPT](ai.com), logo font：[Montserrat Alternates
|
||||
](https://fonts.google.com/specimen/Montserrat+Alternates?preview.text=SadTalker&preview.text_type=custom&query=mont).
|
||||
|
||||
All the copyright of the demo images and audio are from communities users or the geneartion from stable diffusion. Free free to contact us if you feel uncomfortable.
|
||||
|
||||
@@ -0,0 +1,111 @@
|
||||
import os, sys
|
||||
import gradio as gr
|
||||
from src.gradio_demo import SadTalker
|
||||
|
||||
|
||||
try:
|
||||
import webui # in webui
|
||||
in_webui = True
|
||||
except:
|
||||
in_webui = False
|
||||
|
||||
|
||||
def toggle_audio_file(choice):
|
||||
if choice == False:
|
||||
return gr.update(visible=True), gr.update(visible=False)
|
||||
else:
|
||||
return gr.update(visible=False), gr.update(visible=True)
|
||||
|
||||
def ref_video_fn(path_of_ref_video):
|
||||
if path_of_ref_video is not None:
|
||||
return gr.update(value=True)
|
||||
else:
|
||||
return gr.update(value=False)
|
||||
|
||||
def sadtalker_demo(checkpoint_path='checkpoints', config_path='src/config', warpfn=None):
|
||||
|
||||
sad_talker = SadTalker(checkpoint_path, config_path, lazy_load=True)
|
||||
|
||||
with gr.Blocks(analytics_enabled=False) as sadtalker_interface:
|
||||
gr.Markdown("<div align='center'> <h2> 😭 SadTalker: Learning Realistic 3D Motion Coefficients for Stylized Audio-Driven Single Image Talking Face Animation (CVPR 2023) </span> </h2> \
|
||||
<a style='font-size:18px;color: #efefef' href='https://arxiv.org/abs/2211.12194'>Arxiv</a> \
|
||||
<a style='font-size:18px;color: #efefef' href='https://sadtalker.github.io'>Homepage</a> \
|
||||
<a style='font-size:18px;color: #efefef' href='https://github.com/Winfredy/SadTalker'> Github </div>")
|
||||
|
||||
with gr.Row().style(equal_height=False):
|
||||
with gr.Column(variant='panel'):
|
||||
with gr.Tabs(elem_id="sadtalker_source_image"):
|
||||
with gr.TabItem('Upload image'):
|
||||
with gr.Row():
|
||||
source_image = gr.Image(label="Source image", source="upload", type="filepath", elem_id="img2img_image").style(width=512)
|
||||
|
||||
with gr.Tabs(elem_id="sadtalker_driven_audio"):
|
||||
with gr.TabItem('Upload OR TTS'):
|
||||
with gr.Column(variant='panel'):
|
||||
driven_audio = gr.Audio(label="Input audio", source="upload", type="filepath")
|
||||
|
||||
if sys.platform != 'win32' and not in_webui:
|
||||
from src.utils.text2speech import TTSTalker
|
||||
tts_talker = TTSTalker()
|
||||
with gr.Column(variant='panel'):
|
||||
input_text = gr.Textbox(label="Generating audio from text", lines=5, placeholder="please enter some text here, we genreate the audio from text using @Coqui.ai TTS.")
|
||||
tts = gr.Button('Generate audio',elem_id="sadtalker_audio_generate", variant='primary')
|
||||
tts.click(fn=tts_talker.test, inputs=[input_text], outputs=[driven_audio])
|
||||
|
||||
with gr.Column(variant='panel'):
|
||||
with gr.Tabs(elem_id="sadtalker_checkbox"):
|
||||
with gr.TabItem('Settings'):
|
||||
gr.Markdown("need help? please visit our [best practice page](https://github.com/OpenTalker/SadTalker/blob/main/docs/best_practice.md) for more detials")
|
||||
with gr.Column(variant='panel'):
|
||||
# width = gr.Slider(minimum=64, elem_id="img2img_width", maximum=2048, step=8, label="Manually Crop Width", value=512) # img2img_width
|
||||
# height = gr.Slider(minimum=64, elem_id="img2img_height", maximum=2048, step=8, label="Manually Crop Height", value=512) # img2img_width
|
||||
pose_style = gr.Slider(minimum=0, maximum=46, step=1, label="Pose style", value=0) #
|
||||
size_of_image = gr.Radio([256, 512], value=256, label='face model resolution', info="use 256/512 model?") #
|
||||
preprocess_type = gr.Radio(['crop', 'resize','full', 'extcrop', 'extfull'], value='crop', label='preprocess', info="How to handle input image?")
|
||||
is_still_mode = gr.Checkbox(label="Still Mode (fewer hand motion, works with preprocess `full`)")
|
||||
batch_size = gr.Slider(label="batch size in generation", step=1, maximum=10, value=2)
|
||||
enhancer = gr.Checkbox(label="GFPGAN as Face enhancer")
|
||||
submit = gr.Button('Generate', elem_id="sadtalker_generate", variant='primary')
|
||||
|
||||
with gr.Tabs(elem_id="sadtalker_genearted"):
|
||||
gen_video = gr.Video(label="Generated video", format="mp4").style(width=256)
|
||||
|
||||
if warpfn:
|
||||
submit.click(
|
||||
fn=warpfn(sad_talker.test),
|
||||
inputs=[source_image,
|
||||
driven_audio,
|
||||
preprocess_type,
|
||||
is_still_mode,
|
||||
enhancer,
|
||||
batch_size,
|
||||
size_of_image,
|
||||
pose_style
|
||||
],
|
||||
outputs=[gen_video]
|
||||
)
|
||||
else:
|
||||
submit.click(
|
||||
fn=sad_talker.test,
|
||||
inputs=[source_image,
|
||||
driven_audio,
|
||||
preprocess_type,
|
||||
is_still_mode,
|
||||
enhancer,
|
||||
batch_size,
|
||||
size_of_image,
|
||||
pose_style
|
||||
],
|
||||
outputs=[gen_video]
|
||||
)
|
||||
|
||||
return sadtalker_interface
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
demo = sadtalker_demo()
|
||||
demo.queue()
|
||||
demo.launch()
|
||||
|
||||
|
||||
@@ -0,0 +1,35 @@
|
||||
build:
|
||||
gpu: true
|
||||
cuda: "11.3"
|
||||
python_version: "3.8"
|
||||
system_packages:
|
||||
- "ffmpeg"
|
||||
- "libgl1-mesa-glx"
|
||||
- "libglib2.0-0"
|
||||
python_packages:
|
||||
- "torch==1.12.1"
|
||||
- "torchvision==0.13.1"
|
||||
- "torchaudio==0.12.1"
|
||||
- "joblib==1.1.0"
|
||||
- "scikit-image==0.19.3"
|
||||
- "basicsr==1.4.2"
|
||||
- "facexlib==0.3.0"
|
||||
- "resampy==0.3.1"
|
||||
- "pydub==0.25.1"
|
||||
- "scipy==1.10.1"
|
||||
- "kornia==0.6.8"
|
||||
- "face_alignment==1.3.5"
|
||||
- "imageio==2.19.3"
|
||||
- "imageio-ffmpeg==0.4.7"
|
||||
- "librosa==0.9.2" #
|
||||
- "tqdm==4.65.0"
|
||||
- "yacs==0.1.8"
|
||||
- "gfpgan==1.3.8"
|
||||
- "dlib-bin==19.24.1"
|
||||
- "av==10.0.0"
|
||||
- "trimesh==3.9.20"
|
||||
run:
|
||||
- mkdir -p /root/.cache/torch/hub/checkpoints/ && wget --output-document "/root/.cache/torch/hub/checkpoints/s3fd-619a316812.pth" "https://www.adrianbulat.com/downloads/python-fan/s3fd-619a316812.pth"
|
||||
- mkdir -p /root/.cache/torch/hub/checkpoints/ && wget --output-document "/root/.cache/torch/hub/checkpoints/2DFAN4-cd938726ad.zip" "https://www.adrianbulat.com/downloads/python-fan/2DFAN4-cd938726ad.zip"
|
||||
|
||||
predict: "predict.py:Predictor"
|
||||
@@ -0,0 +1,46 @@
|
||||
|
||||
## Frequency Asked Question
|
||||
|
||||
**Q: `ffmpeg` is not recognized as an internal or external command**
|
||||
|
||||
In Linux, you can install the ffmpeg via `conda install ffmpeg`. Or on Mac OS X, try to install ffmpeg via `brew install ffmpeg`. On windows, make sure you have `ffmpeg` in the `%PATH%` as suggested in [#54](https://github.com/Winfredy/SadTalker/issues/54), then, following [this](https://www.geeksforgeeks.org/how-to-install-ffmpeg-on-windows/) installation to install `ffmpeg`.
|
||||
|
||||
**Q: Running Requirments.**
|
||||
|
||||
Please refer to the discussion here: https://github.com/Winfredy/SadTalker/issues/124#issuecomment-1508113989
|
||||
|
||||
|
||||
**Q: ModuleNotFoundError: No module named 'ai'**
|
||||
|
||||
please check the checkpoint's size of the `epoch_20.pth`. (https://github.com/Winfredy/SadTalker/issues/167, https://github.com/Winfredy/SadTalker/issues/113)
|
||||
|
||||
**Q: Illegal Hardware Error: Mac M1**
|
||||
|
||||
please reinstall the `dlib` by `pip install dlib` individually. (https://github.com/Winfredy/SadTalker/issues/129, https://github.com/Winfredy/SadTalker/issues/109)
|
||||
|
||||
|
||||
**Q: FileNotFoundError: [Errno 2] No such file or directory: checkpoints\BFM_Fitting\similarity_Lm3D_all.mat**
|
||||
|
||||
Make sure you have downloaded the checkpoints and gfpgan as [here](https://github.com/Winfredy/SadTalker#-2-download-trained-models) and placed them in the right place.
|
||||
|
||||
**Q: RuntimeError: unexpected EOF, expected 237192 more bytes. The file might be corrupted.**
|
||||
|
||||
The files are not automatically downloaded. Please update the code and download the gfpgan folders as [here](https://github.com/Winfredy/SadTalker#-2-download-trained-models).
|
||||
|
||||
**Q: CUDA out of memory error**
|
||||
|
||||
please refer to https://stackoverflow.com/questions/73747731/runtimeerror-cuda-out-of-memory-how-setting-max-split-size-mb
|
||||
|
||||
```
|
||||
# windows
|
||||
set PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:128
|
||||
python inference.py ...
|
||||
|
||||
# linux
|
||||
export PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:128
|
||||
python inference.py ...
|
||||
```
|
||||
|
||||
**Q: Error while decoding stream #0:0: Invalid data found when processing input [mp3float @ 0000015037628c00] Header missing**
|
||||
|
||||
Our method only support wav or mp3 files as input, please make sure the feeded audios are in these formats.
|
||||
@@ -0,0 +1,94 @@
|
||||
# Best Practice and Tips for configuration
|
||||
|
||||
> Our model only works on REAL person's photo or the portrait image similar to REAL person. The anime talking head genreation method will be released in future.
|
||||
|
||||
Advanced confiurations for `inference.py`:
|
||||
|
||||
| Name | Configuration | default | Explaination |
|
||||
|:------------- |:------------- |:----- | :------------- |
|
||||
| Enhance Mode | `--enhancer` | None | Using `gfpgan` or `RestoreFormer` to enhance the generated face via face restoration network
|
||||
| Background Enhancer | `--background_enhancer` | None | Using `realesrgan` to enhance the full video.
|
||||
| Still Mode | ` --still` | False | Using the same pose parameters as the original image, fewer head motion.
|
||||
| Expressive Mode | `--expression_scale` | 1.0 | a larger value will make the expression motion stronger.
|
||||
| save path | `--result_dir` |`./results` | The file will be save in the newer location.
|
||||
| preprocess | `--preprocess` | `crop` | Run and produce the results in the croped input image. Other choices: `resize`, where the images will be resized to the specific resolution. `full` Run the full image animation, use with `--still` to get better results.
|
||||
| ref Mode (eye) | `--ref_eyeblink` | None | A video path, where we borrow the eyeblink from this reference video to provide more natural eyebrow movement.
|
||||
| ref Mode (pose) | `--ref_pose` | None | A video path, where we borrow the pose from the head reference video.
|
||||
| 3D Mode | `--face3dvis` | False | Need additional installation. More details to generate the 3d face can be founded [here](docs/face3d.md).
|
||||
| free-view Mode | `--input_yaw`,<br> `--input_pitch`,<br> `--input_roll` | None | Genearting novel view or free-view 4D talking head from a single image. More details can be founded [here](https://github.com/Winfredy/SadTalker#generating-4d-free-view-talking-examples-from-audio-and-a-single-image).
|
||||
|
||||
|
||||
### About `--preprocess`
|
||||
|
||||
Our method automatically handle the input images via `crop`, `resize` and `full`.
|
||||
|
||||
In `crop` mode, we only generate the croped image via the facial keypoints and generated the facial anime avator. The animation of both expression and head pose are realistic.
|
||||
|
||||
> still mode will stop the eyeblink and head pose movement.
|
||||
|
||||
| [input image @bagbag1815](https://twitter.com/bagbag1815/status/1642754319094108161) | crop | crop w/still |
|
||||
|:--------------------: |:--------------------: | :----: |
|
||||
| <img src='../examples/source_image/full_body_2.png' width='380'> |  |  |
|
||||
|
||||
|
||||
In `resize` mode, we resize the whole images to generate the fully talking head video. Thus, an image similar to the ID photo can be produced. ⚠️ It will produce bad results for full person images.
|
||||
|
||||
|
||||
|
||||
|
||||
| <img src='../examples/source_image/full_body_2.png' width='380'> | <img src='../examples/source_image/full4.jpeg' width='380'> |
|
||||
|:--------------------: |:--------------------: |
|
||||
| ❌ not suitable for resize mode | ✅ good for resize mode |
|
||||
| <img src='resize_no.gif'> | <img src='resize_good.gif' width='380'> |
|
||||
|
||||
In `full` mode, our model will automatically process the croped region and paste back to the original image. Remember to use `--still` to keep the original head pose.
|
||||
|
||||
| input | `--still` | `--still` & `enhancer` |
|
||||
|:--------------------: |:--------------------: | :--:|
|
||||
| <img src='../examples/source_image/full_body_2.png' width='380'> | <img src='./example_full.gif' width='380'> | <img src='./example_full_enhanced.gif' width='380'>
|
||||
|
||||
|
||||
### About `--enhancer`
|
||||
|
||||
For better facial quality, we intergate [gfpgan](https://github.com/TencentARC/GFPGAN) and [real-esrgan](https://github.com/xinntao/Real-ESRGAN) for different purpose. Just adding `--enhancer <gfpgan or RestoreFormer>` or `--background_enhancer <realesrgan>` for the enhancement of the face and the full image.
|
||||
|
||||
```bash
|
||||
# make sure above packages are available:
|
||||
pip install gfpgan
|
||||
pip install realesrgan
|
||||
```
|
||||
|
||||
### About `--face3dvis`
|
||||
|
||||
This flag indicate that we can generated the 3d-rendered face and it's 3d facial landmarks. More details can be founded [here](face3d.md).
|
||||
|
||||
| Input | Animated 3d face |
|
||||
|:-------------: | :-------------: |
|
||||
| <img src='../examples/source_image/art_0.png' width='200px'> | <video src="https://user-images.githubusercontent.com/4397546/226856847-5a6a0a4d-a5ec-49e2-9b05-3206db65e8e3.mp4"></video> |
|
||||
|
||||
> Kindly ensure to activate the audio as the default audio playing is incompatible with GitHub.
|
||||
|
||||
|
||||
|
||||
#### reference eye-link mode.
|
||||
|
||||
| Input, w/ reference video , reference video |
|
||||
|:-------------: |
|
||||
| |
|
||||
| If the reference video is shorter than the input audio, we will loop the reference video .
|
||||
|
||||
|
||||
|
||||
#### Generating 4D free-view talking examples from audio and a single image
|
||||
|
||||
We use `input_yaw`, `input_pitch`, `input_roll` to control head pose. For example, `--input_yaw -20 30 10` means the input head yaw degree changes from -20 to 30 and then changes from 30 to 10.
|
||||
```bash
|
||||
python inference.py --driven_audio <audio.wav> \
|
||||
--source_image <video.mp4 or picture.png> \
|
||||
--result_dir <a file to store results> \
|
||||
--input_yaw -20 30 10
|
||||
```
|
||||
|
||||
| Results, Free-view results, Novel view results |
|
||||
|:-------------: |
|
||||
| |
|
||||
@@ -0,0 +1,29 @@
|
||||
## changelogs
|
||||
|
||||
|
||||
- __[2023.04.06]__: stable-diffiusion webui extension is release.
|
||||
|
||||
- __[2023.04.03]__: Enable TTS in huggingface and gradio local demo.
|
||||
|
||||
- __[2023.03.30]__: Launch beta version of the full body mode.
|
||||
|
||||
- __[2023.03.30]__: Launch new feature: through using reference videos, our algorithm can generate videos with more natural eye blinking and some eyebrow movement.
|
||||
|
||||
- __[2023.03.29]__: `resize mode` is online by `python infererence.py --preprocess resize`! Where we can produce a larger crop of the image as discussed in https://github.com/Winfredy/SadTalker/issues/35.
|
||||
|
||||
- __[2023.03.29]__: local gradio demo is online! `python app.py` to start the demo. New `requirments.txt` is used to avoid the bugs in `librosa`.
|
||||
|
||||
- __[2023.03.28]__: Online demo is launched in [](https://huggingface.co/spaces/vinthony/SadTalker), thanks AK!
|
||||
|
||||
- __[2023.03.22]__: Launch new feature: generating the 3d face animation from a single image. New applications about it will be updated.
|
||||
|
||||
- __[2023.03.22]__: Launch new feature: `still mode`, where only a small head pose will be produced via `python inference.py --still`.
|
||||
|
||||
- __[2023.03.18]__: Support `expression intensity`, now you can change the intensity of the generated motion: `python inference.py --expression_scale 1.3 (some value > 1)`.
|
||||
|
||||
- __[2023.03.18]__: Reconfig the data folders, now you can download the checkpoint automatically using `bash scripts/download_models.sh`.
|
||||
- __[2023.03.18]__: We have offically integrate the [GFPGAN](https://github.com/TencentARC/GFPGAN) for face enhancement, using `python inference.py --enhancer gfpgan` for better visualization performance.
|
||||
- __[2023.03.14]__: Specify the version of package `joblib` to remove the errors in using `librosa`, [](https://colab.research.google.com/github/Winfredy/SadTalker/blob/main/quick_demo.ipynb) is online!
|
||||
- __[2023.03.06]__: Solve some bugs in code and errors in installation
|
||||
- __[2023.03.03]__: Release the test code for audio-driven single image animation!
|
||||
- __[2023.02.28]__: SadTalker has been accepted by CVPR 2023!
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 1.5 MiB |
Binary file not shown.
|
After Width: | Height: | Size: 1.2 MiB |
Binary file not shown.
|
After Width: | Height: | Size: 1.4 MiB |
Binary file not shown.
|
After Width: | Height: | Size: 817 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 5.5 MiB |
@@ -0,0 +1,48 @@
|
||||
## 3D Face visualization
|
||||
|
||||
We use pytorch3d to visualize the produced 3d face from a single image.
|
||||
|
||||
Since it is not easy to install, we produce a new install guidence here:
|
||||
|
||||
```bash
|
||||
git clone https://github.com/Winfredy/SadTalker.git
|
||||
cd SadTalker
|
||||
conda create -n sadtalker3d python=3.8
|
||||
source activate sadtalker3d
|
||||
|
||||
conda install ffmpeg
|
||||
conda install -c fvcore -c iopath -c conda-forge fvcore iopath
|
||||
conda install libgcc gmp
|
||||
|
||||
pip install torch==1.11.0+cu113 torchvision==0.12.0+cu113 torchaudio==0.11.0 --extra-index-url https://download.pytorch.org/whl/cu113
|
||||
|
||||
# insintall pytorch3d
|
||||
pip install --no-index --no-cache-dir pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/py38_cu113_pyt1110/download.html
|
||||
|
||||
pip install -r requirements3d.txt
|
||||
|
||||
### install gpfgan for enhancer
|
||||
pip install git+https://github.com/TencentARC/GFPGAN
|
||||
|
||||
|
||||
### when occurs gcc version problem `from pytorch import _C` from pytorch3d, add the anaconda path to LD_LIBRARY_PATH
|
||||
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/$YOUR_ANACONDA_PATH/lib/
|
||||
|
||||
```
|
||||
|
||||
|
||||
Then, generating the result via:
|
||||
|
||||
```bash
|
||||
|
||||
|
||||
python inference.py --driven_audio <audio.wav> \
|
||||
--source_image <video.mp4 or picture.png> \
|
||||
--result_dir <a file to store results> \
|
||||
--face3dvis
|
||||
|
||||
```
|
||||
|
||||
Then, the results will be given in the folders with the file name of `face3d.mp4`.
|
||||
|
||||
More applications about 3d face will be released.
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 5.4 MiB |
@@ -0,0 +1,47 @@
|
||||
|
||||
### Mac (Tested on M1 Mac OS 13.3)
|
||||
|
||||
```
|
||||
git clone https://github.com/Winfredy/SadTalker.git
|
||||
|
||||
cd SadTalker
|
||||
|
||||
conda create -n sadtalker python=3.8
|
||||
|
||||
conda activate sadtalker
|
||||
|
||||
# install pytorch 2.0
|
||||
pip install torch torchvision torchaudio
|
||||
|
||||
conda install ffmpeg
|
||||
|
||||
pip install -r requirements.txt
|
||||
|
||||
pip install dlib # mac need to install the original dlib.
|
||||
|
||||
```
|
||||
|
||||
|
||||
|
||||
### Windows Native
|
||||
|
||||
- Make sure you have `ffmpeg` in the `%PATH%` as suggested in [#54](https://github.com/Winfredy/SadTalker/issues/54), following [this](https://www.geeksforgeeks.org/how-to-install-ffmpeg-on-windows/) installation to install `ffmpeg`.
|
||||
|
||||
|
||||
### Windows WSL
|
||||
- Make sure the environment: `export LD_LIBRARY_PATH=/usr/lib/wsl/lib:$LD_LIBRARY_PATH`
|
||||
|
||||
|
||||
### Docker installnation
|
||||
|
||||
A dockerfile are also provided by [@thegenerativegeneration](https://github.com/thegenerativegeneration) in [docker hub](https://hub.docker.com/repository/docker/wawa9000/sadtalker), which can be used directly as:
|
||||
|
||||
```bash
|
||||
docker run --gpus "all" --rm -v $(pwd):/host_dir wawa9000/sadtalker \
|
||||
--driven_audio /host_dir/deyu.wav \
|
||||
--source_image /host_dir/image.jpg \
|
||||
--expression_scale 1.0 \
|
||||
--still \
|
||||
--result_dir /host_dir
|
||||
```
|
||||
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 1.6 MiB |
Binary file not shown.
|
After Width: | Height: | Size: 2.0 MiB |
Binary file not shown.
|
After Width: | Height: | Size: 34 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 7.7 MiB |
@@ -0,0 +1,50 @@
|
||||
|
||||
## Run SadTalker as a Stable Diffusion WebUI Extension.
|
||||
|
||||
1. Installing the lastest version of [stable-diffusion-webui](https://github.com/AUTOMATIC1111/stable-diffusion-webui) and install the sadtalker via `extension`.
|
||||
<img width="726" alt="image" src="https://user-images.githubusercontent.com/4397546/230698519-267d1d1f-6e99-4dd4-81e1-7b889259efbd.png">
|
||||
|
||||
2. Download the checkpoints manually, for Linux and Mac:
|
||||
|
||||
```bash
|
||||
|
||||
cd SOMEWHERE_YOU_LIKE
|
||||
|
||||
bash <(wget -qO- https://raw.githubusercontent.com/Winfredy/SadTalker/main/scripts/download_models.sh)
|
||||
```
|
||||
|
||||
For windows, you can download all the checkpoints from [google drive](https://drive.google.com/drive/folders/1Wd88VDoLhVzYsQ30_qDVluQr_Xm46yHT?usp=sharing) or [百度云盘](https://pan.baidu.com/s/1nXuVNd0exUl37ISwWqbFGA?pwd=sadt) 提取码: sadt.
|
||||
|
||||
3.1. options 1: put the checkpoint in `stable-diffusion-webui/models/SadTalker` or `stable-diffusion-webui/extensions/SadTalker/checkpoints/`, the checkpoints will be detected automatically.
|
||||
|
||||
3.2. Options 2: Set the path of `SADTALKTER_CHECKPOINTS` in `webui_user.sh`(linux) or `webui_user.bat`(windows) by:
|
||||
|
||||
> only works if you are directly starting webui from `webui_user.sh` or `webui_user.bat`.
|
||||
|
||||
```bash
|
||||
# windows (webui_user.bat)
|
||||
set SADTALKER_CHECKPOINTS=D:\SadTalker\checkpoints
|
||||
|
||||
# linux (webui_user.sh)
|
||||
export SADTALKER_CHECKPOINTS=/path/to/SadTalker/checkpoints
|
||||
```
|
||||
|
||||
4. Then, starting the webui via `webui.sh or webui_user.sh(linux)` or `webui_user.bat(windows)` or any other methods, the SadTalker can be used in stable-diffusion-webui directly.
|
||||
|
||||
<img width="726" alt="image" src="https://user-images.githubusercontent.com/4397546/230698614-58015182-2916-4240-b324-e69022ef75b3.png">
|
||||
|
||||
## Questsions
|
||||
|
||||
1. if you are running on CPU, you need to specific `--disable-safe-unpickle` in `webui_user.sh` or `webui_user.bat`.
|
||||
|
||||
```bash
|
||||
# windows (webui_user.bat)
|
||||
set COMMANDLINE_ARGS="--disable-safe-unpickle"
|
||||
|
||||
# linux (webui_user.sh)
|
||||
export COMMANDLINE_ARGS="--disable-safe-unpickle"
|
||||
```
|
||||
|
||||
|
||||
|
||||
(Some [important discussion](https://github.com/Winfredy/SadTalker/issues/78) if you are unable to use `full` mode).
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
|
After Width: | Height: | Size: 733 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 478 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 556 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 478 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 704 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 617 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 635 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 657 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 1.4 MiB |
Binary file not shown.
|
After Width: | Height: | Size: 2.0 MiB |
Binary file not shown.
|
After Width: | Height: | Size: 115 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 462 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 812 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 694 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 1.3 MiB |
Binary file not shown.
|
After Width: | Height: | Size: 3.5 MiB |
Binary file not shown.
|
After Width: | Height: | Size: 1.2 MiB |
Binary file not shown.
|
After Width: | Height: | Size: 98 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 509 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 3.0 MiB |
Binary file not shown.
|
After Width: | Height: | Size: 1.2 MiB |
Binary file not shown.
|
After Width: | Height: | Size: 617 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 26 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 122 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 134 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 108 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 43 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 238 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 108 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 44 KiB |
@@ -0,0 +1,145 @@
|
||||
from glob import glob
|
||||
import shutil
|
||||
import torch
|
||||
from time import strftime
|
||||
import os, sys, time
|
||||
from argparse import ArgumentParser
|
||||
|
||||
from src.utils.preprocess import CropAndExtract
|
||||
from src.test_audio2coeff import Audio2Coeff
|
||||
from src.facerender.animate import AnimateFromCoeff
|
||||
from src.generate_batch import get_data
|
||||
from src.generate_facerender_batch import get_facerender_data
|
||||
from src.utils.init_path import init_path
|
||||
|
||||
def main(args):
|
||||
#torch.backends.cudnn.enabled = False
|
||||
|
||||
pic_path = args.source_image
|
||||
audio_path = args.driven_audio
|
||||
save_dir = os.path.join(args.result_dir, strftime("%Y_%m_%d_%H.%M.%S"))
|
||||
os.makedirs(save_dir, exist_ok=True)
|
||||
pose_style = args.pose_style
|
||||
device = args.device
|
||||
batch_size = args.batch_size
|
||||
input_yaw_list = args.input_yaw
|
||||
input_pitch_list = args.input_pitch
|
||||
input_roll_list = args.input_roll
|
||||
ref_eyeblink = args.ref_eyeblink
|
||||
ref_pose = args.ref_pose
|
||||
|
||||
current_root_path = os.path.split(sys.argv[0])[0]
|
||||
|
||||
sadtalker_paths = init_path(args.checkpoint_dir, os.path.join(current_root_path, 'src/config'), args.size, args.old_version, args.preprocess)
|
||||
|
||||
#init model
|
||||
preprocess_model = CropAndExtract(sadtalker_paths, device)
|
||||
|
||||
audio_to_coeff = Audio2Coeff(sadtalker_paths, device)
|
||||
|
||||
animate_from_coeff = AnimateFromCoeff(sadtalker_paths, device)
|
||||
|
||||
#crop image and extract 3dmm from image
|
||||
first_frame_dir = os.path.join(save_dir, 'first_frame_dir')
|
||||
os.makedirs(first_frame_dir, exist_ok=True)
|
||||
print('3DMM Extraction for source image')
|
||||
first_coeff_path, crop_pic_path, crop_info = preprocess_model.generate(pic_path, first_frame_dir, args.preprocess,\
|
||||
source_image_flag=True, pic_size=args.size)
|
||||
if first_coeff_path is None:
|
||||
print("Can't get the coeffs of the input")
|
||||
return
|
||||
|
||||
if ref_eyeblink is not None:
|
||||
ref_eyeblink_videoname = os.path.splitext(os.path.split(ref_eyeblink)[-1])[0]
|
||||
ref_eyeblink_frame_dir = os.path.join(save_dir, ref_eyeblink_videoname)
|
||||
os.makedirs(ref_eyeblink_frame_dir, exist_ok=True)
|
||||
print('3DMM Extraction for the reference video providing eye blinking')
|
||||
ref_eyeblink_coeff_path, _, _ = preprocess_model.generate(ref_eyeblink, ref_eyeblink_frame_dir, args.preprocess, source_image_flag=False)
|
||||
else:
|
||||
ref_eyeblink_coeff_path=None
|
||||
|
||||
if ref_pose is not None:
|
||||
if ref_pose == ref_eyeblink:
|
||||
ref_pose_coeff_path = ref_eyeblink_coeff_path
|
||||
else:
|
||||
ref_pose_videoname = os.path.splitext(os.path.split(ref_pose)[-1])[0]
|
||||
ref_pose_frame_dir = os.path.join(save_dir, ref_pose_videoname)
|
||||
os.makedirs(ref_pose_frame_dir, exist_ok=True)
|
||||
print('3DMM Extraction for the reference video providing pose')
|
||||
ref_pose_coeff_path, _, _ = preprocess_model.generate(ref_pose, ref_pose_frame_dir, args.preprocess, source_image_flag=False)
|
||||
else:
|
||||
ref_pose_coeff_path=None
|
||||
|
||||
#audio2ceoff
|
||||
batch = get_data(first_coeff_path, audio_path, device, ref_eyeblink_coeff_path, still=args.still)
|
||||
coeff_path = audio_to_coeff.generate(batch, save_dir, pose_style, ref_pose_coeff_path)
|
||||
|
||||
# 3dface render
|
||||
if args.face3dvis:
|
||||
from src.face3d.visualize import gen_composed_video
|
||||
gen_composed_video(args, device, first_coeff_path, coeff_path, audio_path, os.path.join(save_dir, '3dface.mp4'))
|
||||
|
||||
#coeff2video
|
||||
data = get_facerender_data(coeff_path, crop_pic_path, first_coeff_path, audio_path,
|
||||
batch_size, input_yaw_list, input_pitch_list, input_roll_list,
|
||||
expression_scale=args.expression_scale, still_mode=args.still, preprocess=args.preprocess, size=args.size)
|
||||
|
||||
result = animate_from_coeff.generate(data, save_dir, pic_path, crop_info, \
|
||||
enhancer=args.enhancer, background_enhancer=args.background_enhancer, preprocess=args.preprocess, img_size=args.size)
|
||||
|
||||
shutil.move(result, save_dir+'.mp4')
|
||||
print('The generated video is named:', save_dir+'.mp4')
|
||||
|
||||
if not args.verbose:
|
||||
shutil.rmtree(save_dir)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
|
||||
parser = ArgumentParser()
|
||||
parser.add_argument("--driven_audio", default='./examples/driven_audio/bus_chinese.wav', help="path to driven audio")
|
||||
parser.add_argument("--source_image", default='./examples/source_image/full_body_1.png', help="path to source image")
|
||||
parser.add_argument("--ref_eyeblink", default=None, help="path to reference video providing eye blinking")
|
||||
parser.add_argument("--ref_pose", default=None, help="path to reference video providing pose")
|
||||
parser.add_argument("--checkpoint_dir", default='./checkpoints', help="path to output")
|
||||
parser.add_argument("--result_dir", default='./results', help="path to output")
|
||||
parser.add_argument("--pose_style", type=int, default=0, help="input pose style from [0, 46)")
|
||||
parser.add_argument("--batch_size", type=int, default=2, help="the batch size of facerender")
|
||||
parser.add_argument("--size", type=int, default=256, help="the image size of the facerender")
|
||||
parser.add_argument("--expression_scale", type=float, default=1., help="the batch size of facerender")
|
||||
parser.add_argument('--input_yaw', nargs='+', type=int, default=None, help="the input yaw degree of the user ")
|
||||
parser.add_argument('--input_pitch', nargs='+', type=int, default=None, help="the input pitch degree of the user")
|
||||
parser.add_argument('--input_roll', nargs='+', type=int, default=None, help="the input roll degree of the user")
|
||||
parser.add_argument('--enhancer', type=str, default=None, help="Face enhancer, [gfpgan, RestoreFormer]")
|
||||
parser.add_argument('--background_enhancer', type=str, default=None, help="background enhancer, [realesrgan]")
|
||||
parser.add_argument("--cpu", dest="cpu", action="store_true")
|
||||
parser.add_argument("--face3dvis", action="store_true", help="generate 3d face and 3d landmarks")
|
||||
parser.add_argument("--still", action="store_true", help="can crop back to the original videos for the full body aniamtion")
|
||||
parser.add_argument("--preprocess", default='crop', choices=['crop', 'extcrop', 'resize', 'full', 'extfull'], help="how to preprocess the images" )
|
||||
parser.add_argument("--verbose",action="store_true", help="saving the intermedia output or not" )
|
||||
parser.add_argument("--old_version",action="store_true", help="use the pth other than safetensor version" )
|
||||
|
||||
|
||||
# net structure and parameters
|
||||
parser.add_argument('--net_recon', type=str, default='resnet50', choices=['resnet18', 'resnet34', 'resnet50'], help='useless')
|
||||
parser.add_argument('--init_path', type=str, default=None, help='Useless')
|
||||
parser.add_argument('--use_last_fc',default=False, help='zero initialize the last fc')
|
||||
parser.add_argument('--bfm_folder', type=str, default='./checkpoints/BFM_Fitting/')
|
||||
parser.add_argument('--bfm_model', type=str, default='BFM_model_front.mat', help='bfm model')
|
||||
|
||||
# default renderer parameters
|
||||
parser.add_argument('--focal', type=float, default=1015.)
|
||||
parser.add_argument('--center', type=float, default=112.)
|
||||
parser.add_argument('--camera_d', type=float, default=10.)
|
||||
parser.add_argument('--z_near', type=float, default=5.)
|
||||
parser.add_argument('--z_far', type=float, default=15.)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
if torch.cuda.is_available() and not args.cpu:
|
||||
args.device = "cuda"
|
||||
else:
|
||||
args.device = "cpu"
|
||||
|
||||
main(args)
|
||||
|
||||
@@ -0,0 +1,204 @@
|
||||
# this scripts installs necessary requirements and launches main program in webui.py
|
||||
# borrow from : https://github.com/AUTOMATIC1111/stable-diffusion-webui/blob/master/launch.py
|
||||
import subprocess
|
||||
import os
|
||||
import sys
|
||||
import importlib.util
|
||||
import shlex
|
||||
import platform
|
||||
import json
|
||||
|
||||
python = sys.executable
|
||||
git = os.environ.get('GIT', "git")
|
||||
index_url = os.environ.get('INDEX_URL', "")
|
||||
stored_commit_hash = None
|
||||
skip_install = False
|
||||
dir_repos = "repositories"
|
||||
script_path = os.path.dirname(os.path.dirname(os.path.realpath(__file__)))
|
||||
|
||||
if 'GRADIO_ANALYTICS_ENABLED' not in os.environ:
|
||||
os.environ['GRADIO_ANALYTICS_ENABLED'] = 'False'
|
||||
|
||||
|
||||
def check_python_version():
|
||||
is_windows = platform.system() == "Windows"
|
||||
major = sys.version_info.major
|
||||
minor = sys.version_info.minor
|
||||
micro = sys.version_info.micro
|
||||
|
||||
if is_windows:
|
||||
supported_minors = [10]
|
||||
else:
|
||||
supported_minors = [7, 8, 9, 10, 11]
|
||||
|
||||
if not (major == 3 and minor in supported_minors):
|
||||
|
||||
raise (f"""
|
||||
INCOMPATIBLE PYTHON VERSION
|
||||
This program is tested with 3.10.6 Python, but you have {major}.{minor}.{micro}.
|
||||
If you encounter an error with "RuntimeError: Couldn't install torch." message,
|
||||
or any other error regarding unsuccessful package (library) installation,
|
||||
please downgrade (or upgrade) to the latest version of 3.10 Python
|
||||
and delete current Python and "venv" folder in WebUI's directory.
|
||||
You can download 3.10 Python from here: https://www.python.org/downloads/release/python-3109/
|
||||
{"Alternatively, use a binary release of WebUI: https://github.com/AUTOMATIC1111/stable-diffusion-webui/releases" if is_windows else ""}
|
||||
Use --skip-python-version-check to suppress this warning.
|
||||
""")
|
||||
|
||||
|
||||
def commit_hash():
|
||||
global stored_commit_hash
|
||||
|
||||
if stored_commit_hash is not None:
|
||||
return stored_commit_hash
|
||||
|
||||
try:
|
||||
stored_commit_hash = run(f"{git} rev-parse HEAD").strip()
|
||||
except Exception:
|
||||
stored_commit_hash = "<none>"
|
||||
|
||||
return stored_commit_hash
|
||||
|
||||
|
||||
def run(command, desc=None, errdesc=None, custom_env=None, live=False):
|
||||
if desc is not None:
|
||||
print(desc)
|
||||
|
||||
if live:
|
||||
result = subprocess.run(command, shell=True, env=os.environ if custom_env is None else custom_env)
|
||||
if result.returncode != 0:
|
||||
raise RuntimeError(f"""{errdesc or 'Error running command'}.
|
||||
Command: {command}
|
||||
Error code: {result.returncode}""")
|
||||
|
||||
return ""
|
||||
|
||||
result = subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True, env=os.environ if custom_env is None else custom_env)
|
||||
|
||||
if result.returncode != 0:
|
||||
|
||||
message = f"""{errdesc or 'Error running command'}.
|
||||
Command: {command}
|
||||
Error code: {result.returncode}
|
||||
stdout: {result.stdout.decode(encoding="utf8", errors="ignore") if len(result.stdout)>0 else '<empty>'}
|
||||
stderr: {result.stderr.decode(encoding="utf8", errors="ignore") if len(result.stderr)>0 else '<empty>'}
|
||||
"""
|
||||
raise RuntimeError(message)
|
||||
|
||||
return result.stdout.decode(encoding="utf8", errors="ignore")
|
||||
|
||||
|
||||
def check_run(command):
|
||||
result = subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True)
|
||||
return result.returncode == 0
|
||||
|
||||
|
||||
def is_installed(package):
|
||||
try:
|
||||
spec = importlib.util.find_spec(package)
|
||||
except ModuleNotFoundError:
|
||||
return False
|
||||
|
||||
return spec is not None
|
||||
|
||||
|
||||
def repo_dir(name):
|
||||
return os.path.join(script_path, dir_repos, name)
|
||||
|
||||
|
||||
def run_python(code, desc=None, errdesc=None):
|
||||
return run(f'"{python}" -c "{code}"', desc, errdesc)
|
||||
|
||||
|
||||
def run_pip(args, desc=None):
|
||||
if skip_install:
|
||||
return
|
||||
|
||||
index_url_line = f' --index-url {index_url}' if index_url != '' else ''
|
||||
return run(f'"{python}" -m pip {args} --prefer-binary{index_url_line}', desc=f"Installing {desc}", errdesc=f"Couldn't install {desc}")
|
||||
|
||||
|
||||
def check_run_python(code):
|
||||
return check_run(f'"{python}" -c "{code}"')
|
||||
|
||||
|
||||
def git_clone(url, dir, name, commithash=None):
|
||||
# TODO clone into temporary dir and move if successful
|
||||
|
||||
if os.path.exists(dir):
|
||||
if commithash is None:
|
||||
return
|
||||
|
||||
current_hash = run(f'"{git}" -C "{dir}" rev-parse HEAD', None, f"Couldn't determine {name}'s hash: {commithash}").strip()
|
||||
if current_hash == commithash:
|
||||
return
|
||||
|
||||
run(f'"{git}" -C "{dir}" fetch', f"Fetching updates for {name}...", f"Couldn't fetch {name}")
|
||||
run(f'"{git}" -C "{dir}" checkout {commithash}', f"Checking out commit for {name} with hash: {commithash}...", f"Couldn't checkout commit {commithash} for {name}")
|
||||
return
|
||||
|
||||
run(f'"{git}" clone "{url}" "{dir}"', f"Cloning {name} into {dir}...", f"Couldn't clone {name}")
|
||||
|
||||
if commithash is not None:
|
||||
run(f'"{git}" -C "{dir}" checkout {commithash}', None, "Couldn't checkout {name}'s hash: {commithash}")
|
||||
|
||||
|
||||
def git_pull_recursive(dir):
|
||||
for subdir, _, _ in os.walk(dir):
|
||||
if os.path.exists(os.path.join(subdir, '.git')):
|
||||
try:
|
||||
output = subprocess.check_output([git, '-C', subdir, 'pull', '--autostash'])
|
||||
print(f"Pulled changes for repository in '{subdir}':\n{output.decode('utf-8').strip()}\n")
|
||||
except subprocess.CalledProcessError as e:
|
||||
print(f"Couldn't perform 'git pull' on repository in '{subdir}':\n{e.output.decode('utf-8').strip()}\n")
|
||||
|
||||
|
||||
def run_extension_installer(extension_dir):
|
||||
path_installer = os.path.join(extension_dir, "install.py")
|
||||
if not os.path.isfile(path_installer):
|
||||
return
|
||||
|
||||
try:
|
||||
env = os.environ.copy()
|
||||
env['PYTHONPATH'] = os.path.abspath(".")
|
||||
|
||||
print(run(f'"{python}" "{path_installer}"', errdesc=f"Error running install.py for extension {extension_dir}", custom_env=env))
|
||||
except Exception as e:
|
||||
print(e, file=sys.stderr)
|
||||
|
||||
|
||||
def prepare_environment():
|
||||
global skip_install
|
||||
|
||||
torch_command = os.environ.get('TORCH_COMMAND', "pip install torch==1.12.1+cu113 torchvision==0.13.1+cu113 torchaudio==0.12.1 --extra-index-url https://download.pytorch.org/whl/cu113")
|
||||
|
||||
## check windows
|
||||
if sys.platform != 'win32':
|
||||
requirements_file = os.environ.get('REQS_FILE', "req.txt")
|
||||
else:
|
||||
requirements_file = os.environ.get('REQS_FILE', "requirements.txt")
|
||||
|
||||
commit = commit_hash()
|
||||
|
||||
print(f"Python {sys.version}")
|
||||
print(f"Commit hash: {commit}")
|
||||
|
||||
if not is_installed("torch") or not is_installed("torchvision"):
|
||||
run(f'"{python}" -m {torch_command}', "Installing torch and torchvision", "Couldn't install torch", live=True)
|
||||
|
||||
run_pip(f"install -r \"{requirements_file}\"", "requirements for SadTalker WebUI (may take longer time in first time)")
|
||||
|
||||
if sys.platform != 'win32' and not is_installed('tts'):
|
||||
run_pip(f"install TTS", "install TTS individually in SadTalker, which might not work on windows.")
|
||||
|
||||
|
||||
def start():
|
||||
print(f"Launching SadTalker Web UI")
|
||||
from app_sadtalker import sadtalker_demo
|
||||
demo = sadtalker_demo()
|
||||
demo.queue()
|
||||
demo.launch()
|
||||
|
||||
if __name__ == "__main__":
|
||||
prepare_environment()
|
||||
start()
|
||||
@@ -0,0 +1,192 @@
|
||||
"""run bash scripts/download_models.sh first to prepare the weights file"""
|
||||
import os
|
||||
import shutil
|
||||
from argparse import Namespace
|
||||
from src.utils.preprocess import CropAndExtract
|
||||
from src.test_audio2coeff import Audio2Coeff
|
||||
from src.facerender.animate import AnimateFromCoeff
|
||||
from src.generate_batch import get_data
|
||||
from src.generate_facerender_batch import get_facerender_data
|
||||
from src.utils.init_path import init_path
|
||||
from cog import BasePredictor, Input, Path
|
||||
|
||||
checkpoints = "checkpoints"
|
||||
|
||||
|
||||
class Predictor(BasePredictor):
|
||||
def setup(self):
|
||||
"""Load the model into memory to make running multiple predictions efficient"""
|
||||
device = "cuda"
|
||||
|
||||
|
||||
sadtalker_paths = init_path(checkpoints,os.path.join("src","config"))
|
||||
|
||||
# init model
|
||||
self.preprocess_model = CropAndExtract(sadtalker_paths, device
|
||||
)
|
||||
|
||||
self.audio_to_coeff = Audio2Coeff(
|
||||
sadtalker_paths,
|
||||
device,
|
||||
)
|
||||
|
||||
self.animate_from_coeff = {
|
||||
"full": AnimateFromCoeff(
|
||||
sadtalker_paths,
|
||||
device,
|
||||
),
|
||||
"others": AnimateFromCoeff(
|
||||
sadtalker_paths,
|
||||
device,
|
||||
),
|
||||
}
|
||||
|
||||
def predict(
|
||||
self,
|
||||
source_image: Path = Input(
|
||||
description="Upload the source image, it can be video.mp4 or picture.png",
|
||||
),
|
||||
driven_audio: Path = Input(
|
||||
description="Upload the driven audio, accepts .wav and .mp4 file",
|
||||
),
|
||||
enhancer: str = Input(
|
||||
description="Choose a face enhancer",
|
||||
choices=["gfpgan", "RestoreFormer"],
|
||||
default="gfpgan",
|
||||
),
|
||||
preprocess: str = Input(
|
||||
description="how to preprocess the images",
|
||||
choices=["crop", "resize", "full"],
|
||||
default="full",
|
||||
),
|
||||
ref_eyeblink: Path = Input(
|
||||
description="path to reference video providing eye blinking",
|
||||
default=None,
|
||||
),
|
||||
ref_pose: Path = Input(
|
||||
description="path to reference video providing pose",
|
||||
default=None,
|
||||
),
|
||||
still: bool = Input(
|
||||
description="can crop back to the original videos for the full body aniamtion when preprocess is full",
|
||||
default=True,
|
||||
),
|
||||
) -> Path:
|
||||
"""Run a single prediction on the model"""
|
||||
|
||||
animate_from_coeff = (
|
||||
self.animate_from_coeff["full"]
|
||||
if preprocess == "full"
|
||||
else self.animate_from_coeff["others"]
|
||||
)
|
||||
|
||||
args = load_default()
|
||||
args.pic_path = str(source_image)
|
||||
args.audio_path = str(driven_audio)
|
||||
device = "cuda"
|
||||
args.still = still
|
||||
args.ref_eyeblink = None if ref_eyeblink is None else str(ref_eyeblink)
|
||||
args.ref_pose = None if ref_pose is None else str(ref_pose)
|
||||
|
||||
# crop image and extract 3dmm from image
|
||||
results_dir = "results"
|
||||
if os.path.exists(results_dir):
|
||||
shutil.rmtree(results_dir)
|
||||
os.makedirs(results_dir)
|
||||
first_frame_dir = os.path.join(results_dir, "first_frame_dir")
|
||||
os.makedirs(first_frame_dir)
|
||||
|
||||
print("3DMM Extraction for source image")
|
||||
first_coeff_path, crop_pic_path, crop_info = self.preprocess_model.generate(
|
||||
args.pic_path, first_frame_dir, preprocess, source_image_flag=True
|
||||
)
|
||||
if first_coeff_path is None:
|
||||
print("Can't get the coeffs of the input")
|
||||
return
|
||||
|
||||
if ref_eyeblink is not None:
|
||||
ref_eyeblink_videoname = os.path.splitext(os.path.split(ref_eyeblink)[-1])[
|
||||
0
|
||||
]
|
||||
ref_eyeblink_frame_dir = os.path.join(results_dir, ref_eyeblink_videoname)
|
||||
os.makedirs(ref_eyeblink_frame_dir, exist_ok=True)
|
||||
print("3DMM Extraction for the reference video providing eye blinking")
|
||||
ref_eyeblink_coeff_path, _, _ = self.preprocess_model.generate(
|
||||
ref_eyeblink, ref_eyeblink_frame_dir
|
||||
)
|
||||
else:
|
||||
ref_eyeblink_coeff_path = None
|
||||
|
||||
if ref_pose is not None:
|
||||
if ref_pose == ref_eyeblink:
|
||||
ref_pose_coeff_path = ref_eyeblink_coeff_path
|
||||
else:
|
||||
ref_pose_videoname = os.path.splitext(os.path.split(ref_pose)[-1])[0]
|
||||
ref_pose_frame_dir = os.path.join(results_dir, ref_pose_videoname)
|
||||
os.makedirs(ref_pose_frame_dir, exist_ok=True)
|
||||
print("3DMM Extraction for the reference video providing pose")
|
||||
ref_pose_coeff_path, _, _ = self.preprocess_model.generate(
|
||||
ref_pose, ref_pose_frame_dir
|
||||
)
|
||||
else:
|
||||
ref_pose_coeff_path = None
|
||||
|
||||
# audio2ceoff
|
||||
batch = get_data(
|
||||
first_coeff_path,
|
||||
args.audio_path,
|
||||
device,
|
||||
ref_eyeblink_coeff_path,
|
||||
still=still,
|
||||
)
|
||||
coeff_path = self.audio_to_coeff.generate(
|
||||
batch, results_dir, args.pose_style, ref_pose_coeff_path
|
||||
)
|
||||
# coeff2video
|
||||
print("coeff2video")
|
||||
data = get_facerender_data(
|
||||
coeff_path,
|
||||
crop_pic_path,
|
||||
first_coeff_path,
|
||||
args.audio_path,
|
||||
args.batch_size,
|
||||
args.input_yaw,
|
||||
args.input_pitch,
|
||||
args.input_roll,
|
||||
expression_scale=args.expression_scale,
|
||||
still_mode=still,
|
||||
preprocess=preprocess,
|
||||
)
|
||||
animate_from_coeff.generate(
|
||||
data, results_dir, args.pic_path, crop_info,
|
||||
enhancer=enhancer, background_enhancer=args.background_enhancer,
|
||||
preprocess=preprocess)
|
||||
|
||||
output = "/tmp/out.mp4"
|
||||
mp4_path = os.path.join(results_dir, [f for f in os.listdir(results_dir) if "enhanced.mp4" in f][0])
|
||||
shutil.copy(mp4_path, output)
|
||||
|
||||
return Path(output)
|
||||
|
||||
|
||||
def load_default():
|
||||
return Namespace(
|
||||
pose_style=0,
|
||||
batch_size=2,
|
||||
expression_scale=1.0,
|
||||
input_yaw=None,
|
||||
input_pitch=None,
|
||||
input_roll=None,
|
||||
background_enhancer=None,
|
||||
face3dvis=False,
|
||||
net_recon="resnet50",
|
||||
init_path=None,
|
||||
use_last_fc=False,
|
||||
bfm_folder="./src/config/",
|
||||
bfm_model="BFM_model_front.mat",
|
||||
focal=1015.0,
|
||||
center=112.0,
|
||||
camera_d=10.0,
|
||||
z_near=5.0,
|
||||
z_far=15.0,
|
||||
)
|
||||
@@ -0,0 +1,213 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"attachments": {},
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "M74Gs_TjYl_B"
|
||||
},
|
||||
"source": [
|
||||
"[](https://colab.research.google.com/github/Winfredy/SadTalker/blob/main/quick_demo.ipynb)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"attachments": {},
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "view-in-github"
|
||||
},
|
||||
"source": [
|
||||
"### SadTalker：Learning Realistic 3D Motion Coefficients for Stylized Audio-Driven Single Image Talking Face Animation \n",
|
||||
"\n",
|
||||
"[arxiv](https://arxiv.org/abs/2211.12194) | [project](https://sadtalker.github.io) | [Github](https://github.com/Winfredy/SadTalker)\n",
|
||||
"\n",
|
||||
"Wenxuan Zhang, Xiaodong Cun, Xuan Wang, Yong Zhang, Xi Shen, Yu Guo, Ying Shan, Fei Wang.\n",
|
||||
"\n",
|
||||
"Xi'an Jiaotong University, Tencent AI Lab, Ant Group\n",
|
||||
"\n",
|
||||
"CVPR 2023\n",
|
||||
"\n",
|
||||
"TL;DR: A realistic and stylized talking head video generation method from a single image and audio\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"attachments": {},
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "kA89DV-sKS4i"
|
||||
},
|
||||
"source": [
|
||||
"Installation (around 5 mins)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "qJ4CplXsYl_E"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"### make sure that CUDA is available in Edit -> Nootbook settings -> GPU\n",
|
||||
"!nvidia-smi --query-gpu=name,memory.total,memory.free --format=csv,noheader"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "Mdq6j4E5KQAR"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"!update-alternatives --install /usr/local/bin/python3 python3 /usr/bin/python3.8 2 \n",
|
||||
"!update-alternatives --install /usr/local/bin/python3 python3 /usr/bin/python3.9 1 \n",
|
||||
"!python --version \n",
|
||||
"!apt-get update\n",
|
||||
"!apt install software-properties-common\n",
|
||||
"!sudo dpkg --remove --force-remove-reinstreq python3-pip python3-setuptools python3-wheel\n",
|
||||
"!apt-get install python3-pip\n",
|
||||
"\n",
|
||||
"print('Git clone project and install requirements...')\n",
|
||||
"!git clone https://github.com/Winfredy/SadTalker &> /dev/null\n",
|
||||
"%cd SadTalker \n",
|
||||
"!export PYTHONPATH=/content/SadTalker:$PYTHONPATH \n",
|
||||
"!python3.8 -m pip install torch==1.12.1+cu113 torchvision==0.13.1+cu113 torchaudio==0.12.1 --extra-index-url https://download.pytorch.org/whl/cu113\n",
|
||||
"!apt update\n",
|
||||
"!apt install ffmpeg &> /dev/null \n",
|
||||
"!python3.8 -m pip install -r requirements.txt"
|
||||
]
|
||||
},
|
||||
{
|
||||
"attachments": {},
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "DddcKB_nKsnk"
|
||||
},
|
||||
"source": [
|
||||
"Download models (1 mins)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "eDw3_UN8K2xa"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"print('Download pre-trained models...')\n",
|
||||
"!rm -rf checkpoints\n",
|
||||
"!bash scripts/download_models.sh"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "kK7DYeo7Yl_H"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# borrow from makeittalk\n",
|
||||
"import ipywidgets as widgets\n",
|
||||
"import glob\n",
|
||||
"import matplotlib.pyplot as plt\n",
|
||||
"print(\"Choose the image name to animate: (saved in folder 'examples/')\")\n",
|
||||
"img_list = glob.glob1('examples/source_image', '*.png')\n",
|
||||
"img_list.sort()\n",
|
||||
"img_list = [item.split('.')[0] for item in img_list]\n",
|
||||
"default_head_name = widgets.Dropdown(options=img_list, value='full3')\n",
|
||||
"def on_change(change):\n",
|
||||
" if change['type'] == 'change' and change['name'] == 'value':\n",
|
||||
" plt.imshow(plt.imread('examples/source_image/{}.png'.format(default_head_name.value)))\n",
|
||||
" plt.axis('off')\n",
|
||||
" plt.show()\n",
|
||||
"default_head_name.observe(on_change)\n",
|
||||
"display(default_head_name)\n",
|
||||
"plt.imshow(plt.imread('examples/source_image/{}.png'.format(default_head_name.value)))\n",
|
||||
"plt.axis('off')\n",
|
||||
"plt.show()"
|
||||
]
|
||||
},
|
||||
{
|
||||
"attachments": {},
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "-khNZcnGK4UK"
|
||||
},
|
||||
"source": [
|
||||
"Animation"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "ToBlDusjK5sS"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# selected audio from exmaple/driven_audio\n",
|
||||
"img = 'examples/source_image/{}.png'.format(default_head_name.value)\n",
|
||||
"print(img)\n",
|
||||
"!python3.8 inference.py --driven_audio ./examples/driven_audio/RD_Radio31_000.wav \\\n",
|
||||
" --source_image {img} \\\n",
|
||||
" --result_dir ./results --still --preprocess full --enhancer gfpgan"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "fAjwGmKKYl_I"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# visualize code from makeittalk\n",
|
||||
"from IPython.display import HTML\n",
|
||||
"from base64 import b64encode\n",
|
||||
"import os, sys\n",
|
||||
"\n",
|
||||
"# get the last from results\n",
|
||||
"\n",
|
||||
"results = sorted(os.listdir('./results/'))\n",
|
||||
"\n",
|
||||
"mp4_name = glob.glob('./results/*.mp4')[0]\n",
|
||||
"\n",
|
||||
"mp4 = open('{}'.format(mp4_name),'rb').read()\n",
|
||||
"data_url = \"data:video/mp4;base64,\" + b64encode(mp4).decode()\n",
|
||||
"\n",
|
||||
"print('Display animation: {}'.format(mp4_name), file=sys.stderr)\n",
|
||||
"display(HTML(\"\"\"\n",
|
||||
" <video width=256 controls>\n",
|
||||
" <source src=\"%s\" type=\"video/mp4\">\n",
|
||||
" </video>\n",
|
||||
" \"\"\" % data_url))\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"accelerator": "GPU",
|
||||
"colab": {
|
||||
"provenance": []
|
||||
},
|
||||
"gpuClass": "standard",
|
||||
"kernelspec": {
|
||||
"display_name": "base",
|
||||
"language": "python",
|
||||
"name": "python3"
|
||||
},
|
||||
"language_info": {
|
||||
"name": "python",
|
||||
"version": "3.9.7"
|
||||
},
|
||||
"vscode": {
|
||||
"interpreter": {
|
||||
"hash": "db5031b3636a3f037ea48eb287fd3d023feb9033aefc2a9652a92e470fb0851b"
|
||||
}
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 0
|
||||
}
|
||||
@@ -0,0 +1,22 @@
|
||||
llvmlite==0.38.1
|
||||
numpy==1.21.6
|
||||
face_alignment==1.3.5
|
||||
imageio==2.19.3
|
||||
imageio-ffmpeg==0.4.7
|
||||
librosa==0.10.0.post2
|
||||
numba==0.55.1
|
||||
resampy==0.3.1
|
||||
pydub==0.25.1
|
||||
scipy==1.10.1
|
||||
kornia==0.6.8
|
||||
tqdm
|
||||
yacs==0.1.8
|
||||
pyyaml
|
||||
joblib==1.1.0
|
||||
scikit-image==0.19.3
|
||||
basicsr==1.4.2
|
||||
facexlib==0.3.0
|
||||
gradio
|
||||
gfpgan
|
||||
av
|
||||
safetensors
|
||||
@@ -0,0 +1,21 @@
|
||||
numpy==1.23.4
|
||||
face_alignment==1.3.5
|
||||
imageio==2.19.3
|
||||
imageio-ffmpeg==0.4.7
|
||||
librosa==0.9.2 #
|
||||
numba
|
||||
resampy==0.3.1
|
||||
pydub==0.25.1
|
||||
scipy==1.10.1
|
||||
kornia==0.6.8
|
||||
tqdm
|
||||
yacs==0.1.8
|
||||
pyyaml
|
||||
joblib==1.1.0
|
||||
scikit-image==0.19.3
|
||||
basicsr==1.4.2
|
||||
facexlib==0.3.0
|
||||
gradio
|
||||
gfpgan
|
||||
av
|
||||
safetensors
|
||||
@@ -0,0 +1,21 @@
|
||||
numpy==1.23.4
|
||||
face_alignment==1.3.5
|
||||
imageio==2.19.3
|
||||
imageio-ffmpeg==0.4.7
|
||||
librosa==0.9.2 #
|
||||
numba
|
||||
resampy==0.3.1
|
||||
pydub==0.25.1
|
||||
scipy==1.5.3
|
||||
kornia==0.6.8
|
||||
tqdm
|
||||
yacs==0.1.8
|
||||
pyyaml
|
||||
joblib==1.1.0
|
||||
scikit-image==0.19.3
|
||||
basicsr==1.4.2
|
||||
facexlib==0.3.0
|
||||
trimesh==3.9.20
|
||||
gradio
|
||||
gfpgan
|
||||
safetensors
|
||||
@@ -0,0 +1,32 @@
|
||||
mkdir ./checkpoints
|
||||
|
||||
# lagency download link
|
||||
# wget -nc https://github.com/Winfredy/SadTalker/releases/download/v0.0.2/auido2exp_00300-model.pth -O ./checkpoints/auido2exp_00300-model.pth
|
||||
# wget -nc https://github.com/Winfredy/SadTalker/releases/download/v0.0.2/auido2pose_00140-model.pth -O ./checkpoints/auido2pose_00140-model.pth
|
||||
# wget -nc https://github.com/Winfredy/SadTalker/releases/download/v0.0.2/epoch_20.pth -O ./checkpoints/epoch_20.pth
|
||||
# wget -nc https://github.com/Winfredy/SadTalker/releases/download/v0.0.2/facevid2vid_00189-model.pth.tar -O ./checkpoints/facevid2vid_00189-model.pth.tar
|
||||
# wget -nc https://github.com/Winfredy/SadTalker/releases/download/v0.0.2/shape_predictor_68_face_landmarks.dat -O ./checkpoints/shape_predictor_68_face_landmarks.dat
|
||||
# wget -nc https://github.com/Winfredy/SadTalker/releases/download/v0.0.2/wav2lip.pth -O ./checkpoints/wav2lip.pth
|
||||
# wget -nc https://github.com/Winfredy/SadTalker/releases/download/v0.0.2/mapping_00229-model.pth.tar -O ./checkpoints/mapping_00229-model.pth.tar
|
||||
# wget -nc https://github.com/Winfredy/SadTalker/releases/download/v0.0.2/mapping_00109-model.pth.tar -O ./checkpoints/mapping_00109-model.pth.tar
|
||||
# wget -nc https://github.com/Winfredy/SadTalker/releases/download/v0.0.2/hub.zip -O ./checkpoints/hub.zip
|
||||
# unzip -n ./checkpoints/hub.zip -d ./checkpoints/
|
||||
|
||||
|
||||
#### download the new links.
|
||||
wget -nc https://github.com/OpenTalker/SadTalker/releases/download/v0.0.2-rc/mapping_00109-model.pth.tar -O ./checkpoints/mapping_00109-model.pth.tar
|
||||
wget -nc https://github.com/OpenTalker/SadTalker/releases/download/v0.0.2-rc/mapping_00229-model.pth.tar -O ./checkpoints/mapping_00229-model.pth.tar
|
||||
wget -nc https://github.com/OpenTalker/SadTalker/releases/download/v0.0.2-rc/SadTalker_V0.0.2_256.safetensors -O ./checkpoints/SadTalker_V0.0.2_256.safetensors
|
||||
wget -nc https://github.com/OpenTalker/SadTalker/releases/download/v0.0.2-rc/SadTalker_V0.0.2_512.safetensors -O ./checkpoints/SadTalker_V0.0.2_512.safetensors
|
||||
|
||||
|
||||
# wget -nc https://github.com/Winfredy/SadTalker/releases/download/v0.0.2/BFM_Fitting.zip -O ./checkpoints/BFM_Fitting.zip
|
||||
# unzip -n ./checkpoints/BFM_Fitting.zip -d ./checkpoints/
|
||||
|
||||
### enhancer
|
||||
mkdir -p ./gfpgan/weights
|
||||
wget -nc https://github.com/xinntao/facexlib/releases/download/v0.1.0/alignment_WFLW_4HG.pth -O ./gfpgan/weights/alignment_WFLW_4HG.pth
|
||||
wget -nc https://github.com/xinntao/facexlib/releases/download/v0.1.0/detection_Resnet50_Final.pth -O ./gfpgan/weights/detection_Resnet50_Final.pth
|
||||
wget -nc https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth -O ./gfpgan/weights/GFPGANv1.4.pth
|
||||
wget -nc https://github.com/xinntao/facexlib/releases/download/v0.2.2/parsing_parsenet.pth -O ./gfpgan/weights/parsing_parsenet.pth
|
||||
|
||||
@@ -0,0 +1,189 @@
|
||||
import os, sys
|
||||
from pathlib import Path
|
||||
import tempfile
|
||||
import gradio as gr
|
||||
from modules.call_queue import wrap_gradio_gpu_call, wrap_queued_call
|
||||
from modules.shared import opts, OptionInfo
|
||||
from modules import shared, paths, script_callbacks
|
||||
import launch
|
||||
import glob
|
||||
from huggingface_hub import snapshot_download
|
||||
|
||||
|
||||
|
||||
def check_all_files_safetensor(current_dir):
|
||||
kv = {
|
||||
"SadTalker_V0.0.2_256.safetensors": "sadtalker-256",
|
||||
"SadTalker_V0.0.2_512.safetensors": "sadtalker-512",
|
||||
"mapping_00109-model.pth.tar" : "mapping-109" ,
|
||||
"mapping_00229-model.pth.tar" : "mapping-229" ,
|
||||
}
|
||||
|
||||
if not os.path.isdir(current_dir):
|
||||
return False
|
||||
|
||||
dirs = os.listdir(current_dir)
|
||||
|
||||
for f in dirs:
|
||||
if f in kv.keys():
|
||||
del kv[f]
|
||||
|
||||
return len(kv.keys()) == 0
|
||||
|
||||
def check_all_files(current_dir):
|
||||
kv = {
|
||||
"auido2exp_00300-model.pth": "audio2exp",
|
||||
"auido2pose_00140-model.pth": "audio2pose",
|
||||
"epoch_20.pth": "face_recon",
|
||||
"facevid2vid_00189-model.pth.tar": "face-render",
|
||||
"mapping_00109-model.pth.tar" : "mapping-109" ,
|
||||
"mapping_00229-model.pth.tar" : "mapping-229" ,
|
||||
"wav2lip.pth": "wav2lip",
|
||||
"shape_predictor_68_face_landmarks.dat": "dlib",
|
||||
}
|
||||
|
||||
if not os.path.isdir(current_dir):
|
||||
return False
|
||||
|
||||
dirs = os.listdir(current_dir)
|
||||
|
||||
for f in dirs:
|
||||
if f in kv.keys():
|
||||
del kv[f]
|
||||
|
||||
return len(kv.keys()) == 0
|
||||
|
||||
|
||||
|
||||
def download_model(local_dir='./checkpoints'):
|
||||
REPO_ID = 'vinthony/SadTalker'
|
||||
snapshot_download(repo_id=REPO_ID, local_dir=local_dir, local_dir_use_symlinks=False)
|
||||
|
||||
def get_source_image(image):
|
||||
return image
|
||||
|
||||
def get_img_from_txt2img(x):
|
||||
talker_path = Path(paths.script_path) / "outputs"
|
||||
imgs_from_txt_dir = str(talker_path / "txt2img-images/")
|
||||
imgs = glob.glob(imgs_from_txt_dir+'/*/*.png')
|
||||
imgs.sort(key=lambda x:os.path.getmtime(os.path.join(imgs_from_txt_dir, x)))
|
||||
img_from_txt_path = os.path.join(imgs_from_txt_dir, imgs[-1])
|
||||
return img_from_txt_path, img_from_txt_path
|
||||
|
||||
def get_img_from_img2img(x):
|
||||
talker_path = Path(paths.script_path) / "outputs"
|
||||
imgs_from_img_dir = str(talker_path / "img2img-images/")
|
||||
imgs = glob.glob(imgs_from_img_dir+'/*/*.png')
|
||||
imgs.sort(key=lambda x:os.path.getmtime(os.path.join(imgs_from_img_dir, x)))
|
||||
img_from_img_path = os.path.join(imgs_from_img_dir, imgs[-1])
|
||||
return img_from_img_path, img_from_img_path
|
||||
|
||||
def get_default_checkpoint_path():
|
||||
# check the path of models/checkpoints and extensions/
|
||||
checkpoint_path = Path(paths.script_path) / "models"/ "SadTalker"
|
||||
extension_checkpoint_path = Path(paths.script_path) / "extensions"/ "SadTalker" / "checkpoints"
|
||||
|
||||
if check_all_files_safetensor(checkpoint_path):
|
||||
# print('founding sadtalker checkpoint in ' + str(checkpoint_path))
|
||||
return checkpoint_path
|
||||
|
||||
if check_all_files_safetensor(extension_checkpoint_path):
|
||||
# print('founding sadtalker checkpoint in ' + str(extension_checkpoint_path))
|
||||
return extension_checkpoint_path
|
||||
|
||||
if check_all_files(checkpoint_path):
|
||||
# print('founding sadtalker checkpoint in ' + str(checkpoint_path))
|
||||
return checkpoint_path
|
||||
|
||||
if check_all_files(extension_checkpoint_path):
|
||||
# print('founding sadtalker checkpoint in ' + str(extension_checkpoint_path))
|
||||
return extension_checkpoint_path
|
||||
|
||||
return None
|
||||
|
||||
|
||||
|
||||
def install():
|
||||
|
||||
kv = {
|
||||
"face_alignment": "face-alignment==1.3.5",
|
||||
"imageio": "imageio==2.19.3",
|
||||
"imageio_ffmpeg": "imageio-ffmpeg==0.4.7",
|
||||
"librosa":"librosa==0.8.0",
|
||||
"pydub":"pydub==0.25.1",
|
||||
"scipy":"scipy==1.8.1",
|
||||
"tqdm": "tqdm",
|
||||
"yacs":"yacs==0.1.8",
|
||||
"yaml": "pyyaml",
|
||||
"av":"av",
|
||||
"gfpgan": "gfpgan",
|
||||
}
|
||||
|
||||
# # dlib is not necessary currently
|
||||
# if 'darwin' in sys.platform:
|
||||
# kv['dlib'] = "dlib"
|
||||
# else:
|
||||
# kv['dlib'] = 'dlib-bin'
|
||||
|
||||
# #### we need to have a newer version of imageio for our method.
|
||||
# launch.run_pip("install imageio==2.19.3", "requirements for SadTalker")
|
||||
|
||||
for k,v in kv.items():
|
||||
if not launch.is_installed(k):
|
||||
print(k, launch.is_installed(k))
|
||||
launch.run_pip("install "+ v, "requirements for SadTalker")
|
||||
|
||||
if os.getenv('SADTALKER_CHECKPOINTS'):
|
||||
print('load Sadtalker Checkpoints from '+ os.getenv('SADTALKER_CHECKPOINTS'))
|
||||
|
||||
elif get_default_checkpoint_path() is not None:
|
||||
os.environ['SADTALKER_CHECKPOINTS'] = str(get_default_checkpoint_path())
|
||||
else:
|
||||
|
||||
print(
|
||||
""""
|
||||
SadTalker will not support download all the files from hugging face, which will take a long time.
|
||||
|
||||
please manually set the SADTALKER_CHECKPOINTS in `webui_user.bat`(windows) or `webui_user.sh`(linux)
|
||||
"""
|
||||
)
|
||||
|
||||
# python = sys.executable
|
||||
|
||||
# launch.run(f'"{python}" -m pip uninstall -y huggingface_hub', live=True)
|
||||
# launch.run(f'"{python}" -m pip install --upgrade git+https://github.com/huggingface/huggingface_hub@main', live=True)
|
||||
# ### run the scripts to downlod models to correct localtion.
|
||||
# # print('download models for SadTalker')
|
||||
# # launch.run("cd " + paths.script_path+"/extensions/SadTalker && bash ./scripts/download_models.sh", live=True)
|
||||
# # print('SadTalker is successfully installed!')
|
||||
# download_model(paths.script_path+'/extensions/SadTalker/checkpoints')
|
||||
|
||||
|
||||
def on_ui_tabs():
|
||||
install()
|
||||
|
||||
sys.path.extend([paths.script_path+'/extensions/SadTalker'])
|
||||
|
||||
repo_dir = paths.script_path+'/extensions/SadTalker/'
|
||||
|
||||
result_dir = opts.sadtalker_result_dir
|
||||
os.makedirs(result_dir, exist_ok=True)
|
||||
|
||||
from app_sadtalker import sadtalker_demo
|
||||
|
||||
if os.getenv('SADTALKER_CHECKPOINTS'):
|
||||
checkpoint_path = os.getenv('SADTALKER_CHECKPOINTS')
|
||||
else:
|
||||
checkpoint_path = repo_dir+'checkpoints/'
|
||||
|
||||
audio_to_video = sadtalker_demo(checkpoint_path=checkpoint_path, config_path=repo_dir+'src/config', warpfn = wrap_queued_call)
|
||||
|
||||
return [(audio_to_video, "SadTalker", "extension")]
|
||||
|
||||
def on_ui_settings():
|
||||
talker_path = Path(paths.script_path) / "outputs"
|
||||
section = ('extension', "SadTalker")
|
||||
opts.add_option("sadtalker_result_dir", OptionInfo(str(talker_path / "SadTalker/"), "Path to save results of sadtalker", section=section))
|
||||
|
||||
script_callbacks.on_ui_settings(on_ui_settings)
|
||||
script_callbacks.on_ui_tabs(on_ui_tabs)
|
||||
@@ -0,0 +1,21 @@
|
||||
# ### some test command before commit.
|
||||
# python inference.py --preprocess crop --size 256
|
||||
# python inference.py --preprocess crop --size 512
|
||||
|
||||
# python inference.py --preprocess extcrop --size 256
|
||||
# python inference.py --preprocess extcrop --size 512
|
||||
|
||||
# python inference.py --preprocess resize --size 256
|
||||
# python inference.py --preprocess resize --size 512
|
||||
|
||||
# python inference.py --preprocess full --size 256
|
||||
# python inference.py --preprocess full --size 512
|
||||
|
||||
# python inference.py --preprocess extfull --size 256
|
||||
# python inference.py --preprocess extfull --size 512
|
||||
|
||||
python inference.py --preprocess full --size 256 --enhancer gfpgan
|
||||
python inference.py --preprocess full --size 512 --enhancer gfpgan
|
||||
|
||||
python inference.py --preprocess full --size 256 --enhancer gfpgan --still
|
||||
python inference.py --preprocess full --size 512 --enhancer gfpgan --still
|
||||
@@ -0,0 +1,41 @@
|
||||
from tqdm import tqdm
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
|
||||
class Audio2Exp(nn.Module):
|
||||
def __init__(self, netG, cfg, device, prepare_training_loss=False):
|
||||
super(Audio2Exp, self).__init__()
|
||||
self.cfg = cfg
|
||||
self.device = device
|
||||
self.netG = netG.to(device)
|
||||
|
||||
def test(self, batch):
|
||||
|
||||
mel_input = batch['indiv_mels'] # bs T 1 80 16
|
||||
bs = mel_input.shape[0]
|
||||
T = mel_input.shape[1]
|
||||
|
||||
exp_coeff_pred = []
|
||||
|
||||
for i in tqdm(range(0, T, 10),'audio2exp:'): # every 10 frames
|
||||
|
||||
current_mel_input = mel_input[:,i:i+10]
|
||||
|
||||
#ref = batch['ref'][:, :, :64].repeat((1,current_mel_input.shape[1],1)) #bs T 64
|
||||
ref = batch['ref'][:, :, :64][:, i:i+10]
|
||||
ratio = batch['ratio_gt'][:, i:i+10] #bs T
|
||||
|
||||
audiox = current_mel_input.view(-1, 1, 80, 16) # bs*T 1 80 16
|
||||
|
||||
curr_exp_coeff_pred = self.netG(audiox, ref, ratio) # bs T 64
|
||||
|
||||
exp_coeff_pred += [curr_exp_coeff_pred]
|
||||
|
||||
# BS x T x 64
|
||||
results_dict = {
|
||||
'exp_coeff_pred': torch.cat(exp_coeff_pred, axis=1)
|
||||
}
|
||||
return results_dict
|
||||
|
||||
|
||||
@@ -0,0 +1,74 @@
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch import nn
|
||||
|
||||
class Conv2d(nn.Module):
|
||||
def __init__(self, cin, cout, kernel_size, stride, padding, residual=False, use_act = True, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
self.conv_block = nn.Sequential(
|
||||
nn.Conv2d(cin, cout, kernel_size, stride, padding),
|
||||
nn.BatchNorm2d(cout)
|
||||
)
|
||||
self.act = nn.ReLU()
|
||||
self.residual = residual
|
||||
self.use_act = use_act
|
||||
|
||||
def forward(self, x):
|
||||
out = self.conv_block(x)
|
||||
if self.residual:
|
||||
out += x
|
||||
|
||||
if self.use_act:
|
||||
return self.act(out)
|
||||
else:
|
||||
return out
|
||||
|
||||
class SimpleWrapperV2(nn.Module):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.audio_encoder = nn.Sequential(
|
||||
Conv2d(1, 32, kernel_size=3, stride=1, padding=1),
|
||||
Conv2d(32, 32, kernel_size=3, stride=1, padding=1, residual=True),
|
||||
Conv2d(32, 32, kernel_size=3, stride=1, padding=1, residual=True),
|
||||
|
||||
Conv2d(32, 64, kernel_size=3, stride=(3, 1), padding=1),
|
||||
Conv2d(64, 64, kernel_size=3, stride=1, padding=1, residual=True),
|
||||
Conv2d(64, 64, kernel_size=3, stride=1, padding=1, residual=True),
|
||||
|
||||
Conv2d(64, 128, kernel_size=3, stride=3, padding=1),
|
||||
Conv2d(128, 128, kernel_size=3, stride=1, padding=1, residual=True),
|
||||
Conv2d(128, 128, kernel_size=3, stride=1, padding=1, residual=True),
|
||||
|
||||
Conv2d(128, 256, kernel_size=3, stride=(3, 2), padding=1),
|
||||
Conv2d(256, 256, kernel_size=3, stride=1, padding=1, residual=True),
|
||||
|
||||
Conv2d(256, 512, kernel_size=3, stride=1, padding=0),
|
||||
Conv2d(512, 512, kernel_size=1, stride=1, padding=0),
|
||||
)
|
||||
|
||||
#### load the pre-trained audio_encoder
|
||||
#self.audio_encoder = self.audio_encoder.to(device)
|
||||
'''
|
||||
wav2lip_state_dict = torch.load('/apdcephfs_cq2/share_1290939/wenxuazhang/checkpoints/wav2lip.pth')['state_dict']
|
||||
state_dict = self.audio_encoder.state_dict()
|
||||
|
||||
for k,v in wav2lip_state_dict.items():
|
||||
if 'audio_encoder' in k:
|
||||
print('init:', k)
|
||||
state_dict[k.replace('module.audio_encoder.', '')] = v
|
||||
self.audio_encoder.load_state_dict(state_dict)
|
||||
'''
|
||||
|
||||
self.mapping1 = nn.Linear(512+64+1, 64)
|
||||
#self.mapping2 = nn.Linear(30, 64)
|
||||
#nn.init.constant_(self.mapping1.weight, 0.)
|
||||
nn.init.constant_(self.mapping1.bias, 0.)
|
||||
|
||||
def forward(self, x, ref, ratio):
|
||||
x = self.audio_encoder(x).view(x.size(0), -1)
|
||||
ref_reshape = ref.reshape(x.size(0), -1)
|
||||
ratio = ratio.reshape(x.size(0), -1)
|
||||
|
||||
y = self.mapping1(torch.cat([x, ref_reshape, ratio], dim=1))
|
||||
out = y.reshape(ref.shape[0], ref.shape[1], -1) #+ ref # resudial
|
||||
return out
|
||||
@@ -0,0 +1,94 @@
|
||||
import torch
|
||||
from torch import nn
|
||||
from src.audio2pose_models.cvae import CVAE
|
||||
from src.audio2pose_models.discriminator import PoseSequenceDiscriminator
|
||||
from src.audio2pose_models.audio_encoder import AudioEncoder
|
||||
|
||||
class Audio2Pose(nn.Module):
|
||||
def __init__(self, cfg, wav2lip_checkpoint, device='cuda'):
|
||||
super().__init__()
|
||||
self.cfg = cfg
|
||||
self.seq_len = cfg.MODEL.CVAE.SEQ_LEN
|
||||
self.latent_dim = cfg.MODEL.CVAE.LATENT_SIZE
|
||||
self.device = device
|
||||
|
||||
self.audio_encoder = AudioEncoder(wav2lip_checkpoint, device)
|
||||
self.audio_encoder.eval()
|
||||
for param in self.audio_encoder.parameters():
|
||||
param.requires_grad = False
|
||||
|
||||
self.netG = CVAE(cfg)
|
||||
self.netD_motion = PoseSequenceDiscriminator(cfg)
|
||||
|
||||
|
||||
def forward(self, x):
|
||||
|
||||
batch = {}
|
||||
coeff_gt = x['gt'].cuda().squeeze(0) #bs frame_len+1 73
|
||||
batch['pose_motion_gt'] = coeff_gt[:, 1:, 64:70] - coeff_gt[:, :1, 64:70] #bs frame_len 6
|
||||
batch['ref'] = coeff_gt[:, 0, 64:70] #bs 6
|
||||
batch['class'] = x['class'].squeeze(0).cuda() # bs
|
||||
indiv_mels= x['indiv_mels'].cuda().squeeze(0) # bs seq_len+1 80 16
|
||||
|
||||
# forward
|
||||
audio_emb_list = []
|
||||
audio_emb = self.audio_encoder(indiv_mels[:, 1:, :, :].unsqueeze(2)) #bs seq_len 512
|
||||
batch['audio_emb'] = audio_emb
|
||||
batch = self.netG(batch)
|
||||
|
||||
pose_motion_pred = batch['pose_motion_pred'] # bs frame_len 6
|
||||
pose_gt = coeff_gt[:, 1:, 64:70].clone() # bs frame_len 6
|
||||
pose_pred = coeff_gt[:, :1, 64:70] + pose_motion_pred # bs frame_len 6
|
||||
|
||||
batch['pose_pred'] = pose_pred
|
||||
batch['pose_gt'] = pose_gt
|
||||
|
||||
return batch
|
||||
|
||||
def test(self, x):
|
||||
|
||||
batch = {}
|
||||
ref = x['ref'] #bs 1 70
|
||||
batch['ref'] = x['ref'][:,0,-6:]
|
||||
batch['class'] = x['class']
|
||||
bs = ref.shape[0]
|
||||
|
||||
indiv_mels= x['indiv_mels'] # bs T 1 80 16
|
||||
indiv_mels_use = indiv_mels[:, 1:] # we regard the ref as the first frame
|
||||
num_frames = x['num_frames']
|
||||
num_frames = int(num_frames) - 1
|
||||
|
||||
#
|
||||
div = num_frames//self.seq_len
|
||||
re = num_frames%self.seq_len
|
||||
audio_emb_list = []
|
||||
pose_motion_pred_list = [torch.zeros(batch['ref'].unsqueeze(1).shape, dtype=batch['ref'].dtype,
|
||||
device=batch['ref'].device)]
|
||||
|
||||
for i in range(div):
|
||||
z = torch.randn(bs, self.latent_dim).to(ref.device)
|
||||
batch['z'] = z
|
||||
audio_emb = self.audio_encoder(indiv_mels_use[:, i*self.seq_len:(i+1)*self.seq_len,:,:,:]) #bs seq_len 512
|
||||
batch['audio_emb'] = audio_emb
|
||||
batch = self.netG.test(batch)
|
||||
pose_motion_pred_list.append(batch['pose_motion_pred']) #list of bs seq_len 6
|
||||
|
||||
if re != 0:
|
||||
z = torch.randn(bs, self.latent_dim).to(ref.device)
|
||||
batch['z'] = z
|
||||
audio_emb = self.audio_encoder(indiv_mels_use[:, -1*self.seq_len:,:,:,:]) #bs seq_len 512
|
||||
if audio_emb.shape[1] != self.seq_len:
|
||||
pad_dim = self.seq_len-audio_emb.shape[1]
|
||||
pad_audio_emb = audio_emb[:, :1].repeat(1, pad_dim, 1)
|
||||
audio_emb = torch.cat([pad_audio_emb, audio_emb], 1)
|
||||
batch['audio_emb'] = audio_emb
|
||||
batch = self.netG.test(batch)
|
||||
pose_motion_pred_list.append(batch['pose_motion_pred'][:,-1*re:,:])
|
||||
|
||||
pose_motion_pred = torch.cat(pose_motion_pred_list, dim = 1)
|
||||
batch['pose_motion_pred'] = pose_motion_pred
|
||||
|
||||
pose_pred = ref[:, :1, -6:] + pose_motion_pred # bs T 6
|
||||
|
||||
batch['pose_pred'] = pose_pred
|
||||
return batch
|
||||
@@ -0,0 +1,64 @@
|
||||
import torch
|
||||
from torch import nn
|
||||
from torch.nn import functional as F
|
||||
|
||||
class Conv2d(nn.Module):
|
||||
def __init__(self, cin, cout, kernel_size, stride, padding, residual=False, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
self.conv_block = nn.Sequential(
|
||||
nn.Conv2d(cin, cout, kernel_size, stride, padding),
|
||||
nn.BatchNorm2d(cout)
|
||||
)
|
||||
self.act = nn.ReLU()
|
||||
self.residual = residual
|
||||
|
||||
def forward(self, x):
|
||||
out = self.conv_block(x)
|
||||
if self.residual:
|
||||
out += x
|
||||
return self.act(out)
|
||||
|
||||
class AudioEncoder(nn.Module):
|
||||
def __init__(self, wav2lip_checkpoint, device):
|
||||
super(AudioEncoder, self).__init__()
|
||||
|
||||
self.audio_encoder = nn.Sequential(
|
||||
Conv2d(1, 32, kernel_size=3, stride=1, padding=1),
|
||||
Conv2d(32, 32, kernel_size=3, stride=1, padding=1, residual=True),
|
||||
Conv2d(32, 32, kernel_size=3, stride=1, padding=1, residual=True),
|
||||
|
||||
Conv2d(32, 64, kernel_size=3, stride=(3, 1), padding=1),
|
||||
Conv2d(64, 64, kernel_size=3, stride=1, padding=1, residual=True),
|
||||
Conv2d(64, 64, kernel_size=3, stride=1, padding=1, residual=True),
|
||||
|
||||
Conv2d(64, 128, kernel_size=3, stride=3, padding=1),
|
||||
Conv2d(128, 128, kernel_size=3, stride=1, padding=1, residual=True),
|
||||
Conv2d(128, 128, kernel_size=3, stride=1, padding=1, residual=True),
|
||||
|
||||
Conv2d(128, 256, kernel_size=3, stride=(3, 2), padding=1),
|
||||
Conv2d(256, 256, kernel_size=3, stride=1, padding=1, residual=True),
|
||||
|
||||
Conv2d(256, 512, kernel_size=3, stride=1, padding=0),
|
||||
Conv2d(512, 512, kernel_size=1, stride=1, padding=0),)
|
||||
|
||||
#### load the pre-trained audio_encoder, we do not need to load wav2lip model here.
|
||||
# wav2lip_state_dict = torch.load(wav2lip_checkpoint, map_location=torch.device(device))['state_dict']
|
||||
# state_dict = self.audio_encoder.state_dict()
|
||||
|
||||
# for k,v in wav2lip_state_dict.items():
|
||||
# if 'audio_encoder' in k:
|
||||
# state_dict[k.replace('module.audio_encoder.', '')] = v
|
||||
# self.audio_encoder.load_state_dict(state_dict)
|
||||
|
||||
|
||||
def forward(self, audio_sequences):
|
||||
# audio_sequences = (B, T, 1, 80, 16)
|
||||
B = audio_sequences.size(0)
|
||||
|
||||
audio_sequences = torch.cat([audio_sequences[:, i] for i in range(audio_sequences.size(1))], dim=0)
|
||||
|
||||
audio_embedding = self.audio_encoder(audio_sequences) # B, 512, 1, 1
|
||||
dim = audio_embedding.shape[1]
|
||||
audio_embedding = audio_embedding.reshape((B, -1, dim, 1, 1))
|
||||
|
||||
return audio_embedding.squeeze(-1).squeeze(-1) #B seq_len+1 512
|
||||
@@ -0,0 +1,149 @@
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch import nn
|
||||
from src.audio2pose_models.res_unet import ResUnet
|
||||
|
||||
def class2onehot(idx, class_num):
|
||||
|
||||
assert torch.max(idx).item() < class_num
|
||||
onehot = torch.zeros(idx.size(0), class_num).to(idx.device)
|
||||
onehot.scatter_(1, idx, 1)
|
||||
return onehot
|
||||
|
||||
class CVAE(nn.Module):
|
||||
def __init__(self, cfg):
|
||||
super().__init__()
|
||||
encoder_layer_sizes = cfg.MODEL.CVAE.ENCODER_LAYER_SIZES
|
||||
decoder_layer_sizes = cfg.MODEL.CVAE.DECODER_LAYER_SIZES
|
||||
latent_size = cfg.MODEL.CVAE.LATENT_SIZE
|
||||
num_classes = cfg.DATASET.NUM_CLASSES
|
||||
audio_emb_in_size = cfg.MODEL.CVAE.AUDIO_EMB_IN_SIZE
|
||||
audio_emb_out_size = cfg.MODEL.CVAE.AUDIO_EMB_OUT_SIZE
|
||||
seq_len = cfg.MODEL.CVAE.SEQ_LEN
|
||||
|
||||
self.latent_size = latent_size
|
||||
|
||||
self.encoder = ENCODER(encoder_layer_sizes, latent_size, num_classes,
|
||||
audio_emb_in_size, audio_emb_out_size, seq_len)
|
||||
self.decoder = DECODER(decoder_layer_sizes, latent_size, num_classes,
|
||||
audio_emb_in_size, audio_emb_out_size, seq_len)
|
||||
def reparameterize(self, mu, logvar):
|
||||
std = torch.exp(0.5 * logvar)
|
||||
eps = torch.randn_like(std)
|
||||
return mu + eps * std
|
||||
|
||||
def forward(self, batch):
|
||||
batch = self.encoder(batch)
|
||||
mu = batch['mu']
|
||||
logvar = batch['logvar']
|
||||
z = self.reparameterize(mu, logvar)
|
||||
batch['z'] = z
|
||||
return self.decoder(batch)
|
||||
|
||||
def test(self, batch):
|
||||
'''
|
||||
class_id = batch['class']
|
||||
z = torch.randn([class_id.size(0), self.latent_size]).to(class_id.device)
|
||||
batch['z'] = z
|
||||
'''
|
||||
return self.decoder(batch)
|
||||
|
||||
class ENCODER(nn.Module):
|
||||
def __init__(self, layer_sizes, latent_size, num_classes,
|
||||
audio_emb_in_size, audio_emb_out_size, seq_len):
|
||||
super().__init__()
|
||||
|
||||
self.resunet = ResUnet()
|
||||
self.num_classes = num_classes
|
||||
self.seq_len = seq_len
|
||||
|
||||
self.MLP = nn.Sequential()
|
||||
layer_sizes[0] += latent_size + seq_len*audio_emb_out_size + 6
|
||||
for i, (in_size, out_size) in enumerate(zip(layer_sizes[:-1], layer_sizes[1:])):
|
||||
self.MLP.add_module(
|
||||
name="L{:d}".format(i), module=nn.Linear(in_size, out_size))
|
||||
self.MLP.add_module(name="A{:d}".format(i), module=nn.ReLU())
|
||||
|
||||
self.linear_means = nn.Linear(layer_sizes[-1], latent_size)
|
||||
self.linear_logvar = nn.Linear(layer_sizes[-1], latent_size)
|
||||
self.linear_audio = nn.Linear(audio_emb_in_size, audio_emb_out_size)
|
||||
|
||||
self.classbias = nn.Parameter(torch.randn(self.num_classes, latent_size))
|
||||
|
||||
def forward(self, batch):
|
||||
class_id = batch['class']
|
||||
pose_motion_gt = batch['pose_motion_gt'] #bs seq_len 6
|
||||
ref = batch['ref'] #bs 6
|
||||
bs = pose_motion_gt.shape[0]
|
||||
audio_in = batch['audio_emb'] # bs seq_len audio_emb_in_size
|
||||
|
||||
#pose encode
|
||||
pose_emb = self.resunet(pose_motion_gt.unsqueeze(1)) #bs 1 seq_len 6
|
||||
pose_emb = pose_emb.reshape(bs, -1) #bs seq_len*6
|
||||
|
||||
#audio mapping
|
||||
print(audio_in.shape)
|
||||
audio_out = self.linear_audio(audio_in) # bs seq_len audio_emb_out_size
|
||||
audio_out = audio_out.reshape(bs, -1)
|
||||
|
||||
class_bias = self.classbias[class_id] #bs latent_size
|
||||
x_in = torch.cat([ref, pose_emb, audio_out, class_bias], dim=-1) #bs seq_len*(audio_emb_out_size+6)+latent_size
|
||||
x_out = self.MLP(x_in)
|
||||
|
||||
mu = self.linear_means(x_out)
|
||||
logvar = self.linear_means(x_out) #bs latent_size
|
||||
|
||||
batch.update({'mu':mu, 'logvar':logvar})
|
||||
return batch
|
||||
|
||||
class DECODER(nn.Module):
|
||||
def __init__(self, layer_sizes, latent_size, num_classes,
|
||||
audio_emb_in_size, audio_emb_out_size, seq_len):
|
||||
super().__init__()
|
||||
|
||||
self.resunet = ResUnet()
|
||||
self.num_classes = num_classes
|
||||
self.seq_len = seq_len
|
||||
|
||||
self.MLP = nn.Sequential()
|
||||
input_size = latent_size + seq_len*audio_emb_out_size + 6
|
||||
for i, (in_size, out_size) in enumerate(zip([input_size]+layer_sizes[:-1], layer_sizes)):
|
||||
self.MLP.add_module(
|
||||
name="L{:d}".format(i), module=nn.Linear(in_size, out_size))
|
||||
if i+1 < len(layer_sizes):
|
||||
self.MLP.add_module(name="A{:d}".format(i), module=nn.ReLU())
|
||||
else:
|
||||
self.MLP.add_module(name="sigmoid", module=nn.Sigmoid())
|
||||
|
||||
self.pose_linear = nn.Linear(6, 6)
|
||||
self.linear_audio = nn.Linear(audio_emb_in_size, audio_emb_out_size)
|
||||
|
||||
self.classbias = nn.Parameter(torch.randn(self.num_classes, latent_size))
|
||||
|
||||
def forward(self, batch):
|
||||
|
||||
z = batch['z'] #bs latent_size
|
||||
bs = z.shape[0]
|
||||
class_id = batch['class']
|
||||
ref = batch['ref'] #bs 6
|
||||
audio_in = batch['audio_emb'] # bs seq_len audio_emb_in_size
|
||||
#print('audio_in: ', audio_in[:, :, :10])
|
||||
|
||||
audio_out = self.linear_audio(audio_in) # bs seq_len audio_emb_out_size
|
||||
#print('audio_out: ', audio_out[:, :, :10])
|
||||
audio_out = audio_out.reshape([bs, -1]) # bs seq_len*audio_emb_out_size
|
||||
class_bias = self.classbias[class_id] #bs latent_size
|
||||
|
||||
z = z + class_bias
|
||||
x_in = torch.cat([ref, z, audio_out], dim=-1)
|
||||
x_out = self.MLP(x_in) # bs layer_sizes[-1]
|
||||
x_out = x_out.reshape((bs, self.seq_len, -1))
|
||||
|
||||
#print('x_out: ', x_out)
|
||||
|
||||
pose_emb = self.resunet(x_out.unsqueeze(1)) #bs 1 seq_len 6
|
||||
|
||||
pose_motion_pred = self.pose_linear(pose_emb.squeeze(1)) #bs seq_len 6
|
||||
|
||||
batch.update({'pose_motion_pred':pose_motion_pred})
|
||||
return batch
|
||||
@@ -0,0 +1,76 @@
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch import nn
|
||||
|
||||
class ConvNormRelu(nn.Module):
|
||||
def __init__(self, conv_type='1d', in_channels=3, out_channels=64, downsample=False,
|
||||
kernel_size=None, stride=None, padding=None, norm='BN', leaky=False):
|
||||
super().__init__()
|
||||
if kernel_size is None:
|
||||
if downsample:
|
||||
kernel_size, stride, padding = 4, 2, 1
|
||||
else:
|
||||
kernel_size, stride, padding = 3, 1, 1
|
||||
|
||||
if conv_type == '2d':
|
||||
self.conv = nn.Conv2d(
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size,
|
||||
stride,
|
||||
padding,
|
||||
bias=False,
|
||||
)
|
||||
if norm == 'BN':
|
||||
self.norm = nn.BatchNorm2d(out_channels)
|
||||
elif norm == 'IN':
|
||||
self.norm = nn.InstanceNorm2d(out_channels)
|
||||
else:
|
||||
raise NotImplementedError
|
||||
elif conv_type == '1d':
|
||||
self.conv = nn.Conv1d(
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size,
|
||||
stride,
|
||||
padding,
|
||||
bias=False,
|
||||
)
|
||||
if norm == 'BN':
|
||||
self.norm = nn.BatchNorm1d(out_channels)
|
||||
elif norm == 'IN':
|
||||
self.norm = nn.InstanceNorm1d(out_channels)
|
||||
else:
|
||||
raise NotImplementedError
|
||||
nn.init.kaiming_normal_(self.conv.weight)
|
||||
|
||||
self.act = nn.LeakyReLU(negative_slope=0.2, inplace=False) if leaky else nn.ReLU(inplace=True)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.conv(x)
|
||||
if isinstance(self.norm, nn.InstanceNorm1d):
|
||||
x = self.norm(x.permute((0, 2, 1))).permute((0, 2, 1)) # normalize on [C]
|
||||
else:
|
||||
x = self.norm(x)
|
||||
x = self.act(x)
|
||||
return x
|
||||
|
||||
|
||||
class PoseSequenceDiscriminator(nn.Module):
|
||||
def __init__(self, cfg):
|
||||
super().__init__()
|
||||
self.cfg = cfg
|
||||
leaky = self.cfg.MODEL.DISCRIMINATOR.LEAKY_RELU
|
||||
|
||||
self.seq = nn.Sequential(
|
||||
ConvNormRelu('1d', cfg.MODEL.DISCRIMINATOR.INPUT_CHANNELS, 256, downsample=True, leaky=leaky), # B, 256, 64
|
||||
ConvNormRelu('1d', 256, 512, downsample=True, leaky=leaky), # B, 512, 32
|
||||
ConvNormRelu('1d', 512, 1024, kernel_size=3, stride=1, padding=1, leaky=leaky), # B, 1024, 16
|
||||
nn.Conv1d(1024, 1, kernel_size=3, stride=1, padding=1, bias=True) # B, 1, 16
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
x = x.reshape(x.size(0), x.size(1), -1).transpose(1, 2)
|
||||
x = self.seq(x)
|
||||
x = x.squeeze(1)
|
||||
return x
|
||||
@@ -0,0 +1,140 @@
|
||||
import torch.nn as nn
|
||||
import torch
|
||||
|
||||
|
||||
class ResidualConv(nn.Module):
|
||||
def __init__(self, input_dim, output_dim, stride, padding):
|
||||
super(ResidualConv, self).__init__()
|
||||
|
||||
self.conv_block = nn.Sequential(
|
||||
nn.BatchNorm2d(input_dim),
|
||||
nn.ReLU(),
|
||||
nn.Conv2d(
|
||||
input_dim, output_dim, kernel_size=3, stride=stride, padding=padding
|
||||
),
|
||||
nn.BatchNorm2d(output_dim),
|
||||
nn.ReLU(),
|
||||
nn.Conv2d(output_dim, output_dim, kernel_size=3, padding=1),
|
||||
)
|
||||
self.conv_skip = nn.Sequential(
|
||||
nn.Conv2d(input_dim, output_dim, kernel_size=3, stride=stride, padding=1),
|
||||
nn.BatchNorm2d(output_dim),
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
|
||||
return self.conv_block(x) + self.conv_skip(x)
|
||||
|
||||
|
||||
class Upsample(nn.Module):
|
||||
def __init__(self, input_dim, output_dim, kernel, stride):
|
||||
super(Upsample, self).__init__()
|
||||
|
||||
self.upsample = nn.ConvTranspose2d(
|
||||
input_dim, output_dim, kernel_size=kernel, stride=stride
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
return self.upsample(x)
|
||||
|
||||
|
||||
class Squeeze_Excite_Block(nn.Module):
|
||||
def __init__(self, channel, reduction=16):
|
||||
super(Squeeze_Excite_Block, self).__init__()
|
||||
self.avg_pool = nn.AdaptiveAvgPool2d(1)
|
||||
self.fc = nn.Sequential(
|
||||
nn.Linear(channel, channel // reduction, bias=False),
|
||||
nn.ReLU(inplace=True),
|
||||
nn.Linear(channel // reduction, channel, bias=False),
|
||||
nn.Sigmoid(),
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
b, c, _, _ = x.size()
|
||||
y = self.avg_pool(x).view(b, c)
|
||||
y = self.fc(y).view(b, c, 1, 1)
|
||||
return x * y.expand_as(x)
|
||||
|
||||
|
||||
class ASPP(nn.Module):
|
||||
def __init__(self, in_dims, out_dims, rate=[6, 12, 18]):
|
||||
super(ASPP, self).__init__()
|
||||
|
||||
self.aspp_block1 = nn.Sequential(
|
||||
nn.Conv2d(
|
||||
in_dims, out_dims, 3, stride=1, padding=rate[0], dilation=rate[0]
|
||||
),
|
||||
nn.ReLU(inplace=True),
|
||||
nn.BatchNorm2d(out_dims),
|
||||
)
|
||||
self.aspp_block2 = nn.Sequential(
|
||||
nn.Conv2d(
|
||||
in_dims, out_dims, 3, stride=1, padding=rate[1], dilation=rate[1]
|
||||
),
|
||||
nn.ReLU(inplace=True),
|
||||
nn.BatchNorm2d(out_dims),
|
||||
)
|
||||
self.aspp_block3 = nn.Sequential(
|
||||
nn.Conv2d(
|
||||
in_dims, out_dims, 3, stride=1, padding=rate[2], dilation=rate[2]
|
||||
),
|
||||
nn.ReLU(inplace=True),
|
||||
nn.BatchNorm2d(out_dims),
|
||||
)
|
||||
|
||||
self.output = nn.Conv2d(len(rate) * out_dims, out_dims, 1)
|
||||
self._init_weights()
|
||||
|
||||
def forward(self, x):
|
||||
x1 = self.aspp_block1(x)
|
||||
x2 = self.aspp_block2(x)
|
||||
x3 = self.aspp_block3(x)
|
||||
out = torch.cat([x1, x2, x3], dim=1)
|
||||
return self.output(out)
|
||||
|
||||
def _init_weights(self):
|
||||
for m in self.modules():
|
||||
if isinstance(m, nn.Conv2d):
|
||||
nn.init.kaiming_normal_(m.weight)
|
||||
elif isinstance(m, nn.BatchNorm2d):
|
||||
m.weight.data.fill_(1)
|
||||
m.bias.data.zero_()
|
||||
|
||||
|
||||
class Upsample_(nn.Module):
|
||||
def __init__(self, scale=2):
|
||||
super(Upsample_, self).__init__()
|
||||
|
||||
self.upsample = nn.Upsample(mode="bilinear", scale_factor=scale)
|
||||
|
||||
def forward(self, x):
|
||||
return self.upsample(x)
|
||||
|
||||
|
||||
class AttentionBlock(nn.Module):
|
||||
def __init__(self, input_encoder, input_decoder, output_dim):
|
||||
super(AttentionBlock, self).__init__()
|
||||
|
||||
self.conv_encoder = nn.Sequential(
|
||||
nn.BatchNorm2d(input_encoder),
|
||||
nn.ReLU(),
|
||||
nn.Conv2d(input_encoder, output_dim, 3, padding=1),
|
||||
nn.MaxPool2d(2, 2),
|
||||
)
|
||||
|
||||
self.conv_decoder = nn.Sequential(
|
||||
nn.BatchNorm2d(input_decoder),
|
||||
nn.ReLU(),
|
||||
nn.Conv2d(input_decoder, output_dim, 3, padding=1),
|
||||
)
|
||||
|
||||
self.conv_attn = nn.Sequential(
|
||||
nn.BatchNorm2d(output_dim),
|
||||
nn.ReLU(),
|
||||
nn.Conv2d(output_dim, 1, 1),
|
||||
)
|
||||
|
||||
def forward(self, x1, x2):
|
||||
out = self.conv_encoder(x1) + self.conv_decoder(x2)
|
||||
out = self.conv_attn(out)
|
||||
return out * x2
|
||||
@@ -0,0 +1,65 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from src.audio2pose_models.networks import ResidualConv, Upsample
|
||||
|
||||
|
||||
class ResUnet(nn.Module):
|
||||
def __init__(self, channel=1, filters=[32, 64, 128, 256]):
|
||||
super(ResUnet, self).__init__()
|
||||
|
||||
self.input_layer = nn.Sequential(
|
||||
nn.Conv2d(channel, filters[0], kernel_size=3, padding=1),
|
||||
nn.BatchNorm2d(filters[0]),
|
||||
nn.ReLU(),
|
||||
nn.Conv2d(filters[0], filters[0], kernel_size=3, padding=1),
|
||||
)
|
||||
self.input_skip = nn.Sequential(
|
||||
nn.Conv2d(channel, filters[0], kernel_size=3, padding=1)
|
||||
)
|
||||
|
||||
self.residual_conv_1 = ResidualConv(filters[0], filters[1], stride=(2,1), padding=1)
|
||||
self.residual_conv_2 = ResidualConv(filters[1], filters[2], stride=(2,1), padding=1)
|
||||
|
||||
self.bridge = ResidualConv(filters[2], filters[3], stride=(2,1), padding=1)
|
||||
|
||||
self.upsample_1 = Upsample(filters[3], filters[3], kernel=(2,1), stride=(2,1))
|
||||
self.up_residual_conv1 = ResidualConv(filters[3] + filters[2], filters[2], stride=1, padding=1)
|
||||
|
||||
self.upsample_2 = Upsample(filters[2], filters[2], kernel=(2,1), stride=(2,1))
|
||||
self.up_residual_conv2 = ResidualConv(filters[2] + filters[1], filters[1], stride=1, padding=1)
|
||||
|
||||
self.upsample_3 = Upsample(filters[1], filters[1], kernel=(2,1), stride=(2,1))
|
||||
self.up_residual_conv3 = ResidualConv(filters[1] + filters[0], filters[0], stride=1, padding=1)
|
||||
|
||||
self.output_layer = nn.Sequential(
|
||||
nn.Conv2d(filters[0], 1, 1, 1),
|
||||
nn.Sigmoid(),
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
# Encode
|
||||
x1 = self.input_layer(x) + self.input_skip(x)
|
||||
x2 = self.residual_conv_1(x1)
|
||||
x3 = self.residual_conv_2(x2)
|
||||
# Bridge
|
||||
x4 = self.bridge(x3)
|
||||
|
||||
# Decode
|
||||
x4 = self.upsample_1(x4)
|
||||
x5 = torch.cat([x4, x3], dim=1)
|
||||
|
||||
x6 = self.up_residual_conv1(x5)
|
||||
|
||||
x6 = self.upsample_2(x6)
|
||||
x7 = torch.cat([x6, x2], dim=1)
|
||||
|
||||
x8 = self.up_residual_conv2(x7)
|
||||
|
||||
x8 = self.upsample_3(x8)
|
||||
x9 = torch.cat([x8, x1], dim=1)
|
||||
|
||||
x10 = self.up_residual_conv3(x9)
|
||||
|
||||
output = self.output_layer(x10)
|
||||
|
||||
return output
|
||||
@@ -0,0 +1,58 @@
|
||||
DATASET:
|
||||
TRAIN_FILE_LIST: /apdcephfs_cq2/share_1290939/wenxuazhang/code/file_list/train.txt
|
||||
EVAL_FILE_LIST: /apdcephfs_cq2/share_1290939/wenxuazhang/code/file_list/val.txt
|
||||
TRAIN_BATCH_SIZE: 32
|
||||
EVAL_BATCH_SIZE: 32
|
||||
EXP: True
|
||||
EXP_DIM: 64
|
||||
FRAME_LEN: 32
|
||||
COEFF_LEN: 73
|
||||
NUM_CLASSES: 46
|
||||
AUDIO_ROOT_PATH: /apdcephfs_cq2/share_1290939/wenxuazhang/voxceleb1/wav
|
||||
COEFF_ROOT_PATH: /apdcephfs_cq2/share_1290939/wenxuazhang/voxceleb1/wav2lip_3dmm
|
||||
LMDB_PATH: /apdcephfs_cq2/share_1290939/shadowcun/datasets/VoxCeleb/v1/imdb
|
||||
DEBUG: True
|
||||
NUM_REPEATS: 2
|
||||
T: 40
|
||||
|
||||
|
||||
MODEL:
|
||||
FRAMEWORK: V2
|
||||
AUDIOENCODER:
|
||||
LEAKY_RELU: True
|
||||
NORM: 'IN'
|
||||
DISCRIMINATOR:
|
||||
LEAKY_RELU: False
|
||||
INPUT_CHANNELS: 6
|
||||
CVAE:
|
||||
AUDIO_EMB_IN_SIZE: 512
|
||||
AUDIO_EMB_OUT_SIZE: 128
|
||||
SEQ_LEN: 32
|
||||
LATENT_SIZE: 256
|
||||
ENCODER_LAYER_SIZES: [192, 1024]
|
||||
DECODER_LAYER_SIZES: [1024, 192]
|
||||
|
||||
|
||||
TRAIN:
|
||||
MAX_EPOCH: 300
|
||||
GENERATOR:
|
||||
LR: 2.0e-5
|
||||
DISCRIMINATOR:
|
||||
LR: 1.0e-5
|
||||
LOSS:
|
||||
W_FEAT: 0
|
||||
W_COEFF_EXP: 2
|
||||
W_LM: 1.0e-2
|
||||
W_LM_MOUTH: 0
|
||||
W_REG: 0
|
||||
W_SYNC: 0
|
||||
W_COLOR: 0
|
||||
W_EXPRESSION: 0
|
||||
W_LIPREADING: 0.01
|
||||
W_LIPREADING_VV: 0
|
||||
W_EYE_BLINK: 4
|
||||
|
||||
TAG:
|
||||
NAME: small_dataset
|
||||
|
||||
|
||||
@@ -0,0 +1,49 @@
|
||||
DATASET:
|
||||
TRAIN_FILE_LIST: /apdcephfs_cq2/share_1290939/wenxuazhang/code/audio2pose_unet_noAudio/dataset/train_33.txt
|
||||
EVAL_FILE_LIST: /apdcephfs_cq2/share_1290939/wenxuazhang/code/audio2pose_unet_noAudio/dataset/val.txt
|
||||
TRAIN_BATCH_SIZE: 64
|
||||
EVAL_BATCH_SIZE: 1
|
||||
EXP: True
|
||||
EXP_DIM: 64
|
||||
FRAME_LEN: 32
|
||||
COEFF_LEN: 73
|
||||
NUM_CLASSES: 46
|
||||
AUDIO_ROOT_PATH: /apdcephfs_cq2/share_1290939/wenxuazhang/voxceleb1/wav
|
||||
COEFF_ROOT_PATH: /apdcephfs_cq2/share_1290939/shadowcun/datasets/VoxCeleb/v1/imdb
|
||||
DEBUG: True
|
||||
|
||||
|
||||
MODEL:
|
||||
AUDIOENCODER:
|
||||
LEAKY_RELU: True
|
||||
NORM: 'IN'
|
||||
DISCRIMINATOR:
|
||||
LEAKY_RELU: False
|
||||
INPUT_CHANNELS: 6
|
||||
CVAE:
|
||||
AUDIO_EMB_IN_SIZE: 512
|
||||
AUDIO_EMB_OUT_SIZE: 6
|
||||
SEQ_LEN: 32
|
||||
LATENT_SIZE: 64
|
||||
ENCODER_LAYER_SIZES: [192, 128]
|
||||
DECODER_LAYER_SIZES: [128, 192]
|
||||
|
||||
|
||||
TRAIN:
|
||||
MAX_EPOCH: 150
|
||||
GENERATOR:
|
||||
LR: 1.0e-4
|
||||
DISCRIMINATOR:
|
||||
LR: 1.0e-4
|
||||
LOSS:
|
||||
LAMBDA_REG: 1
|
||||
LAMBDA_LANDMARKS: 0
|
||||
LAMBDA_VERTICES: 0
|
||||
LAMBDA_GAN_MOTION: 0.7
|
||||
LAMBDA_GAN_COEFF: 0
|
||||
LAMBDA_KL: 1
|
||||
|
||||
TAG:
|
||||
NAME: cvae_UNET_useAudio_usewav2lipAudioEncoder
|
||||
|
||||
|
||||
@@ -0,0 +1,45 @@
|
||||
model_params:
|
||||
common_params:
|
||||
num_kp: 15
|
||||
image_channel: 3
|
||||
feature_channel: 32
|
||||
estimate_jacobian: False # True
|
||||
kp_detector_params:
|
||||
temperature: 0.1
|
||||
block_expansion: 32
|
||||
max_features: 1024
|
||||
scale_factor: 0.25 # 0.25
|
||||
num_blocks: 5
|
||||
reshape_channel: 16384 # 16384 = 1024 * 16
|
||||
reshape_depth: 16
|
||||
he_estimator_params:
|
||||
block_expansion: 64
|
||||
max_features: 2048
|
||||
num_bins: 66
|
||||
generator_params:
|
||||
block_expansion: 64
|
||||
max_features: 512
|
||||
num_down_blocks: 2
|
||||
reshape_channel: 32
|
||||
reshape_depth: 16 # 512 = 32 * 16
|
||||
num_resblocks: 6
|
||||
estimate_occlusion_map: True
|
||||
dense_motion_params:
|
||||
block_expansion: 32
|
||||
max_features: 1024
|
||||
num_blocks: 5
|
||||
reshape_depth: 16
|
||||
compress: 4
|
||||
discriminator_params:
|
||||
scales: [1]
|
||||
block_expansion: 32
|
||||
max_features: 512
|
||||
num_blocks: 4
|
||||
sn: True
|
||||
mapping_params:
|
||||
coeff_nc: 70
|
||||
descriptor_nc: 1024
|
||||
layer: 3
|
||||
num_kp: 15
|
||||
num_bins: 66
|
||||
|
||||
@@ -0,0 +1,45 @@
|
||||
model_params:
|
||||
common_params:
|
||||
num_kp: 15
|
||||
image_channel: 3
|
||||
feature_channel: 32
|
||||
estimate_jacobian: False # True
|
||||
kp_detector_params:
|
||||
temperature: 0.1
|
||||
block_expansion: 32
|
||||
max_features: 1024
|
||||
scale_factor: 0.25 # 0.25
|
||||
num_blocks: 5
|
||||
reshape_channel: 16384 # 16384 = 1024 * 16
|
||||
reshape_depth: 16
|
||||
he_estimator_params:
|
||||
block_expansion: 64
|
||||
max_features: 2048
|
||||
num_bins: 66
|
||||
generator_params:
|
||||
block_expansion: 64
|
||||
max_features: 512
|
||||
num_down_blocks: 2
|
||||
reshape_channel: 32
|
||||
reshape_depth: 16 # 512 = 32 * 16
|
||||
num_resblocks: 6
|
||||
estimate_occlusion_map: True
|
||||
dense_motion_params:
|
||||
block_expansion: 32
|
||||
max_features: 1024
|
||||
num_blocks: 5
|
||||
reshape_depth: 16
|
||||
compress: 4
|
||||
discriminator_params:
|
||||
scales: [1]
|
||||
block_expansion: 32
|
||||
max_features: 512
|
||||
num_blocks: 4
|
||||
sn: True
|
||||
mapping_params:
|
||||
coeff_nc: 73
|
||||
descriptor_nc: 1024
|
||||
layer: 3
|
||||
num_kp: 15
|
||||
num_bins: 66
|
||||
|
||||
Binary file not shown.
@@ -0,0 +1,116 @@
|
||||
"""This package includes all the modules related to data loading and preprocessing
|
||||
|
||||
To add a custom dataset class called 'dummy', you need to add a file called 'dummy_dataset.py' and define a subclass 'DummyDataset' inherited from BaseDataset.
|
||||
You need to implement four functions:
|
||||
-- <__init__>: initialize the class, first call BaseDataset.__init__(self, opt).
|
||||
-- <__len__>: return the size of dataset.
|
||||
-- <__getitem__>: get a data point from data loader.
|
||||
-- <modify_commandline_options>: (optionally) add dataset-specific options and set default options.
|
||||
|
||||
Now you can use the dataset class by specifying flag '--dataset_mode dummy'.
|
||||
See our template dataset class 'template_dataset.py' for more details.
|
||||
"""
|
||||
import numpy as np
|
||||
import importlib
|
||||
import torch.utils.data
|
||||
from face3d.data.base_dataset import BaseDataset
|
||||
|
||||
|
||||
def find_dataset_using_name(dataset_name):
|
||||
"""Import the module "data/[dataset_name]_dataset.py".
|
||||
|
||||
In the file, the class called DatasetNameDataset() will
|
||||
be instantiated. It has to be a subclass of BaseDataset,
|
||||
and it is case-insensitive.
|
||||
"""
|
||||
dataset_filename = "data." + dataset_name + "_dataset"
|
||||
datasetlib = importlib.import_module(dataset_filename)
|
||||
|
||||
dataset = None
|
||||
target_dataset_name = dataset_name.replace('_', '') + 'dataset'
|
||||
for name, cls in datasetlib.__dict__.items():
|
||||
if name.lower() == target_dataset_name.lower() \
|
||||
and issubclass(cls, BaseDataset):
|
||||
dataset = cls
|
||||
|
||||
if dataset is None:
|
||||
raise NotImplementedError("In %s.py, there should be a subclass of BaseDataset with class name that matches %s in lowercase." % (dataset_filename, target_dataset_name))
|
||||
|
||||
return dataset
|
||||
|
||||
|
||||
def get_option_setter(dataset_name):
|
||||
"""Return the static method <modify_commandline_options> of the dataset class."""
|
||||
dataset_class = find_dataset_using_name(dataset_name)
|
||||
return dataset_class.modify_commandline_options
|
||||
|
||||
|
||||
def create_dataset(opt, rank=0):
|
||||
"""Create a dataset given the option.
|
||||
|
||||
This function wraps the class CustomDatasetDataLoader.
|
||||
This is the main interface between this package and 'train.py'/'test.py'
|
||||
|
||||
Example:
|
||||
>>> from data import create_dataset
|
||||
>>> dataset = create_dataset(opt)
|
||||
"""
|
||||
data_loader = CustomDatasetDataLoader(opt, rank=rank)
|
||||
dataset = data_loader.load_data()
|
||||
return dataset
|
||||
|
||||
class CustomDatasetDataLoader():
|
||||
"""Wrapper class of Dataset class that performs multi-threaded data loading"""
|
||||
|
||||
def __init__(self, opt, rank=0):
|
||||
"""Initialize this class
|
||||
|
||||
Step 1: create a dataset instance given the name [dataset_mode]
|
||||
Step 2: create a multi-threaded data loader.
|
||||
"""
|
||||
self.opt = opt
|
||||
dataset_class = find_dataset_using_name(opt.dataset_mode)
|
||||
self.dataset = dataset_class(opt)
|
||||
self.sampler = None
|
||||
print("rank %d %s dataset [%s] was created" % (rank, self.dataset.name, type(self.dataset).__name__))
|
||||
if opt.use_ddp and opt.isTrain:
|
||||
world_size = opt.world_size
|
||||
self.sampler = torch.utils.data.distributed.DistributedSampler(
|
||||
self.dataset,
|
||||
num_replicas=world_size,
|
||||
rank=rank,
|
||||
shuffle=not opt.serial_batches
|
||||
)
|
||||
self.dataloader = torch.utils.data.DataLoader(
|
||||
self.dataset,
|
||||
sampler=self.sampler,
|
||||
num_workers=int(opt.num_threads / world_size),
|
||||
batch_size=int(opt.batch_size / world_size),
|
||||
drop_last=True)
|
||||
else:
|
||||
self.dataloader = torch.utils.data.DataLoader(
|
||||
self.dataset,
|
||||
batch_size=opt.batch_size,
|
||||
shuffle=(not opt.serial_batches) and opt.isTrain,
|
||||
num_workers=int(opt.num_threads),
|
||||
drop_last=True
|
||||
)
|
||||
|
||||
def set_epoch(self, epoch):
|
||||
self.dataset.current_epoch = epoch
|
||||
if self.sampler is not None:
|
||||
self.sampler.set_epoch(epoch)
|
||||
|
||||
def load_data(self):
|
||||
return self
|
||||
|
||||
def __len__(self):
|
||||
"""Return the number of data in the dataset"""
|
||||
return min(len(self.dataset), self.opt.max_dataset_size)
|
||||
|
||||
def __iter__(self):
|
||||
"""Return a batch of data"""
|
||||
for i, data in enumerate(self.dataloader):
|
||||
if i * self.opt.batch_size >= self.opt.max_dataset_size:
|
||||
break
|
||||
yield data
|
||||
@@ -0,0 +1,125 @@
|
||||
"""This module implements an abstract base class (ABC) 'BaseDataset' for datasets.
|
||||
|
||||
It also includes common transformation functions (e.g., get_transform, __scale_width), which can be later used in subclasses.
|
||||
"""
|
||||
import random
|
||||
import numpy as np
|
||||
import torch.utils.data as data
|
||||
from PIL import Image
|
||||
import torchvision.transforms as transforms
|
||||
from abc import ABC, abstractmethod
|
||||
|
||||
|
||||
class BaseDataset(data.Dataset, ABC):
|
||||
"""This class is an abstract base class (ABC) for datasets.
|
||||
|
||||
To create a subclass, you need to implement the following four functions:
|
||||
-- <__init__>: initialize the class, first call BaseDataset.__init__(self, opt).
|
||||
-- <__len__>: return the size of dataset.
|
||||
-- <__getitem__>: get a data point.
|
||||
-- <modify_commandline_options>: (optionally) add dataset-specific options and set default options.
|
||||
"""
|
||||
|
||||
def __init__(self, opt):
|
||||
"""Initialize the class; save the options in the class
|
||||
|
||||
Parameters:
|
||||
opt (Option class)-- stores all the experiment flags; needs to be a subclass of BaseOptions
|
||||
"""
|
||||
self.opt = opt
|
||||
# self.root = opt.dataroot
|
||||
self.current_epoch = 0
|
||||
|
||||
@staticmethod
|
||||
def modify_commandline_options(parser, is_train):
|
||||
"""Add new dataset-specific options, and rewrite default values for existing options.
|
||||
|
||||
Parameters:
|
||||
parser -- original option parser
|
||||
is_train (bool) -- whether training phase or test phase. You can use this flag to add training-specific or test-specific options.
|
||||
|
||||
Returns:
|
||||
the modified parser.
|
||||
"""
|
||||
return parser
|
||||
|
||||
@abstractmethod
|
||||
def __len__(self):
|
||||
"""Return the total number of images in the dataset."""
|
||||
return 0
|
||||
|
||||
@abstractmethod
|
||||
def __getitem__(self, index):
|
||||
"""Return a data point and its metadata information.
|
||||
|
||||
Parameters:
|
||||
index - - a random integer for data indexing
|
||||
|
||||
Returns:
|
||||
a dictionary of data with their names. It ususally contains the data itself and its metadata information.
|
||||
"""
|
||||
pass
|
||||
|
||||
|
||||
def get_transform(grayscale=False):
|
||||
transform_list = []
|
||||
if grayscale:
|
||||
transform_list.append(transforms.Grayscale(1))
|
||||
transform_list += [transforms.ToTensor()]
|
||||
return transforms.Compose(transform_list)
|
||||
|
||||
def get_affine_mat(opt, size):
|
||||
shift_x, shift_y, scale, rot_angle, flip = 0., 0., 1., 0., False
|
||||
w, h = size
|
||||
|
||||
if 'shift' in opt.preprocess:
|
||||
shift_pixs = int(opt.shift_pixs)
|
||||
shift_x = random.randint(-shift_pixs, shift_pixs)
|
||||
shift_y = random.randint(-shift_pixs, shift_pixs)
|
||||
if 'scale' in opt.preprocess:
|
||||
scale = 1 + opt.scale_delta * (2 * random.random() - 1)
|
||||
if 'rot' in opt.preprocess:
|
||||
rot_angle = opt.rot_angle * (2 * random.random() - 1)
|
||||
rot_rad = -rot_angle * np.pi/180
|
||||
if 'flip' in opt.preprocess:
|
||||
flip = random.random() > 0.5
|
||||
|
||||
shift_to_origin = np.array([1, 0, -w//2, 0, 1, -h//2, 0, 0, 1]).reshape([3, 3])
|
||||
flip_mat = np.array([-1 if flip else 1, 0, 0, 0, 1, 0, 0, 0, 1]).reshape([3, 3])
|
||||
shift_mat = np.array([1, 0, shift_x, 0, 1, shift_y, 0, 0, 1]).reshape([3, 3])
|
||||
rot_mat = np.array([np.cos(rot_rad), np.sin(rot_rad), 0, -np.sin(rot_rad), np.cos(rot_rad), 0, 0, 0, 1]).reshape([3, 3])
|
||||
scale_mat = np.array([scale, 0, 0, 0, scale, 0, 0, 0, 1]).reshape([3, 3])
|
||||
shift_to_center = np.array([1, 0, w//2, 0, 1, h//2, 0, 0, 1]).reshape([3, 3])
|
||||
|
||||
affine = shift_to_center @ scale_mat @ rot_mat @ shift_mat @ flip_mat @ shift_to_origin
|
||||
affine_inv = np.linalg.inv(affine)
|
||||
return affine, affine_inv, flip
|
||||
|
||||
def apply_img_affine(img, affine_inv, method=Image.BICUBIC):
|
||||
return img.transform(img.size, Image.AFFINE, data=affine_inv.flatten()[:6], resample=Image.BICUBIC)
|
||||
|
||||
def apply_lm_affine(landmark, affine, flip, size):
|
||||
_, h = size
|
||||
lm = landmark.copy()
|
||||
lm[:, 1] = h - 1 - lm[:, 1]
|
||||
lm = np.concatenate((lm, np.ones([lm.shape[0], 1])), -1)
|
||||
lm = lm @ np.transpose(affine)
|
||||
lm[:, :2] = lm[:, :2] / lm[:, 2:]
|
||||
lm = lm[:, :2]
|
||||
lm[:, 1] = h - 1 - lm[:, 1]
|
||||
if flip:
|
||||
lm_ = lm.copy()
|
||||
lm_[:17] = lm[16::-1]
|
||||
lm_[17:22] = lm[26:21:-1]
|
||||
lm_[22:27] = lm[21:16:-1]
|
||||
lm_[31:36] = lm[35:30:-1]
|
||||
lm_[36:40] = lm[45:41:-1]
|
||||
lm_[40:42] = lm[47:45:-1]
|
||||
lm_[42:46] = lm[39:35:-1]
|
||||
lm_[46:48] = lm[41:39:-1]
|
||||
lm_[48:55] = lm[54:47:-1]
|
||||
lm_[55:60] = lm[59:54:-1]
|
||||
lm_[60:65] = lm[64:59:-1]
|
||||
lm_[65:68] = lm[67:64:-1]
|
||||
lm = lm_
|
||||
return lm
|
||||
@@ -0,0 +1,125 @@
|
||||
"""This script defines the custom dataset for Deep3DFaceRecon_pytorch
|
||||
"""
|
||||
|
||||
import os.path
|
||||
from data.base_dataset import BaseDataset, get_transform, get_affine_mat, apply_img_affine, apply_lm_affine
|
||||
from data.image_folder import make_dataset
|
||||
from PIL import Image
|
||||
import random
|
||||
import util.util as util
|
||||
import numpy as np
|
||||
import json
|
||||
import torch
|
||||
from scipy.io import loadmat, savemat
|
||||
import pickle
|
||||
from util.preprocess import align_img, estimate_norm
|
||||
from util.load_mats import load_lm3d
|
||||
|
||||
|
||||
def default_flist_reader(flist):
|
||||
"""
|
||||
flist format: impath label\nimpath label\n ...(same to caffe's filelist)
|
||||
"""
|
||||
imlist = []
|
||||
with open(flist, 'r') as rf:
|
||||
for line in rf.readlines():
|
||||
impath = line.strip()
|
||||
imlist.append(impath)
|
||||
|
||||
return imlist
|
||||
|
||||
def jason_flist_reader(flist):
|
||||
with open(flist, 'r') as fp:
|
||||
info = json.load(fp)
|
||||
return info
|
||||
|
||||
def parse_label(label):
|
||||
return torch.tensor(np.array(label).astype(np.float32))
|
||||
|
||||
|
||||
class FlistDataset(BaseDataset):
|
||||
"""
|
||||
It requires one directories to host training images '/path/to/data/train'
|
||||
You can train the model with the dataset flag '--dataroot /path/to/data'.
|
||||
"""
|
||||
|
||||
def __init__(self, opt):
|
||||
"""Initialize this dataset class.
|
||||
|
||||
Parameters:
|
||||
opt (Option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions
|
||||
"""
|
||||
BaseDataset.__init__(self, opt)
|
||||
|
||||
self.lm3d_std = load_lm3d(opt.bfm_folder)
|
||||
|
||||
msk_names = default_flist_reader(opt.flist)
|
||||
self.msk_paths = [os.path.join(opt.data_root, i) for i in msk_names]
|
||||
|
||||
self.size = len(self.msk_paths)
|
||||
self.opt = opt
|
||||
|
||||
self.name = 'train' if opt.isTrain else 'val'
|
||||
if '_' in opt.flist:
|
||||
self.name += '_' + opt.flist.split(os.sep)[-1].split('_')[0]
|
||||
|
||||
|
||||
def __getitem__(self, index):
|
||||
"""Return a data point and its metadata information.
|
||||
|
||||
Parameters:
|
||||
index (int) -- a random integer for data indexing
|
||||
|
||||
Returns a dictionary that contains A, B, A_paths and B_paths
|
||||
img (tensor) -- an image in the input domain
|
||||
msk (tensor) -- its corresponding attention mask
|
||||
lm (tensor) -- its corresponding 3d landmarks
|
||||
im_paths (str) -- image paths
|
||||
aug_flag (bool) -- a flag used to tell whether its raw or augmented
|
||||
"""
|
||||
msk_path = self.msk_paths[index % self.size] # make sure index is within then range
|
||||
img_path = msk_path.replace('mask/', '')
|
||||
lm_path = '.'.join(msk_path.replace('mask', 'landmarks').split('.')[:-1]) + '.txt'
|
||||
|
||||
raw_img = Image.open(img_path).convert('RGB')
|
||||
raw_msk = Image.open(msk_path).convert('RGB')
|
||||
raw_lm = np.loadtxt(lm_path).astype(np.float32)
|
||||
|
||||
_, img, lm, msk = align_img(raw_img, raw_lm, self.lm3d_std, raw_msk)
|
||||
|
||||
aug_flag = self.opt.use_aug and self.opt.isTrain
|
||||
if aug_flag:
|
||||
img, lm, msk = self._augmentation(img, lm, self.opt, msk)
|
||||
|
||||
_, H = img.size
|
||||
M = estimate_norm(lm, H)
|
||||
transform = get_transform()
|
||||
img_tensor = transform(img)
|
||||
msk_tensor = transform(msk)[:1, ...]
|
||||
lm_tensor = parse_label(lm)
|
||||
M_tensor = parse_label(M)
|
||||
|
||||
|
||||
return {'imgs': img_tensor,
|
||||
'lms': lm_tensor,
|
||||
'msks': msk_tensor,
|
||||
'M': M_tensor,
|
||||
'im_paths': img_path,
|
||||
'aug_flag': aug_flag,
|
||||
'dataset': self.name}
|
||||
|
||||
def _augmentation(self, img, lm, opt, msk=None):
|
||||
affine, affine_inv, flip = get_affine_mat(opt, img.size)
|
||||
img = apply_img_affine(img, affine_inv)
|
||||
lm = apply_lm_affine(lm, affine, flip, img.size)
|
||||
if msk is not None:
|
||||
msk = apply_img_affine(msk, affine_inv, method=Image.BILINEAR)
|
||||
return img, lm, msk
|
||||
|
||||
|
||||
|
||||
|
||||
def __len__(self):
|
||||
"""Return the total number of images in the dataset.
|
||||
"""
|
||||
return self.size
|
||||
@@ -0,0 +1,66 @@
|
||||
"""A modified image folder class
|
||||
|
||||
We modify the official PyTorch image folder (https://github.com/pytorch/vision/blob/master/torchvision/datasets/folder.py)
|
||||
so that this class can load images from both current directory and its subdirectories.
|
||||
"""
|
||||
import numpy as np
|
||||
import torch.utils.data as data
|
||||
|
||||
from PIL import Image
|
||||
import os
|
||||
import os.path
|
||||
|
||||
IMG_EXTENSIONS = [
|
||||
'.jpg', '.JPG', '.jpeg', '.JPEG',
|
||||
'.png', '.PNG', '.ppm', '.PPM', '.bmp', '.BMP',
|
||||
'.tif', '.TIF', '.tiff', '.TIFF',
|
||||
]
|
||||
|
||||
|
||||
def is_image_file(filename):
|
||||
return any(filename.endswith(extension) for extension in IMG_EXTENSIONS)
|
||||
|
||||
|
||||
def make_dataset(dir, max_dataset_size=float("inf")):
|
||||
images = []
|
||||
assert os.path.isdir(dir) or os.path.islink(dir), '%s is not a valid directory' % dir
|
||||
|
||||
for root, _, fnames in sorted(os.walk(dir, followlinks=True)):
|
||||
for fname in fnames:
|
||||
if is_image_file(fname):
|
||||
path = os.path.join(root, fname)
|
||||
images.append(path)
|
||||
return images[:min(max_dataset_size, len(images))]
|
||||
|
||||
|
||||
def default_loader(path):
|
||||
return Image.open(path).convert('RGB')
|
||||
|
||||
|
||||
class ImageFolder(data.Dataset):
|
||||
|
||||
def __init__(self, root, transform=None, return_paths=False,
|
||||
loader=default_loader):
|
||||
imgs = make_dataset(root)
|
||||
if len(imgs) == 0:
|
||||
raise(RuntimeError("Found 0 images in: " + root + "\n"
|
||||
"Supported image extensions are: " + ",".join(IMG_EXTENSIONS)))
|
||||
|
||||
self.root = root
|
||||
self.imgs = imgs
|
||||
self.transform = transform
|
||||
self.return_paths = return_paths
|
||||
self.loader = loader
|
||||
|
||||
def __getitem__(self, index):
|
||||
path = self.imgs[index]
|
||||
img = self.loader(path)
|
||||
if self.transform is not None:
|
||||
img = self.transform(img)
|
||||
if self.return_paths:
|
||||
return img, path
|
||||
else:
|
||||
return img
|
||||
|
||||
def __len__(self):
|
||||
return len(self.imgs)
|
||||
@@ -0,0 +1,75 @@
|
||||
"""Dataset class template
|
||||
|
||||
This module provides a template for users to implement custom datasets.
|
||||
You can specify '--dataset_mode template' to use this dataset.
|
||||
The class name should be consistent with both the filename and its dataset_mode option.
|
||||
The filename should be <dataset_mode>_dataset.py
|
||||
The class name should be <Dataset_mode>Dataset.py
|
||||
You need to implement the following functions:
|
||||
-- <modify_commandline_options>: Add dataset-specific options and rewrite default values for existing options.
|
||||
-- <__init__>: Initialize this dataset class.
|
||||
-- <__getitem__>: Return a data point and its metadata information.
|
||||
-- <__len__>: Return the number of images.
|
||||
"""
|
||||
from data.base_dataset import BaseDataset, get_transform
|
||||
# from data.image_folder import make_dataset
|
||||
# from PIL import Image
|
||||
|
||||
|
||||
class TemplateDataset(BaseDataset):
|
||||
"""A template dataset class for you to implement custom datasets."""
|
||||
@staticmethod
|
||||
def modify_commandline_options(parser, is_train):
|
||||
"""Add new dataset-specific options, and rewrite default values for existing options.
|
||||
|
||||
Parameters:
|
||||
parser -- original option parser
|
||||
is_train (bool) -- whether training phase or test phase. You can use this flag to add training-specific or test-specific options.
|
||||
|
||||
Returns:
|
||||
the modified parser.
|
||||
"""
|
||||
parser.add_argument('--new_dataset_option', type=float, default=1.0, help='new dataset option')
|
||||
parser.set_defaults(max_dataset_size=10, new_dataset_option=2.0) # specify dataset-specific default values
|
||||
return parser
|
||||
|
||||
def __init__(self, opt):
|
||||
"""Initialize this dataset class.
|
||||
|
||||
Parameters:
|
||||
opt (Option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions
|
||||
|
||||
A few things can be done here.
|
||||
- save the options (have been done in BaseDataset)
|
||||
- get image paths and meta information of the dataset.
|
||||
- define the image transformation.
|
||||
"""
|
||||
# save the option and dataset root
|
||||
BaseDataset.__init__(self, opt)
|
||||
# get the image paths of your dataset;
|
||||
self.image_paths = [] # You can call sorted(make_dataset(self.root, opt.max_dataset_size)) to get all the image paths under the directory self.root
|
||||
# define the default transform function. You can use <base_dataset.get_transform>; You can also define your custom transform function
|
||||
self.transform = get_transform(opt)
|
||||
|
||||
def __getitem__(self, index):
|
||||
"""Return a data point and its metadata information.
|
||||
|
||||
Parameters:
|
||||
index -- a random integer for data indexing
|
||||
|
||||
Returns:
|
||||
a dictionary of data with their names. It usually contains the data itself and its metadata information.
|
||||
|
||||
Step 1: get a random image path: e.g., path = self.image_paths[index]
|
||||
Step 2: load your data from the disk: e.g., image = Image.open(path).convert('RGB').
|
||||
Step 3: convert your data to a PyTorch tensor. You can use helpder functions such as self.transform. e.g., data = self.transform(image)
|
||||
Step 4: return a data point as a dictionary.
|
||||
"""
|
||||
path = 'temp' # needs to be a string
|
||||
data_A = None # needs to be a tensor
|
||||
data_B = None # needs to be a tensor
|
||||
return {'data_A': data_A, 'data_B': data_B, 'path': path}
|
||||
|
||||
def __len__(self):
|
||||
"""Return the total number of images."""
|
||||
return len(self.image_paths)
|
||||
@@ -0,0 +1,108 @@
|
||||
import os
|
||||
import cv2
|
||||
import time
|
||||
import glob
|
||||
import argparse
|
||||
import face_alignment
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
from tqdm import tqdm
|
||||
from itertools import cycle
|
||||
|
||||
from torch.multiprocessing import Pool, Process, set_start_method
|
||||
|
||||
class KeypointExtractor():
|
||||
def __init__(self, device):
|
||||
self.detector = face_alignment.FaceAlignment(face_alignment.LandmarksType._2D,
|
||||
device=device)
|
||||
|
||||
def extract_keypoint(self, images, name=None, info=True):
|
||||
if isinstance(images, list):
|
||||
keypoints = []
|
||||
if info:
|
||||
i_range = tqdm(images,desc='landmark Det:')
|
||||
else:
|
||||
i_range = images
|
||||
|
||||
for image in i_range:
|
||||
current_kp = self.extract_keypoint(image)
|
||||
if np.mean(current_kp) == -1 and keypoints:
|
||||
keypoints.append(keypoints[-1])
|
||||
else:
|
||||
keypoints.append(current_kp[None])
|
||||
|
||||
keypoints = np.concatenate(keypoints, 0)
|
||||
np.savetxt(os.path.splitext(name)[0]+'.txt', keypoints.reshape(-1))
|
||||
return keypoints
|
||||
else:
|
||||
while True:
|
||||
try:
|
||||
keypoints = self.detector.get_landmarks_from_image(np.array(images))[0]
|
||||
break
|
||||
except RuntimeError as e:
|
||||
if str(e).startswith('CUDA'):
|
||||
print("Warning: out of memory, sleep for 1s")
|
||||
time.sleep(1)
|
||||
else:
|
||||
print(e)
|
||||
break
|
||||
except TypeError:
|
||||
print('No face detected in this image')
|
||||
shape = [68, 2]
|
||||
keypoints = -1. * np.ones(shape)
|
||||
break
|
||||
if name is not None:
|
||||
np.savetxt(os.path.splitext(name)[0]+'.txt', keypoints.reshape(-1))
|
||||
return keypoints
|
||||
|
||||
def read_video(filename):
|
||||
frames = []
|
||||
cap = cv2.VideoCapture(filename)
|
||||
while cap.isOpened():
|
||||
ret, frame = cap.read()
|
||||
if ret:
|
||||
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
|
||||
frame = Image.fromarray(frame)
|
||||
frames.append(frame)
|
||||
else:
|
||||
break
|
||||
cap.release()
|
||||
return frames
|
||||
|
||||
def run(data):
|
||||
filename, opt, device = data
|
||||
os.environ['CUDA_VISIBLE_DEVICES'] = device
|
||||
kp_extractor = KeypointExtractor()
|
||||
images = read_video(filename)
|
||||
name = filename.split('/')[-2:]
|
||||
os.makedirs(os.path.join(opt.output_dir, name[-2]), exist_ok=True)
|
||||
kp_extractor.extract_keypoint(
|
||||
images,
|
||||
name=os.path.join(opt.output_dir, name[-2], name[-1])
|
||||
)
|
||||
|
||||
if __name__ == '__main__':
|
||||
set_start_method('spawn')
|
||||
parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
|
||||
parser.add_argument('--input_dir', type=str, help='the folder of the input files')
|
||||
parser.add_argument('--output_dir', type=str, help='the folder of the output files')
|
||||
parser.add_argument('--device_ids', type=str, default='0,1')
|
||||
parser.add_argument('--workers', type=int, default=4)
|
||||
|
||||
opt = parser.parse_args()
|
||||
filenames = list()
|
||||
VIDEO_EXTENSIONS_LOWERCASE = {'mp4'}
|
||||
VIDEO_EXTENSIONS = VIDEO_EXTENSIONS_LOWERCASE.union({f.upper() for f in VIDEO_EXTENSIONS_LOWERCASE})
|
||||
extensions = VIDEO_EXTENSIONS
|
||||
|
||||
for ext in extensions:
|
||||
os.listdir(f'{opt.input_dir}')
|
||||
print(f'{opt.input_dir}/*.{ext}')
|
||||
filenames = sorted(glob.glob(f'{opt.input_dir}/*.{ext}'))
|
||||
print('Total number of videos:', len(filenames))
|
||||
pool = Pool(opt.workers)
|
||||
args_list = cycle([opt])
|
||||
device_ids = opt.device_ids.split(",")
|
||||
device_ids = cycle(device_ids)
|
||||
for data in tqdm(pool.imap_unordered(run, zip(filenames, args_list, device_ids))):
|
||||
None
|
||||
@@ -0,0 +1,151 @@
|
||||
import os
|
||||
import cv2
|
||||
import time
|
||||
import glob
|
||||
import argparse
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
import torch
|
||||
from tqdm import tqdm
|
||||
from itertools import cycle
|
||||
from torch.multiprocessing import Pool, Process, set_start_method
|
||||
|
||||
from facexlib.alignment import landmark_98_to_68
|
||||
from facexlib.detection import init_detection_model
|
||||
|
||||
from facexlib.utils import load_file_from_url
|
||||
from src.face3d.util.my_awing_arch import FAN
|
||||
|
||||
def init_alignment_model(model_name, half=False, device='cuda', model_rootpath=None):
|
||||
if model_name == 'awing_fan':
|
||||
model = FAN(num_modules=4, num_landmarks=98, device=device)
|
||||
model_url = 'https://github.com/xinntao/facexlib/releases/download/v0.1.0/alignment_WFLW_4HG.pth'
|
||||
else:
|
||||
raise NotImplementedError(f'{model_name} is not implemented.')
|
||||
|
||||
model_path = load_file_from_url(
|
||||
url=model_url, model_dir='facexlib/weights', progress=True, file_name=None, save_dir=model_rootpath)
|
||||
model.load_state_dict(torch.load(model_path, map_location=device)['state_dict'], strict=True)
|
||||
model.eval()
|
||||
model = model.to(device)
|
||||
return model
|
||||
|
||||
|
||||
class KeypointExtractor():
|
||||
def __init__(self, device='cuda'):
|
||||
|
||||
### gfpgan/weights
|
||||
try:
|
||||
import webui # in webui
|
||||
root_path = 'extensions/SadTalker/gfpgan/weights'
|
||||
|
||||
except:
|
||||
root_path = 'gfpgan/weights'
|
||||
|
||||
self.detector = init_alignment_model('awing_fan',device=device, model_rootpath=root_path)
|
||||
self.det_net = init_detection_model('retinaface_resnet50', half=False,device=device, model_rootpath=root_path)
|
||||
|
||||
def extract_keypoint(self, images, name=None, info=True):
|
||||
if isinstance(images, list):
|
||||
keypoints = []
|
||||
if info:
|
||||
i_range = tqdm(images,desc='landmark Det:')
|
||||
else:
|
||||
i_range = images
|
||||
|
||||
for image in i_range:
|
||||
current_kp = self.extract_keypoint(image)
|
||||
# current_kp = self.detector.get_landmarks(np.array(image))
|
||||
if np.mean(current_kp) == -1 and keypoints:
|
||||
keypoints.append(keypoints[-1])
|
||||
else:
|
||||
keypoints.append(current_kp[None])
|
||||
|
||||
keypoints = np.concatenate(keypoints, 0)
|
||||
np.savetxt(os.path.splitext(name)[0]+'.txt', keypoints.reshape(-1))
|
||||
return keypoints
|
||||
else:
|
||||
while True:
|
||||
try:
|
||||
with torch.no_grad():
|
||||
# face detection -> face alignment.
|
||||
img = np.array(images)
|
||||
bboxes = self.det_net.detect_faces(images, 0.97)
|
||||
|
||||
bboxes = bboxes[0]
|
||||
img = img[int(bboxes[1]):int(bboxes[3]), int(bboxes[0]):int(bboxes[2]), :]
|
||||
|
||||
keypoints = landmark_98_to_68(self.detector.get_landmarks(img)) # [0]
|
||||
|
||||
#### keypoints to the original location
|
||||
keypoints[:,0] += int(bboxes[0])
|
||||
keypoints[:,1] += int(bboxes[1])
|
||||
|
||||
break
|
||||
except RuntimeError as e:
|
||||
if str(e).startswith('CUDA'):
|
||||
print("Warning: out of memory, sleep for 1s")
|
||||
time.sleep(1)
|
||||
else:
|
||||
print(e)
|
||||
break
|
||||
except TypeError:
|
||||
print('No face detected in this image')
|
||||
shape = [68, 2]
|
||||
keypoints = -1. * np.ones(shape)
|
||||
break
|
||||
if name is not None:
|
||||
np.savetxt(os.path.splitext(name)[0]+'.txt', keypoints.reshape(-1))
|
||||
return keypoints
|
||||
|
||||
def read_video(filename):
|
||||
frames = []
|
||||
cap = cv2.VideoCapture(filename)
|
||||
while cap.isOpened():
|
||||
ret, frame = cap.read()
|
||||
if ret:
|
||||
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
|
||||
frame = Image.fromarray(frame)
|
||||
frames.append(frame)
|
||||
else:
|
||||
break
|
||||
cap.release()
|
||||
return frames
|
||||
|
||||
def run(data):
|
||||
filename, opt, device = data
|
||||
os.environ['CUDA_VISIBLE_DEVICES'] = device
|
||||
kp_extractor = KeypointExtractor()
|
||||
images = read_video(filename)
|
||||
name = filename.split('/')[-2:]
|
||||
os.makedirs(os.path.join(opt.output_dir, name[-2]), exist_ok=True)
|
||||
kp_extractor.extract_keypoint(
|
||||
images,
|
||||
name=os.path.join(opt.output_dir, name[-2], name[-1])
|
||||
)
|
||||
|
||||
if __name__ == '__main__':
|
||||
set_start_method('spawn')
|
||||
parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
|
||||
parser.add_argument('--input_dir', type=str, help='the folder of the input files')
|
||||
parser.add_argument('--output_dir', type=str, help='the folder of the output files')
|
||||
parser.add_argument('--device_ids', type=str, default='0,1')
|
||||
parser.add_argument('--workers', type=int, default=4)
|
||||
|
||||
opt = parser.parse_args()
|
||||
filenames = list()
|
||||
VIDEO_EXTENSIONS_LOWERCASE = {'mp4'}
|
||||
VIDEO_EXTENSIONS = VIDEO_EXTENSIONS_LOWERCASE.union({f.upper() for f in VIDEO_EXTENSIONS_LOWERCASE})
|
||||
extensions = VIDEO_EXTENSIONS
|
||||
|
||||
for ext in extensions:
|
||||
os.listdir(f'{opt.input_dir}')
|
||||
print(f'{opt.input_dir}/*.{ext}')
|
||||
filenames = sorted(glob.glob(f'{opt.input_dir}/*.{ext}'))
|
||||
print('Total number of videos:', len(filenames))
|
||||
pool = Pool(opt.workers)
|
||||
args_list = cycle([opt])
|
||||
device_ids = opt.device_ids.split(",")
|
||||
device_ids = cycle(device_ids)
|
||||
for data in tqdm(pool.imap_unordered(run, zip(filenames, args_list, device_ids))):
|
||||
None
|
||||
@@ -0,0 +1,67 @@
|
||||
"""This package contains modules related to objective functions, optimizations, and network architectures.
|
||||
|
||||
To add a custom model class called 'dummy', you need to add a file called 'dummy_model.py' and define a subclass DummyModel inherited from BaseModel.
|
||||
You need to implement the following five functions:
|
||||
-- <__init__>: initialize the class; first call BaseModel.__init__(self, opt).
|
||||
-- <set_input>: unpack data from dataset and apply preprocessing.
|
||||
-- <forward>: produce intermediate results.
|
||||
-- <optimize_parameters>: calculate loss, gradients, and update network weights.
|
||||
-- <modify_commandline_options>: (optionally) add model-specific options and set default options.
|
||||
|
||||
In the function <__init__>, you need to define four lists:
|
||||
-- self.loss_names (str list): specify the training losses that you want to plot and save.
|
||||
-- self.model_names (str list): define networks used in our training.
|
||||
-- self.visual_names (str list): specify the images that you want to display and save.
|
||||
-- self.optimizers (optimizer list): define and initialize optimizers. You can define one optimizer for each network. If two networks are updated at the same time, you can use itertools.chain to group them. See cycle_gan_model.py for an usage.
|
||||
|
||||
Now you can use the model class by specifying flag '--model dummy'.
|
||||
See our template model class 'template_model.py' for more details.
|
||||
"""
|
||||
|
||||
import importlib
|
||||
from src.face3d.models.base_model import BaseModel
|
||||
|
||||
|
||||
def find_model_using_name(model_name):
|
||||
"""Import the module "models/[model_name]_model.py".
|
||||
|
||||
In the file, the class called DatasetNameModel() will
|
||||
be instantiated. It has to be a subclass of BaseModel,
|
||||
and it is case-insensitive.
|
||||
"""
|
||||
model_filename = "face3d.models." + model_name + "_model"
|
||||
modellib = importlib.import_module(model_filename)
|
||||
model = None
|
||||
target_model_name = model_name.replace('_', '') + 'model'
|
||||
for name, cls in modellib.__dict__.items():
|
||||
if name.lower() == target_model_name.lower() \
|
||||
and issubclass(cls, BaseModel):
|
||||
model = cls
|
||||
|
||||
if model is None:
|
||||
print("In %s.py, there should be a subclass of BaseModel with class name that matches %s in lowercase." % (model_filename, target_model_name))
|
||||
exit(0)
|
||||
|
||||
return model
|
||||
|
||||
|
||||
def get_option_setter(model_name):
|
||||
"""Return the static method <modify_commandline_options> of the model class."""
|
||||
model_class = find_model_using_name(model_name)
|
||||
return model_class.modify_commandline_options
|
||||
|
||||
|
||||
def create_model(opt):
|
||||
"""Create a model given the option.
|
||||
|
||||
This function warps the class CustomDatasetDataLoader.
|
||||
This is the main interface between this package and 'train.py'/'test.py'
|
||||
|
||||
Example:
|
||||
>>> from models import create_model
|
||||
>>> model = create_model(opt)
|
||||
"""
|
||||
model = find_model_using_name(opt.model)
|
||||
instance = model(opt)
|
||||
print("model [%s] was created" % type(instance).__name__)
|
||||
return instance
|
||||
@@ -0,0 +1,164 @@
|
||||
# Distributed Arcface Training in Pytorch
|
||||
|
||||
This is a deep learning library that makes face recognition efficient, and effective, which can train tens of millions
|
||||
identity on a single server.
|
||||
|
||||
## Requirements
|
||||
|
||||
- Install [pytorch](http://pytorch.org) (torch>=1.6.0), our doc for [install.md](docs/install.md).
|
||||
- `pip install -r requirements.txt`.
|
||||
- Download the dataset
|
||||
from [https://github.com/deepinsight/insightface/tree/master/recognition/_datasets_](https://github.com/deepinsight/insightface/tree/master/recognition/_datasets_)
|
||||
.
|
||||
|
||||
## How to Training
|
||||
|
||||
To train a model, run `train.py` with the path to the configs:
|
||||
|
||||
### 1. Single node, 8 GPUs:
|
||||
|
||||
```shell
|
||||
python -m torch.distributed.launch --nproc_per_node=8 --nnodes=1 --node_rank=0 --master_addr="127.0.0.1" --master_port=1234 train.py configs/ms1mv3_r50
|
||||
```
|
||||
|
||||
### 2. Multiple nodes, each node 8 GPUs:
|
||||
|
||||
Node 0:
|
||||
|
||||
```shell
|
||||
python -m torch.distributed.launch --nproc_per_node=8 --nnodes=2 --node_rank=0 --master_addr="ip1" --master_port=1234 train.py train.py configs/ms1mv3_r50
|
||||
```
|
||||
|
||||
Node 1:
|
||||
|
||||
```shell
|
||||
python -m torch.distributed.launch --nproc_per_node=8 --nnodes=2 --node_rank=1 --master_addr="ip1" --master_port=1234 train.py train.py configs/ms1mv3_r50
|
||||
```
|
||||
|
||||
### 3.Training resnet2060 with 8 GPUs:
|
||||
|
||||
```shell
|
||||
python -m torch.distributed.launch --nproc_per_node=8 --nnodes=1 --node_rank=0 --master_addr="127.0.0.1" --master_port=1234 train.py configs/ms1mv3_r2060.py
|
||||
```
|
||||
|
||||
## Model Zoo
|
||||
|
||||
- The models are available for non-commercial research purposes only.
|
||||
- All models can be found in here.
|
||||
- [Baidu Yun Pan](https://pan.baidu.com/s/1CL-l4zWqsI1oDuEEYVhj-g): e8pw
|
||||
- [onedrive](https://1drv.ms/u/s!AswpsDO2toNKq0lWY69vN58GR6mw?e=p9Ov5d)
|
||||
|
||||
### Performance on [**ICCV2021-MFR**](http://iccv21-mfr.com/)
|
||||
|
||||
ICCV2021-MFR testset consists of non-celebrities so we can ensure that it has very few overlap with public available face
|
||||
recognition training set, such as MS1M and CASIA as they mostly collected from online celebrities.
|
||||
As the result, we can evaluate the FAIR performance for different algorithms.
|
||||
|
||||
For **ICCV2021-MFR-ALL** set, TAR is measured on all-to-all 1:1 protocal, with FAR less than 0.000001(e-6). The
|
||||
globalised multi-racial testset contains 242,143 identities and 1,624,305 images.
|
||||
|
||||
For **ICCV2021-MFR-MASK** set, TAR is measured on mask-to-nonmask 1:1 protocal, with FAR less than 0.0001(e-4).
|
||||
Mask testset contains 6,964 identities, 6,964 masked images and 13,928 non-masked images.
|
||||
There are totally 13,928 positive pairs and 96,983,824 negative pairs.
|
||||
|
||||
| Datasets | backbone | Training throughout | Size / MB | **ICCV2021-MFR-MASK** | **ICCV2021-MFR-ALL** |
|
||||
| :---: | :--- | :--- | :--- |:--- |:--- |
|
||||
| MS1MV3 | r18 | - | 91 | **47.85** | **68.33** |
|
||||
| Glint360k | r18 | 8536 | 91 | **53.32** | **72.07** |
|
||||
| MS1MV3 | r34 | - | 130 | **58.72** | **77.36** |
|
||||
| Glint360k | r34 | 6344 | 130 | **65.10** | **83.02** |
|
||||
| MS1MV3 | r50 | 5500 | 166 | **63.85** | **80.53** |
|
||||
| Glint360k | r50 | 5136 | 166 | **70.23** | **87.08** |
|
||||
| MS1MV3 | r100 | - | 248 | **69.09** | **84.31** |
|
||||
| Glint360k | r100 | 3332 | 248 | **75.57** | **90.66** |
|
||||
| MS1MV3 | mobilefacenet | 12185 | 7.8 | **41.52** | **65.26** |
|
||||
| Glint360k | mobilefacenet | 11197 | 7.8 | **44.52** | **66.48** |
|
||||
|
||||
### Performance on IJB-C and Verification Datasets
|
||||
|
||||
| Datasets | backbone | IJBC(1e-05) | IJBC(1e-04) | agedb30 | cfp_fp | lfw | log |
|
||||
| :---: | :--- | :--- | :--- | :--- |:--- |:--- |:--- |
|
||||
| MS1MV3 | r18 | 92.07 | 94.66 | 97.77 | 97.73 | 99.77 |[log](https://raw.githubusercontent.com/anxiangsir/insightface_arcface_log/master/ms1mv3_arcface_r18_fp16/training.log)|
|
||||
| MS1MV3 | r34 | 94.10 | 95.90 | 98.10 | 98.67 | 99.80 |[log](https://raw.githubusercontent.com/anxiangsir/insightface_arcface_log/master/ms1mv3_arcface_r34_fp16/training.log)|
|
||||
| MS1MV3 | r50 | 94.79 | 96.46 | 98.35 | 98.96 | 99.83 |[log](https://raw.githubusercontent.com/anxiangsir/insightface_arcface_log/master/ms1mv3_arcface_r50_fp16/training.log)|
|
||||
| MS1MV3 | r100 | 95.31 | 96.81 | 98.48 | 99.06 | 99.85 |[log](https://raw.githubusercontent.com/anxiangsir/insightface_arcface_log/master/ms1mv3_arcface_r100_fp16/training.log)|
|
||||
| MS1MV3 | **r2060**| 95.34 | 97.11 | 98.67 | 99.24 | 99.87 |[log](https://raw.githubusercontent.com/anxiangsir/insightface_arcface_log/master/ms1mv3_arcface_r2060_fp16/training.log)|
|
||||
| Glint360k |r18-0.1 | 93.16 | 95.33 | 97.72 | 97.73 | 99.77 |[log](https://raw.githubusercontent.com/anxiangsir/insightface_arcface_log/master/glint360k_cosface_r18_fp16_0.1/training.log)|
|
||||
| Glint360k |r34-0.1 | 95.16 | 96.56 | 98.33 | 98.78 | 99.82 |[log](https://raw.githubusercontent.com/anxiangsir/insightface_arcface_log/master/glint360k_cosface_r34_fp16_0.1/training.log)|
|
||||
| Glint360k |r50-0.1 | 95.61 | 96.97 | 98.38 | 99.20 | 99.83 |[log](https://raw.githubusercontent.com/anxiangsir/insightface_arcface_log/master/glint360k_cosface_r50_fp16_0.1/training.log)|
|
||||
| Glint360k |r100-0.1 | 95.88 | 97.32 | 98.48 | 99.29 | 99.82 |[log](https://raw.githubusercontent.com/anxiangsir/insightface_arcface_log/master/glint360k_cosface_r100_fp16_0.1/training.log)|
|
||||
|
||||
[comment]: <> (More details see [model.md](docs/modelzoo.md) in docs.)
|
||||
|
||||
|
||||
## [Speed Benchmark](docs/speed_benchmark.md)
|
||||
|
||||
**Arcface Torch** can train large-scale face recognition training set efficiently and quickly. When the number of
|
||||
classes in training sets is greater than 300K and the training is sufficient, partial fc sampling strategy will get same
|
||||
accuracy with several times faster training performance and smaller GPU memory.
|
||||
Partial FC is a sparse variant of the model parallel architecture for large sacle face recognition. Partial FC use a
|
||||
sparse softmax, where each batch dynamicly sample a subset of class centers for training. In each iteration, only a
|
||||
sparse part of the parameters will be updated, which can reduce a lot of GPU memory and calculations. With Partial FC,
|
||||
we can scale trainset of 29 millions identities, the largest to date. Partial FC also supports multi-machine distributed
|
||||
training and mixed precision training.
|
||||
|
||||
|
||||
|
||||
More details see
|
||||
[speed_benchmark.md](docs/speed_benchmark.md) in docs.
|
||||
|
||||
### 1. Training speed of different parallel methods (samples / second), Tesla V100 32GB * 8. (Larger is better)
|
||||
|
||||
`-` means training failed because of gpu memory limitations.
|
||||
|
||||
| Number of Identities in Dataset | Data Parallel | Model Parallel | Partial FC 0.1 |
|
||||
| :--- | :--- | :--- | :--- |
|
||||
|125000 | 4681 | 4824 | 5004 |
|
||||
|1400000 | **1672** | 3043 | 4738 |
|
||||
|5500000 | **-** | **1389** | 3975 |
|
||||
|8000000 | **-** | **-** | 3565 |
|
||||
|16000000 | **-** | **-** | 2679 |
|
||||
|29000000 | **-** | **-** | **1855** |
|
||||
|
||||
### 2. GPU memory cost of different parallel methods (MB per GPU), Tesla V100 32GB * 8. (Smaller is better)
|
||||
|
||||
| Number of Identities in Dataset | Data Parallel | Model Parallel | Partial FC 0.1 |
|
||||
| :--- | :--- | :--- | :--- |
|
||||
|125000 | 7358 | 5306 | 4868 |
|
||||
|1400000 | 32252 | 11178 | 6056 |
|
||||
|5500000 | **-** | 32188 | 9854 |
|
||||
|8000000 | **-** | **-** | 12310 |
|
||||
|16000000 | **-** | **-** | 19950 |
|
||||
|29000000 | **-** | **-** | 32324 |
|
||||
|
||||
## Evaluation ICCV2021-MFR and IJB-C
|
||||
|
||||
More details see [eval.md](docs/eval.md) in docs.
|
||||
|
||||
## Test
|
||||
|
||||
We tested many versions of PyTorch. Please create an issue if you are having trouble.
|
||||
|
||||
- [x] torch 1.6.0
|
||||
- [x] torch 1.7.1
|
||||
- [x] torch 1.8.0
|
||||
- [x] torch 1.9.0
|
||||
|
||||
## Citation
|
||||
|
||||
```
|
||||
@inproceedings{deng2019arcface,
|
||||
title={Arcface: Additive angular margin loss for deep face recognition},
|
||||
author={Deng, Jiankang and Guo, Jia and Xue, Niannan and Zafeiriou, Stefanos},
|
||||
booktitle={Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition},
|
||||
pages={4690--4699},
|
||||
year={2019}
|
||||
}
|
||||
@inproceedings{an2020partical_fc,
|
||||
title={Partial FC: Training 10 Million Identities on a Single Machine},
|
||||
author={An, Xiang and Zhu, Xuhan and Xiao, Yang and Wu, Lan and Zhang, Ming and Gao, Yuan and Qin, Bin and
|
||||
Zhang, Debing and Fu Ying},
|
||||
booktitle={Arxiv 2010.05222},
|
||||
year={2020}
|
||||
}
|
||||
```
|
||||
@@ -0,0 +1,25 @@
|
||||
from .iresnet import iresnet18, iresnet34, iresnet50, iresnet100, iresnet200
|
||||
from .mobilefacenet import get_mbf
|
||||
|
||||
|
||||
def get_model(name, **kwargs):
|
||||
# resnet
|
||||
if name == "r18":
|
||||
return iresnet18(False, **kwargs)
|
||||
elif name == "r34":
|
||||
return iresnet34(False, **kwargs)
|
||||
elif name == "r50":
|
||||
return iresnet50(False, **kwargs)
|
||||
elif name == "r100":
|
||||
return iresnet100(False, **kwargs)
|
||||
elif name == "r200":
|
||||
return iresnet200(False, **kwargs)
|
||||
elif name == "r2060":
|
||||
from .iresnet2060 import iresnet2060
|
||||
return iresnet2060(False, **kwargs)
|
||||
elif name == "mbf":
|
||||
fp16 = kwargs.get("fp16", False)
|
||||
num_features = kwargs.get("num_features", 512)
|
||||
return get_mbf(fp16=fp16, num_features=num_features)
|
||||
else:
|
||||
raise ValueError()
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user