HeadCountAlgo/detector/components/yolocrowd/source/YoloCrowdApp.cpp

#include "YoloCrowdApp.h"

#include "ilogger.hpp"
#include "trt_infer.hpp"

NAMESPACE_MAS_BEGIN

NAMESPACE_YOLOCROWD_BEGIN

tzc::Mutex YoloCrowdApp::_mutex;
YoloCrowdApp* YoloCrowdApp::_instance = nullptr;

YoloCrowdApp* YoloCrowdApp::Instance()
{
    if (!_instance)
	{
        _mutex.Lock();
        if (!_instance)
		{
            _instance = new YoloCrowdApp();
            TZLogInfo("YoloCrowdApp Instance created~~~");
        }
        _mutex.Unlock();
    }

    ++(_instance->m_usecnt);
    return _instance;
}

void YoloCrowdApp::DestroyInstance()
{
    tzc::ScopedLock lock(_mutex);
    if ((_instance->m_usecnt) < 0)
	{
        TZLogWarn("YoloCrowdApp Instance does not exist!!!");
		_instance->m_usecnt = 0;
    }
    else if ((_instance->m_usecnt) == 0)
	{
        TZ_delete(_instance);
        TZLogInfo("YoloCrowdApp Instance destroyed~~~");
    }
    else
	{
        TZLogWarn("YoloCrowdApp still in use by %d objects!", _instance->m_usecnt);
		TZ_delete(_instance);
		TZLogInfo("YoloCrowdApp Instance destroyed~~~");
    }
}

TZ_INT YoloCrowdApp::Initialize(const std::string& initParam)
{
    if (m_inited)
	{
        TZLogInfo("YoloCrowdApp has been initialized~~~");
        return MEC_OK;
    }

    TRT::set_device(m_gpuid);

    YoloCrowdBuild param;
    YoloCrowdBuild::fromJson(initParam, param);

    m_gpuid = param.gpu_id;
    m_maxObjects = param.max_objects;
    m_confidenceThreshold = param.confidence_threshold;
    m_NMSThreshold = param.nms_threshold;
    m_modelPath = param.model_path;

    if (!iLogger::exists(m_modelPath))
	{
        TZLogWarn("YoloCrowd model file does not exist!!!");
        return MEC_FAILED;
    }

    m_engine = YoloCrowd::create_infer(
        m_modelPath,
        YoloCrowd::Type::V5,
        m_gpuid,
        m_confidenceThreshold,
        m_NMSThreshold,
        YoloCrowd::NMSMethod::FastGPU,
        m_maxObjects,
        false
    );

    if (m_engine == nullptr)
	{
        TZLogWarn("YoloCrowd Engine is nullptr!!!");
        return MEC_NULL_OBJ;
    }

    // this->warmup();

    m_inited = TRUE;

    TZLogInfo("YoloCrowdApp Initialized~~~");

    return MEC_OK;
}

TZ_INT YoloCrowdApp::Dispose()
{
    if (!m_inited)
	{
        return MEC_NOT_INITED;
    }

    _mutex.Lock();

    --(m_usecnt);
    TZLogInfo("YoloCrowdApp usecnt -1, now = %d~~~", m_usecnt);
    TZ_BOOL empty = (m_usecnt == 0);

    _mutex.Unlock();

    if (empty)
	{
        m_inited = FALSE;
        m_engine.reset();
        TZLogInfo("YoloCrowdApp Dispose~~~");
        this->DestroyInstance();
    }

    return MEC_OK;
}

TZ_INT YoloCrowdApp::DoDetect(cv::Mat& input,
                         TargetInfo& targets,
                         std::vector<int>& objClass,
                         std::vector<std::vector<float>>& objPos)
						 {
    auto boxes = m_engine->commit(input).get();

    for (auto& obj : boxes)
	{
        if (obj.confidence < targets.target_threshold) continue;

        objClass.push_back(obj.class_label);
        objPos.push_back({obj.left, obj.top, 
			obj.right, obj.bottom, obj.confidence});
    }

    return MEC_OK;
}

YoloCrowdApp::YoloCrowdApp()
    : m_inited(FALSE),
      m_usecnt(0),
      m_gpuid(0),
      m_maxObjects(1024),
      m_confidenceThreshold(0.10),
      m_NMSThreshold(0.45),
      m_modelPath(""),
      m_engine(nullptr) {}

YoloCrowdApp::~YoloCrowdApp()
{
    this->Dispose();
}

void YoloCrowdApp::warmup()
{
    auto files = iLogger::find_files(PROGRAM_ROOT_PATH "models/yolo_warmup/", "*.jpg;*.jpeg;*.png");

    std::vector<cv::Mat> images;
    for (size_t i = 0; i < files.size(); ++i)
	{
        auto image = cv::imread(files[i]);
        images.emplace_back(image);
    }

    std::vector<std::shared_future<YoloCrowd::BoxArray>> boxes_array;
    for (size_t i = 0; i < 10; ++i)
	{
        boxes_array = m_engine->commits(images);
    }

	boxes_array.back().get();
    boxes_array.clear();
}

NAMESPACE_YOLOCROWD_END

NAMESPACE_MAS_END