all lib

handler.py

@@ -14,6 +14,10 @@ from diffusers.pipelines.controlnet.multicontrolnet import MultiControlNetModel
 
 from huggingface_hub import hf_hub_download
 
+import sys
+root_local = './'
+sys.path.insert(0, root_local)
+
 from insightface.app import FaceAnalysis
 
 from style_template import styles
@@ -48,7 +52,7 @@ class EndpointHandler():
         #     providers=["CPUExecutionProvider"],
         # )
         self.app = FaceAnalysis(
-            name="antelopev2",
+            name="buffalo_l",
             root="./",
             providers=["CPUExecutionProvider"],
         )

insightface/.gitignore

@@ -0,0 +1,103 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+env/
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+.hypothesis/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# pyenv
+.python-version
+
+# celery beat schedule file
+celerybeat-schedule
+
+# SageMath parsed files
+*.sage.py
+
+# dotenv
+.env
+
+# virtualenv
+.venv
+venv/
+ENV/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+
+.DS_Store

insightface/CODE_OF_CONDUCT.md

@@ -0,0 +1,128 @@
+# Contributor Covenant Code of Conduct
+
+## Our Pledge
+
+We as members, contributors, and leaders pledge to make participation in our
+community a harassment-free experience for everyone, regardless of age, body
+size, visible or invisible disability, ethnicity, sex characteristics, gender
+identity and expression, level of experience, education, socio-economic status,
+nationality, personal appearance, race, religion, or sexual identity
+and orientation.
+
+We pledge to act and interact in ways that contribute to an open, welcoming,
+diverse, inclusive, and healthy community.
+
+## Our Standards
+
+Examples of behavior that contributes to a positive environment for our
+community include:
+
+* Demonstrating empathy and kindness toward other people
+* Being respectful of differing opinions, viewpoints, and experiences
+* Giving and gracefully accepting constructive feedback
+* Accepting responsibility and apologizing to those affected by our mistakes,
+  and learning from the experience
+* Focusing on what is best not just for us as individuals, but for the
+  overall community
+
+Examples of unacceptable behavior include:
+
+* The use of sexualized language or imagery, and sexual attention or
+  advances of any kind
+* Trolling, insulting or derogatory comments, and personal or political attacks
+* Public or private harassment
+* Publishing others' private information, such as a physical or email
+  address, without their explicit permission
+* Other conduct which could reasonably be considered inappropriate in a
+  professional setting
+
+## Enforcement Responsibilities
+
+Community leaders are responsible for clarifying and enforcing our standards of
+acceptable behavior and will take appropriate and fair corrective action in
+response to any behavior that they deem inappropriate, threatening, offensive,
+or harmful.
+
+Community leaders have the right and responsibility to remove, edit, or reject
+comments, commits, code, wiki edits, issues, and other contributions that are
+not aligned to this Code of Conduct, and will communicate reasons for moderation
+decisions when appropriate.
+
+## Scope
+
+This Code of Conduct applies within all community spaces, and also applies when
+an individual is officially representing the community in public spaces.
+Examples of representing our community include using an official e-mail address,
+posting via an official social media account, or acting as an appointed
+representative at an online or offline event.
+
+## Enforcement
+
+Instances of abusive, harassing, or otherwise unacceptable behavior may be
+reported to the community leaders responsible for enforcement at
+[email protected].
+All complaints will be reviewed and investigated promptly and fairly.
+
+All community leaders are obligated to respect the privacy and security of the
+reporter of any incident.
+
+## Enforcement Guidelines
+
+Community leaders will follow these Community Impact Guidelines in determining
+the consequences for any action they deem in violation of this Code of Conduct:
+
+### 1. Correction
+
+**Community Impact**: Use of inappropriate language or other behavior deemed
+unprofessional or unwelcome in the community.
+
+**Consequence**: A private, written warning from community leaders, providing
+clarity around the nature of the violation and an explanation of why the
+behavior was inappropriate. A public apology may be requested.
+
+### 2. Warning
+
+**Community Impact**: A violation through a single incident or series
+of actions.
+
+**Consequence**: A warning with consequences for continued behavior. No
+interaction with the people involved, including unsolicited interaction with
+those enforcing the Code of Conduct, for a specified period of time. This
+includes avoiding interactions in community spaces as well as external channels
+like social media. Violating these terms may lead to a temporary or
+permanent ban.
+
+### 3. Temporary Ban
+
+**Community Impact**: A serious violation of community standards, including
+sustained inappropriate behavior.
+
+**Consequence**: A temporary ban from any sort of interaction or public
+communication with the community for a specified period of time. No public or
+private interaction with the people involved, including unsolicited interaction
+with those enforcing the Code of Conduct, is allowed during this period.
+Violating these terms may lead to a permanent ban.
+
+### 4. Permanent Ban
+
+**Community Impact**: Demonstrating a pattern of violation of community
+standards, including sustained inappropriate behavior,  harassment of an
+individual, or aggression toward or disparagement of classes of individuals.
+
+**Consequence**: A permanent ban from any sort of public interaction within
+the community.
+
+## Attribution
+
+This Code of Conduct is adapted from the [Contributor Covenant][homepage],
+version 2.0, available at
+https://www.contributor-covenant.org/version/2/0/code_of_conduct.html.
+
+Community Impact Guidelines were inspired by [Mozilla's code of conduct
+enforcement ladder](https://github.com/mozilla/diversity).
+
+[homepage]: https://www.contributor-covenant.org
+
+For answers to common questions about this code of conduct, see the FAQ at
+https://www.contributor-covenant.org/faq. Translations are available at
+https://www.contributor-covenant.org/translations.

insightface/README.md

@@ -0,0 +1,258 @@
+
+# InsightFace: 2D and 3D Face Analysis Project
+
+<div align="left">
+  <img src="https://insightface.ai/assets/img/custom/logo3.jpg" width="240"/>
+</div>
+
+InsightFace project is mainly maintained By [Jia Guo](mailto:[email protected]?subject=[GitHub]%20InsightFace%20Project) and [Jiankang Deng](https://jiankangdeng.github.io/). 
+
+For all main contributors, please check [contributing](#contributing).
+
+## License
+
+The code of InsightFace is released under the MIT License. There is no limitation for both academic and commercial usage.
+
+The training data containing the annotation (and the models trained with these data) are available for non-commercial research purposes only.
+
+Both manual-downloading models from our github repo and auto-downloading models with our [python-library](python-package) follow the above license policy(which is for non-commercial research purposes only).
+
+## Top News
+
+**`2023-08-08`**: We released the implementation of [Generalizing Gaze Estimation with Weak-Supervision from Synthetic Views](https://arxiv.org/abs/2212.02997) at [reconstruction/gaze](reconstruction/gaze).
+
+**`2023-05-03`**: We have launched the ongoing version of wild face anti-spoofing challenge. See details [here](https://github.com/deepinsight/insightface/tree/master/challenges/cvpr23-fas-wild#updates).
+
+**`2023-04-01`**: We move the swapping demo to Discord bot, which support editing on Midjourney generated images, see detail at [web-demos/swapping_discord](web-demos/swapping_discord).
+
+**`2023-02-13`**: We launch a large scale in the wild face anti-spoofing challenge on CVPR23 Workshop, see details at [challenges/cvpr23-fas-wild](challenges/cvpr23-fas-wild).
+
+**`2022-11-28`**: Single line code for facial identity swapping in our python packge ver 0.7, please check the example [here](examples/in_swapper).
+
+**`2022-10-28`**: [MFR-Ongoing](http://iccv21-mfr.com) website is refactored, please create issues if there's any bug.
+
+**`2022-09-22`**: Now we have [web-demos](web-demos): [face-localization](http://demo.insightface.ai:7007/), [face-recognition](http://demo.insightface.ai:7008/), and [face-swapping](http://demo.insightface.ai:7009/).
+
+**`2022-08-12`**: We achieved Rank-1st of 
+[Perspective Projection Based Monocular 3D Face Reconstruction Challenge](https://tianchi.aliyun.com/competition/entrance/531961/introduction)
+of [ECCV-2022 WCPA Workshop](https://sites.google.com/view/wcpa2022), [paper](https://arxiv.org/abs/2208.07142) and [code](reconstruction/jmlr).
+
+**`2022-03-30`**: [Partial FC](https://arxiv.org/abs/2203.15565) accepted by CVPR-2022.
+
+**`2022-02-23`**: [SCRFD](detection/scrfd) accepted by [ICLR-2022](https://iclr.cc/Conferences/2022).
+
+**`2021-11-30`**: [MFR-Ongoing](challenges/mfr) challenge launched(same with IFRT), which is an extended version of [iccv21-mfr](challenges/iccv21-mfr).
+
+**`2021-10-29`**: We achieved 1st place on the [VISA track](https://pages.nist.gov/frvt/plots/11/visa.html) of [NIST-FRVT 1:1](https://pages.nist.gov/frvt/html/frvt11.html) by using Partial FC (Xiang An, Jiankang Deng, Jia Guo).
+
+**`2021-10-11`**: [Leaderboard](https://insightface.ai/mfr21) of [ICCV21 - Masked Face Recognition Challenge](challenges/iccv21-mfr) released. Video: [Youtube](https://www.youtube.com/watch?v=lL-7l5t6x2w), [Bilibili](https://www.bilibili.com/video/BV15b4y1h79N/).
+
+**`2021-06-05`**: We launch a [Masked Face Recognition Challenge & Workshop](challenges/iccv21-mfr) on ICCV 2021.
+
+
+
+## Introduction
+
+[InsightFace](https://insightface.ai) is an open source 2D&3D deep face analysis toolbox, mainly based on PyTorch and MXNet. 
+
+Please check our [website](https://insightface.ai) for detail.
+
+The master branch works with **PyTorch 1.6+** and/or **MXNet=1.6-1.8**, with **Python 3.x**.
+
+InsightFace efficiently implements a rich variety of state of the art algorithms of face recognition, face detection and face alignment, which optimized for both training and deployment.
+
+## Quick Start
+
+Please start with our [python-package](python-package/), for testing detection, recognition and alignment models on input images.
+
+
+### ArcFace Video Demo
+
+
+[<img src=https://insightface.ai/assets/img/github/facerecognitionfromvideo.PNG width="760" />](https://www.youtube.com/watch?v=y-D1tReryGA&t=81s)
+
+
+Please click the image to watch the Youtube video. For Bilibili users, click [here](https://www.bilibili.com/video/av38041494?from=search&seid=11501833604850032313).
+
+
+
+## Projects
+
+The [page](https://insightface.ai/projects) on InsightFace website also describes all supported projects in InsightFace.
+
+You may also interested in some [challenges](https://insightface.ai/challenges) hold by InsightFace.
+
+
+
+## Face Recognition
+
+### Introduction
+
+In this module, we provide training data, network settings and loss designs for deep face recognition.
+
+The supported methods are as follows:
+
+- [x] [ArcFace_mxnet (CVPR'2019)](recognition/arcface_mxnet)
+- [x] [ArcFace_torch (CVPR'2019)](recognition/arcface_torch)
+- [x] [SubCenter ArcFace (ECCV'2020)](recognition/subcenter_arcface)
+- [x] [PartialFC_mxnet (CVPR'2022)](recognition/partial_fc)
+- [x] [PartialFC_torch (CVPR'2022)](recognition/arcface_torch)
+- [x] [VPL (CVPR'2021)](recognition/vpl)
+- [x] [Arcface_oneflow](recognition/arcface_oneflow)
+- [x] [ArcFace_Paddle (CVPR'2019)](recognition/arcface_paddle)
+
+Commonly used network backbones are included in most of the methods, such as IResNet, MobilefaceNet, MobileNet, InceptionResNet_v2, DenseNet, etc..
+
+
+### Datasets
+
+The training data includes, but not limited to the cleaned MS1M, VGG2 and CASIA-Webface datasets, which were already packed in MXNet binary format. Please [dataset](recognition/_datasets_) page for detail.
+
+### Evaluation
+
+We provide standard IJB and Megaface evaluation pipelines in [evaluation](recognition/_evaluation_)
+
+
+### Pretrained Models
+
+**Please check [Model-Zoo](https://github.com/deepinsight/insightface/wiki/Model-Zoo) for more pretrained models.**
+
+### Third-party Re-implementation of ArcFace
+
+- TensorFlow: [InsightFace_TF](https://github.com/auroua/InsightFace_TF)
+- TensorFlow: [tf-insightface](https://github.com/AIInAi/tf-insightface)
+- TensorFlow:[insightface](https://github.com/Fei-Wang/insightface)
+- PyTorch: [InsightFace_Pytorch](https://github.com/TreB1eN/InsightFace_Pytorch)
+- PyTorch: [arcface-pytorch](https://github.com/ronghuaiyang/arcface-pytorch)
+- Caffe: [arcface-caffe](https://github.com/xialuxi/arcface-caffe)
+- Caffe: [CombinedMargin-caffe](https://github.com/gehaocool/CombinedMargin-caffe)
+- Tensorflow: [InsightFace-tensorflow](https://github.com/luckycallor/InsightFace-tensorflow)
+- TensorRT: [wang-xinyu/tensorrtx](https://github.com/wang-xinyu/tensorrtx)  
+- TensorRT: [InsightFace-REST](https://github.com/SthPhoenix/InsightFace-REST)
+- ONNXRuntime C++: [ArcFace-ONNXRuntime](https://github.com/DefTruth/lite.ai.toolkit/blob/main/lite/ort/cv/glint_arcface.cpp)
+- ONNXRuntime Go: [arcface-go](https://github.com/jack139/arcface-go)
+- MNN: [ArcFace-MNN](https://github.com/DefTruth/lite.ai.toolkit/blob/main/lite/mnn/cv/mnn_glint_arcface.cpp)
+- TNN: [ArcFace-TNN](https://github.com/DefTruth/lite.ai.toolkit/blob/main/lite/tnn/cv/tnn_glint_arcface.cpp)
+- NCNN: [ArcFace-NCNN](https://github.com/DefTruth/lite.ai.toolkit/blob/main/lite/ncnn/cv/ncnn_glint_arcface.cpp)
+
+## Face Detection
+
+### Introduction
+
+<div align="left">
+  <img src="https://insightface.ai/assets/img/github/11513D05.jpg" width="640"/>
+</div>
+
+In this module, we provide training data with annotation, network settings and loss designs for face detection training, evaluation and inference.
+
+The supported methods are as follows:
+
+- [x] [RetinaFace (CVPR'2020)](detection/retinaface)
+- [x] [SCRFD (Arxiv'2021)](detection/scrfd)
+- [x] [blazeface_paddle](detection/blazeface_paddle)
+
+[RetinaFace](detection/retinaface) is a practical single-stage face detector which is accepted by [CVPR 2020](https://openaccess.thecvf.com/content_CVPR_2020/html/Deng_RetinaFace_Single-Shot_Multi-Level_Face_Localisation_in_the_Wild_CVPR_2020_paper.html). We provide training code, training dataset, pretrained models and evaluation scripts. 
+
+[SCRFD](detection/scrfd) is an efficient high accuracy face detection approach which is initialy described in [Arxiv](https://arxiv.org/abs/2105.04714). We provide an easy-to-use pipeline to train high efficiency face detectors with NAS supporting.
+
+
+## Face Alignment
+
+### Introduction
+
+<div align="left">
+  <img src="https://insightface.ai/assets/img/custom/thumb_sdunet.png" width="600"/>
+</div>
+
+In this module, we provide datasets and training/inference pipelines for face alignment.
+
+Supported methods:
+
+- [x] [SDUNets (BMVC'2018)](alignment/heatmap)
+- [x] [SimpleRegression](alignment/coordinate_reg)
+
+
+[SDUNets](alignment/heatmap) is a heatmap based method which accepted on [BMVC](http://bmvc2018.org/contents/papers/0051.pdf).
+
+[SimpleRegression](alignment/coordinate_reg) provides very lightweight facial landmark models with fast coordinate regression. The input of these models is loose cropped face image while the output is the direct landmark coordinates.
+
+
+## Citation
+
+If you find *InsightFace* useful in your research, please consider to cite the following related papers:
+
+```
+@inproceedings{ren2023pbidr,
+  title={Facial Geometric Detail Recovery via Implicit Representation},
+  author={Ren, Xingyu and Lattas, Alexandros and Gecer, Baris and Deng, Jiankang and Ma, Chao and Yang, Xiaokang},
+  booktitle={2023 IEEE 17th International Conference on Automatic Face and Gesture Recognition (FG)},  
+  year={2023}
+ }
+
+@article{guo2021sample,
+  title={Sample and Computation Redistribution for Efficient Face Detection},
+  author={Guo, Jia and Deng, Jiankang and Lattas, Alexandros and Zafeiriou, Stefanos},
+  journal={arXiv preprint arXiv:2105.04714},
+  year={2021}
+}
+
+@inproceedings{gecer2021ostec,
+  title={OSTeC: One-Shot Texture Completion},
+  author={Gecer, Baris and Deng, Jiankang and Zafeiriou, Stefanos},
+  booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
+  year={2021}
+}
+
+@inproceedings{an2020partical_fc,
+  title={Partial FC: Training 10 Million Identities on a Single Machine},
+  author={An, Xiang and Zhu, Xuhan and Xiao, Yang and Wu, Lan and Zhang, Ming and Gao, Yuan and Qin, Bin and
+  Zhang, Debing and Fu Ying},
+  booktitle={Arxiv 2010.05222},
+  year={2020}
+}
+
+@inproceedings{deng2020subcenter,
+  title={Sub-center ArcFace: Boosting Face Recognition by Large-scale Noisy Web Faces},
+  author={Deng, Jiankang and Guo, Jia and Liu, Tongliang and Gong, Mingming and Zafeiriou, Stefanos},
+  booktitle={Proceedings of the IEEE Conference on European Conference on Computer Vision},
+  year={2020}
+}
+
+@inproceedings{Deng2020CVPR,
+title = {RetinaFace: Single-Shot Multi-Level Face Localisation in the Wild},
+author = {Deng, Jiankang and Guo, Jia and Ververas, Evangelos and Kotsia, Irene and Zafeiriou, Stefanos},
+booktitle = {CVPR},
+year = {2020}
+}
+
+@inproceedings{guo2018stacked,
+  title={Stacked Dense U-Nets with Dual Transformers for Robust Face Alignment},
+  author={Guo, Jia and Deng, Jiankang and Xue, Niannan and Zafeiriou, Stefanos},
+  booktitle={BMVC},
+  year={2018}
+}
+
+@article{deng2018menpo,
+  title={The Menpo benchmark for multi-pose 2D and 3D facial landmark localisation and tracking},
+  author={Deng, Jiankang and Roussos, Anastasios and Chrysos, Grigorios and Ververas, Evangelos and Kotsia, Irene and Shen, Jie and Zafeiriou, Stefanos},
+  journal={IJCV},
+  year={2018}
+}
+
+@inproceedings{deng2018arcface,
+title={ArcFace: Additive Angular Margin Loss for Deep Face Recognition},
+author={Deng, Jiankang and Guo, Jia and Niannan, Xue and Zafeiriou, Stefanos},
+booktitle={CVPR},
+year={2019}
+}
+```
+
+## Contributing
+
+Main contributors:
+
+- [Jia Guo](https://github.com/nttstar), ``guojia[at]gmail.com``
+- [Jiankang Deng](https://github.com/jiankangdeng) ``jiankangdeng[at]gmail.com``
+- [Xiang An](https://github.com/anxiangsir) ``anxiangsir[at]gmail.com``
+- [Jack Yu](https://github.com/szad670401) ``jackyu961127[at]gmail.com``
+- [Baris Gecer](https://barisgecer.github.io/) ``barisgecer[at]msn.com``

insightface/alignment/README.md

@@ -0,0 +1,42 @@
+## Face Alignment
+
+
+<div align="left">
+  <img src="https://insightface.ai/assets/img/custom/logo3.jpg" width="240"/>
+</div>
+
+
+## Introduction
+
+These are the face alignment methods of [InsightFace](https://insightface.ai)
+
+
+<div align="left">
+  <img src="https://insightface.ai/assets/img/custom/thumb_sdunet.png" width="600"/>
+</div>
+
+
+### Datasets
+
+  Please refer to [datasets](_datasets_) page for the details of face alignment datasets used for training and evaluation.
+
+### Evaluation
+
+  Please refer to [evaluation](_evaluation_) page for the details of face alignment evaluation.
+
+
+## Methods
+
+
+Supported methods:
+
+- [x] [SDUNets (BMVC'2018)](heatmap)
+- [x] [SimpleRegression](coordinate_reg)
+- [x] [Alignment By Face Synthetics](synthetics)
+
+
+## Contributing
+
+We appreciate all contributions to improve the face alignment model zoo of InsightFace. 
+
+

insightface/alignment/_datasets_/README.md

@@ -0,0 +1,57 @@
+# Face Alignment Datasets
+
+(Updating)
+
+## Training Datasets
+
+### Menpo2D-Train
+
+https://ibug.doc.ic.ac.uk/resources/2nd-facial-landmark-tracking-competition-menpo-ben/
+
+### 300W-Train
+
+https://ibug.doc.ic.ac.uk/resources/300-W/
+
+
+### LFPW
+
+https://neerajkumar.org/databases/lfpw/
+
+### Helen
+
+http://www.ifp.illinois.edu/~vuongle2/helen/
+
+### AFW
+
+### AFLW
+
+https://www.tugraz.at/institute/icg/research/team-bischof/lrs/downloads/aflw/
+
+### FDDB
+
+
+### Face Synthetics
+
+https://github.com/microsoft/FaceSynthetics
+
+### 300W-LP (3D annotation)
+
+http://www.cbsr.ia.ac.cn/users/xiangyuzhu/projects/3DDFA/main.htm
+
+## Test Datasets
+
+### 300W-Test
+
+https://ibug.doc.ic.ac.uk/resources/300-W/
+
+### COFW
+
+http://www.vision.caltech.edu/xpburgos/ICCV13/#dataset
+
+### Menpo2D-Test
+
+https://ibug.doc.ic.ac.uk/resources/2nd-facial-landmark-tracking-competition-menpo-ben/
+
+### AFLW2000-3D (3D annotation)
+
+http://www.cbsr.ia.ac.cn/users/xiangyuzhu/projects/3DDFA/main.htm

insightface/alignment/coordinate_reg/README.md

@@ -0,0 +1,58 @@
+### Introduction
+
+ Here we provide some lightweight facial landmark models with fast coordinate regression.
+ The input of these models is loose cropped face image while the output is the direct landmark coordinates. 
+
+
+### Pretrained Models
+
+- **Model ``2d106det``**
+
+  **2021.07: We now support model inference by our `insightface` python package, please check [image_infer.py](image_infer.py) for detail.**
+
+  Given face detection bounding box, predict 2d-106 landmarks. Mainly used for static image inference. 
+  
+  Backbone: MobileNet-0.5, size 5MB.
+  
+  Input: size 192x192, loose cropped detection bounding-box.
+  
+  Download link:
+  
+  [baidu cloud](https://pan.baidu.com/s/10m5GmtNV5snynDrq3KqIdg) (code: ``lqvv``)
+  
+  [google drive](https://drive.google.com/file/d/13Pz8mH-a1s7RXpq_jFUXxaqCpDUE0oSr/view?usp=sharing)
+  
+  
+
+- **Model ``2d106track``** 
+
+  Given landmarks bounding box, predict 2d-106 landmarks. Used for video landmarks tracking.
+  
+  Download link: coming soon
+
+### Visualization
+
+
+<p align="center">Points mark-up(ordered by point names):</p>
+
+<div align="center">
+	<img src="https://github.com/nttstar/insightface-resources/blob/master/alignment/images/2d106markup.jpg" alt="markup" width="320">
+</div>
+
+
+<p align="center">Image result:</p>
+
+<div align="center">
+	<img src="https://github.com/nttstar/insightface-resources/blob/master/alignment/images/t1_out.jpg" alt="imagevis" width="800">
+</div>
+
+
+<p align="center">Video result:</p>
+
+<div align="center">
+	<img src="https://github.com/nttstar/insightface-resources/blob/master/alignment/images/C_jiaguo.gif" alt="videovis" width="240">
+</div>
+
+
+### FAQ
+

insightface/alignment/coordinate_reg/image_infer.py

@@ -0,0 +1,23 @@
+import cv2
+import numpy as np
+import os
+import insightface
+from insightface.app import FaceAnalysis
+from insightface.data import get_image as ins_get_image
+
+if __name__ == '__main__':
+    app = FaceAnalysis(allowed_modules=['detection', 'landmark_2d_106'])
+    app.prepare(ctx_id=0, det_size=(640, 640))
+    img = ins_get_image('t1')
+    faces = app.get(img)
+    #assert len(faces)==6
+    tim = img.copy()
+    color = (200, 160, 75)
+    for face in faces:
+        lmk = face.landmark_2d_106
+        lmk = np.round(lmk).astype(np.int)
+        for i in range(lmk.shape[0]):
+            p = tuple(lmk[i])
+            cv2.circle(tim, p, 1, color, 1, cv2.LINE_AA)
+    cv2.imwrite('./test_out.jpg', tim)
+

insightface/alignment/heatmap/README.md

@@ -0,0 +1,10 @@
+We provide our implementation of ``Stacked Dense U-Nets with Dual Transformers for Robust Face Alignment`` here at [BMVC](http://bmvc2018.org/contents/papers/0051.pdf) or link at [Arxiv](https://arxiv.org/abs/1812.01936).
+
+We also provide some popular heatmap based approaches like stacked hourglass, etc..  You can define different loss-type/network structure/dataset in ``config.py``(from ``sample_config.py``).
+
+For example, by default, you can train our approach by ``train.py --network sdu`` or train hourglass network by ``train.py --network hourglass``.
+
+2D training/validation dataset is now available at [baiducloud](https://pan.baidu.com/s/1kdquiIGTlK7l26SPWO_cmw) or [dropbox](https://www.dropbox.com/s/por6mbguegmywo6/bmvc_sdu_data2d.zip?dl=0)
+
+3D training/validation dataset is now available at [baiducloud](https://pan.baidu.com/s/1VjFWm6eEtIqGKk92GE2rgw) or [dropbox](https://www.dropbox.com/s/tjze176lh76nciw/bmvc_sdu_data3d.zip?dl=0)
+

insightface/alignment/heatmap/data.py

@@ -0,0 +1,354 @@
+# pylint: skip-file
+import mxnet as mx
+import numpy as np
+import sys, os
+import random
+import math
+import scipy.misc
+import cv2
+import logging
+import sklearn
+import datetime
+import img_helper
+from mxnet.io import DataIter
+from mxnet import ndarray as nd
+from mxnet import io
+from mxnet import recordio
+from PIL import Image
+from config import config
+from skimage import transform as tf
+
+
+class FaceSegIter(DataIter):
+    def __init__(self,
+                 batch_size,
+                 per_batch_size=0,
+                 path_imgrec=None,
+                 aug_level=0,
+                 force_mirror=False,
+                 exf=1,
+                 use_coherent=0,
+                 args=None,
+                 data_name="data",
+                 label_name="softmax_label"):
+        self.aug_level = aug_level
+        self.force_mirror = force_mirror
+        self.use_coherent = use_coherent
+        self.exf = exf
+        self.batch_size = batch_size
+        self.per_batch_size = per_batch_size
+        self.data_name = data_name
+        self.label_name = label_name
+        assert path_imgrec
+        logging.info('loading recordio %s...', path_imgrec)
+        path_imgidx = path_imgrec[0:-4] + ".idx"
+        self.imgrec = mx.recordio.MXIndexedRecordIO(path_imgidx, path_imgrec,
+                                                    'r')  # pylint: disable=redefined-variable-type
+        self.oseq = list(self.imgrec.keys)
+        print('train size', len(self.oseq))
+        self.cur = 0
+        self.reset()
+        self.data_shape = (3, config.input_img_size, config.input_img_size)
+        self.num_classes = config.num_classes
+        self.input_img_size = config.input_img_size
+        #self.label_classes = self.num_classes
+        if config.losstype == 'heatmap':
+            if aug_level > 0:
+                self.output_label_size = config.output_label_size
+                self.label_shape = (self.num_classes, self.output_label_size,
+                                    self.output_label_size)
+            else:
+                self.output_label_size = self.input_img_size
+                #self.label_shape = (self.num_classes, 2)
+                self.label_shape = (self.num_classes, self.output_label_size,
+                                    self.output_label_size)
+        else:
+            if aug_level > 0:
+                self.output_label_size = config.output_label_size
+                self.label_shape = (self.num_classes, 2)
+            else:
+                self.output_label_size = self.input_img_size
+                #self.label_shape = (self.num_classes, 2)
+                self.label_shape = (self.num_classes, 2)
+        self.provide_data = [(data_name, (batch_size, ) + self.data_shape)]
+        self.provide_label = [(label_name, (batch_size, ) + self.label_shape)]
+        self.img_num = 0
+        self.invalid_num = 0
+        self.mode = 1
+        self.vis = 0
+        self.stats = [0, 0]
+        self.flip_order = [
+            16, 15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0, 26, 25,
+            24, 23, 22, 21, 20, 19, 18, 17, 27, 28, 29, 30, 35, 34, 33, 32, 31,
+            45, 44, 43, 42, 47, 46, 39, 38, 37, 36, 41, 40, 54, 53, 52, 51, 50,
+            49, 48, 59, 58, 57, 56, 55, 64, 63, 62, 61, 60, 67, 66, 65
+        ]
+        #self.mirror_set = [
+        #        (22,23),
+        #        (21,24),
+        #        (20,25),
+        #        (19,26),
+        #        (18,27),
+        #        (40,43),
+        #        (39,44),
+        #        (38,45),
+        #        (37,46),
+        #        (42,47),
+        #        (41,48),
+        #        (33,35),
+        #        (32,36),
+        #        (51,53),
+        #        (50,54),
+        #        (62,64),
+        #        (61,65),
+        #        (49,55),
+        #        (49,55),
+        #        (68,66),
+        #        (60,56),
+        #        (59,57),
+        #        (1,17),
+        #        (2,16),
+        #        (3,15),
+        #        (4,14),
+        #        (5,13),
+        #        (6,12),
+        #        (7,11),
+        #        (8,10),
+        #    ]
+
+    def get_data_shape(self):
+        return self.data_shape
+
+    #def get_label_shape(self):
+    #    return self.label_shape
+
+    def get_shape_dict(self):
+        D = {}
+        for (k, v) in self.provide_data:
+            D[k] = v
+        for (k, v) in self.provide_label:
+            D[k] = v
+        return D
+
+    def get_label_names(self):
+        D = []
+        for (k, v) in self.provide_label:
+            D.append(k)
+        return D
+
+    def reset(self):
+        #print('reset')
+        if self.aug_level == 0:
+            self.seq = self.oseq
+        else:
+            self.seq = []
+            for _ in range(self.exf):
+                _seq = self.oseq[:]
+                random.shuffle(_seq)
+                self.seq += _seq
+            print('train size after reset', len(self.seq))
+        self.cur = 0
+
+    def next_sample(self):
+        """Helper function for reading in next sample."""
+        if self.cur >= len(self.seq):
+            raise StopIteration
+        idx = self.seq[self.cur]
+        self.cur += 1
+        s = self.imgrec.read_idx(idx)
+        header, img = recordio.unpack(s)
+        img = mx.image.imdecode(img).asnumpy()
+        hlabel = np.array(header.label).reshape((self.num_classes, 2))
+        if not config.label_xfirst:
+            hlabel = hlabel[:, ::-1]  #convert to X/W first
+        annot = {'scale': config.base_scale}
+
+        #ul = np.array( (50000,50000), dtype=np.int32)
+        #br = np.array( (0,0), dtype=np.int32)
+        #for i in range(hlabel.shape[0]):
+        #  h = int(hlabel[i][0])
+        #  w = int(hlabel[i][1])
+        #  key = np.array((h,w))
+        #  ul = np.minimum(key, ul)
+        #  br = np.maximum(key, br)
+
+        return img, hlabel, annot
+
+    def get_flip(self, data, label):
+        data_flip = np.zeros_like(data)
+        label_flip = np.zeros_like(label)
+        for k in range(data_flip.shape[2]):
+            data_flip[:, :, k] = np.fliplr(data[:, :, k])
+        for k in range(label_flip.shape[0]):
+            label_flip[k, :] = np.fliplr(label[k, :])
+        #print(label[0,:].shape)
+        label_flip = label_flip[self.flip_order, :]
+        return data_flip, label_flip
+
+    def get_data(self, data, label, annot):
+        if self.vis:
+            self.img_num += 1
+            #if self.img_num<=self.vis:
+            #  filename = './vis/raw_%d.jpg' % (self.img_num)
+            #  print('save', filename)
+            #  draw = data.copy()
+            #  for i in range(label.shape[0]):
+            #    cv2.circle(draw, (label[i][1], label[i][0]), 1, (0, 0, 255), 2)
+            #  scipy.misc.imsave(filename, draw)
+
+        rotate = 0
+        #scale = 1.0
+        if 'scale' in annot:
+            scale = annot['scale']
+        else:
+            scale = max(data.shape[0], data.shape[1])
+        if 'center' in annot:
+            center = annot['center']
+        else:
+            center = np.array((data.shape[1] / 2, data.shape[0] / 2))
+        max_retry = 3
+        if self.aug_level == 0:  #validation mode
+            max_retry = 6
+        retry = 0
+        found = False
+        base_scale = scale
+        while retry < max_retry:
+            retry += 1
+            succ = True
+            _scale = base_scale
+            if self.aug_level > 0:
+                rotate = np.random.randint(-40, 40)
+                scale_config = 0.2
+                #rotate = 0
+                #scale_config = 0.0
+                scale_ratio = min(
+                    1 + scale_config,
+                    max(1 - scale_config,
+                        (np.random.randn() * scale_config) + 1))
+                _scale = int(base_scale * scale_ratio)
+                #translate = np.random.randint(-5, 5, size=(2,))
+                #center += translate
+            data_out, trans = img_helper.transform(data, center,
+                                                   self.input_img_size, _scale,
+                                                   rotate)
+            #data_out = img_helper.crop2(data, center, _scale, (self.input_img_size, self.input_img_size), rot=rotate)
+            label_out = np.zeros(self.label_shape, dtype=np.float32)
+            #print('out shapes', data_out.shape, label_out.shape)
+            for i in range(label.shape[0]):
+                pt = label[i].copy()
+                #pt = pt[::-1]
+                npt = img_helper.transform_pt(pt, trans)
+                if npt[0] >= data_out.shape[1] or npt[1] >= data_out.shape[
+                        0] or npt[0] < 0 or npt[1] < 0:
+                    succ = False
+                    #print('err npt', npt)
+                    break
+                if config.losstype == 'heatmap':
+                    pt_scale = float(
+                        self.output_label_size) / self.input_img_size
+                    npt *= pt_scale
+                    npt = npt.astype(np.int32)
+                    img_helper.gaussian(label_out[i], npt, config.gaussian)
+                else:
+                    label_out[i] = (npt / self.input_img_size)
+                #print('before gaussian', label_out[i].shape, pt.shape)
+                #trans = img_helper.transform(pt, center, _scale, (self.output_label_size, self.output_label_size), rot=rotate)
+                #print(trans.shape)
+                #if not img_helper.gaussian(label_out[i], trans, _g):
+                #    succ = False
+                #    break
+            if not succ:
+                if self.aug_level == 0:
+                    base_scale += 20
+                continue
+
+            flip_data_out = None
+            flip_label_out = None
+            if config.net_coherent:
+                flip_data_out, flip_label_out = self.get_flip(
+                    data_out, label_out)
+            elif ((self.aug_level > 0 and np.random.rand() < 0.5)
+                  or self.force_mirror):  #flip aug
+                flip_data_out, flip_label_out = self.get_flip(
+                    data_out, label_out)
+                data_out, label_out = flip_data_out, flip_label_out
+
+            found = True
+            break
+
+        #self.stats[0]+=1
+        if not found:
+            #self.stats[1]+=1
+            #print('find aug error', retry)
+            #print(self.stats)
+            #print('!!!ERR')
+            return None
+        #print('found with scale', _scale, rotate)
+
+        if self.vis > 0 and self.img_num <= self.vis:
+            print('crop', data.shape, center, _scale, rotate, data_out.shape)
+            filename = './vis/cropped_%d.jpg' % (self.img_num)
+            print('save', filename)
+            draw = data_out.copy()
+            alabel = label_out.copy()
+            for i in range(label.shape[0]):
+                a = cv2.resize(alabel[i],
+                               (self.input_img_size, self.input_img_size))
+                ind = np.unravel_index(np.argmax(a, axis=None), a.shape)
+                cv2.circle(draw, (ind[1], ind[0]), 1, (0, 0, 255), 2)
+            scipy.misc.imsave(filename, draw)
+            filename = './vis/raw_%d.jpg' % (self.img_num)
+            scipy.misc.imsave(filename, data)
+
+        return data_out, label_out, flip_data_out, flip_label_out
+
+    def next(self):
+        """Returns the next batch of data."""
+        #print('next')
+        batch_size = self.batch_size
+        batch_data = nd.empty((batch_size, ) + self.data_shape)
+        batch_label = nd.empty((batch_size, ) + self.label_shape)
+        i = 0
+        #self.cutoff = random.randint(800,1280)
+        try:
+            while i < batch_size:
+                #print('N', i)
+                data, label, annot = self.next_sample()
+                R = self.get_data(data, label, annot)
+                if R is None:
+                    continue
+                data_out, label_out, flip_data_out, flip_label_out = R
+                if not self.use_coherent:
+                    data = nd.array(data_out)
+                    data = nd.transpose(data, axes=(2, 0, 1))
+                    label = nd.array(label_out)
+                    #print(data.shape, label.shape)
+                    batch_data[i][:] = data
+                    batch_label[i][:] = label
+                    i += 1
+                else:
+                    data = nd.array(data_out)
+                    data = nd.transpose(data, axes=(2, 0, 1))
+                    label = nd.array(label_out)
+                    data2 = nd.array(flip_data_out)
+                    data2 = nd.transpose(data2, axes=(2, 0, 1))
+                    label2 = nd.array(flip_label_out)
+                    #M = nd.array(M)
+                    #print(data.shape, label.shape)
+                    batch_data[i][:] = data
+                    batch_label[i][:] = label
+                    #i+=1
+                    j = i + self.per_batch_size // 2
+                    batch_data[j][:] = data2
+                    batch_label[j][:] = label2
+                    i += 1
+                    if j % self.per_batch_size == self.per_batch_size - 1:
+                        i = j + 1
+        except StopIteration:
+            if i < batch_size:
+                raise StopIteration
+
+        #return {self.data_name  :  batch_data,
+        #        self.label_name :  batch_label}
+        #print(batch_data.shape, batch_label.shape)
+        return mx.io.DataBatch([batch_data], [batch_label], batch_size - i)

insightface/alignment/heatmap/img_helper.py

@@ -0,0 +1,86 @@
+import numpy as np
+import math
+import cv2
+from skimage import transform as stf
+
+
+def transform(data, center, output_size, scale, rotation):
+    scale_ratio = float(output_size) / scale
+    rot = float(rotation) * np.pi / 180.0
+    #translation = (output_size/2-center[0]*scale_ratio, output_size/2-center[1]*scale_ratio)
+    t1 = stf.SimilarityTransform(scale=scale_ratio)
+    cx = center[0] * scale_ratio
+    cy = center[1] * scale_ratio
+    t2 = stf.SimilarityTransform(translation=(-1 * cx, -1 * cy))
+    t3 = stf.SimilarityTransform(rotation=rot)
+    t4 = stf.SimilarityTransform(translation=(output_size / 2,
+                                              output_size / 2))
+    t = t1 + t2 + t3 + t4
+    trans = t.params[0:2]
+    #print('M', scale, rotation, trans)
+    cropped = cv2.warpAffine(data,
+                             trans, (output_size, output_size),
+                             borderValue=0.0)
+    return cropped, trans
+
+
+def transform_pt(pt, trans):
+    new_pt = np.array([pt[0], pt[1], 1.]).T
+    new_pt = np.dot(trans, new_pt)
+    #print('new_pt', new_pt.shape, new_pt)
+    return new_pt[:2]
+
+
+def gaussian(img, pt, sigma):
+    # Draw a 2D gaussian
+    assert (sigma >= 0)
+    if sigma == 0:
+        img[pt[1], pt[0]] = 1.0
+        return True
+    #assert pt[0]<=img.shape[1]
+    #assert pt[1]<=img.shape[0]
+
+    # Check that any part of the gaussian is in-bounds
+    ul = [int(pt[0] - 3 * sigma), int(pt[1] - 3 * sigma)]
+    br = [int(pt[0] + 3 * sigma + 1), int(pt[1] + 3 * sigma + 1)]
+    if (ul[0] > img.shape[1] or ul[1] >= img.shape[0] or br[0] < 0
+            or br[1] < 0):
+        # If not, just return the image as is
+        #print('gaussian error')
+        return False
+        #return img
+
+    # Generate gaussian
+    size = 6 * sigma + 1
+    x = np.arange(0, size, 1, float)
+    y = x[:, np.newaxis]
+    x0 = y0 = size // 2
+    # The gaussian is not normalized, we want the center value to equal 1
+    g = np.exp(-((x - x0)**2 + (y - y0)**2) / (2 * sigma**2))
+
+    # Usable gaussian range
+    g_x = max(0, -ul[0]), min(br[0], img.shape[1]) - ul[0]
+    g_y = max(0, -ul[1]), min(br[1], img.shape[0]) - ul[1]
+    # Image range
+    img_x = max(0, ul[0]), min(br[0], img.shape[1])
+    img_y = max(0, ul[1]), min(br[1], img.shape[0])
+
+    img[img_y[0]:img_y[1], img_x[0]:img_x[1]] = g[g_y[0]:g_y[1], g_x[0]:g_x[1]]
+    return True
+    #return img
+
+
+def estimate_trans_bbox(face, input_size, s=2.0):
+    w = face[2] - face[0]
+    h = face[3] - face[1]
+    wc = int((face[2] + face[0]) / 2)
+    hc = int((face[3] + face[1]) / 2)
+    im_size = max(w, h)
+    #size = int(im_size*1.2)
+    scale = input_size / (max(w, h) * s)
+    M = [
+        [scale, 0, input_size / 2 - wc * scale],
+        [0, scale, input_size / 2 - hc * scale],
+    ]
+    M = np.array(M)
+    return M

insightface/alignment/heatmap/metric.py

@@ -0,0 +1,107 @@
+import mxnet as mx
+import numpy as np
+import math
+import cv2
+from config import config
+
+
+class LossValueMetric(mx.metric.EvalMetric):
+    def __init__(self):
+        self.axis = 1
+        super(LossValueMetric, self).__init__('lossvalue',
+                                              axis=self.axis,
+                                              output_names=None,
+                                              label_names=None)
+        self.losses = []
+
+    def update(self, labels, preds):
+        loss = preds[0].asnumpy()[0]
+        self.sum_metric += loss
+        self.num_inst += 1.0
+
+
+class NMEMetric(mx.metric.EvalMetric):
+    def __init__(self):
+        self.axis = 1
+        super(NMEMetric, self).__init__('NME',
+                                        axis=self.axis,
+                                        output_names=None,
+                                        label_names=None)
+        #self.losses = []
+        self.count = 0
+
+    def cal_nme(self, label, pred_label):
+        nme = []
+        for b in range(pred_label.shape[0]):
+            record = [None] * 6
+            item = []
+            if label.ndim == 4:
+                _heatmap = label[b][36]
+                if np.count_nonzero(_heatmap) == 0:
+                    continue
+            else:  #ndim==3
+                #print(label[b])
+                if np.count_nonzero(label[b]) == 0:
+                    continue
+            for p in range(pred_label.shape[1]):
+                if label.ndim == 4:
+                    heatmap_gt = label[b][p]
+                    ind_gt = np.unravel_index(np.argmax(heatmap_gt, axis=None),
+                                              heatmap_gt.shape)
+                    ind_gt = np.array(ind_gt)
+                else:
+                    ind_gt = label[b][p]
+                    #ind_gt = ind_gt.astype(np.int)
+                    #print(ind_gt)
+                heatmap_pred = pred_label[b][p]
+                heatmap_pred = cv2.resize(
+                    heatmap_pred,
+                    (config.input_img_size, config.input_img_size))
+                ind_pred = np.unravel_index(np.argmax(heatmap_pred, axis=None),
+                                            heatmap_pred.shape)
+                ind_pred = np.array(ind_pred)
+                #print(ind_gt.shape)
+                #print(ind_pred)
+                if p == 36:
+                    #print('b', b, p, ind_gt, np.count_nonzero(heatmap_gt))
+                    record[0] = ind_gt
+                elif p == 39:
+                    record[1] = ind_gt
+                elif p == 42:
+                    record[2] = ind_gt
+                elif p == 45:
+                    record[3] = ind_gt
+                if record[4] is None or record[5] is None:
+                    record[4] = ind_gt
+                    record[5] = ind_gt
+                else:
+                    record[4] = np.minimum(record[4], ind_gt)
+                    record[5] = np.maximum(record[5], ind_gt)
+                #print(ind_gt.shape, ind_pred.shape)
+                value = np.sqrt(np.sum(np.square(ind_gt - ind_pred)))
+                item.append(value)
+            _nme = np.mean(item)
+            if config.landmark_type == '2d':
+                left_eye = (record[0] + record[1]) / 2
+                right_eye = (record[2] + record[3]) / 2
+                _dist = np.sqrt(np.sum(np.square(left_eye - right_eye)))
+                #print('eye dist', _dist, left_eye, right_eye)
+                _nme /= _dist
+            else:
+                #_dist = np.sqrt(float(label.shape[2]*label.shape[3]))
+                _dist = np.sqrt(np.sum(np.square(record[5] - record[4])))
+                #print(_dist)
+                _nme /= _dist
+            nme.append(_nme)
+        return np.mean(nme)
+
+    def update(self, labels, preds):
+        self.count += 1
+        label = labels[0].asnumpy()
+        pred_label = preds[-1].asnumpy()
+        nme = self.cal_nme(label, pred_label)
+
+        #print('nme', nme)
+        #nme = np.mean(nme)
+        self.sum_metric += np.mean(nme)
+        self.num_inst += 1.0

insightface/alignment/heatmap/optimizer.py

@@ -0,0 +1,65 @@
+import mxnet as mx
+import mxnet.optimizer as optimizer
+from mxnet.ndarray import (NDArray, zeros, clip, sqrt, cast, maximum, abs as
+                           NDabs)
+#from mxnet.ndarray import (sgd_update, sgd_mom_update, adam_update, rmsprop_update, rmspropalex_update,
+#                      mp_sgd_update, mp_sgd_mom_update, square, ftrl_update)
+
+
+class ONadam(optimizer.Optimizer):
+    def __init__(self,
+                 learning_rate=0.001,
+                 beta1=0.9,
+                 beta2=0.999,
+                 epsilon=1e-8,
+                 schedule_decay=0.004,
+                 **kwargs):
+        super(ONadam, self).__init__(learning_rate=learning_rate, **kwargs)
+        self.beta1 = beta1
+        self.beta2 = beta2
+        self.epsilon = epsilon
+        self.schedule_decay = schedule_decay
+        self.m_schedule = 1.
+
+    def create_state(self, index, weight):
+        return (
+            zeros(weight.shape, weight.context, dtype=weight.dtype),  # mean
+            zeros(weight.shape, weight.context,
+                  dtype=weight.dtype))  # variance
+
+    def update(self, index, weight, grad, state):
+        assert (isinstance(weight, NDArray))
+        assert (isinstance(grad, NDArray))
+        self._update_count(index)
+        lr = self._get_lr(index)
+        wd = self._get_wd(index)
+
+        t = self._index_update_count[index]
+
+        # preprocess grad
+        #grad = grad * self.rescale_grad + wd * weight
+        grad *= self.rescale_grad + wd * weight
+        if self.clip_gradient is not None:
+            grad = clip(grad, -self.clip_gradient, self.clip_gradient)
+
+        # warming momentum schedule
+        momentum_t = self.beta1 * (1. - 0.5 *
+                                   (pow(0.96, t * self.schedule_decay)))
+        momentum_t_1 = self.beta1 * (1. - 0.5 *
+                                     (pow(0.96,
+                                          (t + 1) * self.schedule_decay)))
+        self.m_schedule = self.m_schedule * momentum_t
+        m_schedule_next = self.m_schedule * momentum_t_1
+
+        # update m_t and v_t
+        m_t, v_t = state
+        m_t[:] = self.beta1 * m_t + (1. - self.beta1) * grad
+        v_t[:] = self.beta2 * v_t + (1. - self.beta2) * grad * grad
+
+        grad_prime = grad / (1. - self.m_schedule)
+        m_t_prime = m_t / (1. - m_schedule_next)
+        v_t_prime = v_t / (1. - pow(self.beta2, t))
+        m_t_bar = (1. - momentum_t) * grad_prime + momentum_t_1 * m_t_prime
+
+        # update weight
+        weight[:] -= lr * m_t_bar / (sqrt(v_t_prime) + self.epsilon)

insightface/alignment/heatmap/sample_config.py

@@ -0,0 +1,98 @@
+import numpy as np
+from easydict import EasyDict as edict
+
+config = edict()
+
+#default training/dataset config
+config.num_classes = 68
+config.record_img_size = 384
+config.base_scale = 256
+config.input_img_size = 128
+config.output_label_size = 64
+config.label_xfirst = False
+config.losstype = 'heatmap'
+config.net_coherent = False
+config.multiplier = 1.0
+
+config.gaussian = 0
+
+# network settings
+network = edict()
+
+network.hourglass = edict()
+network.hourglass.net_coherent = False
+network.hourglass.net_sta = 0
+network.hourglass.net_n = 3
+network.hourglass.net_dcn = 0
+network.hourglass.net_stacks = 2
+network.hourglass.net_block = 'resnet'
+network.hourglass.net_binarize = False
+network.hourglass.losstype = 'heatmap'
+
+network.sdu = edict()
+network.sdu.net_coherent = False
+network.sdu.net_sta = 1
+network.sdu.net_n = 3
+network.sdu.net_dcn = 3
+network.sdu.net_stacks = 2
+network.sdu.net_block = 'cab'
+network.sdu.net_binarize = False
+network.sdu.losstype = 'heatmap'
+
+# dataset settings
+dataset = edict()
+
+dataset.i2d = edict()
+dataset.i2d.dataset = '2D'
+dataset.i2d.landmark_type = '2d'
+dataset.i2d.dataset_path = './data_2d'
+dataset.i2d.num_classes = 68
+dataset.i2d.record_img_size = 384
+dataset.i2d.base_scale = 256
+dataset.i2d.input_img_size = 128
+dataset.i2d.output_label_size = 64
+dataset.i2d.label_xfirst = False
+dataset.i2d.val_targets = ['ibug', 'cofw_testset', '300W']
+
+dataset.i3d = edict()
+dataset.i3d.dataset = '3D'
+dataset.i3d.landmark_type = '3d'
+dataset.i3d.dataset_path = './data_3d'
+dataset.i3d.num_classes = 68
+dataset.i3d.record_img_size = 384
+dataset.i3d.base_scale = 256
+dataset.i3d.input_img_size = 128
+dataset.i3d.output_label_size = 64
+dataset.i3d.label_xfirst = False
+dataset.i3d.val_targets = ['AFLW2000-3D']
+
+# default settings
+default = edict()
+
+# default network
+default.network = 'hourglass'
+default.pretrained = ''
+default.pretrained_epoch = 0
+# default dataset
+default.dataset = 'i2d'
+default.frequent = 20
+default.verbose = 200
+default.kvstore = 'device'
+
+default.prefix = 'model/A'
+default.end_epoch = 10000
+default.lr = 0.00025
+default.wd = 0.0
+default.per_batch_size = 20
+default.lr_step = '16000,24000,30000'
+
+
+def generate_config(_network, _dataset):
+    for k, v in network[_network].items():
+        config[k] = v
+        default[k] = v
+    for k, v in dataset[_dataset].items():
+        config[k] = v
+        default[k] = v
+    config.network = _network
+    config.dataset = _dataset

insightface/alignment/heatmap/symbol/sym_heatmap.py

@@ -0,0 +1,1085 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+import mxnet as mx
+import numpy as np
+from config import config
+
+ACT_BIT = 1
+bn_mom = 0.9
+workspace = 256
+memonger = False
+
+
+def Conv(**kwargs):
+    body = mx.sym.Convolution(**kwargs)
+    return body
+
+
+def Act(data, act_type, name):
+    if act_type == 'prelu':
+        body = mx.sym.LeakyReLU(data=data, act_type='prelu', name=name)
+    else:
+        body = mx.symbol.Activation(data=data, act_type=act_type, name=name)
+    return body
+
+
+#def lin(data, num_filter, workspace, name, binarize, dcn):
+#  bit = 1
+#  if not binarize:
+#    if not dcn:
+#        conv1 = Conv(data=data, num_filter=num_filter, kernel=(1,1), stride=(1,1), pad=(0,0),
+#                                      no_bias=True, workspace=workspace, name=name + '_conv')
+#        bn1 = mx.sym.BatchNorm(data=conv1, fix_gamma=False, momentum=bn_mom, eps=2e-5, name=name + '_bn')
+#        act1 = Act(data=bn1, act_type='relu', name=name + '_relu')
+#        return act1
+#    else:
+#        bn1 = mx.sym.BatchNorm(data=data, fix_gamma=False, momentum=bn_mom, eps=2e-5, name=name + '_bn')
+#        act1 = Act(data=bn1, act_type='relu', name=name + '_relu')
+#        conv1_offset = mx.symbol.Convolution(name=name+'_conv_offset', data = act1,
+#                num_filter=18, pad=(1, 1), kernel=(3, 3), stride=(1, 1))
+#        conv1 = mx.contrib.symbol.DeformableConvolution(name=name+"_conv", data=act1, offset=conv1_offset,
+#                num_filter=num_filter, pad=(1,1), kernel=(3, 3), num_deformable_group=1, stride=(1, 1), dilate=(1, 1), no_bias=False)
+#        #conv1 = Conv(data=act1, num_filter=num_filter, kernel=(3,3), stride=(1,1), pad=(1,1),
+#        #                              no_bias=False, workspace=workspace, name=name + '_conv')
+#        return conv1
+#  else:
+#    bn1 = mx.sym.BatchNorm(data=data, fix_gamma=False, momentum=bn_mom, eps=2e-5, name=name + '_bn')
+#    act1 = Act(data=bn1, act_type='relu', name=name + '_relu')
+#    conv1 = mx.sym.QConvolution_v1(data=act1, num_filter=num_filter, kernel=(1,1), stride=(1,1), pad=(0,0),
+#                               no_bias=True, workspace=workspace, name=name + '_conv', act_bit=ACT_BIT, weight_bit=bit)
+#    conv1 = mx.sym.BatchNorm(data=conv1, fix_gamma=False, momentum=bn_mom, eps=2e-5, name=name + '_bn2')
+#    return conv1
+
+
+def lin3(data, num_filter, workspace, name, k, g=1, d=1):
+    if k != 3:
+        conv1 = Conv(data=data,
+                     num_filter=num_filter,
+                     kernel=(k, k),
+                     stride=(1, 1),
+                     pad=((k - 1) // 2, (k - 1) // 2),
+                     num_group=g,
+                     no_bias=True,
+                     workspace=workspace,
+                     name=name + '_conv')
+    else:
+        conv1 = Conv(data=data,
+                     num_filter=num_filter,
+                     kernel=(k, k),
+                     stride=(1, 1),
+                     pad=(d, d),
+                     num_group=g,
+                     dilate=(d, d),
+                     no_bias=True,
+                     workspace=workspace,
+                     name=name + '_conv')
+    bn1 = mx.sym.BatchNorm(data=conv1,
+                           fix_gamma=False,
+                           momentum=bn_mom,
+                           eps=2e-5,
+                           name=name + '_bn')
+    act1 = Act(data=bn1, act_type='relu', name=name + '_relu')
+    ret = act1
+    return ret
+
+
+def ConvFactory(data,
+                num_filter,
+                kernel,
+                stride=(1, 1),
+                pad=(0, 0),
+                act_type="relu",
+                mirror_attr={},
+                with_act=True,
+                dcn=False,
+                name=''):
+    if not dcn:
+        conv = mx.symbol.Convolution(data=data,
+                                     num_filter=num_filter,
+                                     kernel=kernel,
+                                     stride=stride,
+                                     pad=pad,
+                                     no_bias=True,
+                                     workspace=workspace,
+                                     name=name + '_conv')
+    else:
+        conv_offset = mx.symbol.Convolution(name=name + '_conv_offset',
+                                            data=data,
+                                            num_filter=18,
+                                            pad=(1, 1),
+                                            kernel=(3, 3),
+                                            stride=(1, 1))
+        conv = mx.contrib.symbol.DeformableConvolution(name=name + "_conv",
+                                                       data=data,
+                                                       offset=conv_offset,
+                                                       num_filter=num_filter,
+                                                       pad=(1, 1),
+                                                       kernel=(3, 3),
+                                                       num_deformable_group=1,
+                                                       stride=stride,
+                                                       dilate=(1, 1),
+                                                       no_bias=False)
+    bn = mx.symbol.BatchNorm(data=conv,
+                             fix_gamma=False,
+                             momentum=bn_mom,
+                             eps=2e-5,
+                             name=name + '_bn')
+    if with_act:
+        act = Act(bn, act_type, name=name + '_relu')
+        #act = mx.symbol.Activation(
+        #    data=bn, act_type=act_type, attr=mirror_attr, name=name+'_relu')
+        return act
+    else:
+        return bn
+
+
+class CAB:
+    def __init__(self, data, nFilters, nModules, n, workspace, name, dilate,
+                 group):
+        self.data = data
+        self.nFilters = nFilters
+        self.nModules = nModules
+        self.n = n
+        self.workspace = workspace
+        self.name = name
+        self.dilate = dilate
+        self.group = group
+        self.sym_map = {}
+
+    def get_output(self, w, h):
+        key = (w, h)
+        if key in self.sym_map:
+            return self.sym_map[key]
+        ret = None
+        if h == self.n:
+            if w == self.n:
+                ret = (self.data, self.nFilters)
+            else:
+                x = self.get_output(w + 1, h)
+                f = int(x[1] * 0.5)
+                if w != self.n - 1:
+                    body = lin3(x[0], f, self.workspace,
+                                "%s_w%d_h%d_1" % (self.name, w, h), 3,
+                                self.group, 1)
+                else:
+                    body = lin3(x[0], f, self.workspace,
+                                "%s_w%d_h%d_1" % (self.name, w, h), 3,
+                                self.group, self.dilate)
+                ret = (body, f)
+        else:
+            x = self.get_output(w + 1, h + 1)
+            y = self.get_output(w, h + 1)
+            if h % 2 == 1 and h != w:
+                xbody = lin3(x[0], x[1], self.workspace,
+                             "%s_w%d_h%d_2" % (self.name, w, h), 3, x[1])
+                #xbody = xbody+x[0]
+            else:
+                xbody = x[0]
+            #xbody = x[0]
+            #xbody = lin3(x[0], x[1], self.workspace, "%s_w%d_h%d_2"%(self.name, w, h), 3, x[1])
+            if w == 0:
+                ybody = lin3(y[0], y[1], self.workspace,
+                             "%s_w%d_h%d_3" % (self.name, w, h), 3, self.group)
+            else:
+                ybody = y[0]
+            ybody = mx.sym.concat(y[0], ybody, dim=1)
+            body = mx.sym.add_n(xbody,
+                                ybody,
+                                name="%s_w%d_h%d_add" % (self.name, w, h))
+            body = body / 2
+            ret = (body, x[1])
+        self.sym_map[key] = ret
+        return ret
+
+    def get(self):
+        return self.get_output(1, 1)[0]
+
+
+def conv_resnet(data, num_filter, stride, dim_match, name, binarize, dcn,
+                dilate, **kwargs):
+    bit = 1
+    #print('in unit2')
+    # the same as https://github.com/facebook/fb.resnet.torch#notes, a bit difference with origin paper
+    bn1 = mx.sym.BatchNorm(data=data,
+                           fix_gamma=False,
+                           eps=2e-5,
+                           momentum=bn_mom,
+                           name=name + '_bn1')
+    if not binarize:
+        act1 = Act(data=bn1, act_type='relu', name=name + '_relu1')
+        conv1 = Conv(data=act1,
+                     num_filter=int(num_filter * 0.5),
+                     kernel=(1, 1),
+                     stride=(1, 1),
+                     pad=(0, 0),
+                     no_bias=True,
+                     workspace=workspace,
+                     name=name + '_conv1')
+    else:
+        act1 = mx.sym.QActivation(data=bn1,
+                                  act_bit=ACT_BIT,
+                                  name=name + '_relu1',
+                                  backward_only=True)
+        conv1 = mx.sym.QConvolution(data=act1,
+                                    num_filter=int(num_filter * 0.5),
+                                    kernel=(1, 1),
+                                    stride=(1, 1),
+                                    pad=(0, 0),
+                                    no_bias=True,
+                                    workspace=workspace,
+                                    name=name + '_conv1',
+                                    act_bit=ACT_BIT,
+                                    weight_bit=bit)
+    bn2 = mx.sym.BatchNorm(data=conv1,
+                           fix_gamma=False,
+                           eps=2e-5,
+                           momentum=bn_mom,
+                           name=name + '_bn2')
+    if not binarize:
+        act2 = Act(data=bn2, act_type='relu', name=name + '_relu2')
+        conv2 = Conv(data=act2,
+                     num_filter=int(num_filter * 0.5),
+                     kernel=(3, 3),
+                     stride=(1, 1),
+                     pad=(1, 1),
+                     no_bias=True,
+                     workspace=workspace,
+                     name=name + '_conv2')
+    else:
+        act2 = mx.sym.QActivation(data=bn2,
+                                  act_bit=ACT_BIT,
+                                  name=name + '_relu2',
+                                  backward_only=True)
+        conv2 = mx.sym.QConvolution(data=act2,
+                                    num_filter=int(num_filter * 0.5),
+                                    kernel=(3, 3),
+                                    stride=(1, 1),
+                                    pad=(1, 1),
+                                    no_bias=True,
+                                    workspace=workspace,
+                                    name=name + '_conv2',
+                                    act_bit=ACT_BIT,
+                                    weight_bit=bit)
+    bn3 = mx.sym.BatchNorm(data=conv2,
+                           fix_gamma=False,
+                           eps=2e-5,
+                           momentum=bn_mom,
+                           name=name + '_bn3')
+    if not binarize:
+        act3 = Act(data=bn3, act_type='relu', name=name + '_relu3')
+        conv3 = Conv(data=act3,
+                     num_filter=num_filter,
+                     kernel=(1, 1),
+                     stride=(1, 1),
+                     pad=(0, 0),
+                     no_bias=True,
+                     workspace=workspace,
+                     name=name + '_conv3')
+    else:
+        act3 = mx.sym.QActivation(data=bn3,
+                                  act_bit=ACT_BIT,
+                                  name=name + '_relu3',
+                                  backward_only=True)
+        conv3 = mx.sym.QConvolution(data=act3,
+                                    num_filter=num_filter,
+                                    kernel=(1, 1),
+                                    stride=(1, 1),
+                                    pad=(0, 0),
+                                    no_bias=True,
+                                    workspace=workspace,
+                                    name=name + '_conv3',
+                                    act_bit=ACT_BIT,
+                                    weight_bit=bit)
+    #if binarize:
+    #  conv3 = mx.sym.BatchNorm(data=conv3, fix_gamma=False, eps=2e-5, momentum=bn_mom, name=name + '_bn4')
+    if dim_match:
+        shortcut = data
+    else:
+        if not binarize:
+            shortcut = Conv(data=act1,
+                            num_filter=num_filter,
+                            kernel=(1, 1),
+                            stride=stride,
+                            no_bias=True,
+                            workspace=workspace,
+                            name=name + '_sc')
+        else:
+            shortcut = mx.sym.QConvolution(data=act1,
+                                           num_filter=num_filter,
+                                           kernel=(1, 1),
+                                           stride=stride,
+                                           pad=(0, 0),
+                                           no_bias=True,
+                                           workspace=workspace,
+                                           name=name + '_sc',
+                                           act_bit=ACT_BIT,
+                                           weight_bit=bit)
+    if memonger:
+        shortcut._set_attr(mirror_stage='True')
+    return conv3 + shortcut
+
+
+def conv_hpm(data, num_filter, stride, dim_match, name, binarize, dcn,
+             dilation, **kwargs):
+    bit = 1
+    #print('in unit2')
+    # the same as https://github.com/facebook/fb.resnet.torch#notes, a bit difference with origin paper
+    bn1 = mx.sym.BatchNorm(data=data,
+                           fix_gamma=False,
+                           eps=2e-5,
+                           momentum=bn_mom,
+                           name=name + '_bn1')
+    if not binarize:
+        act1 = Act(data=bn1, act_type='relu', name=name + '_relu1')
+        if not dcn:
+            conv1 = Conv(data=act1,
+                         num_filter=int(num_filter * 0.5),
+                         kernel=(3, 3),
+                         stride=(1, 1),
+                         pad=(dilation, dilation),
+                         dilate=(dilation, dilation),
+                         no_bias=True,
+                         workspace=workspace,
+                         name=name + '_conv1')
+        else:
+            conv1_offset = mx.symbol.Convolution(name=name + '_conv1_offset',
+                                                 data=act1,
+                                                 num_filter=18,
+                                                 pad=(1, 1),
+                                                 kernel=(3, 3),
+                                                 stride=(1, 1))
+            conv1 = mx.contrib.symbol.DeformableConvolution(
+                name=name + '_conv1',
+                data=act1,
+                offset=conv1_offset,
+                num_filter=int(num_filter * 0.5),
+                pad=(1, 1),
+                kernel=(3, 3),
+                num_deformable_group=1,
+                stride=(1, 1),
+                dilate=(1, 1),
+                no_bias=True)
+    else:
+        act1 = mx.sym.QActivation(data=bn1,
+                                  act_bit=ACT_BIT,
+                                  name=name + '_relu1',
+                                  backward_only=True)
+        conv1 = mx.sym.QConvolution_v1(data=act1,
+                                       num_filter=int(num_filter * 0.5),
+                                       kernel=(3, 3),
+                                       stride=(1, 1),
+                                       pad=(1, 1),
+                                       no_bias=True,
+                                       workspace=workspace,
+                                       name=name + '_conv1',
+                                       act_bit=ACT_BIT,
+                                       weight_bit=bit)
+    bn2 = mx.sym.BatchNorm(data=conv1,
+                           fix_gamma=False,
+                           eps=2e-5,
+                           momentum=bn_mom,
+                           name=name + '_bn2')
+    if not binarize:
+        act2 = Act(data=bn2, act_type='relu', name=name + '_relu2')
+        if not dcn:
+            conv2 = Conv(data=act2,
+                         num_filter=int(num_filter * 0.25),
+                         kernel=(3, 3),
+                         stride=(1, 1),
+                         pad=(dilation, dilation),
+                         dilate=(dilation, dilation),
+                         no_bias=True,
+                         workspace=workspace,
+                         name=name + '_conv2')
+        else:
+            conv2_offset = mx.symbol.Convolution(name=name + '_conv2_offset',
+                                                 data=act2,
+                                                 num_filter=18,
+                                                 pad=(1, 1),
+                                                 kernel=(3, 3),
+                                                 stride=(1, 1))
+            conv2 = mx.contrib.symbol.DeformableConvolution(
+                name=name + '_conv2',
+                data=act2,
+                offset=conv2_offset,
+                num_filter=int(num_filter * 0.25),
+                pad=(1, 1),
+                kernel=(3, 3),
+                num_deformable_group=1,
+                stride=(1, 1),
+                dilate=(1, 1),
+                no_bias=True)
+    else:
+        act2 = mx.sym.QActivation(data=bn2,
+                                  act_bit=ACT_BIT,
+                                  name=name + '_relu2',
+                                  backward_only=True)
+        conv2 = mx.sym.QConvolution_v1(data=act2,
+                                       num_filter=int(num_filter * 0.25),
+                                       kernel=(3, 3),
+                                       stride=(1, 1),
+                                       pad=(1, 1),
+                                       no_bias=True,
+                                       workspace=workspace,
+                                       name=name + '_conv2',
+                                       act_bit=ACT_BIT,
+                                       weight_bit=bit)
+    bn3 = mx.sym.BatchNorm(data=conv2,
+                           fix_gamma=False,
+                           eps=2e-5,
+                           momentum=bn_mom,
+                           name=name + '_bn3')
+    if not binarize:
+        act3 = Act(data=bn3, act_type='relu', name=name + '_relu3')
+        if not dcn:
+            conv3 = Conv(data=act3,
+                         num_filter=int(num_filter * 0.25),
+                         kernel=(3, 3),
+                         stride=(1, 1),
+                         pad=(dilation, dilation),
+                         dilate=(dilation, dilation),
+                         no_bias=True,
+                         workspace=workspace,
+                         name=name + '_conv3')
+        else:
+            conv3_offset = mx.symbol.Convolution(name=name + '_conv3_offset',
+                                                 data=act3,
+                                                 num_filter=18,
+                                                 pad=(1, 1),
+                                                 kernel=(3, 3),
+                                                 stride=(1, 1))
+            conv3 = mx.contrib.symbol.DeformableConvolution(
+                name=name + '_conv3',
+                data=act3,
+                offset=conv3_offset,
+                num_filter=int(num_filter * 0.25),
+                pad=(1, 1),
+                kernel=(3, 3),
+                num_deformable_group=1,
+                stride=(1, 1),
+                dilate=(1, 1),
+                no_bias=True)
+    else:
+        act3 = mx.sym.QActivation(data=bn3,
+                                  act_bit=ACT_BIT,
+                                  name=name + '_relu3',
+                                  backward_only=True)
+        conv3 = mx.sym.QConvolution_v1(data=act3,
+                                       num_filter=int(num_filter * 0.25),
+                                       kernel=(3, 3),
+                                       stride=(1, 1),
+                                       pad=(1, 1),
+                                       no_bias=True,
+                                       workspace=workspace,
+                                       name=name + '_conv3',
+                                       act_bit=ACT_BIT,
+                                       weight_bit=bit)
+    conv4 = mx.symbol.Concat(*[conv1, conv2, conv3])
+    if binarize:
+        conv4 = mx.sym.BatchNorm(data=conv4,
+                                 fix_gamma=False,
+                                 eps=2e-5,
+                                 momentum=bn_mom,
+                                 name=name + '_bn4')
+    if dim_match:
+        shortcut = data
+    else:
+        if not binarize:
+            shortcut = Conv(data=act1,
+                            num_filter=num_filter,
+                            kernel=(1, 1),
+                            stride=stride,
+                            no_bias=True,
+                            workspace=workspace,
+                            name=name + '_sc')
+        else:
+            #assert(False)
+            shortcut = mx.sym.QConvolution_v1(data=act1,
+                                              num_filter=num_filter,
+                                              kernel=(1, 1),
+                                              stride=stride,
+                                              pad=(0, 0),
+                                              no_bias=True,
+                                              workspace=workspace,
+                                              name=name + '_sc',
+                                              act_bit=ACT_BIT,
+                                              weight_bit=bit)
+            shortcut = mx.sym.BatchNorm(data=shortcut,
+                                        fix_gamma=False,
+                                        eps=2e-5,
+                                        momentum=bn_mom,
+                                        name=name + '_sc_bn')
+    if memonger:
+        shortcut._set_attr(mirror_stage='True')
+    return conv4 + shortcut
+    #return bn4 + shortcut
+    #return act4 + shortcut
+
+
+def block17(net,
+            input_num_channels,
+            scale=1.0,
+            with_act=True,
+            act_type='relu',
+            mirror_attr={},
+            name=''):
+    tower_conv = ConvFactory(net, 192, (1, 1), name=name + '_conv')
+    tower_conv1_0 = ConvFactory(net, 129, (1, 1), name=name + '_conv1_0')
+    tower_conv1_1 = ConvFactory(tower_conv1_0,
+                                160, (1, 7),
+                                pad=(1, 2),
+                                name=name + '_conv1_1')
+    tower_conv1_2 = ConvFactory(tower_conv1_1,
+                                192, (7, 1),
+                                pad=(2, 1),
+                                name=name + '_conv1_2')
+    tower_mixed = mx.symbol.Concat(*[tower_conv, tower_conv1_2])
+    tower_out = ConvFactory(tower_mixed,
+                            input_num_channels, (1, 1),
+                            with_act=False,
+                            name=name + '_conv_out')
+    net = net + scale * tower_out
+    if with_act:
+        act = mx.symbol.Activation(data=net,
+                                   act_type=act_type,
+                                   attr=mirror_attr)
+        return act
+    else:
+        return net
+
+
+def block35(net,
+            input_num_channels,
+            scale=1.0,
+            with_act=True,
+            act_type='relu',
+            mirror_attr={},
+            name=''):
+    M = 1.0
+    tower_conv = ConvFactory(net,
+                             int(input_num_channels * 0.25 * M), (1, 1),
+                             name=name + '_conv')
+    tower_conv1_0 = ConvFactory(net,
+                                int(input_num_channels * 0.25 * M), (1, 1),
+                                name=name + '_conv1_0')
+    tower_conv1_1 = ConvFactory(tower_conv1_0,
+                                int(input_num_channels * 0.25 * M), (3, 3),
+                                pad=(1, 1),
+                                name=name + '_conv1_1')
+    tower_conv2_0 = ConvFactory(net,
+                                int(input_num_channels * 0.25 * M), (1, 1),
+                                name=name + '_conv2_0')
+    tower_conv2_1 = ConvFactory(tower_conv2_0,
+                                int(input_num_channels * 0.375 * M), (3, 3),
+                                pad=(1, 1),
+                                name=name + '_conv2_1')
+    tower_conv2_2 = ConvFactory(tower_conv2_1,
+                                int(input_num_channels * 0.5 * M), (3, 3),
+                                pad=(1, 1),
+                                name=name + '_conv2_2')
+    tower_mixed = mx.symbol.Concat(*[tower_conv, tower_conv1_1, tower_conv2_2])
+    tower_out = ConvFactory(tower_mixed,
+                            input_num_channels, (1, 1),
+                            with_act=False,
+                            name=name + '_conv_out')
+
+    net = net + scale * tower_out
+    if with_act:
+        act = mx.symbol.Activation(data=net,
+                                   act_type=act_type,
+                                   attr=mirror_attr)
+        return act
+    else:
+        return net
+
+
+def conv_inception(data, num_filter, stride, dim_match, name, binarize, dcn,
+                   dilate, **kwargs):
+    assert not binarize
+    if stride[0] > 1 or not dim_match:
+        return conv_resnet(data, num_filter, stride, dim_match, name, binarize,
+                           dcn, dilate, **kwargs)
+    conv4 = block35(data, num_filter, name=name + '_block35')
+    return conv4
+
+
+def conv_cab(data, num_filter, stride, dim_match, name, binarize, dcn, dilate,
+             **kwargs):
+    if stride[0] > 1 or not dim_match:
+        return conv_hpm(data, num_filter, stride, dim_match, name, binarize,
+                        dcn, dilate, **kwargs)
+    cab = CAB(data, num_filter, 1, 4, workspace, name, dilate, 1)
+    return cab.get()
+
+
+def conv_block(data, num_filter, stride, dim_match, name, binarize, dcn,
+               dilate):
+    if config.net_block == 'resnet':
+        return conv_resnet(data, num_filter, stride, dim_match, name, binarize,
+                           dcn, dilate)
+    elif config.net_block == 'inception':
+        return conv_inception(data, num_filter, stride, dim_match, name,
+                              binarize, dcn, dilate)
+    elif config.net_block == 'hpm':
+        return conv_hpm(data, num_filter, stride, dim_match, name, binarize,
+                        dcn, dilate)
+    elif config.net_block == 'cab':
+        return conv_cab(data, num_filter, stride, dim_match, name, binarize,
+                        dcn, dilate)
+
+
+def hourglass(data, nFilters, nModules, n, workspace, name, binarize, dcn):
+    s = 2
+    _dcn = False
+    up1 = data
+    for i in range(nModules):
+        up1 = conv_block(up1, nFilters, (1, 1), True, "%s_up1_%d" % (name, i),
+                         binarize, _dcn, 1)
+    low1 = mx.sym.Pooling(data=data,
+                          kernel=(s, s),
+                          stride=(s, s),
+                          pad=(0, 0),
+                          pool_type='max')
+    for i in range(nModules):
+        low1 = conv_block(low1, nFilters, (1, 1), True,
+                          "%s_low1_%d" % (name, i), binarize, _dcn, 1)
+    if n > 1:
+        low2 = hourglass(low1, nFilters, nModules, n - 1, workspace,
+                         "%s_%d" % (name, n - 1), binarize, dcn)
+    else:
+        low2 = low1
+        for i in range(nModules):
+            low2 = conv_block(low2, nFilters, (1, 1), True,
+                              "%s_low2_%d" % (name, i), binarize, _dcn,
+                              1)  #TODO
+    low3 = low2
+    for i in range(nModules):
+        low3 = conv_block(low3, nFilters, (1, 1), True,
+                          "%s_low3_%d" % (name, i), binarize, _dcn, 1)
+    up2 = mx.symbol.UpSampling(low3,
+                               scale=s,
+                               sample_type='nearest',
+                               workspace=512,
+                               name='%s_upsampling_%s' % (name, n),
+                               num_args=1)
+    return mx.symbol.add_n(up1, up2)
+
+
+class STA:
+    def __init__(self, data, nFilters, nModules, n, workspace, name):
+        self.data = data
+        self.nFilters = nFilters
+        self.nModules = nModules
+        self.n = n
+        self.workspace = workspace
+        self.name = name
+        self.sym_map = {}
+
+    def get_conv(self, data, name, dilate=1, group=1):
+        cab = CAB(data, self.nFilters, self.nModules, 4, self.workspace, name,
+                  dilate, group)
+        return cab.get()
+
+    def get_output(self, w, h):
+        #print(w,h)
+        assert w >= 1 and w <= config.net_n + 1
+        assert h >= 1 and h <= config.net_n + 1
+        s = 2
+        bn_mom = 0.9
+        key = (w, h)
+        if key in self.sym_map:
+            return self.sym_map[key]
+        ret = None
+        if h == self.n:
+            if w == self.n:
+                ret = self.data, 64
+            else:
+                x = self.get_output(w + 1, h)
+                body = self.get_conv(x[0], "%s_w%d_h%d_1" % (self.name, w, h))
+                body = mx.sym.Pooling(data=body,
+                                      kernel=(s, s),
+                                      stride=(s, s),
+                                      pad=(0, 0),
+                                      pool_type='max')
+                body = self.get_conv(body, "%s_w%d_h%d_2" % (self.name, w, h))
+                ret = body, x[1] // 2
+        else:
+            x = self.get_output(w + 1, h + 1)
+            y = self.get_output(w, h + 1)
+
+            HC = False
+
+            if h % 2 == 1 and h != w:
+                xbody = lin3(x[0], self.nFilters, self.workspace,
+                             "%s_w%d_h%d_x" % (self.name, w, h), 3,
+                             self.nFilters, 1)
+                HC = True
+                #xbody = x[0]
+            else:
+                xbody = x[0]
+            if x[1] // y[1] == 2:
+                if w > 1:
+                    ybody = mx.symbol.Deconvolution(
+                        data=y[0],
+                        num_filter=self.nFilters,
+                        kernel=(s, s),
+                        stride=(s, s),
+                        name='%s_upsampling_w%d_h%d' % (self.name, w, h),
+                        attr={'lr_mult': '1.0'},
+                        workspace=self.workspace)
+                    ybody = mx.sym.BatchNorm(data=ybody,
+                                             fix_gamma=False,
+                                             momentum=bn_mom,
+                                             eps=2e-5,
+                                             name="%s_w%d_h%d_y_bn" %
+                                             (self.name, w, h))
+                    ybody = Act(data=ybody,
+                                act_type='relu',
+                                name="%s_w%d_h%d_y_act" % (self.name, w, h))
+                else:
+                    if h >= 1:
+                        ybody = mx.symbol.UpSampling(
+                            y[0],
+                            scale=s,
+                            sample_type='nearest',
+                            workspace=512,
+                            name='%s_upsampling_w%d_h%d' % (self.name, w, h),
+                            num_args=1)
+                        ybody = self.get_conv(
+                            ybody, "%s_w%d_h%d_4" % (self.name, w, h))
+                    else:
+                        ybody = mx.symbol.Deconvolution(
+                            data=y[0],
+                            num_filter=self.nFilters,
+                            kernel=(s, s),
+                            stride=(s, s),
+                            name='%s_upsampling_w%d_h%d' % (self.name, w, h),
+                            attr={'lr_mult': '1.0'},
+                            workspace=self.workspace)
+                        ybody = mx.sym.BatchNorm(data=ybody,
+                                                 fix_gamma=False,
+                                                 momentum=bn_mom,
+                                                 eps=2e-5,
+                                                 name="%s_w%d_h%d_y_bn" %
+                                                 (self.name, w, h))
+                        ybody = Act(data=ybody,
+                                    act_type='relu',
+                                    name="%s_w%d_h%d_y_act" %
+                                    (self.name, w, h))
+                        ybody = Conv(data=ybody,
+                                     num_filter=self.nFilters,
+                                     kernel=(3, 3),
+                                     stride=(1, 1),
+                                     pad=(1, 1),
+                                     no_bias=True,
+                                     name="%s_w%d_h%d_y_conv2" %
+                                     (self.name, w, h),
+                                     workspace=self.workspace)
+                        ybody = mx.sym.BatchNorm(data=ybody,
+                                                 fix_gamma=False,
+                                                 momentum=bn_mom,
+                                                 eps=2e-5,
+                                                 name="%s_w%d_h%d_y_bn2" %
+                                                 (self.name, w, h))
+                        ybody = Act(data=ybody,
+                                    act_type='relu',
+                                    name="%s_w%d_h%d_y_act2" %
+                                    (self.name, w, h))
+            else:
+                ybody = self.get_conv(y[0], "%s_w%d_h%d_5" % (self.name, w, h))
+            #if not HC:
+            if config.net_sta == 2 and h == 3 and w == 2:
+                z = self.get_output(w + 1, h)
+                zbody = z[0]
+                zbody = mx.sym.Pooling(data=zbody,
+                                       kernel=(z[1], z[1]),
+                                       stride=(z[1], z[1]),
+                                       pad=(0, 0),
+                                       pool_type='avg')
+                body = xbody + ybody
+                body = body / 2
+                body = mx.sym.broadcast_mul(body, zbody)
+            else:  #sta==1
+                body = xbody + ybody
+                body = body / 2
+            ret = body, x[1]
+
+        assert ret is not None
+        self.sym_map[key] = ret
+        return ret
+
+    def get(self):
+        return self.get_output(1, 1)[0]
+
+
+class SymCoherent:
+    def __init__(self, per_batch_size):
+        self.per_batch_size = per_batch_size
+        self.flip_order = [
+            16, 15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0, 26, 25,
+            24, 23, 22, 21, 20, 19, 18, 17, 27, 28, 29, 30, 35, 34, 33, 32, 31,
+            45, 44, 43, 42, 47, 46, 39, 38, 37, 36, 41, 40, 54, 53, 52, 51, 50,
+            49, 48, 59, 58, 57, 56, 55, 64, 63, 62, 61, 60, 67, 66, 65
+        ]
+
+    def get(self, data):
+        #data.shape[0]==per_batch_size
+        b = self.per_batch_size // 2
+        ux = mx.sym.slice_axis(data, axis=0, begin=0, end=b)
+        dx = mx.sym.slice_axis(data, axis=0, begin=b, end=b * 2)
+        ux = mx.sym.flip(ux, axis=3)
+        #ux = mx.sym.take(ux, indices = self.flip_order, axis=0)
+        ux_list = []
+        for o in self.flip_order:
+            _ux = mx.sym.slice_axis(ux, axis=1, begin=o, end=o + 1)
+            ux_list.append(_ux)
+        ux = mx.sym.concat(*ux_list, dim=1)
+        return ux, dx
+
+
+def l2_loss(x, y):
+    loss = x - y
+    loss = mx.symbol.smooth_l1(loss, scalar=1.0)
+    #loss = loss*loss
+    loss = mx.symbol.mean(loss)
+    return loss
+
+
+def ce_loss(x, y):
+    #loss = mx.sym.SoftmaxOutput(data = x, label = y, normalization='valid', multi_output=True)
+    x_max = mx.sym.max(x, axis=[2, 3], keepdims=True)
+    x = mx.sym.broadcast_minus(x, x_max)
+    body = mx.sym.exp(x)
+    sums = mx.sym.sum(body, axis=[2, 3], keepdims=True)
+    body = mx.sym.broadcast_div(body, sums)
+    loss = mx.sym.log(body)
+    loss = loss * y * -1.0
+    loss = mx.symbol.mean(loss, axis=[1, 2, 3])
+    #loss = mx.symbol.mean(loss)
+    return loss
+
+
+def get_symbol(num_classes):
+    m = config.multiplier
+    sFilters = max(int(64 * m), 32)
+    mFilters = max(int(128 * m), 32)
+    nFilters = int(256 * m)
+
+    nModules = 1
+    nStacks = config.net_stacks
+    binarize = config.net_binarize
+    input_size = config.input_img_size
+    label_size = config.output_label_size
+    use_coherent = config.net_coherent
+    use_STA = config.net_sta
+    N = config.net_n
+    DCN = config.net_dcn
+    per_batch_size = config.per_batch_size
+    print('binarize', binarize)
+    print('use_coherent', use_coherent)
+    print('use_STA', use_STA)
+    print('use_N', N)
+    print('use_DCN', DCN)
+    print('per_batch_size', per_batch_size)
+    #assert(label_size==64 or label_size==32)
+    #assert(input_size==128 or input_size==256)
+    coherentor = SymCoherent(per_batch_size)
+    D = input_size // label_size
+    print(input_size, label_size, D)
+    data = mx.sym.Variable(name='data')
+    data = data - 127.5
+    data = data * 0.0078125
+    gt_label = mx.symbol.Variable(name='softmax_label')
+    losses = []
+    closses = []
+    ref_label = gt_label
+    if D == 4:
+        body = Conv(data=data,
+                    num_filter=sFilters,
+                    kernel=(7, 7),
+                    stride=(2, 2),
+                    pad=(3, 3),
+                    no_bias=True,
+                    name="conv0",
+                    workspace=workspace)
+    else:
+        body = Conv(data=data,
+                    num_filter=sFilters,
+                    kernel=(3, 3),
+                    stride=(1, 1),
+                    pad=(1, 1),
+                    no_bias=True,
+                    name="conv0",
+                    workspace=workspace)
+    body = mx.sym.BatchNorm(data=body,
+                            fix_gamma=False,
+                            eps=2e-5,
+                            momentum=bn_mom,
+                            name='bn0')
+    body = Act(data=body, act_type='relu', name='relu0')
+
+    dcn = False
+    body = conv_block(body, mFilters, (1, 1), sFilters == mFilters, 'res0',
+                      False, dcn, 1)
+
+    body = mx.sym.Pooling(data=body,
+                          kernel=(2, 2),
+                          stride=(2, 2),
+                          pad=(0, 0),
+                          pool_type='max')
+
+    body = conv_block(body, mFilters, (1, 1), True, 'res1', False, dcn,
+                      1)  #TODO
+    body = conv_block(body, nFilters, (1, 1), mFilters == nFilters, 'res2',
+                      binarize, dcn, 1)  #binarize=True?
+
+    heatmap = None
+
+    for i in range(nStacks):
+        shortcut = body
+        if config.net_sta > 0:
+            sta = STA(body, nFilters, nModules, config.net_n + 1, workspace,
+                      'sta%d' % (i))
+            body = sta.get()
+        else:
+            body = hourglass(body, nFilters, nModules, config.net_n, workspace,
+                             'stack%d_hg' % (i), binarize, dcn)
+        for j in range(nModules):
+            body = conv_block(body, nFilters, (1, 1), True,
+                              'stack%d_unit%d' % (i, j), binarize, dcn, 1)
+        _dcn = True if config.net_dcn >= 2 else False
+        ll = ConvFactory(body,
+                         nFilters, (1, 1),
+                         dcn=_dcn,
+                         name='stack%d_ll' % (i))
+        _name = "heatmap%d" % (i) if i < nStacks - 1 else "heatmap"
+        _dcn = True if config.net_dcn >= 2 else False
+        if not _dcn:
+            out = Conv(data=ll,
+                       num_filter=num_classes,
+                       kernel=(1, 1),
+                       stride=(1, 1),
+                       pad=(0, 0),
+                       name=_name,
+                       workspace=workspace)
+        else:
+            out_offset = mx.symbol.Convolution(name=_name + '_offset',
+                                               data=ll,
+                                               num_filter=18,
+                                               pad=(1, 1),
+                                               kernel=(3, 3),
+                                               stride=(1, 1))
+            out = mx.contrib.symbol.DeformableConvolution(
+                name=_name,
+                data=ll,
+                offset=out_offset,
+                num_filter=num_classes,
+                pad=(1, 1),
+                kernel=(3, 3),
+                num_deformable_group=1,
+                stride=(1, 1),
+                dilate=(1, 1),
+                no_bias=False)
+            #out = Conv(data=ll, num_filter=num_classes, kernel=(3,3), stride=(1,1), pad=(1,1),
+            #                          name=_name, workspace=workspace)
+        if i == nStacks - 1:
+            heatmap = out
+        loss = ce_loss(out, ref_label)
+        #loss = loss/nStacks
+        #loss = l2_loss(out, ref_label)
+        losses.append(loss)
+        if config.net_coherent > 0:
+            ux, dx = coherentor.get(out)
+            closs = l2_loss(ux, dx)
+            closs = closs / nStacks
+            closses.append(closs)
+
+        if i < nStacks - 1:
+            ll2 = Conv(data=ll,
+                       num_filter=nFilters,
+                       kernel=(1, 1),
+                       stride=(1, 1),
+                       pad=(0, 0),
+                       name="stack%d_ll2" % (i),
+                       workspace=workspace)
+            out2 = Conv(data=out,
+                        num_filter=nFilters,
+                        kernel=(1, 1),
+                        stride=(1, 1),
+                        pad=(0, 0),
+                        name="stack%d_out2" % (i),
+                        workspace=workspace)
+            body = mx.symbol.add_n(shortcut, ll2, out2)
+            _dcn = True if (config.net_dcn == 1
+                            or config.net_dcn == 3) else False
+            if _dcn:
+                _name = "stack%d_out3" % (i)
+                out3_offset = mx.symbol.Convolution(name=_name + '_offset',
+                                                    data=body,
+                                                    num_filter=18,
+                                                    pad=(1, 1),
+                                                    kernel=(3, 3),
+                                                    stride=(1, 1))
+                out3 = mx.contrib.symbol.DeformableConvolution(
+                    name=_name,
+                    data=body,
+                    offset=out3_offset,
+                    num_filter=nFilters,
+                    pad=(1, 1),
+                    kernel=(3, 3),
+                    num_deformable_group=1,
+                    stride=(1, 1),
+                    dilate=(1, 1),
+                    no_bias=False)
+                body = out3
+
+    pred = mx.symbol.BlockGrad(heatmap)
+    #loss = mx.symbol.add_n(*losses)
+    #loss = mx.symbol.MakeLoss(loss)
+    #syms = [loss]
+    syms = []
+    for loss in losses:
+        loss = mx.symbol.MakeLoss(loss)
+        syms.append(loss)
+    if len(closses) > 0:
+        coherent_weight = 0.0001
+        closs = mx.symbol.add_n(*closses)
+        closs = mx.symbol.MakeLoss(closs, grad_scale=coherent_weight)
+        syms.append(closs)
+    syms.append(pred)
+    sym = mx.symbol.Group(syms)
+    return sym
+
+
+def init_weights(sym, data_shape_dict):
+    #print('in hg')
+    arg_name = sym.list_arguments()
+    aux_name = sym.list_auxiliary_states()
+    arg_shape, _, aux_shape = sym.infer_shape(**data_shape_dict)
+    arg_shape_dict = dict(zip(arg_name, arg_shape))
+    aux_shape_dict = dict(zip(aux_name, aux_shape))
+    #print(aux_shape)
+    #print(aux_params)
+    #print(arg_shape_dict)
+    arg_params = {}
+    aux_params = {}
+    for k in arg_shape_dict:
+        v = arg_shape_dict[k]
+        #print(k,v)
+        if k.endswith('offset_weight') or k.endswith('offset_bias'):
+            print('initializing', k)
+            arg_params[k] = mx.nd.zeros(shape=v)
+        elif k.startswith('fc6_'):
+            if k.endswith('_weight'):
+                print('initializing', k)
+                arg_params[k] = mx.random.normal(0, 0.01, shape=v)
+            elif k.endswith('_bias'):
+                print('initializing', k)
+                arg_params[k] = mx.nd.zeros(shape=v)
+        elif k.find('upsampling') >= 0:
+            print('initializing upsampling_weight', k)
+            arg_params[k] = mx.nd.zeros(shape=arg_shape_dict[k])
+            init = mx.init.Initializer()
+            init._init_bilinear(k, arg_params[k])
+    return arg_params, aux_params

insightface/alignment/heatmap/test.py

@@ -0,0 +1,100 @@
+import argparse
+import cv2
+import sys
+import numpy as np
+import os
+import mxnet as mx
+import datetime
+import img_helper
+sys.path.append(os.path.join(os.path.dirname(__file__), '..', 'deploy'))
+from mtcnn_detector import MtcnnDetector
+
+
+class Handler:
+    def __init__(self, prefix, epoch, ctx_id=0):
+        print('loading', prefix, epoch)
+        if ctx_id >= 0:
+            ctx = mx.gpu(ctx_id)
+        else:
+            ctx = mx.cpu()
+        sym, arg_params, aux_params = mx.model.load_checkpoint(prefix, epoch)
+        all_layers = sym.get_internals()
+        sym = all_layers['heatmap_output']
+        image_size = (128, 128)
+        self.image_size = image_size
+        model = mx.mod.Module(symbol=sym, context=ctx, label_names=None)
+        #model = mx.mod.Module(symbol=sym, context=ctx)
+        model.bind(for_training=False,
+                   data_shapes=[('data', (1, 3, image_size[0], image_size[1]))
+                                ])
+        model.set_params(arg_params, aux_params)
+        self.model = model
+        mtcnn_path = os.path.join(os.path.dirname(__file__), '..', 'deploy',
+                                  'mtcnn-model')
+        self.det_threshold = [0.6, 0.7, 0.8]
+        self.detector = MtcnnDetector(model_folder=mtcnn_path,
+                                      ctx=ctx,
+                                      num_worker=1,
+                                      accurate_landmark=True,
+                                      threshold=self.det_threshold)
+
+    def get(self, img):
+        ret = self.detector.detect_face(img, det_type=0)
+        if ret is None:
+            return None
+        bbox, points = ret
+        if bbox.shape[0] == 0:
+            return None
+        bbox = bbox[0, 0:4]
+        points = points[0, :].reshape((2, 5)).T
+        M = img_helper.estimate_trans_bbox(bbox, self.image_size[0], s=2.0)
+        rimg = cv2.warpAffine(img, M, self.image_size, borderValue=0.0)
+        img = cv2.cvtColor(rimg, cv2.COLOR_BGR2RGB)
+        img = np.transpose(img, (2, 0, 1))  #3*112*112, RGB
+        input_blob = np.zeros((1, 3, self.image_size[1], self.image_size[0]),
+                              dtype=np.uint8)
+        input_blob[0] = img
+        ta = datetime.datetime.now()
+        data = mx.nd.array(input_blob)
+        db = mx.io.DataBatch(data=(data, ))
+        self.model.forward(db, is_train=False)
+        alabel = self.model.get_outputs()[-1].asnumpy()[0]
+        tb = datetime.datetime.now()
+        print('module time cost', (tb - ta).total_seconds())
+        ret = np.zeros((alabel.shape[0], 2), dtype=np.float32)
+        for i in range(alabel.shape[0]):
+            a = cv2.resize(alabel[i], (self.image_size[1], self.image_size[0]))
+            ind = np.unravel_index(np.argmax(a, axis=None), a.shape)
+            #ret[i] = (ind[0], ind[1]) #h, w
+            ret[i] = (ind[1], ind[0])  #w, h
+        return ret, M
+
+
+ctx_id = 4
+img_path = '../deploy/Tom_Hanks_54745.png'
+img = cv2.imread(img_path)
+#img = np.zeros( (128,128,3), dtype=np.uint8 )
+
+handler = Handler('./model/HG', 1, ctx_id)
+for _ in range(10):
+    ta = datetime.datetime.now()
+    landmark, M = handler.get(img)
+    tb = datetime.datetime.now()
+    print('get time cost', (tb - ta).total_seconds())
+#visualize landmark
+IM = cv2.invertAffineTransform(M)
+for i in range(landmark.shape[0]):
+    p = landmark[i]
+    point = np.ones((3, ), dtype=np.float32)
+    point[0:2] = p
+    point = np.dot(IM, point)
+    landmark[i] = point[0:2]
+
+for i in range(landmark.shape[0]):
+    p = landmark[i]
+    point = (int(p[0]), int(p[1]))
+    cv2.circle(img, point, 1, (0, 255, 0), 2)
+
+filename = './landmark_test.png'
+print('writing', filename)
+cv2.imwrite(filename, img)

insightface/alignment/heatmap/test_rec_nme.py

@@ -0,0 +1,71 @@
+import argparse
+import cv2
+import sys
+import numpy as np
+import os
+import mxnet as mx
+import datetime
+import img_helper
+from config import config
+from data import FaceSegIter
+from metric import LossValueMetric, NMEMetric
+
+parser = argparse.ArgumentParser(description='test nme on rec data')
+# general
+parser.add_argument('--rec',
+                    default='./data_2d/ibug.rec',
+                    help='rec data path')
+parser.add_argument('--prefix', default='', help='model prefix')
+parser.add_argument('--epoch', type=int, default=1, help='model epoch')
+parser.add_argument('--gpu', type=int, default=0, help='')
+parser.add_argument('--landmark-type', default='2d', help='')
+parser.add_argument('--image-size', type=int, default=128, help='')
+args = parser.parse_args()
+
+rec_path = args.rec
+ctx_id = args.gpu
+prefix = args.prefix
+epoch = args.epoch
+image_size = (args.image_size, args.image_size)
+config.landmark_type = args.landmark_type
+config.input_img_size = image_size[0]
+
+if ctx_id >= 0:
+    ctx = mx.gpu(ctx_id)
+else:
+    ctx = mx.cpu()
+sym, arg_params, aux_params = mx.model.load_checkpoint(prefix, epoch)
+all_layers = sym.get_internals()
+sym = all_layers['heatmap_output']
+#model = mx.mod.Module(symbol=sym, context=ctx, data_names=['data'], label_names=['softmax_label'])
+model = mx.mod.Module(symbol=sym,
+                      context=ctx,
+                      data_names=['data'],
+                      label_names=None)
+#model = mx.mod.Module(symbol=sym, context=ctx)
+model.bind(for_training=False,
+           data_shapes=[('data', (1, 3, image_size[0], image_size[1]))])
+model.set_params(arg_params, aux_params)
+
+val_iter = FaceSegIter(
+    path_imgrec=rec_path,
+    batch_size=1,
+    aug_level=0,
+)
+_metric = NMEMetric()
+#val_metric = mx.metric.create(_metric)
+#val_metric.reset()
+#val_iter.reset()
+nme = []
+for i, eval_batch in enumerate(val_iter):
+    if i % 10 == 0:
+        print('processing', i)
+    #print(eval_batch.data[0].shape, eval_batch.label[0].shape)
+    batch_data = mx.io.DataBatch(eval_batch.data)
+    model.forward(batch_data, is_train=False)
+    #model.update_metric(val_metric, eval_batch.label, True)
+    pred_label = model.get_outputs()[-1].asnumpy()
+    label = eval_batch.label[0].asnumpy()
+    _nme = _metric.cal_nme(label, pred_label)
+    nme.append(_nme)
+print(np.mean(nme))

insightface/alignment/heatmap/train.py

@@ -0,0 +1,236 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import logging
+import argparse
+from data import FaceSegIter
+import mxnet as mx
+import mxnet.optimizer as optimizer
+import numpy as np
+import os
+import sys
+import math
+import random
+import cv2
+from config import config, default, generate_config
+from optimizer import ONadam
+from metric import LossValueMetric, NMEMetric
+sys.path.append(os.path.join(os.path.dirname(__file__), 'symbol'))
+import sym_heatmap
+#import sym_fc
+#from symbol import fc
+
+args = None
+logger = logging.getLogger()
+logger.setLevel(logging.INFO)
+
+
+def main(args):
+    _seed = 727
+    random.seed(_seed)
+    np.random.seed(_seed)
+    mx.random.seed(_seed)
+    ctx = []
+    cvd = os.environ['CUDA_VISIBLE_DEVICES'].strip()
+    if len(cvd) > 0:
+        for i in range(len(cvd.split(','))):
+            ctx.append(mx.gpu(i))
+    if len(ctx) == 0:
+        ctx = [mx.cpu()]
+        print('use cpu')
+    else:
+        print('gpu num:', len(ctx))
+    #ctx = [mx.gpu(0)]
+    args.ctx_num = len(ctx)
+
+    args.batch_size = args.per_batch_size * args.ctx_num
+    config.per_batch_size = args.per_batch_size
+
+    print('Call with', args, config)
+    train_iter = FaceSegIter(
+        path_imgrec=os.path.join(config.dataset_path, 'train.rec'),
+        batch_size=args.batch_size,
+        per_batch_size=args.per_batch_size,
+        aug_level=1,
+        exf=args.exf,
+        args=args,
+    )
+
+    data_shape = train_iter.get_data_shape()
+    #label_shape = train_iter.get_label_shape()
+    sym = sym_heatmap.get_symbol(num_classes=config.num_classes)
+    if len(args.pretrained) == 0:
+        #data_shape_dict = {'data' : (args.per_batch_size,)+data_shape, 'softmax_label' : (args.per_batch_size,)+label_shape}
+        data_shape_dict = train_iter.get_shape_dict()
+        arg_params, aux_params = sym_heatmap.init_weights(sym, data_shape_dict)
+    else:
+        vec = args.pretrained.split(',')
+        print('loading', vec)
+        _, arg_params, aux_params = mx.model.load_checkpoint(
+            vec[0], int(vec[1]))
+        #sym, arg_params, aux_params = get_symbol(args, arg_params, aux_params)
+
+    model = mx.mod.Module(
+        context=ctx,
+        symbol=sym,
+        label_names=train_iter.get_label_names(),
+    )
+    #lr = 1.0e-3
+    #lr = 2.5e-4
+    _rescale_grad = 1.0 / args.ctx_num
+    #_rescale_grad = 1.0/args.batch_size
+    #lr = args.lr
+    #opt = optimizer.Nadam(learning_rate=args.lr, wd=args.wd, rescale_grad=_rescale_grad, clip_gradient=5.0)
+    if args.optimizer == 'onadam':
+        opt = ONadam(learning_rate=args.lr,
+                     wd=args.wd,
+                     rescale_grad=_rescale_grad,
+                     clip_gradient=5.0)
+    elif args.optimizer == 'nadam':
+        opt = optimizer.Nadam(learning_rate=args.lr,
+                              rescale_grad=_rescale_grad)
+    elif args.optimizer == 'rmsprop':
+        opt = optimizer.RMSProp(learning_rate=args.lr,
+                                rescale_grad=_rescale_grad)
+    elif args.optimizer == 'adam':
+        opt = optimizer.Adam(learning_rate=args.lr, rescale_grad=_rescale_grad)
+    else:
+        opt = optimizer.SGD(learning_rate=args.lr,
+                            momentum=0.9,
+                            wd=args.wd,
+                            rescale_grad=_rescale_grad)
+    initializer = mx.init.Xavier(rnd_type='gaussian',
+                                 factor_type="in",
+                                 magnitude=2)
+    _cb = mx.callback.Speedometer(args.batch_size, args.frequent)
+    _metric = LossValueMetric()
+    #_metric = NMEMetric()
+    #_metric2 = AccMetric()
+    #eval_metrics = [_metric, _metric2]
+    eval_metrics = [_metric]
+    lr_steps = [int(x) for x in args.lr_step.split(',')]
+    print('lr-steps', lr_steps)
+    global_step = [0]
+
+    def val_test():
+        all_layers = sym.get_internals()
+        vsym = all_layers['heatmap_output']
+        vmodel = mx.mod.Module(symbol=vsym, context=ctx, label_names=None)
+        #model.bind(data_shapes=[('data', (args.batch_size, 3, image_size[0], image_size[1]))], label_shapes=[('softmax_label', (args.batch_size,))])
+        vmodel.bind(data_shapes=[('data', (args.batch_size, ) + data_shape)])
+        arg_params, aux_params = model.get_params()
+        vmodel.set_params(arg_params, aux_params)
+        for target in config.val_targets:
+            _file = os.path.join(config.dataset_path, '%s.rec' % target)
+            if not os.path.exists(_file):
+                continue
+            val_iter = FaceSegIter(
+                path_imgrec=_file,
+                batch_size=args.batch_size,
+                #batch_size = 4,
+                aug_level=0,
+                args=args,
+            )
+            _metric = NMEMetric()
+            val_metric = mx.metric.create(_metric)
+            val_metric.reset()
+            val_iter.reset()
+            for i, eval_batch in enumerate(val_iter):
+                #print(eval_batch.data[0].shape, eval_batch.label[0].shape)
+                batch_data = mx.io.DataBatch(eval_batch.data)
+                model.forward(batch_data, is_train=False)
+                model.update_metric(val_metric, eval_batch.label)
+            nme_value = val_metric.get_name_value()[0][1]
+            print('[%d][%s]NME: %f' % (global_step[0], target, nme_value))
+
+    def _batch_callback(param):
+        _cb(param)
+        global_step[0] += 1
+        mbatch = global_step[0]
+        for _lr in lr_steps:
+            if mbatch == _lr:
+                opt.lr *= 0.2
+                print('lr change to', opt.lr)
+                break
+        if mbatch % 1000 == 0:
+            print('lr-batch-epoch:', opt.lr, param.nbatch, param.epoch)
+        if mbatch > 0 and mbatch % args.verbose == 0:
+            val_test()
+            if args.ckpt == 1:
+                msave = mbatch // args.verbose
+                print('saving', msave)
+                arg, aux = model.get_params()
+                mx.model.save_checkpoint(args.prefix, msave, model.symbol, arg,
+                                         aux)
+        if mbatch == lr_steps[-1]:
+            if args.ckpt == 2:
+                #msave = mbatch//args.verbose
+                msave = 1
+                print('saving', msave)
+                arg, aux = model.get_params()
+                mx.model.save_checkpoint(args.prefix, msave, model.symbol, arg,
+                                         aux)
+            sys.exit(0)
+
+    train_iter = mx.io.PrefetchingIter(train_iter)
+
+    model.fit(
+        train_iter,
+        begin_epoch=0,
+        num_epoch=9999,
+        #eval_data          = val_iter,
+        eval_data=None,
+        eval_metric=eval_metrics,
+        kvstore='device',
+        optimizer=opt,
+        initializer=initializer,
+        arg_params=arg_params,
+        aux_params=aux_params,
+        allow_missing=True,
+        batch_end_callback=_batch_callback,
+        epoch_end_callback=None,
+    )
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='Train face alignment')
+    # general
+    parser.add_argument('--network',
+                        help='network name',
+                        default=default.network,
+                        type=str)
+    parser.add_argument('--dataset',
+                        help='dataset name',
+                        default=default.dataset,
+                        type=str)
+    args, rest = parser.parse_known_args()
+    generate_config(args.network, args.dataset)
+    parser.add_argument('--prefix',
+                        default=default.prefix,
+                        help='directory to save model.')
+    parser.add_argument('--pretrained', default=default.pretrained, help='')
+    parser.add_argument('--optimizer', default='nadam', help='')
+    parser.add_argument('--lr', type=float, default=default.lr, help='')
+    parser.add_argument('--wd', type=float, default=default.wd, help='')
+    parser.add_argument('--per-batch-size',
+                        type=int,
+                        default=default.per_batch_size,
+                        help='')
+    parser.add_argument('--lr-step',
+                        help='learning rate steps (in epoch)',
+                        default=default.lr_step,
+                        type=str)
+    parser.add_argument('--ckpt', type=int, default=1, help='')
+    parser.add_argument('--norm', type=int, default=0, help='')
+    parser.add_argument('--exf', type=int, default=1, help='')
+    parser.add_argument('--frequent',
+                        type=int,
+                        default=default.frequent,
+                        help='')
+    parser.add_argument('--verbose',
+                        type=int,
+                        default=default.verbose,
+                        help='')
+    args = parser.parse_args()
+    main(args)

insightface/alignment/synthetics/README.md

@@ -0,0 +1,63 @@
+# Introduction
+
+We provide training and testing tools on synthetics data.
+
+
+## Dataset
+
+### Training dataset
+
+Download `Face Synthetics dataset` from [https://github.com/microsoft/FaceSynthetics](https://github.com/microsoft/FaceSynthetics) and put it somewhere.
+
+<div align="left">
+  <img src="https://github.com/microsoft/FaceSynthetics/raw/main/docs/img/dataset_samples_2.jpg" width="640"/>
+</div>
+<br/>
+
+Then use [tools/prepare_synthetics.py](tools/prepare_synthetics.py) for training data preparation.
+
+
+### Testing dataset
+
+[300-W](https://ibug.doc.ic.ac.uk/resources/300-W/)
+
+
+## Pretrained Model
+
+[ResNet50d](https://drive.google.com/file/d/1kNP7qEl3AYNbaHFUg_ZiyRB1CtfDWXR4/view?usp=sharing)
+
+
+## Train and Test
+
+### Prerequisites
+
+- pytorch_lightning
+- timm
+- albumentations
+
+### Training
+
+`` python -u trainer_synthetics.py ``
+
+which uses `resnet50d` as backbone by default, please check the [code](trainer_synthetics.py) for detail.
+
+### Testing
+
+Please check [test_synthetics.py](test_synthetics.py) for detail.
+
+
+## Result Visualization(3D 68 Keypoints)
+
+<div align="left">
+  <img src="https://github.com/nttstar/insightface-resources/blob/master/alignment/images/image_008_1.jpg?raw=true" width="320"/>
+</div>
+
+<div align="left">
+  <img src="https://github.com/nttstar/insightface-resources/blob/master/alignment/images/image_017_1.jpg?raw=true" width="320"/>
+</div>
+
+<div align="left">
+  <img src="https://github.com/nttstar/insightface-resources/blob/master/alignment/images/image_039.jpg?raw=true" width="320"/>
+</div>
+
+

insightface/alignment/synthetics/datasets/augs.py

@@ -0,0 +1,40 @@
+import numpy as np
+import albumentations as A
+from albumentations.core.transforms_interface import ImageOnlyTransform
+
+class RectangleBorderAugmentation(ImageOnlyTransform):
+
+    def __init__(
+            self,
+            fill_value = 0,
+            limit = 0.3,
+            always_apply=False,
+            p=1.0,
+            ):
+        super(RectangleBorderAugmentation, self).__init__(always_apply, p)
+        assert limit>0.0 and limit<1.0
+        self.fill_value = 0
+        self.limit = limit
+
+
+    def apply(self, image, border_size_limit, **params):
+        assert len(border_size_limit)==4
+        border_size = border_size_limit.copy()
+        border_size[0] *= image.shape[1]
+        border_size[2] *= image.shape[1]
+        border_size[1] *= image.shape[0]
+        border_size[3] *= image.shape[0]
+        border_size = border_size.astype(np.int)
+        image[:,:border_size[0],:] = self.fill_value
+        image[:border_size[1],:,:] = self.fill_value
+        image[:,-border_size[2]:,:] = self.fill_value
+        image[-border_size[3]:,:,:] = self.fill_value
+        return image
+
+    def get_params(self):
+        border_size_limit = np.random.uniform(0.0, self.limit, size=4)
+        return {'border_size_limit': border_size_limit}
+
+    def get_transform_init_args_names(self):
+        return ('fill_value', 'limit')
+

insightface/alignment/synthetics/datasets/dataset_synthetics.py

@@ -0,0 +1,163 @@
+import os
+import os.path as osp
+import queue as Queue
+import pickle
+import threading
+import logging
+import numpy as np
+import torch
+from torch.utils.data import DataLoader, Dataset
+from torchvision import transforms
+import cv2
+import albumentations as A
+from albumentations.pytorch import ToTensorV2
+from .augs import RectangleBorderAugmentation
+
+class BackgroundGenerator(threading.Thread):
+    def __init__(self, generator, local_rank, max_prefetch=6):
+        super(BackgroundGenerator, self).__init__()
+        self.queue = Queue.Queue(max_prefetch)
+        self.generator = generator
+        self.local_rank = local_rank
+        self.daemon = True
+        self.start()
+
+    def run(self):
+        torch.cuda.set_device(self.local_rank)
+        for item in self.generator:
+            self.queue.put(item)
+        self.queue.put(None)
+
+    def next(self):
+        next_item = self.queue.get()
+        if next_item is None:
+            raise StopIteration
+        return next_item
+
+    def __next__(self):
+        return self.next()
+
+    def __iter__(self):
+        return self
+
+
+class DataLoaderX(DataLoader):
+    def __init__(self, local_rank, **kwargs):
+        super(DataLoaderX, self).__init__(**kwargs)
+        self.stream = torch.cuda.Stream(local_rank)
+        self.local_rank = local_rank
+
+    def __iter__(self):
+        self.iter = super(DataLoaderX, self).__iter__()
+        self.iter = BackgroundGenerator(self.iter, self.local_rank)
+        self.preload()
+        return self
+
+    def preload(self):
+        self.batch = next(self.iter, None)
+        if self.batch is None:
+            return None
+        with torch.cuda.stream(self.stream):
+            for k in range(len(self.batch)):
+                self.batch[k] = self.batch[k].to(device=self.local_rank,
+                                                 non_blocking=True)
+
+    def __next__(self):
+        torch.cuda.current_stream().wait_stream(self.stream)
+        batch = self.batch
+        if batch is None:
+            raise StopIteration
+        self.preload()
+        return batch
+
+
+
+class FaceDataset(Dataset):
+    def __init__(self, root_dir, is_train):
+        super(FaceDataset, self).__init__()
+
+        #self.local_rank = local_rank
+        self.is_train = is_train
+        self.input_size = 256
+        self.num_kps = 68
+        transform_list = []
+        if is_train:
+            transform_list += \
+                [
+                    A.ColorJitter(brightness=0.8, contrast=0.5, p=0.5),
+                    A.ToGray(p=0.1),
+                    A.ISONoise(p=0.1),
+                    A.MedianBlur(blur_limit=(1,7), p=0.1),
+                    A.GaussianBlur(blur_limit=(1,7), p=0.1),
+                    A.MotionBlur(blur_limit=(5,12), p=0.1),
+                    A.ImageCompression(quality_lower=50, quality_upper=90, p=0.05),
+                    A.ShiftScaleRotate(shift_limit=0.1, scale_limit=0.2, rotate_limit=40, interpolation=cv2.INTER_LINEAR, 
+                        border_mode=cv2.BORDER_CONSTANT, value=0, mask_value=0, p=0.8),
+                    A.HorizontalFlip(p=0.5),
+                    RectangleBorderAugmentation(limit=0.33, fill_value=0, p=0.2),
+                ]
+        transform_list += \
+            [
+                A.geometric.resize.Resize(self.input_size, self.input_size, interpolation=cv2.INTER_LINEAR, always_apply=True),
+                A.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
+                ToTensorV2(),
+            ]
+        self.transform = A.ReplayCompose(
+            transform_list,
+            keypoint_params=A.KeypointParams(format='xy', remove_invisible=False)
+        )
+        self.root_dir = root_dir
+        with open(osp.join(root_dir, 'annot.pkl'), 'rb') as f:
+            annot = pickle.load(f)
+            self.X, self.Y = annot
+        train_size = int(len(self.X)*0.99)
+        
+        if is_train:
+            self.X = self.X[:train_size]
+            self.Y = self.Y[:train_size]
+        else:
+            self.X = self.X[train_size:]
+            self.Y = self.Y[train_size:]
+        #if local_rank==0:
+        #    logging.info('data_transform_list:%s'%transform_list)
+        flip_parts = ([1, 17], [2, 16], [3, 15], [4, 14], [5, 13], [6, 12], [7, 11], [8, 10],
+            [18, 27], [19, 26], [20, 25], [21, 24], [22, 23],
+            [32, 36], [33, 35],
+            [37, 46], [38, 45], [39, 44], [40, 43], [41, 48], [42, 47],
+            [49, 55], [50, 54], [51, 53], [62, 64], [61, 65], [68, 66], [59, 57], [60, 56])
+        self.flip_order = np.arange(self.num_kps)
+        for pair in flip_parts:
+            self.flip_order[pair[1]-1] = pair[0]-1
+            self.flip_order[pair[0]-1] = pair[1]-1
+        logging.info('len:%d'%len(self.X))
+        print('!!!len:%d'%len(self.X))
+
+    def __getitem__(self, index):
+        x = self.X[index]
+        y = self.Y[index]
+        image_path = os.path.join(self.root_dir, x)
+        img = cv2.imread(image_path)[:,:,::-1]
+        label = y
+        if self.transform is not None:
+            t = self.transform(image=img, keypoints=label)
+            flipped = False
+            for trans in t["replay"]["transforms"]:
+                if trans["__class_fullname__"].endswith('HorizontalFlip'):
+                    if trans["applied"]:
+                        flipped = True
+            img = t['image']
+            label = t['keypoints']
+            label = np.array(label, dtype=np.float32)
+            #print(img.shape)
+            if flipped:
+                #label[:, 0] = self.input_size - 1 - label[:, 0]  #already applied in horizantal flip aug
+                label = label[self.flip_order,:]
+            label /= (self.input_size/2)
+            label -= 1.0
+            label = label.flatten()
+            label = torch.tensor(label, dtype=torch.float32)
+        return img, label
+
+    def __len__(self):
+        return len(self.X)
+

insightface/alignment/synthetics/test_synthetics.py

@@ -0,0 +1,104 @@
+
+from trainer_synthetics import FaceSynthetics
+import sys
+import glob
+import torch
+import os
+import numpy as np
+import cv2
+import os.path as osp
+import insightface
+from insightface.app import FaceAnalysis
+from insightface.utils import face_align
+
+flip_parts = ([1, 17], [2, 16], [3, 15], [4, 14], [5, 13], [6, 12], [7, 11], [8, 10],
+    [18, 27], [19, 26], [20, 25], [21, 24], [22, 23],
+    [32, 36], [33, 35],
+    [37, 46], [38, 45], [39, 44], [40, 43], [41, 48], [42, 47],
+    [49, 55], [50, 54], [51, 53], [62, 64], [61, 65], [68, 66], [59, 57], [60, 56])
+
+app = FaceAnalysis()
+app.prepare(ctx_id=0, det_size=(224, 224))
+input_size = 256
+USE_FLIP = False
+
+root = 'data/300W/Validation'
+output_dir = 'outputs/'
+
+if not osp.exists(output_dir):
+    os.makedirs(output_dir)
+
+outf = open(osp.join(output_dir, 'pred.txt'), 'w')
+
+model = FaceSynthetics.load_from_checkpoint(sys.argv[1]).cuda()
+model.eval()
+for line in open(osp.join(root, '300W_validation.txt'), 'r'):
+    line = line.strip().split()
+    img_path = osp.join(root, line[0])
+    gt = line[1:]
+    #print(len(gt))
+    name = img_path.split('/')[-1]
+    img = cv2.imread(img_path)
+    dimg = img.copy()
+    faces = app.get(img, max_num=1)
+    if len(faces)!=1:
+        continue
+    bbox = faces[0].bbox
+    w, h = (bbox[2] - bbox[0]), (bbox[3] - bbox[1])
+    center = (bbox[2] + bbox[0]) / 2, (bbox[3] + bbox[1]) / 2
+    rotate = 0
+    _scale = input_size  / (max(w, h)*1.5)
+    aimg, M = face_align.transform(img, center, input_size, _scale, rotate)
+    #cv2.imwrite("outputs/a_%s"%name, aimg)
+    aimg = cv2.cvtColor(aimg, cv2.COLOR_BGR2RGB)
+    kps = None
+    flips = [0, 1] if USE_FLIP else [0]
+    for flip in flips:
+        input = aimg.copy()
+        if flip:
+            input = input[:,::-1,:].copy()
+        input = np.transpose(input, (2, 0, 1))
+        input = np.expand_dims(input, 0)
+        imgs = torch.Tensor(input).cuda()
+        imgs.div_(255).sub_(0.5).div_(0.5)
+        pred = model(imgs).detach().cpu().numpy().flatten().reshape( (-1, 2) )
+        pred[:, 0:2] += 1
+        pred[:, 0:2] *= (input_size // 2)
+        if flip:
+            pred_flip = pred.copy()
+            pred_flip[:, 0] = input_size - 1 - pred_flip[:, 0] 
+            for pair in flip_parts:
+                tmp = pred_flip[pair[0] - 1, :].copy()
+                pred_flip[pair[0] - 1, :] = pred_flip[pair[1] - 1, :]
+                pred_flip[pair[1] - 1, :] = tmp
+            pred = pred_flip
+        if kps is None:
+            kps = pred
+        else:
+            kps += pred
+            kps /= 2.0
+    #print(pred.shape)
+
+    IM = cv2.invertAffineTransform(M)
+    kps = face_align.trans_points(kps, IM)
+    outf.write(line[0])
+    outf.write(' ')
+    outf.write(' '.join(["%.5f"%x for x in kps.flatten()]))
+    outf.write("\n")
+    box = bbox.astype(np.int)
+    color = (0, 0, 255)
+    cv2.rectangle(dimg, (box[0], box[1]), (box[2], box[3]), color, 2)
+    kps = kps.astype(np.int)
+    #print(landmark.shape)
+    for l in range(kps.shape[0]):
+        color = (0, 0, 255)
+        cv2.circle(dimg, (kps[l][0], kps[l][1]), 1, color, 2)
+
+    cv2.imwrite("outputs/%s"%name, dimg)
+
+    #ret = np.argmax(feat)
+    #print(feat)
+    #outf.write("%s %.4f %.4f %.4f\n"%(line[0], feat[0], feat[1], feat[2]))
+
+outf.close()
+

insightface/alignment/synthetics/tools/prepare_synthetics.py

@@ -0,0 +1,70 @@
+
+import sys
+import glob
+import torch
+import pickle
+import os
+import numpy as np
+import cv2
+import os.path as osp
+import insightface
+from insightface.app import FaceAnalysis
+from insightface.utils import face_align
+
+app = FaceAnalysis()
+app.prepare(ctx_id=0, det_size=(224, 224))
+output_size = 384
+
+input_dir = '/root/codebase/FaceSynthetics'
+output_dir = 'data/synthetics'
+
+if not osp.exists(output_dir):
+    os.makedirs(output_dir)
+
+X = []
+Y = []
+
+for i in range(0, 100000):
+    if i%1000==0:
+        print('loading', i)
+    x = "%06d.png"%i
+    img_path = osp.join(input_dir, x)
+    img = cv2.imread(img_path)
+    dimg = img.copy()
+    ylines = open(osp.join(input_dir, "%06d_ldmks.txt"%i)).readlines()
+    ylines = ylines[:68]
+    y = []
+    for yline in ylines:
+        lmk = [float(x) for x in yline.strip().split()]
+        y.append( tuple(lmk) )
+    pred = np.array(y)
+    faces = app.get(img, max_num=1)
+    if len(faces)!=1:
+        continue
+    bbox = faces[0].bbox
+    w, h = (bbox[2] - bbox[0]), (bbox[3] - bbox[1])
+    center = (bbox[2] + bbox[0]) / 2, (bbox[3] + bbox[1]) / 2
+    rotate = 0
+    _scale = output_size  / (max(w, h)*1.5)
+    aimg, M = face_align.transform(dimg, center, output_size, _scale, rotate)
+    pred = face_align.trans_points(pred, M)
+    #box = bbox.astype(np.int)
+    #color = (0, 0, 255)
+    #cv2.rectangle(dimg, (box[0], box[1]), (box[2], box[3]), color, 2)
+
+    #kps = pred.astype(np.int)
+    #for l in range(kps.shape[0]):
+    #    color = (0, 0, 255)
+    #    cv2.circle(aimg, (kps[l][0], kps[l][1]), 1, color, 2)
+    x = x.replace('png', 'jpg')
+    X.append(x)
+    y = []
+    for k in range(pred.shape[0]):
+        y.append( (pred[k][0], pred[k][1]) )
+    Y.append(y)
+    cv2.imwrite("%s/%s"%(output_dir, x), aimg)
+
+
+with open(osp.join(output_dir, 'annot.pkl'), 'wb') as pfile:
+    pickle.dump((X, Y), pfile, protocol=pickle.HIGHEST_PROTOCOL)
+

insightface/alignment/synthetics/trainer_synthetics.py

@@ -0,0 +1,140 @@
+from argparse import ArgumentParser
+
+import os
+import os.path as osp
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+from torch.utils.data import DataLoader
+import pytorch_lightning as pl
+from pytorch_lightning.callbacks import ModelCheckpoint
+from pytorch_lightning.callbacks import LearningRateMonitor
+from pytorch_lightning.loggers import TensorBoardLogger
+import timm
+from datasets.dataset_synthetics import FaceDataset, DataLoaderX
+
+
+class FaceSynthetics(pl.LightningModule):
+    def __init__(self, backbone):
+        super().__init__()
+        self.save_hyperparameters()
+        backbone = timm.create_model(backbone, num_classes=68*2)
+        self.backbone = backbone
+        self.loss = nn.L1Loss(reduction='mean')
+        self.hard_mining = False
+
+    def forward(self, x):
+        # use forward for inference/predictions
+        y = self.backbone(x)
+        return y
+
+    def training_step(self, batch, batch_idx):
+        x, y = batch
+        y_hat = self.backbone(x)
+        if self.hard_mining:
+            loss = torch.abs(y_hat - y) #(B,K)
+            loss = torch.mean(loss, dim=1) #(B,)
+            B = len(loss)
+            S = int(B*0.5)
+            loss, _ = torch.sort(loss, descending=True)
+            loss = loss[:S]
+            loss = torch.mean(loss) * 5.0
+        else:
+            loss = self.loss(y_hat, y) * 5.0
+        self.log('train_loss', loss, on_epoch=True)
+        return loss
+
+    def validation_step(self, batch, batch_idx):
+        x, y = batch
+        y_hat = self.backbone(x)
+        loss = self.loss(y_hat, y)
+        self.log('val_loss', loss, on_step=True)
+
+    def test_step(self, batch, batch_idx):
+        x, y = batch
+        y_hat = self.backbone(x)
+        loss = self.loss(y_hat, y)
+        self.log('test_loss', loss)
+
+    def configure_optimizers(self):
+        #return torch.optim.Adam(self.parameters(), lr=0.0002)
+        opt = torch.optim.SGD(self.parameters(), lr = 0.1, momentum=0.9, weight_decay = 0.0005)
+        def lr_step_func(epoch):
+            return 0.1 ** len([m for m in [15, 25, 28] if m <= epoch])
+        scheduler = torch.optim.lr_scheduler.LambdaLR(
+                optimizer=opt, lr_lambda=lr_step_func)
+        lr_scheduler = {
+                'scheduler': scheduler,
+                'name': 'learning_rate',
+                'interval':'epoch',
+                'frequency': 1}
+        return [opt], [lr_scheduler]
+
+
+
+def cli_main():
+    pl.seed_everything(727)
+
+    # ------------
+    # args
+    # ------------
+    parser = ArgumentParser()
+    parser.add_argument('--backbone', default='resnet50d', type=str)
+    parser.add_argument('--batch_size', default=64, type=int)
+    parser.add_argument('--root', default='data/synthetics', type=str)
+    parser.add_argument('--num-gpus', default=2, type=int)
+    parser.add_argument('--tf32', action='store_true')
+    parser = pl.Trainer.add_argparse_args(parser)
+    args = parser.parse_args()
+
+    if not args.tf32:
+        torch.backends.cuda.matmul.allow_tf32 = False
+        torch.backends.cudnn.allow_tf32 = False
+    else:
+        torch.backends.cuda.matmul.allow_tf32 = True
+        torch.backends.cudnn.allow_tf32 = True
+    torch.backends.cudnn.benchmark = True
+
+    # ------------
+    # data
+    # ------------
+    train_set = FaceDataset(root_dir=args.root, is_train=True)
+    val_set = FaceDataset(root_dir=args.root, is_train=False)
+
+    train_loader = DataLoader(train_set, batch_size=args.batch_size, shuffle=True, num_workers=3, pin_memory=True)
+    val_loader = DataLoader(val_set, batch_size=args.batch_size, shuffle=False)
+
+    # ------------
+    # model
+    # ------------
+    model = FaceSynthetics(backbone=args.backbone)
+    ckpt_path = 'work_dirs/synthetics'
+    if not os.path.exists(ckpt_path):
+        os.makedirs(ckpt_path)
+
+    # ------------
+    # training
+    # ------------
+    checkpoint_callback = ModelCheckpoint(
+            monitor='val_loss',
+            dirpath=ckpt_path,
+            filename='{epoch:02d}-{val_loss:.6f}',
+            save_top_k=10,
+            mode='min',
+            )
+    lr_monitor = LearningRateMonitor(logging_interval='step')
+    trainer = pl.Trainer(
+        gpus = args.num_gpus,
+        accelerator="ddp",
+        benchmark=True,
+		logger=TensorBoardLogger(osp.join(ckpt_path, 'logs')),
+        callbacks=[checkpoint_callback, lr_monitor],
+        check_val_every_n_epoch=1,
+        progress_bar_refresh_rate=1,
+        max_epochs=30,
+    )
+    trainer.fit(model, train_loader, val_loader)
+
+if __name__ == '__main__':
+    cli_main()
+

insightface/attribute/README.md

@@ -0,0 +1,33 @@
+## Face Attribute
+
+
+<div align="left">
+  <img src="https://insightface.ai/assets/img/custom/logo3.jpg" width="320"/>
+</div>
+
+
+## Introduction
+
+These are the face attribute methods of [InsightFace](https://insightface.ai)
+
+
+<div align="left">
+  <img src="https://insightface.ai/assets/img/github/t1_genderage.jpg" width="600"/>
+</div>
+
+
+
+## Methods
+
+
+Supported methods:
+
+- [x] [Gender_Age](gender_age)
+
+
+
+## Contributing
+
+We appreciate all contributions to improve the face attribute module of InsightFace. 
+
+

insightface/attribute/_datasets_/README.md

@@ -0,0 +1,15 @@
+# Face Attribute Datasets
+
+(Updating)
+
+## Training Datasets
+
+### CelebA
+
+https://mmlab.ie.cuhk.edu.hk/projects/CelebA.html
+
+
+
+## Test Datasets
+
+

insightface/attribute/gender_age/test.py

@@ -0,0 +1,24 @@
+import argparse
+import cv2
+import sys
+import numpy as np
+import insightface
+from insightface.app import FaceAnalysis
+from insightface.data import get_image as ins_get_image
+
+
+parser = argparse.ArgumentParser(description='insightface gender-age test')
+# general
+parser.add_argument('--ctx', default=0, type=int, help='ctx id, <0 means using cpu')
+args = parser.parse_args()
+
+app = FaceAnalysis(allowed_modules=['detection', 'genderage'])
+app.prepare(ctx_id=args.ctx, det_size=(640,640))
+
+img = ins_get_image('t1')
+faces = app.get(img)
+assert len(faces)==6
+for face in faces:
+    print(face.bbox)
+    print(face.sex, face.age)
+

insightface/benchmarks/train/nvidia_a10.md

@@ -0,0 +1,48 @@
+# Training performance report on NVIDIA A10
+
+[NVIDIA A10 Tensor Core GPU](https://www.nvidia.com/en-us/data-center/products/a10-gpu/) 
+
+We can use A10 to train deep learning models by its FP16 and TF32 supports.
+
+
+
+## Test Server Spec
+
+| Key          | Value                                            |
+| ------------ | ------------------------------------------------ |
+| System       | ServMax G408-X2 Rackmountable Server             |
+| CPU          | 2 x Intel(R) Xeon(R) Gold 5220R CPU @ 2.20GHz    |
+| Memory       | 384GB, 12 x Samsung 32GB DDR4-2933               |
+| GPU          | 8 x NVIDIA A10 22GB                              |
+| Cooling      | 2x Customized GPU Kit for GPU support FAN-1909L2 |
+| Hard Drive   | Intel SSD S4500 1.9TB/SATA/TLC/2.5"              |
+| OS           | Ubuntu 16.04.7 LTS                               |
+| Installation | CUDA 11.1, cuDNN 8.0.5                           |
+| Installation | Python 3.7.10                                    |
+| Installation | PyTorch 1.9.0 (conda)                            |
+
+This server is donated by [AMAX](https://www.amaxchina.com/), many thanks!
+
+
+
+## Experiments on arcface_torch
+
+We report training speed in following table, please also note that:
+
+1. The training dataset is in mxnet record format and located on SSD hard drive.
+
+2. Embedding-size are all set to 512.
+
+3. We use a large dataset which contains about 618K identities to simulate real cases.
+
+| Dataset     | Classes | Backbone    | Batch-size | FP16 | TF32 | Samples/sec |
+| ----------- | ------- | ----------- | ---------- | ---- | ---- | ----------- |
+| WebFace600K | 618K    | IResNet-50  | 1024       | ×    | ×    | ~2040       |
+| WebFace600K | 618K    | IResNet-50  | 1024       | ×    | √    | ~2255       |
+| WebFace600K | 618K    | IResNet-50  | 1024       | √    | ×    | ~3300       |
+| WebFace600K | 618K    | IResNet-50  | 1024       | √    | √    | ~3360       |
+| WebFace600K | 618K    | IResNet-50  | 2048       | √    | √    | ~3940       |
+| WebFace600K | 618K    | IResNet-100 | 1024       | √    | √    | ~2210       |
+| WebFace600K | 618K    | IResNet-180 | 1024       | √    | √    | ~1410       |
+
+

insightface/benchmarks/train/nvidia_a100.md

@@ -0,0 +1,53 @@
+# Training performance report on NVIDIA A100
+
+[NVIDIA A100 Tensor Core GPU](https://www.nvidia.com/en-us/data-center/a100/) 
+
+
+
+## Test Server Spec
+
+| Key          | Value                                            |
+| ------------ | ------------------------------------------------ |
+| System       | ServMax G408-X2 Rackmountable Server             |
+| CPU          | 2 x Intel(R) Xeon(R) Gold 5220R CPU @ 2.20GHz    |
+| Memory       | 384GB, 12 x Samsung 32GB DDR4-2933               |
+| GPU          | 8 x NVIDIA A100 80GB                             |
+| Cooling      | 2x Customized GPU Kit for GPU support FAN-1909L2 |
+| Hard Drive   | Intel SSD S4500 1.9TB/SATA/TLC/2.5"              |
+| OS           | Ubuntu 16.04.7 LTS                               |
+| Installation | CUDA 11.1, cuDNN 8.0.5                           |
+| Installation | Python 3.7.10                                    |
+| Installation | PyTorch 1.9.0 (conda)                            |
+
+This server is donated by [AMAX](https://www.amaxchina.com/), many thanks!
+
+
+
+## Experiments on arcface_torch
+
+We report training speed in following table, please also note that:
+
+1. The training dataset is in mxnet record format and located on SSD hard drive.
+2. Embedding-size are all set to 512.
+3. We use large datasets with about 618K/2M identities to simulate real cases.
+4. We test the 10K batch-size on real dataset to take the full advantage of 80GB memory.
+5. We also test on huge synthetic datasets which include 50M~80M classes.
+
+| Dataset     | Classes | Backbone    | Batch-size | PFC  | FP16 | TF32 | Samples/sec | GPU Mem(GB) |
+| ----------- | ------- | ----------- | ---------- | ---- | ---- | ---- | ----------- | ----------- |
+| WebFace600K | 618K    | IResNet-50  | 1024       | ×    | ×    | ×    | ~3670       | ~18.2       |
+| WebFace600K | 618K    | IResNet-50  | 1024       | ×    | ×    | √    | ~4760       | ~15.0       |
+| WebFace600K | 618K    | IResNet-50  | 1024       | ×    | √    | ×    | ~5170       | ~10.1       |
+| WebFace600K | 618K    | IResNet-50  | 1024       | ×    | √    | √    | ~5400       | ~10.1       |
+| WebFace600K | 618K    | IResNet-50  | 2048       | ×    | √    | √    | ~7780       | ~16.4       |
+| WebFace600K | 618K    | IResNet-50  | 10240      | ×    | √    | √    | ~9400       | ~66.7       |
+| WebFace600K | 618K    | IResNet-100 | 1024       | ×    | √    | √    | ~3700       | ~13.1       |
+| WebFace600K | 618K    | IResNet-180 | 1024       | ×    | √    | √    | ~2380       | ~17.5       |
+| WebFace2M   | 2M      | IResNet-100 | 1024       | ×    | √    | √    | ~3480       | ~20.5       |
+| WebFace2M   | 2M      | IResNet-180 | 1024       | ×    | √    | √    | ~2350       | ~25.0       |
+| WebFace2M   | 2M      | IResNet-300 | 1024       | ×    | √    | √    | ~1541       | ~32.6       |
+| Virtual     | 50M     | IResNet-50  | 1024       | 0.1  | √    | √    | ~2700       | ~54.1       |
+| Virtual     | 70M     | IResNet-50  | 1024       | 0.1  | √    | √    | ~2170       | ~73.7       |
+| Virtual     | 80M     | IResNet-50  | 1024       | 0.1  | √    | √    | ~1080       | ~79.6       |
+
+

insightface/benchmarks/train/nvidia_a30.md

@@ -0,0 +1,52 @@
+# Training performance report on NVIDIA A30
+
+[NVIDIA A30 Tensor Core GPU](https://www.nvidia.com/en-us/data-center/products/a30-gpu/) is the most versatile mainstream
+compute GPU for AI inference and mainstream enterprise
+workloads. 
+
+Besides, we can also use A30 to train deep learning models by its FP16 and TF32 supports.
+
+
+
+## Test Server Spec
+
+| Key          | Value                                            |
+| ------------ | ------------------------------------------------ |
+| System       | ServMax G408-X2 Rackmountable Server             |
+| CPU          | 2 x Intel(R) Xeon(R) Gold 5220R CPU @ 2.20GHz    |
+| Memory       | 384GB, 12 x Samsung 32GB DDR4-2933               |
+| GPU          | 8 x NVIDIA A30 24GB                              |
+| Cooling      | 2x Customized GPU Kit for GPU support FAN-1909L2 |
+| Hard Drive   | Intel SSD S4500 1.9TB/SATA/TLC/2.5"              |
+| OS           | Ubuntu 16.04.7 LTS                               |
+| Installation | CUDA 11.1, cuDNN 8.0.5                           |
+| Installation | Python 3.7.10                                    |
+| Installation | PyTorch 1.9 (conda)                              |
+
+This server is donated by [AMAX](https://www.amaxchina.com/), many thanks!
+
+
+
+## Experiments on arcface_torch
+
+We report training speed in following table, please also note that:
+
+1. The training dataset is in mxnet record format and located on SSD hard drive.
+
+2. Embedding-size are all set to 512.
+
+3. We use a large dataset which contains about 618K identities to simulate real cases.
+
+| Dataset     | Classes | Backbone    | Batch-size | FP16 | TF32 | Samples/sec |
+| ----------- | ------- | ----------- | ---------- | ---- | ---- | ----------- |
+| WebFace600K | 618K    | IResNet-50  | 1024       | ×    | ×    | ~2230       |
+| WebFace600K | 618K    | IResNet-50  | 1024       | ×    | √    | ~3200       |
+| WebFace600K | 618K    | IResNet-50  | 1024       | √    | ×    | ~3940       |
+| WebFace600K | 618K    | IResNet-50  | 1024       | √    | √    | ~4350       |
+| WebFace600K | 618K    | IResNet-50  | 2048       | √    | √    | ~5100       |
+| WebFace600K | 618K    | IResNet-100 | 1024       | √    | √    | ~2810       |
+| WebFace600K | 618K    | IResNet-180 | 1024       | √    | √    | ~1800       |
+
+
+
+

insightface/benchmarks/train/nvidia_rtx3080.md

@@ -0,0 +1,58 @@
+# Training performance report on NVIDIA RTX3080
+
+[GeForce RTX 3080](https://www.nvidia.com/en-us/geforce/graphics-cards/30-series/rtx-3080-3080ti/) 
+The GeForce RTX™ 3080 Ti and RTX 3080 graphics cards deliver the ultra performance that gamers crave, powered by Ampere—NVIDIA’s 2nd gen RTX architecture. They are built with enhanced RT Cores and Tensor Cores, new streaming multiprocessors, and superfast G6X memory for an amazing gaming experience.
+
+Besides, we can also use GeForce RTX™ 3080 to train deep learning models by its FP16 and TF32 supports.
+
+
+
+## Test Server Spec
+
+| Key          | Value                                             |
+|--------------|---------------------------------------------------|
+| CPU          | 2 x Intel(R) Xeon(R) Platinum 8255C CPU @ 2.50GHz |
+| Memory       | 384GB                                             |
+| GPU          | 8 x GeForce RTX™ 3080                             |
+| OS           | Ubuntu 18.04.4 LTS                                |
+| Installation | CUDA 11.1,                                        |
+| Installation | Python 3.7.3                                      |
+| Installation | PyTorch 1.9.0 (pip)                               |
+
+
+## Experiments on arcface_torch
+
+We report training speed in following table, please also note that:
+
+1. The training dataset is SyntheticDataset.
+
+2. Embedding-size are all set to 512.
+
+
+### 1. 2 Million Identities
+
+We use a large dataset which contains about 2 millions identities to simulate real cases.
+
+
+| Dataset    | Classes    | Backbone   | Batch-size | FP16 | Partial FC | Samples/sec |
+|------------|------------|------------|------------|------|------------|-------------|
+| WebFace40M | 2 Millions | IResNet-50 | 512        | ×    | ×          | Fail        |
+| WebFace40M | 2 Millions | IResNet-50 | 512        | x    | √          | ~2190       |
+| WebFace40M | 2 Millions | IResNet-50 | 512        | √    | ×          | Fail        |
+| WebFace40M | 2 Millions | IResNet-50 | 512        | √    | √          | ~2620       |
+| WebFace40M | 2 Millions | IResNet-50 | 1024       | ×    | ×          | Fail        |
+| WebFace40M | 2 Millions | IResNet-50 | 1024       | x    | √          | Fail        |
+| WebFace40M | 2 Millions | IResNet-50 | 1024       | √    | ×          | Fail        |
+| WebFace40M | 2 Millions | IResNet-50 | 1024       | √    | √          | ~3800       |
+
+### 2. 600K Identities
+
+We use a large dataset which contains about 600k identities to simulate real cases.
+
+| Dataset     | Classes | Backbone   | Batch-size | Partial FC | FP16 | Samples/sec |
+|-------------|---------|------------|------------|------------|------|-------------|
+| WebFace600K | 618K    | IResNet-50 | 512        | ×          | ×    | ~2023       |
+| WebFace600K | 618K    | IResNet-50 | 512        | ×          | √    | ~2392       |
+| WebFace600K | 618K    | IResNet-50 | 1024       | ×          | ×    | Fail        |
+| WebFace600K | 618K    | IResNet-50 | 1024       | ×          | √    | Fail        |
+| WebFace600K | 618K    | IResNet-50 | 1024       | √          | √    | ~4010       |

insightface/benchmarks/train/nvidia_rtx3090.md

@@ -0,0 +1,57 @@
+# Training performance report on NVIDIA RTX3090
+
+[GEFORCE RTX 3090](https://www.nvidia.com/en-us/geforce/graphics-cards/30-series/rtx-3090/) 
+The GeForce RTX™ 3090 is a big ferocious GPU (BFGPU) with TITAN class performance.
+
+Besides, we can also use GeForce RTX™ 3090 to train deep learning models by its FP16 and TF32 supports.
+
+
+
+## Test Server Spec
+
+| Key          | Value                                             |
+|--------------|---------------------------------------------------|
+| CPU          | 2 x Intel(R) Xeon(R) Platinum 8255C CPU @ 2.50GHz |
+| Memory       | 384GB                                             |
+| GPU          | 8 x GeForce RTX™ 3090                             |
+| OS           | Ubuntu 18.04.4 LTS                                |
+| Installation | CUDA 11.1,                                        |
+| Installation | Python 3.7.3                                      |
+| Installation | PyTorch 1.9.0 (pip)                               |
+
+
+## Experiments on arcface_torch
+
+We report training speed in following table, please also note that:
+
+1. The training dataset is SyntheticDataset.
+
+2. Embedding-size are all set to 512.
+
+
+### 1. 2 Million Identities
+
+We use a large dataset which contains about 2 millions identities to simulate real cases.
+
+| Dataset    | Classes    | Backbone   | Batch-size | FP16 | TF32 | Partial FC | Samples/sec |
+|------------|------------|------------|------------|------|------|------------|-------------|
+| WebFace40M | 2 Millions | IResNet-50 | 512        | ×    | ×    | ×          | ~1750       |
+| WebFace40M | 2 Millions | IResNet-50 | 512        | ×    | √    | ×          | ~1810       |
+| WebFace40M | 2 Millions | IResNet-50 | 512        | √    | √    | ×          | ~2056       |
+| WebFace40M | 2 Millions | IResNet-50 | 512        | √    | √    | √          | ~2850       |
+| WebFace40M | 2 Millions | IResNet-50 | 1024       | √    | √    | ×          | ~2810       |
+| WebFace40M | 2 Millions | IResNet-50 | 1024       | √    | √    | √          | ~4220       |
+| WebFace40M | 2 Millions | IResNet-50 | 2048       | √    | √    | √          | ~5330       |
+
+
+### 2. 600K Identities
+
+We use a large dataset which contains about 600k identities to simulate real cases.
+
+| Dataset     | Classes | Backbone   | Batch-size | FP16 | Samples/sec |
+|-------------|---------|------------|------------|------|-------------|
+| WebFace600K | 618K    | IResNet-50 | 512        | ×    | ~2220       |
+| WebFace600K | 618K    | IResNet-50 | 512        | √    | ~2610       |
+| WebFace600K | 618K    | IResNet-50 | 1024       | ×    | ~2940       |
+| WebFace600K | 618K    | IResNet-50 | 1024       | √    | ~3790       |
+| WebFace600K | 618K    | IResNet-50 | 2048       | √    | ~4680       |

insightface/benchmarks/train/nvidia_v100.md

@@ -0,0 +1,54 @@
+# Training performance report on NVIDIA® V100
+
+[NVIDIA® V100](https://www.nvidia.com/en-us/data-center/v100/) 
+NVIDIA® V100 Tensor Core is the most advanced data center GPU ever built to accelerate AI, high performance computing (HPC), data science and graphics. It’s powered by NVIDIA Volta architecture, comes in 16 and 32GB configurations, and offers the performance of up to 32 CPUs in a single GPU.
+
+Besides, we can also use NVIDIA® V100 to train deep learning models by its FP16 and FP32 supports.
+
+## Test Server Spec
+
+| Key          | Value                                        |
+|--------------|----------------------------------------------|
+| CPU          | 2 x Intel(R) Xeon(R) Gold 6133 CPU @ 2.50GHz |
+| Memory       | 384GB                                        |
+| GPU          | 8 x Tesla V100-SXM2-32GB                     |
+| OS           | Ubuntu 16.04 LTS                             |
+| Installation | CUDA 10.2                                    |
+| Installation | Python 3.7.3                                 |
+| Installation | PyTorch 1.9.0 (pip)                          |
+
+## Experiments on arcface_torch
+
+We report training speed in following table, please also note that:
+
+1. The training dataset is SyntheticDataset.
+
+2. Embedding-size are all set to 512.
+
+### 1. 2 Million Identities
+
+We use a large dataset which contains about 2 millions identities to simulate real cases.
+
+| Dataset    | Classes    | Backbone   | Batch-size | FP16 | Partial FC | Samples/sec |
+|------------|------------|------------|------------|------|------------|-------------|
+| WebFace40M | 2 Millions | IResNet-50 | 512        | ×    | ×          | ~1868       |
+| WebFace40M | 2 Millions | IResNet-50 | 512        | x    | √          | ~2712       |
+| WebFace40M | 2 Millions | IResNet-50 | 512        | √    | ×          | ~2576       |
+| WebFace40M | 2 Millions | IResNet-50 | 512        | √    | √          | ~4501       |
+| WebFace40M | 2 Millions | IResNet-50 | 1024       | ×    | ×          | ~1960       |
+| WebFace40M | 2 Millions | IResNet-50 | 1024       | x    | √          | ~2922       |
+| WebFace40M | 2 Millions | IResNet-50 | 1024       | √    | ×          | ~2810       |
+| WebFace40M | 2 Millions | IResNet-50 | 1024       | √    | √          | ~5430       |
+| WebFace40M | 2 Millions | IResNet-50 | 2048       | √    | √          | ~6095       |
+
+### 2. 600K Identities
+
+We use a large dataset which contains about 600k identities to simulate real cases.
+
+| Dataset     | Classes | Backbone   | Batch-size | FP16 | Samples/sec |
+|-------------|---------|------------|------------|------|-------------|
+| WebFace600K | 618K    | IResNet-50 | 512        | ×    | ~2430       |
+| WebFace600K | 618K    | IResNet-50 | 512        | √    | ~3889       |
+| WebFace600K | 618K    | IResNet-50 | 1024       | ×    | ~2607       |
+| WebFace600K | 618K    | IResNet-50 | 1024       | √    | ~4322       |
+| WebFace600K | 618K    | IResNet-50 | 2048       | √    | ~4921       |

insightface/body/human_pose/ambiguity_aware/README.md

@@ -0,0 +1,94 @@
+# Towards Alleviating the Modeling Ambiguity of Unsupervised Monocular 3D Human Pose Estimation
+
+## Introduction
+
+**Ambiguity-Aware** studies the ambiguity problem in the task of unsupervised 3D human pose estimation from 2D counterpart, please refer to [ICCV2022](https://openaccess.thecvf.com/content/ICCV2021/papers/Yu_Towards_Alleviating_the_Modeling_Ambiguity_of_Unsupervised_Monocular_3D_Human_ICCV_2021_paper.pdf) for more details.
+
+
+<div align="center">
+ <img src="https://github.com/yuzhenbo/yuzhenbo.github.io/raw/main/assets/extra/ICCV2022/all.gif" alt="videovis" width="800">
+</div>
+
+
+## Installation 
+```
+conda create -n uvhpe python=3.6 
+conda activate uvhpe
+pip install -r requirements.txt
+# for output,  tensorboard, visualization  
+mkdir log output vis models data
+```
+
+## Dataset And Pretrained Models 
+Download our preprocessed dataset into `data` and pretrained models into `models` from [webpage](https://sites.google.com/view/ambiguity-aware-hpe)
+
+This part will be updated soon.
+## Inference 
+We put some samples with preprocessed 2d keypoints at `scripts/demo_input`. Run inference with command `sh demo.sh` and output can be found at `scripts/demo_output`. 
+
+## Evaluation
+### Evaluation on Human3.6M 
+##### 2D ground-truth as inputs
+* baseline `python main.py --cfg ../cfg/h36m_gt_adv.yaml --pretrain ../models/adv.pth.tar --gpu 0 --eval `
+* scale `python main.py --cfg ../cfg/h36m_gt_scale.yaml --pretrain ../models/tmc_klbone.pth.tar  --eval --gpu 0`
+
+##### 2D predictions as inputs
+* baseline `python main.py --cfg ../cfg/pre_adv.yaml --pretrain ../models/pre_adv.pth.tar --gpu 0 --eval `
+* scale `python main.py --cfg ../cfg/pre_tmc_klbone.yaml --pretrain ../models/pre_tmc_klbone.pth.tar --gpu 0 --eval `
+
+**Note:** baseline is our reproduced version fo "Unsupervised 3d pose estimation with geometric self-supervision"
+
+### Evaluation on LSP
+use the pretrained model from Human3.6M
+
+`python eval_lsp.py --cfg ../cfg/h36m_gt_scale.yaml --pretrain ../models/tmc_klbone.pth.tar`
+
+### Results
+
+The expected **MPJPE** and **P-MPJPE**  results on **Human36M** dataset are shown here:
+
+| Input  | Model                         |     MPJPE     |     PMPJPE     | 
+| :--------- | :------------                  | :------------: | :------------: | 
+| GT | baseline                              |      105.0      |       46.0    |   
+| GT | best                             |      87.85      |       42.0     |     
+| Pre | baseline                             |      113.3     |    54.9     | 
+| Pre | best                            |      93.1       |    52.3     | 
+
+
+**Note:**  MPJPE from the evaluation is slightly different from the performance we release in the paper. This is because MPJPE in the paper is the best MPJPE during training process.
+
+
+
+## Training 
+### Human3.6M 
+* Using ground-truth 2D as inputs: 
+    
+    baseline `python main.py --cfg ../cfg/h36m_gt_adv.yaml --gpu 0 `
+
+    best `python main.py --cfg ../cfg/h36m_gt_scale.yaml --gpu 0`
+
+* Using predicted 2D as inputs: 
+
+    baseline `python main.py --cfg ../cfg/pre_adv.yaml --gpu 0 `
+
+    best `python main.py --cfg ../cfg/pre_tmc_klbone.yaml --gpu 0`
+
+## Visualization
+
+### Human3.6M
+<div align="center">
+ <img src="https://github.com/yuzhenbo/yuzhenbo.github.io/raw/main/assets/extra/ICCV2022/S9_Discussion 1.gif" width="200"/><img src="https://github.com/yuzhenbo/yuzhenbo.github.io/raw/main/assets/extra/ICCV2022/S9_Phoning 1.gif" width="200"/><img src="https://github.com/yuzhenbo/yuzhenbo.github.io/raw/main/assets/extra/ICCV2022/S9_Photo.gif" width="200"/><img src="https://github.com/yuzhenbo/yuzhenbo.github.io/raw/main/assets/extra/ICCV2022/S9_WalkTogether 1.gif" width="200"/>
+</div>
+
+### Sureal
+<div align="center">
+ <img src="https://github.com/yuzhenbo/yuzhenbo.github.io/raw/main/assets/extra/ICCV2022/surreal1.gif" width="200"/><img src="https://github.com/yuzhenbo/yuzhenbo.github.io/raw/main/assets/extra/ICCV2022/surreal2.gif" width="200"/><img src="https://github.com/yuzhenbo/yuzhenbo.github.io/raw/main/assets/extra/ICCV2022/surreal3.gif" width="200"/><img src="https://github.com/yuzhenbo/yuzhenbo.github.io/raw/main/assets/extra/ICCV2022/surreal4.gif" width="200"/>
+</div>
+
+### MPI-3DHP
+<div align="center">
+ <img src="https://github.com/yuzhenbo/yuzhenbo.github.io/raw/main/assets/extra/ICCV2022/TS1.gif" width="200"/><img src="https://github.com/yuzhenbo/yuzhenbo.github.io/raw/main/assets/extra/ICCV2022/TS2.gif" width="200"/><img src="https://github.com/yuzhenbo/yuzhenbo.github.io/raw/main/assets/extra/ICCV2022/TS3.gif" width="200"/><img src="https://github.com/yuzhenbo/yuzhenbo.github.io/raw/main/assets/extra/ICCV2022/TS6.gif" width="200"/>
+</div>
+
+
+### The code of our another paper in ICCV2022 Skeleton2Mesh will be coming soon! 

insightface/body/human_pose/ambiguity_aware/cfg/h36m_gt_adv.yaml

@@ -0,0 +1,36 @@
+BATCH_SIZE: 512
+DATA:
+  NUM_FRAMES: 1
+  SCALE_MID_MEAN: 0.720643
+  SCALE_MID_STD: 0.058
+  USE_RANDOM_DIFF: true
+NETWORK:
+  DIS_RES_BLOCKS: 2
+  DIS_TEMP_RES_BLOCKS: 2
+  DIS_USE_SPECTRAL_NORM: false
+  SCALER_INPUT_SIZE: 34
+TRAIN:
+  BOUND_AZIM: 2.44346
+  BOUND_ELEV: 0.34906585
+  DIS_LR: 0.0002
+  LOSS_TYPE: ss_adv
+  LOSS_WEIGHTS:
+  - 1.0
+  - 1.0
+  - 1.0
+  - 1.0
+  MAINNET_CRITICS: 4
+  NUM_CRITICS: 3
+  NUM_CRITICS_TEMP: 3
+  POSE_LR: 0.0002
+  PRETRAIN_LIFTER: false
+  SCALE_LOSS_WEIGHTS:
+  - 0.001
+  - 1.0
+  SUBNET_CRITICS: 1
+  TEMP_LR: 0.0002
+  USE_CYCLE: false
+  USE_NEW_ROT: false
+  USE_NEW_TEMP: false
+  USE_SCALER: false
+USE_GT: true

insightface/body/human_pose/ambiguity_aware/cfg/h36m_gt_scale.yaml

@@ -0,0 +1,40 @@
+BATCH_SIZE: 512
+DATA:
+  NUM_FRAMES: 1
+  SCALE_MID_MEAN: 0.720643
+  SCALE_MID_STD: 0.058
+  USE_RANDOM_DIFF: true
+NETWORK:
+  DIS_RES_BLOCKS: 2
+  DIS_TEMP_RES_BLOCKS: 2
+  DIS_USE_SPECTRAL_NORM: false
+  SCALER_INPUT_SIZE: 34
+TRAIN:
+  BOUND_AZIM: 2.44346
+  BOUND_ELEV: 0.34906585
+  DIS_LR: 0.0001
+  LOSS_TYPE: ss_adv
+  LOSS_WEIGHTS:
+  - 0.5
+  - 5.0
+  - 1.0
+  - 1.0
+  MAINNET_CRITICS: 4
+  NUM_CRITICS: 3
+  NUM_CRITICS_TEMP: 3
+  POSE_LR: 0.00015
+  PRETRAIN_LIFTER: false
+  SCALE_LOSS_WEIGHTS:
+  - 0.001
+  - 1.0
+  SUBNET_CRITICS: 1
+  TEMP_LR: 0.0001
+  SCHEDULER_STEP_SIZE: 5
+  USE_CYCLE: true
+  USE_NEW_ROT: false
+  USE_NEW_TEMP: true
+  USE_SCALER: true
+USE_GT: true
+FIX:
+  FIX_TRAJ: true 
+  FIX_TRAJ_BY_ROT: false

insightface/body/human_pose/ambiguity_aware/cfg/pre_adv.yaml

@@ -0,0 +1,36 @@
+BATCH_SIZE: 512
+DATA:
+  NUM_FRAMES: 1
+  SCALE_MID_MEAN: 0.720643
+  SCALE_MID_STD: 0.058
+  USE_RANDOM_DIFF: true
+NETWORK:
+  DIS_RES_BLOCKS: 2
+  DIS_TEMP_RES_BLOCKS: 2
+  DIS_USE_SPECTRAL_NORM: false
+  SCALER_INPUT_SIZE: 34
+TRAIN:
+  BOUND_AZIM: 2.44346
+  BOUND_ELEV: 0.34906585
+  DIS_LR: 0.0001
+  LOSS_TYPE: ss_adv
+  LOSS_WEIGHTS:
+  - 0.5
+  - 5.0
+  - 1.0
+  - 1.0
+  MAINNET_CRITICS: 4
+  NUM_CRITICS: 3
+  NUM_CRITICS_TEMP: 3
+  POSE_LR: 0.0001
+  PRETRAIN_LIFTER: false
+  SCALE_LOSS_WEIGHTS:
+  - 0.001
+  - 1.0
+  SUBNET_CRITICS: 1
+  TEMP_LR: 0.0002
+  USE_CYCLE: false
+  USE_NEW_ROT: false
+  USE_NEW_TEMP: false
+  USE_SCALER: false
+USE_GT: false

insightface/body/human_pose/ambiguity_aware/cfg/pre_tmc_klbone.yaml

@@ -0,0 +1,39 @@
+BATCH_SIZE: 512
+DATA:
+  EXP_TMC: true
+  EXP_TMC_DETERMINISTIC: true
+  EXP_TMC_INTERVAL: 3
+  NUM_FRAMES: 1
+  SCALE_MID_MEAN: 0.720643
+  SCALE_MID_STD: 0.058
+  USE_RANDOM_DIFF: true
+NETWORK:
+  DIS_RES_BLOCKS: 2
+  DIS_TEMP_RES_BLOCKS: 2
+  DIS_USE_SPECTRAL_NORM: false
+  SCALER_INPUT_SIZE: 34
+TRAIN:
+  BOUND_AZIM: 2.44346
+  BOUND_ELEV: 0.34906585
+  DIS_LR: 0.0001
+  LOSS_TYPE: ss_adv
+  LOSS_WEIGHTS:
+  - 0.5
+  - 5.0
+  - 1.0
+  - 1.0
+  MAINNET_CRITICS: 4
+  NUM_CRITICS: 3
+  NUM_CRITICS_TEMP: 3
+  POSE_LR: 0.0001
+  PRETRAIN_LIFTER: false
+  SCALE_LOSS_WEIGHTS:
+  - 0.01
+  - 1.0
+  SUBNET_CRITICS: 1
+  TEMP_LR: 0.0002
+  USE_CYCLE: true
+  USE_NEW_ROT: false
+  USE_NEW_TEMP: false
+  USE_SCALER: true
+USE_GT: false

insightface/body/human_pose/ambiguity_aware/requirements.txt

@@ -0,0 +1,12 @@
+torch==1.13.1
+torchvision==0.5.0
+tqdm 
+opencv-python 
+scikit-learn 
+matplotlib
+h5py 
+pyyaml
+seaborn 
+imageio 
+easydict
+tensorboardX

insightface/body/human_pose/ambiguity_aware/scripts/_init_paths.py

@@ -0,0 +1,13 @@
+import os.path as osp
+import sys
+
+
+def add_path(path):
+    if path not in sys.path:
+        sys.path.insert(0, path)
+
+
+this_dir = osp.dirname(__file__)
+
+lib_path = osp.join(this_dir, '..')
+add_path(lib_path)

insightface/body/human_pose/ambiguity_aware/scripts/demo.sh

@@ -0,0 +1,5 @@
+#!/bin/bash
+
+rm -rf demo_output
+
+python inference.py --indir demo_input --outdir demo_output --cfg ../cfg/h36m_gt_scale.yaml --pretrain ../models/tmc_klbone.pth.tar

insightface/body/human_pose/ambiguity_aware/scripts/demo_input/0.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:15e90d8305cb0946a6f3c08c625dad617032102e520eb9c5d4c17d0af1609482
+size 513

insightface/body/human_pose/ambiguity_aware/scripts/demo_input/1.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6aa3178631c3dd2613b3def95a2c17e284bb262ee62cefdf647fe768a8c6efc6
+size 513

insightface/body/human_pose/ambiguity_aware/scripts/demo_input/10.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:01f9a4cc1b06719e65b4af28a3ac62e27afaadacc1028c9ae8e8e96d5724e23e
+size 513

insightface/body/human_pose/ambiguity_aware/scripts/demo_input/11.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:88c1b12e23512a66e8cb7802c06b447924ad5efae06098caf9de42ec996cd1ef
+size 513