GrabBag/App/ParkingSpaceGuide/ParkingSpaceGuideApp/Presenter/Src/PersionModelDetector.cpp

#include "PersionModelDetector.h"

#include <cmath>
#include <mutex>

#include <QCoreApplication>
#include <QDir>
#include <QFile>
#include <QFileInfo>
#include <QLibrary>
#include <QStringList>

#include "../../../../../AppAlgo/Persion_model/include/yolo_runtime.h"

namespace {

const char kModelConfigRelativePath[] = "models/yolov8n/yolov8n.yaml";
const int kPersionRuntimeAbiVersion = 4;

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());
    return QDir::cleanPath(sourceDir.absoluteFilePath(
        QStringLiteral("../../../../../AppAlgo/Persion_model")));
}

QString PackageRootForConfig(const QString& configPath)
{
    QDir directory(QFileInfo(configPath).absolutePath());
    for (int level = 0; level < 5; ++level) {
        if (QFileInfo::exists(directory.filePath(
                QStringLiteral("lib/libyolo_runtime.so.3")))) {
            return directory.absolutePath();
        }
        if (!directory.cdUp()) {
            break;
        }
    }
    return QFileInfo(configPath).absolutePath();
}

bool FindModelConfig(QString& configPath,
                     QString& packageRoot,
                     QString& errorMessage)
{
    const QString configuredPath = EnvironmentPath("PERSION_MODEL_CONFIG");
    if (!configuredPath.isEmpty()) {
        QFileInfo configuredFile(configuredPath);
        if (configuredFile.isDir()) {
            const QDir configuredDir(configuredFile.absoluteFilePath());
            const QFileInfo directConfig(configuredDir.filePath(
                QStringLiteral("yolov8n.yaml")));
            configuredFile.setFile(directConfig.isFile()
                ? directConfig.absoluteFilePath()
                : configuredDir.filePath(QString::fromLatin1(
                      kModelConfigRelativePath)));
        }
        if (!configuredFile.isFile()) {
            errorMessage = QStringLiteral("Persion人员检测模型配置不存在：%1")
                               .arg(configuredPath);
            return false;
        }
        configPath = configuredFile.absoluteFilePath();
        packageRoot = PackageRootForConfig(configPath);
        return true;
    }

    QStringList roots;
    const QString configuredRoot = EnvironmentPath("PERSION_MODEL_ROOT");
    if (!configuredRoot.isEmpty()) {
        AppendUniquePath(roots, configuredRoot);
    }
    const QString applicationDir = QCoreApplication::applicationDirPath();
    AppendUniquePath(roots, QDir(applicationDir).filePath(QStringLiteral("Persion_model")));
    AppendUniquePath(roots, QDir(applicationDir).filePath(QStringLiteral("../Persion_model")));
    AppendUniquePath(roots, QDir::current().filePath(QStringLiteral("Persion_model")));
    AppendUniquePath(roots, QDir::current().filePath(QStringLiteral("AppAlgo/Persion_model")));
    AppendUniquePath(roots, SourcePackageRoot());
    AppendUniquePath(roots, QStringLiteral("/opt/rk3588-ai/Persion_model"));
    AppendUniquePath(roots, QStringLiteral("/usr/lib/Persion_model"));
    AppendUniquePath(roots, QStringLiteral("/usr/local/lib/Persion_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(
        "未找到Persion人员检测模型配置，请部署到"
        "/opt/rk3588-ai/Persion_model/models/yolov8n/yolov8n.yaml，"
        "或设置PERSION_MODEL_CONFIG");
    return false;
}

QString ResultLabel(const yolo_result_t& item)
{
    int length = 0;
    while (length < YOLO_LABEL_LENGTH && item.label[length] != '\0') {
        ++length;
    }
    return QString::fromUtf8(item.label, length).trimmed();
}

} // namespace

class PersionModelDetector::Impl
{
public:
    using GetAbiVersionFunction = int (*)();
    using CreateFunction = int (*)(const char*, yolo_runtime_t**);
    using GetCapabilitiesFunction = unsigned int (*)(const yolo_runtime_t*);
    using InferImageFunction = int (*)(yolo_runtime_t*, const yolo_image_view_t*,
                                       yolo_result_list_t*);
    using LastErrorFunction = const char* (*)(const yolo_runtime_t*);
    using DestroyFunction = void (*)(yolo_runtime_t*);

    ~Impl()
    {
        UnloadRuntimeLibrary();
    }

    bool Analyze(const QImage& frame,
                 Analysis& result,
                 QString& errorMessage)
    {
        std::lock_guard<std::mutex> lock(m_mutex);
        result = Analysis();
        errorMessage.clear();

        if (frame.isNull() || frame.width() <= 0 || frame.height() <= 0) {
            errorMessage = QStringLiteral("Persion人员检测图像无效");
            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("Persion人员检测图像转换失败");
            return false;
        }

        yolo_image_view_t image{};
        image.struct_size = sizeof(image);
        image.abi_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_result_list_t output{};
        const int status = m_inferImage(m_runtime, &image, &output);
        if (status != YOLO_STATUS_OK) {
            const QString runtimeError = RuntimeError();
            if (status == YOLO_STATUS_INITIALIZATION_FAILED ||
                status == YOLO_STATUS_BACKEND_FAILED) {
                UnloadRuntimeLibrary();
            }
            errorMessage = QStringLiteral("Persion人员检测推理失败(%1)：%2")
                               .arg(status)
                               .arg(runtimeError);
            return false;
        }
        if (output.task != YOLO_TASK_DETECT) {
            errorMessage = QStringLiteral("Persion模型任务不是目标检测任务：%1")
                               .arg(static_cast<int>(output.task));
            return false;
        }
        if (output.count < 0 || output.count > YOLO_MAX_RESULTS) {
            errorMessage = QStringLiteral("Persion检测结果数量异常：%1")
                               .arg(output.count);
            return false;
        }

        for (int index = 0; index < output.count; ++index) {
            const yolo_result_t& item = output.results[index];
            if (!std::isfinite(item.score) || item.score < 0.0f ||
                item.score > 1.0f || !std::isfinite(item.bbox.x) ||
                !std::isfinite(item.bbox.y) ||
                !std::isfinite(item.bbox.width) ||
                !std::isfinite(item.bbox.height) ||
                item.bbox.width <= 0.0f || item.bbox.height <= 0.0f) {
                continue;
            }

            Detection detection;
            detection.classId = item.class_id;
            detection.label = ResultLabel(item);
            detection.confidence = static_cast<double>(item.score);
            detection.boundingBox = QRectF(item.bbox.x,
                                           item.bbox.y,
                                           item.bbox.width,
                                           item.bbox.height);
            result.detections.push_back(detection);
        }
        return true;
    }

private:
    bool EnsureInitialized(QString& errorMessage)
    {
        if (m_runtime) {
            return true;
        }
#if !defined(Q_OS_LINUX) || !defined(Q_PROCESSOR_ARM_64)
        errorMessage = QStringLiteral(
            "Persion人员检测仅支持RK3588 Linux AArch64");
        return false;
#else
        QString configPath;
        QString packageRoot;
        if (!FindModelConfig(configPath, packageRoot, errorMessage) ||
            !LoadRuntimeLibrary(packageRoot, errorMessage)) {
            return false;
        }

        const int runtimeAbi = m_getAbiVersion();
        if (runtimeAbi != kPersionRuntimeAbiVersion) {
            errorMessage = QStringLiteral("Persion运行库ABI不匹配：期望%1，实际%2")
                               .arg(kPersionRuntimeAbiVersion)
                               .arg(runtimeAbi);
            UnloadRuntimeLibrary();
            return false;
        }

        const QByteArray encodedConfigPath = QFile::encodeName(configPath);
        const int status = m_create(encodedConfigPath.constData(), &m_runtime);
        if (status != YOLO_STATUS_OK || !m_runtime) {
            const QString runtimeError = RuntimeError();
            UnloadRuntimeLibrary();
            errorMessage = QStringLiteral("Persion模型初始化失败(%1)：%2；配置：%3")
                               .arg(status)
                               .arg(runtimeError)
                               .arg(configPath);
            return false;
        }
        if ((m_getCapabilities(m_runtime) & YOLO_CAPABILITY_INFER_IMAGE) == 0) {
            UnloadRuntimeLibrary();
            errorMessage = QStringLiteral("Persion运行库不支持内存图像推理");
            return false;
        }
        return true;
#endif
    }

    bool LoadRuntimeLibrary(const QString& packageRoot, QString& errorMessage)
    {
        QStringList candidates;
        const QString configuredLibrary = EnvironmentPath("PERSION_RUNTIME_LIBRARY");
        if (!configuredLibrary.isEmpty()) {
            AppendUniquePath(candidates, configuredLibrary);
        }
        AppendUniquePath(candidates, QDir(packageRoot).filePath(
            QStringLiteral("lib/libyolo_runtime.so.3")));
        AppendUniquePath(candidates, QDir(packageRoot).filePath(
            QStringLiteral("lib/libyolo_runtime.so")));
        AppendUniquePath(candidates, QStringLiteral("libyolo_runtime.so.3"));

        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("Persion运行库缺少必需的C ABI符号");
            UnloadRuntimeLibrary();
        }
        errorMessage = QStringLiteral("加载Persion运行库失败：%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_getCapabilities = reinterpret_cast<GetCapabilitiesFunction>(
            m_library.resolve("yolo_runtime_get_capabilities"));
        m_inferImage = reinterpret_cast<InferImageFunction>(
            m_library.resolve("yolo_runtime_infer_image"));
        m_lastError = reinterpret_cast<LastErrorFunction>(
            m_library.resolve("yolo_runtime_last_error"));
        m_destroy = reinterpret_cast<DestroyFunction>(
            m_library.resolve("yolo_runtime_destroy"));
        return m_getAbiVersion && m_create && m_getCapabilities &&
               m_inferImage && m_lastError && m_destroy;
    }

    QString RuntimeError() const
    {
        if (!m_runtime || !m_lastError) {
            return QStringLiteral("未知错误");
        }
        const char* error = m_lastError(m_runtime);
        const QString text = error ? QString::fromUtf8(error).trimmed() : QString();
        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_getCapabilities = nullptr;
        m_inferImage = nullptr;
        m_lastError = 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;
    GetCapabilitiesFunction m_getCapabilities = nullptr;
    InferImageFunction m_inferImage = nullptr;
    LastErrorFunction m_lastError = nullptr;
    DestroyFunction m_destroy = nullptr;
};

PersionModelDetector::PersionModelDetector()
    : m_impl(std::make_unique<Impl>())
{
}

PersionModelDetector::~PersionModelDetector() = default;

bool PersionModelDetector::Analyze(const QImage& frame,
                                   Analysis& result,
                                   QString& errorMessage)
{
    return m_impl && m_impl->Analyze(frame, result, errorMessage);
}