当前位置：首页 > news >正文

人脸识别Adaface之libpytorch部署

news 2024/12/14 0:42:10

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 = Falsemodel.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
转换为tensor
N H W C->N C H W
reshape 1，3，112，112（模型输入shape）

4.2 推理

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

#include <torch/script.h>
#include <iostream>
#include <memory>
#include <opencv2/opencv.hpp>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<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 OpenCVcv::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);auto img_tensor = to_input(img);// Inference 推理std::vector<torch::jit::IValue> inputs;inputs.push_back(img_tensor);auto output = model.forward(inputs);// Check if the output is a tupleif (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})