Skip to content

unofficial implementation of Few-Shot Head Swapping in the Wild

License

Notifications You must be signed in to change notification settings

JackZhouSz/HeSer.Pytorch

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

21 Commits
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

HeSer.Pytorch

unofficial implementation of Few-Shot Head Swapping in the Wild
you can find official version here
I did not use the discriminator from the paper and just follow DCT-NET

enviroment

  • torch
  • opencv-python
  • tensorboardX
  • imgaug
  • face-alignment
# download pretrain model
cd process
bash download_weight.sh

How to RUN

train

I only train one ID for driving

Data Process

  1. download voxceleb2
    a. I just download voxceleb2 test dataset, you can use this website
    b. You can unzip this file like this:
    +--- dataset
    |   +--- vox2_test_txt
    |   |   +--- txt
    |   |   |   +--- id00017
    |   |   |   |   +--- 01dfn2spqyE
    |   |   |   |   |   +--- 00001.txt
    |   |   |   |   +--- 5MkXgwdrmJw
    |   |   |   |   |   +--- 00002.txt
    |   |   |   |   +--- 7t6lfzvVaTM
    |   |   |   |   |   +--- 00003.txt
    |   |   |   |   |   +--- 00004.txt
    |   |   |   |   |   +--- 00005.txt
    |   |   |   |   |   +--- 00006.txt
    |   |   |   |   |   +--- 00007.txt
    
    
    c. Install yt-dlp and aria2c by yourself. I think you can do that through internet.
    cd process
    python download_and_process.py
    
    d. the dataset is like:
    voceleb2/
    |-- id00017
    |   |-- 01dfn2spqyE
    |   |   `-- 00.npy
    |   |-- 5MkXgwdrmJw
    |   |   |-- 00.npy
    |   |   `-- 5MkXgwdrmJw.mp4
    |   |-- 7t6lfzvVaTM
    |   |   |-- 00.npy
    |   |   |-- 01.npy
    |   |   |-- 02.npy
    |   |   |-- 03.npy
    |   |   |-- 04.npy
    |   |   |-- 05.npy
    |   |   |-- 06.npy
    |   |   |-- 07.npy
    |   |   |-- 08.npy
    |   |   |-- 09.npy
    |   |   `-- 7t6lfzvVaTM.mp4
    
  2. crop and align
    cd process 
    python process_raw_video.py
    
    the dataset is like:
    process/
    |-- img
    |   |-- id00017
    |   |   |-- 5MkXgwdrmJw-0000
    |   |   |   |-- 1273.png
    |   |   |   |-- 1274.png
    |   |   |   |-- 1275.png
    |   |   |   |-- 1276.png
    |   |   |   |-- 1277.png
    |   |   |   |-- 1278.png
    |   |   |   |-- 1279.png
    |   |   |   |-- 1280.png
    |   |   |   |-- 1281.png
    |   |   |   |-- 1282.png
    |   |   |   |-- 1283.png
    |   |   |   |-- 1284.png
    |   |   |   |-- 1285.png
    |   |   |   |-- 1286.png
    |   |   |   |-- 1287.png
    |   |   |   |-- 1288.png
    |   |   |   |-- 1289.png
    
  3. Remove data below threshold
    cd process
    python filter_idfiles.py
    
  4. face parsing
    follow LVT to get face parsing
    the mask data is like:
    process/mask/
    |-- id00017
    |   |-- 5MkXgwdrmJw-0000
    |   |   |-- 1273.png
    |   |   |-- 1274.png
    |   |   |-- 1275.png
    |   |   |-- 1276.png
    |   |   |-- 1277.png
    |   |   |-- 1278.png
    |   |   |-- 1279.png
    |   |   |-- 1280.png
    

Train Align

I just use id00061 to train align
check model/AlignModule/config.py to put your own path and params
for single gpu

python  train.py --model align --batch_size 8 --checkpoint_path checkpoint --lr 2e-4 --print_interval 100 --save_interval 100 --dist

for multi gpu

python -m torch.distributed.launch train.py --model align --batch_size 8 --checkpoint_path checkpoint --lr 2e-4 --print_interval 100 --save_interval 100

Train Blend

check model/BlendModule/config.py to put your own path and params
for single gpu

python  train.py --model blend --batch_size 8 --checkpoint_path checkpoint --lr 2e-4 --print_interval 100 --save_interval 100 --dist

for multi gpu

python -m torch.distributed.launch train.py --model blend --batch_size 8 --checkpoint_path checkpoint --lr 2e-4 --print_interval 100 --save_interval 100

Inference

follow inference.py, change your own model path and input images

python inference.py

Show

The result is just overfitting

Credits

DCT-NET model and implementation:
https://github.com/LeslieZhoa/DCT-NET.Pytorch Copyright © 2022, LeslieZhoa.
License https://github.com/LeslieZhoa/DCT-NET.Pytorch/blob/main/LICENSE

latent-pose-reenactment model and implementation:
https://github.com/shrubb/latent-pose-reenactment Copyright © 2020, shrubb.
License https://github.com/shrubb/latent-pose-reenactment/blob/master/LICENSE.txt

arcface pytorch model pytorch model and implementation:
https://github.com/ronghuaiyang/arcface-pytorch Copyright © 2018, ronghuaiyang.

LVT model and implementation:
https://github.com/LeslieZhoa/LVT Copyright © 2022, LeslieZhoa.

face-parsing model and implementation:
https://github.com/zllrunning/face-parsing.PyTorch Copyright © 2019, zllrunning.
License https://github.com/zllrunning/face-parsing.PyTorch/blob/master/LICENSE

About

unofficial implementation of Few-Shot Head Swapping in the Wild

Resources

License

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published