482 lines
17 KiB
C++
482 lines
17 KiB
C++
#include "AapgsModelClassifier.h"
|
||
|
||
#include <cmath>
|
||
#include <cstdint>
|
||
#include <mutex>
|
||
|
||
#include <QCoreApplication>
|
||
#include <QDir>
|
||
#include <QFile>
|
||
#include <QFileInfo>
|
||
#include <QLibrary>
|
||
#include <QStringList>
|
||
|
||
#include "yolo_runtime.h"
|
||
|
||
namespace {
|
||
|
||
const char kModelConfigRelativePath[] =
|
||
"models/primary/model.yaml";
|
||
const char kModelConfigFileName[] = "model.yaml";
|
||
|
||
void AppendUniquePath(QStringList& paths, const QString& path)
|
||
{
|
||
const QString cleanPath = QDir::cleanPath(path.trimmed());
|
||
if (!cleanPath.isEmpty() && !paths.contains(cleanPath)) {
|
||
paths.push_back(cleanPath);
|
||
}
|
||
}
|
||
|
||
QString EnvironmentPath(const char* name)
|
||
{
|
||
return QString::fromLocal8Bit(qgetenv(name)).trimmed();
|
||
}
|
||
|
||
QString SourcePackageRoot()
|
||
{
|
||
QDir sourceDir(QFileInfo(QString::fromUtf8(__FILE__)).absolutePath());
|
||
#ifdef Q_OS_WIN
|
||
const QString packageDirectory = QStringLiteral("Windows");
|
||
#else
|
||
const QString packageDirectory = QStringLiteral("Arm");
|
||
#endif
|
||
return QDir::cleanPath(sourceDir.absoluteFilePath(
|
||
QStringLiteral("../../../../../AppAlgo/AAPGS_model/%1")
|
||
.arg(packageDirectory)));
|
||
}
|
||
|
||
QString PackageRootForConfig(const QString& configPath)
|
||
{
|
||
QDir directory(QFileInfo(configPath).absolutePath());
|
||
for (int level = 0; level < 5; ++level) {
|
||
const bool hasRuntime =
|
||
QFileInfo::exists(directory.filePath(
|
||
QStringLiteral("lib/libyolo_runtime.so.5"))) ||
|
||
QFileInfo::exists(directory.filePath(
|
||
QStringLiteral("lib/yolo_runtime.dll")));
|
||
if (hasRuntime) {
|
||
return directory.absolutePath();
|
||
}
|
||
if (!directory.cdUp()) {
|
||
break;
|
||
}
|
||
}
|
||
return QFileInfo(configPath).absolutePath();
|
||
}
|
||
|
||
bool FindModelConfig(QString& configPath,
|
||
QString& packageRoot,
|
||
QString& errorMessage)
|
||
{
|
||
const QString configuredPath = EnvironmentPath("AAPGS_MODEL_CONFIG");
|
||
if (!configuredPath.isEmpty()) {
|
||
QFileInfo configuredFile(configuredPath);
|
||
if (configuredFile.isDir()) {
|
||
const QDir configuredDir(configuredFile.absoluteFilePath());
|
||
const QFileInfo directConfig(configuredDir.filePath(
|
||
QString::fromLatin1(kModelConfigFileName)));
|
||
configuredFile.setFile(directConfig.isFile()
|
||
? directConfig.absoluteFilePath()
|
||
: configuredDir.filePath(QString::fromLatin1(
|
||
kModelConfigRelativePath)));
|
||
}
|
||
if (!configuredFile.isFile()) {
|
||
errorMessage = QStringLiteral("AAPGS模型配置不存在：%1")
|
||
.arg(configuredPath);
|
||
return false;
|
||
}
|
||
configPath = configuredFile.absoluteFilePath();
|
||
packageRoot = PackageRootForConfig(configPath);
|
||
return true;
|
||
}
|
||
|
||
QStringList roots;
|
||
const QString configuredRoot = EnvironmentPath("AAPGS_MODEL_ROOT");
|
||
if (!configuredRoot.isEmpty()) {
|
||
AppendUniquePath(roots, configuredRoot);
|
||
}
|
||
|
||
const QString applicationDir = QCoreApplication::applicationDirPath();
|
||
AppendUniquePath(roots, QDir(applicationDir).filePath(QStringLiteral("AAPGS_model")));
|
||
AppendUniquePath(roots, QDir(applicationDir).filePath(QStringLiteral("aapgs")));
|
||
AppendUniquePath(roots, QDir(applicationDir).filePath(QStringLiteral("../AAPGS_model")));
|
||
AppendUniquePath(roots, QDir::current().filePath(QStringLiteral("AAPGS_model")));
|
||
AppendUniquePath(roots, QDir::current().filePath(QStringLiteral("AppAlgo/AAPGS_model")));
|
||
AppendUniquePath(roots, SourcePackageRoot());
|
||
AppendUniquePath(roots, QStringLiteral("/opt/rk3588-ai/AAPGS_model"));
|
||
AppendUniquePath(roots, QStringLiteral("/usr/lib/AAPGS_model"));
|
||
AppendUniquePath(roots, QStringLiteral("/usr/local/lib/AAPGS_model"));
|
||
|
||
for (const QString& root : roots) {
|
||
const QFileInfo candidate(QDir(root).filePath(
|
||
QString::fromLatin1(kModelConfigRelativePath)));
|
||
if (candidate.isFile()) {
|
||
configPath = candidate.absoluteFilePath();
|
||
packageRoot = QDir(root).absolutePath();
|
||
return true;
|
||
}
|
||
}
|
||
|
||
errorMessage = QStringLiteral(
|
||
"未找到AAPGS机型分类配置，请部署到"
|
||
"/opt/rk3588-ai/AAPGS_model/"
|
||
"models/primary/model.yaml，"
|
||
"或设置AAPGS_MODEL_CONFIG");
|
||
return false;
|
||
}
|
||
|
||
QString ModelName(const yolo_result_common_v1_t& item,
|
||
QString& errorMessage)
|
||
{
|
||
int length = 0;
|
||
while (length < YOLO_LABEL_LENGTH && item.label[length] != '\0') {
|
||
++length;
|
||
}
|
||
const QString label = QString::fromUtf8(item.label, length).trimmed();
|
||
if (label.compare(QStringLiteral("a320"), Qt::CaseInsensitive) == 0 ||
|
||
(label.isEmpty() && item.class_id == 0)) {
|
||
return QStringLiteral("A320");
|
||
}
|
||
if (label.compare(QStringLiteral("b737"), Qt::CaseInsensitive) == 0 ||
|
||
(label.isEmpty() && item.class_id == 1)) {
|
||
return QStringLiteral("B737");
|
||
}
|
||
if (label.compare(QStringLiteral("other"), Qt::CaseInsensitive) == 0 ||
|
||
label == QStringLiteral("其他") ||
|
||
(label.isEmpty() && item.class_id == 2)) {
|
||
return QStringLiteral("其他");
|
||
}
|
||
errorMessage = label.isEmpty()
|
||
? QStringLiteral("AAPGS返回未知机型类别：%1").arg(item.class_id)
|
||
: QStringLiteral("AAPGS返回未知机型类别：%1/%2")
|
||
.arg(label)
|
||
.arg(item.class_id);
|
||
return QString();
|
||
}
|
||
|
||
} // namespace
|
||
|
||
class AapgsModelClassifier::Impl
|
||
{
|
||
public:
|
||
using GetAbiVersionFunction = uint32_t (*)();
|
||
using CreateFunction = yolo_status_t (*)(const char*, yolo_runtime_t**,
|
||
yolo_error_info_v1_t*);
|
||
using GetInfoFunction = yolo_status_t (*)(const yolo_runtime_t*,
|
||
yolo_runtime_info_v1_t*);
|
||
using InferImageFunction = yolo_status_t (*)(
|
||
yolo_runtime_t*, const yolo_image_view_v1_t*,
|
||
const yolo_infer_options_v1_t*, yolo_result_set_t**,
|
||
yolo_error_info_v1_t*);
|
||
using ResultGetCountFunction = yolo_status_t (*)(
|
||
const yolo_result_set_t*, uint32_t*);
|
||
using ResultGetCommonFunction = yolo_status_t (*)(
|
||
const yolo_result_set_t*, uint32_t, yolo_result_common_v1_t*);
|
||
using ResultReleaseFunction = void (*)(yolo_result_set_t*);
|
||
using DestroyFunction = void (*)(yolo_runtime_t*);
|
||
|
||
~Impl()
|
||
{
|
||
UnloadRuntimeLibrary();
|
||
}
|
||
|
||
bool Classify(const QImage& frame,
|
||
Classification& result,
|
||
QString& errorMessage)
|
||
{
|
||
std::lock_guard<std::mutex> lock(m_mutex);
|
||
result = Classification();
|
||
errorMessage.clear();
|
||
|
||
if (frame.isNull() || frame.width() <= 0 || frame.height() <= 0) {
|
||
errorMessage = QStringLiteral("AAPGS机型识别图像无效");
|
||
return false;
|
||
}
|
||
if (!EnsureInitialized(errorMessage)) {
|
||
return false;
|
||
}
|
||
|
||
const QImage rgbFrame = frame.format() == QImage::Format_RGB888
|
||
? frame
|
||
: frame.convertToFormat(QImage::Format_RGB888);
|
||
if (rgbFrame.isNull()) {
|
||
errorMessage = QStringLiteral("AAPGS机型识别图像转换失败");
|
||
return false;
|
||
}
|
||
|
||
yolo_image_view_v1_t image{};
|
||
image.struct_size = sizeof(image);
|
||
image.version = YOLO_IMAGE_VIEW_ABI_VERSION;
|
||
image.data = rgbFrame.constBits();
|
||
image.width = rgbFrame.width();
|
||
image.height = rgbFrame.height();
|
||
image.row_stride_bytes = rgbFrame.bytesPerLine();
|
||
image.pixel_format = YOLO_PIXEL_FORMAT_RGB888;
|
||
|
||
yolo_infer_options_v1_t options{};
|
||
options.struct_size = sizeof(options);
|
||
options.version = YOLO_INFER_OPTIONS_VERSION;
|
||
options.seg_geometry_mode = YOLO_SEG_GEOMETRY_NONE;
|
||
|
||
yolo_error_info_v1_t error = MakeErrorInfo();
|
||
yolo_result_set_t* rawResults = nullptr;
|
||
const yolo_status_t status = m_inferImage(
|
||
m_runtime, &image, &options, &rawResults, &error);
|
||
if (status != YOLO_STATUS_OK || !rawResults) {
|
||
const QString runtimeError = RuntimeError(error);
|
||
if (status == YOLO_STATUS_INITIALIZATION_FAILED ||
|
||
status == YOLO_STATUS_BACKEND_FAILED) {
|
||
UnloadRuntimeLibrary();
|
||
}
|
||
errorMessage = QStringLiteral("AAPGS模型推理失败(%1)：%2")
|
||
.arg(status)
|
||
.arg(runtimeError);
|
||
return false;
|
||
}
|
||
|
||
const auto releaseResults = [this](yolo_result_set_t* results) {
|
||
if (results && m_resultRelease) {
|
||
m_resultRelease(results);
|
||
}
|
||
};
|
||
std::unique_ptr<yolo_result_set_t, decltype(releaseResults)> results(
|
||
rawResults, releaseResults);
|
||
|
||
uint32_t resultCount = 0;
|
||
if (m_resultGetCount(results.get(), &resultCount) != YOLO_STATUS_OK ||
|
||
resultCount == 0 || resultCount > 4096) {
|
||
errorMessage = QStringLiteral("AAPGS分类结果数量异常：%1")
|
||
.arg(resultCount);
|
||
return false;
|
||
}
|
||
|
||
yolo_result_common_v1_t bestResult{};
|
||
bool hasBestResult = false;
|
||
for (uint32_t index = 0; index < resultCount; ++index) {
|
||
yolo_result_common_v1_t item{};
|
||
item.struct_size = sizeof(item);
|
||
item.version = YOLO_RESULT_COMMON_VERSION;
|
||
if (m_resultGetCommon(results.get(), index, &item) !=
|
||
YOLO_STATUS_OK) {
|
||
continue;
|
||
}
|
||
if (item.task != YOLO_TASK_CLS) {
|
||
continue;
|
||
}
|
||
if (!std::isfinite(item.score) || item.score < 0.0f ||
|
||
item.score > 1.0f) {
|
||
continue;
|
||
}
|
||
if (!hasBestResult || item.score > bestResult.score) {
|
||
bestResult = item;
|
||
hasBestResult = true;
|
||
}
|
||
}
|
||
if (!hasBestResult) {
|
||
errorMessage = QStringLiteral("AAPGS未返回有效分类结果");
|
||
return false;
|
||
}
|
||
result.modelType = ModelName(bestResult, errorMessage);
|
||
result.confidence = static_cast<double>(bestResult.score);
|
||
if (result.modelType.isEmpty()) {
|
||
result = Classification();
|
||
return false;
|
||
}
|
||
return true;
|
||
}
|
||
|
||
private:
|
||
bool EnsureInitialized(QString& errorMessage)
|
||
{
|
||
if (m_runtime) {
|
||
return true;
|
||
}
|
||
|
||
QString configPath;
|
||
QString packageRoot;
|
||
if (!FindModelConfig(configPath, packageRoot, errorMessage) ||
|
||
!LoadRuntimeLibrary(packageRoot, errorMessage)) {
|
||
return false;
|
||
}
|
||
|
||
const uint32_t runtimeAbi = m_getAbiVersion();
|
||
if (runtimeAbi != YOLO_RUNTIME_ABI_VERSION) {
|
||
errorMessage = QStringLiteral("AAPGS运行库ABI不匹配：期望%1，实际%2")
|
||
.arg(YOLO_RUNTIME_ABI_VERSION)
|
||
.arg(runtimeAbi);
|
||
UnloadRuntimeLibrary();
|
||
return false;
|
||
}
|
||
|
||
const QByteArray encodedConfigPath = QFile::encodeName(configPath);
|
||
yolo_error_info_v1_t error = MakeErrorInfo();
|
||
const yolo_status_t status = m_create(
|
||
encodedConfigPath.constData(), &m_runtime, &error);
|
||
if (status != YOLO_STATUS_OK || !m_runtime) {
|
||
const QString runtimeError = RuntimeError(error);
|
||
UnloadRuntimeLibrary();
|
||
errorMessage = QStringLiteral("AAPGS模型初始化失败(%1)：%2；配置：%3")
|
||
.arg(status)
|
||
.arg(runtimeError)
|
||
.arg(configPath);
|
||
return false;
|
||
}
|
||
|
||
yolo_runtime_info_v1_t runtimeInfo{};
|
||
runtimeInfo.struct_size = sizeof(runtimeInfo);
|
||
runtimeInfo.version = YOLO_RUNTIME_INFO_VERSION;
|
||
if (m_getInfo(m_runtime, &runtimeInfo) != YOLO_STATUS_OK) {
|
||
UnloadRuntimeLibrary();
|
||
errorMessage = QStringLiteral("读取AAPGS运行库信息失败");
|
||
return false;
|
||
}
|
||
if (runtimeInfo.task != YOLO_TASK_CLS) {
|
||
const int task = static_cast<int>(runtimeInfo.task);
|
||
UnloadRuntimeLibrary();
|
||
errorMessage = QStringLiteral("AAPGS模型任务不是分类任务：%1")
|
||
.arg(task);
|
||
return false;
|
||
}
|
||
if ((runtimeInfo.capabilities & YOLO_CAPABILITY_INFER_IMAGE) == 0) {
|
||
UnloadRuntimeLibrary();
|
||
errorMessage = QStringLiteral("AAPGS运行库不支持内存图像推理");
|
||
return false;
|
||
}
|
||
return true;
|
||
}
|
||
|
||
bool LoadRuntimeLibrary(const QString& packageRoot, QString& errorMessage)
|
||
{
|
||
if (m_library.isLoaded()) {
|
||
if (ResolveFunctions()) {
|
||
return true;
|
||
}
|
||
UnloadRuntimeLibrary();
|
||
}
|
||
|
||
QStringList candidates;
|
||
const QString configuredLibrary = EnvironmentPath("AAPGS_RUNTIME_LIBRARY");
|
||
if (!configuredLibrary.isEmpty()) {
|
||
AppendUniquePath(candidates, configuredLibrary);
|
||
}
|
||
#ifdef Q_OS_WIN
|
||
AppendUniquePath(candidates, QDir(packageRoot).filePath(
|
||
QStringLiteral("lib/yolo_runtime.dll")));
|
||
AppendUniquePath(candidates, QStringLiteral("yolo_runtime"));
|
||
#else
|
||
AppendUniquePath(candidates, QDir(packageRoot).filePath(
|
||
QStringLiteral("lib/libyolo_runtime.so.5")));
|
||
AppendUniquePath(candidates, QDir(packageRoot).filePath(
|
||
QStringLiteral("lib/libyolo_runtime.so")));
|
||
AppendUniquePath(candidates, QStringLiteral("libyolo_runtime.so.5"));
|
||
AppendUniquePath(candidates, QStringLiteral("yolo_runtime"));
|
||
#endif
|
||
|
||
QString lastLoadError;
|
||
for (const QString& candidate : candidates) {
|
||
m_library.setFileName(candidate);
|
||
if (!m_library.load()) {
|
||
lastLoadError = m_library.errorString();
|
||
continue;
|
||
}
|
||
if (ResolveFunctions()) {
|
||
return true;
|
||
}
|
||
lastLoadError = QStringLiteral("AAPGS运行库缺少必需的C ABI符号");
|
||
UnloadRuntimeLibrary();
|
||
}
|
||
|
||
errorMessage = QStringLiteral("加载AAPGS运行库失败：%1")
|
||
.arg(lastLoadError);
|
||
return false;
|
||
}
|
||
|
||
bool ResolveFunctions()
|
||
{
|
||
m_getAbiVersion = reinterpret_cast<GetAbiVersionFunction>(
|
||
m_library.resolve("yolo_runtime_get_abi_version"));
|
||
m_create = reinterpret_cast<CreateFunction>(
|
||
m_library.resolve("yolo_runtime_create"));
|
||
m_getInfo = reinterpret_cast<GetInfoFunction>(
|
||
m_library.resolve("yolo_runtime_get_info"));
|
||
m_inferImage = reinterpret_cast<InferImageFunction>(
|
||
m_library.resolve("yolo_runtime_infer_image"));
|
||
m_resultGetCount = reinterpret_cast<ResultGetCountFunction>(
|
||
m_library.resolve("yolo_result_set_get_count"));
|
||
m_resultGetCommon = reinterpret_cast<ResultGetCommonFunction>(
|
||
m_library.resolve("yolo_result_set_get_common"));
|
||
m_resultRelease = reinterpret_cast<ResultReleaseFunction>(
|
||
m_library.resolve("yolo_result_set_release"));
|
||
m_destroy = reinterpret_cast<DestroyFunction>(
|
||
m_library.resolve("yolo_runtime_destroy"));
|
||
return m_getAbiVersion && m_create && m_getInfo && m_inferImage &&
|
||
m_resultGetCount && m_resultGetCommon && m_resultRelease &&
|
||
m_destroy;
|
||
}
|
||
|
||
static yolo_error_info_v1_t MakeErrorInfo()
|
||
{
|
||
yolo_error_info_v1_t error{};
|
||
error.struct_size = sizeof(error);
|
||
error.version = YOLO_ERROR_INFO_VERSION;
|
||
return error;
|
||
}
|
||
|
||
static QString RuntimeError(const yolo_error_info_v1_t& error)
|
||
{
|
||
int length = 0;
|
||
while (length < static_cast<int>(sizeof(error.message)) &&
|
||
error.message[length] != '\0') {
|
||
++length;
|
||
}
|
||
const QString text = QString::fromUtf8(error.message, length).trimmed();
|
||
return text.isEmpty() ? QStringLiteral("未知错误") : text;
|
||
}
|
||
|
||
void UnloadRuntimeLibrary()
|
||
{
|
||
if (m_runtime && m_destroy) {
|
||
m_destroy(m_runtime);
|
||
}
|
||
m_runtime = nullptr;
|
||
if (m_library.isLoaded()) {
|
||
m_library.unload();
|
||
}
|
||
m_getAbiVersion = nullptr;
|
||
m_create = nullptr;
|
||
m_getInfo = nullptr;
|
||
m_inferImage = nullptr;
|
||
m_resultGetCount = nullptr;
|
||
m_resultGetCommon = nullptr;
|
||
m_resultRelease = nullptr;
|
||
m_destroy = nullptr;
|
||
}
|
||
|
||
private:
|
||
std::mutex m_mutex;
|
||
QLibrary m_library;
|
||
yolo_runtime_t* m_runtime = nullptr;
|
||
GetAbiVersionFunction m_getAbiVersion = nullptr;
|
||
CreateFunction m_create = nullptr;
|
||
GetInfoFunction m_getInfo = nullptr;
|
||
InferImageFunction m_inferImage = nullptr;
|
||
ResultGetCountFunction m_resultGetCount = nullptr;
|
||
ResultGetCommonFunction m_resultGetCommon = nullptr;
|
||
ResultReleaseFunction m_resultRelease = nullptr;
|
||
DestroyFunction m_destroy = nullptr;
|
||
};
|
||
|
||
AapgsModelClassifier::AapgsModelClassifier()
|
||
: m_impl(std::make_unique<Impl>())
|
||
{
|
||
}
|
||
|
||
AapgsModelClassifier::~AapgsModelClassifier() = default;
|
||
|
||
bool AapgsModelClassifier::Classify(const QImage& frame,
|
||
Classification& result,
|
||
QString& errorMessage)
|
||
{
|
||
return m_impl && m_impl->Classify(frame, result, errorMessage);
|
||
}
|