lesliezhoa / heser.pytorch

unofficial implementation of Few-Shot Head Swapping in the Wild

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

!!!!NEWS!!!!!

HeadSwap is Back. You can use it in colab. Try Now!!!

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

heser.pytorch's Issues

测试了一下，效果不好

weights

Hello LeslieZhoa, I checked another git you contributed :https://github.com/keshav1990/HeadSwap/blob/main/process/download_weight.sh and it's missing some weights related to this project, can you provide them?

https://github.com/LeslieZhoa/HeadSwap/releases/download/v0.0/sr_cf.onnx
https://github.com/LeslieZhoa/HeadSwap/releases/download/v0.0/Blender-401-00012900.pth

thank you so much

Weights?

Hi, can you post your weights?

How to start

Hello, I have only recently started to be interested in neural networks and it is not clear to me how to replace the head using your code. Can you write in more detail in which folders to drop the files 073-00000000.pth, 417-00000000.pth, parsing.pth, vgg19-d01eb7cb.pth?
In which folders should there be a photo source and a photo from which the head is taken?

About face parsing picture

Thanks for your work! I'd like to ask about the format of face parsing. The original format is from LVT with plot axis and colored. Is this format directly useable for training? And I save another format of gary images using cv2, is this more appropriate for training?