人脸识别Adaface之libpytorch部署

目录

1. libpytorch下载2. Adaface模型下载3. 模型转换4. c++推理4.1 前处理4.2 推理4.3 编译运行4.3.1 写CMakeLists.txt4.3.2 编译4.3.3 运行

1. libpytorch下载

参考： https://blog.csdn.net/liang_baikai/article/details/127849577 下载完成后，将其解压到/usr/local下

2. Adaface模型下载

https://github.com/mk-minchul/AdaFace?tab=readme-ov-file WebFace4M模型准确率最高，R50 WebFace4M和R100 WebFace12M的准确率十分接近，但耗时却低了不少，所以建议使用R50 WebFace4M

3. 模型转换

下载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)

# 转换为 TorchScript

traced_model = torch.jit.trace(model_cpu, example_input)

# 保存模型

traced_model.save('adaface.pt')

# 导出为 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'}})

4. c++推理

4.1 前处理

resize人脸图片为112x112归一化BGR->RGB转换为tensorN H W C->N C H Wreshape 1，3，112，112（模型输入shape）

4.2 推理

load model读取图片人脸检测对齐前处理model.forward推理

#include

#include

#include

#include

torch::Tensor to_input(const cv::Mat& pil_rgb_image) {

cv::Mat brg_img;

cv::resize(pil_rgb_image, brg_img, cv::Size(112, 112));

brg_img.convertTo(brg_img, CV_32FC3, 1.0 / 255.0);

brg_img = (brg_img - 0.5) / 0.5;

cv::cvtColor(brg_img, brg_img, cv::COLOR_BGR2RGB);

torch::Tensor tensor = torch::from_blob(brg_img.data, {1, brg_img.rows, brg_img.cols, 3}, torch::kFloat32);

tensor = tensor.permute({0, 3, 1, 2});

tensor = tensor.reshape({1, 3, 112, 112});

tensor = tensor.to(at::kCPU);

return tensor;

}

int main() {

// 模型加载

torch::jit::script::Module model;

try {

model = torch::jit::load("./adaface.pt");

//model.eval();

model.to(at::kCPU);

} catch (const c10::Error& e) {

std::cerr << "Error loading the model\n";

return -1;

}

// 读取图片

std::vector 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 aligned_faces;

// 人脸检测对齐

open_cv_face.detectAndAlign(orig_img, aligned_faces);

//std::cout<<"detect use time is "<< (GetCurTimestamp() - detect_start)<

for (const auto &face:aligned_faces)

{

cv::Mat img(face);

auto img_tensor = to_input(img);

// Inference 推理

std::vector inputs;

inputs.push_back(img_tensor);

auto output = model.forward(inputs);

// Check if the output is a tuple

if (output.isTuple()) {

auto output_tuple = output.toTuple();

if (output_tuple->elements().size() > 0) {

at::Tensor output_tensor = output_tuple->elements()[0].toTensor();

//std::cout << output_tensor << std::endl;

} else {

std::cerr << "Output tuple is empty\n";

return -1;

}

} else {

at::Tensor output_tensor = output.toTensor();

//std::cout << output_tensor << std::endl;

}

}

}

return 0;

}

注意：本代码的人脸检测和对齐使用opencv的Yunet和SFace实现, 地址

4.3 编译运行

4.3.1 写CMakeLists.txt

本工程依赖opencv和libtorch，一并下载解压到/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})

if (AARCH64)

set(Torch_DIR /usr/local/libtorch/lib/python3.10/site-packages/torch/share/cmake/Torch)

else ()

set(Torch_DIR /usr/local/libtorch/share/cmake/Torch)

endif ()

find_package(Torch REQUIRED)

include_directories(${TORCH_INCLUDE_DIRS})

AUX_SOURCE_DIRECTORY(./src DIR_SRCS)

add_executable(adaface-demo ${DIR_SRCS})

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

4.3.2 编译

mkdir build

cd build

cmake ..

4.3.3 运行

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

cd ../bin

./main