人脸识别Adaface之onnx部署

1. onnxruntime下载

下载地址

2. Adaface模型下载

https://github.com/mk-minchul/AdaFace?tab=readme-ov-file

WebFace4M模型准确率最高，R50 WebFace4M和R100 WebFace12M的准确率十分接近，但耗时却低了不少，所以建议使用R50 WebFace4M

3. 创建工程

mkdir adaface-demo
cd adaface-demo
mkdir 3rdparty
mkdir bin

4. 模型转换

下载Adaface源码，并将下面代码放到其目录下执行即可

model_trans.py

import torch
import torch.nn as nn
from head import AdaFace 
import net
import onnxruntime as ort
import numpy as np
import onnx


# 加载模型
adaface_models = {
#    'ir_101':"./adaface_ir101_ms1mv2.ckpt",
    'ir_50':"./adaface_ir50_webface4m.ckpt",
}
architecture = 'ir_50'

model = net.build_model(architecture)
#model = AdaFace()
statedict = torch.load(adaface_models[architecture],map_location=torch.device('cpu'),weights_only=True)['state_dict']
model_statedict = {key[6:]:val for key, val in statedict.items() if key.startswith('model.')}

model.load_state_dict(model_statedict, strict=True)

for p in model.parameters():
    p.requires_grad = False

model.eval()
device = torch.device("cpu");
model_cpu = model.to(device)

# 创建一个示例输入
example_input = torch.rand(1, 3, 112, 112)  # 假设输入大小为 (1, 3, 112, 112)


# 导出为 ONNX 格式
onnx_file_path = 'adaface.onnx'  # 输出文件名
torch.onnx.export(model, example_input, onnx_file_path,
                  export_params=True)
                  #opset_version=11,  # ONNX 版本
                  #do_constant_folding=True,  # 是否进行常量折叠
                  #input_names=['input'],  # 输入名称
                  #output_names=['output'],  # 输出名称
                  #dynamic_axes={'input': {0: 'batch_size'},  # 动态 batch size
                  #              'output': {0: 'batch_size'}})

拷贝转换好的模型到工程目录下

cp adaface.onnx adaface-demo/bin

4. c++推理

#include <iostream>
#include <memory>
#include <opencv2/opencv.hpp>
#include <chrono>
#include <dirent.h>
#include <sys/stat.h>
#include <sys/time.h>
#include "onnxruntime_cxx_api.h"
#include "cpu_provider_factory.h"
#include "OpenCVFace.h"

using namespace std;

bool check_path(std::string file_path,
                std::vector<std::string> correct_postfixes)
{
    auto index = file_path.rfind('.');
    std::string postfix = file_path.substr(index + 1);
    if (find(correct_postfixes.begin(), correct_postfixes.end(), postfix) !=
        correct_postfixes.end())
    {
        return true;
    }
    else
    {
        printf("skipping path: %s , please check your dataset!", file_path.c_str());
        return false;
    }
}


void getAllFiles(std::string path, std::vector<std::string> &files,
                 std::vector<std::string> correct_postfixes)
{
    DIR *dir;
    struct dirent *ptr;
    if ((dir = opendir(path.c_str())) == NULL)
    {
        perror("Open dri error...");
        return;
    }
    while ((ptr = readdir(dir)) != NULL)
    {
        if (strcmp(ptr->d_name, ".") == 0 || strcmp(ptr->d_name, "..") == 0)
            continue;
        else if (ptr->d_type == 8 &&
                 check_path(path + "/" + ptr->d_name, correct_postfixes)) // file
                     files.push_back(path + "/" + ptr->d_name);
        else if (ptr->d_type == 10) // link file
            continue;
        else if (ptr->d_type == 4)
        {
            // files.push_back(ptr->d_name);//dir
            getAllFiles(path + "/" + ptr->d_name, files, correct_postfixes);
        }
    }
    closedir(dir);
}

auto GetCurTimestamp()
{
	struct timeval tv;
	gettimeofday(&tv, NULL);
	return ((unsigned long long)tv.tv_sec * 1000 + (unsigned long long)tv.tv_usec / 1000);
};

Ort::Session onnx_init(const string &model_path)
{
	// 初始化ONNXRuntime环境
	Ort::Env env = Ort::Env(ORT_LOGGING_LEVEL_ERROR, "ir_101");

	// 设置会话选项
	Ort::SessionOptions session_options;
	// 优化器级别:基本的图优化级别
	session_options.SetGraphOptimizationLevel(ORT_ENABLE_BASIC);
	// 线程数:4
	session_options.SetIntraOpNumThreads(4);
	// 设备使用优先使用GPU而是才是CPU
	std::cout << "onnxruntime inference try to use GPU Device" << std::endl;
	OrtSessionOptionsAppendExecutionProvider_CUDA(session_options, 0);
	OrtSessionOptionsAppendExecutionProvider_CPU(session_options, 1);

	// onnx训练模型文件
	std::string onnxpath = model_path;
	std::wstring modelPath = std::wstring(onnxpath.begin(), onnxpath.end());

	// 加载模型并创建会话
	Ort::Session session_(env, onnxpath.c_str(), session_options);

	return move(session_);
}

int onnx_infer(const cv::Mat &image, Ort::Session &session_)
{
	// 获取模型输入输出信息
	int input_nodes_num = session_.GetInputCount();			// 输入节点输
	int output_nodes_num = session_.GetOutputCount();		// 输出节点数
	std::vector<std::string> input_node_names;				// 输入节点名称
	std::vector<std::string> output_node_names;				// 输出节点名称
	Ort::AllocatorWithDefaultOptions allocator;
	// 输入图像尺寸
	int input_h = 0;
	int input_w = 0;

	// 获取模型输入信息
	for (int i = 0; i < input_nodes_num; i++) {
		// 获得输入节点的名称并存储
		auto input_name = session_.GetInputNameAllocated(i, allocator);
		input_node_names.push_back(input_name.get());
		// 显示输入图像的形状
		auto inputShapeInfo = session_.GetInputTypeInfo(i).GetTensorTypeAndShapeInfo().GetShape();
		int ch = inputShapeInfo[1];
		input_h = inputShapeInfo[2];
		input_w = inputShapeInfo[3];
		//std::cout << "input format: " << ch << "x" << input_h << "x" << input_w << std::endl;
	}

	// 获取模型输出信息
	int num = 0;
	int nc = 0;
	for (int i = 0; i < output_nodes_num; i++) {
		// 获得输出节点的名称并存储
		auto output_name = session_.GetOutputNameAllocated(i, allocator);
		output_node_names.push_back(output_name.get());
		// 显示输出结果的形状
		auto outShapeInfo = session_.GetOutputTypeInfo(i).GetTensorTypeAndShapeInfo().GetShape();
		num = outShapeInfo[0];
		nc = outShapeInfo[1];
		//std::cout << "output format: " << num << "x" << nc << std::endl;
	}

	// 预处理输入数据
	cv::Mat rgb, blob;
	// 默认是BGR需要转化成RGB
	cv::cvtColor(image, rgb, cv::COLOR_BGR2RGB);
	// 对图像尺寸进行缩放
	cv::resize(rgb, blob, cv::Size(input_w, input_h));
	blob.convertTo(blob, CV_32F);
	// 对图像进行标准化处理
	blob = blob / 255.0;	// 归一化
	cv::subtract(blob, cv::Scalar(0.485, 0.456, 0.406), blob);	// 减去均值
	cv::divide(blob, cv::Scalar(0.229, 0.224, 0.225), blob);	//除以方差
	// CHW-->NCHW 维度扩展
	cv::Mat timg = cv::dnn::blobFromImage(blob);
	//std::cout << timg.size[0] << "x" << timg.size[1] << "x" << timg.size[2] << "x" << timg.size[3] << std::endl;
	// 占用内存大小,后续计算是总像素*数据类型大小
	size_t tpixels = input_h * input_w * 3;
	std::array<int64_t, 4> input_shape_info{ 1, 3, input_h, input_w };

	// 准备数据输入
	auto allocator_info = Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU);
	Ort::Value input_tensor_ = Ort::Value::CreateTensor<float>(allocator_info, timg.ptr<float>(), tpixels, input_shape_info.data(), input_shape_info.size());

	// 模型输入输出所需数据(名称及其数量),模型只认这种类型的数组
	const std::array<const char*, 1> inputNames = { input_node_names[0].c_str() };
	const std::array<const char*, 1> outNames = { output_node_names[0].c_str() };

	// 模型推理
	std::vector<Ort::Value> ort_outputs;
	try {
		auto start = GetCurTimestamp();
		ort_outputs = session_.Run(Ort::RunOptions{ nullptr }, inputNames.data(), &input_tensor_, 1, outNames.data(), outNames.size());
		std::cout<< "use time:"<<GetCurTimestamp()-start<<std::endl;
	}
	catch (std::exception e) {
		std::cout << e.what() << std::endl;
	}
	// 1x5 获取输出数据并包装成一个cv::Mat对象,为了方便后处理
	const float* pdata = ort_outputs[0].GetTensorMutableData<float>();
	cv::Mat prob(num, nc, CV_32F, (float*)pdata);

    return 0;
}

int main() {
	Ort::Session session = onnx_init("./adaface.onnx");

    std::vector<std::string> images;
    getAllFiles("./images", images, {"jpg", "jpeg", "png"});
	OpenCVFace open_cv_face;
	open_cv_face.Init("./models/face_detection_yunet_2023mar.onnx",
		"./models/face_recognition_sface_2021dec.onnx", 0.9, 0.5);

    for (const auto &image_path : images)
    {
        // Load an image using OpenCV
        cv::Mat orig_img = cv::imread(image_path);
        if (orig_img.empty()) {
            std::cerr << "Could not read the image\n";
            return -1;
        }

        auto detect_start = GetCurTimestamp();
        std::vector<cv::Mat> aligned_faces;
        
    	open_cv_face.detectAndAlign(orig_img, aligned_faces);
    	
        //std::cout<<"detect use time is  "<< (GetCurTimestamp() - detect_start)<<std::endl;
        for (const auto &face:aligned_faces)
        {
            cv::Mat img(face);
            
        	std::cout<<image_path<<" ";
        	onnx_infer(img, session);
        }
    }

    return 0;
}

注意：本代码的人脸检测和对齐使用opencv的Yunet和SFace实现, 地址
本代码参考：https://blog.csdn.net/yangyu0515/article/details/142093965

5. 编译运行

5.1 写CMakeLists.txt

本工程依赖opencv4.9，下载后解压到/usr/local下即可。

cmake_minimum_required(VERSION 3.22.1)
project(adaface-demo)

set(QMAKE_CXXFLAGS "-std=c++17")
set(EXECUTABLE_OUTPUT_PATH ${PROJECT_SOURCE_DIR}/bin)

include_directories(/usr/local/include)
link_directories(/usr/local/lib)

set(OPENCV_VERSION "4.9.0")
set(OPENCV_INSTALLATION_PATH "/usr/local/opencv4" CACHE PATH "Where to look for OpenCV installation")

# Find OpenCV
find_package(OpenCV ${OPENCV_VERSION} REQUIRED HINTS ${OPENCV_INSTALLATION_PATH})

set(ONNXRUNTIME_DIR ${PROJECT_SOURCE_DIR}/3rdparty/onnxruntime)
include_directories(${ONNXRUNTIME_DIR}/include)
link_directories(${ONNXRUNTIME_DIR}/lib)
set(ONNX_LIB ${ONNXRUNTIME_DIR}/lib/libonnxruntime.so)

AUX_SOURCE_DIRECTORY(./src DIR_SRCS)
add_executable(adaface-demo ${DIR_SRCS})

target_link_libraries(adaface-demo ${OpenCV_LIBS} ${ONNX_LIB})

5.2 编译

mkdir build
cd build
cmake ..

5.3 运行

将模型文件adaface.py拷贝到bin目录下

cd ../bin
./main