diff --git a/.gitignore b/.gitignore index 62acda3a..40a29305 100644 --- a/.gitignore +++ b/.gitignore @@ -12,3 +12,5 @@ test.py .cache .ipython .local +uv.lock +.venv/ diff --git a/.python-version b/.python-version new file mode 100644 index 00000000..c8cfe395 --- /dev/null +++ b/.python-version @@ -0,0 +1 @@ +3.10 diff --git a/README_cn.md b/README_cn.md deleted file mode 100644 index b11c9eca..00000000 --- a/README_cn.md +++ /dev/null @@ -1,282 +0,0 @@ -# 使用 Pytorch 进行人脸识别 - -*Click [here](README.md) to return to the English document* - -> 译者注： -> -> 本项目 [facenet-pytorch](https://github.com/timesler/facenet-pytorch) 是一个十分方便的人脸识别库，可以通过 [pip](https://pypi.org/project/facenet-pytorch/) 直接安装。 -> -> 库中包含了两个重要功能 -> -> - 人脸检测：使用MTCNN算法 -> - 人脸识别：使用FaceNet算法 -> -> 利用这个库，可以轻松实现人脸检测和人脸向量映射操作。 -> -> 为了方便中文开发者研究学习人脸识别相关任务、贡献代码，我将本项目的README文件以及位于 `examples` 里面的几个示例脚本中必要的部分翻译成了中文，以供参考。 -> -> 向本项目的所有贡献者致敬。 -> -> 英译汉：[远哥挺乐](https://github.com/yuan2001425) -> -> Translator's Note: -> -> This project [facenet-pytorch](https://github.com/timesler/facenet-pytorch) is a very convenient face recognition library that can be installed directly via [pip](https://pypi.org/project/facenet-pytorch/). -> -> The library contains two important features: -> -> - Face detection: using the MTCNN algorithm -> - Face recognition: using the FaceNet algorithm -> -> With this library, one can easily carry out face detection and face vector mapping operations. -> -> In order to facilitate Chinese developers in studying face recognition and contributing code, I have translated the README file of this project and some necessary parts of several example scripts located in the `examples` directory into Chinese. -> -> Salute to all contributors to this project. -> -> Translated from English to Chinese by [远哥挺乐](https://github.com/yuan2001425). - -[![下载](https://pepy.tech/badge/facenet-pytorch)](https://pepy.tech/project/facenet-pytorch) - -[![代码覆盖率](https://img.shields.io/codecov/c/github/timesler/facenet-pytorch.svg)](https://codecov.io/gh/timesler/facenet-pytorch) - -|Python | 3.10 | 3.10 3.9 | 3.9 3.8 | -| :---: | :---: | :---: | :---: | -| 测试结果 | [![测试状态](https://github.com/timesler/facenet-pytorch/actions/workflows/python-3.10.yml/badge.svg?branch=master)](https://github.com/timesler/facenet-pytorch/actions?query=workflow%3A%22Python+3.10%22+branch%3Amaster) | [![测试状态](https://github.com/timesler/facenet-pytorch/actions/workflows/python-3.9.yml/badge.svg?branch=master)](https://github.com/timesler/facenet-pytorch/actions?query=workflow%3A%22Python+3.9%22+branch%3Amaster) | [![测试状态](https://github.com/timesler/facenet-pytorch/actions/workflows/python-3.8.yml/badge.svg?branch=master)](https://github.com/timesler/facenet-pytorch/actions?query=workflow%3A%22Python+3.8%22+branch%3Amaster) | - -[![xscode](https://img.shields.io/badge/Available%20on-xs%3Acode-blue?style=?style=plastic&logo=appveyor&logo=data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAEAAAABACAMAAACdt4HsAAAAGXRFWHRTb2Z0d2FyZQBBZG9iZSBJbWFnZVJlYWR5ccllPAAAAAZQTFRF////////VXz1bAAAAAJ0Uk5T/wDltzBKAAAAlUlEQVR42uzXSwqAMAwE0Mn9L+3Ggtgkk35QwcnSJo9S+yGwM9DCooCbgn4YrJ4CIPUcQF7/XSBbx2TEz4sAZ2q1RAECBAiYBlCtvwN+KiYAlG7UDGj59MViT9hOwEqAhYCtAsUZvL6I6W8c2wcbd+LIWSCHSTeSAAECngN4xxIDSK9f4B9t377Wd7H5Nt7/Xz8eAgwAvesLRjYYPuUAAAAASUVORK5CYII=)](https://xscode.com/timesler/facenet-pytorch) - -这是 pytorch 中 Inception Resnet (V1) 模型的存储库，在 VGGFace2 和 CASIA-Webface 上进行了预训练。 - -Pytorch 模型权重使用从 David Sandberg 的 [tensorflow Facenet repo](https://github.com/davidsandberg/facenet) 移植的参数进行初始化。 - -该存储库中还包含 MTCNN 的高效 pytorch 实现，用于推理之前的人脸检测。这些模型也是经过预训练的。据我们所知，这是最快的 MTCNN 实现。 - -## 目录 - -* [目录](#table-of-contents) -* [快速启动](#quick-start) -* [预训练模型](#pretrained-models) -* [示例笔记本](#example-notebooks) - + [*完整的检测和识别流程*](#complete-detection-and-recognition-pipeline) - + [*视频流中的人脸跟踪*](#face-tracking-in-video-streams) - + [*使用新数据微调预训练模型*](#finetuning-pretrained-models-with-new-data) - + [*facenet-pytorch 中的 MTCNN 指南*](#guide-to-mtcnn-in-facenet-pytorch) - + [*人脸检测包的性能比较*](#performance-comparison-of-face-detection-packages) - + [*FastMTCNN 算法*](#the-fastmtcnn-algorithm) -* [使用 docker 运行](#running-with-docker) -* [在您自己的 git 项目中使用此存储库](#use-this-repo-in-your-own-git-project) -* [Tensorflow 到 Pytorch 的参数转换](#conversion-of-parameters-from-tensorflow-to-pytorch) -* [参考资料](#references) - -## 快速启动 - -1. 安装： - -````bash -# 使用pip安装： -pip install facenet-pytorch - -# 或克隆此存储库，删除“-”以允许 python 导入： -git clone https://github.com/timesler/facenet-pytorch.git facenet_pytorch - -# 或使用 docker 容器（参见 https://github.com/timesler/docker-jupyter-dl-gpu）： -docker run -it --rm timesler/jupyter-dl-gpu pip install facenet-pytorch && ipython -```` - -2. 在python中，导入 facenet-pytorch 并实例化模型： - -````python -from facenet_pytorch import MTCNN, InceptionResnetV1 - -# 如果需要，使用 MTCNN 创建人脸检测模型： -mtcnn = MTCNN(image_size=, margin=) - -# 创建一个 inception resnet（在 eval 模式下）： -resnet = InceptionResnetV1(pretrained='vggface2').eval() -```` - -3. 处理图像： - -````python -from PIL import Image - -img = Image.open() - -# 获取裁剪和预白化的图像张量 -img_cropped = mtcnn(img, save_path=) - -# 计算嵌入（解压缩以添加批量维度） -img_embedding = resnet(img_cropped.unsqueeze(0)) - -# 或者，如果用于 VGGFace2 分类 -resnet.classify = True -img_probs = resnet(img_cropped.unsqueeze(0)) -```` - -有关使用和实现详细信息，请参阅 `help(MTCNN)` 和 `help(InceptionResnetV1)` 。 - -## 预训练模型 - -请参阅：[models/inception_resnet_v1.py](models/inception_resnet_v1.py) - -以下模型已移植到 pytorch（包含下载 pytorch state_dict 的链接）： - -|模型名称|LFW 准确度（如[此处](https://github.com/davidsandberg/facenet)列出）|训练数据集| -| :- | :-: | -: | -|[20180408-102900](https://github.com/timesler/facenet-pytorch/releases/download/v2.2.9/20180408-102900-casia-webface.pt) (111MB)|0.9905|CASIA-Webface| -|[20180402-114759](https://github.com/timesler/facenet-pytorch/releases/download/v2.2.9/20180402-114759-vggface2.pt) (107MB)|0.9965|VGGFace2| - -无需手动下载预训练的state_dict；它们会在模型实例化时自动下载，并缓存在 torch 缓存中以供将来使用。要在 pytorch 中使用 Inception Resnet (V1) 模型进行面部识别/识别，请使用： - -````python -from facenet_pytorch import InceptionResnetV1 - -# 对于在 VGGFace2 上预训练的模型 -model = InceptionResnetV1(pretrained='vggface2').eval() - -# 对于在 CASIA-Webface 上预训练的模型 -model = InceptionResnetV1(pretrained='casia-webface').eval() - -# 对于具有 100 个类的未经训练的模型 -model = InceptionResnetV1(num_classes=100).eval() - -# 对于未经训练的 1001 类分类器 -model = InceptionResnetV1(classify=True, num_classes=1001).eval() -```` - -两个预训练模型均在 160x160 像素图像上进行训练，因此如果应用于调整为该形状的图像，则效果最佳。为了获得最佳结果，还应该使用 MTCNN 将图像裁剪到脸部（见下文）。 - -默认情况下，上述模型将返回 512 维图像嵌入。要启用分类，请将 `classify=True` 传递给模型构造函数，或者您可以随后使用 `model.classify = True` 设置对象属性。对于 VGGFace2，预训练模型将输出长度为 8631 的 logit 向量，对于 CASIA-Webface 则输出长度为 10575 的 logit 向量。 - -## 示例笔记本 - -### *完整的检测和识别流程* - -通过首先使用 MTCNN 检测人脸，然后使用 Inception Resnet 模型计算嵌入或概率，可以轻松地将人脸识别应用于原始图像。 [examples/infer_cn.ipynb](examples/infer_cn.ipynb) 中的示例代码提供了一个利用数据集、数据加载器和可选 GPU 处理的完整示例流程。 - -### *视频流中的人脸跟踪* - -MTCNN 可用于构建人脸跟踪系统（使用 `MTCNN.detect()` 方法）。完整的面部跟踪示例可以在 [examples/face_tracking_cn.ipynb](examples/face_tracking_cn.ipynb) 中找到。 - -![](examples/tracked.gif) - -### *使用新数据微调预训练模型* - -在大多数情况下，实现人脸识别的最佳方法是直接使用预训练模型，通过聚类算法或简单的距离度量来确定人脸的身份。但是，如果需要微调（即，如果您想根据模型的输出 logits 选择标识），可以在 [examples/finetune_cn.ipynb](examples/finetune_cn.ipynb) 中找到示例。 - -### *facenet-pytorch 中的 MTCNN 指南* - -本指南演示了 MTCNN 模块的功能。涵盖的主题有： - -* 基本用法 -* 图像标准化 -* 面边距 -* 单张图像中的多个面孔 -* 批量检测 -* 边界框和面部标志 -* 保存人脸数据集 - -请参阅[kaggle 笔记本](https://www.kaggle.com/timesler/guide-to-mtcnn-in-facenet-pytorch)。 - -### *人脸检测包的性能比较* - -本笔记本演示了三个人脸检测包的使用： - -1. facenet-pytorch -2. mtcnn -3. dlib - -每个包都经过测试，测试其在启用 GPU 支持的情况下检测一组 300 张图像（来自一个视频的所有帧）中的面部的速度。性能基于 Kaggle 的 P100 笔记本内核。结果总结如下。 - -|套餐|FPS (1080x1920)|FPS (720x1280)|FPS (540x960)| -|---|---|---|---| -|facenet-pytorch|12.97|20.32|25.50| -|facenet-pytorch（非批处理）|9.75|14.81|19.68| -|dlib|3.80|8.39|14.53| -|mtcnn|3.04|5.70|8.23| - -![](examples/performance-comparison.png) - -请参阅[kaggle 笔记本](https://www.kaggle.com/timesler/comparison-of-face-detection-packages)。 - -### *FastMTCNN 算法* - -该算法演示了如何通过利用相邻帧之间的相似性来实现极其高效的人脸检测，特别是在视频中。 - -请参阅[kaggle 笔记本](https://www.kaggle.com/timesler/fast-mtcnn- detector-55-fps-at-full-resolution)。 - -## 使用 docker 运行 - -该包和任何示例笔记本都可以使用 docker（或 nvidia-docker）运行： - -````bash -docker run --rm -p 8888:8888 - -v ./facenet-pytorch:/home/jovyan timesler/jupyter-dl-gpu \ - -v :/home/jovyan/data - pip install facenet-pytorch && jupyter lab -```` - -导航到 example/ 目录并运行任何 ipython 笔记本。 - -有关 docker 容器的详细信息，请参阅 [timesler/jupyter-dl-gpu](https://github.com/timesler/docker-jupyter-dl-gpu)。 - -## 在您自己的 git 项目中使用此存储库 - -要在您自己的 git 存储库中使用此代码，我建议首先将此存储库添加为子模块。请注意，当克隆为子模块时，应删除存储库名称中的破折号（“-”），因为它会在导入时破坏 python： - -`git submodule add https://github.com/timesler/facenet-pytorch.git facenet_pytorch` - -或者，可以使用 pip 将代码安装为包： - -`pip install facenet-pytorch` - -## Tensorflow 到 Pytorch 的参数转换 - -请参阅：[models/utils/tensorflow2pytorch.py](models/tensorflow2pytorch.py) - -请注意，使用此存储库中的模型不需要此功能，该功能仅依赖于pytorch保存的 `state_dict`。 - -实例化 pytorch 模型后，每层的权重均从 [davidsandberg/facenet](https://github.com/davidsandberg/facenet) 的预训练Tensorflow模型中的等效层加载。 - -原始Tensorflow模型和 pytorch 移植模型的输出的等效性已经过测试并且是相同的： - ---- - -`>>> compare_model_outputs(mdl, sess, torch.randn(5, 160, 160, 3).detach())` - -```` -Passing test data through TF model （通过TF模型传递测试数据） - -tensor([[-0.0142, 0.0615, 0.0057, ..., 0.0497, 0.0375, -0.0838], - [-0.0139, 0.0611, 0.0054, ..., 0.0472, 0.0343, -0.0850], - [-0.0238, 0.0619, 0.0124, ..., 0.0598, 0.0334, -0.0852], - [-0.0089, 0.0548, 0.0032, ..., 0.0506, 0.0337, -0.0881], - [-0.0173, 0.0630, -0.0042, ..., 0.0487, 0.0295, -0.0791]]) - -Passing test data through PT model （通过PT模型传递测试数据） - -tensor([[-0.0142, 0.0615, 0.0057, ..., 0.0497, 0.0375, -0.0838], - [-0.0139, 0.0611, 0.0054, ..., 0.0472, 0.0343, -0.0850], - [-0.0238, 0.0619, 0.0124, ..., 0.0598, 0.0334, -0.0852], - [-0.0089, 0.0548, 0.0032, ..., 0.0506, 0.0337, -0.0881], - [-0.0173, 0.0630, -0.0042, ..., 0.0487, 0.0295, -0.0791]], - grad_fn=) - -Distance 1.2874517096861382e-06 （距离1.2874517096861382e-06） -```` - ---- - -为了重新运行Tensorflow参数到 pytorch 模型的转换，请确保使用子模块克隆此存储库，因为 davidsandberg/facenet 存储库作为子模块包含在内，并且转换需要其中的一部分。 - -## 参考资料 - -1. David Sandberg's facenet repo: [https://github.com/davidsandberg/facenet](https://github.com/davidsandberg/facenet) -2. F. Schroff, D. Kalenichenko, J. Philbin. _FaceNet: A Unified Embedding for Face Recognition and Clustering_, arXiv:1503.03832, 2015. [PDF](https://arxiv.org/pdf/1503.03832) - -3. Q. Cao, L. Shen, W. Xie, O. M. Parkhi, A. Zisserman. _VGGFace2: A dataset for recognising face across pose and age_, International Conference on Automatic Face and Gesture Recognition, 2018. [PDF](http://www.robots.ox.ac.uk/~vgg/publications/2018/Cao18/cao18.pdf) - -4. D. Yi, Z. Lei, S. Liao and S. Z. Li. _CASIAWebface: Learning Face Representation from Scratch_, arXiv:1411.7923, 2014. [PDF](https://arxiv.org/pdf/1411.7923) - -5. K. Zhang, Z. Zhang, Z. Li and Y. Qiao. _Joint Face Detection and Alignment Using Multitask Cascaded Convolutional Networks_, IEEE Signal Processing Letters, 2016. [PDF](https://kpzhang93.github.io/MTCNN_face_detection_alignment/paper/spl.pdf) diff --git a/__init__.py b/__init__.py index 2c00284f..d84af8bc 100644 --- a/__init__.py +++ b/__init__.py @@ -1,11 +1,19 @@ from .models.inception_resnet_v1 import InceptionResnetV1 -from .models.mtcnn import MTCNN, PNet, RNet, ONet, prewhiten, fixed_image_standardization +from .models.mtcnn import ( + MTCNN, + PNet, + RNet, + ONet, + prewhiten, + fixed_image_standardization, +) from .models.utils.detect_face import extract_face from .models.utils import training import warnings + warnings.filterwarnings( - action="ignore", - message="This overload of nonzero is deprecated:\n\tnonzero()", - category=UserWarning -) \ No newline at end of file + action="ignore", + message="This overload of nonzero is deprecated:\n\tnonzero()", + category=UserWarning, +) diff --git a/models/inception_resnet_v1.py b/models/inception_resnet_v1.py index 3b7de0dd..532cb95d 100644 --- a/models/inception_resnet_v1.py +++ b/models/inception_resnet_v1.py @@ -10,19 +10,21 @@ class BasicConv2d(nn.Module): - def __init__(self, in_planes, out_planes, kernel_size, stride, padding=0): super().__init__() self.conv = nn.Conv2d( - in_planes, out_planes, - kernel_size=kernel_size, stride=stride, - padding=padding, bias=False - ) # verify bias false + in_planes, + out_planes, + kernel_size=kernel_size, + stride=stride, + padding=padding, + bias=False, + ) # verify bias false self.bn = nn.BatchNorm2d( out_planes, - eps=0.001, # value found in tensorflow - momentum=0.1, # default pytorch value - affine=True + eps=0.001, # value found in tensorflow + momentum=0.1, # default pytorch value + affine=True, ) self.relu = nn.ReLU(inplace=False) @@ -34,7 +36,6 @@ def forward(self, x): class Block35(nn.Module): - def __init__(self, scale=1.0): super().__init__() @@ -44,13 +45,13 @@ def __init__(self, scale=1.0): self.branch1 = nn.Sequential( BasicConv2d(256, 32, kernel_size=1, stride=1), - BasicConv2d(32, 32, kernel_size=3, stride=1, padding=1) + BasicConv2d(32, 32, kernel_size=3, stride=1, padding=1), ) self.branch2 = nn.Sequential( BasicConv2d(256, 32, kernel_size=1, stride=1), BasicConv2d(32, 32, kernel_size=3, stride=1, padding=1), - BasicConv2d(32, 32, kernel_size=3, stride=1, padding=1) + BasicConv2d(32, 32, kernel_size=3, stride=1, padding=1), ) self.conv2d = nn.Conv2d(96, 256, kernel_size=1, stride=1) @@ -68,7 +69,6 @@ def forward(self, x): class Block17(nn.Module): - def __init__(self, scale=1.0): super().__init__() @@ -78,8 +78,8 @@ def __init__(self, scale=1.0): self.branch1 = nn.Sequential( BasicConv2d(896, 128, kernel_size=1, stride=1), - BasicConv2d(128, 128, kernel_size=(1,7), stride=1, padding=(0,3)), - BasicConv2d(128, 128, kernel_size=(7,1), stride=1, padding=(3,0)) + BasicConv2d(128, 128, kernel_size=(1, 7), stride=1, padding=(0, 3)), + BasicConv2d(128, 128, kernel_size=(7, 1), stride=1, padding=(3, 0)), ) self.conv2d = nn.Conv2d(256, 896, kernel_size=1, stride=1) @@ -96,7 +96,6 @@ def forward(self, x): class Block8(nn.Module): - def __init__(self, scale=1.0, noReLU=False): super().__init__() @@ -107,8 +106,8 @@ def __init__(self, scale=1.0, noReLU=False): self.branch1 = nn.Sequential( BasicConv2d(1792, 192, kernel_size=1, stride=1), - BasicConv2d(192, 192, kernel_size=(1,3), stride=1, padding=(0,1)), - BasicConv2d(192, 192, kernel_size=(3,1), stride=1, padding=(1,0)) + BasicConv2d(192, 192, kernel_size=(1, 3), stride=1, padding=(0, 1)), + BasicConv2d(192, 192, kernel_size=(3, 1), stride=1, padding=(1, 0)), ) self.conv2d = nn.Conv2d(384, 1792, kernel_size=1, stride=1) @@ -127,7 +126,6 @@ def forward(self, x): class Mixed_6a(nn.Module): - def __init__(self): super().__init__() @@ -136,7 +134,7 @@ def __init__(self): self.branch1 = nn.Sequential( BasicConv2d(256, 192, kernel_size=1, stride=1), BasicConv2d(192, 192, kernel_size=3, stride=1, padding=1), - BasicConv2d(192, 256, kernel_size=3, stride=2) + BasicConv2d(192, 256, kernel_size=3, stride=2), ) self.branch2 = nn.MaxPool2d(3, stride=2) @@ -150,24 +148,23 @@ def forward(self, x): class Mixed_7a(nn.Module): - def __init__(self): super().__init__() self.branch0 = nn.Sequential( BasicConv2d(896, 256, kernel_size=1, stride=1), - BasicConv2d(256, 384, kernel_size=3, stride=2) + BasicConv2d(256, 384, kernel_size=3, stride=2), ) self.branch1 = nn.Sequential( BasicConv2d(896, 256, kernel_size=1, stride=1), - BasicConv2d(256, 256, kernel_size=3, stride=2) + BasicConv2d(256, 256, kernel_size=3, stride=2), ) self.branch2 = nn.Sequential( BasicConv2d(896, 256, kernel_size=1, stride=1), BasicConv2d(256, 256, kernel_size=3, stride=1, padding=1), - BasicConv2d(256, 256, kernel_size=3, stride=2) + BasicConv2d(256, 256, kernel_size=3, stride=2), ) self.branch3 = nn.MaxPool2d(3, stride=2) @@ -199,7 +196,15 @@ class InceptionResnetV1(nn.Module): initialized. (default: {None}) dropout_prob {float} -- Dropout probability. (default: {0.6}) """ - def __init__(self, pretrained=None, classify=False, num_classes=None, dropout_prob=0.6, device=None): + + def __init__( + self, + pretrained=None, + classify=False, + num_classes=None, + dropout_prob=0.6, + device=None, + ): super().__init__() # Set simple attributes @@ -207,13 +212,14 @@ def __init__(self, pretrained=None, classify=False, num_classes=None, dropout_pr self.classify = classify self.num_classes = num_classes - if pretrained == 'vggface2': + if pretrained == "vggface2": tmp_classes = 8631 - elif pretrained == 'casia-webface': + elif pretrained == "casia-webface": tmp_classes = 10575 elif pretrained is None and self.classify and self.num_classes is None: - raise Exception('If "pretrained" is not specified and "classify" is True, "num_classes" must be specified') - + raise Exception( + 'If "pretrained" is not specified and "classify" is True, "num_classes" must be specified' + ) # Define layers self.conv2d_1a = BasicConv2d(3, 32, kernel_size=3, stride=2) @@ -264,7 +270,7 @@ def __init__(self, pretrained=None, classify=False, num_classes=None, dropout_pr if self.classify and self.num_classes is not None: self.logits = nn.Linear(512, self.num_classes) - self.device = torch.device('cpu') + self.device = torch.device("cpu") if device is not None: self.device = device self.to(device) @@ -312,14 +318,16 @@ def load_weights(mdl, name): Raises: ValueError: If 'pretrained' not equal to 'vggface2' or 'casia-webface'. """ - if name == 'vggface2': - path = 'https://github.com/timesler/facenet-pytorch/releases/download/v2.2.9/20180402-114759-vggface2.pt' - elif name == 'casia-webface': - path = 'https://github.com/timesler/facenet-pytorch/releases/download/v2.2.9/20180408-102900-casia-webface.pt' + if name == "vggface2": + path = "https://github.com/timesler/facenet-pytorch/releases/download/v2.2.9/20180402-114759-vggface2.pt" + elif name == "casia-webface": + path = "https://github.com/timesler/facenet-pytorch/releases/download/v2.2.9/20180408-102900-casia-webface.pt" else: - raise ValueError('Pretrained models only exist for "vggface2" and "casia-webface"') + raise ValueError( + 'Pretrained models only exist for "vggface2" and "casia-webface"' + ) - model_dir = os.path.join(get_torch_home(), 'checkpoints') + model_dir = os.path.join(get_torch_home(), "checkpoints") os.makedirs(model_dir, exist_ok=True) cached_file = os.path.join(model_dir, os.path.basename(path)) @@ -333,8 +341,7 @@ def load_weights(mdl, name): def get_torch_home(): torch_home = os.path.expanduser( os.getenv( - 'TORCH_HOME', - os.path.join(os.getenv('XDG_CACHE_HOME', '~/.cache'), 'torch') + "TORCH_HOME", os.path.join(os.getenv("XDG_CACHE_HOME", "~/.cache"), "torch") ) ) return torch_home diff --git a/models/mtcnn.py b/models/mtcnn.py index 13b33a93..1b66a985 100644 --- a/models/mtcnn.py +++ b/models/mtcnn.py @@ -8,7 +8,7 @@ class PNet(nn.Module): """MTCNN PNet. - + Keyword Arguments: pretrained {bool} -- Whether or not to load saved pretrained weights (default: {True}) """ @@ -30,7 +30,7 @@ def __init__(self, pretrained=True): self.training = False if pretrained: - state_dict_path = os.path.join(os.path.dirname(__file__), '../data/pnet.pt') + state_dict_path = os.path.join(os.path.dirname(__file__), "../data/pnet.pt") state_dict = torch.load(state_dict_path) self.load_state_dict(state_dict) @@ -50,7 +50,7 @@ def forward(self, x): class RNet(nn.Module): """MTCNN RNet. - + Keyword Arguments: pretrained {bool} -- Whether or not to load saved pretrained weights (default: {True}) """ @@ -75,7 +75,7 @@ def __init__(self, pretrained=True): self.training = False if pretrained: - state_dict_path = os.path.join(os.path.dirname(__file__), '../data/rnet.pt') + state_dict_path = os.path.join(os.path.dirname(__file__), "../data/rnet.pt") state_dict = torch.load(state_dict_path) self.load_state_dict(state_dict) @@ -99,7 +99,7 @@ def forward(self, x): class ONet(nn.Module): """MTCNN ONet. - + Keyword Arguments: pretrained {bool} -- Whether or not to load saved pretrained weights (default: {True}) """ @@ -128,7 +128,7 @@ def __init__(self, pretrained=True): self.training = False if pretrained: - state_dict_path = os.path.join(os.path.dirname(__file__), '../data/onet.pt') + state_dict_path = os.path.join(os.path.dirname(__file__), "../data/onet.pt") state_dict = torch.load(state_dict_path) self.load_state_dict(state_dict) @@ -163,10 +163,10 @@ class MTCNN(nn.Module): - numpy.ndarray (uint8) representing either a single image (3D) or a batch of images (4D). Cropped faces can optionally be saved to file also. - + Keyword Arguments: image_size {int} -- Output image size in pixels. The image will be square. (default: {160}) - margin {int} -- Margin to add to bounding box, in terms of pixels in the final image. + margin {int} -- Margin to add to bounding box, in terms of pixels in the final image. Note that the application of the margin differs slightly from the davidsandberg/facenet repo, which applies the margin to the original image before resizing, making the margin dependent on the original image size (this is a bug in davidsandberg/facenet). @@ -195,9 +195,17 @@ class MTCNN(nn.Module): """ def __init__( - self, image_size=160, margin=0, min_face_size=20, - thresholds=[0.6, 0.7, 0.7], factor=0.709, post_process=True, - select_largest=True, selection_method=None, keep_all=False, device=None + self, + image_size=160, + margin=0, + min_face_size=20, + thresholds=[0.6, 0.7, 0.7], + factor=0.709, + post_process=True, + select_largest=True, + selection_method=None, + keep_all=False, + device=None, ): super().__init__() @@ -215,22 +223,22 @@ def __init__( self.rnet = RNet() self.onet = ONet() - self.device = torch.device('cpu') + self.device = torch.device("cpu") if device is not None: self.device = device self.to(device) if not self.selection_method: - self.selection_method = 'largest' if self.select_largest else 'probability' + self.selection_method = "largest" if self.select_largest else "probability" def forward(self, img, save_path=None, return_prob=False): """Run MTCNN face detection on a PIL image or numpy array. This method performs both detection and extraction of faces, returning tensors representing detected faces rather than the bounding boxes. To access bounding boxes, see the MTCNN.detect() method below. - + Arguments: img {PIL.Image, np.ndarray, or list} -- A PIL image, np.ndarray, torch.Tensor, or list. - + Keyword Arguments: save_path {str} -- An optional save path for the cropped image. Note that when self.post_process=True, although the returned tensor is post processed, the saved @@ -239,13 +247,13 @@ def forward(self, img, save_path=None, return_prob=False): (default: {None}) return_prob {bool} -- Whether or not to return the detection probability. (default: {False}) - + Returns: Union[torch.Tensor, tuple(torch.tensor, float)] -- If detected, cropped image of a face with dimensions 3 x image_size x image_size. Optionally, the probability that a face was detected. If self.keep_all is True, n detected faces are returned in an n x 3 x image_size x image_size tensor with an optional list of detection - probabilities. If `img` is a list of images, the item(s) returned have an extra + probabilities. If `img` is a list of images, the item(s) returned have an extra dimension (batch) as the first dimension. Example: @@ -259,7 +267,11 @@ def forward(self, img, save_path=None, return_prob=False): # Select faces if not self.keep_all: batch_boxes, batch_probs, batch_points = self.select_boxes( - batch_boxes, batch_probs, batch_points, img, method=self.selection_method + batch_boxes, + batch_probs, + batch_points, + img, + method=self.selection_method, ) # Extract faces faces = self.extract(img, batch_boxes, save_path) @@ -276,14 +288,14 @@ def detect(self, img, landmarks=False): that require lower-level handling of bounding boxes and facial landmarks (e.g., face tracking). The functionality of the forward function can be emulated by using this method followed by the extract_face() function. - + Arguments: img {PIL.Image, np.ndarray, or list} -- A PIL image, np.ndarray, torch.Tensor, or list. Keyword Arguments: landmarks {bool} -- Whether to return facial landmarks in addition to bounding boxes. (default: {False}) - + Returns: tuple(numpy.ndarray, list) -- For N detected faces, a tuple containing an Nx4 array of bounding boxes and a length N list of detection probabilities. @@ -311,10 +323,14 @@ def detect(self, img, landmarks=False): with torch.no_grad(): batch_boxes, batch_points = detect_face( - img, self.min_face_size, - self.pnet, self.rnet, self.onet, - self.thresholds, self.factor, - self.device + img, + self.min_face_size, + self.pnet, + self.rnet, + self.onet, + self.thresholds, + self.factor, + self.device, ) boxes, probs, points = [], [], [] @@ -326,7 +342,9 @@ def detect(self, img, landmarks=False): probs.append([None]) points.append(None) elif self.select_largest: - box_order = np.argsort((box[:, 2] - box[:, 0]) * (box[:, 3] - box[:, 1]))[::-1] + box_order = np.argsort( + (box[:, 2] - box[:, 0]) * (box[:, 3] - box[:, 1]) + )[::-1] box = box[box_order] point = point[box_order] boxes.append(box[:, :4]) @@ -341,9 +359,9 @@ def detect(self, img, landmarks=False): points = np.array(points, dtype=object) if ( - not isinstance(img, (list, tuple)) and - not (isinstance(img, np.ndarray) and len(img.shape) == 4) and - not (isinstance(img, torch.Tensor) and len(img.shape) == 4) + not isinstance(img, (list, tuple)) + and not (isinstance(img, np.ndarray) and len(img.shape) == 4) + and not (isinstance(img, torch.Tensor) and len(img.shape) == 4) ): boxes = boxes[0] probs = probs[0] @@ -355,8 +373,14 @@ def detect(self, img, landmarks=False): return boxes, probs def select_boxes( - self, all_boxes, all_probs, all_points, imgs, method='probability', threshold=0.9, - center_weight=2.0 + self, + all_boxes, + all_probs, + all_points, + imgs, + method="probability", + threshold=0.9, + center_weight=2.0, ): """Selects a single box from multiple for a given image using one of multiple heuristics. @@ -385,12 +409,12 @@ def select_boxes( for n images. Ix0 array of probabilities for each box, array of landmark points. """ - #copying batch detection from extract, but would be easier to ensure detect creates consistent output. + # copying batch detection from extract, but would be easier to ensure detect creates consistent output. batch_mode = True if ( - not isinstance(imgs, (list, tuple)) and - not (isinstance(imgs, np.ndarray) and len(imgs.shape) == 4) and - not (isinstance(imgs, torch.Tensor) and len(imgs.shape) == 4) + not isinstance(imgs, (list, tuple)) + and not (isinstance(imgs, np.ndarray) and len(imgs.shape) == 4) + and not (isinstance(imgs, torch.Tensor) and len(imgs.shape) == 4) ): imgs = [imgs] all_boxes = [all_boxes] @@ -400,33 +424,46 @@ def select_boxes( selected_boxes, selected_probs, selected_points = [], [], [] for boxes, points, probs, img in zip(all_boxes, all_points, all_probs, imgs): - + if boxes is None: selected_boxes.append(None) selected_probs.append([None]) selected_points.append(None) continue - + # If at least 1 box found boxes = np.array(boxes) probs = np.array(probs) points = np.array(points) - - if method == 'largest': - box_order = np.argsort((boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1]))[::-1] - elif method == 'probability': + + if method == "largest": + box_order = np.argsort( + (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1]) + )[::-1] + elif method == "probability": box_order = np.argsort(probs)[::-1] - elif method == 'center_weighted_size': + elif method == "center_weighted_size": box_sizes = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1]) - img_center = (img.width / 2, img.height/2) - box_centers = np.array(list(zip((boxes[:, 0] + boxes[:, 2]) / 2, (boxes[:, 1] + boxes[:, 3]) / 2))) + img_center = (img.width / 2, img.height / 2) + box_centers = np.array( + list( + zip( + (boxes[:, 0] + boxes[:, 2]) / 2, + (boxes[:, 1] + boxes[:, 3]) / 2, + ) + ) + ) offsets = box_centers - img_center offset_dist_squared = np.sum(np.power(offsets, 2.0), 1) - box_order = np.argsort(box_sizes - offset_dist_squared * center_weight)[::-1] - elif method == 'largest_over_threshold': + box_order = np.argsort(box_sizes - offset_dist_squared * center_weight)[ + ::-1 + ] + elif method == "largest_over_threshold": box_mask = probs > threshold boxes = boxes[box_mask] - box_order = np.argsort((boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1]))[::-1] + box_order = np.argsort( + (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1]) + )[::-1] if sum(box_mask) == 0: selected_boxes.append(None) selected_probs.append([None]) @@ -455,9 +492,9 @@ def extract(self, img, batch_boxes, save_path): # Determine if a batch or single image was passed batch_mode = True if ( - not isinstance(img, (list, tuple)) and - not (isinstance(img, np.ndarray) and len(img.shape) == 4) and - not (isinstance(img, torch.Tensor) and len(img.shape) == 4) + not isinstance(img, (list, tuple)) + and not (isinstance(img, np.ndarray) and len(img.shape) == 4) + and not (isinstance(img, torch.Tensor) and len(img.shape) == 4) ): img = [img] batch_boxes = [batch_boxes] @@ -485,7 +522,7 @@ def extract(self, img, batch_boxes, save_path): face_path = path_im if path_im is not None and i > 0: save_name, ext = os.path.splitext(path_im) - face_path = save_name + '_' + str(i + 1) + ext + face_path = save_name + "_" + str(i + 1) + ext face = extract_face(im, box, self.image_size, self.margin, face_path) if self.post_process: @@ -513,7 +550,6 @@ def fixed_image_standardization(image_tensor): def prewhiten(x): mean = x.mean() std = x.std() - std_adj = std.clamp(min=1.0/(float(x.numel())**0.5)) + std_adj = std.clamp(min=1.0 / (float(x.numel()) ** 0.5)) y = (x - mean) / std_adj return y - diff --git a/models/utils/detect_face.py b/models/utils/detect_face.py index 5d144864..57da98b9 100644 --- a/models/utils/detect_face.py +++ b/models/utils/detect_face.py @@ -13,21 +13,23 @@ except: pass + def fixed_batch_process(im_data, model): batch_size = 512 out = [] for i in range(0, len(im_data), batch_size): - batch = im_data[i:(i+batch_size)] + batch = im_data[i : (i + batch_size)] out.append(model(batch)) return tuple(torch.cat(v, dim=0) for v in zip(*out)) + def detect_face(imgs, minsize, pnet, rnet, onet, threshold, factor, device): if isinstance(imgs, (np.ndarray, torch.Tensor)): - if isinstance(imgs,np.ndarray): + if isinstance(imgs, np.ndarray): imgs = torch.as_tensor(imgs.copy(), device=device) - if isinstance(imgs,torch.Tensor): + if isinstance(imgs, torch.Tensor): imgs = torch.as_tensor(imgs, device=device) if len(imgs.shape) == 3: @@ -36,12 +38,12 @@ def detect_face(imgs, minsize, pnet, rnet, onet, threshold, factor, device): if not isinstance(imgs, (list, tuple)): imgs = [imgs] if any(img.size != imgs[0].size for img in imgs): - raise Exception("MTCNN batch processing only compatible with equal-dimension images.") + raise Exception( + "MTCNN batch processing only compatible with equal-dimension images." + ) imgs = np.stack([np.uint8(img) for img in imgs]) imgs = torch.as_tensor(imgs.copy(), device=device) - - model_dtype = next(pnet.parameters()).dtype imgs = imgs.permute(0, 3, 1, 2).type(model_dtype) @@ -71,8 +73,10 @@ def detect_face(imgs, minsize, pnet, rnet, onet, threshold, factor, device): im_data = imresample(imgs, (int(h * scale + 1), int(w * scale + 1))) im_data = (im_data - 127.5) * 0.0078125 reg, probs = pnet(im_data) - - boxes_scale, image_inds_scale = generateBoundingBox(reg, probs[:, 1], scale, threshold[0]) + + boxes_scale, image_inds_scale = generateBoundingBox( + reg, probs[:, 1], scale, threshold[0] + ) boxes.append(boxes_scale) image_inds.append(image_inds_scale) @@ -88,7 +92,6 @@ def detect_face(imgs, minsize, pnet, rnet, onet, threshold, factor, device): # NMS within each scale + image boxes, image_inds = boxes[scale_picks], image_inds[scale_picks] - # NMS within each image pick = batched_nms(boxes[:, :4], boxes[:, 4], image_inds, 0.7) boxes, image_inds = boxes[pick], image_inds[pick] @@ -102,13 +105,15 @@ def detect_face(imgs, minsize, pnet, rnet, onet, threshold, factor, device): boxes = torch.stack([qq1, qq2, qq3, qq4, boxes[:, 4]]).permute(1, 0) boxes = rerec(boxes) y, ey, x, ex = pad(boxes, w, h) - + # Second stage if len(boxes) > 0: im_data = [] for k in range(len(y)): if ey[k] > (y[k] - 1) and ex[k] > (x[k] - 1): - img_k = imgs[image_inds[k], :, (y[k] - 1):ey[k], (x[k] - 1):ex[k]].unsqueeze(0) + img_k = imgs[ + image_inds[k], :, (y[k] - 1) : ey[k], (x[k] - 1) : ex[k] + ].unsqueeze(0) im_data.append(imresample(img_k, (24, 24))) im_data = torch.cat(im_data, dim=0) im_data = (im_data - 127.5) * 0.0078125 @@ -137,11 +142,13 @@ def detect_face(imgs, minsize, pnet, rnet, onet, threshold, factor, device): im_data = [] for k in range(len(y)): if ey[k] > (y[k] - 1) and ex[k] > (x[k] - 1): - img_k = imgs[image_inds[k], :, (y[k] - 1):ey[k], (x[k] - 1):ex[k]].unsqueeze(0) + img_k = imgs[ + image_inds[k], :, (y[k] - 1) : ey[k], (x[k] - 1) : ex[k] + ].unsqueeze(0) im_data.append(imresample(img_k, (48, 48))) im_data = torch.cat(im_data, dim=0) im_data = (im_data - 127.5) * 0.0078125 - + # This is equivalent to out = onet(im_data) to avoid GPU out of memory. out = fixed_batch_process(im_data, onet) @@ -165,7 +172,7 @@ def detect_face(imgs, minsize, pnet, rnet, onet, threshold, factor, device): # NMS within each image using "Min" strategy # pick = batched_nms(boxes[:, :4], boxes[:, 4], image_inds, 0.7) - pick = batched_nms_numpy(boxes[:, :4], boxes[:, 4], image_inds, 0.7, 'Min') + pick = batched_nms_numpy(boxes[:, :4], boxes[:, 4], image_inds, 0.7, "Min") boxes, image_inds, points = boxes[pick], image_inds[pick], points[pick] boxes = boxes.cpu().numpy() @@ -180,7 +187,9 @@ def detect_face(imgs, minsize, pnet, rnet, onet, threshold, factor, device): batch_boxes.append(boxes[b_i_inds].copy()) batch_points.append(points[b_i_inds].copy()) - batch_boxes, batch_points = np.array(batch_boxes, dtype=object), np.array(batch_points, dtype=object) + batch_boxes, batch_points = np.array(batch_boxes, dtype=object), np.array( + batch_points, dtype=object + ) return batch_boxes, batch_points @@ -247,7 +256,7 @@ def nms_numpy(boxes, scores, threshold, method): h = np.maximum(0.0, yy2 - yy1 + 1).copy() inter = w * h - if method == 'Min': + if method == "Min": o = inter / np.minimum(area[i], area[idx]) else: o = inter / (area[i] + area[idx] - inter) @@ -292,7 +301,7 @@ def pad(boxes, w, h): def rerec(bboxA): h = bboxA[:, 3] - bboxA[:, 1] w = bboxA[:, 2] - bboxA[:, 0] - + l = torch.max(w, h) bboxA[:, 0] = bboxA[:, 0] + w * 0.5 - l * 0.5 bboxA[:, 1] = bboxA[:, 1] + h * 0.5 - l * 0.5 @@ -308,18 +317,20 @@ def imresample(img, sz): def crop_resize(img, box, image_size): if isinstance(img, np.ndarray): - img = img[box[1]:box[3], box[0]:box[2]] + img = img[box[1] : box[3], box[0] : box[2]] out = cv2.resize( - img, - (image_size, image_size), - interpolation=cv2.INTER_AREA + img, (image_size, image_size), interpolation=cv2.INTER_AREA ).copy() elif isinstance(img, torch.Tensor): - img = img[box[1]:box[3], box[0]:box[2]] - out = imresample( - img.permute(2, 0, 1).unsqueeze(0).float(), - (image_size, image_size) - ).byte().squeeze(0).permute(1, 2, 0) + img = img[box[1] : box[3], box[0] : box[2]] + out = ( + imresample( + img.permute(2, 0, 1).unsqueeze(0).float(), (image_size, image_size) + ) + .byte() + .squeeze(0) + .permute(1, 2, 0) + ) else: out = img.crop(box).copy().resize((image_size, image_size), Image.BILINEAR) return out @@ -341,17 +352,17 @@ def get_size(img): def extract_face(img, box, image_size=160, margin=0, save_path=None): """Extract face + margin from PIL Image given bounding box. - + Arguments: img {PIL.Image} -- A PIL Image. box {numpy.ndarray} -- Four-element bounding box. image_size {int} -- Output image size in pixels. The image will be square. - margin {int} -- Margin to add to bounding box, in terms of pixels in the final image. + margin {int} -- Margin to add to bounding box, in terms of pixels in the final image. Note that the application of the margin differs slightly from the davidsandberg/facenet repo, which applies the margin to the original image before resizing, making the margin dependent on the original image size. save_path {str} -- Save path for extracted face image. (default: {None}) - + Returns: torch.tensor -- tensor representing the extracted face. """ diff --git a/models/utils/download.py b/models/utils/download.py index 328684f2..2a269d94 100644 --- a/models/utils/download.py +++ b/models/utils/download.py @@ -7,7 +7,9 @@ from urllib.request import urlopen, Request try: - from tqdm.auto import tqdm # automatically select proper tqdm submodule if available + from tqdm.auto import ( + tqdm, + ) # automatically select proper tqdm submodule if available except ImportError: from tqdm import tqdm @@ -30,7 +32,7 @@ def download_url_to_file(url, dst, hash_prefix=None, progress=True): req = Request(url, headers={"User-Agent": "torch.hub"}) u = urlopen(req) meta = u.info() - if hasattr(meta, 'getheaders'): + if hasattr(meta, "getheaders"): content_length = meta.getheaders("Content-Length") else: content_length = meta.get_all("Content-Length") @@ -47,8 +49,13 @@ def download_url_to_file(url, dst, hash_prefix=None, progress=True): try: if hash_prefix is not None: sha256 = hashlib.sha256() - with tqdm(total=file_size, disable=not progress, - unit='B', unit_scale=True, unit_divisor=1024) as pbar: + with tqdm( + total=file_size, + disable=not progress, + unit="B", + unit_scale=True, + unit_divisor=1024, + ) as pbar: while True: buffer = u.read(8192) if len(buffer) == 0: @@ -61,9 +68,12 @@ def download_url_to_file(url, dst, hash_prefix=None, progress=True): f.close() if hash_prefix is not None: digest = sha256.hexdigest() - if digest[:len(hash_prefix)] != hash_prefix: - raise RuntimeError('invalid hash value (expected "{}", got "{}")' - .format(hash_prefix, digest)) + if digest[: len(hash_prefix)] != hash_prefix: + raise RuntimeError( + 'invalid hash value (expected "{}", got "{}")'.format( + hash_prefix, digest + ) + ) shutil.move(f.name, dst) finally: f.close() diff --git a/models/utils/tensorflow2pytorch.py b/models/utils/tensorflow2pytorch.py index 7e3644c4..a74ef3cb 100644 --- a/models/utils/tensorflow2pytorch.py +++ b/models/utils/tensorflow2pytorch.py @@ -13,22 +13,22 @@ def import_tf_params(tf_mdl_dir, sess): """Import tensorflow model from save directory. - + Arguments: tf_mdl_dir {str} -- Location of protobuf, checkpoint, meta files. sess {tensorflow.Session} -- Tensorflow session object. - + Returns: (list, list, list) -- Tuple of lists containing the layer names, parameter arrays as numpy ndarrays, parameter shapes. """ - print('\nLoading tensorflow model\n') + print("\nLoading tensorflow model\n") if callable(tf_mdl_dir): tf_mdl_dir(sess) else: facenet.load_model(tf_mdl_dir) - print('\nGetting model weights\n') + print("\nGetting model weights\n") tf_layers = tf.trainable_variables() tf_params = sess.run(tf_layers) @@ -36,10 +36,10 @@ def import_tf_params(tf_mdl_dir, sess): tf_layers = [l.name for l in tf_layers] if not callable(tf_mdl_dir): - path = os.path.join(tf_mdl_dir, 'layer_description.json') + path = os.path.join(tf_mdl_dir, "layer_description.json") else: - path = 'data/layer_description.json' - with open(path, 'w') as f: + path = "data/layer_description.json" + with open(path, "w") as f: json.dump({l: s for l, s in zip(tf_layers, tf_shapes)}, f) return tf_layers, tf_params, tf_shapes @@ -48,25 +48,27 @@ def import_tf_params(tf_mdl_dir, sess): def get_layer_indices(layer_lookup, tf_layers): """Giving a lookup of model layer attribute names and tensorflow variable names, find matching parameters. - + Arguments: layer_lookup {dict} -- Dictionary mapping pytorch attribute names to (partial) tensorflow variable names. Expects dict of the form {'attr': ['tf_name', ...]} where the '...'s are ignored. tf_layers {list} -- List of tensorflow variable names. - + Returns: list -- The input dictionary with the list of matching inds appended to each item. """ layer_inds = {} for name, value in layer_lookup.items(): - layer_inds[name] = value + [[i for i, n in enumerate(tf_layers) if value[0] in n]] + layer_inds[name] = value + [ + [i for i, n in enumerate(tf_layers) if value[0] in n] + ] return layer_inds def load_tf_batchNorm(weights, layer): """Load tensorflow weights into nn.BatchNorm object. - + Arguments: weights {list} -- Tensorflow parameters. layer {torch.nn.Module} -- nn.BatchNorm. @@ -79,28 +81,23 @@ def load_tf_batchNorm(weights, layer): def load_tf_conv2d(weights, layer, transpose=False): """Load tensorflow weights into nn.Conv2d object. - + Arguments: weights {list} -- Tensorflow parameters. layer {torch.nn.Module} -- nn.Conv2d. """ if isinstance(weights, list): if len(weights) == 2: - layer.bias.data = ( - torch.tensor(weights[1]) - .view(layer.bias.data.shape) - ) + layer.bias.data = torch.tensor(weights[1]).view(layer.bias.data.shape) weights = weights[0] - + if transpose: dim_order = (3, 2, 1, 0) else: dim_order = (3, 2, 0, 1) layer.weight.data = ( - torch.tensor(weights) - .permute(dim_order) - .view(layer.weight.data.shape) + torch.tensor(weights).permute(dim_order).view(layer.weight.data.shape) ) @@ -110,7 +107,7 @@ def load_tf_conv2d_trans(weights, layer): def load_tf_basicConv2d(weights, layer): """Load tensorflow weights into grouped Conv2d+BatchNorm object. - + Arguments: weights {list} -- Tensorflow parameters. layer {torch.nn.Module} -- Object containing Conv2d+BatchNorm. @@ -121,27 +118,23 @@ def load_tf_basicConv2d(weights, layer): def load_tf_linear(weights, layer): """Load tensorflow weights into nn.Linear object. - + Arguments: weights {list} -- Tensorflow parameters. layer {torch.nn.Module} -- nn.Linear. """ if isinstance(weights, list): if len(weights) == 2: - layer.bias.data = ( - torch.tensor(weights[1]) - .view(layer.bias.data.shape) - ) + layer.bias.data = torch.tensor(weights[1]).view(layer.bias.data.shape) weights = weights[0] layer.weight.data = ( - torch.tensor(weights) - .transpose(-1, 0) - .view(layer.weight.data.shape) + torch.tensor(weights).transpose(-1, 0).view(layer.weight.data.shape) ) # High-level parameter-loading functions: + def load_tf_block35(weights, layer): load_tf_basicConv2d(weights[:4], layer.branch0) load_tf_basicConv2d(weights[4:8], layer.branch1[0]) @@ -162,7 +155,7 @@ def load_tf_block17_8(weights, layer): def load_tf_mixed6a(weights, layer): if len(weights) != 16: - raise ValueError(f'Number of weight arrays ({len(weights)}) not equal to 16') + raise ValueError(f"Number of weight arrays ({len(weights)}) not equal to 16") load_tf_basicConv2d(weights[:4], layer.branch0) load_tf_basicConv2d(weights[4:8], layer.branch1[0]) load_tf_basicConv2d(weights[8:12], layer.branch1[1]) @@ -171,7 +164,7 @@ def load_tf_mixed6a(weights, layer): def load_tf_mixed7a(weights, layer): if len(weights) != 28: - raise ValueError(f'Number of weight arrays ({len(weights)}) not equal to 28') + raise ValueError(f"Number of weight arrays ({len(weights)}) not equal to 28") load_tf_basicConv2d(weights[:4], layer.branch0[0]) load_tf_basicConv2d(weights[4:8], layer.branch0[1]) load_tf_basicConv2d(weights[8:12], layer.branch1[0]) @@ -183,8 +176,10 @@ def load_tf_mixed7a(weights, layer): def load_tf_repeats(weights, layer, rptlen, subfun): if len(weights) % rptlen != 0: - raise ValueError(f'Number of weight arrays ({len(weights)}) not divisible by {rptlen}') - weights_split = [weights[i:i+rptlen] for i in range(0, len(weights), rptlen)] + raise ValueError( + f"Number of weight arrays ({len(weights)}) not divisible by {rptlen}" + ) + weights_split = [weights[i : i + rptlen] for i in range(0, len(weights), rptlen)] for i, w in enumerate(weights_split): subfun(w, getattr(layer, str(i))) @@ -204,7 +199,7 @@ def load_tf_repeat_3(weights, layer): def test_loaded_params(mdl, tf_params, tf_layers): """Check each parameter in a pytorch model for an equivalent parameter in a list of tensorflow variables. - + Arguments: mdl {torch.nn.Module} -- Pytorch model. tf_params {list} -- List of ndarrays representing tensorflow variables. @@ -214,67 +209,72 @@ def test_loaded_params(mdl, tf_params, tf_layers): for name, param in mdl.named_parameters(): pt_mean = param.data.mean() matching_inds = ((tf_means - pt_mean).abs() < 1e-8).nonzero() - print(f'{name} equivalent to {[tf_layers[i] for i in matching_inds]}') + print(f"{name} equivalent to {[tf_layers[i] for i in matching_inds]}") def compare_model_outputs(pt_mdl, sess, test_data): """Given some testing data, compare the output of pytorch and tensorflow models. - + Arguments: pt_mdl {torch.nn.Module} -- Pytorch model. sess {tensorflow.Session} -- Tensorflow session object. test_data {torch.Tensor} -- Pytorch tensor. """ - print('\nPassing test data through TF model\n') + print("\nPassing test data through TF model\n") if isinstance(sess, tf.Session): images_placeholder = tf.get_default_graph().get_tensor_by_name("input:0") - phase_train_placeholder = tf.get_default_graph().get_tensor_by_name("phase_train:0") + phase_train_placeholder = tf.get_default_graph().get_tensor_by_name( + "phase_train:0" + ) embeddings = tf.get_default_graph().get_tensor_by_name("embeddings:0") - feed_dict = {images_placeholder: test_data.numpy(), phase_train_placeholder: False} + feed_dict = { + images_placeholder: test_data.numpy(), + phase_train_placeholder: False, + } tf_output = torch.tensor(sess.run(embeddings, feed_dict=feed_dict)) else: tf_output = sess(test_data) print(tf_output) - print('\nPassing test data through PT model\n') + print("\nPassing test data through PT model\n") pt_output = pt_mdl(test_data.permute(0, 3, 1, 2)) print(pt_output) distance = (tf_output - pt_output).norm() - print(f'\nDistance {distance}\n') + print(f"\nDistance {distance}\n") def compare_mtcnn(pt_mdl, tf_fun, sess, ind, test_data): tf_mdls = tf_fun(sess) tf_mdl = tf_mdls[ind] - print('\nPassing test data through TF model\n') + print("\nPassing test data through TF model\n") tf_output = tf_mdl(test_data.numpy()) tf_output = [torch.tensor(out) for out in tf_output] - print('\n'.join([str(o.view(-1)[:10]) for o in tf_output])) + print("\n".join([str(o.view(-1)[:10]) for o in tf_output])) - print('\nPassing test data through PT model\n') + print("\nPassing test data through PT model\n") with torch.no_grad(): pt_output = pt_mdl(test_data.permute(0, 3, 2, 1)) pt_output = [torch.tensor(out) for out in pt_output] for i in range(len(pt_output)): if len(pt_output[i].shape) == 4: pt_output[i] = pt_output[i].permute(0, 3, 2, 1).contiguous() - print('\n'.join([str(o.view(-1)[:10]) for o in pt_output])) + print("\n".join([str(o.view(-1)[:10]) for o in pt_output])) distance = [(tf_o - pt_o).norm() for tf_o, pt_o in zip(tf_output, pt_output)] - print(f'\nDistance {distance}\n') + print(f"\nDistance {distance}\n") def load_tf_model_weights(mdl, layer_lookup, tf_mdl_dir, is_resnet=True, arg_num=None): """Load tensorflow parameters into a pytorch model. - + Arguments: mdl {torch.nn.Module} -- Pytorch model. layer_lookup {[type]} -- Dictionary mapping pytorch attribute names to (partial) tensorflow variable names, and a function suitable for loading weights. - Expects dict of the form {'attr': ['tf_name', function]}. + Expects dict of the form {'attr': ['tf_name', function]}. tf_mdl_dir {str} -- Location of protobuf, checkpoint, meta files. """ tf.reset_default_graph() @@ -283,7 +283,7 @@ def load_tf_model_weights(mdl, layer_lookup, tf_mdl_dir, is_resnet=True, arg_num layer_info = get_layer_indices(layer_lookup, tf_layers) for layer_name, info in layer_info.items(): - print(f'Loading {info[0]}/* into {layer_name}') + print(f"Loading {info[0]}/* into {layer_name}") weights = [tf_params[i] for i in info[2]] layer = getattr(mdl, layer_name) info[1](weights, layer) @@ -296,121 +296,121 @@ def load_tf_model_weights(mdl, layer_lookup, tf_mdl_dir, is_resnet=True, arg_num def tensorflow2pytorch(): lookup_inception_resnet_v1 = { - 'conv2d_1a': ['InceptionResnetV1/Conv2d_1a_3x3', load_tf_basicConv2d], - 'conv2d_2a': ['InceptionResnetV1/Conv2d_2a_3x3', load_tf_basicConv2d], - 'conv2d_2b': ['InceptionResnetV1/Conv2d_2b_3x3', load_tf_basicConv2d], - 'conv2d_3b': ['InceptionResnetV1/Conv2d_3b_1x1', load_tf_basicConv2d], - 'conv2d_4a': ['InceptionResnetV1/Conv2d_4a_3x3', load_tf_basicConv2d], - 'conv2d_4b': ['InceptionResnetV1/Conv2d_4b_3x3', load_tf_basicConv2d], - 'repeat_1': ['InceptionResnetV1/Repeat/block35', load_tf_repeat_1], - 'mixed_6a': ['InceptionResnetV1/Mixed_6a', load_tf_mixed6a], - 'repeat_2': ['InceptionResnetV1/Repeat_1/block17', load_tf_repeat_2], - 'mixed_7a': ['InceptionResnetV1/Mixed_7a', load_tf_mixed7a], - 'repeat_3': ['InceptionResnetV1/Repeat_2/block8', load_tf_repeat_3], - 'block8': ['InceptionResnetV1/Block8', load_tf_block17_8], - 'last_linear': ['InceptionResnetV1/Bottleneck/weights', load_tf_linear], - 'last_bn': ['InceptionResnetV1/Bottleneck/BatchNorm', load_tf_batchNorm], - 'logits': ['Logits', load_tf_linear], + "conv2d_1a": ["InceptionResnetV1/Conv2d_1a_3x3", load_tf_basicConv2d], + "conv2d_2a": ["InceptionResnetV1/Conv2d_2a_3x3", load_tf_basicConv2d], + "conv2d_2b": ["InceptionResnetV1/Conv2d_2b_3x3", load_tf_basicConv2d], + "conv2d_3b": ["InceptionResnetV1/Conv2d_3b_1x1", load_tf_basicConv2d], + "conv2d_4a": ["InceptionResnetV1/Conv2d_4a_3x3", load_tf_basicConv2d], + "conv2d_4b": ["InceptionResnetV1/Conv2d_4b_3x3", load_tf_basicConv2d], + "repeat_1": ["InceptionResnetV1/Repeat/block35", load_tf_repeat_1], + "mixed_6a": ["InceptionResnetV1/Mixed_6a", load_tf_mixed6a], + "repeat_2": ["InceptionResnetV1/Repeat_1/block17", load_tf_repeat_2], + "mixed_7a": ["InceptionResnetV1/Mixed_7a", load_tf_mixed7a], + "repeat_3": ["InceptionResnetV1/Repeat_2/block8", load_tf_repeat_3], + "block8": ["InceptionResnetV1/Block8", load_tf_block17_8], + "last_linear": ["InceptionResnetV1/Bottleneck/weights", load_tf_linear], + "last_bn": ["InceptionResnetV1/Bottleneck/BatchNorm", load_tf_batchNorm], + "logits": ["Logits", load_tf_linear], } - print('\nLoad VGGFace2-trained weights and save\n') + print("\nLoad VGGFace2-trained weights and save\n") mdl = InceptionResnetV1(num_classes=8631).eval() - tf_mdl_dir = 'data/20180402-114759' - data_name = 'vggface2' + tf_mdl_dir = "data/20180402-114759" + data_name = "vggface2" load_tf_model_weights(mdl, lookup_inception_resnet_v1, tf_mdl_dir) state_dict = mdl.state_dict() - torch.save(state_dict, f'{tf_mdl_dir}-{data_name}.pt') + torch.save(state_dict, f"{tf_mdl_dir}-{data_name}.pt") torch.save( { - 'logits.weight': state_dict['logits.weight'], - 'logits.bias': state_dict['logits.bias'], + "logits.weight": state_dict["logits.weight"], + "logits.bias": state_dict["logits.bias"], }, - f'{tf_mdl_dir}-{data_name}-logits.pt' + f"{tf_mdl_dir}-{data_name}-logits.pt", ) - state_dict.pop('logits.weight') - state_dict.pop('logits.bias') - torch.save(state_dict, f'{tf_mdl_dir}-{data_name}-features.pt') - - print('\nLoad CASIA-Webface-trained weights and save\n') + state_dict.pop("logits.weight") + state_dict.pop("logits.bias") + torch.save(state_dict, f"{tf_mdl_dir}-{data_name}-features.pt") + + print("\nLoad CASIA-Webface-trained weights and save\n") mdl = InceptionResnetV1(num_classes=10575).eval() - tf_mdl_dir = 'data/20180408-102900' - data_name = 'casia-webface' + tf_mdl_dir = "data/20180408-102900" + data_name = "casia-webface" load_tf_model_weights(mdl, lookup_inception_resnet_v1, tf_mdl_dir) state_dict = mdl.state_dict() - torch.save(state_dict, f'{tf_mdl_dir}-{data_name}.pt') + torch.save(state_dict, f"{tf_mdl_dir}-{data_name}.pt") torch.save( { - 'logits.weight': state_dict['logits.weight'], - 'logits.bias': state_dict['logits.bias'], + "logits.weight": state_dict["logits.weight"], + "logits.bias": state_dict["logits.bias"], }, - f'{tf_mdl_dir}-{data_name}-logits.pt' + f"{tf_mdl_dir}-{data_name}-logits.pt", ) - state_dict.pop('logits.weight') - state_dict.pop('logits.bias') - torch.save(state_dict, f'{tf_mdl_dir}-{data_name}-features.pt') - + state_dict.pop("logits.weight") + state_dict.pop("logits.bias") + torch.save(state_dict, f"{tf_mdl_dir}-{data_name}-features.pt") + lookup_pnet = { - 'conv1': ['pnet/conv1', load_tf_conv2d_trans], - 'prelu1': ['pnet/PReLU1', load_tf_linear], - 'conv2': ['pnet/conv2', load_tf_conv2d_trans], - 'prelu2': ['pnet/PReLU2', load_tf_linear], - 'conv3': ['pnet/conv3', load_tf_conv2d_trans], - 'prelu3': ['pnet/PReLU3', load_tf_linear], - 'conv4_1': ['pnet/conv4-1', load_tf_conv2d_trans], - 'conv4_2': ['pnet/conv4-2', load_tf_conv2d_trans], + "conv1": ["pnet/conv1", load_tf_conv2d_trans], + "prelu1": ["pnet/PReLU1", load_tf_linear], + "conv2": ["pnet/conv2", load_tf_conv2d_trans], + "prelu2": ["pnet/PReLU2", load_tf_linear], + "conv3": ["pnet/conv3", load_tf_conv2d_trans], + "prelu3": ["pnet/PReLU3", load_tf_linear], + "conv4_1": ["pnet/conv4-1", load_tf_conv2d_trans], + "conv4_2": ["pnet/conv4-2", load_tf_conv2d_trans], } lookup_rnet = { - 'conv1': ['rnet/conv1', load_tf_conv2d_trans], - 'prelu1': ['rnet/prelu1', load_tf_linear], - 'conv2': ['rnet/conv2', load_tf_conv2d_trans], - 'prelu2': ['rnet/prelu2', load_tf_linear], - 'conv3': ['rnet/conv3', load_tf_conv2d_trans], - 'prelu3': ['rnet/prelu3', load_tf_linear], - 'dense4': ['rnet/conv4', load_tf_linear], - 'prelu4': ['rnet/prelu4', load_tf_linear], - 'dense5_1': ['rnet/conv5-1', load_tf_linear], - 'dense5_2': ['rnet/conv5-2', load_tf_linear], + "conv1": ["rnet/conv1", load_tf_conv2d_trans], + "prelu1": ["rnet/prelu1", load_tf_linear], + "conv2": ["rnet/conv2", load_tf_conv2d_trans], + "prelu2": ["rnet/prelu2", load_tf_linear], + "conv3": ["rnet/conv3", load_tf_conv2d_trans], + "prelu3": ["rnet/prelu3", load_tf_linear], + "dense4": ["rnet/conv4", load_tf_linear], + "prelu4": ["rnet/prelu4", load_tf_linear], + "dense5_1": ["rnet/conv5-1", load_tf_linear], + "dense5_2": ["rnet/conv5-2", load_tf_linear], } lookup_onet = { - 'conv1': ['onet/conv1', load_tf_conv2d_trans], - 'prelu1': ['onet/prelu1', load_tf_linear], - 'conv2': ['onet/conv2', load_tf_conv2d_trans], - 'prelu2': ['onet/prelu2', load_tf_linear], - 'conv3': ['onet/conv3', load_tf_conv2d_trans], - 'prelu3': ['onet/prelu3', load_tf_linear], - 'conv4': ['onet/conv4', load_tf_conv2d_trans], - 'prelu4': ['onet/prelu4', load_tf_linear], - 'dense5': ['onet/conv5', load_tf_linear], - 'prelu5': ['onet/prelu5', load_tf_linear], - 'dense6_1': ['onet/conv6-1', load_tf_linear], - 'dense6_2': ['onet/conv6-2', load_tf_linear], - 'dense6_3': ['onet/conv6-3', load_tf_linear], + "conv1": ["onet/conv1", load_tf_conv2d_trans], + "prelu1": ["onet/prelu1", load_tf_linear], + "conv2": ["onet/conv2", load_tf_conv2d_trans], + "prelu2": ["onet/prelu2", load_tf_linear], + "conv3": ["onet/conv3", load_tf_conv2d_trans], + "prelu3": ["onet/prelu3", load_tf_linear], + "conv4": ["onet/conv4", load_tf_conv2d_trans], + "prelu4": ["onet/prelu4", load_tf_linear], + "dense5": ["onet/conv5", load_tf_linear], + "prelu5": ["onet/prelu5", load_tf_linear], + "dense6_1": ["onet/conv6-1", load_tf_linear], + "dense6_2": ["onet/conv6-2", load_tf_linear], + "dense6_3": ["onet/conv6-3", load_tf_linear], } - print('\nLoad PNet weights and save\n') + print("\nLoad PNet weights and save\n") tf_mdl_dir = lambda sess: detect_face.create_mtcnn(sess, None) mdl = PNet() - data_name = 'pnet' + data_name = "pnet" load_tf_model_weights(mdl, lookup_pnet, tf_mdl_dir, is_resnet=False, arg_num=0) - torch.save(mdl.state_dict(), f'data/{data_name}.pt') + torch.save(mdl.state_dict(), f"data/{data_name}.pt") tf.reset_default_graph() with tf.Session() as sess: compare_mtcnn(mdl, tf_mdl_dir, sess, 0, torch.randn(1, 256, 256, 3).detach()) - print('\nLoad RNet weights and save\n') + print("\nLoad RNet weights and save\n") mdl = RNet() - data_name = 'rnet' + data_name = "rnet" load_tf_model_weights(mdl, lookup_rnet, tf_mdl_dir, is_resnet=False, arg_num=1) - torch.save(mdl.state_dict(), f'data/{data_name}.pt') + torch.save(mdl.state_dict(), f"data/{data_name}.pt") tf.reset_default_graph() with tf.Session() as sess: compare_mtcnn(mdl, tf_mdl_dir, sess, 1, torch.randn(1, 24, 24, 3).detach()) - print('\nLoad ONet weights and save\n') + print("\nLoad ONet weights and save\n") mdl = ONet() - data_name = 'onet' + data_name = "onet" load_tf_model_weights(mdl, lookup_onet, tf_mdl_dir, is_resnet=False, arg_num=2) - torch.save(mdl.state_dict(), f'data/{data_name}.pt') + torch.save(mdl.state_dict(), f"data/{data_name}.pt") tf.reset_default_graph() with tf.Session() as sess: compare_mtcnn(mdl, tf_mdl_dir, sess, 2, torch.randn(1, 48, 48, 3).detach()) diff --git a/models/utils/training.py b/models/utils/training.py index 20f81df9..34afc8aa 100644 --- a/models/utils/training.py +++ b/models/utils/training.py @@ -4,7 +4,6 @@ class Logger(object): - def __init__(self, mode, length, calculate_mean=False): self.mode = mode self.length = length @@ -15,18 +14,20 @@ def __init__(self, mode, length, calculate_mean=False): self.fn = lambda x, i: x def __call__(self, loss, metrics, i): - track_str = '\r{} | {:5d}/{:<5d}| '.format(self.mode, i + 1, self.length) - loss_str = 'loss: {:9.4f} | '.format(self.fn(loss, i)) - metric_str = ' | '.join('{}: {:9.4f}'.format(k, self.fn(v, i)) for k, v in metrics.items()) - print(track_str + loss_str + metric_str + ' ', end='') + track_str = "\r{} | {:5d}/{:<5d}| ".format(self.mode, i + 1, self.length) + loss_str = "loss: {:9.4f} | ".format(self.fn(loss, i)) + metric_str = " | ".join( + "{}: {:9.4f}".format(k, self.fn(v, i)) for k, v in metrics.items() + ) + print(track_str + loss_str + metric_str + " ", end="") if i + 1 == self.length: - print('') + print("") class BatchTimer(object): """Batch timing class. Use this class for tracking training and testing time/rate per batch or per sample. - + Keyword Arguments: rate {bool} -- Whether to report a rate (batches or samples per second) or a time (seconds per batch or sample). (default: {True}) @@ -60,17 +61,23 @@ def accuracy(logits, y): def pass_epoch( - model, loss_fn, loader, optimizer=None, scheduler=None, - batch_metrics={'time': BatchTimer()}, show_running=True, - device='cpu', writer=None + model, + loss_fn, + loader, + optimizer=None, + scheduler=None, + batch_metrics={"time": BatchTimer()}, + show_running=True, + device="cpu", + writer=None, ): """Train or evaluate over a data epoch. - + Arguments: model {torch.nn.Module} -- Pytorch model. loss_fn {callable} -- A function to compute (scalar) loss. loader {torch.utils.data.DataLoader} -- A pytorch data loader. - + Keyword Arguments: optimizer {torch.optim.Optimizer} -- A pytorch optimizer. scheduler {torch.optim.lr_scheduler._LRScheduler} -- LR scheduler (default: {None}) @@ -81,13 +88,13 @@ def pass_epoch( or rolling averages. (default: {False}) device {str or torch.device} -- Device for pytorch to use. (default: {'cpu'}) writer {torch.utils.tensorboard.SummaryWriter} -- Tensorboard SummaryWriter. (default: {None}) - + Returns: tuple(torch.Tensor, dict) -- A tuple of the average loss and a dictionary of average metric values across the epoch. """ - - mode = 'Train' if model.training else 'Valid' + + mode = "Train" if model.training else "Valid" logger = Logger(mode, length=len(loader), calculate_mean=show_running) loss = 0 metrics = {} @@ -106,39 +113,45 @@ def pass_epoch( metrics_batch = {} for metric_name, metric_fn in batch_metrics.items(): metrics_batch[metric_name] = metric_fn(y_pred, y).detach().cpu() - metrics[metric_name] = metrics.get(metric_name, 0) + metrics_batch[metric_name] - + metrics[metric_name] = ( + metrics.get(metric_name, 0) + metrics_batch[metric_name] + ) + if writer is not None and model.training: if writer.iteration % writer.interval == 0: - writer.add_scalars('loss', {mode: loss_batch.detach().cpu()}, writer.iteration) + writer.add_scalars( + "loss", {mode: loss_batch.detach().cpu()}, writer.iteration + ) for metric_name, metric_batch in metrics_batch.items(): - writer.add_scalars(metric_name, {mode: metric_batch}, writer.iteration) + writer.add_scalars( + metric_name, {mode: metric_batch}, writer.iteration + ) writer.iteration += 1 - + loss_batch = loss_batch.detach().cpu() loss += loss_batch if show_running: logger(loss, metrics, i_batch) else: logger(loss_batch, metrics_batch, i_batch) - + if model.training and scheduler is not None: scheduler.step() loss = loss / (i_batch + 1) metrics = {k: v / (i_batch + 1) for k, v in metrics.items()} - + if writer is not None and not model.training: - writer.add_scalars('loss', {mode: loss.detach()}, writer.iteration) + writer.add_scalars("loss", {mode: loss.detach()}, writer.iteration) for metric_name, metric in metrics.items(): writer.add_scalars(metric_name, {mode: metric}) return loss, metrics -def collate_pil(x): - out_x, out_y = [], [] - for xx, yy in x: - out_x.append(xx) - out_y.append(yy) - return out_x, out_y +def collate_pil(x): + out_x, out_y = [], [] + for xx, yy in x: + out_x.append(xx) + out_y.append(yy) + return out_x, out_y diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 00000000..62f62f8d --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,49 @@ +[build-system] +requires = ["setuptools>=61.0"] +build-backend = "setuptools.build_meta" + +[project] +name = "facenet-pytorch" +version = "2.5.3" +description = "Pretrained Pytorch face detection and recognition models" +readme = "README.md" +requires-python = ">=3.10" +authors = [ + { name = "Tim Esler", email = "tim.esler@gmail.com" }, + { name = "Vitthal Gupta", email = "vitthal.gupta@pi-labs.ai"} +] +license = { text = "MIT" } +urls = { Homepage = "https://github.com/timesler/facenet-pytorch" } + +dependencies = [ + "numpy>=1.24.0,<2.0.0", + "Pillow>=10.2.0,<10.3.0", + "requests>=2.0.0,<3.0.0", + "torch>=2.2.2", + "torchvision==0.22.1", + "tqdm>=4.0.0,<5.0.0", +] + +classifiers = [ + "Programming Language :: Python :: 3", + "License :: OSI Approved :: MIT License", + "Operating System :: OS Independent", +] + +[tool.setuptools] +package-dir = { "facenet_pytorch" = "." } + +[tool.setuptools.packages.find] +include = ["facenet_pytorch", "facenet_pytorch.models", "facenet_pytorch.models.utils", "facenet_pytorch.data"] + +[tool.setuptools.package-data] +facenet_pytorch = ["*net.pt"] + +[tool.uv] +link-mode = "copy" +compile-bytecode = true + +[dependency-groups] +dev = [ + "black>=25.1.0", +] diff --git a/setup.py b/setup.py deleted file mode 100644 index 1864989b..00000000 --- a/setup.py +++ /dev/null @@ -1,46 +0,0 @@ -import setuptools, os - -PACKAGE_NAME = 'facenet-pytorch' -VERSION = '2.5.2' -AUTHOR = 'Tim Esler' -EMAIL = 'tim.esler@gmail.com' -DESCRIPTION = 'Pretrained Pytorch face detection and recognition models' -GITHUB_URL = 'https://github.com/timesler/facenet-pytorch' - -parent_dir = os.path.dirname(os.path.realpath(__file__)) -import_name = os.path.basename(parent_dir) - -with open('{}/README.md'.format(parent_dir), 'r') as f: - long_description = f.read() - -setuptools.setup( - name=PACKAGE_NAME, - version=VERSION, - author=AUTHOR, - author_email=EMAIL, - description=DESCRIPTION, - long_description=long_description, - long_description_content_type='text/markdown', - url=GITHUB_URL, - packages=[ - 'facenet_pytorch', - 'facenet_pytorch.models', - 'facenet_pytorch.models.utils', - 'facenet_pytorch.data', - ], - package_dir={'facenet_pytorch':'.'}, - package_data={'': ['*net.pt']}, - classifiers=[ - "Programming Language :: Python :: 3", - "License :: OSI Approved :: MIT License", - "Operating System :: OS Independent", - ], - install_requires=[ - 'numpy>=1.24.0,<2.0.0', - 'Pillow>=10.2.0,<10.3.0', - 'requests>=2.0.0,<3.0.0', - 'torch>=2.2.0,<=2.3.0', - 'torchvision>=0.17.0,<=0.18.0', - 'tqdm>=4.0.0,<5.0.0', - ], -) diff --git a/tests/actions_test.py b/tests/actions_test.py index 520e4bd3..b96764c3 100644 --- a/tests/actions_test.py +++ b/tests/actions_test.py @@ -19,49 +19,49 @@ #### CLEAR ALL OUTPUT FILES #### -checkpoints = glob.glob(os.path.join(get_torch_home(), 'checkpoints/*')) +checkpoints = glob.glob(os.path.join(get_torch_home(), "checkpoints/*")) for c in checkpoints: - print('Removing {}'.format(c)) + print("Removing {}".format(c)) os.remove(c) -crop_files = glob.glob('data/test_images_aligned/**/*.png') +crop_files = glob.glob("data/test_images_aligned/**/*.png") for c in crop_files: - print('Removing {}'.format(c)) + print("Removing {}".format(c)) os.remove(c) #### TEST EXAMPLE IPYNB'S #### -os.system('jupyter nbconvert --to script --stdout examples/infer.ipynb examples/finetune.ipynb > examples/tmptest.py') -os.chdir('examples') +os.system( + "jupyter nbconvert --to script --stdout examples/infer.ipynb examples/finetune.ipynb > examples/tmptest.py" +) +os.chdir("examples") try: import examples.tmptest except: import tmptest -os.chdir('..') +os.chdir("..") #### TEST MTCNN #### + def get_image(path, trans): img = Image.open(path) img = trans(img) return img -trans = transforms.Compose([ - transforms.Resize(512) -]) -trans_cropped = transforms.Compose([ - np.float32, - transforms.ToTensor(), - fixed_image_standardization -]) +trans = transforms.Compose([transforms.Resize(512)]) + +trans_cropped = transforms.Compose( + [np.float32, transforms.ToTensor(), fixed_image_standardization] +) -dataset = datasets.ImageFolder('data/test_images', transform=trans) +dataset = datasets.ImageFolder("data/test_images", transform=trans) dataset.idx_to_class = {k: v for v, k in dataset.class_to_idx.items()} -mtcnn_pt = MTCNN(device=torch.device('cpu')) +mtcnn_pt = MTCNN(device=torch.device("cpu")) names = [] aligned = [] @@ -69,9 +69,11 @@ def get_image(path, trans): for img, idx in dataset: name = dataset.idx_to_class[idx] start = time() - img_align = mtcnn_pt(img, save_path='data/test_images_aligned/{}/1.png'.format(name)) - print('MTCNN time: {:6f} seconds'.format(time() - start)) - + img_align = mtcnn_pt( + img, save_path="data/test_images_aligned/{}/1.png".format(name) + ) + print("MTCNN time: {:6f} seconds".format(time() - start)) + # Comparison between types img_box = mtcnn_pt.detect(img)[0] assert (img_box - mtcnn_pt.detect(np.array(img))[0]).sum() < 1e-2 @@ -79,23 +81,29 @@ def get_image(path, trans): # Batching test assert (img_box - mtcnn_pt.detect([img, img])[0]).sum() < 1e-2 - assert (img_box - mtcnn_pt.detect(np.array([np.array(img), np.array(img)]))[0]).sum() < 1e-2 - assert (img_box - mtcnn_pt.detect(torch.as_tensor([np.array(img), np.array(img)]))[0]).sum() < 1e-2 + assert ( + img_box - mtcnn_pt.detect(np.array([np.array(img), np.array(img)]))[0] + ).sum() < 1e-2 + assert ( + img_box - mtcnn_pt.detect(torch.as_tensor([np.array(img), np.array(img)]))[0] + ).sum() < 1e-2 # Box selection - mtcnn_pt.selection_method = 'probability' - print('\nprobability - ', mtcnn_pt.detect(img)) - mtcnn_pt.selection_method = 'largest' - print('largest - ', mtcnn_pt.detect(img)) - mtcnn_pt.selection_method = 'largest_over_theshold' - print('largest_over_theshold - ', mtcnn_pt.detect(img)) - mtcnn_pt.selection_method = 'center_weighted_size' - print('center_weighted_size - ', mtcnn_pt.detect(img)) + mtcnn_pt.selection_method = "probability" + print("\nprobability - ", mtcnn_pt.detect(img)) + mtcnn_pt.selection_method = "largest" + print("largest - ", mtcnn_pt.detect(img)) + mtcnn_pt.selection_method = "largest_over_theshold" + print("largest_over_theshold - ", mtcnn_pt.detect(img)) + mtcnn_pt.selection_method = "center_weighted_size" + print("center_weighted_size - ", mtcnn_pt.detect(img)) if img_align is not None: names.append(name) aligned.append(img_align) - aligned_fromfile.append(get_image('data/test_images_aligned/{}/1.png'.format(name), trans_cropped)) + aligned_fromfile.append( + get_image("data/test_images_aligned/{}/1.png".format(name), trans_cropped) + ) aligned = torch.stack(aligned) aligned_fromfile = torch.stack(aligned_fromfile) @@ -109,42 +117,46 @@ def get_image(path, trans): [1.482895, 0.000000, 1.345686, 1.029880, 1.061939], [0.886342, 1.345686, 0.000000, 1.363125, 1.338803], [1.438450, 1.029880, 1.363125, 0.000000, 1.066040], - [1.437583, 1.061939, 1.338803, 1.066040, 0.000000] + [1.437583, 1.061939, 1.338803, 1.066040, 0.000000], ], [ [0.000000, 1.430769, 0.992931, 1.414197, 1.329544], [1.430769, 0.000000, 1.253911, 1.144899, 1.079755], [0.992931, 1.253911, 0.000000, 1.358875, 1.337322], [1.414197, 1.144899, 1.358875, 0.000000, 1.204118], - [1.329544, 1.079755, 1.337322, 1.204118, 0.000000] - ] + [1.329544, 1.079755, 1.337322, 1.204118, 0.000000], + ], ] -for i, ds in enumerate(['vggface2', 'casia-webface']): +for i, ds in enumerate(["vggface2", "casia-webface"]): resnet_pt = InceptionResnetV1(pretrained=ds).eval() start = time() embs = resnet_pt(aligned) - print('\nResnet time: {:6f} seconds\n'.format(time() - start)) + print("\nResnet time: {:6f} seconds\n".format(time() - start)) embs_fromfile = resnet_pt(aligned_fromfile) dists = [[(emb - e).norm().item() for e in embs] for emb in embs] - dists_fromfile = [[(emb - e).norm().item() for e in embs_fromfile] for emb in embs_fromfile] + dists_fromfile = [ + [(emb - e).norm().item() for e in embs_fromfile] for emb in embs_fromfile + ] - print('\nOutput:') + print("\nOutput:") print(pd.DataFrame(dists, columns=names, index=names)) - print('\nOutput (from file):') + print("\nOutput (from file):") print(pd.DataFrame(dists_fromfile, columns=names, index=names)) - print('\nExpected:') + print("\nExpected:") print(pd.DataFrame(expected[i], columns=names, index=names)) total_error = (torch.tensor(dists) - torch.tensor(expected[i])).norm() - total_error_fromfile = (torch.tensor(dists_fromfile) - torch.tensor(expected[i])).norm() + total_error_fromfile = ( + torch.tensor(dists_fromfile) - torch.tensor(expected[i]) + ).norm() - print('\nTotal error: {}, {}'.format(total_error, total_error_fromfile)) + print("\nTotal error: {}, {}".format(total_error, total_error_fromfile)) - if sys.platform != 'win32': + if sys.platform != "win32": assert total_error < 1e-2 assert total_error_fromfile < 1e-2 @@ -158,19 +170,19 @@ def get_image(path, trans): #### MULTI-FACE TEST #### mtcnn = MTCNN(keep_all=True) -img = Image.open('data/multiface.jpg') +img = Image.open("data/multiface.jpg") boxes, probs = mtcnn.detect(img) draw = ImageDraw.Draw(img) for i, box in enumerate(boxes): draw.rectangle(box.tolist()) -mtcnn(img, save_path='data/tmp.png') +mtcnn(img, save_path="data/tmp.png") #### MTCNN TYPES TEST #### -img = Image.open('data/multiface.jpg') +img = Image.open("data/multiface.jpg") mtcnn = MTCNN(keep_all=True) boxes_ref, _ = mtcnn.detect(img) @@ -180,9 +192,11 @@ def get_image(path, trans): boxes_test, _ = mtcnn.detect(img) _ = mtcnn(img) -box_diff = boxes_ref[np.argsort(boxes_ref[:,1])] - boxes_test[np.argsort(boxes_test[:,1])] +box_diff = ( + boxes_ref[np.argsort(boxes_ref[:, 1])] - boxes_test[np.argsort(boxes_test[:, 1])] +) total_error = np.sum(np.abs(box_diff)) -print('\nfp64 Total box error: {}'.format(total_error)) +print("\nfp64 Total box error: {}".format(total_error)) assert total_error < 1e-2 @@ -190,41 +204,44 @@ def get_image(path, trans): # half is not supported on CPUs, only GPUs if torch.cuda.is_available(): - mtcnn = MTCNN(keep_all=True, device='cuda').half() + mtcnn = MTCNN(keep_all=True, device="cuda").half() boxes_test, _ = mtcnn.detect(img) _ = mtcnn(img) - box_diff = boxes_ref[np.argsort(boxes_ref[:,1])] - boxes_test[np.argsort(boxes_test[:,1])] - print('fp16 Total box error: {}'.format(np.sum(np.abs(box_diff)))) + box_diff = ( + boxes_ref[np.argsort(boxes_ref[:, 1])] + - boxes_test[np.argsort(boxes_test[:, 1])] + ) + print("fp16 Total box error: {}".format(np.sum(np.abs(box_diff)))) # test new automatic multi precision to compare - if hasattr(torch.cuda, 'amp'): + if hasattr(torch.cuda, "amp"): with torch.cuda.amp.autocast(): - mtcnn = MTCNN(keep_all=True, device='cuda') + mtcnn = MTCNN(keep_all=True, device="cuda") boxes_test, _ = mtcnn.detect(img) _ = mtcnn(img) - box_diff = boxes_ref[np.argsort(boxes_ref[:,1])] - boxes_test[np.argsort(boxes_test[:,1])] - print('AMP total box error: {}'.format(np.sum(np.abs(box_diff)))) + box_diff = ( + boxes_ref[np.argsort(boxes_ref[:, 1])] + - boxes_test[np.argsort(boxes_test[:, 1])] + ) + print("AMP total box error: {}".format(np.sum(np.abs(box_diff)))) + - #### MULTI-IMAGE TEST #### mtcnn = MTCNN(keep_all=True) -img = [ - Image.open('data/multiface.jpg'), - Image.open('data/multiface.jpg') -] +img = [Image.open("data/multiface.jpg"), Image.open("data/multiface.jpg")] batch_boxes, batch_probs = mtcnn.detect(img) -mtcnn(img, save_path=['data/tmp1.png', 'data/tmp1.png']) -tmp_files = glob.glob('data/tmp*') +mtcnn(img, save_path=["data/tmp1.png", "data/tmp1.png"]) +tmp_files = glob.glob("data/tmp*") for f in tmp_files: os.remove(f) #### NO-FACE TEST #### -img = Image.new('RGB', (512, 512)) +img = Image.new("RGB", (512, 512)) mtcnn(img) mtcnn(img, return_prob=True) diff --git a/tests/perf_test.py b/tests/perf_test.py index 363c5696..55a806e6 100644 --- a/tests/perf_test.py +++ b/tests/perf_test.py @@ -7,7 +7,7 @@ def main(): - device = 'cuda' if torch.cuda.is_available() else 'cpu' + device = "cuda" if torch.cuda.is_available() else "cpu" print(f'Running on device "{device}"') mtcnn = MTCNN(device=device) @@ -16,8 +16,7 @@ def main(): # Generate data loader ds = datasets.ImageFolder( - root='data/test_images/', - transform=transforms.Resize((512, 512)) + root="data/test_images/", transform=transforms.Resize((512, 512)) ) dl = DataLoader( dataset=ds, @@ -32,8 +31,8 @@ def main(): for x, _ in tqdm(dl): faces.extend(mtcnn(x)) elapsed = time.time() - start - print(f'Elapsed: {elapsed} | EPS: {len(dl) * batch_size / elapsed}') + print(f"Elapsed: {elapsed} | EPS: {len(dl) * batch_size / elapsed}") -if __name__ == '__main__': +if __name__ == "__main__": main()