Большинство туториалов по созданию и использованию моделей нейросетей написаны на Python. Однако для какого-нибудь проекта рано или поздно может понадобится использовать более быстрый и надежный язык для этих задач, например, C++. Эта статья о том, как на C++ можно использовать модель нейросети на примере модели YOLOv8 для детектирования лиц и библиотеки PyTorch.
Библиотеки
PyTorch предоставляет C++-версию своего модуля — библиотеку LibTorch. Её можно скачать (cxx11 ABI) и расположить рядом с будущей программой:
wget https://download.pytorch.org/libtorch/cpu/libtorch-cxx11-abi-shared-with-deps-2.0.0%2Bcpu.zip
unzip libtorch-cxx11-abi-shared-with-deps-2.0.0%2Bcpu.zip
Для примера также понадобится библиотека OpenCV. Её стоит скачать и установить:
wget https://github.com/opencv/opencv/archive/4.7.0.zip
unzip opencv-4.7.0.zip && cd opencv-4.7.0
mkdir build && cd build
cmake ..
cmake --build .
sudo make install
Веса
Обычный формат весов модели PyTorch необходимо экспортировать в формат TorchScript
. Для большинства моделей достаточно написать следующее:
# model_export.py
import torch
MODEL_PATH = 'model.pt'
MODEL_INPUT_SHAPE = (1,3,640,640)
model = torch.load(MODEL_PATH)
example = torch.randn(MODEL_INPUT_SHAPE)
model(example) # model warmup
model_traced = torch.jit.trace(model, example)
model_traced.save('model_traced.torchscript')
Однако для примера я использую модель YOLOv8 для детектирования лиц от Ultralytics, её можно экспортировать с помощью готовых модулей:
pip3 install ultralytics
pip3 install gdown
gdown 1jG8_C_P0SbnzYROZORe7CiJO6oMgp7eZ # yolov8n-face.pt
# model_export_yolov8face.py
import torch
from ultralytics import YOLO
model = YOLO('yolov8n-face.pt')
model.export(format='torchscript')
Использование
Инклюды, либы и линковка
Для компилятора будущей программы необходимо указать путь до icnlude- и lib-файлов библиотек LibTorch и OpenCV, а также линковать необходимые lib-файлы (они понадобятся во время компиляции; для автоматизации стоит использовать IDE, например, Visual Studio Code):
# ...
-I/usr/local/include/opencv4
-I/path/to/libtorch/include
-L/path/to/libtorch/lib
# ...
-o main.o
# ...
-lopencv_core
-lopencv_highgui
-lopencv_imgcodecs
-lopencv_imgproc
-lopencv_videoio
-ltorch
-ltorch_cpu
-lc10
Загрузка модели
Загрузить модель достаточно просто:
// main.cpp
#include <torch/script.h>
#include <opencv2/opencv.hpp>
int main() {
torch::jit::script::Module
model = torch::jit::load("yolov8n-face.torchscript");
return 0;
}
Загрузка данных на вход модели
LibTorch принимает на вход моделей вектор из torch::jit::IValue
и возвращает значения torch::jit::IValue
, которые затем необходимо экспортировать в привычные типы (тензор .toTensor()
, кортеж .toTuple()
или любой другой доступный).
Модель-пример принимает на вход тензор с формой (B, 3, 640, 640)
нормализованного изображения, а возвращает тензор с формой (1, 5, 8400)
, где 5 тензоров содержат 8400 целых значений cx, cy, w, h и вещественный conf для каждого бокса детектирования соответственно. Реализуем загрузку данных на вход модели и получение выходов:
// main.cpp
// ...
cv::Mat img = cv::imread(source);
cv::cvtColor(img, img, cv::COLOR_BGR2RGB);
cv::normalize(img, img, 0.0, 1.0, cv::NORM_MINMAX, CV_32F);
vector<torch::jit::IValue> inputs = {
torch::from_blob(
imgNorm.data,
{640,640,3},
torch::kFloat32
).permute({2,0,1}).unsqueeze(0)
};
at::Tensor outputs = model.forward(inputs).toTensor();
// ...
Детектирование лиц
Далее реализуем использование конкретно модели-примера.
Создадим класс Box
для боксов детектирования:
// detect.h
// ...
class Box {
public:
int x1, y1, x2, y2;
float conf;
Box(int x1, int y1, int x2, int y2, float conf) {
this->x1 = x1;
this->y1 = y1;
this->x2 = x2;
this->y2 = y2;
this->conf = conf;
}
};
// ...
Создадим функцию getBoxes
, которая будет возвращать боксы отфильтрованные по уверенности conf и с помощью будущего Non Maximum Supression (NMS):
// detect.h
// ...
vector<Box> getBoxes(
at::Tensor &outputs,
float confThres = 0.25,
float iouThres = 0.15
) {
vector<Box> candidates;
for (unsigned short ibatch = 0; ibatch < outputs.sizes()[0]; ibatch++) {
for (unsigned short ibox = 0; ibox < outputs.sizes()[2]; ibox++) {
float conf = outputs[ibatch][4][ibox].item<float>();
if (conf >= confThres) {
unsigned short
cx = outputs[ibatch][0][ibox].item<int>(),
cy = outputs[ibatch][1][ibox].item<int>(),
w = outputs[ibatch][2][ibox].item<int>(),
h = outputs[ibatch][3][ibox].item<int>();
unsigned short
x1 = cx - w / 2,
y1 = cy - h / 2,
x2 = cx + w / 2,
y2 = cy + h / 2;
candidates.push_back(Box(x1,y1,x2,y2,conf));
}
}
}
sort(candidates.begin(), candidates.end(), [](Box b1, Box b2){return b1.conf > b2.conf;});
vector<Box> boxes = nms(candidates, iouThres);
return boxes;
}
// ...
Создадим функцию nms
, которая будет фильтровать боксы-кандитаты и сопутствующую iou
, которая будет вычислять для них метрику IoU (Intersection Over Union):
// detect.h
// ...
float iou(Box &fb, Box &sb) {
float inter = max(min(fb.x2, sb.x2) - min(fb.x1, sb.x1), 0) * max(min(fb.y2, sb.y2) - min(fb.y1, sb.y1), 0);
float union_ = (fb.x2-fb.x1)*(fb.y2-fb.y1) + (sb.x2-sb.x1)*(sb.y2-sb.y1) - inter;
return inter / union_;
}
vector<Box> nms(vector<Box> &boxes, float iouThres) {
vector<Box> supBoxes;
for (Box box: boxes) {
bool valid = true;
for (Box supBox: supBoxes) {
if (iou(box, supBox) > iouThres) {
valid = false;
break;
}
}
if (valid == true) {
supBoxes.push_back(box);
}
}
return supBoxes;
}
// getBoxes
Наконец, получим боксы детектирования:
// main.cpp
// ...
vector<Box> boxes = getBoxes(outputs);
// ...
А также напишем процедуру для выделения боксов на оригинальном изображении:
// detect.h
// ...
void highlightBoxes(cv::Mat &img, vector<Box> &boxes) {
cv::Scalar rectColor(0,192,0);
unsigned short fontScale = 2, confPrecis = 2;
for (Box box: boxes) {
string text = to_string(box.conf);
cv::rectangle(img, {box.x1,box.y1}, {box.x2,box.y2}, rectColor, 2);
cv::rectangle(
img,
{box.x1, box.y1 - fontScale * 12},
{box.x1 + (unsigned short)text.length() * fontScale * 9, box.y1},
rectColor,
-1
);
cv::putText(img, text, {box.x1,box.y1}, cv::FONT_HERSHEY_PLAIN, fontScale, {255,255,255}, 2);
}
}
// ...
И выведем результат на экран:
// main.cpp
// ...
highlightBoxes(img, boxes);
cv::imshow("Result", img);
// return
Заключение
На более подробный и структурированный код вы можете взглянуть в одной из веток репозитория студенческого проекта по распознаванию лиц Распли. Надеюсь, вы сможете извлечь для себя из этого туториала что-то полезное.