waterman 发表于 2024-12-4 02:12

AI挑战营(进阶):2. retinaface+facenet测试及导出onnx

在上一节中我们完成了使用insigtface进行多人人脸识别,其本质是分别使用人脸检测模型进行人脸检测,之后使用人脸识别模型对检测出的人脸进行识别。在本节,我们将分别使用retinaface模型和facenet模型在PC端分别完成人脸检测和人脸识别,并导出相应的onnx模型。
# retinaface
下载源码
```bash
git clone https://github.com/bubbliiiing/retinaface-pytorch.git
```
下载完库后解压,运行predict.py,输入需要预测的图片,在img文件夹下提供了两张测试图片可以使用。
```bash
img/timg.jpg
```
若在服务器端测试,无法显示图片的话,可以注释掉cv2.imshow和cv2.waitKey,否则会报错。

若想要将识别到的结果输出至图片,则可以对predict.py进行如下修改:
首先在文件开头添加
```python
import os
```
之后添加如下代码

```python
if not os.path.exists(dir_save_path):
    os.makedirs(dir_save_path)
```
第一部分的代码用于创建不存在的dir_save_path文件夹
```python
cv2.imwrite(os.path.join(dir_save_path, 'predict.jpg'), r_image)
```
第二部分的代码用于将识别后的结果保存至dir_save_path文件夹下的predict.jpg文件中。
最终结果如下:

# facenet
下载源码
```bash
git clone https://github.com/bubbliiiing/facenet-pytorch.git
```
在model_data目录中已包含facenet_mobilenet.pth权重文件,可直接运行根目录的predict.py进行人脸识别
```python
from PIL import Image

from facenet import Facenet

if __name__ == "__main__":
    model = Facenet()
      
    while True:
      image_1 = input('Input image_1 filename:')
      try:
            image_1 = Image.open(image_1)
      except:
            print('Image_1 Open Error! Try again!')
            continue

      image_2 = input('Input image_2 filename:')
      try:
            image_2 = Image.open(image_2)
      except:
            print('Image_2 Open Error! Try again!')
            continue
      
      probability = model.detect_image(image_1,image_2)
      print(probability)
```
其中在根目录的img目录下已存在三张图片,1_001和1_002为一类,2_001为单独一类。可直接进行测试




测试结果如下:


从结果中可以看到,1_001和1_002的距离很近,为0.699;1_001和2_001的距离较远,为1.336。说明模型能够正确对人脸进行比对识别。
# 导出onnx
完成pytorch模型的测试之后,我们需要将它们转换为onnx格式的模型备用。
以下分别为retinaface和facenet导出onnx模型的代码。
- retinaface_export_onnx.py
```python
from nets.retinaface import RetinaFace
from utils.config import cfg_mnet
import torch

model_path='model_data/Retinaface_mobilenet0.25.pth' #模型路径
model=RetinaFace(cfg=cfg_mnet,pretrained = False) #模型初始化
device = torch.device('cpu')
model.load_state_dict(torch.load(model_path,map_location=device),strict=False) #模型加载
net=model.eval()
example=torch.rand(1,3,640,640) #给定输入
torch.onnx.export(model,(example),'model_data/retinaface.onnx',verbose=True,opset_version=9) #导出
```

- facenet_export_onnx.py
```python
from nets.facenet import Facenet
from torch import onnx
import torch

model_path='model_data/facenet_mobilenet.pth' #模型路径
model = Facenet(backbone="mobilenet",mode="predict",pretrained=True) #模型初始化
device = torch.device('cpu')
model.load_state_dict(torch.load(model_path, map_location=device), strict=False)
example=torch.rand(1,3,160,160) #给定一个输入
torch.onnx.export(model,example,'model_data/facenet.onnx',verbose=True,opset_version=9) #导出
```
执行以上命令后,就能够在model_data目录下生成相应的onnx模型如下图:


使用 Netron 工具查看 ONNX 模型的结构,确定是否存在 RV1103/RV1106 暂时无法支持的操作符(如Layer Normalization),参考手册 RKNN支持操作待列表(https://github.com/rockchip-linux/rknn-toolkit2/blob/master/doc/05_RKNN_Compiler_Support_Operator_List_v1.6.0.pdf) 。如果不支持的操作符位于模型的最后几层,可以考虑使用 CPU 进行实现。使用 Netron 工具观察 facenet 源码的模型,可以发现模型在输出前使用了 RKNPU 不支持的 ReduceSum 算子,我们可以对模型进行裁减,自行编写最后输出部分的运算。

ONNX模型的最后输出部分如下图所示:

下面我们进行模型的裁剪,参考源码路径 facenet-pytorch/nets/facenet.py 中的 forward 函数源码可以获知在模型推理阶段输出 128 维特征向量前进行了一个标准化的过程。

我们将最后一层的实现分配给 CPU 来运行,将最后一层的代码注释后,重新导出 ONNX 模型,得到的模型结构如下:

页: [1]
查看完整版本: AI挑战营(进阶):2. retinaface+facenet测试及导出onnx