handleaiuinew/HanddleAI/HandleAiUi-new/handleAi/handleAiModel.h

#include <handleai.h>
#include <io.h>
#include <mutex>
#include <future>
#include <qobject.h>
#include <simpleTools/fileManager.h>
#include <unordered_map>
#include <qdebug.h>
#include <appCore.hpp>
#include <serviceModule.hpp>
class appConfig;
class handleAiModelSeg 
{
public:
    std::string path;
    int gpuIndex =0;
    std::mutex* gpuMutex;
    handleAiModelSeg(std::string pathIn, int gpuIndexIn, std::mutex* mutex)
    {
        path = pathIn;
        gpuIndex = gpuIndexIn;
        gpuMutex = mutex;
    };
    std::string init()
    {
        segPtr = handdleAI::CreateSegObj((char*)path.c_str(),gpuIndex);
        return path;
    }
    bool destroy()
    {
        if (segPtr == nullptr)
            return false;
        handdleAI::DestroySegObj(segPtr);
        return true;
    }
    cv::Mat Evaluate(cv::Mat& in)
    {
        gpuMutex->lock();
        cv::Mat res = segPtr->Evaluate(in);
        gpuMutex->unlock();
        return res;
    }
private:
    handdleAI::pSegmentation* segPtr = nullptr;
};
class handleAiModelCls
{
public:
    std::string path;
    int gpuIndex = 0;
    std::mutex* gpuMutex;
    handleAiModelCls(std::string pathIn, int gpuIndexIn, std::mutex* mutex)
    {
        path = pathIn;
        gpuIndex = gpuIndexIn;
        gpuMutex = mutex;
    };
    std::string init()
    {
        clsPtr = handdleAI::CreateClsObjV3((char*)path.c_str(), gpuIndex,1);
        return path;
    }
    bool destroy()
    {
        if (clsPtr == nullptr)
            return false;
        handdleAI::DestroyClsObj(clsPtr);
        return true;
    }
    int Evaluate(cv::Mat& in)
    {
        gpuMutex->lock();
        int res = clsPtr->Evaluate(in);
        gpuMutex->unlock();
        return res;
    }
private:
    handdleAI::pClassification* clsPtr;
};
class handleAiModelDet
{
public:
    std::string path;
    int gpuIndex = 0;
    std::mutex* gpuMutex;
    handleAiModelDet(std::string pathIn, int gpuIndexIn, std::mutex* mutex)
    {
        path = pathIn;
        gpuIndex = gpuIndexIn;
        gpuMutex = mutex;
    };
    std::string init()
    {
        detPtr = handdleAI::CreateDetObjS((char*)path.c_str(), 1, gpuIndex);
        return path;
    }
    bool destroy()
    {
        if (detPtr == nullptr || path=="")
            return false;
        handdleAI::DestroyDetObj(detPtr);
        return true;
    }
    bool Evaluate(cv::Mat& in, std::vector<handdleAI::Bbox>& result)
    {
        gpuMutex->lock();
        result=detPtr->Evaluate(in);
        gpuMutex->unlock();
        return true;
    }
private:
    handdleAI::pDetection* detPtr;
};
class handleAiModelManager :public serviceModule
{
public:
    enum type
    {
        seg, cls, det
    };
    bool initFlag = false;
    std::string basePath;
    handleAiModelManager() 
    {
        _core._appConfig.productDatasMutex.lock();
        this->basePath = _core._appConfig.handleAiModelPath + "/";
        _core._appConfig.productDatasMutex.unlock();
    };
    bool insertSegModelPath(std::string modelName,int gpuIndex)
    {
        auto mapResult = segModels.find(modelName);
        if (mapResult != segModels.end())//已有数据
            return true;
        handleAiModelSeg* segModel = new handleAiModelSeg(basePath + modelName, gpuIndex, gpuMutex[gpuIndex]);
        segModels.insert(std::make_pair(modelName, segModel));
    }
    bool insertDetModelPath(std::string modelName, int gpuIndex)
    {
        auto mapResult = detModels.find(modelName);
        if (mapResult != detModels.end())//已有数据
            return true;
        handleAiModelDet* detModel = new handleAiModelDet(basePath + modelName, gpuIndex, gpuMutex[gpuIndex]);
        detModels.insert(std::make_pair(modelName, detModel));
    }
    bool insertClsModelPath(std::string modelName, int gpuIndex)
    {
        auto mapResult = clsModels.find(modelName);
        if (mapResult != clsModels.end())//已有数据
            return true;
        handleAiModelCls* clsModel = new handleAiModelCls(basePath + modelName, gpuIndex, gpuMutex[gpuIndex]);
        clsModels.insert(std::make_pair(modelName, clsModel));
    }
    //
    handleAiModelSeg* getSegModel(std::string modelName)
    {
        auto mapResult = segModels.find(modelName);
        if (mapResult == segModels.end())//已有数据
            return nullptr;
        return  mapResult->second;
    }
    handleAiModelDet* getDetModel(std::string modelName)
    {
        auto mapResult = detModels.find(modelName);
        if (mapResult == detModels.end())//已有数据
            return nullptr;
        return mapResult->second;
    }
    handleAiModelCls* getClsModel(std::string modelName)
    {
        auto mapResult = clsModels.find(modelName);
        if (mapResult == clsModels.end())//已有数据
            return nullptr;
        return mapResult->second;
    }
    void initialization()
    {

        insertSegModelPath("APEX-Seg.BHD", 0);
        insertClsModelPath("APEX-Cls.BHD", 0);

    }
    void run()
    {
        if (init == true)
            return;
#if enableHandleAiModel
        std::vector<std::future<std::string>> futures;
        for (auto &model : segModels)
        {
            futures.push_back(std::async(std::launch::async, &handleAiModelSeg::init, model.second));
            emit _core._qtFunc.sendInitString(QString::fromStdString(fileManager::getFileName(model.second->path) + "初始化中..."));
        } 
        for (auto &model : clsModels)
        {
            futures.push_back(std::async(std::launch::async, &handleAiModelCls::init, model.second));
            emit _core._qtFunc.sendInitString(QString::fromStdString(fileManager::getFileName(model.second->path) + "初始化中..."));
        }
        for (auto &model : detModels)
        {
            futures.push_back(std::async(std::launch::async, &handleAiModelDet::init, model.second));
            emit _core._qtFunc.sendInitString(QString::fromStdString(fileManager::getFileName(model.second->path) + "初始化中..."));
        }

        for (auto& func : futures)
        {
            func.wait();
            std::string result = func.get();
            emit _core._qtFunc.sendInitString(QString::fromStdString(fileManager::getFileName(result) + "成功"));
            qDebug() << QString::fromStdString(fileManager::getFileName(result) + "成功");
            std::cout << fileManager::getFileName(result) + "成功" << std::endl;
        }
        init = true;
        status = InitOK;
#endif
    }
    ~handleAiModelManager()
    {
#if enableHandleAiModel
        for (auto& model : detModels)
            model.second->destroy();
        for (auto& model : segModels)
            model.second->destroy();
        for (auto& model : clsModels)
            model.second->destroy();
#endif
    }
private:
    std::unordered_map<std::string, handleAiModelSeg*> segModels;
    std::unordered_map<std::string, handleAiModelCls*> clsModels;
    std::unordered_map<std::string, handleAiModelDet*> detModels;
    std::vector<std::mutex*>gpuMutex = std::vector<std::mutex*>(6, new std::mutex());
    int gpuCnt = 1;//gpu数量
    bool init = false;
};

#pragma once