├── .gitignore ├── LICENSE ├── README.md ├── condor.sh ├── condor ├── config.sub ├── dataset.sh └── job.sh ├── configs ├── __init__.py ├── config.py └── mica.yml ├── data └── FLAME2020 │ ├── FLAME_masks │ ├── FLAME_masks.gif │ ├── FLAME_masks.pkl │ └── readme │ ├── head_template.obj │ └── landmark_embedding.npy ├── datasets ├── README.md ├── __init__.py ├── base.py ├── creation │ ├── __init__.py │ ├── generator.py │ ├── instances │ │ ├── __init__.py │ │ ├── bu3dfe.py │ │ ├── d3dfacs.py │ │ ├── facewarehouse.py │ │ ├── florence.py │ │ ├── frgc.py │ │ ├── instance.py │ │ ├── lyhm.py │ │ ├── pb4d.py │ │ └── stirling.py │ ├── main.py │ └── util.py └── image_paths │ ├── BP4D.npy │ ├── BU3DFE.npy │ ├── D3DFACS.npy │ ├── FACEWAREHOUSE.npy │ ├── FLORENCE.npy │ ├── FRGC.npy │ ├── LYHM.npy │ └── STIRLING.npy ├── demo.py ├── demo └── input │ ├── carell.jpg │ ├── connelly.jpg │ ├── justin.png │ └── lawrence.jpg ├── documents ├── BP4D.gif ├── D3DFACS.gif ├── FACEWAREHOUSE.gif ├── FLORENCE.gif ├── FRGC.gif ├── LYHM.gif ├── STIRLING.gif ├── teaser.jpg └── voxceleb.gif ├── environment.yml ├── install.sh ├── jobs.py ├── micalib ├── __init__.py ├── base_model.py ├── models │ ├── __init__.py │ └── mica.py ├── renderer.py ├── tester.py ├── trainer.py └── validator.py ├── models ├── __init__.py ├── arcface.py ├── flame.py ├── generator.py └── lbs.py ├── render_dataset.py ├── test.py ├── testing ├── now │ ├── now.py │ └── template.sh └── stirling │ ├── stirling.py │ └── template.sh ├── train.py └── utils ├── __init__.py ├── best_model.py ├── landmark_detector.py ├── masking.py └── util.py /.gitignore: -------------------------------------------------------------------------------- 1 | # Compiled source # 2 | ################### 3 | .idea 4 | 5 | datasets/creation/template/* 6 | statistics/* 7 | *.o 8 | *.so 9 | 10 | # Packages # 11 | ############ 12 | # it's better to unpack these files and commit the raw source 13 | # git has its own built in compression methods 14 | *.7z 15 | *.dmg 16 | *.gz 17 | *.iso 18 | *.jar 19 | *.rar 20 | *.tar 21 | *.zip 22 | 23 | # OS generated files # 24 | ###################### 25 | .DS_Store 26 | .DS_Store? 27 | ._* 28 | .Spotlight-V100 29 | .Trashes 30 | ehthumbs.db 31 | Thumbs.db 32 | 33 | # 3D data # 34 | ############ 35 | *.mat 36 | *.obj 37 | *.dat 38 | *.npz 39 | *.pkl 40 | 41 | # python file # 42 | ############ 43 | *.pyc 44 | __pycache__ 45 | 46 | ## deca data 47 | data/FLAME2020/generic_model.pkl 48 | data/FLAME2020/female_model.pkl 49 | data/FLAME2020/male_model.pkl 50 | data/FLAME2020/FLAME_albedo_from_BFM.npz 51 | results 52 | output 53 | TestSamples 54 | 55 | ## dump files 56 | __dump 57 | 58 | ## visual code files 59 | .vscode 60 | render_dataset.py 61 | shapes.pt 62 | partial 63 | images 64 | *.pt 65 | testing/now/jobs 66 | testing/now/logs 67 | testing/stirling/logs 68 | testing/stirling/jobs 69 | 70 | demo/arcface 71 | demo/output 72 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | License 2 | 3 | Software Copyright License for non-commercial scientific research purposes 4 | Please read carefully the following terms and conditions and any accompanying documentation before you download 5 | and/or use the MICA model, data and software, (the "Model & Software"), including 3D meshes, software, and scripts. 6 | By downloading and/or using the Model & Software (including downloading, cloning, installing, and any other use 7 | of this github repository), you acknowledge that you have read these terms and conditions, understand them, and 8 | agree to be bound by them. If you do not agree with these terms and conditions, you must not download and/or use 9 | the Model & Software. Any infringement of the terms of this agreement will automatically terminate your rights 10 | under this License 11 | 12 | Ownership / Licensees 13 | The Model & Software and the associated materials has been developed at the 14 | Max Planck Institute for Intelligent Systems (hereinafter "MPI"). 15 | 16 | Any copyright or patent right is owned by and proprietary material of the 17 | Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (hereinafter “MPG”; MPI and MPG hereinafter 18 | collectively “Max-Planck”) hereinafter the “Licensor”. 19 | 20 | License Grant 21 | Licensor grants you (Licensee) personally a single-user, non-exclusive, non-transferable, free of charge right: 22 | 23 | • To install the Model & Software on computers owned, leased or otherwise controlled by you and/or your organization. 24 | • To use the Model & Software for the sole purpose of performing peaceful non-commercial scientific research, 25 | non-commercial education, or non-commercial artistic projects. 26 | 27 | Any other use, in particular any use for commercial, pornographic, military, or surveillance purposes is prohibited. 28 | This includes, without limitation, incorporation in a commercial product, use in a commercial service, 29 | or production of other artefacts for commercial purposes. 30 | 31 | The Model & Software may not be used to create fake, libelous, misleading, or defamatory content of any kind, excluding 32 | analyses in peer-reviewed scientific research. 33 | 34 | The Model & Software may not be reproduced, modified and/or made available in any form to any third party 35 | without Max-Planck’s prior written permission. 36 | 37 | The Model & Software may not be used for pornographic purposes or to generate pornographic material whether 38 | commercial or not. This license also prohibits the use of the Model & Software to train methods/algorithms/neural 39 | networks/etc. for commercial use of any kind. By downloading the Model & Software, you agree not to reverse engineer it. 40 | 41 | No Distribution 42 | The Model & Software and the license herein granted shall not be copied, shared, distributed, re-sold, offered 43 | for re-sale, transferred or sub-licensed in whole or in part except that you may make one copy for archive 44 | purposes only. 45 | 46 | Disclaimer of Representations and Warranties 47 | You expressly acknowledge and agree that the Model & Software results from basic research, is provided “AS IS”, 48 | may contain errors, and that any use of the Model & Software is at your sole risk. 49 | LICENSOR MAKES NO REPRESENTATIONS 50 | OR WARRANTIES OF ANY KIND CONCERNING THE MODEL & SOFTWARE, NEITHER EXPRESS NOR IMPLIED, AND THE ABSENCE OF ANY 51 | LEGAL OR ACTUAL DEFECTS, WHETHER DISCOVERABLE OR NOT. Specifically, and not to limit the foregoing, licensor 52 | makes no representations or warranties (i) regarding the merchantability or fitness for a particular purpose of 53 | the Model & Software, (ii) that the use of the Model & Software will not infringe any patents, copyrights or other 54 | intellectual property rights of a third party, and (iii) that the use of the Model & Software will not cause any 55 | damage of any kind to you or a third party. 56 | 57 | Limitation of Liability 58 | Because this Model & Software License Agreement qualifies as a donation, according to Section 521 of the German 59 | Civil Code (Bürgerliches Gesetzbuch – BGB) Licensor as a donor is liable for intent and gross negligence only. 60 | If the Licensor fraudulently conceals a legal or material defect, they are obliged to compensate the Licensee 61 | for the resulting damage. 62 | 63 | Licensor shall be liable for loss of data only up to the amount of typical recovery costs which would have 64 | arisen had proper and regular data backup measures been taken. For the avoidance of doubt Licensor shall be 65 | liable in accordance with the German Product Liability Act in the event of product liability. The foregoing 66 | applies also to Licensor’s legal representatives or assistants in performance. Any further liability shall 67 | be excluded. Patent claims generated through the usage of the Model & Software cannot be directed towards the copyright holders. 68 | The Model & Software is provided in the state of development the licensor defines. If modified or extended by 69 | Licensee, the Licensor makes no claims about the fitness of the Model & Software and is not responsible 70 | for any problems such modifications cause. 71 | 72 | No Maintenance Services 73 | You understand and agree that Licensor is under no obligation to provide either maintenance services, 74 | update services, notices of latent defects, or corrections of defects with regard to the Model & Software. 75 | Licensor nevertheless reserves the right to update, modify, or discontinue the Model & Software at any time. 76 | 77 | Defects of the Model & Software must be notified in writing to the Licensor with a comprehensible description 78 | of the error symptoms. The notification of the defect should enable the reproduction of the error. 79 | The Licensee is encouraged to communicate any use, results, modification or publication. 80 | 81 | Publications using the Model & Software 82 | You acknowledge that the Model & Software is a valuable scientific resource and agree to appropriately reference 83 | the following paper in any publication making use of the Model & Software. 84 | 85 | Commercial licensing opportunities 86 | For commercial uses of the Model & Software, please send email to justus.thies@tuebingen.mpg.de 87 | 88 | This Agreement shall be governed by the laws of the Federal Republic of Germany except for the UN Sales Convention. 89 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 |

MICA - Towards Metrical Reconstruction of Human Faces

2 | 3 |

Wojciech Zielonka, Timo Bolkart, Justus Thies

4 | 5 |
Max Planck Institute for Intelligent Systems, Tübingen, Germany
6 | 7 |

8 | Video  9 | Paper  10 | Project Website  11 | Face Tracker  12 | Dataset  13 | Supplemental  14 | Email 15 |

16 | 17 |
18 | 19 | Official Repository for ECCV 2022 paper Towards Metrical Reconstruction of Human Faces 20 |
21 |
22 | 23 |
24 | ⚠ The face tracker is now available under Metrical Photometric Tracker  ⚠ 25 |
26 | 27 | ### Installation 28 | 29 | After cloning the repository please install the environment by using attached conda `environment.yml` file with the command 30 | ``conda env create -f environment.yml``. Additionally, the FLAME2020 model is needed. To obtain it please create an account at the [website](https://flame.is.tue.mpg.de/) download the model and place it in the `/data/pretrained/FLAME2020/` folder. 31 | 32 | You can also simply run the `install.sh` script: 33 | 34 | ```shell 35 | git clone https://github.com/Zielon/MICA.git 36 | cd MICA 37 | ./install.sh 38 | ``` 39 | you will be asked to provide `{flame_user}` and `{flame_password}` for your FLAME account in order to access the file server. 40 | 41 | ### Pre-trained Models 42 | 43 | If you decide to not use the installation script, the pretrained model can be found under the [link](https://drive.google.com/file/d/1bYsI_spptzyuFmfLYqYkcJA6GZWZViNt/view?usp=sharing). After downloading, please place it in the `/data/pretrained/mica.tar` location. Additionally, you will need to provide models for `inisghtface`: 44 | 1) [antelopev2](https://drive.google.com/file/d/16PWKI_RjjbE4_kqpElG-YFqe8FpXjads/view?usp=sharing) 45 | 2) [buffalo_l](https://drive.google.com/file/d/1navJMy0DTr1_DHjLWu1i48owCPvXWfYc/view?usp=sharing) 46 | 47 | then you need to unzip them and place in `~/.insightface/models/`. The `install.sh` script does it for you. 48 | 49 | ### How To Use 50 | 51 | To use MICA you can simply run the `demo.py` file. It will process all the images from `demo/input/` folder and create the output destination for each subject with `.ply` mesh, rendered image, and `.npy` FLAME parameters. 52 | 53 | ### Dataset and Training 54 | 55 | The MICA dataset consists of eight smaller datasets for about 2300 subjects under a common FLAME topology. Read more information about how to obtain and use it under the [link](https://github.com/Zielon/MICA/tree/master/datasets/). To train MICA the images from all eight datasets are needed. The repository contains scripts how to generate the Arcface input images as well as the complete list of all the images used for the training. More information can be found [here](https://github.com/Zielon/MICA/tree/master/datasets). 56 | 57 | When you train from scratch for Arcface model initialization please download [Glint360K](https://github.com/deepinsight/insightface/tree/master/recognition/arcface_torch) and specify the path to it in the config as `cfg.model.arcface_pretrained_model`. 58 | 59 | ### Testing 60 | 61 | The testing was done using two datasets, [Stirling](http://pics.stir.ac.uk/ESRC/) and [NoW](https://now.is.tue.mpg.de/). In the [model folder](https://github.com/Zielon/MICA/tree/master/models) you can find the corresponding scripts to run testing routine, which generates the meshes. To calculate the NoW challenge error you can use the following [repository](https://github.com/soubhiksanyal/now_evaluation). 62 | 63 | ### Citation 64 | If you use this project in your research please cite MICA: 65 | ```bibtex 66 | @proceedings{zielonka22mica, 67 | author = {Zielonka, Wojciech and Bolkart, Timo and Thies, Justus}, 68 | title = {Towards Metrical Reconstruction of Human Faces}, 69 | journal = {European Conference on Computer Vision}, 70 | year = {2022} 71 | } 72 | ``` 73 | -------------------------------------------------------------------------------- /condor.sh: -------------------------------------------------------------------------------- 1 | #!/bin/sh 2 | 3 | # bash condor.sh 100 ./configs/mica.yml 1 4 | 5 | # default parameters 6 | BID=3 7 | CONFIG=./configs/mica.yml 8 | NODE_CONFIG=condor/config.sub 9 | NODE_SCRIPT=./condor/job.sh 10 | GPUS=1 11 | GPU_TYPE=0 12 | 13 | # set parameters 14 | if [ -n "$1" ]; then BID=${1}; fi 15 | if [ -n "$2" ]; then CONFIG=${2}; fi 16 | if [ -n "$3" ]; then GPU_TYPE=${3}; fi 17 | if [ -n "$4" ]; then GPUS=${4}; fi 18 | if [ -n "$5" ]; then NODE_CONFIG=${5}; fi 19 | if [ -n "$6" ]; then NODE_SCRIPT=${6}; fi 20 | 21 | mkdir -p output/condor_logs 22 | cp -nf ${NODE_CONFIG}{,.bak} 23 | 24 | GPU_NAME=Error 25 | 26 | if [ $GPU_TYPE -eq 0 ]; then GPU_NAME='Quadro RTX 6000'; fi 27 | if [ $GPU_TYPE -eq 1 ]; then GPU_NAME='Tesla V100-SXM2-32GB'; fi 28 | if [ $GPU_TYPE -eq 2 ]; then GPU_NAME='NVIDIA GeForce RTX 2080 Ti'; fi 29 | 30 | NAME=$(basename ${CONFIG} .yml) 31 | sed -i "s/{errorfile}/${NAME}/" ${NODE_CONFIG}.bak 32 | sed -i "s/{outfile}/${NAME}/" ${NODE_CONFIG}.bak 33 | sed -i "s/{logfile}/${NAME}/" ${NODE_CONFIG}.bak 34 | sed -i "s/{gpus}/${GPUS}/" ${NODE_CONFIG}.bak 35 | sed -i "s/{gpu_name}/${GPU_NAME}/" ${NODE_CONFIG}.bak 36 | 37 | # start node and execute script 38 | echo 'Executing:' ${NODE_SCRIPT} ${CONFIG} 39 | echo '# BID:' ${BID} 40 | echo '# GPUS:' ${GPUS} 41 | echo '# GPU NAME:' ${GPU_NAME} 42 | 43 | condor_submit_bid ${BID} ${NODE_CONFIG}.bak -append "arguments = ${NODE_SCRIPT} ${CONFIG}" 44 | rm ${NODE_CONFIG}.bak 45 | -------------------------------------------------------------------------------- /condor/config.sub: -------------------------------------------------------------------------------- 1 | executable = /bin/bash 2 | error = ./output/condor_logs/{errorfile}_$(ClusterId).$(ProcId).err 3 | output = ./output/condor_logs/{outfile}_$(ClusterId).$(ProcId).out 4 | log = ./output/condor_logs/{logfile}_$(ClusterId).$(ProcId).log 5 | request_memory = 32768 6 | request_cpus = 6 7 | request_gpus = {gpus} 8 | +WantGPUStats = true 9 | requirements = (TARGET.CUDADeviceName=="{gpu_name}") 10 | # EXIT SETTINGS 11 | on_exit_hold = (ExitCode =?= 3) 12 | on_exit_hold_reason = "Checkpointed, will resume" 13 | on_exit_hold_subcode = 2 14 | periodic_release = ( (JobStatus =?= 5) && (HoldReasonCode =?= 3) && (HoldReasonSubCode =?= 2) ) 15 | +RunningPriceExceededAction = "kill" 16 | queue 17 | -------------------------------------------------------------------------------- /condor/dataset.sh: -------------------------------------------------------------------------------- 1 | #!/bin/sh 2 | 3 | CONFIG=${1} 4 | 5 | PYTHON_ENV=/home/wzielonka/miniconda3/etc/profile.d/conda.sh 6 | SCRIPT_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )" 7 | export PATH=/usr/local/bin:/usr/bin:/bin:/usr/sbin:$PATH 8 | export PYTHONPATH="${PYTHONPATH}:/home/wzielonka/projects/OnFlame-internal/" 9 | 10 | echo 'START JOB (dataset generation)' 11 | 12 | module load cuda/10.1 13 | module load gcc/4.9 14 | 15 | echo 'ACTIVATE MICA' 16 | source ${PYTHON_ENV} 17 | conda activate MICA 18 | 19 | echo 'RUN SCRIPT' 20 | cd ${SCRIPT_DIR}/../datasets/creation 21 | python ./main.py -------------------------------------------------------------------------------- /condor/job.sh: -------------------------------------------------------------------------------- 1 | #!/bin/sh 2 | 3 | CONFIG=${1} 4 | 5 | PYTHON_ENV=/home/wzielonka/miniconda3/etc/profile.d/conda.sh 6 | SCRIPT_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )" 7 | export PATH=/usr/local/bin:/usr/bin:/bin:/usr/sbin:$PATH 8 | 9 | echo 'START JOB (MICA training)' 10 | 11 | module load cuda/10.1 12 | module load gcc/4.9 13 | 14 | echo 'ACTIVATE MICA' 15 | source ${PYTHON_ENV} 16 | conda activate MICA 17 | 18 | echo 'RUN SCRIPT' 19 | #echo 'ScriptDir' ${SCRIPT_DIR} 20 | echo 'CONFIG: ' ${CONFIG} 21 | cd ${SCRIPT_DIR}/.. 22 | python ./train.py --cfg ${CONFIG} -------------------------------------------------------------------------------- /configs/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Zielon/MICA/af22e7a5810d474bc28a1433db533723d6bd2b07/configs/__init__.py -------------------------------------------------------------------------------- /configs/config.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2023 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: mica@tue.mpg.de 16 | 17 | 18 | import argparse 19 | import os 20 | 21 | from yacs.config import CfgNode as CN 22 | 23 | cfg = CN() 24 | 25 | abs_mica_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), '..')) 26 | cfg.mica_dir = abs_mica_dir 27 | cfg.device = 'cuda' 28 | cfg.device_id = '0' 29 | cfg.pretrained_model_path = os.path.join(cfg.mica_dir, 'data/pretrained', 'mica.tar') 30 | cfg.output_dir = '' 31 | 32 | # ---------------------------------------------------------------------------- # 33 | # Options for Face model 34 | # ---------------------------------------------------------------------------- # 35 | cfg.model = CN() 36 | cfg.model.testing = False 37 | cfg.model.name = 'mica' 38 | 39 | cfg.model.topology_path = os.path.join(cfg.mica_dir, 'data/FLAME2020', 'head_template.obj') 40 | cfg.model.flame_model_path = os.path.join(cfg.mica_dir, 'data/FLAME2020', 'generic_model.pkl') 41 | cfg.model.flame_lmk_embedding_path = os.path.join(cfg.mica_dir, 'data/FLAME2020', 'landmark_embedding.npy') 42 | cfg.model.n_shape = 300 43 | cfg.model.layers = 8 44 | cfg.model.hidden_layers_size = 256 45 | cfg.model.mapping_layers = 3 46 | cfg.model.use_pretrained = True 47 | cfg.model.arcface_pretrained_model = '/scratch/is-rg-ncs/models_weights/arcface-torch/backbone100.pth' 48 | 49 | # ---------------------------------------------------------------------------- # 50 | # Options for Dataset 51 | # ---------------------------------------------------------------------------- # 52 | cfg.dataset = CN() 53 | cfg.dataset.training_data = ['LYHM'] 54 | cfg.dataset.eval_data = ['FLORENCE'] 55 | cfg.dataset.batch_size = 2 56 | cfg.dataset.K = 4 57 | cfg.dataset.n_train = 100000 58 | cfg.dataset.num_workers = 4 59 | cfg.dataset.root = '/datasets/MICA/' 60 | 61 | # ---------------------------------------------------------------------------- # 62 | # Mask weights 63 | # ---------------------------------------------------------------------------- # 64 | cfg.mask_weights = CN() 65 | cfg.mask_weights.face = 150.0 66 | cfg.mask_weights.nose = 50.0 67 | cfg.mask_weights.lips = 50.0 68 | cfg.mask_weights.forehead = 50.0 69 | cfg.mask_weights.lr_eye_region = 50.0 70 | cfg.mask_weights.eye_region = 50.0 71 | 72 | cfg.mask_weights.whole = 1.0 73 | cfg.mask_weights.ears = 0.01 74 | cfg.mask_weights.eyes = 0.01 75 | 76 | cfg.running_average = 7 77 | 78 | # ---------------------------------------------------------------------------- # 79 | # Options for training 80 | # ---------------------------------------------------------------------------- # 81 | cfg.train = CN() 82 | cfg.train.use_mask = False 83 | cfg.train.max_epochs = 50 84 | cfg.train.max_steps = 100000 85 | cfg.train.lr = 1e-4 86 | cfg.train.arcface_lr = 1e-3 87 | cfg.train.weight_decay = 0.0 88 | cfg.train.lr_update_step = 100000000 89 | cfg.train.log_dir = 'logs' 90 | cfg.train.log_steps = 10 91 | cfg.train.vis_dir = 'train_images' 92 | cfg.train.vis_steps = 200 93 | cfg.train.write_summary = True 94 | cfg.train.checkpoint_steps = 1000 95 | cfg.train.checkpoint_epochs_steps = 2 96 | cfg.train.val_steps = 1000 97 | cfg.train.val_vis_dir = 'val_images' 98 | cfg.train.eval_steps = 5000 99 | cfg.train.reset_optimizer = False 100 | cfg.train.val_save_img = 5000 101 | cfg.test_dataset = 'now' 102 | 103 | 104 | def get_cfg_defaults(): 105 | return cfg.clone() 106 | 107 | 108 | def update_cfg(cfg, cfg_file): 109 | cfg.merge_from_file(cfg_file) 110 | return cfg.clone() 111 | 112 | 113 | def parse_args(): 114 | parser = argparse.ArgumentParser() 115 | parser.add_argument('--cfg', type=str, help='cfg file path', required=True) 116 | parser.add_argument('--test_dataset', type=str, help='Test dataset type', default='') 117 | parser.add_argument('--checkpoint', type=str, help='Checkpoint to load', default='') 118 | 119 | args = parser.parse_args() 120 | print(args, end='\n\n') 121 | 122 | cfg = get_cfg_defaults() 123 | if args.cfg is not None: 124 | cfg_file = args.cfg 125 | cfg = update_cfg(cfg, args.cfg) 126 | cfg.cfg_file = cfg_file 127 | 128 | return cfg, args 129 | -------------------------------------------------------------------------------- /configs/mica.yml: -------------------------------------------------------------------------------- 1 | # Mica config 2 | 3 | pretrained_model_path: '' 4 | 5 | dataset: 6 | root: '/scratch/NFC/MICA/dataset/' 7 | training_data: [ 'LYHM', 'D3DFACS', 'BU3DFE', 'FRGC', 'Stirling', 'FaceWarehouse', 'BP4D' ] 8 | eval_data: [ 'FLORENCE' ] 9 | num_workers: 4 10 | batch_size: 8 11 | K: 2 12 | 13 | train: 14 | lr: 1e-5 15 | arcface_lr: 1e-5 16 | weight_decay: 2e-4 17 | use_mask: True 18 | reset_optimizer: False 19 | max_steps: 160000 20 | log_steps: 50 21 | val_steps: 300 22 | vis_steps: 1200 23 | val_save_img: 1200 24 | checkpoint_steps: 1000 25 | checkpoint_epochs_steps: 10000 26 | 27 | model: 28 | use_pretrained: False 29 | n_shape: 300 30 | name: 'mica' 31 | -------------------------------------------------------------------------------- /data/FLAME2020/FLAME_masks/FLAME_masks.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Zielon/MICA/af22e7a5810d474bc28a1433db533723d6bd2b07/data/FLAME2020/FLAME_masks/FLAME_masks.gif -------------------------------------------------------------------------------- /data/FLAME2020/FLAME_masks/readme: -------------------------------------------------------------------------------- 1 | Dictionary with vertex indices for different masks for the publicly available FLAME head model (https://flame.is.tue.mpg.de/). 2 | See the gif for a visualization of all masks. 3 | -------------------------------------------------------------------------------- /data/FLAME2020/landmark_embedding.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Zielon/MICA/af22e7a5810d474bc28a1433db533723d6bd2b07/data/FLAME2020/landmark_embedding.npy -------------------------------------------------------------------------------- /datasets/README.md: -------------------------------------------------------------------------------- 1 |

MICA - Dataset

2 | 3 | The MICA dataset consists of eight smaller datasets for about 2315 subjects, built by unifying existing small- and medium-scale datasets under a common FLAME topology. It consists of shape geometry only, therefore, to obtain images for each subject please refer to the primary dataset. 4 | 5 | This dataset contains registration meshes together with corresponding fitted FLAME parameters. Actors are split between individual folders with a unique identifier based on the original dataset. The folder name of the parameters and mesh is the same as the ones in the analogous dataset with images. 6 | 7 | To obtain the dataset please follow each link separately and request the given subset. 8 | 9 | In the case of any questions feel free to email us. 10 | 11 |
12 | 13 | 14 | 15 | 16 | 17 | 18 | 19 | 20 | 21 | 22 | 23 | 24 | 25 | 26 | 27 | 28 | 29 | 30 | 31 | 32 | 33 | 34 | 35 | 36 | 37 | 38 | 39 | 40 | 41 | 42 | 43 | 44 | 45 | 46 | 47 | 48 | 49 | 50 | 51 | 52 | 53 | 54 | 55 | 56 | 57 | 58 | 59 | 61 | 62 | 63 | 64 | 65 | 66 | 67 | 68 |
FLAME RenderingsDatasetSubjectsSource
Stirling133pics.stir.ac.uk
D3DFACS10cs.bath.ac.uk
Florence 2D/3D53micc.unifi.it
LYHM1211cs.york.ac.uk
FaceWarehouse150kunzhou.net
*FRGC531cvrl.nd.edu 60 |
*BP4D+127cs.binghamton.edu
69 |
70 | 71 | Each subset zip file has the following structure: 72 | ```shell 73 | root\ 74 | FLAME_parameters\ 75 | actor_id\ 76 | *.npz 77 | registrations\ 78 | actor_id\ 79 | *.obj 80 | ``` 81 | 82 | To retrieve FLAME2020 parameters you can simply do: 83 | ```python 84 | import numpy as np 85 | import torch 86 | 87 | params = np.load('path.npz', allow_pickle=True) 88 | pose = torch.tensor(params['pose']).float() 89 | betas = torch.tensor(params['betas']).float() 90 | 91 | flame = { 92 | 'shape_params': betas[:300], 93 | 'expression_params': betas[300:], 94 | 'pose_params': torch.cat([pose[:3], pose[6:9]]), 95 | } 96 | ``` 97 | 98 | ### MICA Training Dataset Preparation 99 | 100 | To prepare the MICA training dataset you can follow the scripts from the [creation](https://github.com/Zielon/MICA/tree/master/datasets/creation) folder. Additionally, the complete list of images used for the training can be found in [image_paths](https://github.com/Zielon/MICA/tree/master/datasets/image_paths) folder. It contains the name of the FLAME parameters file `.npz` and a list of all images used for the training with their original name stored as dictionary. 101 | 102 | ### Licence 103 | This dataset is for academic, non-commercial usage only. Moreover, it is an extension of already existing datasets, therefore, the license is shared and applies equivalently to both, the original and the corresponding derived one. Please read the original license of each original dataset for more information, especially in the context of data privacy. 104 | 105 | ### Additional Information 106 | 107 | Please note that some subsets do not have an identical amount of subjects due to registration errors. 108 | 109 | #### BP4D+ Subset 110 | 111 | The original dataset contains 140 subjects from where we selected scans in neutral pose and successfully registered 127 FLAME meshes for them. 112 | 113 | #### FRGC Subset 114 | 115 | 1) A group desiring to obtain FLAME results on the FRGC data must license the FRGC 2.0 data set from CVRL using the licensing procedure at the website: http://cvrl.nd.edu . 116 | 2) Once the license form is completed and CVRL has approved it, an email will be sent from Globus indicating that you have been authorized to retrieve FRGC 2.0. 117 | 3) That authorization will serve as authorization to receive any derivative work, therefore, please send a confirmation to mica [AT] tue.mpg.de address to receive the MICA dataset. 118 | 119 | ### Citation 120 | If you use this dataset in your research please cite MICA: 121 | ```bibtex 122 | @proceedings{MICA:ECCV2022, 123 | author = {Zielonka, Wojciech and Bolkart, Timo and Thies, Justus}, 124 | title = {Towards Metrical Reconstruction of Human Faces}, 125 | journal = {European Conference on Computer Vision}, 126 | year = {2022} 127 | } 128 | ``` 129 | -------------------------------------------------------------------------------- /datasets/__init__.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2023 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: mica@tue.mpg.de 16 | 17 | 18 | import numpy as np 19 | from torch.utils.data import ConcatDataset 20 | 21 | from datasets.base import BaseDataset 22 | 23 | 24 | def build_train(config, device): 25 | data_list = [] 26 | total_images = 0 27 | for dataset in config.training_data: 28 | dataset_name = dataset.upper() 29 | config.n_train = np.Inf 30 | if type(dataset) is list: 31 | dataset_name, n_train = dataset 32 | config.n_train = n_train 33 | 34 | dataset = BaseDataset(name=dataset_name, config=config, device=device, isEval=False) 35 | data_list.append(dataset) 36 | total_images += dataset.total_images 37 | 38 | return ConcatDataset(data_list), total_images 39 | 40 | 41 | def build_val(config, device): 42 | data_list = [] 43 | total_images = 0 44 | for dataset in config.eval_data: 45 | dataset_name = dataset.upper() 46 | config.n_train = np.Inf 47 | if type(dataset) is list: 48 | dataset_name, n_train = dataset 49 | config.n_train = n_train 50 | 51 | dataset = BaseDataset(name=dataset_name, config=config, device=device, isEval=True) 52 | data_list.append(dataset) 53 | total_images += dataset.total_images 54 | 55 | return ConcatDataset(data_list), total_images 56 | -------------------------------------------------------------------------------- /datasets/base.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2023 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: mica@tue.mpg.de 16 | 17 | 18 | import os 19 | import re 20 | from abc import ABC 21 | from functools import reduce 22 | from pathlib import Path 23 | 24 | import loguru 25 | import numpy as np 26 | import torch 27 | from loguru import logger 28 | from skimage.io import imread 29 | from torch.utils.data import Dataset 30 | from torchvision import transforms 31 | 32 | 33 | class BaseDataset(Dataset, ABC): 34 | def __init__(self, name, config, device, isEval): 35 | self.K = config.K 36 | self.isEval = isEval 37 | self.n_train = np.Inf 38 | self.imagepaths = [] 39 | self.face_dict = {} 40 | self.name = name 41 | self.device = device 42 | self.min_max_K = 0 43 | self.cluster = False 44 | self.dataset_root = config.root 45 | self.total_images = 0 46 | self.image_folder = 'arcface_input' 47 | self.flame_folder = 'FLAME_parameters' 48 | self.initialize() 49 | 50 | def initialize(self): 51 | logger.info(f'[{self.name}] Initialization') 52 | image_list = f'{os.path.abspath(os.path.dirname(__file__))}/image_paths/{self.name}.npy' 53 | logger.info(f'[{self.name}] Load cached file list: ' + image_list) 54 | self.face_dict = np.load(image_list, allow_pickle=True).item() 55 | self.imagepaths = list(self.face_dict.keys()) 56 | logger.info(f'[Dataset {self.name}] Total {len(self.imagepaths)} actors loaded!') 57 | self.set_smallest_k() 58 | 59 | def set_smallest_k(self): 60 | self.min_max_K = np.Inf 61 | max_min_k = -np.Inf 62 | for key in self.face_dict.keys(): 63 | length = len(self.face_dict[key][0]) 64 | if length < self.min_max_K: 65 | self.min_max_K = length 66 | if length > max_min_k: 67 | max_min_k = length 68 | 69 | self.total_images = reduce(lambda k, l: l + k, map(lambda e: len(self.face_dict[e][0]), self.imagepaths)) 70 | loguru.logger.info(f'Dataset {self.name} with min K = {self.min_max_K} max K = {max_min_k} length = {len(self.face_dict)} total images = {self.total_images}') 71 | return self.min_max_K 72 | 73 | def compose_transforms(self, *args): 74 | self.transforms = transforms.Compose([t for t in args]) 75 | 76 | def get_arcface_path(self, image_path): 77 | return re.sub('png|jpg', 'npy', str(image_path)) 78 | 79 | def __len__(self): 80 | return len(self.imagepaths) 81 | 82 | def __getitem__(self, index): 83 | actor = self.imagepaths[index] 84 | images, params_path = self.face_dict[actor] 85 | images = [Path(self.dataset_root, self.name, self.image_folder, path) for path in images] 86 | sample_list = np.array(np.random.choice(range(len(images)), size=self.K, replace=False)) 87 | 88 | K = self.K 89 | if self.isEval: 90 | K = max(0, min(200, self.min_max_K)) 91 | sample_list = np.array(range(len(images))[:K]) 92 | 93 | params = np.load(os.path.join(self.dataset_root, self.name, self.flame_folder, params_path), allow_pickle=True) 94 | pose = torch.tensor(params['pose']).float() 95 | betas = torch.tensor(params['betas']).float() 96 | 97 | flame = { 98 | 'shape_params': torch.cat(K * [betas[:300][None]], dim=0), 99 | 'expression_params': torch.cat(K * [betas[300:][None]], dim=0), 100 | 'pose_params': torch.cat(K * [torch.cat([pose[:3], pose[6:9]])[None]], dim=0), 101 | } 102 | 103 | images_list = [] 104 | arcface_list = [] 105 | 106 | for i in sample_list: 107 | image_path = images[i] 108 | image = np.array(imread(image_path)) 109 | image = image / 255. 110 | image = image.transpose(2, 0, 1) 111 | arcface_image = np.load(self.get_arcface_path(image_path), allow_pickle=True) 112 | 113 | images_list.append(image) 114 | arcface_list.append(torch.tensor(arcface_image)) 115 | 116 | images_array = torch.from_numpy(np.array(images_list)).float() 117 | arcface_array = torch.stack(arcface_list).float() 118 | 119 | return { 120 | 'image': images_array, 121 | 'arcface': arcface_array, 122 | 'imagename': actor, 123 | 'dataset': self.name, 124 | 'flame': flame, 125 | } 126 | -------------------------------------------------------------------------------- /datasets/creation/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Zielon/MICA/af22e7a5810d474bc28a1433db533723d6bd2b07/datasets/creation/__init__.py -------------------------------------------------------------------------------- /datasets/creation/generator.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2023 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: mica@tue.mpg.de 16 | 17 | 18 | import os 19 | from glob import glob 20 | from multiprocessing import Pool 21 | from pathlib import Path 22 | from typing import List 23 | 24 | import cv2 25 | import numpy as np 26 | from insightface.app import FaceAnalysis 27 | from insightface.app.common import Face 28 | from insightface.utils import face_align 29 | from loguru import logger 30 | from tqdm import tqdm 31 | 32 | from datasets.creation.instances.instance import Instance 33 | from datasets.creation.util import get_image, get_center, get_arcface_input 34 | 35 | 36 | def _transfer(src, dst): 37 | src.parent.mkdir(parents=True, exist_ok=True) 38 | dst.parent.mkdir(parents=True, exist_ok=True) 39 | os.system(f'cp {str(src)} {str(dst)}') 40 | 41 | 42 | def _copy(payload): 43 | instance, func, target, transform_path = payload 44 | files = func() 45 | for actor in files.keys(): 46 | for file in files[actor]: 47 | _transfer(Path(file), Path(instance.get_dst(), target, actor, transform_path(file))) 48 | 49 | 50 | class Generator: 51 | def __init__(self, instances): 52 | self.instances: List[Instance] = instances 53 | self.ARCFACE = 'arcface_input' 54 | 55 | def copy(self): 56 | logger.info('Start copying...') 57 | for instance in tqdm(self.instances): 58 | payloads = [(instance, instance.get_images, 'images', instance.transform_path)] 59 | with Pool(processes=len(payloads)) as pool: 60 | for _ in tqdm(pool.imap_unordered(_copy, payloads), total=len(payloads)): 61 | pass 62 | 63 | def preprocess(self): 64 | logger.info('Start preprocessing...') 65 | for instance in tqdm(self.instances): 66 | instance.preprocess() 67 | 68 | def arcface(self): 69 | app = FaceAnalysis(name='antelopev2', providers=['CUDAExecutionProvider']) 70 | app.prepare(ctx_id=0, det_size=(224, 224)) 71 | 72 | logger.info('Start arcface...') 73 | for instance in tqdm(self.instances): 74 | src = instance.get_dst() 75 | for image_path in tqdm(sorted(glob(f'{src}/images/*/*'))): 76 | dst = image_path.replace('images', self.ARCFACE) 77 | Path(dst).parent.mkdir(exist_ok=True, parents=True) 78 | for img in instance.transform_image(get_image(image_path[0:-4])): 79 | bboxes, kpss = app.det_model.detect(img, max_num=0, metric='default') 80 | if bboxes.shape[0] == 0: 81 | continue 82 | i = get_center(bboxes, img) 83 | bbox = bboxes[i, 0:4] 84 | det_score = bboxes[i, 4] 85 | if det_score < instance.get_min_det_score(): 86 | continue 87 | kps = None 88 | if kpss is not None: 89 | kps = kpss[i] 90 | face = Face(bbox=bbox, kps=kps, det_score=det_score) 91 | blob, aimg = get_arcface_input(face, img) 92 | np.save(dst[0:-4], blob) 93 | cv2.imwrite(dst, face_align.norm_crop(img, landmark=face.kps, image_size=224)) 94 | 95 | def run(self): 96 | self.copy() 97 | self.preprocess() 98 | self.arcface() 99 | -------------------------------------------------------------------------------- /datasets/creation/instances/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Zielon/MICA/af22e7a5810d474bc28a1433db533723d6bd2b07/datasets/creation/instances/__init__.py -------------------------------------------------------------------------------- /datasets/creation/instances/bu3dfe.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2023 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: mica@tue.mpg.de 16 | 17 | 18 | from abc import ABC 19 | from glob import glob 20 | from pathlib import Path 21 | 22 | from pytorch3d.io import load_objs_as_meshes 23 | 24 | from datasets.creation.instances.instance import Instance 25 | 26 | 27 | class BU3DFE(Instance, ABC): 28 | def __init__(self): 29 | super(BU3DFE, self).__init__() 30 | self.dst = '/scratch/NFC/OnFlame/BU3DFE/' 31 | self.src = '/scratch/NFC/BU-3DFE/' 32 | 33 | def get_images(self): 34 | images = {} 35 | for actor in sorted(glob(self.get_src().replace('BU-3DFE', 'BU-3DFE_clean') + 'images/*')): 36 | images[Path(actor).name] = glob(f'{actor}/*.jpg') 37 | 38 | return images 39 | 40 | def get_flame_params(self): 41 | prams = {} 42 | for actor in sorted(glob(self.get_src() + 'FLAME_parameters/iter2/*')): 43 | prams[Path(actor).name] = glob(f'{actor}/*.npz') 44 | 45 | return prams 46 | 47 | def get_registrations(self): 48 | registrations = {} 49 | for actor in sorted(glob(self.get_src() + 'registrations/iter2/neutral_align/*')): 50 | registrations[Path(actor).name] = glob(f'{actor}/*.obj') 51 | 52 | return registrations 53 | 54 | def get_meshes(self): 55 | meshes = {} 56 | files = sorted(glob(self.get_src() + 'raw_ne_data/*')) 57 | actors = set(map(lambda f: Path(f).name[0:5], files)) 58 | for actor in actors: 59 | meshes[Path(actor).name] = next(filter(lambda f: actor in f and 'obj' in f, files)) 60 | 61 | return meshes 62 | 63 | def transform_mesh(self, path): 64 | self.update_obj(path) 65 | mesh = load_objs_as_meshes([path], device=self.device) 66 | vertices = mesh._verts_list[0] 67 | center = vertices.mean(0) 68 | mesh._verts_list = [vertices - center] 69 | mesh.scale_verts_(0.01) 70 | 71 | return mesh.clone() 72 | -------------------------------------------------------------------------------- /datasets/creation/instances/d3dfacs.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2023 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: mica@tue.mpg.de 16 | 17 | 18 | from abc import ABC 19 | from glob import glob 20 | from pathlib import Path 21 | 22 | from datasets.creation.instances.instance import Instance 23 | 24 | 25 | class D3DFACS(Instance, ABC): 26 | def __init__(self): 27 | super(D3DFACS, self).__init__() 28 | self.dst = '/scratch/NFC/OnFlame/D3DFACS/' 29 | self.src = '/home/wzielonka/datasets/D3DFACS/' 30 | 31 | def get_images(self): 32 | images = {} 33 | for file in sorted(glob(self.get_src() + 'processed/images/*')): 34 | actor = Path(file).stem 35 | images[actor] = glob(f'{file}/*.jpg') 36 | 37 | return images 38 | 39 | def get_flame_params(self): 40 | params = {} 41 | for file in sorted(glob(self.get_src() + 'processed/FLAME/*.npz')): 42 | actor = Path(file).stem 43 | params[actor] = [file] 44 | 45 | return params 46 | 47 | def get_registrations(self): 48 | registrations = {} 49 | for file in sorted(glob(self.get_src() + 'processed/registrations/*')): 50 | actor = Path(file).stem.split('_')[0] 51 | registrations[actor] = [file] 52 | 53 | return registrations 54 | -------------------------------------------------------------------------------- /datasets/creation/instances/facewarehouse.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2023 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: mica@tue.mpg.de 16 | 17 | 18 | from abc import ABC 19 | from glob import glob 20 | from pathlib import Path 21 | 22 | from datasets.creation.instances.instance import Instance 23 | 24 | 25 | class FaceWarehouse(Instance, ABC): 26 | def __init__(self): 27 | super(FaceWarehouse, self).__init__() 28 | self.dst = '/scratch/NFC/OnFlame/FACEWAREHOUSE/' 29 | self.src = '/scratch/NFC/FaceWarehouse/' 30 | 31 | def get_images(self): 32 | images = {} 33 | for actor in sorted(glob(self.get_src() + 'Images/*')): 34 | images[Path(actor).stem] = glob(f'{actor}/*.png') 35 | 36 | return images 37 | 38 | def get_flame_params(self): 39 | params = {} 40 | for actor in sorted(glob(self.get_src() + 'FLAME_fits/*')): 41 | params[Path(actor).stem] = [sorted(glob(f'{actor}/*.npz'))[0]] 42 | 43 | return params 44 | 45 | def get_registrations(self): 46 | registrations = {} 47 | for actor in sorted(glob(self.get_src() + 'FLAME_fits/*')): 48 | registrations[Path(actor).stem] = [f'{actor}/tmp/pose_0__def_trafo_fit.obj'] 49 | 50 | return registrations 51 | -------------------------------------------------------------------------------- /datasets/creation/instances/florence.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2023 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: mica@tue.mpg.de 16 | 17 | 18 | from abc import ABC 19 | from glob import glob 20 | from pathlib import Path 21 | 22 | import numpy as np 23 | 24 | from datasets.creation.instances.instance import Instance 25 | 26 | 27 | class Florence(Instance, ABC): 28 | def __init__(self): 29 | super(Florence, self).__init__() 30 | self.dst = '/scratch/NFC/OnFlame/FLORENCE/' 31 | self.src = '/scratch/NFC/MICC_Florence/' 32 | 33 | def get_min_det_score(self): 34 | return 0.85 35 | 36 | def get_images(self): 37 | images = {} 38 | for actor in sorted(glob(self.get_src() + 'images/*')): 39 | imgs = sorted(list(filter(lambda f: 'PTZ-Outdoor' not in f, glob(f'{actor}/*/*.jpg')))) 40 | indecies = np.random.choice(len(imgs), 1000, replace=False) 41 | images[Path(actor).stem] = [imgs[i] for i in indecies] 42 | 43 | return images 44 | 45 | def get_flame_params(self): 46 | params = {} 47 | for actor in sorted(glob(self.get_src() + 'FLAME_parameters/iter1/*')): 48 | params[Path(actor).stem] = glob(f'{actor}/*.npz') 49 | 50 | return params 51 | 52 | def get_registrations(self): 53 | registrations = {} 54 | for actor in sorted(glob(self.get_src() + 'registrations/iter1/*')): 55 | if 'rendering' in actor: 56 | continue 57 | registrations[Path(actor).stem] = glob(f'{actor}/*.obj') 58 | 59 | return registrations 60 | -------------------------------------------------------------------------------- /datasets/creation/instances/frgc.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2023 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: mica@tue.mpg.de 16 | 17 | 18 | from abc import ABC 19 | from glob import glob 20 | from pathlib import Path 21 | 22 | import numpy as np 23 | from pytorch3d.io import load_objs_as_meshes 24 | 25 | from datasets.creation.instances.instance import Instance 26 | 27 | 28 | class FRGC(Instance, ABC): 29 | def __init__(self): 30 | super(FRGC, self).__init__() 31 | self.dst = '/scratch/NFC/OnFlame/FRGC/' 32 | self.src = '/scratch/NFC/FRGC_v2/' 33 | 34 | def get_images(self): 35 | images = {} 36 | for actor in sorted(glob(self.get_src() + 'images/*')): 37 | imgs = list(filter(lambda f: 'Spring2003range' not in f, glob(f'/{actor}/*/*.jpg'))) 38 | images[Path(actor).name] = imgs 39 | 40 | return images 41 | 42 | def get_flame_params(self): 43 | prams = {} 44 | for actor in sorted(glob(self.get_src() + 'FLAME_parameters/*')): 45 | prams[Path(actor).name] = glob(f'/{actor}/*.npz') 46 | 47 | return prams 48 | 49 | def get_registrations(self): 50 | registrations = {} 51 | for actor in sorted(glob(self.get_src() + 'registrations/*')): 52 | registrations[Path(actor).name] = glob(f'/{actor}/*.obj') 53 | 54 | return registrations 55 | 56 | def get_meshes(self): 57 | meshes = {} 58 | for file in sorted(glob(self.get_src() + 'registrations_tmp_new/*')): 59 | meshes[Path(file).name] = glob(f'/{file}/*.obj') 60 | 61 | sessions = np.load('/home/wzielonka/documents/scans_to_session.npy', allow_pickle=True)[()] 62 | valid = [] 63 | for key in sessions.keys(): 64 | if 'Spring2003range' not in sessions[key]: 65 | valid.append(key) 66 | 67 | filtered = {} 68 | for actor in meshes.keys(): 69 | files = meshes[actor] 70 | selected = list(filter(lambda f: Path(f).stem in valid, files)) 71 | if len(selected) > 0: 72 | filtered[actor] = selected 73 | 74 | return filtered 75 | 76 | def transform_mesh(self, path): 77 | self.update_obj(path[0]) 78 | mesh = load_objs_as_meshes(path, device=self.device) 79 | mesh.scale_verts_(10.0) 80 | vertices = mesh._verts_list[0] 81 | center = vertices.mean(0) 82 | mesh._verts_list = [vertices - center] 83 | 84 | return mesh.clone() 85 | -------------------------------------------------------------------------------- /datasets/creation/instances/instance.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2023 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: mica@tue.mpg.de 16 | 17 | 18 | import os 19 | from abc import abstractmethod 20 | from pathlib import Path 21 | 22 | from pytorch3d.transforms import RotateAxisAngle 23 | 24 | 25 | class Instance: 26 | def __init__(self): 27 | self.mount = '/home/wzielonka/Cluster/lustre' 28 | self.dst = 'empty' 29 | self.src = 'empty' 30 | self.device = 'cuda:0' 31 | self.actors = [] 32 | self.use_mount = os.path.exists(self.mount) 33 | 34 | def get_dst(self): 35 | return self.dst if not self.use_mount else self.mount + self.dst 36 | 37 | def get_src(self): 38 | return self.src if not self.use_mount else self.mount + self.src 39 | 40 | @abstractmethod 41 | def get_min_det_score(self): 42 | return 0 43 | 44 | @abstractmethod 45 | def preprocess(self): 46 | pass 47 | 48 | @abstractmethod 49 | def get_images(self): 50 | return {} 51 | 52 | @abstractmethod 53 | def get_flame_params(self): 54 | return {} 55 | 56 | @abstractmethod 57 | def get_registrations(self): 58 | return {} 59 | 60 | @abstractmethod 61 | def get_meshes(self): 62 | return {} 63 | 64 | @abstractmethod 65 | def transform_mesh(self, path): 66 | return None 67 | 68 | @abstractmethod 69 | def transform_image(self, img): 70 | return [img] 71 | 72 | @abstractmethod 73 | def transform_path(self, file): 74 | return Path(file).name 75 | 76 | @abstractmethod 77 | def get_rotations(self): 78 | rots = {} 79 | degree = 2.5 80 | step = int(15 / degree / 2) 81 | X = range(-step, step + 1) 82 | degree = 8.0 83 | step = int(144 / degree / 2) 84 | Y = range(-step, step + 1) 85 | for a, angles in [('X', X), ('Y', Y)]: 86 | r = [] 87 | for i in angles: 88 | r.append((RotateAxisAngle(float(degree * i), axis=a, device=self.device), float(degree * i))) 89 | rots[a] = r 90 | return rots 91 | 92 | @abstractmethod 93 | def update_obj(self, path, fix_mtl=False): 94 | mesh = Path(path).stem 95 | with open(path, 'r') as file: 96 | filedata = file.readlines() 97 | 98 | input = [] 99 | for line in filedata: 100 | if 'usemtl' in line or 'newmtl' in line: 101 | continue 102 | input.append(line) 103 | 104 | output = [] 105 | for line in input: 106 | if 'mtllib' in line: 107 | mtl = line.split(' ')[-1].split('.')[0] 108 | line += f'usemtl {mtl}\n' 109 | output.append(line) 110 | with open(path, 'w') as file: 111 | file_lines = "".join(output) 112 | file.write(file_lines) 113 | 114 | if not fix_mtl: 115 | return 116 | 117 | with open(path.replace('obj', 'mtl'), 'r') as file: 118 | filedata = file.readlines() 119 | 120 | output = [] 121 | for line in filedata: 122 | if 'newmtl' in line: 123 | line = 'newmtl ' + mesh + '\n' 124 | output.append(line) 125 | with open(path.replace('obj', 'mtl'), 'w') as file: 126 | file_lines = "".join(output) 127 | file.write(file_lines) 128 | -------------------------------------------------------------------------------- /datasets/creation/instances/lyhm.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2023 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: mica@tue.mpg.de 16 | 17 | 18 | from abc import ABC 19 | from glob import glob 20 | from pathlib import Path 21 | 22 | from PIL import ImageFile 23 | 24 | ImageFile.LOAD_TRUNCATED_IMAGES = True 25 | from pytorch3d.io import load_objs_as_meshes 26 | from pytorch3d.transforms import RotateAxisAngle 27 | 28 | from datasets.creation.instances.instance import Instance 29 | 30 | 31 | class LYHM(Instance, ABC): 32 | def __init__(self): 33 | super(LYHM, self).__init__() 34 | self.dst = '/scratch/NFC/MICA/LYHM/' 35 | self.src = '/scratch/NFC/LYHM/' 36 | 37 | def get_images(self): 38 | images = {} 39 | for actor in sorted(glob(self.get_src() + '/*')): 40 | images[Path(actor).name] = glob(f'/{actor}/*.png') 41 | 42 | return images 43 | 44 | def get_flame_params(self): 45 | prams = {} 46 | for actor in sorted(glob(self.get_src() + '/*')): 47 | prams[Path(actor).name] = glob(f'/{actor}/*.npz') 48 | 49 | return prams 50 | 51 | def get_registrations(self): 52 | registrations = {} 53 | for actor in sorted(glob(self.get_src() + '/*')): 54 | all = glob(f'/{actor}/*.obj') 55 | registrations[Path(actor).name] = list(filter(lambda m: 'model_fit' not in m, all)) 56 | 57 | return registrations 58 | 59 | def get_meshes(self): 60 | meshes = {} 61 | for actor in sorted(glob(self.get_src() + '/*')): 62 | meshes[Path(actor).name] = glob(f'/{actor}/scan/*.obj') 63 | 64 | return meshes 65 | 66 | def transform_mesh(self, path): 67 | mesh = load_objs_as_meshes(path, device=self.device) 68 | vertices = mesh._verts_list[0] 69 | center = vertices.mean(0) 70 | mesh._verts_list = [vertices - center] 71 | mesh.scale_verts_(0.01) 72 | 73 | rot = RotateAxisAngle(-45, axis='X', device=self.device) 74 | mesh._verts_list = [rot.transform_points(mesh.verts_list()[0])] 75 | rot = RotateAxisAngle(-45, axis='Y', device=self.device) 76 | mesh._verts_list = [rot.transform_points(mesh.verts_list()[0])] 77 | 78 | return mesh.clone() 79 | -------------------------------------------------------------------------------- /datasets/creation/instances/pb4d.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2023 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: mica@tue.mpg.de 16 | 17 | 18 | from abc import ABC 19 | from glob import glob 20 | from pathlib import Path 21 | 22 | import numpy as np 23 | from pytorch3d.io import load_objs_as_meshes 24 | 25 | from datasets.creation.instances.instance import Instance 26 | 27 | 28 | class PB4D(Instance, ABC): 29 | def __init__(self): 30 | super(PB4D, self).__init__() 31 | self.dst = '/scratch/NFC/OnFlame/BP4D/' 32 | self.src = '/scratch/NFC/BP4D/' 33 | 34 | def get_images(self): 35 | images = {} 36 | for actor in sorted(glob(self.get_src() + 'images/*')): 37 | imgs = sorted(glob(f'/{actor}/*.jpg')) 38 | indecies = np.random.choice(len(imgs), 100, replace=False) 39 | images[Path(actor).name] = [imgs[i] for i in indecies] 40 | 41 | return images 42 | 43 | def get_flame_params(self): 44 | prams = {} 45 | for file in sorted(glob(self.get_src() + 'FLAME_parameters/*.npz')): 46 | prams[Path(file).stem] = [file] 47 | 48 | return prams 49 | 50 | def get_registrations(self): 51 | registrations = {} 52 | for file in sorted(glob(self.get_src() + 'registrations/*')): 53 | registrations[Path(file).stem] = [file] 54 | 55 | return registrations 56 | 57 | def get_meshes(self): 58 | meshes = {} 59 | for file in sorted(glob(self.get_src() + 'scans/*.obj')): 60 | meshes[Path(file).stem] = [file] 61 | 62 | return meshes 63 | 64 | def transform_mesh(self, path): 65 | mesh = load_objs_as_meshes(path, device=self.device) 66 | mesh.scale_verts_(0.01) 67 | vertices = mesh._verts_list[0] 68 | center = vertices.mean(0) 69 | mesh._verts_list = [vertices - center] 70 | 71 | return mesh.clone() 72 | -------------------------------------------------------------------------------- /datasets/creation/instances/stirling.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2023 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: mica@tue.mpg.de 16 | 17 | 18 | from abc import ABC 19 | from glob import glob 20 | from pathlib import Path 21 | 22 | from pytorch3d.io import load_objs_as_meshes 23 | 24 | from datasets.creation.instances.instance import Instance 25 | 26 | 27 | class Stirling(Instance, ABC): 28 | def __init__(self): 29 | super(Stirling, self).__init__() 30 | self.dst = '/scratch/NFC/OnFlame/STIRLING/' 31 | self.src = '/scratch/NFC/Stirling/' 32 | 33 | def get_min_det_score(self): 34 | return 0.75 35 | 36 | def get_images(self): 37 | images = {} 38 | for file in sorted(glob(self.get_src() + 'images/Real_images__Subset_2D_FG2018/HQ/*')): 39 | actor = Path(file).stem.split('_')[0].upper() 40 | if actor not in images: 41 | images[actor] = [] 42 | images[actor].append(file) 43 | 44 | return images 45 | 46 | def get_flame_params(self): 47 | prams = {} 48 | for file in sorted(glob(self.get_src() + 'FLAME_parameters/iter1/*/*.npz')): 49 | actor = Path(file).stem[0:5].upper() 50 | prams[Path(actor).name] = [file] 51 | 52 | return prams 53 | 54 | def get_registrations(self): 55 | registrations = {} 56 | for file in sorted(glob(self.get_src() + 'registrations/iter1/*/*')): 57 | if 'obj' not in file: 58 | continue 59 | actor = Path(file).stem[0:5].upper() 60 | registrations[Path(actor).name] = [file] 61 | 62 | return registrations 63 | 64 | def get_meshes(self): 65 | meshes = {} 66 | for file in sorted(glob(self.get_src() + 'scans/*/*.obj')): 67 | actor = Path(file).stem[0:5].upper() 68 | if 'obj' in file: 69 | meshes[actor] = file 70 | 71 | return meshes 72 | 73 | def transform_mesh(self, path): 74 | self.update_obj(path, fix_mtl=True) 75 | mesh = load_objs_as_meshes([path], device=self.device) 76 | vertices = mesh._verts_list[0] 77 | center = vertices.mean(0) 78 | mesh._verts_list = [vertices - center] 79 | mesh.scale_verts_(0.01) 80 | 81 | return mesh.clone() 82 | 83 | def transform_path(self, file): 84 | name = Path(file).name 85 | return name 86 | -------------------------------------------------------------------------------- /datasets/creation/main.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2023 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: mica@tue.mpg.de 16 | 17 | 18 | import numpy as np 19 | import torch 20 | 21 | from datasets.creation.generator import Generator 22 | from datasets.creation.instances.bu3dfe import BU3DFE 23 | from datasets.creation.instances.d3dfacs import D3DFACS 24 | from datasets.creation.instances.facewarehouse import FaceWarehouse 25 | from datasets.creation.instances.florence import Florence 26 | from datasets.creation.instances.frgc import FRGC 27 | from datasets.creation.instances.lyhm import LYHM 28 | from datasets.creation.instances.pb4d import PB4D 29 | from datasets.creation.instances.stirling import Stirling 30 | 31 | np.random.seed(42) 32 | 33 | if __name__ == '__main__': 34 | torch.multiprocessing.set_start_method('spawn') 35 | 36 | datasets = [FaceWarehouse(), LYHM(), D3DFACS(), FRGC(), Florence(), Stirling(), BU3DFE(), PB4D()] 37 | generator = Generator([FaceWarehouse()]) 38 | 39 | generator.run() 40 | -------------------------------------------------------------------------------- /datasets/creation/util.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2023 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: mica@tue.mpg.de 16 | 17 | 18 | import os 19 | import os.path as osp 20 | from pathlib import Path 21 | 22 | import cv2 23 | import numpy as np 24 | from insightface.utils import face_align 25 | from numpy.lib import math 26 | 27 | input_mean = 127.5 28 | input_std = 127.5 29 | 30 | 31 | def create_folders(folders): 32 | if not type(folders) is list: 33 | folders = folders.split('/') 34 | parents = '/' 35 | for folder in folders: 36 | parents = os.path.join(parents, folder) 37 | if os.path.exists(parents): 38 | continue 39 | Path(parents).mkdir(exist_ok=True) 40 | 41 | 42 | def get_arcface_input(face, img): 43 | aimg = face_align.norm_crop(img, landmark=face.kps) 44 | blob = cv2.dnn.blobFromImages([aimg], 1.0 / input_std, (112, 112), (input_mean, input_mean, input_mean), swapRB=True) 45 | return blob[0], aimg 46 | 47 | 48 | def get_image(name, to_rgb=False): 49 | images_dir = osp.join(Path(__file__).parent.absolute(), '../images') 50 | ext_names = ['.jpg', '.png', '.jpeg'] 51 | image_file = None 52 | for ext_name in ext_names: 53 | _image_file = osp.join(images_dir, "%s%s" % (name, ext_name)) 54 | if osp.exists(_image_file): 55 | image_file = _image_file 56 | break 57 | assert image_file is not None, '%s not found' % name 58 | img = cv2.imread(image_file) 59 | if to_rgb: 60 | img = img[:, :, ::-1] 61 | return img 62 | 63 | 64 | # from the original insightface.app.face_analysis.py file 65 | def draw_on(img, faces): 66 | import cv2 67 | dimg = img.copy() 68 | for i in range(len(faces)): 69 | face = faces[i] 70 | box = face.bbox.astype(np.int) 71 | color = (0, 0, 255) 72 | cv2.rectangle(dimg, (box[0], box[1]), (box[2], box[3]), color, 2) 73 | if face.kps is not None: 74 | kps = face.kps.astype(np.int) 75 | # print(landmark.shape) 76 | for l in range(kps.shape[0]): 77 | color = (0, 0, 255) 78 | if l == 0 or l == 3: 79 | color = (0, 255, 0) 80 | cv2.circle(dimg, (kps[l][0], kps[l][1]), 1, color, 81 | 2) 82 | if face.gender is not None and face.age is not None: 83 | cv2.putText(dimg, '%s,%d' % (face.sex, face.age), (box[0] - 1, box[1] - 4), cv2.FONT_HERSHEY_COMPLEX, 0.7, (0, 255, 0), 1) 84 | 85 | return dimg 86 | 87 | 88 | def dist(p1, p2): 89 | return math.sqrt(((p1[0] - p2[0]) ** 2) + ((p1[1] - p2[1]) ** 2)) 90 | 91 | 92 | def get_center(bboxes, img): 93 | img_center = img.shape[1] // 2, img.shape[0] // 2 94 | size = bboxes.shape[0] 95 | distance = np.Inf 96 | j = 0 97 | for i in range(size): 98 | x1, y1, x2, y2 = bboxes[i, 0:4] 99 | dx = abs(x2 - x1) / 2.0 100 | dy = abs(y2 - y1) / 2.0 101 | current = dist((x1 + dx, y1 + dy), img_center) 102 | if current < distance: 103 | distance = current 104 | j = i 105 | 106 | return j 107 | 108 | 109 | def bbox2point(left, right, top, bottom, type='bbox'): 110 | if type == 'kpt68': 111 | old_size = (right - left + bottom - top) / 2 * 1.1 112 | center = np.array([right - (right - left) / 2.0, bottom - (bottom - top) / 2.0]) 113 | elif type == 'bbox': 114 | old_size = (right - left + bottom - top) / 2 115 | center = np.array([right - (right - left) / 2.0, bottom - (bottom - top) / 2.0 + old_size * 0.12]) 116 | else: 117 | raise NotImplementedError 118 | return old_size, center 119 | 120 | 121 | def get_bbox(image, lmks, bb_scale=1.0): 122 | h, w, c = image.shape 123 | bbox = [] 124 | for i in range(lmks.shape[0]): 125 | lmks = lmks.astype(np.int32) 126 | x_min, x_max, y_min, y_max = np.min(lmks[i, :, 0]), np.max(lmks[i, :, 0]), np.min(lmks[i, :, 1]), np.max(lmks[i, :, 1]) 127 | x_center, y_center = int((x_max + x_min) / 2.0), int((y_max + y_min) / 2.0) 128 | size = int(bb_scale * 2 * max(x_center - x_min, y_center - y_min)) 129 | xb_min, xb_max, yb_min, yb_max = max(x_center - size // 2, 0), min(x_center + size // 2, w - 1), \ 130 | max(y_center - size // 2, 0), min(y_center + size // 2, h - 1) 131 | 132 | yb_max = min(yb_max, h - 1) 133 | xb_max = min(xb_max, w - 1) 134 | yb_min = max(yb_min, 0) 135 | xb_min = max(xb_min, 0) 136 | 137 | if (xb_max - xb_min) % 2 != 0: 138 | xb_min += 1 139 | 140 | if (yb_max - yb_min) % 2 != 0: 141 | yb_min += 1 142 | 143 | # x1, y1, x2, y2 144 | bbox.append(np.array([xb_min, yb_min, xb_max, yb_max, 0])) 145 | 146 | return np.stack(bbox) 147 | -------------------------------------------------------------------------------- /datasets/image_paths/BP4D.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Zielon/MICA/af22e7a5810d474bc28a1433db533723d6bd2b07/datasets/image_paths/BP4D.npy -------------------------------------------------------------------------------- /datasets/image_paths/BU3DFE.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Zielon/MICA/af22e7a5810d474bc28a1433db533723d6bd2b07/datasets/image_paths/BU3DFE.npy -------------------------------------------------------------------------------- /datasets/image_paths/D3DFACS.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Zielon/MICA/af22e7a5810d474bc28a1433db533723d6bd2b07/datasets/image_paths/D3DFACS.npy -------------------------------------------------------------------------------- /datasets/image_paths/FACEWAREHOUSE.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Zielon/MICA/af22e7a5810d474bc28a1433db533723d6bd2b07/datasets/image_paths/FACEWAREHOUSE.npy -------------------------------------------------------------------------------- /datasets/image_paths/FLORENCE.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Zielon/MICA/af22e7a5810d474bc28a1433db533723d6bd2b07/datasets/image_paths/FLORENCE.npy -------------------------------------------------------------------------------- /datasets/image_paths/FRGC.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Zielon/MICA/af22e7a5810d474bc28a1433db533723d6bd2b07/datasets/image_paths/FRGC.npy -------------------------------------------------------------------------------- /datasets/image_paths/LYHM.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Zielon/MICA/af22e7a5810d474bc28a1433db533723d6bd2b07/datasets/image_paths/LYHM.npy -------------------------------------------------------------------------------- /datasets/image_paths/STIRLING.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Zielon/MICA/af22e7a5810d474bc28a1433db533723d6bd2b07/datasets/image_paths/STIRLING.npy -------------------------------------------------------------------------------- /demo.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2023 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: mica@tue.mpg.de 16 | 17 | 18 | import argparse 19 | import os 20 | import random 21 | from glob import glob 22 | from pathlib import Path 23 | 24 | import cv2 25 | import numpy as np 26 | import torch 27 | import torch.backends.cudnn as cudnn 28 | import trimesh 29 | from insightface.app.common import Face 30 | from insightface.utils import face_align 31 | from loguru import logger 32 | from skimage.io import imread 33 | from tqdm import tqdm 34 | 35 | from configs.config import get_cfg_defaults 36 | from datasets.creation.util import get_arcface_input, get_center, draw_on 37 | from utils import util 38 | from utils.landmark_detector import LandmarksDetector, detectors 39 | 40 | 41 | def deterministic(rank): 42 | torch.manual_seed(rank) 43 | torch.cuda.manual_seed(rank) 44 | np.random.seed(rank) 45 | random.seed(rank) 46 | 47 | cudnn.deterministic = True 48 | cudnn.benchmark = False 49 | 50 | 51 | def process(args, app, image_size=224, draw_bbox=False): 52 | dst = Path(args.a) 53 | dst.mkdir(parents=True, exist_ok=True) 54 | processes = [] 55 | image_paths = sorted(glob(args.i + '/*.*')) 56 | for image_path in tqdm(image_paths): 57 | name = Path(image_path).stem 58 | img = cv2.imread(image_path) 59 | bboxes, kpss = app.detect(img) 60 | if bboxes.shape[0] == 0: 61 | logger.error(f'[ERROR] Face not detected for {image_path}') 62 | continue 63 | i = get_center(bboxes, img) 64 | bbox = bboxes[i, 0:4] 65 | det_score = bboxes[i, 4] 66 | kps = None 67 | if kpss is not None: 68 | kps = kpss[i] 69 | face = Face(bbox=bbox, kps=kps, det_score=det_score) 70 | blob, aimg = get_arcface_input(face, img) 71 | file = str(Path(dst, name)) 72 | np.save(file, blob) 73 | processes.append(file + '.npy') 74 | cv2.imwrite(file + '.jpg', face_align.norm_crop(img, landmark=face.kps, image_size=image_size)) 75 | if draw_bbox: 76 | dimg = draw_on(img, [face]) 77 | cv2.imwrite(file + '_bbox.jpg', dimg) 78 | 79 | return processes 80 | 81 | 82 | def to_batch(path): 83 | src = path.replace('npy', 'jpg') 84 | if not os.path.exists(src): 85 | src = path.replace('npy', 'png') 86 | 87 | image = imread(src)[:, :, :3] 88 | image = image / 255. 89 | image = cv2.resize(image, (224, 224)).transpose(2, 0, 1) 90 | image = torch.tensor(image).cuda()[None] 91 | 92 | arcface = np.load(path) 93 | arcface = torch.tensor(arcface).cuda()[None] 94 | 95 | return image, arcface 96 | 97 | 98 | def load_checkpoint(args, mica): 99 | checkpoint = torch.load(args.m) 100 | if 'arcface' in checkpoint: 101 | mica.arcface.load_state_dict(checkpoint['arcface']) 102 | if 'flameModel' in checkpoint: 103 | mica.flameModel.load_state_dict(checkpoint['flameModel']) 104 | 105 | 106 | def main(cfg, args): 107 | device = 'cuda:0' 108 | cfg.model.testing = True 109 | mica = util.find_model_using_name(model_dir='micalib.models', model_name=cfg.model.name)(cfg, device) 110 | load_checkpoint(args, mica) 111 | mica.eval() 112 | 113 | faces = mica.flameModel.generator.faces_tensor.cpu() 114 | Path(args.o).mkdir(exist_ok=True, parents=True) 115 | 116 | app = LandmarksDetector(model=detectors.RETINAFACE) 117 | 118 | with torch.no_grad(): 119 | logger.info(f'Processing has started...') 120 | paths = process(args, app, draw_bbox=False) 121 | for path in tqdm(paths): 122 | name = Path(path).stem 123 | images, arcface = to_batch(path) 124 | codedict = mica.encode(images, arcface) 125 | opdict = mica.decode(codedict) 126 | meshes = opdict['pred_canonical_shape_vertices'] 127 | code = opdict['pred_shape_code'] 128 | lmk = mica.flame.compute_landmarks(meshes) 129 | 130 | mesh = meshes[0] 131 | landmark_51 = lmk[0, 17:] 132 | landmark_7 = landmark_51[[19, 22, 25, 28, 16, 31, 37]] 133 | 134 | dst = Path(args.o, name) 135 | dst.mkdir(parents=True, exist_ok=True) 136 | trimesh.Trimesh(vertices=mesh.cpu() * 1000.0, faces=faces, process=False).export(f'{dst}/mesh.ply') # save in millimeters 137 | trimesh.Trimesh(vertices=mesh.cpu() * 1000.0, faces=faces, process=False).export(f'{dst}/mesh.obj') 138 | np.save(f'{dst}/identity', code[0].cpu().numpy()) 139 | np.save(f'{dst}/kpt7', landmark_7.cpu().numpy() * 1000.0) 140 | np.save(f'{dst}/kpt68', lmk.cpu().numpy() * 1000.0) 141 | 142 | logger.info(f'Processing finished. Results has been saved in {args.o}') 143 | 144 | 145 | if __name__ == '__main__': 146 | parser = argparse.ArgumentParser(description='MICA - Towards Metrical Reconstruction of Human Faces') 147 | parser.add_argument('-i', default='demo/input', type=str, help='Input folder with images') 148 | parser.add_argument('-o', default='demo/output', type=str, help='Output folder') 149 | parser.add_argument('-a', default='demo/arcface', type=str, help='Processed images for MICA input') 150 | parser.add_argument('-m', default='data/pretrained/mica.tar', type=str, help='Pretrained model path') 151 | 152 | args = parser.parse_args() 153 | cfg = get_cfg_defaults() 154 | 155 | deterministic(42) 156 | main(cfg, args) 157 | -------------------------------------------------------------------------------- /demo/input/carell.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Zielon/MICA/af22e7a5810d474bc28a1433db533723d6bd2b07/demo/input/carell.jpg -------------------------------------------------------------------------------- /demo/input/connelly.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Zielon/MICA/af22e7a5810d474bc28a1433db533723d6bd2b07/demo/input/connelly.jpg -------------------------------------------------------------------------------- /demo/input/justin.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Zielon/MICA/af22e7a5810d474bc28a1433db533723d6bd2b07/demo/input/justin.png -------------------------------------------------------------------------------- /demo/input/lawrence.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Zielon/MICA/af22e7a5810d474bc28a1433db533723d6bd2b07/demo/input/lawrence.jpg -------------------------------------------------------------------------------- /documents/BP4D.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Zielon/MICA/af22e7a5810d474bc28a1433db533723d6bd2b07/documents/BP4D.gif -------------------------------------------------------------------------------- /documents/D3DFACS.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Zielon/MICA/af22e7a5810d474bc28a1433db533723d6bd2b07/documents/D3DFACS.gif -------------------------------------------------------------------------------- /documents/FACEWAREHOUSE.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Zielon/MICA/af22e7a5810d474bc28a1433db533723d6bd2b07/documents/FACEWAREHOUSE.gif -------------------------------------------------------------------------------- /documents/FLORENCE.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Zielon/MICA/af22e7a5810d474bc28a1433db533723d6bd2b07/documents/FLORENCE.gif -------------------------------------------------------------------------------- /documents/FRGC.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Zielon/MICA/af22e7a5810d474bc28a1433db533723d6bd2b07/documents/FRGC.gif -------------------------------------------------------------------------------- /documents/LYHM.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Zielon/MICA/af22e7a5810d474bc28a1433db533723d6bd2b07/documents/LYHM.gif -------------------------------------------------------------------------------- /documents/STIRLING.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Zielon/MICA/af22e7a5810d474bc28a1433db533723d6bd2b07/documents/STIRLING.gif -------------------------------------------------------------------------------- /documents/teaser.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Zielon/MICA/af22e7a5810d474bc28a1433db533723d6bd2b07/documents/teaser.jpg -------------------------------------------------------------------------------- /documents/voxceleb.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Zielon/MICA/af22e7a5810d474bc28a1433db533723d6bd2b07/documents/voxceleb.gif -------------------------------------------------------------------------------- /environment.yml: -------------------------------------------------------------------------------- 1 | name: MICA 2 | channels: 3 | - pytorch 4 | - nvidia 5 | - defaults 6 | dependencies: 7 | - _libgcc_mutex=0.1=main 8 | - _openmp_mutex=5.1=1_gnu 9 | - blas=1.0=mkl 10 | - brotlipy=0.7.0=py39h27cfd23_1003 11 | - bzip2=1.0.8=h7b6447c_0 12 | - ca-certificates=2022.10.11=h06a4308_0 13 | - certifi=2022.9.24=py39h06a4308_0 14 | - cffi=1.15.1=py39h5eee18b_2 15 | - charset-normalizer=2.0.4=pyhd3eb1b0_0 16 | - cryptography=38.0.1=py39h9ce1e76_0 17 | - cuda=11.6.2=0 18 | - cuda-cccl=11.6.55=hf6102b2_0 19 | - cuda-command-line-tools=11.6.2=0 20 | - cuda-compiler=11.6.2=0 21 | - cuda-cudart=11.6.55=he381448_0 22 | - cuda-cudart-dev=11.6.55=h42ad0f4_0 23 | - cuda-cuobjdump=11.6.124=h2eeebcb_0 24 | - cuda-cupti=11.6.124=h86345e5_0 25 | - cuda-cuxxfilt=11.6.124=hecbf4f6_0 26 | - cuda-driver-dev=11.6.55=0 27 | - cuda-gdb=11.8.86=0 28 | - cuda-libraries=11.6.2=0 29 | - cuda-libraries-dev=11.6.2=0 30 | - cuda-memcheck=11.8.86=0 31 | - cuda-nsight=11.8.86=0 32 | - cuda-nsight-compute=11.8.0=0 33 | - cuda-nvcc=11.6.124=hbba6d2d_0 34 | - cuda-nvdisasm=11.8.86=0 35 | - cuda-nvml-dev=11.6.55=haa9ef22_0 36 | - cuda-nvprof=11.8.87=0 37 | - cuda-nvprune=11.6.124=he22ec0a_0 38 | - cuda-nvrtc=11.6.124=h020bade_0 39 | - cuda-nvrtc-dev=11.6.124=h249d397_0 40 | - cuda-nvtx=11.6.124=h0630a44_0 41 | - cuda-nvvp=11.8.87=0 42 | - cuda-runtime=11.6.2=0 43 | - cuda-samples=11.6.101=h8efea70_0 44 | - cuda-sanitizer-api=11.8.86=0 45 | - cuda-toolkit=11.6.2=0 46 | - cuda-tools=11.6.2=0 47 | - cuda-visual-tools=11.6.2=0 48 | - ffmpeg=4.3=hf484d3e_0 49 | - freetype=2.12.1=h4a9f257_0 50 | - gds-tools=1.4.0.31=0 51 | - giflib=5.2.1=h7b6447c_0 52 | - gmp=6.2.1=h295c915_3 53 | - gnutls=3.6.15=he1e5248_0 54 | - idna=3.4=py39h06a4308_0 55 | - intel-openmp=2021.4.0=h06a4308_3561 56 | - jpeg=9e=h7f8727e_0 57 | - lame=3.100=h7b6447c_0 58 | - lcms2=2.12=h3be6417_0 59 | - ld_impl_linux-64=2.38=h1181459_1 60 | - lerc=3.0=h295c915_0 61 | - libcublas=11.11.3.6=0 62 | - libcublas-dev=11.11.3.6=0 63 | - libcufft=10.9.0.58=0 64 | - libcufft-dev=10.9.0.58=0 65 | - libcufile=1.4.0.31=0 66 | - libcufile-dev=1.4.0.31=0 67 | - libcurand=10.3.0.86=0 68 | - libcurand-dev=10.3.0.86=0 69 | - libcusolver=11.4.1.48=0 70 | - libcusolver-dev=11.4.1.48=0 71 | - libcusparse=11.7.5.86=0 72 | - libcusparse-dev=11.7.5.86=0 73 | - libdeflate=1.8=h7f8727e_5 74 | - libffi=3.4.2=h6a678d5_6 75 | - libgcc-ng=11.2.0=h1234567_1 76 | - libgomp=11.2.0=h1234567_1 77 | - libiconv=1.16=h7f8727e_2 78 | - libidn2=2.3.2=h7f8727e_0 79 | - libnpp=11.8.0.86=0 80 | - libnpp-dev=11.8.0.86=0 81 | - libnvjpeg=11.9.0.86=0 82 | - libnvjpeg-dev=11.9.0.86=0 83 | - libpng=1.6.37=hbc83047_0 84 | - libstdcxx-ng=11.2.0=h1234567_1 85 | - libtasn1=4.16.0=h27cfd23_0 86 | - libtiff=4.4.0=hecacb30_2 87 | - libunistring=0.9.10=h27cfd23_0 88 | - libwebp=1.2.4=h11a3e52_0 89 | - libwebp-base=1.2.4=h5eee18b_0 90 | - lz4-c=1.9.3=h295c915_1 91 | - mkl=2021.4.0=h06a4308_640 92 | - mkl-service=2.4.0=py39h7f8727e_0 93 | - mkl_fft=1.3.1=py39hd3c417c_0 94 | - mkl_random=1.2.2=py39h51133e4_0 95 | - ncurses=6.3=h5eee18b_3 96 | - nettle=3.7.3=hbbd107a_1 97 | - nsight-compute=2022.3.0.22=0 98 | - numpy=1.23.4=py39h14f4228_0 99 | - numpy-base=1.23.4=py39h31eccc5_0 100 | - openh264=2.1.1=h4ff587b_0 101 | - openssl=1.1.1s=h7f8727e_0 102 | - pillow=9.2.0=py39hace64e9_1 103 | - pip=22.2.2=py39h06a4308_0 104 | - pycparser=2.21=pyhd3eb1b0_0 105 | - pyopenssl=22.0.0=pyhd3eb1b0_0 106 | - pysocks=1.7.1=py39h06a4308_0 107 | - python=3.9.15=h7a1cb2a_2 108 | - pytorch=1.13.0=py3.9_cuda11.6_cudnn8.3.2_0 109 | - pytorch-cuda=11.6=h867d48c_0 110 | - pytorch-mutex=1.0=cuda 111 | - readline=8.2=h5eee18b_0 112 | - requests=2.28.1=py39h06a4308_0 113 | - setuptools=65.5.0=py39h06a4308_0 114 | - six=1.16.0=pyhd3eb1b0_1 115 | - sqlite=3.40.0=h5082296_0 116 | - tk=8.6.12=h1ccaba5_0 117 | - torchaudio=0.13.0=py39_cu116 118 | - torchvision=0.14.0=py39_cu116 119 | - typing_extensions=4.3.0=py39h06a4308_0 120 | - tzdata=2022f=h04d1e81_0 121 | - urllib3=1.26.12=py39h06a4308_0 122 | - wheel=0.37.1=pyhd3eb1b0_0 123 | - xz=5.2.6=h5eee18b_0 124 | - zlib=1.2.13=h5eee18b_0 125 | - zstd=1.5.2=ha4553b6_0 126 | - pip: 127 | - albumentations==1.3.0 128 | - cachetools==5.2.0 129 | - chumpy==0.70 130 | - coloredlogs==15.0.1 131 | - contourpy==1.0.6 132 | - cycler==0.11.0 133 | - cython==0.29.32 134 | - easydict==1.10 135 | - face-alignment==1.3.5 136 | - falcon==3.1.1 137 | - falcon-multipart==0.2.0 138 | - flatbuffers==22.11.23 139 | - fonttools==4.38.0 140 | - google-api-core==2.11.0 141 | - google-api-python-client==2.69.0 142 | - google-auth==2.15.0 143 | - google-auth-httplib2==0.1.0 144 | - googleapis-common-protos==1.57.0 145 | - gunicorn==20.1.0 146 | - httplib2==0.21.0 147 | - humanfriendly==10.0 148 | - imageio==2.22.4 149 | - insightface==0.7 150 | - joblib==1.2.0 151 | - kiwisolver==1.4.4 152 | - llvmlite==0.39.1 153 | - loguru==0.6.0 154 | - matplotlib==3.6.2 155 | - mpmath==1.2.1 156 | - networkx==2.8.8 157 | - numba==0.56.4 158 | - oauth2client==4.1.3 159 | - onnx==1.13.0 160 | - onnxruntime==1.13.1 161 | - opencv-python==4.7.0.72 162 | - opencv-python-headless==4.6.0.66 163 | - packaging==21.3 164 | - prettytable==3.5.0 165 | - protobuf==3.20.2 166 | - pyasn1==0.4.8 167 | - pyasn1-modules==0.2.8 168 | - pydrive2==1.15.0 169 | - pyparsing==3.0.9 170 | - python-datauri==1.1.0 171 | - python-dateutil==2.8.2 172 | - pywavelets==1.4.1 173 | - pyyaml==6.0 174 | - qudida==0.0.4 175 | - rsa==4.9 176 | - scikit-image==0.19.3 177 | - scikit-learn==1.1.3 178 | - scipy==1.9.3 179 | - sympy==1.11.1 180 | - threadpoolctl==3.1.0 181 | - tifffile==2022.10.10 182 | - tqdm==4.64.1 183 | - trimesh==3.16.4 184 | - uritemplate==4.1.1 185 | - wcwidth==0.2.5 186 | - yacs==0.1.8 187 | prefix: /home/wzielonka/miniconda3/envs/MICA 188 | -------------------------------------------------------------------------------- /install.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | urle () { [[ "${1}" ]] || return 1; local LANG=C i x; for (( i = 0; i < ${#1}; i++ )); do x="${1:i:1}"; [[ "${x}" == [a-zA-Z0-9.~-] ]] && echo -n "${x}" || printf '%%%02X' "'${x}"; done; echo; } 3 | 4 | # username and password input 5 | echo -e "\nIf you do not have an account you can register at https://flame.is.tue.mpg.de/ following the installation instruction." 6 | read -p "Username (FLAME):" username 7 | read -p "Password (FLAME):" password 8 | username=$(urle $username) 9 | password=$(urle $password) 10 | 11 | echo -e "\nDownloading FLAME..." 12 | mkdir -p data/FLAME2020/ 13 | wget --post-data "username=$username&password=$password" 'https://download.is.tue.mpg.de/download.php?domain=flame&sfile=FLAME2020.zip&resume=1' -O './FLAME2020.zip' --no-check-certificate --continue 14 | unzip FLAME2020.zip -d data/FLAME2020/ 15 | rm -rf FLAME2020.zip 16 | 17 | # Install gdown if not installed 18 | if ! command -v gdown &> /dev/null; then 19 | echo "Installing gdown..." 20 | pip install gdown 21 | fi 22 | 23 | echo -e "\nDownloading MICA..." 24 | mkdir -p data/pretrained/ 25 | gdown --id 1bYsI_spptzyuFmfLYqYkcJA6GZWZViNt -O data/pretrained/mica.tar 26 | 27 | # https://github.com/deepinsight/insightface/issues/1896 28 | # Insightface has problems with hosting the models 29 | echo -e "\nDownloading insightface models..." 30 | mkdir -p ~/.insightface/models/ 31 | if [ ! -d ~/.insightface/models/antelopev2 ]; then 32 | gdown --id 16PWKI_RjjbE4_kqpElG-YFqe8FpXjads -O ~/.insightface/models/antelopev2.zip 33 | unzip ~/.insightface/models/antelopev2.zip -d ~/.insightface/models/antelopev2 34 | fi 35 | if [ ! -d ~/.insightface/models/buffalo_l ]; then 36 | gdown --id 1navJMy0DTr1_DHjLWu1i48owCPvXWfYc -O ~/.insightface/models/buffalo_l.zip 37 | unzip ~/.insightface/models/buffalo_l.zip -d ~/.insightface/models/buffalo_l 38 | fi 39 | 40 | echo -e "\nInstalling conda env..." 41 | conda env create -f environment.yml 42 | 43 | echo -e "\nInstallation has finished!" 44 | -------------------------------------------------------------------------------- /jobs.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2023 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: mica@tue.mpg.de 16 | 17 | 18 | import os 19 | import random 20 | import sys 21 | 22 | import numpy as np 23 | import torch 24 | import torch.backends.cudnn as cudnn 25 | import torch.distributed as dist 26 | import yaml 27 | from loguru import logger 28 | 29 | from micalib.tester import Tester 30 | from micalib.trainer import Trainer 31 | from utils import util 32 | 33 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '.'))) 34 | 35 | 36 | def setup(rank, world_size, port): 37 | os.environ['MASTER_ADDR'] = 'localhost' 38 | os.environ['MASTER_PORT'] = str(port) 39 | dist.init_process_group("nccl", rank=rank, world_size=world_size, init_method="env://") 40 | 41 | 42 | def deterministic(rank): 43 | torch.manual_seed(rank) 44 | torch.cuda.manual_seed(rank) 45 | np.random.seed(rank) 46 | random.seed(rank) 47 | 48 | cudnn.deterministic = True 49 | cudnn.benchmark = False 50 | 51 | 52 | def test(rank, world_size, cfg, args): 53 | port = np.random.randint(low=0, high=2000) 54 | setup(rank, world_size, 12310 + port) 55 | 56 | deterministic(rank) 57 | 58 | cfg.model.testing = True 59 | mica = util.find_model_using_name(model_dir='micalib.models', model_name=cfg.model.name)(cfg, rank) 60 | tester = Tester(nfc_model=mica, config=cfg, device=rank) 61 | tester.render_mesh = True 62 | 63 | if args.test_dataset.upper() == 'STIRLING': 64 | tester.test_stirling(args.checkpoint) 65 | elif args.test_dataset.upper() == 'NOW': 66 | tester.test_now(args.checkpoint) 67 | else: 68 | logger.error('[TESTER] Test dataset was not specified!') 69 | 70 | dist.destroy_process_group() 71 | 72 | 73 | def train(rank, world_size, cfg): 74 | port = np.random.randint(low=0, high=2000) 75 | setup(rank, world_size, 12310 + port) 76 | 77 | logger.info(f'[MAIN] output_dir: {cfg.output_dir}') 78 | os.makedirs(os.path.join(cfg.output_dir, cfg.train.log_dir), exist_ok=True) 79 | os.makedirs(os.path.join(cfg.output_dir, cfg.train.vis_dir), exist_ok=True) 80 | os.makedirs(os.path.join(cfg.output_dir, cfg.train.val_vis_dir), exist_ok=True) 81 | 82 | with open(os.path.join(cfg.output_dir, cfg.train.log_dir, 'full_config.yaml'), 'w') as f: 83 | yaml.dump(cfg, f, default_flow_style=False) 84 | # shutil.copy(cfg.cfg_file, os.path.join(cfg.output_dir, 'config.yaml')) 85 | 86 | deterministic(rank) 87 | 88 | nfc = util.find_model_using_name(model_dir='micalib.models', model_name=cfg.model.name)(cfg, rank) 89 | trainer = Trainer(nfc_model=nfc, config=cfg, device=rank) 90 | trainer.fit() 91 | 92 | dist.destroy_process_group() 93 | -------------------------------------------------------------------------------- /micalib/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Zielon/MICA/af22e7a5810d474bc28a1433db533723d6bd2b07/micalib/__init__.py -------------------------------------------------------------------------------- /micalib/base_model.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2023 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: mica@tue.mpg.de 16 | 17 | 18 | from abc import abstractmethod 19 | 20 | import numpy as np 21 | import torch 22 | import torch.nn as nn 23 | 24 | from configs.config import cfg 25 | from models.flame import FLAME 26 | from utils.masking import Masking 27 | 28 | 29 | class BaseModel(nn.Module): 30 | def __init__(self, config=None, device=None, tag=''): 31 | super(BaseModel, self).__init__() 32 | if config is None: 33 | self.cfg = cfg 34 | else: 35 | self.cfg = config 36 | 37 | self.tag = tag 38 | self.use_mask = self.cfg.train.use_mask 39 | self.device = device 40 | self.masking = Masking(config) 41 | self.testing = self.cfg.model.testing 42 | 43 | def initialize(self): 44 | self.create_flame(self.cfg.model) 45 | self.create_model(self.cfg.model) 46 | self.load_model() 47 | self.setup_renderer(self.cfg.model) 48 | 49 | self.create_weights() 50 | 51 | def create_flame(self, model_cfg): 52 | self.flame = FLAME(model_cfg).to(self.device) 53 | self.average_face = self.flame.v_template.clone()[None] 54 | 55 | self.flame.eval() 56 | 57 | @abstractmethod 58 | def create_model(self): 59 | return 60 | 61 | @abstractmethod 62 | def create_load(self): 63 | return 64 | 65 | @abstractmethod 66 | def model_dict(self): 67 | return 68 | 69 | @abstractmethod 70 | def parameters_to_optimize(self): 71 | return 72 | 73 | @abstractmethod 74 | def encode(self, images, arcface_images): 75 | return 76 | 77 | @abstractmethod 78 | def decode(self, codedict, epoch): 79 | pass 80 | 81 | @abstractmethod 82 | def compute_losses(self, input, encoder_output, decoder_output): 83 | pass 84 | 85 | @abstractmethod 86 | def compute_masks(self, input, decoder_output): 87 | pass 88 | 89 | def setup_renderer(self, model_cfg): 90 | self.verts_template_neutral = self.flame.v_template[None] 91 | self.verts_template = None 92 | self.verts_template_uv = None 93 | 94 | def create_weights(self): 95 | self.vertices_mask = self.masking.get_weights_per_vertex().to(self.device) 96 | self.triangle_mask = self.masking.get_weights_per_triangle().to(self.device) 97 | 98 | def create_template(self, B): 99 | with torch.no_grad(): 100 | if self.verts_template is None: 101 | self.verts_template_neutral = self.flame.v_template[None] 102 | pose = torch.zeros(B, self.cfg.model.n_pose, device=self.device) 103 | pose[:, 3] = 10.0 * np.pi / 180.0 # 48 104 | self.verts_template, _, _ = self.flame(shape_params=torch.zeros(B, self.cfg.model.n_shape, device=self.device), expression_params=torch.zeros(B, self.cfg.model.n_exp, device=self.device), pose_params=pose) # use template mesh with open mouth 105 | 106 | if self.verts_template.shape[0] != B: 107 | self.verts_template_neutral = self.verts_template_neutral[0:1].repeat(B, 1, 1) 108 | self.verts_template = self.verts_template[0:1].repeat(B, 1, 1) 109 | -------------------------------------------------------------------------------- /micalib/models/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Zielon/MICA/af22e7a5810d474bc28a1433db533723d6bd2b07/micalib/models/__init__.py -------------------------------------------------------------------------------- /micalib/models/mica.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2023 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: mica@tue.mpg.de 16 | 17 | 18 | import os 19 | import sys 20 | 21 | sys.path.append("./nfclib") 22 | 23 | import torch 24 | import torch.nn.functional as F 25 | 26 | from models.arcface import Arcface 27 | from models.generator import Generator 28 | from micalib.base_model import BaseModel 29 | 30 | from loguru import logger 31 | 32 | 33 | class MICA(BaseModel): 34 | def __init__(self, config=None, device=None, tag='MICA'): 35 | super(MICA, self).__init__(config, device, tag) 36 | 37 | self.initialize() 38 | 39 | def create_model(self, model_cfg): 40 | mapping_layers = model_cfg.mapping_layers 41 | pretrained_path = None 42 | if not model_cfg.use_pretrained: 43 | pretrained_path = model_cfg.arcface_pretrained_model 44 | self.arcface = Arcface(pretrained_path=pretrained_path).to(self.device) 45 | self.flameModel = Generator(512, 300, self.cfg.model.n_shape, mapping_layers, model_cfg, self.device) 46 | 47 | def load_model(self): 48 | model_path = os.path.join(self.cfg.output_dir, 'model.tar') 49 | if os.path.exists(self.cfg.pretrained_model_path) and self.cfg.model.use_pretrained: 50 | model_path = self.cfg.pretrained_model_path 51 | if os.path.exists(model_path): 52 | logger.info(f'[{self.tag}] Trained model found. Path: {model_path} | GPU: {self.device}') 53 | checkpoint = torch.load(model_path) 54 | if 'arcface' in checkpoint: 55 | self.arcface.load_state_dict(checkpoint['arcface']) 56 | if 'flameModel' in checkpoint: 57 | self.flameModel.load_state_dict(checkpoint['flameModel']) 58 | else: 59 | logger.info(f'[{self.tag}] Checkpoint not available starting from scratch!') 60 | 61 | def model_dict(self): 62 | return { 63 | 'flameModel': self.flameModel.state_dict(), 64 | 'arcface': self.arcface.state_dict() 65 | } 66 | 67 | def parameters_to_optimize(self): 68 | return [ 69 | {'params': self.flameModel.parameters(), 'lr': self.cfg.train.lr}, 70 | {'params': self.arcface.parameters(), 'lr': self.cfg.train.arcface_lr}, 71 | ] 72 | 73 | def encode(self, images, arcface_imgs): 74 | codedict = {} 75 | 76 | codedict['arcface'] = F.normalize(self.arcface(arcface_imgs)) 77 | codedict['images'] = images 78 | 79 | return codedict 80 | 81 | def decode(self, codedict, epoch=0): 82 | self.epoch = epoch 83 | 84 | flame_verts_shape = None 85 | shapecode = None 86 | 87 | if not self.testing: 88 | flame = codedict['flame'] 89 | shapecode = flame['shape_params'].view(-1, flame['shape_params'].shape[2]) 90 | shapecode = shapecode.to(self.device)[:, :self.cfg.model.n_shape] 91 | with torch.no_grad(): 92 | flame_verts_shape, _, _ = self.flame(shape_params=shapecode) 93 | 94 | identity_code = codedict['arcface'] 95 | pred_canonical_vertices, pred_shape_code = self.flameModel(identity_code) 96 | 97 | output = { 98 | 'flame_verts_shape': flame_verts_shape, 99 | 'flame_shape_code': shapecode, 100 | 'pred_canonical_shape_vertices': pred_canonical_vertices, 101 | 'pred_shape_code': pred_shape_code, 102 | 'faceid': codedict['arcface'] 103 | } 104 | 105 | return output 106 | 107 | def compute_losses(self, input, encoder_output, decoder_output): 108 | losses = {} 109 | 110 | pred_verts = decoder_output['pred_canonical_shape_vertices'] 111 | gt_verts = decoder_output['flame_verts_shape'].detach() 112 | 113 | pred_verts_shape_canonical_diff = (pred_verts - gt_verts).abs() 114 | 115 | if self.use_mask: 116 | pred_verts_shape_canonical_diff *= self.vertices_mask 117 | 118 | losses['pred_verts_shape_canonical_diff'] = torch.mean(pred_verts_shape_canonical_diff) * 1000.0 119 | 120 | return losses 121 | -------------------------------------------------------------------------------- /micalib/renderer.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2023 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: mica@tue.mpg.de 16 | 17 | 18 | import pytorch3d 19 | import torch 20 | import torch.nn as nn 21 | from pytorch3d.io import load_obj 22 | from pytorch3d.renderer import ( 23 | FoVPerspectiveCameras, look_at_view_transform, 24 | RasterizationSettings, MeshRenderer, MeshRasterizer, SoftPhongShader, TexturesVertex 25 | ) 26 | 27 | 28 | class MeshShapeRenderer(nn.Module): 29 | def __init__(self, obj_filename): 30 | super().__init__() 31 | 32 | verts, faces, aux = load_obj(obj_filename) 33 | faces = faces.verts_idx[None, ...].cuda() 34 | self.register_buffer('faces', faces) 35 | 36 | R, T = look_at_view_transform(2.7, 10.0, 10.0) 37 | self.cameras = FoVPerspectiveCameras(device='cuda:0', R=R, T=T, fov=6) 38 | raster_settings = RasterizationSettings( 39 | image_size=512, 40 | blur_radius=0.0, 41 | faces_per_pixel=1, 42 | perspective_correct=True 43 | ) 44 | 45 | lights = pytorch3d.renderer.DirectionalLights( 46 | device='cuda:0', 47 | direction=((0, 0, 1),), 48 | ambient_color=((0.4, 0.4, 0.4),), 49 | diffuse_color=((0.35, 0.35, 0.35),), 50 | specular_color=((0.05, 0.05, 0.05),)) 51 | 52 | self.renderer = MeshRenderer( 53 | rasterizer=MeshRasterizer(cameras=self.cameras, raster_settings=raster_settings), 54 | shader=SoftPhongShader(device='cuda:0', cameras=self.cameras, lights=lights) 55 | ) 56 | 57 | def render_mesh(self, vertices, faces=None, verts_rgb=None): 58 | B, N, V = vertices.shape 59 | if faces is None: 60 | faces = self.faces.repeat(B, 1, 1) 61 | else: 62 | faces = faces.repeat(B, 1, 1) 63 | 64 | if verts_rgb is None: 65 | verts_rgb = torch.ones_like(vertices) 66 | textures = TexturesVertex(verts_features=verts_rgb.cuda()) 67 | meshes = pytorch3d.structures.Meshes(verts=vertices, faces=faces, textures=textures) 68 | 69 | rendering = self.renderer(meshes).permute(0, 3, 1, 2) 70 | color = rendering[:, 0:3, ...] 71 | 72 | return color 73 | -------------------------------------------------------------------------------- /micalib/tester.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2023 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: mica@tue.mpg.de 16 | 17 | 18 | import os 19 | from glob import glob 20 | 21 | import cv2 22 | import numpy as np 23 | import torch 24 | import torch.distributed as dist 25 | from insightface.app import FaceAnalysis 26 | from insightface.app.common import Face 27 | from insightface.utils import face_align 28 | from loguru import logger 29 | from pytorch3d.io import save_ply 30 | from skimage.io import imread 31 | from skimage.transform import estimate_transform, warp 32 | from tqdm import tqdm 33 | 34 | from configs.config import cfg 35 | from utils import util 36 | 37 | input_mean = 127.5 38 | input_std = 127.5 39 | 40 | NOW_SCANS = '/home/wzielonka/datasets/NoWDataset/final_release_version/scans/' 41 | NOW_PICTURES = '/home/wzielonka/datasets/NoWDataset/final_release_version/iphone_pictures/' 42 | NOW_BBOX = '/home/wzielonka/datasets/NoWDataset/final_release_version/detected_face/' 43 | STIRLING_PICTURES = '/home/wzielonka/datasets/Stirling/images/' 44 | 45 | 46 | class Tester(object): 47 | def __init__(self, nfc_model, config=None, device=None): 48 | if config is None: 49 | self.cfg = cfg 50 | else: 51 | self.cfg = config 52 | 53 | self.device = device 54 | self.batch_size = self.cfg.dataset.batch_size 55 | self.K = self.cfg.dataset.K 56 | self.render_mesh = True 57 | self.embeddings_lyhm = {} 58 | 59 | # deca model 60 | self.nfc = nfc_model.to(self.device) 61 | self.nfc.testing = True 62 | 63 | logger.info(f'[INFO] {torch.cuda.get_device_name(device)}') 64 | 65 | def load_checkpoint(self, model_path): 66 | dist.barrier() 67 | map_location = {'cuda:%d' % 0: 'cuda:%d' % self.device} 68 | 69 | checkpoint = torch.load(model_path, map_location) 70 | 71 | if 'arcface' in checkpoint: 72 | self.nfc.arcface.load_state_dict(checkpoint['arcface']) 73 | if 'flameModel' in checkpoint: 74 | self.nfc.flameModel.load_state_dict(checkpoint['flameModel']) 75 | 76 | logger.info(f"[TESTER] Resume from {model_path}") 77 | 78 | def load_model_dict(self, model_dict): 79 | dist.barrier() 80 | 81 | self.nfc.canonicalModel.load_state_dict(model_dict['canonicalModel']) 82 | self.nfc.arcface.load_state_dict(model_dict['arcface']) 83 | 84 | def process_image(self, img, app): 85 | images = [] 86 | bboxes, kpss = app.det_model.detect(img, max_num=0, metric='default') 87 | if bboxes.shape[0] != 1: 88 | logger.error('Face not detected!') 89 | return images 90 | i = 0 91 | bbox = bboxes[i, 0:4] 92 | det_score = bboxes[i, 4] 93 | kps = None 94 | if kpss is not None: 95 | kps = kpss[i] 96 | face = Face(bbox=bbox, kps=kps, det_score=det_score) 97 | aimg = face_align.norm_crop(img, landmark=face.kps) 98 | blob = cv2.dnn.blobFromImages([aimg], 1.0 / input_std, (112, 112), (input_mean, input_mean, input_mean), swapRB=True) 99 | 100 | images.append(torch.tensor(blob[0])[None]) 101 | 102 | return images 103 | 104 | def process_folder(self, folder, app): 105 | images = [] 106 | arcface = [] 107 | files_actor = sorted(sorted(os.listdir(folder))) 108 | for file in files_actor: 109 | image_path = folder + '/' + file 110 | logger.info(image_path) 111 | 112 | ### NOW CROPPING 113 | scale = 1.6 114 | # scale = np.random.rand() * (1.8 - 1.1) + 1.1 115 | bbx_path = image_path.replace('.jpg', '.npy').replace('iphone_pictures', 'detected_face') 116 | bbx_data = np.load(bbx_path, allow_pickle=True, encoding='latin1').item() 117 | left = bbx_data['left'] 118 | right = bbx_data['right'] 119 | top = bbx_data['top'] 120 | bottom = bbx_data['bottom'] 121 | 122 | image = imread(image_path)[:, :, :3] 123 | 124 | h, w, _ = image.shape 125 | old_size = (right - left + bottom - top) / 2 126 | center = np.array([right - (right - left) / 2.0, bottom - (bottom - top) / 2.0]) 127 | size = int(old_size * scale) 128 | 129 | crop_size = 224 130 | # crop image 131 | src_pts = np.array([[center[0] - size / 2, center[1] - size / 2], [center[0] - size / 2, center[1] + size / 2], [center[0] + size / 2, center[1] - size / 2]]) 132 | DST_PTS = np.array([[0, 0], [0, crop_size - 1], [crop_size - 1, 0]]) 133 | tform = estimate_transform('similarity', src_pts, DST_PTS) 134 | 135 | image = image / 255. 136 | dst_image = warp(image, tform.inverse, output_shape=(crop_size, crop_size)) 137 | 138 | arcface += self.process_image(cv2.cvtColor(dst_image.astype(np.float32) * 255.0, cv2.COLOR_RGB2BGR), app) 139 | 140 | dst_image = dst_image.transpose(2, 0, 1) 141 | images.append(torch.tensor(dst_image)[None]) 142 | 143 | images = torch.cat(images, dim=0).float() 144 | arcface = torch.cat(arcface, dim=0).float() 145 | 146 | return images, arcface 147 | 148 | def get_name(self, best_model, id): 149 | if '_' in best_model: 150 | name = id if id is not None else best_model.split('_')[-1][0:-4] 151 | else: 152 | name = id if id is not None else best_model.split('/')[-1][0:-4] 153 | return name 154 | 155 | def test_now(self, best_model, id=None): 156 | self.load_checkpoint(best_model) 157 | name = self.get_name(best_model, id) 158 | self.now(name) 159 | 160 | def test_stirling(self, best_model, id=None): 161 | self.load_checkpoint(best_model) 162 | name = self.get_name(best_model, id) 163 | self.stirling(name) 164 | 165 | def save_mesh(self, file, vertices): 166 | scaled = vertices * 1000.0 167 | save_ply(file, scaled.cpu(), self.nfc.render.faces[0].cpu()) 168 | 169 | # mask = self.nfc.masking.get_triangle_whole_mask() 170 | # v, f = self.nfc.masking.get_masked_mesh(vertices, mask) 171 | # save_obj(file, v[0], f[0]) 172 | 173 | def cache_to_cuda(self, cache): 174 | for key in cache.keys(): 175 | i, a = cache[key] 176 | cache[key] = (i.to(self.device), a.to(self.device)) 177 | return cache 178 | 179 | def create_now_cache(self): 180 | if os.path.exists('test_now_cache.pt'): 181 | cache = self.cache_to_cuda(torch.load('test_now_cache.pt')) 182 | return cache 183 | else: 184 | cache = {} 185 | 186 | app = FaceAnalysis(name='antelopev2', providers=['CUDAExecutionProvider']) 187 | app.prepare(ctx_id=0, det_size=(224, 224), det_thresh=0.4) 188 | 189 | for actor in tqdm(sorted(os.listdir(NOW_PICTURES))): 190 | image_paths = sorted(glob(NOW_PICTURES + actor + '/*')) 191 | for folder in image_paths: 192 | images, arcface = self.process_folder(folder, app) 193 | cache[folder] = (images, arcface) 194 | 195 | torch.save(cache, 'test_now_cache.pt') 196 | return self.cache_to_cuda(cache) 197 | 198 | def create_stirling_cache(self): 199 | if os.path.exists('test_stirling_cache.pt'): 200 | cache = torch.load('test_stirling_cache.pt') 201 | return cache 202 | else: 203 | cache = {} 204 | 205 | app = FaceAnalysis(name='antelopev2', providers=['CUDAExecutionProvider']) 206 | app.prepare(ctx_id=0, det_size=(224, 224), det_thresh=0.1) 207 | 208 | cache['HQ'] = {} 209 | cache['LQ'] = {} 210 | 211 | for folder in ['Real_images__Subset_2D_FG2018']: 212 | for quality in ['HQ', 'LQ']: 213 | for path in tqdm(sorted(glob(STIRLING_PICTURES + folder + '/' + quality + '/*.jpg'))): 214 | actor = path.split('/')[-1][:9].upper() 215 | image = imread(path)[:, :, :3] 216 | blobs = self.process_image(cv2.cvtColor(image, cv2.COLOR_RGB2BGR), app) 217 | if len(blobs) == 0: 218 | continue 219 | image = image / 255. 220 | image = cv2.resize(image, (224, 224)).transpose(2, 0, 1) 221 | image = torch.tensor(image).cuda()[None] 222 | 223 | if actor not in cache[quality]: 224 | cache[quality][actor] = [] 225 | cache[quality][actor].append((image, blobs[0])) 226 | 227 | for q in cache.keys(): 228 | for a in cache[q].keys(): 229 | images, arcface = list(zip(*cache[q][a])) 230 | images = torch.cat(images, dim=0).float() 231 | arcface = torch.cat(arcface, dim=0).float() 232 | cache[q][a] = (images, arcface) 233 | 234 | torch.save(cache, 'test_stirling_cache.pt') 235 | return self.cache_to_cuda(cache) 236 | 237 | def update_embeddings(self, actor, arcface): 238 | if actor not in self.embeddings_lyhm: 239 | self.embeddings_lyhm[actor] = [] 240 | self.embeddings_lyhm[actor] += [arcface[i].data.cpu().numpy() for i in range(arcface.shape[0])] 241 | 242 | def stirling(self, best_id): 243 | logger.info(f"[TESTER] Stirling testing has begun!") 244 | self.nfc.eval() 245 | cache = self.create_stirling_cache() 246 | for quality in cache.keys(): 247 | images_processed = 0 248 | self.embeddings_lyhm = {} 249 | for actor in tqdm(cache[quality].keys()): 250 | images, arcface = cache[quality][actor] 251 | with torch.no_grad(): 252 | codedict = self.nfc.encode(images.cuda(), arcface.cuda()) 253 | opdict = self.nfc.decode(codedict, 0) 254 | 255 | self.update_embeddings(actor, codedict['arcface']) 256 | 257 | dst_actor = actor[:5] 258 | os.makedirs(os.path.join(self.cfg.output_dir, f'stirling_test_{best_id}', 'predicted_meshes', quality), exist_ok=True) 259 | dst_folder = os.path.join(self.cfg.output_dir, f'stirling_test_{best_id}', 'predicted_meshes', quality, dst_actor) 260 | os.makedirs(dst_folder, exist_ok=True) 261 | 262 | meshes = opdict['pred_canonical_shape_vertices'] 263 | lmk = self.nfc.flame.compute_landmarks(meshes) 264 | 265 | for m in range(meshes.shape[0]): 266 | v = torch.reshape(meshes[m], (-1, 3)) 267 | savepath = os.path.join(self.cfg.output_dir, f'stirling_test_{best_id}', 'predicted_meshes', quality, dst_actor, f'{actor}.ply') 268 | self.save_mesh(savepath, v) 269 | landmark_51 = lmk[m, 17:] 270 | landmark_7 = landmark_51[[19, 22, 25, 28, 16, 31, 37]] 271 | landmark_7 = landmark_7.cpu().numpy() * 1000.0 272 | np.save(os.path.join(self.cfg.output_dir, f'stirling_test_{best_id}', 'predicted_meshes', quality, dst_actor, f'{actor}.npy'), landmark_7) 273 | images_processed += 1 274 | 275 | pred = self.nfc.render.render_mesh(meshes) 276 | dict = { 277 | 'pred': pred, 278 | 'images': images 279 | } 280 | 281 | savepath = os.path.join(self.cfg.output_dir, f'stirling_test_{best_id}', 'predicted_meshes', quality, dst_actor, f'{actor}.jpg') 282 | util.visualize_grid(dict, savepath, size=512) 283 | 284 | logger.info(f"[TESTER] Stirling dataset {quality} with {images_processed} processed!") 285 | 286 | # util.save_embedding_projection(self.embeddings_lyhm, f'{self.cfg.output_dir}/stirling_test_{best_id}/stirling_{quality}_arcface_embeds.pdf') 287 | 288 | def now(self, best_id): 289 | logger.info(f"[TESTER] NoW testing has begun!") 290 | self.nfc.eval() 291 | cache = self.create_now_cache() 292 | # for actor in tqdm(sorted(os.listdir(NOW_SCANS))): # only 20 293 | for actor in tqdm(sorted(os.listdir(NOW_PICTURES))): # all 100 294 | image_paths = sorted(glob(NOW_PICTURES + actor + '/*')) 295 | for folder in image_paths: 296 | files_actor = sorted(os.listdir(folder)) 297 | images, arcface = cache[folder] 298 | with torch.no_grad(): 299 | codedict = self.nfc.encode(images, arcface) 300 | opdict = self.nfc.decode(codedict, 0) 301 | 302 | self.update_embeddings(actor.split('_')[2], codedict['arcface']) 303 | 304 | type = folder.split('/')[-1] 305 | os.makedirs(os.path.join(self.cfg.output_dir, f'now_test_{best_id}', 'predicted_meshes'), exist_ok=True) 306 | dst_folder = os.path.join(self.cfg.output_dir, f'now_test_{best_id}', 'predicted_meshes', actor, type) 307 | os.makedirs(dst_folder, exist_ok=True) 308 | 309 | meshes = opdict['pred_canonical_shape_vertices'] 310 | lmk = self.nfc.flame.compute_landmarks(meshes) 311 | 312 | for m in range(meshes.shape[0]): 313 | a = files_actor[m] 314 | v = torch.reshape(meshes[m], (-1, 3)) 315 | savepath = os.path.join(self.cfg.output_dir, f'now_test_{best_id}', 'predicted_meshes', actor, type, a.replace('jpg', 'ply')) 316 | self.save_mesh(savepath, v) 317 | landmark_51 = lmk[m, 17:] 318 | landmark_7 = landmark_51[[19, 22, 25, 28, 16, 31, 37]] 319 | landmark_7 = landmark_7.cpu().numpy() * 1000.0 320 | np.save(os.path.join(self.cfg.output_dir, f'now_test_{best_id}', 'predicted_meshes', actor, type, a.replace('jpg', 'npy')), landmark_7) 321 | 322 | if self.render_mesh: 323 | pred = self.nfc.render.render_mesh(meshes) 324 | 325 | dict = { 326 | 'pred': pred, 327 | # 'deca': deca, 328 | 'images': images 329 | } 330 | 331 | savepath = os.path.join(dst_folder, 'render.jpg') 332 | util.visualize_grid(dict, savepath, size=512) 333 | 334 | # util.save_embedding_projection(self.embeddings_lyhm, f'{self.cfg.output_dir}/now_test_{best_id}/now_arcface_embeds.pdf') 335 | -------------------------------------------------------------------------------- /micalib/trainer.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2023 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: mica@tue.mpg.de 16 | 17 | 18 | import os 19 | import random 20 | import sys 21 | from datetime import datetime 22 | 23 | import numpy as np 24 | import torch 25 | import torch.distributed as dist 26 | from loguru import logger 27 | from torch.utils.data import DataLoader 28 | from tqdm import tqdm 29 | 30 | import datasets 31 | from configs.config import cfg 32 | from utils import util 33 | 34 | sys.path.append("./micalib") 35 | from validator import Validator 36 | 37 | 38 | def print_info(rank): 39 | props = torch.cuda.get_device_properties(rank) 40 | 41 | logger.info(f'[INFO] {torch.cuda.get_device_name(rank)}') 42 | logger.info(f'[INFO] Rank: {str(rank)}') 43 | logger.info(f'[INFO] Memory: {round(props.total_memory / 1024 ** 3, 1)} GB') 44 | logger.info(f'[INFO] Allocated: {round(torch.cuda.memory_allocated(rank) / 1024 ** 3, 1)} GB') 45 | logger.info(f'[INFO] Cached: {round(torch.cuda.memory_reserved(rank) / 1024 ** 3, 1)} GB') 46 | 47 | 48 | def seed_worker(worker_id): 49 | worker_seed = torch.initial_seed() % 2 ** 32 50 | np.random.seed(worker_seed) 51 | random.seed(worker_seed) 52 | 53 | 54 | class Trainer(object): 55 | def __init__(self, nfc_model, config=None, device=None): 56 | if config is None: 57 | self.cfg = cfg 58 | else: 59 | self.cfg = config 60 | 61 | logger.add(os.path.join(self.cfg.output_dir, self.cfg.train.log_dir, 'train.log')) 62 | 63 | self.device = device 64 | self.batch_size = self.cfg.dataset.batch_size 65 | self.K = self.cfg.dataset.K 66 | 67 | # deca model 68 | self.nfc = nfc_model.to(self.device) 69 | 70 | self.validator = Validator(self) 71 | self.configure_optimizers() 72 | self.load_checkpoint() 73 | 74 | # reset optimizer if loaded from pretrained model 75 | if self.cfg.train.reset_optimizer: 76 | self.configure_optimizers() # reset optimizer 77 | logger.info(f"[TRAINER] Optimizer was reset") 78 | 79 | if self.cfg.train.write_summary and self.device == 0: 80 | from torch.utils.tensorboard import SummaryWriter 81 | self.writer = SummaryWriter(log_dir=os.path.join(self.cfg.output_dir, self.cfg.train.log_dir)) 82 | 83 | print_info(device) 84 | 85 | def configure_optimizers(self): 86 | self.opt = torch.optim.AdamW( 87 | lr=self.cfg.train.lr, 88 | weight_decay=self.cfg.train.weight_decay, 89 | params=self.nfc.parameters_to_optimize(), 90 | amsgrad=False) 91 | 92 | self.scheduler = torch.optim.lr_scheduler.StepLR(self.opt, step_size=1, gamma=0.1) 93 | 94 | def load_checkpoint(self): 95 | self.epoch = 0 96 | self.global_step = 0 97 | dist.barrier() 98 | map_location = {'cuda:%d' % 0: 'cuda:%d' % self.device} 99 | model_path = os.path.join(self.cfg.output_dir, 'model.tar') 100 | if os.path.exists(self.cfg.pretrained_model_path): 101 | model_path = self.cfg.pretrained_model_path 102 | if os.path.exists(model_path): 103 | checkpoint = torch.load(model_path, map_location) 104 | if 'opt' in checkpoint: 105 | self.opt.load_state_dict(checkpoint['opt']) 106 | if 'scheduler' in checkpoint: 107 | self.scheduler.load_state_dict(checkpoint['scheduler']) 108 | if 'epoch' in checkpoint: 109 | self.epoch = checkpoint['epoch'] 110 | if 'global_step' in checkpoint: 111 | self.global_step = checkpoint['global_step'] 112 | logger.info(f"[TRAINER] Resume training from {model_path}") 113 | logger.info(f"[TRAINER] Start from step {self.global_step}") 114 | logger.info(f"[TRAINER] Start from epoch {self.epoch}") 115 | else: 116 | logger.info('[TRAINER] Model path not found, start training from scratch') 117 | 118 | def save_checkpoint(self, filename): 119 | if self.device == 0: 120 | model_dict = self.nfc.model_dict() 121 | 122 | model_dict['opt'] = self.opt.state_dict() 123 | model_dict['scheduler'] = self.scheduler.state_dict() 124 | model_dict['validator'] = self.validator.state_dict() 125 | model_dict['epoch'] = self.epoch 126 | model_dict['global_step'] = self.global_step 127 | model_dict['batch_size'] = self.batch_size 128 | 129 | torch.save(model_dict, filename) 130 | 131 | def training_step(self, batch): 132 | self.nfc.train() 133 | 134 | images = batch['image'].to(self.device) 135 | images = images.view(-1, images.shape[-3], images.shape[-2], images.shape[-1]) 136 | flame = batch['flame'] 137 | arcface = batch['arcface'] 138 | arcface = arcface.view(-1, arcface.shape[-3], arcface.shape[-2], arcface.shape[-1]).to(self.device) 139 | 140 | inputs = { 141 | 'images': images, 142 | 'dataset': batch['dataset'][0] 143 | } 144 | 145 | encoder_output = self.nfc.encode(images, arcface) 146 | encoder_output['flame'] = flame 147 | 148 | decoder_output = self.nfc.decode(encoder_output, self.epoch) 149 | losses = self.nfc.compute_losses(inputs, encoder_output, decoder_output) 150 | 151 | all_loss = 0. 152 | losses_key = losses.keys() 153 | 154 | for key in losses_key: 155 | all_loss = all_loss + losses[key] 156 | 157 | losses['all_loss'] = all_loss 158 | 159 | opdict = \ 160 | { 161 | 'images': images, 162 | 'flame_verts_shape': decoder_output['flame_verts_shape'], 163 | 'pred_canonical_shape_vertices': decoder_output['pred_canonical_shape_vertices'], 164 | } 165 | 166 | if 'deca' in decoder_output: 167 | opdict['deca'] = decoder_output['deca'] 168 | 169 | return losses, opdict 170 | 171 | def validation_step(self): 172 | self.validator.run() 173 | 174 | def evaluation_step(self): 175 | pass 176 | 177 | def prepare_data(self): 178 | generator = torch.Generator() 179 | generator.manual_seed(self.device) 180 | 181 | self.train_dataset, total_images = datasets.build_train(self.cfg.dataset, self.device) 182 | self.train_dataloader = DataLoader( 183 | self.train_dataset, batch_size=self.batch_size, 184 | num_workers=self.cfg.dataset.num_workers, 185 | shuffle=True, 186 | pin_memory=True, 187 | drop_last=False, 188 | worker_init_fn=seed_worker, 189 | generator=generator) 190 | 191 | self.train_iter = iter(self.train_dataloader) 192 | logger.info(f'[TRAINER] Training dataset is ready with {len(self.train_dataset)} actors and {total_images} images.') 193 | 194 | def fit(self): 195 | self.prepare_data() 196 | iters_every_epoch = int(len(self.train_dataset) / self.batch_size) 197 | max_epochs = int(self.cfg.train.max_steps / iters_every_epoch) 198 | start_epoch = self.epoch 199 | for epoch in range(start_epoch, max_epochs): 200 | for step in tqdm(range(iters_every_epoch), desc=f"Epoch[{epoch + 1}/{max_epochs}]"): 201 | if self.global_step > self.cfg.train.max_steps: 202 | break 203 | try: 204 | batch = next(self.train_iter) 205 | except Exception as e: 206 | self.train_iter = iter(self.train_dataloader) 207 | batch = next(self.train_iter) 208 | 209 | visualizeTraining = self.global_step % self.cfg.train.vis_steps == 0 210 | 211 | self.opt.zero_grad() 212 | losses, opdict = self.training_step(batch) 213 | 214 | all_loss = losses['all_loss'] 215 | all_loss.backward() 216 | self.opt.step() 217 | self.global_step += 1 218 | 219 | if self.global_step % self.cfg.train.log_steps == 0 and self.device == 0: 220 | loss_info = f"\n" \ 221 | f" Epoch: {epoch}\n" \ 222 | f" Step: {self.global_step}\n" \ 223 | f" Iter: {step}/{iters_every_epoch}\n" \ 224 | f" LR: {self.opt.param_groups[0]['lr']}\n" \ 225 | f" Time: {datetime.now().strftime('%Y-%m-%d-%H:%M:%S')}\n" 226 | for k, v in losses.items(): 227 | loss_info = loss_info + f' {k}: {v:.4f}\n' 228 | if self.cfg.train.write_summary: 229 | self.writer.add_scalar('train_loss/' + k, v, global_step=self.global_step) 230 | logger.info(loss_info) 231 | 232 | if visualizeTraining and self.device == 0: 233 | visdict = { 234 | 'input_images': opdict['images'], 235 | } 236 | # add images to tensorboard 237 | for k, v in visdict.items(): 238 | self.writer.add_images(k, np.clip(v.detach().cpu(), 0.0, 1.0), self.global_step) 239 | 240 | pred_canonical_shape_vertices = torch.empty(0, 3, 512, 512).cuda() 241 | flame_verts_shape = torch.empty(0, 3, 512, 512).cuda() 242 | deca_images = torch.empty(0, 3, 512, 512).cuda() 243 | input_images = torch.empty(0, 3, 224, 224).cuda() 244 | L = opdict['pred_canonical_shape_vertices'].shape[0] 245 | S = 4 if L > 4 else L 246 | for n in np.random.choice(range(L), size=S, replace=False): 247 | rendering = self.nfc.render.render_mesh(opdict['pred_canonical_shape_vertices'][n:n + 1, ...]) 248 | pred_canonical_shape_vertices = torch.cat([pred_canonical_shape_vertices, rendering]) 249 | rendering = self.nfc.render.render_mesh(opdict['flame_verts_shape'][n:n + 1, ...]) 250 | flame_verts_shape = torch.cat([flame_verts_shape, rendering]) 251 | input_images = torch.cat([input_images, opdict['images'][n:n + 1, ...]]) 252 | if 'deca' in opdict: 253 | deca = self.nfc.render.render_mesh(opdict['deca'][n:n + 1, ...]) 254 | deca_images = torch.cat([deca_images, deca]) 255 | 256 | visdict = {} 257 | 258 | if 'deca' in opdict: 259 | visdict['deca'] = deca_images 260 | 261 | visdict["pred_canonical_shape_vertices"] = pred_canonical_shape_vertices 262 | visdict["flame_verts_shape"] = flame_verts_shape 263 | visdict["images"] = input_images 264 | 265 | savepath = os.path.join(self.cfg.output_dir, 'train_images/train_' + str(epoch) + '.jpg') 266 | util.visualize_grid(visdict, savepath, size=512) 267 | 268 | if self.global_step % self.cfg.train.val_steps == 0: 269 | self.validation_step() 270 | 271 | if self.global_step % self.cfg.train.lr_update_step == 0: 272 | self.scheduler.step() 273 | 274 | if self.global_step % self.cfg.train.eval_steps == 0: 275 | self.evaluation_step() 276 | 277 | if self.global_step % self.cfg.train.checkpoint_steps == 0: 278 | self.save_checkpoint(os.path.join(self.cfg.output_dir, 'model' + '.tar')) 279 | 280 | if self.global_step % self.cfg.train.checkpoint_epochs_steps == 0: 281 | self.save_checkpoint(os.path.join(self.cfg.output_dir, 'model_' + str(self.global_step) + '.tar')) 282 | 283 | self.epoch += 1 284 | 285 | self.save_checkpoint(os.path.join(self.cfg.output_dir, 'model' + '.tar')) 286 | logger.info(f'[TRAINER] Fitting has ended!') 287 | -------------------------------------------------------------------------------- /micalib/validator.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2023 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: mica@tue.mpg.de 16 | 17 | 18 | import os 19 | import subprocess 20 | from copy import deepcopy 21 | from datetime import datetime 22 | 23 | import numpy as np 24 | import torch 25 | from loguru import logger 26 | from torch.utils.data import DataLoader 27 | 28 | import datasets 29 | from utils import util 30 | from utils.best_model import BestModel 31 | 32 | 33 | class Validator(object): 34 | def __init__(self, trainer): 35 | self.trainer = trainer 36 | self.device = self.trainer.device 37 | self.nfc = self.trainer.nfc 38 | self.cfg = deepcopy(self.trainer.cfg) 39 | self.device = trainer.device 40 | 41 | # Create a separate instance only for predictions 42 | # nfc = util.find_model_using_name(model_dir='nfclib.models', model_name=self.cfg.model.name)(self.cfg, self.device) 43 | # self.tester = Tester(nfc, self.cfg, self.device) 44 | # self.tester.render_mesh = False 45 | 46 | self.embeddings_lyhm = {} 47 | self.best_model = BestModel(trainer) 48 | self.prepare_data() 49 | 50 | def prepare_data(self): 51 | self.val_dataset, total_images = datasets.build_val(self.cfg.dataset, self.device) 52 | self.val_dataloader = DataLoader( 53 | self.val_dataset, 54 | batch_size=2, 55 | shuffle=False, 56 | num_workers=4, 57 | pin_memory=True, 58 | drop_last=False) 59 | 60 | self.val_iter = iter(self.val_dataloader) 61 | logger.info(f'[VALIDATOR] Validation dataset is ready with {len(self.val_dataset)} actors and {total_images} images.') 62 | 63 | def state_dict(self): 64 | return { 65 | 'embeddings_lyhm': self.embeddings_lyhm, 66 | 'best_model': self.best_model.state_dict(), 67 | } 68 | 69 | def load_state_dict(self, dict): 70 | self.embeddings_lyhm = dict['embeddings_lyhm'] 71 | self.best_model.load_state_dict(dict['best_model']) 72 | 73 | def update_embeddings(self, actors, arcface): 74 | B = len(actors) 75 | for i in range(B): 76 | actor = actors[i] 77 | if actor not in self.embeddings_lyhm: 78 | self.embeddings_lyhm[actor] = [] 79 | self.embeddings_lyhm[actor].append(arcface[i].data.cpu().numpy()) 80 | 81 | def run(self): 82 | with torch.no_grad(): 83 | # In the case of using multiple GPUs 84 | if self.trainer.device != 0: 85 | return 86 | 87 | self.nfc.eval() 88 | optdicts = [] 89 | while True: 90 | try: 91 | batch = next(self.val_iter) 92 | except Exception as e: 93 | print(e) 94 | self.val_iter = iter(self.val_dataloader) 95 | break 96 | 97 | actors = batch['imagename'] 98 | dataset = batch['dataset'] 99 | images = batch['image'].cuda() 100 | images = images.view(-1, images.shape[-3], images.shape[-2], images.shape[-1]) 101 | arcface = batch['arcface'].cuda() 102 | arcface = arcface.view(-1, arcface.shape[-3], arcface.shape[-2], arcface.shape[-1]).to(self.device) 103 | flame = batch['flame'] 104 | 105 | codedict = self.nfc.encode(images, arcface) 106 | codedict['flame'] = flame 107 | opdict = self.nfc.decode(codedict, self.trainer.epoch) 108 | self.update_embeddings(actors, opdict['faceid']) 109 | loss = self.nfc.compute_losses(None, None, opdict)['pred_verts_shape_canonical_diff'] 110 | optdicts.append((opdict, images, dataset, actors, loss)) 111 | 112 | # Calculate averages 113 | weighted_average = 0. 114 | average = 0. 115 | avg_per_dataset = {} 116 | for optdict in optdicts: 117 | opdict, images, dataset, actors, loss = optdict 118 | name = dataset[0] 119 | average += loss 120 | if name not in avg_per_dataset: 121 | avg_per_dataset[name] = (loss, 1.) 122 | else: 123 | l, i = avg_per_dataset[name] 124 | avg_per_dataset[name] = (l + loss, i + 1.) 125 | 126 | average = average.item() / len(optdicts) 127 | 128 | loss_info = f"Step: {self.trainer.global_step}, Time: {datetime.now().strftime('%Y-%m-%d-%H:%M:%S')} \n" 129 | loss_info += f' validation loss (average) : {average:.5f} \n' 130 | logger.info(loss_info) 131 | 132 | self.trainer.writer.add_scalar('val/average', average, global_step=self.trainer.global_step) 133 | for key in avg_per_dataset.keys(): 134 | l, i = avg_per_dataset[key] 135 | avg = l.item() / i 136 | self.trainer.writer.add_scalar(f'val/average_{key}', avg, global_step=self.trainer.global_step) 137 | 138 | # Save best model 139 | smoothed_weighted, smoothed = self.best_model(weighted_average, average) 140 | self.trainer.writer.add_scalar(f'val/smoothed_average', smoothed, global_step=self.trainer.global_step) 141 | 142 | # self.now() 143 | 144 | # Print embeddings every nth validation step 145 | if self.trainer.global_step % (self.cfg.train.val_steps * 5) == 0: 146 | lyhm_keys = list(self.embeddings_lyhm.keys()) 147 | embeddings = {**{key: self.embeddings_lyhm[key] for key in lyhm_keys}} 148 | # util.save_embedding_projection(embeddings, os.path.join(self.cfg.output_dir, self.cfg.train.val_vis_dir, f'{self.trainer.global_step:08}_embeddings.jpg')) 149 | self.embeddings_lyhm = {} 150 | 151 | # Render predicted meshes 152 | if self.trainer.global_step % self.cfg.train.val_save_img != 0: 153 | return 154 | 155 | pred_canonical_shape_vertices = torch.empty(0, 3, 512, 512).cuda() 156 | flame_verts_shape = torch.empty(0, 3, 512, 512).cuda() 157 | input_images = torch.empty(0, 3, 224, 224).cuda() 158 | 159 | for i in np.random.choice(range(0, len(optdicts)), size=4, replace=False): 160 | opdict, images, _, _, _ = optdicts[i] 161 | n = np.random.randint(0, len(images) - 1) 162 | rendering = self.nfc.render.render_mesh(opdict['pred_canonical_shape_vertices'][n:n + 1, ...]) 163 | pred_canonical_shape_vertices = torch.cat([pred_canonical_shape_vertices, rendering]) 164 | rendering = self.nfc.render.render_mesh(opdict['flame_verts_shape'][n:n + 1, ...]) 165 | flame_verts_shape = torch.cat([flame_verts_shape, rendering]) 166 | input_images = torch.cat([input_images, images[n:n + 1, ...]]) 167 | 168 | visdict = { 169 | "pred_canonical_shape_vertices": pred_canonical_shape_vertices, 170 | "flame_verts_shape": flame_verts_shape, 171 | "input": input_images 172 | } 173 | 174 | savepath = os.path.join(self.cfg.output_dir, self.cfg.train.val_vis_dir, f'{self.trainer.global_step:08}.jpg') 175 | util.visualize_grid(visdict, savepath, size=512) 176 | 177 | def now(self): 178 | logger.info(f'[Validator] NoW testing has begun...') 179 | # self.tester.test_now('', 'training', self.nfc.model_dict()) 180 | root = os.path.abspath(os.path.join(os.path.dirname(__file__), '../')) 181 | path = f'{root}{self.cfg.output_dir[1:]}/now_test_training/predicted_meshes' 182 | cmd = f'./now_validation.sh {path}' 183 | subprocess.call(cmd, shell=True) 184 | errors = np.load(f'{path}/results/_computed_distances.npy', allow_pickle=True, encoding="latin1").item()['computed_distances'] 185 | median = np.median(np.hstack(errors)) 186 | mean = np.mean(np.hstack(errors)) 187 | std = np.std(np.hstack(errors)) 188 | 189 | self.best_model.now(median, mean, std) 190 | 191 | self.trainer.writer.add_scalar(f'val/now_mean', mean, global_step=self.trainer.global_step) 192 | logger.info(f'[Validator] NoW testing has ended...') 193 | -------------------------------------------------------------------------------- /models/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Zielon/MICA/af22e7a5810d474bc28a1433db533723d6bd2b07/models/__init__.py -------------------------------------------------------------------------------- /models/arcface.py: -------------------------------------------------------------------------------- 1 | ### Taken from https://github.com/deepinsight/insightface/blob/master/recognition/arcface_torch/backbones/iresnet.py 2 | 3 | import os 4 | 5 | import torch 6 | from loguru import logger 7 | from torch import nn 8 | 9 | __all__ = ['iresnet18', 'iresnet34', 'iresnet50', 'iresnet100', 'iresnet200'] 10 | 11 | 12 | def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1): 13 | """3x3 convolution with padding""" 14 | return nn.Conv2d(in_planes, 15 | out_planes, 16 | kernel_size=3, 17 | stride=stride, 18 | padding=dilation, 19 | groups=groups, 20 | bias=False, 21 | dilation=dilation) 22 | 23 | 24 | def conv1x1(in_planes, out_planes, stride=1): 25 | """1x1 convolution""" 26 | return nn.Conv2d(in_planes, 27 | out_planes, 28 | kernel_size=1, 29 | stride=stride, 30 | bias=False) 31 | 32 | 33 | class IBasicBlock(nn.Module): 34 | expansion = 1 35 | 36 | def __init__(self, inplanes, planes, stride=1, downsample=None, 37 | groups=1, base_width=64, dilation=1): 38 | super(IBasicBlock, self).__init__() 39 | if groups != 1 or base_width != 64: 40 | raise ValueError('BasicBlock only supports groups=1 and base_width=64') 41 | if dilation > 1: 42 | raise NotImplementedError("Dilation > 1 not supported in BasicBlock") 43 | self.bn1 = nn.BatchNorm2d(inplanes, eps=1e-05, ) 44 | self.conv1 = conv3x3(inplanes, planes) 45 | self.bn2 = nn.BatchNorm2d(planes, eps=1e-05, ) 46 | self.prelu = nn.PReLU(planes) 47 | self.conv2 = conv3x3(planes, planes, stride) 48 | self.bn3 = nn.BatchNorm2d(planes, eps=1e-05, ) 49 | self.downsample = downsample 50 | self.stride = stride 51 | 52 | def forward(self, x): 53 | identity = x 54 | out = self.bn1(x) 55 | out = self.conv1(out) 56 | out = self.bn2(out) 57 | out = self.prelu(out) 58 | out = self.conv2(out) 59 | out = self.bn3(out) 60 | if self.downsample is not None: 61 | identity = self.downsample(x) 62 | out += identity 63 | return out 64 | 65 | 66 | class IResNet(nn.Module): 67 | fc_scale = 7 * 7 68 | 69 | def __init__(self, 70 | block, layers, dropout=0, num_features=512, zero_init_residual=False, 71 | groups=1, width_per_group=64, replace_stride_with_dilation=None, fp16=False): 72 | super(IResNet, self).__init__() 73 | self.fp16 = fp16 74 | self.inplanes = 64 75 | self.dilation = 1 76 | self.block = block 77 | if replace_stride_with_dilation is None: 78 | replace_stride_with_dilation = [False, False, False] 79 | if len(replace_stride_with_dilation) != 3: 80 | raise ValueError("replace_stride_with_dilation should be None " 81 | "or a 3-element tuple, got {}".format(replace_stride_with_dilation)) 82 | self.groups = groups 83 | self.base_width = width_per_group 84 | self.conv1 = nn.Conv2d(3, self.inplanes, kernel_size=3, stride=1, padding=1, bias=False) 85 | self.bn1 = nn.BatchNorm2d(self.inplanes, eps=1e-05) 86 | self.prelu = nn.PReLU(self.inplanes) 87 | self.layer1 = self._make_layer(block, 64, layers[0], stride=2) 88 | self.layer2 = self._make_layer(block, 89 | 128, 90 | layers[1], 91 | stride=2, 92 | dilate=replace_stride_with_dilation[0]) 93 | self.layer3 = self._make_layer(block, 94 | 256, 95 | layers[2], 96 | stride=2, 97 | dilate=replace_stride_with_dilation[1]) 98 | self.layer4 = self._make_layer(block, 99 | 512, 100 | layers[3], 101 | stride=2, 102 | dilate=replace_stride_with_dilation[2]) 103 | self.bn2 = nn.BatchNorm2d(512 * block.expansion, eps=1e-05, ) 104 | self.dropout = nn.Dropout(p=dropout, inplace=True) 105 | self.fc = nn.Linear(512 * block.expansion * self.fc_scale, num_features) 106 | self.features = nn.BatchNorm1d(num_features, eps=1e-05) 107 | nn.init.constant_(self.features.weight, 1.0) 108 | self.features.weight.requires_grad = False 109 | 110 | for m in self.modules(): 111 | if isinstance(m, nn.Conv2d): 112 | nn.init.normal_(m.weight, 0, 0.1) 113 | elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)): 114 | nn.init.constant_(m.weight, 1) 115 | nn.init.constant_(m.bias, 0) 116 | 117 | if zero_init_residual: 118 | for m in self.modules(): 119 | if isinstance(m, IBasicBlock): 120 | nn.init.constant_(m.bn2.weight, 0) 121 | 122 | def _make_layer(self, block, planes, blocks, stride=1, dilate=False): 123 | downsample = None 124 | previous_dilation = self.dilation 125 | if dilate: 126 | self.dilation *= stride 127 | stride = 1 128 | if stride != 1 or self.inplanes != planes * block.expansion: 129 | downsample = nn.Sequential( 130 | conv1x1(self.inplanes, planes * block.expansion, stride), 131 | nn.BatchNorm2d(planes * block.expansion, eps=1e-05, ), 132 | ) 133 | layers = [] 134 | layers.append( 135 | block(self.inplanes, planes, stride, downsample, self.groups, 136 | self.base_width, previous_dilation)) 137 | self.inplanes = planes * block.expansion 138 | for _ in range(1, blocks): 139 | layers.append( 140 | block(self.inplanes, 141 | planes, 142 | groups=self.groups, 143 | base_width=self.base_width, 144 | dilation=self.dilation)) 145 | 146 | return nn.Sequential(*layers) 147 | 148 | def forward(self, x): 149 | with torch.cuda.amp.autocast(self.fp16): 150 | x = self.conv1(x) 151 | x = self.bn1(x) 152 | x = self.prelu(x) 153 | x = self.layer1(x) 154 | x = self.layer2(x) 155 | x = self.layer3(x) 156 | x = self.layer4(x) 157 | x = self.bn2(x) 158 | x = torch.flatten(x, 1) 159 | x = self.dropout(x) 160 | x = self.fc(x.float() if self.fp16 else x) 161 | x = self.features(x) 162 | return x 163 | 164 | 165 | class Arcface(IResNet): 166 | def __init__(self, pretrained_path=None, **kwargs): 167 | super(Arcface, self).__init__(IBasicBlock, [3, 13, 30, 3], **kwargs) 168 | if pretrained_path is not None and os.path.exists(pretrained_path): 169 | logger.info(f'[Arcface] Initializing from insightface model from {pretrained_path}.') 170 | self.load_state_dict(torch.load(pretrained_path)) 171 | self.freezer([self.layer1, self.layer2, self.layer3, self.conv1, self.bn1, self.prelu]) 172 | 173 | def freezer(self, layers): 174 | for layer in layers: 175 | for block in layer.parameters(): 176 | block.requires_grad = False 177 | 178 | def forward(self, images): 179 | x = self.forward_arcface(images) 180 | return x 181 | 182 | def forward_arcface(self, x): 183 | with torch.cuda.amp.autocast(self.fp16): 184 | ### FROZEN ### 185 | with torch.no_grad(): 186 | x = self.conv1(x) 187 | x = self.bn1(x) 188 | x = self.prelu(x) 189 | x = self.layer1(x) 190 | x = self.layer2(x) 191 | x = self.layer3(x) 192 | 193 | x = self.layer4(x) 194 | x = self.bn2(x) 195 | x = torch.flatten(x, 1) 196 | x = self.dropout(x) 197 | x = self.fc(x.float() if self.fp16 else x) 198 | x = self.features(x) 199 | return x 200 | -------------------------------------------------------------------------------- /models/flame.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2023 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # For commercial licensing contact, please contact ps-license@tuebingen.mpg.de 16 | 17 | 18 | import pickle 19 | 20 | import loguru 21 | import numpy as np 22 | import torch 23 | import torch.nn as nn 24 | 25 | from .lbs import lbs, batch_rodrigues, vertices2landmarks, rot_mat_to_euler 26 | 27 | 28 | def to_tensor(array, dtype=torch.float32): 29 | if 'torch.tensor' not in str(type(array)): 30 | return torch.tensor(array, dtype=dtype) 31 | 32 | 33 | def to_np(array, dtype=np.float32): 34 | if 'scipy.sparse' in str(type(array)): 35 | array = array.todense() 36 | return np.array(array, dtype=dtype) 37 | 38 | 39 | class Struct(object): 40 | def __init__(self, **kwargs): 41 | for key, val in kwargs.items(): 42 | setattr(self, key, val) 43 | 44 | 45 | class FLAME(nn.Module): 46 | """ 47 | borrowed from https://github.com/soubhiksanyal/FLAME_PyTorch/blob/master/FLAME.py 48 | Given flame parameters this class generates a differentiable FLAME function 49 | which outputs the a mesh and 2D/3D facial landmarks 50 | """ 51 | 52 | def __init__(self, config, optimize_basis=False): 53 | super(FLAME, self).__init__() 54 | loguru.logger.info("[FLAME] creating the FLAME Decoder") 55 | with open(config.flame_model_path, 'rb') as f: 56 | ss = pickle.load(f, encoding='latin1') 57 | flame_model = Struct(**ss) 58 | 59 | self.optimize_basis = optimize_basis 60 | self.cfg = config 61 | self.dtype = torch.float32 62 | self.register_buffer('faces_tensor', to_tensor(to_np(flame_model.f, dtype=np.int64), dtype=torch.long)) 63 | # The vertices of the template model 64 | self.register_buffer('v_template', to_tensor(to_np(flame_model.v_template), dtype=self.dtype)) 65 | self.n_vertices = self.v_template.shape[0] 66 | # The shape components and expression 67 | shapedirs = to_tensor(to_np(flame_model.shapedirs), dtype=self.dtype) 68 | shapedirs = torch.cat([shapedirs[:, :, :config.n_shape], shapedirs[:, :, 300:]], 2) 69 | 70 | if optimize_basis: 71 | self.register_parameter('shapedirs', torch.nn.Parameter(shapedirs)) 72 | else: 73 | self.register_buffer('shapedirs', shapedirs) 74 | 75 | self.n_shape = config.n_shape 76 | # The pose components 77 | num_pose_basis = flame_model.posedirs.shape[-1] 78 | posedirs = np.reshape(flame_model.posedirs, [-1, num_pose_basis]).T 79 | self.register_buffer('posedirs', to_tensor(to_np(posedirs), dtype=self.dtype)) 80 | # 81 | self.register_buffer('J_regressor', to_tensor(to_np(flame_model.J_regressor), dtype=self.dtype)) 82 | parents = to_tensor(to_np(flame_model.kintree_table[0])).long(); 83 | parents[0] = -1 84 | self.register_buffer('parents', parents) 85 | self.register_buffer('lbs_weights', to_tensor(to_np(flame_model.weights), dtype=self.dtype)) 86 | 87 | # Fixing Eyeball and neck rotation 88 | default_eyball_pose = torch.zeros([1, 6], dtype=self.dtype, requires_grad=False) 89 | self.register_parameter('eye_pose', nn.Parameter(default_eyball_pose, requires_grad=False)) 90 | default_neck_pose = torch.zeros([1, 3], dtype=self.dtype, requires_grad=False) 91 | self.register_parameter('neck_pose', nn.Parameter(default_neck_pose, requires_grad=False)) 92 | 93 | # Static and Dynamic Landmark embeddings for FLAME 94 | lmk_embeddings = np.load(config.flame_lmk_embedding_path, allow_pickle=True, encoding='latin1') 95 | lmk_embeddings = lmk_embeddings[()] 96 | self.register_buffer('lmk_faces_idx', torch.from_numpy(lmk_embeddings['static_lmk_faces_idx']).long()) 97 | self.register_buffer('lmk_bary_coords', torch.from_numpy(lmk_embeddings['static_lmk_bary_coords']).to(self.dtype)) 98 | self.register_buffer('dynamic_lmk_faces_idx', lmk_embeddings['dynamic_lmk_faces_idx'].long()) 99 | self.register_buffer('dynamic_lmk_bary_coords', lmk_embeddings['dynamic_lmk_bary_coords'].to(self.dtype)) 100 | self.register_buffer('full_lmk_faces_idx', torch.from_numpy(lmk_embeddings['full_lmk_faces_idx']).long()) 101 | self.register_buffer('full_lmk_bary_coords', torch.from_numpy(lmk_embeddings['full_lmk_bary_coords']).to(self.dtype)) 102 | 103 | neck_kin_chain = []; 104 | NECK_IDX = 1 105 | curr_idx = torch.tensor(NECK_IDX, dtype=torch.long) 106 | while curr_idx != -1: 107 | neck_kin_chain.append(curr_idx) 108 | curr_idx = self.parents[curr_idx] 109 | self.register_buffer('neck_kin_chain', torch.stack(neck_kin_chain)) 110 | 111 | def _find_dynamic_lmk_idx_and_bcoords(self, pose, dynamic_lmk_faces_idx, 112 | dynamic_lmk_b_coords, 113 | neck_kin_chain, dtype=torch.float32): 114 | """ 115 | Selects the face contour depending on the reletive position of the head 116 | Input: 117 | vertices: N X num_of_vertices X 3 118 | pose: N X full pose 119 | dynamic_lmk_faces_idx: The list of contour face indexes 120 | dynamic_lmk_b_coords: The list of contour barycentric weights 121 | neck_kin_chain: The tree to consider for the relative rotation 122 | dtype: Data type 123 | return: 124 | The contour face indexes and the corresponding barycentric weights 125 | """ 126 | 127 | batch_size = pose.shape[0] 128 | 129 | aa_pose = torch.index_select(pose.view(batch_size, -1, 3), 1, 130 | neck_kin_chain) 131 | rot_mats = batch_rodrigues( 132 | aa_pose.view(-1, 3), dtype=dtype).view(batch_size, -1, 3, 3) 133 | 134 | rel_rot_mat = torch.eye(3, device=pose.device, 135 | dtype=dtype).unsqueeze_(dim=0).expand(batch_size, -1, -1) 136 | for idx in range(len(neck_kin_chain)): 137 | rel_rot_mat = torch.bmm(rot_mats[:, idx], rel_rot_mat) 138 | 139 | y_rot_angle = torch.round( 140 | torch.clamp(rot_mat_to_euler(rel_rot_mat) * 180.0 / np.pi, 141 | max=39)).to(dtype=torch.long) 142 | 143 | neg_mask = y_rot_angle.lt(0).to(dtype=torch.long) 144 | mask = y_rot_angle.lt(-39).to(dtype=torch.long) 145 | neg_vals = mask * 78 + (1 - mask) * (39 - y_rot_angle) 146 | y_rot_angle = (neg_mask * neg_vals + 147 | (1 - neg_mask) * y_rot_angle) 148 | 149 | dyn_lmk_faces_idx = torch.index_select(dynamic_lmk_faces_idx, 150 | 0, y_rot_angle) 151 | dyn_lmk_b_coords = torch.index_select(dynamic_lmk_b_coords, 152 | 0, y_rot_angle) 153 | return dyn_lmk_faces_idx, dyn_lmk_b_coords 154 | 155 | def _vertices2landmarks(self, vertices, faces, lmk_faces_idx, lmk_bary_coords): 156 | """ 157 | Calculates landmarks by barycentric interpolation 158 | Input: 159 | vertices: torch.tensor NxVx3, dtype = torch.float32 160 | The tensor of input vertices 161 | faces: torch.tensor (N*F)x3, dtype = torch.long 162 | The faces of the mesh 163 | lmk_faces_idx: torch.tensor N X L, dtype = torch.long 164 | The tensor with the indices of the faces used to calculate the 165 | landmarks. 166 | lmk_bary_coords: torch.tensor N X L X 3, dtype = torch.float32 167 | The tensor of barycentric coordinates that are used to interpolate 168 | the landmarks 169 | 170 | Returns: 171 | landmarks: torch.tensor NxLx3, dtype = torch.float32 172 | The coordinates of the landmarks for each mesh in the batch 173 | """ 174 | # Extract the indices of the vertices for each face 175 | # NxLx3 176 | batch_size, num_verts = vertices.shape[:dd2] 177 | lmk_faces = torch.index_select(faces, 0, lmk_faces_idx.view(-1)).view( 178 | 1, -1, 3).view(batch_size, lmk_faces_idx.shape[1], -1) 179 | 180 | lmk_faces += torch.arange(batch_size, dtype=torch.long).view(-1, 1, 1).to( 181 | device=vertices.device) * num_verts 182 | 183 | lmk_vertices = vertices.view(-1, 3)[lmk_faces] 184 | landmarks = torch.einsum('blfi,blf->bli', [lmk_vertices, lmk_bary_coords]) 185 | return landmarks 186 | 187 | # def seletec_3d68(self, vertices): 188 | def compute_landmarks(self, vertices): 189 | landmarks3d = vertices2landmarks(vertices, self.faces_tensor, 190 | self.full_lmk_faces_idx.repeat(vertices.shape[0], 1), 191 | self.full_lmk_bary_coords.repeat(vertices.shape[0], 1, 1)) 192 | return landmarks3d 193 | 194 | def seletec_3d68(self, vertices): 195 | landmarks3d = vertices2landmarks(vertices, self.faces_tensor, 196 | self.full_lmk_faces_idx.repeat(vertices.shape[0], 1), 197 | self.full_lmk_bary_coords.repeat(vertices.shape[0], 1, 1)) 198 | return landmarks3d 199 | 200 | def project_to_shape_basis(self, shape_vector, shape_as_offset=False): 201 | batch_size = shape_vector.shape[0] 202 | n_vertices = self.v_template.shape[0] 203 | n_eigenvectors = self.n_shape 204 | # shape_params = basis dot (shape_vector - average) # uses properties of the PCA 205 | if shape_as_offset: 206 | diff = shape_vector 207 | else: 208 | diff = shape_vector - self.v_template 209 | return torch.matmul(diff.reshape(batch_size, -1), self.shapedirs[:, :, :n_eigenvectors].reshape(3 * n_vertices, n_eigenvectors)) 210 | 211 | def compute_distance_to_basis(self, shape_vector, shape_as_offset=False): 212 | batch_size = shape_vector.shape[0] 213 | n_vertices = self.v_template.shape[0] 214 | n_eigenvectors = self.n_shape 215 | 216 | # shape_vector torch.Size([3, 5023, 3]) 217 | # self.v_template torch.Size([5023, 3]) 218 | # self.shapedirs torch.Size([5023, 3, 150]) 219 | # diff torch.Size([3, 5023, 3]) 220 | # shape_params torch.Size([5023, 15069]) 221 | 222 | # shape_params = basis dot (shape_vector - average) # uses properties of the PCA 223 | if shape_as_offset: 224 | diff = shape_vector 225 | else: 226 | diff = shape_vector - self.v_template 227 | shape_params = torch.matmul(diff.reshape(batch_size, -1), self.shapedirs[:, :, :n_eigenvectors].reshape(3 * n_vertices, n_eigenvectors)) 228 | distance = diff - torch.matmul(shape_params, self.shapedirs[:, :, :n_eigenvectors].reshape(n_vertices * 3, n_eigenvectors).t()).reshape(batch_size, n_vertices, 3) 229 | return distance 230 | 231 | def get_std(self): 232 | n_eigenvectors = self.cfg.n_shape 233 | basis = self.shapedirs[:, :, :n_eigenvectors] 234 | std = torch.norm(basis.reshape(-1, n_eigenvectors), dim=0) 235 | 236 | return std 237 | 238 | def compute_closest_shape(self, shape_vector): 239 | B = shape_vector.shape[0] 240 | N = self.v_template.shape[0] 241 | n_eigenvectors = self.cfg.n_shape 242 | 243 | basis = self.shapedirs[:, :, :n_eigenvectors] 244 | diff = (shape_vector - self.v_template).reshape(B, -1) 245 | std = torch.norm(basis.reshape(-1, n_eigenvectors), dim=0) 246 | inv = 1.0 / std.square() 247 | params = inv * torch.matmul(diff, basis.reshape(3 * N, n_eigenvectors)) 248 | # params = torch.max(torch.min(params, std*-3.0), std*3.0) 249 | 250 | return self.v_template + torch.matmul(params, basis.reshape(N * 3, n_eigenvectors).T).reshape(B, N, 3), params 251 | 252 | def forward(self, shape_params=None, expression_params=None, pose_params=None, eye_pose_params=None, neck_pose_params=None, shape_basis_delta=None): 253 | """ 254 | Input: 255 | shape_params: N X number of shape parameters 256 | expression_params: N X number of expression parameters 257 | pose_params: N X number of pose parameters (6) 258 | return:d 259 | vertices: N X V X 3 260 | landmarks: N X number of landmarks X 3 261 | """ 262 | batch_size = shape_params.shape[0] 263 | if pose_params is None: 264 | pose_params = self.eye_pose.expand(batch_size, -1) 265 | if eye_pose_params is None: 266 | eye_pose_params = self.eye_pose.expand(batch_size, -1) 267 | if neck_pose_params is None: 268 | neck_pose_params = self.neck_pose.expand(batch_size, -1) 269 | if expression_params is None: 270 | expression_params = torch.zeros([1, 100], dtype=self.dtype, requires_grad=False, device=self.neck_pose.device).expand(batch_size, -1) 271 | 272 | betas = torch.cat([shape_params, expression_params], dim=1) 273 | full_pose = torch.cat([pose_params[:, :3], neck_pose_params, pose_params[:, 3:], eye_pose_params], dim=1) 274 | template_vertices = self.v_template.unsqueeze(0).expand(batch_size, -1, -1) 275 | 276 | vertices, _ = lbs(betas, full_pose, template_vertices, 277 | self.shapedirs, self.posedirs, 278 | self.J_regressor, self.parents, 279 | self.lbs_weights, dtype=self.dtype) 280 | 281 | lmk_faces_idx = self.lmk_faces_idx.unsqueeze(dim=0).expand(batch_size, -1) 282 | lmk_bary_coords = self.lmk_bary_coords.unsqueeze(dim=0).expand(batch_size, -1, -1) 283 | 284 | dyn_lmk_faces_idx, dyn_lmk_bary_coords = self._find_dynamic_lmk_idx_and_bcoords( 285 | full_pose, self.dynamic_lmk_faces_idx, 286 | self.dynamic_lmk_bary_coords, 287 | self.neck_kin_chain, dtype=self.dtype) 288 | lmk_faces_idx = torch.cat([dyn_lmk_faces_idx, lmk_faces_idx], 1) 289 | lmk_bary_coords = torch.cat([dyn_lmk_bary_coords, lmk_bary_coords], 1) 290 | 291 | landmarks2d = vertices2landmarks(vertices, self.faces_tensor, 292 | lmk_faces_idx, 293 | lmk_bary_coords) 294 | bz = vertices.shape[0] 295 | landmarks3d = vertices2landmarks(vertices, self.faces_tensor, 296 | self.full_lmk_faces_idx.repeat(bz, 1), 297 | self.full_lmk_bary_coords.repeat(bz, 1, 1)) 298 | return vertices, landmarks2d, landmarks3d 299 | -------------------------------------------------------------------------------- /models/generator.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2023 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: mica@tue.mpg.de 16 | 17 | 18 | import torch 19 | import torch.nn as nn 20 | import torch.nn.functional as Functional 21 | 22 | from models.flame import FLAME 23 | 24 | 25 | def kaiming_leaky_init(m): 26 | classname = m.__class__.__name__ 27 | if classname.find('Linear') != -1: 28 | torch.nn.init.kaiming_normal_(m.weight, a=0.2, mode='fan_in', nonlinearity='leaky_relu') 29 | 30 | 31 | class MappingNetwork(nn.Module): 32 | def __init__(self, z_dim, map_hidden_dim, map_output_dim, hidden=2): 33 | super().__init__() 34 | 35 | if hidden > 5: 36 | self.skips = [int(hidden / 2)] 37 | else: 38 | self.skips = [] 39 | 40 | self.network = nn.ModuleList( 41 | [nn.Linear(z_dim, map_hidden_dim)] + 42 | [nn.Linear(map_hidden_dim, map_hidden_dim) if i not in self.skips else 43 | nn.Linear(map_hidden_dim + z_dim, map_hidden_dim) for i in range(hidden)] 44 | ) 45 | 46 | self.output = nn.Linear(map_hidden_dim, map_output_dim) 47 | self.network.apply(kaiming_leaky_init) 48 | with torch.no_grad(): 49 | self.output.weight *= 0.25 50 | 51 | def forward(self, z): 52 | h = z 53 | for i, l in enumerate(self.network): 54 | h = self.network[i](h) 55 | h = Functional.leaky_relu(h, negative_slope=0.2) 56 | if i in self.skips: 57 | h = torch.cat([z, h], 1) 58 | 59 | output = self.output(h) 60 | return output 61 | 62 | 63 | class Generator(nn.Module): 64 | def __init__(self, z_dim, map_hidden_dim, map_output_dim, hidden, model_cfg, device, regress=True): 65 | super().__init__() 66 | self.device = device 67 | self.cfg = model_cfg 68 | self.regress = regress 69 | 70 | if self.regress: 71 | self.regressor = MappingNetwork(z_dim, map_hidden_dim, map_output_dim, hidden).to(self.device) 72 | self.generator = FLAME(model_cfg).to(self.device) 73 | 74 | def forward(self, arcface): 75 | if self.regress: 76 | shape = self.regressor(arcface) 77 | else: 78 | shape = arcface 79 | 80 | prediction, _, _ = self.generator(shape_params=shape) 81 | 82 | return prediction, shape 83 | -------------------------------------------------------------------------------- /models/lbs.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2023 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # For commercial licensing contact, please contact ps-license@tuebingen.mpg.de 16 | 17 | 18 | from __future__ import absolute_import 19 | from __future__ import division 20 | from __future__ import print_function 21 | 22 | import numpy as np 23 | import torch 24 | import torch.nn.functional as F 25 | 26 | 27 | def rot_mat_to_euler(rot_mats): 28 | # Calculates rotation matrix to euler angles 29 | # Careful for extreme cases of eular angles like [0.0, pi, 0.0] 30 | 31 | sy = torch.sqrt(rot_mats[:, 0, 0] * rot_mats[:, 0, 0] + 32 | rot_mats[:, 1, 0] * rot_mats[:, 1, 0]) 33 | return torch.atan2(-rot_mats[:, 2, 0], sy) 34 | 35 | 36 | def find_dynamic_lmk_idx_and_bcoords(vertices, pose, dynamic_lmk_faces_idx, 37 | dynamic_lmk_b_coords, 38 | neck_kin_chain, dtype=torch.float32): 39 | ''' Compute the faces, barycentric coordinates for the dynamic landmarks 40 | 41 | 42 | To do so, we first compute the rotation of the neck around the y-axis 43 | and then use a pre-computed look-up table to find the faces and the 44 | barycentric coordinates that will be used. 45 | 46 | Special thanks to Soubhik Sanyal (soubhik.sanyal@tuebingen.mpg.de) 47 | for providing the original TensorFlow implementation and for the LUT. 48 | 49 | Parameters 50 | ---------- 51 | vertices: torch.tensor BxVx3, dtype = torch.float32 52 | The tensor of input vertices 53 | pose: torch.tensor Bx(Jx3), dtype = torch.float32 54 | The current pose of the body model 55 | dynamic_lmk_faces_idx: torch.tensor L, dtype = torch.long 56 | The look-up table from neck rotation to faces 57 | dynamic_lmk_b_coords: torch.tensor Lx3, dtype = torch.float32 58 | The look-up table from neck rotation to barycentric coordinates 59 | neck_kin_chain: list 60 | A python list that contains the indices of the joints that form the 61 | kinematic chain of the neck. 62 | dtype: torch.dtype, optional 63 | 64 | Returns 65 | ------- 66 | dyn_lmk_faces_idx: torch.tensor, dtype = torch.long 67 | A tensor of size BxL that contains the indices of the faces that 68 | will be used to compute the current dynamic landmarks. 69 | dyn_lmk_b_coords: torch.tensor, dtype = torch.float32 70 | A tensor of size BxL that contains the indices of the faces that 71 | will be used to compute the current dynamic landmarks. 72 | ''' 73 | 74 | batch_size = vertices.shape[0] 75 | 76 | aa_pose = torch.index_select(pose.view(batch_size, -1, 3), 1, 77 | neck_kin_chain) 78 | rot_mats = batch_rodrigues( 79 | aa_pose.view(-1, 3), dtype=dtype).view(batch_size, -1, 3, 3) 80 | 81 | rel_rot_mat = torch.eye(3, device=vertices.device, 82 | dtype=dtype).unsqueeze_(dim=0) 83 | for idx in range(len(neck_kin_chain)): 84 | rel_rot_mat = torch.bmm(rot_mats[:, idx], rel_rot_mat) 85 | 86 | y_rot_angle = torch.round( 87 | torch.clamp(-rot_mat_to_euler(rel_rot_mat) * 180.0 / np.pi, 88 | max=39)).to(dtype=torch.long) 89 | neg_mask = y_rot_angle.lt(0).to(dtype=torch.long) 90 | mask = y_rot_angle.lt(-39).to(dtype=torch.long) 91 | neg_vals = mask * 78 + (1 - mask) * (39 - y_rot_angle) 92 | y_rot_angle = (neg_mask * neg_vals + 93 | (1 - neg_mask) * y_rot_angle) 94 | 95 | dyn_lmk_faces_idx = torch.index_select(dynamic_lmk_faces_idx, 96 | 0, y_rot_angle) 97 | dyn_lmk_b_coords = torch.index_select(dynamic_lmk_b_coords, 98 | 0, y_rot_angle) 99 | 100 | return dyn_lmk_faces_idx, dyn_lmk_b_coords 101 | 102 | 103 | def vertices2landmarks(vertices, faces, lmk_faces_idx, lmk_bary_coords): 104 | ''' Calculates landmarks by barycentric interpolation 105 | 106 | Parameters 107 | ---------- 108 | vertices: torch.tensor BxVx3, dtype = torch.float32 109 | The tensor of input vertices 110 | faces: torch.tensor Fx3, dtype = torch.long 111 | The faces of the mesh 112 | lmk_faces_idx: torch.tensor L, dtype = torch.long 113 | The tensor with the indices of the faces used to calculate the 114 | landmarks. 115 | lmk_bary_coords: torch.tensor Lx3, dtype = torch.float32 116 | The tensor of barycentric coordinates that are used to interpolate 117 | the landmarks 118 | 119 | Returns 120 | ------- 121 | landmarks: torch.tensor BxLx3, dtype = torch.float32 122 | The coordinates of the landmarks for each mesh in the batch 123 | ''' 124 | # Extract the indices of the vertices for each face 125 | # BxLx3 126 | batch_size, num_verts = vertices.shape[:2] 127 | device = vertices.device 128 | 129 | lmk_faces = torch.index_select(faces, 0, lmk_faces_idx.view(-1)).view( 130 | batch_size, -1, 3) 131 | 132 | lmk_faces += torch.arange( 133 | batch_size, dtype=torch.long, device=device).view(-1, 1, 1) * num_verts 134 | 135 | lmk_vertices = vertices.view(-1, 3)[lmk_faces].view( 136 | batch_size, -1, 3, 3) 137 | 138 | landmarks = torch.einsum('blfi,blf->bli', [lmk_vertices, lmk_bary_coords]) 139 | return landmarks 140 | 141 | 142 | def lbs(betas, pose, v_template, shapedirs, posedirs, J_regressor, parents, 143 | lbs_weights, pose2rot=True, dtype=torch.float32): 144 | ''' Performs Linear Blend Skinning with the given shape and pose parameters 145 | 146 | Parameters 147 | ---------- 148 | betas : torch.tensor BxNB 149 | The tensor of shape parameters 150 | pose : torch.tensor Bx(J + 1) * 3 151 | The pose parameters in axis-angle format 152 | v_template torch.tensor BxVx3 153 | The template mesh that will be deformed 154 | shapedirs : torch.tensor 1xNB 155 | The tensor of PCA shape displacements 156 | posedirs : torch.tensor Px(V * 3) 157 | The pose PCA coefficients 158 | J_regressor : torch.tensor JxV 159 | The regressor array that is used to calculate the joints from 160 | the position of the vertices 161 | parents: torch.tensor J 162 | The array that describes the kinematic tree for the model 163 | lbs_weights: torch.tensor N x V x (J + 1) 164 | The linear blend skinning weights that represent how much the 165 | rotation matrix of each part affects each vertex 166 | pose2rot: bool, optional 167 | Flag on whether to convert the input pose tensor to rotation 168 | matrices. The default value is True. If False, then the pose tensor 169 | should already contain rotation matrices and have a size of 170 | Bx(J + 1)x9 171 | dtype: torch.dtype, optional 172 | 173 | Returns 174 | ------- 175 | verts: torch.tensor BxVx3 176 | The vertices of the mesh after applying the shape and pose 177 | displacements. 178 | joints: torch.tensor BxJx3 179 | The joints of the model 180 | ''' 181 | 182 | batch_size = max(betas.shape[0], pose.shape[0]) 183 | device = betas.device 184 | 185 | # Add shape contribution 186 | v_shaped = v_template + blend_shapes(betas, shapedirs) 187 | 188 | # Get the joints 189 | # NxJx3 array 190 | J = vertices2joints(J_regressor, v_shaped) 191 | 192 | # 3. Add pose blend shapes 193 | # N x J x 3 x 3 194 | ident = torch.eye(3, dtype=dtype, device=device) 195 | if pose2rot: 196 | rot_mats = batch_rodrigues( 197 | pose.view(-1, 3), dtype=dtype).view([batch_size, -1, 3, 3]) 198 | 199 | pose_feature = (rot_mats[:, 1:, :, :] - ident).view([batch_size, -1]) 200 | # (N x P) x (P, V * 3) -> N x V x 3 201 | pose_offsets = torch.matmul(pose_feature, posedirs) \ 202 | .view(batch_size, -1, 3) 203 | else: 204 | pose_feature = pose[:, 1:].view(batch_size, -1, 3, 3) - ident 205 | rot_mats = pose.view(batch_size, -1, 3, 3) 206 | 207 | pose_offsets = torch.matmul(pose_feature.view(batch_size, -1), 208 | posedirs).view(batch_size, -1, 3) 209 | 210 | v_posed = pose_offsets + v_shaped 211 | # 4. Get the global joint location 212 | J_transformed, A = batch_rigid_transform(rot_mats, J, parents, dtype=dtype) 213 | 214 | # 5. Do skinning: 215 | # W is N x V x (J + 1) 216 | W = lbs_weights.unsqueeze(dim=0).expand([batch_size, -1, -1]) 217 | # (N x V x (J + 1)) x (N x (J + 1) x 16) 218 | num_joints = J_regressor.shape[0] 219 | T = torch.matmul(W, A.view(batch_size, num_joints, 16)) \ 220 | .view(batch_size, -1, 4, 4) 221 | 222 | homogen_coord = torch.ones([batch_size, v_posed.shape[1], 1], 223 | dtype=dtype, device=device) 224 | v_posed_homo = torch.cat([v_posed, homogen_coord], dim=2) 225 | v_homo = torch.matmul(T, torch.unsqueeze(v_posed_homo, dim=-1)) 226 | 227 | verts = v_homo[:, :, :3, 0] 228 | 229 | return verts, J_transformed 230 | 231 | 232 | def vertices2joints(J_regressor, vertices): 233 | ''' Calculates the 3D joint locations from the vertices 234 | 235 | Parameters 236 | ---------- 237 | J_regressor : torch.tensor JxV 238 | The regressor array that is used to calculate the joints from the 239 | position of the vertices 240 | vertices : torch.tensor BxVx3 241 | The tensor of mesh vertices 242 | 243 | Returns 244 | ------- 245 | torch.tensor BxJx3 246 | The location of the joints 247 | ''' 248 | 249 | return torch.einsum('bik,ji->bjk', [vertices, J_regressor]) 250 | 251 | 252 | def blend_shapes(betas, shape_disps): 253 | ''' Calculates the per vertex displacement due to the blend shapes 254 | 255 | 256 | Parameters 257 | ---------- 258 | betas : torch.tensor Bx(num_betas) 259 | Blend shape coefficients 260 | shape_disps: torch.tensor Vx3x(num_betas) 261 | Blend shapes 262 | 263 | Returns 264 | ------- 265 | torch.tensor BxVx3 266 | The per-vertex displacement due to shape deformation 267 | ''' 268 | 269 | # Displacement[b, m, k] = sum_{l} betas[b, l] * shape_disps[m, k, l] 270 | # i.e. Multiply each shape displacement by its corresponding beta and 271 | # then sum them. 272 | blend_shape = torch.einsum('bl,mkl->bmk', [betas, shape_disps]) 273 | return blend_shape 274 | 275 | 276 | def batch_rodrigues(rot_vecs, epsilon=1e-8, dtype=torch.float32): 277 | ''' Calculates the rotation matrices for a batch of rotation vectors 278 | Parameters 279 | ---------- 280 | rot_vecs: torch.tensor Nx3 281 | array of N axis-angle vectors 282 | Returns 283 | ------- 284 | R: torch.tensor Nx3x3 285 | The rotation matrices for the given axis-angle parameters 286 | ''' 287 | 288 | batch_size = rot_vecs.shape[0] 289 | device = rot_vecs.device 290 | 291 | angle = torch.norm(rot_vecs + 1e-8, dim=1, keepdim=True) 292 | rot_dir = rot_vecs / angle 293 | 294 | cos = torch.unsqueeze(torch.cos(angle), dim=1) 295 | sin = torch.unsqueeze(torch.sin(angle), dim=1) 296 | 297 | # Bx1 arrays 298 | rx, ry, rz = torch.split(rot_dir, 1, dim=1) 299 | K = torch.zeros((batch_size, 3, 3), dtype=dtype, device=device) 300 | 301 | zeros = torch.zeros((batch_size, 1), dtype=dtype, device=device) 302 | K = torch.cat([zeros, -rz, ry, rz, zeros, -rx, -ry, rx, zeros], dim=1) \ 303 | .view((batch_size, 3, 3)) 304 | 305 | ident = torch.eye(3, dtype=dtype, device=device).unsqueeze(dim=0) 306 | rot_mat = ident + sin * K + (1 - cos) * torch.bmm(K, K) 307 | return rot_mat 308 | 309 | 310 | def transform_mat(R, t): 311 | ''' Creates a batch of transformation matrices 312 | Args: 313 | - R: Bx3x3 array of a batch of rotation matrices 314 | - t: Bx3x1 array of a batch of translation vectors 315 | Returns: 316 | - T: Bx4x4 Transformation matrix 317 | ''' 318 | # No padding left or right, only add an extra row 319 | return torch.cat([F.pad(R, [0, 0, 0, 1]), 320 | F.pad(t, [0, 0, 0, 1], value=1)], dim=2) 321 | 322 | 323 | def batch_rigid_transform(rot_mats, joints, parents, dtype=torch.float32): 324 | """ 325 | Applies a batch of rigid transformations to the joints 326 | 327 | Parameters 328 | ---------- 329 | rot_mats : torch.tensor BxNx3x3 330 | Tensor of rotation matrices 331 | joints : torch.tensor BxNx3 332 | Locations of joints 333 | parents : torch.tensor BxN 334 | The kinematic tree of each object 335 | dtype : torch.dtype, optional: 336 | The data type of the created tensors, the default is torch.float32 337 | 338 | Returns 339 | ------- 340 | posed_joints : torch.tensor BxNx3 341 | The locations of the joints after applying the pose rotations 342 | rel_transforms : torch.tensor BxNx4x4 343 | The relative (with respect to the root joint) rigid transformations 344 | for all the joints 345 | """ 346 | 347 | joints = torch.unsqueeze(joints, dim=-1) 348 | 349 | rel_joints = joints.clone() 350 | rel_joints[:, 1:] -= joints[:, parents[1:]] 351 | 352 | # transforms_mat = transform_mat( 353 | # rot_mats.view(-1, 3, 3), 354 | # rel_joints.view(-1, 3, 1)).view(-1, joints.shape[1], 4, 4) 355 | transforms_mat = transform_mat( 356 | rot_mats.view(-1, 3, 3), 357 | rel_joints.reshape(-1, 3, 1)).reshape(-1, joints.shape[1], 4, 4) 358 | 359 | transform_chain = [transforms_mat[:, 0]] 360 | for i in range(1, parents.shape[0]): 361 | # Subtract the joint location at the rest pose 362 | # No need for rotation, since it's identity when at rest 363 | curr_res = torch.matmul(transform_chain[parents[i]], transforms_mat[:, i]) 364 | transform_chain.append(curr_res) 365 | 366 | transforms = torch.stack(transform_chain, dim=1) 367 | 368 | # The last column of the transformations contains the posed joints 369 | posed_joints = transforms[:, :, :3, 3] 370 | 371 | # The last column of the transformations contains the posed joints 372 | posed_joints = transforms[:, :, :3, 3] 373 | 374 | joints_homogen = F.pad(joints, [0, 0, 0, 1]) 375 | 376 | rel_transforms = transforms - F.pad( 377 | torch.matmul(transforms, joints_homogen), [3, 0, 0, 0, 0, 0, 0, 0]) 378 | 379 | return posed_joints, rel_transforms 380 | -------------------------------------------------------------------------------- /render_dataset.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2023 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: mica@tue.mpg.de 16 | import os 17 | import random 18 | from glob import glob 19 | from pathlib import Path 20 | 21 | import cv2 22 | import numpy as np 23 | import torch 24 | import trimesh 25 | from tqdm import tqdm 26 | 27 | from configs.config import get_cfg_defaults 28 | from micalib.renderer import MeshShapeRenderer 29 | from models.flame import FLAME 30 | 31 | np.random.seed(125) 32 | random.seed(125) 33 | 34 | 35 | def main(): 36 | cfg = get_cfg_defaults() 37 | render = MeshShapeRenderer(obj_filename=cfg.model.topology_path) 38 | flame = FLAME(cfg.model).to('cuda:0') 39 | datasets = sorted(glob('/home/wzielonka/datasets/MICA/*')) 40 | for dataset in tqdm(datasets): 41 | meshes = sorted(glob(f'{dataset}/FLAME_parameters/*/*.npz')) 42 | sample_list = np.array(np.random.choice(range(len(meshes)), size=30 * 5)) 43 | dst = Path('./output', Path(dataset).name) 44 | dst.mkdir(parents=True, exist_ok=True) 45 | j = 0 46 | k = 0 47 | images = np.zeros((512, 512 * 5, 3)) 48 | for i in sample_list: 49 | params = np.load(meshes[i], allow_pickle=True) 50 | betas = torch.tensor(params['betas']).float().cuda() 51 | shape_params = betas[:300][None] 52 | v = flame(shape_params=shape_params)[0] 53 | rendering = render.render_mesh(v) 54 | image = (rendering[0].cpu().numpy().transpose(1, 2, 0).copy() * 255)[:, :, [2, 1, 0]] 55 | image = np.minimum(np.maximum(image, 0), 255).astype(np.uint8) 56 | images[0:512, 512 * j:512 * (j + 1), :] = image 57 | j += 1 58 | 59 | if j % 5 == 0 and j > 0: 60 | dst.mkdir(parents=True, exist_ok=True) 61 | cv2.imwrite(f'{dst}/{str(k).zfill(4)}.png', images) 62 | images = np.zeros((512, 512 * 5, 3)) 63 | j = 0 64 | k += 1 65 | 66 | os.system(f'ffmpeg -y -framerate 1 -pattern_type glob -i \'{dst}/*.png\' -c:v libx264 -pix_fmt yuv420p {dst}/video.mp4') 67 | os.system(f'gifski -o ./output/{Path(dataset).name}.gif {dst}/*.png --quality 100 --fps 1') 68 | 69 | 70 | if __name__ == '__main__': 71 | main() 72 | -------------------------------------------------------------------------------- /test.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2023 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: mica@tue.mpg.de 16 | 17 | 18 | import os 19 | import sys 20 | 21 | import torch 22 | import torch.backends.cudnn as cudnn 23 | import torch.multiprocessing as mp 24 | 25 | from jobs import test 26 | 27 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '.'))) 28 | 29 | if __name__ == '__main__': 30 | from configs.config import parse_args 31 | 32 | cfg, args = parse_args() 33 | 34 | if cfg.cfg_file is not None: 35 | exp_name = cfg.cfg_file.split('/')[-1].split('.')[0] 36 | cfg.output_dir = os.path.join('./output', exp_name) 37 | 38 | cudnn.benchmark = False 39 | cudnn.deterministic = True 40 | torch.cuda.empty_cache() 41 | num_gpus = torch.cuda.device_count() 42 | 43 | mp.spawn(test, args=(num_gpus, cfg, args), nprocs=num_gpus, join=True) 44 | 45 | exit(0) 46 | -------------------------------------------------------------------------------- /testing/now/now.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2023 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: mica@tue.mpg.de 16 | 17 | 18 | import os 19 | import time 20 | from glob import glob 21 | from shutil import copyfile 22 | 23 | logs = '/home/wzielonka/projects/MICA/testing/now/logs/' 24 | jobs = '/home/wzielonka/projects/MICA/testing/now/jobs/' 25 | root = '/home/wzielonka/projects/MICA/output/' 26 | 27 | experiments = [] 28 | 29 | 30 | def test(): 31 | global experiments 32 | if len(experiments) == 0: 33 | experiments = list(filter(lambda f: 'condor' not in f, os.listdir('../../output/'))) 34 | 35 | os.system('rm -rf logs') 36 | os.system('rm -rf jobs') 37 | 38 | os.makedirs('logs', exist_ok=True) 39 | os.makedirs('jobs', exist_ok=True) 40 | 41 | for experiment in sorted(experiments): 42 | print(f'Testing {experiment}') 43 | copyfile(f'{root}{experiment}/model.tar', f'{root}{experiment}/best_models/best_model_last.tar') 44 | for idx, checkpoint in enumerate(glob(root + experiment + f'/best_models/*.tar')): 45 | model_name = checkpoint.split('/')[-1].split('.')[0] 46 | model_name = model_name.replace('best_model_', 'now_test_') 47 | predicted_meshes = f'{root}{experiment}/{model_name}/predicted_meshes/' 48 | run = f'{experiment}_{str(idx).zfill(5)}' 49 | with open(f'{jobs}/{run}.sub', 'w') as fid: 50 | fid.write('executable = /bin/bash\n') 51 | arguments = f'/home/wzielonka/projects/MICA/testing/now/template.sh {experiment} {checkpoint} now {predicted_meshes}' 52 | fid.write(f'arguments = {arguments}\n') 53 | fid.write(f'error = {logs}{run}.err\n') 54 | fid.write(f'output = {logs}{run}.out\n') 55 | fid.write(f'log = {logs}{run}.log\n') 56 | fid.write(f'request_cpus = 4\n') 57 | fid.write(f'request_gpus = 1\n') 58 | fid.write(f'requirements = (TARGET.CUDAGlobalMemoryMb > 5000) && (TARGET.CUDAGlobalMemoryMb < 33000)\n') 59 | fid.write(f'request_memory = 8192\n') 60 | fid.write(f'queue\n') 61 | 62 | os.system(f'condor_submit_bid 512 {jobs}/{run}.sub') 63 | 64 | time.sleep(2) 65 | 66 | 67 | if __name__ == '__main__': 68 | test() 69 | -------------------------------------------------------------------------------- /testing/now/template.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | PYTHON_ENV=/home/wzielonka/miniconda3/etc/profile.d/conda.sh 4 | SCRIPT_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )" 5 | export PATH=/usr/local/bin:/usr/bin:/bin:/usr/sbin:$PATH 6 | export LD_LIBRARY_PATH=/is/software/nvidia/nccl-2.4.8-cuda10.1/lib/ 7 | 8 | source ${PYTHON_ENV} 9 | module load cuda/10.1 10 | module load gcc/4.9 11 | 12 | EXPERIMENT='' 13 | CHECKPOINT='' 14 | BENCHMARK='' 15 | PREDICTED='' 16 | 17 | echo 'Testing has started...' 18 | 19 | if [ -n "$1" ]; then EXPERIMENT=${1}; fi 20 | if [ -n "$2" ]; then CHECKPOINT=${2}; fi 21 | if [ -n "$3" ]; then BENCHMARK=${3}; fi 22 | if [ -n "$4" ]; then PREDICTED=${4}; fi 23 | 24 | ROOT=/home/wzielonka/projects/MICA/output/ 25 | NOW=/home/wzielonka/datasets/NoWDataset/final_release_version/ 26 | 27 | conda activate NFC 28 | 29 | cd /home/wzielonka/projects/MICA 30 | python test.py --cfg /home/wzielonka/projects/MICA/configs/${EXPERIMENT}.yml --test_dataset ${BENCHMARK} --checkpoint ${CHECKPOINT} 31 | 32 | source /home/wzielonka/.virtualenvs/NoW/bin/activate 33 | cd /home/wzielonka/projects/NoW 34 | python compute_error.py ${NOW} ${PREDICTED} true 35 | 36 | # Plot diagram 37 | # source /home/wzielonka/.virtualenvs/NoW/bin/activate 38 | # python cumulative_errors.py -------------------------------------------------------------------------------- /testing/stirling/stirling.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2023 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: mica@tue.mpg.de 16 | 17 | 18 | import os 19 | import time 20 | from glob import glob 21 | from shutil import copyfile 22 | 23 | logs = '/home/wzielonka/projects/MICA/testing/stirling/logs/' 24 | jobs = '/home/wzielonka/projects/MICA/testing/stirling/jobs/' 25 | root = '/home/wzielonka/projects/MICA/output/' 26 | 27 | experiments = [] 28 | 29 | 30 | def test(): 31 | global experiments 32 | if len(experiments) == 0: 33 | experiments = list(filter(lambda f: 'condor' not in f and 'resnet' in f, os.listdir('../../output/'))) 34 | # experiments = list(filter(lambda f: 'experiment_' in f, os.listdir('../../output/'))) 35 | 36 | os.system('rm -rf logs') 37 | os.system('rm -rf jobs') 38 | 39 | os.makedirs('logs', exist_ok=True) 40 | os.makedirs('jobs', exist_ok=True) 41 | 42 | for experiment in sorted(experiments): 43 | print(f'Testing {experiment}') 44 | copyfile(f'{root}{experiment}/model.tar', f'{root}{experiment}/best_models/best_model_last.tar') 45 | for idx, checkpoint in enumerate(glob(root + experiment + f'/best_models/*.tar')): 46 | model_name = checkpoint.split('/')[-1].split('.')[0] 47 | model_name = model_name.replace('best_model_', 'stirling_test_') 48 | predicted_meshes = f'{root}{experiment}/{model_name}/predicted_meshes/' 49 | run = f'{experiment}_{str(idx).zfill(5)}' 50 | with open(f'{jobs}/{run}.sub', 'w') as fid: 51 | fid.write('executable = /bin/bash\n') 52 | arguments = f'/home/wzielonka/projects/MICA/testing/stirling/template.sh {experiment} {checkpoint} stirling {predicted_meshes}' 53 | fid.write(f'arguments = {arguments}\n') 54 | fid.write(f'error = {logs}{run}.err\n') 55 | fid.write(f'output = {logs}{run}.out\n') 56 | fid.write(f'log = {logs}{run}.log\n') 57 | fid.write(f'request_cpus = 8\n') 58 | fid.write(f'request_gpus = 1\n') 59 | fid.write(f'requirements = (TARGET.CUDAGlobalMemoryMb > 5000) && (TARGET.CUDAGlobalMemoryMb < 12000)\n') 60 | fid.write(f'request_memory = 8192\n') 61 | fid.write(f'queue\n') 62 | 63 | os.system(f'condor_submit_bid 512 {jobs}/{run}.sub') 64 | 65 | time.sleep(2) 66 | 67 | 68 | if __name__ == '__main__': 69 | test() 70 | -------------------------------------------------------------------------------- /testing/stirling/template.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | PYTHON_ENV=/home/wzielonka/miniconda3/etc/profile.d/conda.sh 4 | SCRIPT_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )" 5 | export PATH=/usr/local/bin:/usr/bin:/bin:/usr/sbin:$PATH 6 | export LD_LIBRARY_PATH=/is/software/nvidia/nccl-2.4.8-cuda10.1/lib/ 7 | 8 | source ${PYTHON_ENV} 9 | module load cuda/10.1 10 | module load gcc/4.9 11 | 12 | EXPERIMENT='' 13 | CHECKPOINT='' 14 | BENCHMARK='' 15 | PREDICTED='' 16 | 17 | echo 'Testing has started...' 18 | 19 | if [ -n "$1" ]; then EXPERIMENT=${1}; fi 20 | if [ -n "$2" ]; then CHECKPOINT=${2}; fi 21 | if [ -n "$3" ]; then BENCHMARK=${3}; fi 22 | if [ -n "$4" ]; then PREDICTED=${4}; fi 23 | 24 | ROOT=/home/wzielonka/projects/MICA/output/ 25 | NOW=/home/wzielonka/datasets/NoWDataset/final_release_version/ 26 | 27 | conda activate NFC 28 | 29 | cd /home/wzielonka/projects/MICA 30 | python test.py --cfg /home/wzielonka/projects/MICA/configs/${EXPERIMENT}.yml --test_dataset ${BENCHMARK} --checkpoint ${CHECKPOINT} 31 | 32 | source /home/wzielonka/.virtualenvs/NoW/bin/activate 33 | cd /home/wzielonka/projects/NoW 34 | 35 | # Arguments for NoW 36 | # predicted_mesh_folder = sys.argv[1] 37 | # fixed = sys.argv[2] 38 | # now = sys.argv[3] 39 | 40 | # python feng_error.py ${PREDICTED}/HQ true true 41 | # python feng_error.py ${PREDICTED}/LQ true true 42 | 43 | # Plot diagram 44 | # source /home/wzielonka/.virtualenvs/NoW/bin/activate 45 | # python cumulative_errors.py --type hq 46 | -------------------------------------------------------------------------------- /train.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2023 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: mica@tue.mpg.de 16 | 17 | 18 | import os 19 | import sys 20 | 21 | import torch 22 | import torch.backends.cudnn as cudnn 23 | import torch.multiprocessing as mp 24 | 25 | from jobs import train 26 | 27 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '.'))) 28 | 29 | if __name__ == '__main__': 30 | from configs.config import parse_args 31 | 32 | cfg, args = parse_args() 33 | 34 | if cfg.cfg_file is not None: 35 | exp_name = cfg.cfg_file.split('/')[-1].split('.')[0] 36 | cfg.output_dir = os.path.join('./output', exp_name) 37 | 38 | cudnn.benchmark = False 39 | cudnn.deterministic = True 40 | torch.cuda.empty_cache() 41 | num_gpus = torch.cuda.device_count() 42 | 43 | mp.spawn(train, args=(num_gpus, cfg), nprocs=num_gpus, join=True) 44 | 45 | exit(0) 46 | -------------------------------------------------------------------------------- /utils/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Zielon/MICA/af22e7a5810d474bc28a1433db533723d6bd2b07/utils/__init__.py -------------------------------------------------------------------------------- /utils/best_model.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2023 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: mica@tue.mpg.de 16 | 17 | 18 | import os 19 | 20 | import numpy as np 21 | from loguru import logger 22 | 23 | 24 | class BestModel: 25 | def __init__(self, trainer): 26 | self.average = np.Inf 27 | self.weighted_average = np.Inf 28 | self.smoothed_average = np.Inf 29 | self.smoothed_weighted_average = np.Inf 30 | self.running_average = np.Inf 31 | self.running_weighted_average = np.Inf 32 | self.now_mean = None 33 | 34 | self.trainer = trainer 35 | self.counter = None 36 | 37 | self.N = trainer.cfg.running_average 38 | 39 | os.makedirs(os.path.join(self.trainer.cfg.output_dir, 'best_models'), exist_ok=True) 40 | 41 | def state_dict(self): 42 | return { 43 | 'average': self.average, 44 | 'smoothed_average': self.smoothed_average, 45 | 'running_average': self.running_average, 46 | 'now_mean': self.now_mean, 47 | 'counter': self.counter, 48 | } 49 | 50 | def load_state_dict(self, dict): 51 | self.average = dict['average'] 52 | self.smoothed_average = dict['smoothed_average'] 53 | self.running_average = dict['running_average'] 54 | self.now_mean = dict['now_mean'] 55 | self.counter = dict['counter'] 56 | 57 | logger.info(f'[BEST] Best score weighted average: ' 58 | f'NoW mean: {self.now_mean:.6f} | ' 59 | f'average: {self.average:.6f} | ' 60 | f'smoothed average: {self.running_average:.6f}') 61 | 62 | def __call__(self, weighted_average, average): 63 | if self.counter is None: 64 | self.counter = 1 65 | self.average = average 66 | self.weighted_average = weighted_average 67 | self.running_weighted_average = weighted_average 68 | self.running_average = average 69 | 70 | return weighted_average, average 71 | 72 | if weighted_average < self.weighted_average: 73 | delta = self.weighted_average - weighted_average 74 | self.weighted_average = weighted_average 75 | logger.info(f'[BEST] (Average weighted) {self.trainer.global_step} | {delta:.6f} improvement and value: {self.weighted_average:.6f}') 76 | self.trainer.save_checkpoint(os.path.join(self.trainer.cfg.output_dir, 'best_models', f'best_model_0.tar')) 77 | 78 | if average < self.average: 79 | delta = self.average - average 80 | self.average = average 81 | logger.info(f'[BEST] (Average) {self.trainer.global_step} | {delta:.6f} improvement and value: {self.average:.6f}') 82 | self.trainer.save_checkpoint(os.path.join(self.trainer.cfg.output_dir, 'best_models', f'best_model_1.tar')) 83 | 84 | n = self.N 85 | 86 | self.running_average = self.running_average * ((n - 1) / n) + (average / n) 87 | if self.running_average < self.smoothed_average: 88 | delta = self.smoothed_average - self.running_average 89 | self.smoothed_average = self.running_average 90 | logger.info(f'[BEST] (Average Smoothed) {self.trainer.global_step} | {delta:.6f} improvement and value: {self.smoothed_average:.6f} | counter: {self.counter} | window: {n}') 91 | self.trainer.save_checkpoint(os.path.join(self.trainer.cfg.output_dir, 'best_models', f'best_model_3.tar')) 92 | 93 | self.counter += 1 94 | 95 | return self.running_weighted_average, self.running_average 96 | 97 | def now(self, median, mean, std): 98 | if self.now_mean is None: 99 | self.now_mean = mean 100 | return 101 | 102 | if mean < self.now_mean: 103 | delta = self.now_mean - mean 104 | self.now_mean = mean 105 | logger.info(f'[BEST] (NoW) {self.trainer.global_step} | {delta:.6f} improvement and mean: {self.now_mean:.6f} std: {std} median: {median}') 106 | self.trainer.save_checkpoint(os.path.join(self.trainer.cfg.output_dir, 'best_models', f'best_model_now.tar')) 107 | -------------------------------------------------------------------------------- /utils/landmark_detector.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 3 | # holder of all proprietary rights on this computer program. 4 | # You can only use this computer program if you have closed 5 | # a license agreement with MPG or you get the right to use the computer 6 | # program from someone who is authorized to grant you that right. 7 | # Any use of the computer program without a valid license is prohibited and 8 | # liable to prosecution. 9 | # 10 | # Copyright©2023 Max-Planck-Gesellschaft zur Förderung 11 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 12 | # for Intelligent Systems. All rights reserved. 13 | # 14 | # Contact: mica@tue.mpg.de 15 | 16 | 17 | import face_alignment 18 | import numpy as np 19 | from insightface.app import FaceAnalysis 20 | from loguru import logger 21 | 22 | from datasets.creation.util import get_bbox 23 | 24 | 25 | class Detectors: 26 | def __init__(self): 27 | self.RETINAFACE = 'RETINAFACE' 28 | self.FAN = 'FAN' 29 | 30 | 31 | detectors = Detectors() 32 | 33 | 34 | class LandmarksDetector: 35 | def __init__(self, model='retinaface', device='cuda:0'): 36 | model = model.upper() 37 | self.predictor = model 38 | if model == detectors.RETINAFACE: 39 | self._face_detector = FaceAnalysis(name='antelopev2', providers=['CUDAExecutionProvider']) 40 | self._face_detector.prepare(ctx_id=0, det_size=(224, 224)) 41 | elif model == detectors.FAN: 42 | self._face_detector = face_alignment.FaceAlignment(face_alignment.LandmarksType._2D, device=device) 43 | else: 44 | logger.error(f'[ERROR] Landmark predictor not supported {model}') 45 | exit(-1) 46 | 47 | logger.info(f'[DETECTOR] Selected {model} as landmark detector.') 48 | 49 | def detect(self, img): 50 | if self.predictor == detectors.RETINAFACE: 51 | bboxes, kpss = self._face_detector.det_model.detect(img, max_num=0, metric='default') 52 | return bboxes, kpss 53 | 54 | if self.predictor == detectors.FAN: 55 | lmks, scores, detected_faces = self._face_detector.get_landmarks_from_image(img, return_landmark_score=True, return_bboxes=True) 56 | if detected_faces is None: 57 | return np.empty(0), np.empty(0) 58 | bboxes = np.stack(detected_faces) 59 | # bboxes = get_bbox(img, np.stack(lmks)) 60 | # bboxes[:, 4] = detected_faces[:, 4] 61 | # https://github.com/Rubikplayer/flame-fitting/blob/master/data/landmarks_51_annotated.png 62 | lmk51 = np.stack(lmks)[:, 17:, :] 63 | kpss = lmk51[:, [20, 27, 13, 43, 47], :] # left eye, right eye, nose, left mouth, right mouth 64 | kpss[:, 0, :] = lmk51[:, [21, 24], :].mean(1) # center of eye 65 | kpss[:, 1, :] = lmk51[:, [27, 29], :].mean(1) 66 | return bboxes, kpss 67 | 68 | return None, None 69 | -------------------------------------------------------------------------------- /utils/masking.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2023 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: mica@tue.mpg.de 16 | import os 17 | import pickle 18 | 19 | import numpy as np 20 | import torch 21 | import torch.nn as nn 22 | from trimesh import Trimesh 23 | 24 | 25 | def to_tensor(array, dtype=torch.float32): 26 | if 'torch.tensor' not in str(type(array)): 27 | return torch.tensor(array, dtype=dtype) 28 | 29 | 30 | def to_np(array, dtype=np.float32): 31 | if 'scipy.sparse' in str(type(array)): 32 | array = array.todense() 33 | return np.array(array, dtype=dtype) 34 | 35 | 36 | class Struct(object): 37 | def __init__(self, **kwargs): 38 | for key, val in kwargs.items(): 39 | setattr(self, key, val) 40 | 41 | 42 | class Masking(nn.Module): 43 | def __init__(self, config): 44 | super(Masking, self).__init__() 45 | ROOT_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), '..') 46 | with open(f'{ROOT_DIR}/data/FLAME2020/FLAME_masks/FLAME_masks.pkl', 'rb') as f: 47 | ss = pickle.load(f, encoding='latin1') 48 | self.masks = Struct(**ss) 49 | 50 | with open(f'{ROOT_DIR}/data/FLAME2020/generic_model.pkl', 'rb') as f: 51 | ss = pickle.load(f, encoding='latin1') 52 | flame_model = Struct(**ss) 53 | 54 | self.masked_faces = None 55 | 56 | self.cfg = config.mask_weights 57 | self.dtype = torch.float32 58 | self.register_buffer('faces', to_tensor(to_np(flame_model.f, dtype=np.int64), dtype=torch.long)) 59 | self.register_buffer('vertices', to_tensor(to_np(flame_model.v_template), dtype=self.dtype)) 60 | 61 | self.neighbours = {} 62 | for f in self.faces.numpy(): 63 | for v in f: 64 | if str(v) not in self.neighbours: 65 | self.neighbours[str(v)] = set() 66 | for a in list(filter(lambda i: i != v, f)): 67 | self.neighbours[str(v)].add(a) 68 | 69 | def get_faces(self): 70 | return self.faces 71 | 72 | def get_mask_face(self): 73 | return self.masks.face 74 | 75 | def get_mask_eyes(self): 76 | left = self.masks.left_eyeball 77 | right = self.masks.right_eyeball 78 | 79 | return np.unique(np.concatenate((left, right))) 80 | 81 | def get_mask_forehead(self): 82 | return self.masks.forehead 83 | 84 | def get_mask_lips(self): 85 | return self.masks.lips 86 | 87 | def get_mask_eye_region(self): 88 | return self.masks.eye_region 89 | 90 | def get_mask_lr_eye_region(self): 91 | left = self.masks.left_eye_region 92 | right = self.masks.right_eye_region 93 | 94 | return np.unique(np.concatenate((left, right, self.get_mask_eyes()))) 95 | 96 | def get_mask_nose(self): 97 | return self.masks.nose 98 | 99 | def get_mask_ears(self): 100 | left = self.masks.left_ear 101 | right = self.masks.right_ear 102 | 103 | return np.unique(np.concatenate((left, right))) 104 | 105 | def get_triangle_face_mask(self): 106 | m = self.masks.face 107 | return self.get_triangle_mask(m) 108 | 109 | def get_triangle_eyes_mask(self): 110 | m = self.get_mask_eyes() 111 | return self.get_triangle_mask(m) 112 | 113 | def get_triangle_whole_mask(self): 114 | m = self.get_whole_mask() 115 | return self.get_triangle_mask(m) 116 | 117 | def get_triangle_mask(self, m): 118 | f = self.faces.cpu().numpy() 119 | selected = [] 120 | for i in range(f.shape[0]): 121 | l = f[i] 122 | valid = 0 123 | for j in range(3): 124 | if l[j] in m: 125 | valid += 1 126 | if valid == 3: 127 | selected.append(i) 128 | 129 | return np.unique(selected) 130 | 131 | def make_soft(self, mask, value, degree=4): 132 | soft = [] 133 | mask = set(mask) 134 | for ring in range(degree): 135 | soft_ring = [] 136 | for v in mask.copy(): 137 | for n in self.neighbours[str(v)]: 138 | if n in mask: 139 | continue 140 | soft_ring.append(n) 141 | mask.add(n) 142 | 143 | soft.append((soft_ring, value / (ring + 2))) 144 | 145 | return soft 146 | 147 | def get_binary_triangle_mask(self): 148 | mask = self.get_whole_mask() 149 | faces = self.faces.cpu().numpy() 150 | reduced_faces = [] 151 | for f in faces: 152 | valid = 0 153 | for v in f: 154 | if v in mask: 155 | valid += 1 156 | reduced_faces.append(True if valid == 3 else False) 157 | 158 | return reduced_faces 159 | 160 | def get_masked_faces(self): 161 | if self.masked_faces is None: 162 | faces = self.faces.cpu().numpy() 163 | vertices = self.vertices.cpu().numpy() 164 | m = Trimesh(vertices=vertices, faces=faces, process=False) 165 | m.update_faces(self.get_binary_triangle_mask()) 166 | self.masked_faces = torch.from_numpy(np.array(m.faces)).cuda().long()[None] 167 | 168 | return self.masked_faces 169 | 170 | def get_weights_per_triangle(self): 171 | mask = torch.ones_like(self.get_faces()[None]).detach() * self.cfg.whole 172 | 173 | mask[:, self.get_triangle_eyes_mask(), :] = self.cfg.eyes 174 | mask[:, self.get_triangle_face_mask(), :] = self.cfg.face 175 | 176 | return mask[:, :, 0:1] 177 | 178 | def get_weights_per_vertex(self): 179 | mask = torch.ones_like(self.vertices[None]).detach() * self.cfg.whole 180 | 181 | mask[:, self.get_mask_eyes(), :] = self.cfg.eyes 182 | mask[:, self.get_mask_ears(), :] = self.cfg.ears 183 | mask[:, self.get_mask_face(), :] = self.cfg.face 184 | 185 | return mask 186 | 187 | def get_masked_mesh(self, vertices, triangle_mask): 188 | if len(vertices.shape) == 2: 189 | vertices = vertices[None] 190 | B, N, V = vertices.shape 191 | faces = self.faces.cpu().numpy() 192 | masked_vertices = torch.empty(0, 0, 3).cuda() 193 | masked_faces = torch.empty(0, 0, 3).cuda() 194 | for i in range(B): 195 | m = Trimesh(vertices=vertices[i].detach().cpu().numpy(), faces=faces, process=False) 196 | m.update_faces(triangle_mask) 197 | m.process() 198 | f = torch.from_numpy(np.array(m.faces)).cuda()[None] 199 | v = torch.from_numpy(np.array(m.vertices)).cuda()[None].float() 200 | if masked_vertices.shape[1] != v.shape[1]: 201 | masked_vertices = torch.empty(0, v.shape[1], 3).cuda() 202 | if masked_faces.shape[1] != f.shape[1]: 203 | masked_faces = torch.empty(0, f.shape[1], 3).cuda() 204 | masked_vertices = torch.cat([masked_vertices, v]) 205 | masked_faces = torch.cat([masked_faces, f]) 206 | 207 | return masked_vertices, masked_faces 208 | -------------------------------------------------------------------------------- /utils/util.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is 4 | # holder of all proprietary rights on this computer program. 5 | # You can only use this computer program if you have closed 6 | # a license agreement with MPG or you get the right to use the computer 7 | # program from someone who is authorized to grant you that right. 8 | # Any use of the computer program without a valid license is prohibited and 9 | # liable to prosecution. 10 | # 11 | # Copyright©2023 Max-Planck-Gesellschaft zur Förderung 12 | # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute 13 | # for Intelligent Systems. All rights reserved. 14 | # 15 | # Contact: mica@tue.mpg.de 16 | 17 | 18 | import importlib 19 | 20 | import cv2 21 | import numpy as np 22 | import torch 23 | import torch.nn.functional as F 24 | import torchvision 25 | 26 | 27 | def find_model_using_name(model_dir, model_name): 28 | # adapted from pix2pix framework: https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix/blob/master/models/__init__.py#L25 29 | # import "model_dir/modelname.py" 30 | model_filename = model_dir + "." + model_name 31 | modellib = importlib.import_module(model_filename, package=model_dir) 32 | 33 | # In the file, the class called ModelName() will 34 | # be instantiated. It has to be a subclass of BaseModel, 35 | # and it is case-insensitive. 36 | model = None 37 | target_model_name = model_name.replace('_', '') 38 | for name, cls in modellib.__dict__.items(): 39 | # if name.lower() == target_model_name.lower() and issubclass(cls, BaseModel): 40 | if name.lower() == target_model_name.lower(): 41 | model = cls 42 | 43 | if model is None: 44 | print("In %s.py, there should be a class with class name that matches %s in lowercase." % (model_filename, target_model_name)) 45 | exit(0) 46 | 47 | return model 48 | 49 | 50 | def visualize_grid(visdict, savepath=None, size=224, dim=1, return_gird=True): 51 | ''' 52 | image range should be [0,1] 53 | dim: 2 for horizontal. 1 for vertical 54 | ''' 55 | assert dim == 1 or dim == 2 56 | grids = {} 57 | for key in visdict: 58 | b, c, h, w = visdict[key].shape 59 | if dim == 2: 60 | new_h = size 61 | new_w = int(w * size / h) 62 | elif dim == 1: 63 | new_h = int(h * size / w) 64 | new_w = size 65 | grids[key] = torchvision.utils.make_grid(F.interpolate(visdict[key], [new_h, new_w]).detach().cpu(), nrow=b, padding=0) 66 | grid = torch.cat(list(grids.values()), dim) 67 | grid_image = (grid.numpy().transpose(1, 2, 0).copy() * 255)[:, :, [2, 1, 0]] 68 | grid_image = np.minimum(np.maximum(grid_image, 0), 255).astype(np.uint8) 69 | if savepath: 70 | cv2.imwrite(savepath, grid_image) 71 | if return_gird: 72 | return grid_image 73 | --------------------------------------------------------------------------------