Hu Chunming / mvpt_dq_ascend_arm

#include "./interface/DecoderManager.h"
#include <mutex>
#include <thread>
#include <chrono>
#include "./interface/logger.hpp"

#include "vpt.h"
#include "muxer.h"

#include "acl/acl.h"
#include "acl/ops/acl_dvpp.h"

#include <opencv2/opencv.hpp>

using namespace std;
using namespace cv;

struct decode_cbk_userdata{
  string task_id;
  void* opaque;
  void* opaque1;
};

deque<DeviceMemory*> m_RgbDataList;
mutex m_DataListMtx; 

thread* m_pAlgorthimThread{nullptr};

bool m_bfinish{false};
int m_devId = 0;
const char* task_id = "test0";
int skip_frame_ = 5;
int m_batch_size = 20;


void algorthim_process_thread();

void post_decod_cbk(const void * userPtr, DeviceMemory* devFrame){
    do{
        if(m_bfinish){
            break;
        }
        m_DataListMtx.lock();
        if(m_RgbDataList.size() >= 30){
            m_DataListMtx.unlock();
            std::this_thread::sleep_for(std::chrono::milliseconds(3));
            continue;
        }
        m_RgbDataList.push_back(devFrame);
        m_DataListMtx.unlock();
        break;
    }while (true);
}

void decode_finished_cbk(const void * userPtr){
    decode_cbk_userdata* ptr = (decode_cbk_userdata*)userPtr;
    if (ptr!= nullptr){
        printf("task finished: %s \n", ptr->task_id.c_str());
        delete ptr;
        ptr = nullptr;
    }
}

bool create_task(string task_id) {
    
    // 创建解码任务
    DecoderManager* pDecManager = DecoderManager::getInstance();

    MgrDecConfig config;
    config.name = task_id;
    config.cfg.uri = "/home/cmhu/data/爱剪辑-高清事件测试.mp4";
    config.cfg.post_decoded_cbk = post_decod_cbk;
    config.cfg.decode_finished_cbk = decode_finished_cbk;
    config.cfg.force_tcp = true;  // rtsp用tcp
    config.cfg.gpuid = to_string(m_devId);
    config.cfg.skip_frame = skip_frame_;

    config.dec_type = DECODER_TYPE_DVPP;
    
    AbstractDecoder* dec = pDecManager->createDecoder(config);
    if (!dec){
        printf("创建解码器失败 \n");
        return false;
    }

    // decode_cbk_userdata* userPtr = new decode_cbk_userdata;
    // userPtr->task_id = string(task_id);
    // pDecManager->setPostDecArg(config.name, userPtr);
    // pDecManager->setFinishedDecArg(config.name, userPtr);

    dec->setSnapTimeInterval(1000);


    int input_image_width = 0;
    int input_image_height = 0;
    pDecManager->getResolution(config.name, input_image_width, input_image_height);

    pDecManager->startDecodeByName(config.name);
    // pDecManager->closeAllDecoder();

    return true;
}

void* handle = nullptr;
aclrtContext ctx;
int m_max_batchsize = 16;

muxer* pEncoder = nullptr;

int main(){

    set_default_logger(LogLevel(0), "test_decode","logs/main.log", 64 * 1024 * 1024, 64 * 1024 * 1024);
    LOG_INFO("编译时间：{} {}", __DATE__, __TIME__);

    printf("start... \n");


    vpt_param param;
    param.modelNames_b = "./models/vpt/vpt230323_b16_bgr_310p.om";
    param.threshold = 0.35;
    param.devId = 0;
    param.max_batch = m_max_batchsize;
    param.isTrk = false;

    cout << "init start " << endl;
    aclInit(nullptr);
    aclrtSetDevice(param.devId);
    aclrtCreateContext(&ctx, param.devId);

    int ret = vpt_init(&handle, param);
    if (ret != 0) {
        cout << "init vpt failed! " << endl;
        return -1;
    }

    DecoderManager* pDecManager = DecoderManager::getInstance();


    // 创建算法线程
    m_pAlgorthimThread = new thread([](void* arg) {
      algorthim_process_thread();
      return (void*)0;
    }
    , nullptr);


    

    char ch = 'a';
    int task_id = 1;
    create_task(to_string(task_id));  
    task_id++;

    while (ch != 'q') {
        ch = getchar();
        switch (ch)
        {
        case '4':
            for (size_t i = 0; i < 4; i++) 
            {
                create_task(to_string(task_id));  
                task_id++;
                std::this_thread::sleep_for(std::chrono::seconds(1));
            }
            break;
        case 'a':
                create_task(to_string(task_id));  
                task_id++;
                std::this_thread::sleep_for(std::chrono::seconds(1));
            break;
        case 'c':
            pDecManager->closeAllDecoder();
            break;
        default:
            break;
        }
    }

    if (pEncoder != nullptr) {
        pEncoder->close();
        delete pEncoder;
    }
}

void do_work(vector<DeviceMemory*> vec_gpuMem){
    int batchsize = vec_gpuMem.size();
	vector <vector< vector <float>>> detectResult(batchsize);
	int real_index = 0;

	int cycle_time = batchsize / m_max_batchsize;
	cycle_time = (batchsize % m_max_batchsize) == 0 ? cycle_time : (cycle_time + 1);
	vector<sy_img> vec_img;
	for (int i = 0; i < cycle_time; i++) {
		vec_img.clear();

		int start_index = i * m_max_batchsize;
		int end_index = start_index + m_max_batchsize;
		if (end_index >= batchsize) {
			end_index = batchsize;
		}

		for (int j = start_index; j < end_index; j++){
			DeviceMemory* gpuMem = vec_gpuMem[j];
			sy_img img;
			img.set_data(gpuMem->getWidth(), gpuMem->getHeight(), gpuMem->getChannel(), gpuMem->getMem());
			vec_img.push_back(img);
		}

		vpt_result detresult[vec_img.size()];
		for (int b = 0; b < vec_img.size(); b++) {
			detresult[b].obj_results_ = new vpt_obj_result[MAX_DET_COUNT];
		}

 		int res_status = vpt_batchV2(handle, vec_img.data(), vec_img.size(), detresult);

		for (size_t b = 0; b < vec_img.size(); b++) {
			auto& cur_result = detresult[b];

            sy_img image = vec_img[b];
            int data_size = image.w_ * image.h_ * 3;
            unsigned char* pData = (unsigned char*) malloc (data_size);
            int ret = aclrtMemcpy(pData, data_size, image.data_, data_size, ACL_MEMCPY_DEVICE_TO_HOST);
            if(ret != ACL_ERROR_NONE){
                free(pData);
                cout << "aclrtMemcpy failed" << endl;
                continue;
            }

            cv::Mat cvImg(image.h_, image.w_ , CV_8UC3, pData);

            for (int i = 0; i <cur_result.obj_count_; i++) {
                Point lt(cur_result.obj_results_[i].obj_rect.left_,cur_result.obj_results_[i].obj_rect.top_);
                Point rb((cur_result.obj_results_[i].obj_rect.left_ + 
               cur_result.obj_results_[i].obj_rect.width_), (cur_result.obj_results_[i].obj_rect.top_ + 
               cur_result.obj_results_[i].obj_rect.height_));
                rectangle(cvImg, lt, rb, cv::Scalar(0, 0, 255), 4);

                char buffer[50];
                int fontface = cv::FONT_HERSHEY_SIMPLEX;
                double scale = 0.8;
                int thickness = 2;
                int baseline = 0;
                snprintf(buffer, sizeof(buffer), "%d:%.2f",cur_result.obj_results_[i].obj_index,cur_result.obj_results_[i].obj_score);
                cv::Size text = cv::getTextSize(buffer, fontface, scale, thickness, &baseline);
                cv::putText(cvImg, buffer, lt - cv::Point(0, baseline), fontface,
                        scale, cv::Scalar(0, 0, 255), thickness, 4);

            }

            if (pEncoder == nullptr)
            {
                pEncoder = new muxer();
                pEncoder->init(cvImg.cols, cvImg.rows, 25, 4992000, "./vpt.mp4");
            }
            
            pEncoder->write_image(cvImg.data);

            free(pData);

			delete [] cur_result.obj_results_;

			real_index++;
		}
	}
}

void algorthim_process_thread(){

    while (true){
        if(m_bfinish){
            break;
        }

        vector<DeviceMemory*> vec_gpuMem;
        m_DataListMtx.lock();
        while (!m_RgbDataList.empty()){
            DeviceMemory* gpuMem = m_RgbDataList.front();
            if(gpuMem->getMem() == nullptr){
              // 错误数据，直接删除
              printf("mem is null \n");
            } else {
              vec_gpuMem.push_back(gpuMem);
            }
            m_RgbDataList.pop_front();
            if(vec_gpuMem.size() >= m_batch_size){
                break;
            }
        }
        m_DataListMtx.unlock();

        if(vec_gpuMem.size() <= 0){
            std::this_thread::sleep_for(std::chrono::milliseconds(3));
            continue;
        }

        aclrtSetDevice(0);
	    aclrtSetCurrentContext(ctx);
        // do work
        do_work(vec_gpuMem);

        for (size_t i = 0; i < vec_gpuMem.size(); i++) {
            DeviceMemory* gpuMem = vec_gpuMem[i];
            delete gpuMem;
            gpuMem = nullptr;
        }

        vec_gpuMem.clear();
    }

    printf("algorthim_process_thread exit. \n");
}