retinaface+facenet测试及导出onnx

最新推荐文章于 2025-04-07 11:13:30 发布

zxfeng~

最新推荐文章于 2025-04-07 11:13:30 发布

阅读量942

点赞数 12

分类专栏： AI 文章标签： python retinaface facenet 人脸识别人脸检测

本文链接：https://blog.csdn.net/qq_30479517/article/details/145285298

版权

AI 专栏收录该内容

13 篇文章

订阅专栏

在之前的文章中我们完成了使用insigtface进行多人人脸识别，其本质是分别使用人脸检测模型进行人脸检测，之后使用人脸识别模型对检测出的人脸进行识别。在本节，我们将分别使用retinaface模型和facenet模型在PC端分别完成人脸检测和人脸识别，并导出相应的onnx模型。

retinaface

下载源码

git clone https://github.com/bubbliiiing/retinaface-pytorch.git

下载完库后解压，运行predict.py，输入需要预测的图片，在img文件夹下提供了两张测试图片可以使用。

img/timg.jpg

若在服务器端测试，无法显示图片的话，可以注释掉cv2.imshow和cv2.waitKey，否则会报错。
在这里插入图片描述

若想要将识别到的结果输出至图片，则可以对predict.py进行如下修改：
首先在文件开头添加

import os

之后添加如下代码
在这里插入图片描述

if not os.path.exists(dir_save_path):
    os.makedirs(dir_save_path)

第一部分的代码用于创建不存在的dir_save_path文件夹

cv2.imwrite(os.path.join(dir_save_path, 'predict.jpg'), r_image)

第二部分的代码用于将识别后的结果保存至dir_save_path文件夹下的predict.jpg文件中。
最终结果如下：
在这里插入图片描述

facenet

下载源码

git clone https://github.com/bubbliiiing/facenet-pytorch.git

在model_data目录中已包含facenet_mobilenet.pth权重文件，可直接运行根目录的predict.py进行人脸识别

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_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

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

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 算子，我们可以对模型进行裁减，自行编写最后输出部分的运算。
在这里插入图片描述