人脸识别Adaface之libpytorch部署

news2025/1/17 4:03:35

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
转换为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 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);

        	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 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})