jetson c++ facenet 给出实现代码

以下是 Jetson Nano 上使用 C++ 实现 FaceNet 的代码：

#include <iostream>
#include <fstream>
#include <sstream>
#include <iomanip>
#include <vector>

#include <opencv2/core.hpp>
#include <opencv2/highgui.hpp>
#include <opencv2/imgproc.hpp>

#include <NvInfer.h>
#include <NvOnnxParser.h>
#include <NvOnnxConfig.h>
#include <cuda_runtime_api.h>

using namespace std;
using namespace cv;
using namespace nvinfer1;
using namespace nvonnxparser;

// 定义常量
const string MODEL_PATH = "facenet.onnx";
const int INPUT_H = 160;
const int INPUT_W = 160;
const int CHANNELS = 3;
const int BATCH_SIZE = 1;
const int EMBEDDING_SIZE = 128;

// 定义全局变量
IExecutionContext* context;

// 定义结构体，用于存储人脸信息
struct FaceInfo {
    Rect faceRect;  // 人脸位置
    float score;    // 人脸置信度
};

// 定义函数，用于读取图像
Mat readImage(string imagePath) {
    Mat image = imread(imagePath);
    if (image.empty()) {
        cerr << "Failed to read image: " << imagePath << endl;
        exit(EXIT_FAILURE);
    }
    return image;
}

// 定义函数，用于将图像转换为 Tensor
void imageToTensor(Mat& image, float* inputTensor) {
    resize(image, image, Size(INPUT_H, INPUT_W));
    Mat floatImage;
    image.convertTo(floatImage, CV_32FC3, 1.f / 255.f);

    vector<Mat> channels;
    split(floatImage, channels);

    float* ptr = inputTensor;
    for (int i = 0; i < CHANNELS; ++i) {
        Mat channel = channels[i];
        memcpy(ptr, channel.data, INPUT_H * INPUT_W * sizeof(float));
        ptr += INPUT_H * INPUT_W;
    }
}

// 定义函数，用于从 Tensor 中提取特征向量
vector<float> extractFeatures(float* outputTensor) {
    vector<float> features;
    for (int i = 0; i < EMBEDDING_SIZE; ++i) {
        features.push_back(outputTensor[i]);
    }
    return features;
}

// 定义函数，用于检测人脸
vector<FaceInfo> detectFaces(Mat& image) {
    vector<FaceInfo> faceInfos;

    // 构建 TensorRT 的推理引擎
    IHostMemory* modelStream{ nullptr };
    ifstream file(MODEL_PATH, ios::binary);
    if (!file.good()) {
        cerr << "Failed to load model file: " << MODEL_PATH << endl;
        exit(EXIT_FAILURE);
    }
    file.seekg(0, file.end);
    const int modelSize = file.tellg();
    file.seekg(0, file.beg);
    vector<char> modelData(modelSize);
    file.read(modelData.data(), modelSize);
    file.close();

    int maxBatchSize = BATCH_SIZE;
    IRuntime* runtime = createInferRuntime(gLogger);
    ICudaEngine* engine = runtime->deserializeCudaEngine(modelData.data(), modelSize, nullptr);
    context = engine->createExecutionContext();

    // 读取输入 Tensor 的名称
    vector<string> inputTensorNames;
    for (int i = 0; i < engine->getNbBindings(); ++i) {
        if (engine->bindingIsInput(i)) {
            inputTensorNames.push_back(engine->getBindingName(i));
        }
    }
    assert(inputTensorNames.size() == 1);

    // 读取输出 Tensor 的名称
    vector<string> outputTensorNames;
    for (int i = 0; i < engine->getNbBindings(); ++i) {
        if (!engine->bindingIsInput(i)) {
            outputTensorNames.push_back(engine->getBindingName(i));
        }
    }
    assert(outputTensorNames.size() == 1);

    // 创建输入和输出 Tensor
    float* inputTensor;
    cudaMalloc(&inputTensor, BATCH_SIZE * CHANNELS * INPUT_H * INPUT_W * sizeof(float));
    void* outputTensor;
    cudaMalloc(&outputTensor, BATCH_SIZE * EMBEDDING_SIZE * sizeof(float));

    // 检测人脸
    Mat grayImage;
    cvtColor(image, grayImage, COLOR_BGR2GRAY);
    equalizeHist(grayImage, grayImage);

    CascadeClassifier cascade;
    if (!cascade.load("haarcascade_frontalface_alt.xml")) {
        cerr << "Failed to load face detection model" << endl;
        exit(EXIT_FAILURE);
    }

    vector<Rect> faces;
    cascade.detectMultiScale(grayImage, faces, 1.1, 2, 0 | CASCADE_SCALE_IMAGE, Size(30, 30));
    for (size_t i = 0; i < faces.size(); ++i) {
        // 提取人脸信息
        Rect faceRect = faces[i];
        float score = 1.f;

        // 裁剪人脸并将其转换为 Tensor
        Mat faceImage = image(faceRect);
        imageToTensor(faceImage, inputTensor);

        // 进行推理，提取特征向量
        cudaMemcpyAsync(inputTensor, inputTensor, BATCH_SIZE * CHANNELS * INPUT_H * INPUT_W * sizeof(float), cudaMemcpyHostToDevice);
        context->enqueue(BATCH_SIZE, &inputTensor, cudaStreamDefault, nullptr);
        cudaMemcpyAsync(outputTensor, outputTensor, BATCH_SIZE * EMBEDDING_SIZE * sizeof(float), cudaMemcpyDeviceToHost);
        cudaStreamSynchronize(cudaStreamDefault);
        vector<float> features = extractFeatures((float*)outputTensor);

        // 存储人脸信息
        FaceInfo faceInfo = {faceRect, score};
        faceInfos.push_back(faceInfo);
    }

    // 释放资源
    cudaFree(inputTensor);
    cudaFree(outputTensor);
    context->destroy();
    engine->destroy();
    runtime->destroy();

    return faceInfos;
}

int main(int argc, char** argv) {
    if (argc != 2) {
        cerr << "Usage: " << argv[0] << " <image_path>" << endl;
        return EXIT_FAILURE;
    }

    string imagePath = argv[1];
    Mat image = readImage(imagePath);
    vector<FaceInfo> faceInfos = detectFaces(image);

    // 在图像中绘制人脸框
    for (size_t i = 0; i < faceInfos.size(); ++i) {
        Rect faceRect = faceInfos[i].faceRect;
        float score = faceInfos[i].score;
        rectangle(image, faceRect, Scalar(0, 0, 255), 2);
        ostringstream ss;
        ss << fixed << setprecision(2) << score;
        putText(image, ss.str(), Point(faceRect.x, faceRect.y - 5), FONT_HERSHEY_SIMPLEX, 0.5, Scalar(0, 0, 255), 2);
    }

    // 显示图像
    imshow("Face Detection", image);
    waitKey(0);

    return EXIT_SUCCESS;
}

需要注意的是，该代码中使用了 OpenCV 和 TensorRT 库，因此需要进行相应的安装和配置。此外，还需要下载 FaceNet 模型文件 facenet.onnx，并将其放置在与代码文件相同的目录下。