yolov5+bytetrack算法在华为NPU上进行端到端开发

        自从毕业后开始进入了华为曻腾生态圈,现在越来越多的公司开始走国产化路线了,现在国内做AI芯片的厂商比如:寒武纪、地平线等,虽然我了解的不多,但是相对于瑞芯微这样的AI开发板来说,华为曻腾的生态比瑞芯微好太多了,参考文档非常多,学习资料也有很多,也容易上手开发。

华为曻腾官网:昇腾AI应用案例-昇腾社区 (hiascend.com)

        直接步入正题,现在的目标检测已经很成熟了,所以越来越多的公司会用到基于检测的跟踪算法,这样不仅起到了单一检测功能,还有跟踪目标或者计数的功能;

        现在应用较广泛的目标检测算法从最开始的yolov5一直到现在的yolov8,虽然只是简单的看了一下算法的原理,整体来说yolo的更新还是针对神经网络在GPU上的优化加速,而对比曻腾NPU,yolov5的速度还是在其他yolo算法中速度最快的一个;

        目标跟踪算法以前是sort+yolo,deepsort+yolo,bytetrack,fairmot等算法,本章主要介绍如何利用华为的ACL语言+ffmpeg推流进行整个业务的开发流程,大家可以借鉴下面的开发代码,首先你要具备基本的ACL语言知识,以及yolov5的后处理逻辑,跟踪方面直接借鉴开源作者的卡尔曼滤波进行预测更新即可:参考主函数代码如下:

//1.先测试yolov5_nms可以泡桐?
//使用dvpp+aipp编解码再使用opencv进行
 #include
 #include"acl/acl.h"
#include "opencv2/opencv.hpp"
#include "opencv2/imgproc/types_c.h"
#include "acllite/AclLiteUtils.h"
#include "acllite/AclLiteError.h"
#include "acllite/AclLiteResource.h"
#include "acllite/AclLiteModel.h"
#include "acllite/AclLiteImageProc.h"
#include "AclLiteVideoProc.h"
#include "AclLiteVideoCapBase.h"
#include "BYTETracker.h"
#include 
extern"C" {
	#include 
	#include 
	#include "libavcodec/avcodec.h"
	#include "libavformat/avformat.h"
	#include "libswscale/swscale.h"
	#include "libavutil/imgutils.h"
	#include "libavutil/opt.h"
};
using namespace std;
using namespace cv;
typedef struct box {
    float x;
    float y;
    float w;
    float h;
    float score;
    size_t classIndex;
    size_t index; // index of output buffer
} box;
namespace{
    int a  = 0;
}
int main()
{
    //1.定义初始化变量dvpp\model\acl\rtsp解码接口cap
    AclLiteResource aclDev;

    aclrtRunMode g_runMode_;

    AclLiteVideoProc* cap_;

    AclLiteImageProc g_dvpp_;

    AclLiteModel g_model_;

    string streamName_;

    streamName_ = "rtsp://admin:[email protected]/LiveMedia/ch1/Media1";
    //ffmpeg初始化
       AVFormatContext* g_fmtCtx;
   AVCodecContext* g_codecCtx;
   AVStream* g_avStream;
   AVCodec* g_codec;
   AVPacket* g_pkt;
	AVFrame* g_yuvFrame;
	uint8_t* g_yuvBuf;
	AVFrame* g_rgbFrame;
	uint8_t* g_brgBuf;
	int g_yuvSize;
	int g_rgbSize;
	struct SwsContext* g_imgCtx;
//参数初始化
//rtsp初始化

   g_avStream = NULL;
   g_codec = NULL;
   g_codecCtx = NULL;
   g_fmtCtx = NULL;
   g_pkt  = NULL;
   g_imgCtx = NULL;
   g_yuvSize = 0;
	g_rgbSize = 0;
   int picWidth = 416;
   int picHeight = 416;
   string rtsp_url = "rtsp://192.168.3.38:8554/stream";
   int channelId = 0;
   string g_outFile = rtsp_url + to_string(channelId);
//rtsp初始化
   avformat_network_init();

   if (avformat_alloc_output_context2(&g_fmtCtx, NULL, g_avFormat.c_str(), g_outFile.c_str()) < 0) {
       ACLLITE_LOG_ERROR("Cannot alloc output file context");
       return ACLLITE_ERROR;
   }
   av_opt_set(g_fmtCtx->priv_data, "rtsp_transport", "tcp", 0);
   av_opt_set(g_fmtCtx->priv_data, "tune", "zerolatency", 0);
   av_opt_set(g_fmtCtx->priv_data, "preset", "superfast", 0);
   //获取编码器的ID返回一个编码器
   g_codec = avcodec_find_encoder(AV_CODEC_ID_H264);
   if (g_codec == NULL) {
       ACLLITE_LOG_ERROR("Cannot find any endcoder");
       return ACLLITE_ERROR;
   }

   g_codecCtx = avcodec_alloc_context3(g_codec);
   if (g_codecCtx == NULL) {
       ACLLITE_LOG_ERROR("Cannot alloc context");
       return ACLLITE_ERROR;
   }
   //创建流
   g_avStream = avformat_new_stream(g_fmtCtx, g_codec);
   if (g_avStream == NULL) {
       ACLLITE_LOG_ERROR("failed create new video stream");
       return ACLLITE_ERROR;
   }
   //设置帧率
   g_avStream->time_base = AVRational{1, g_frameRate};
   //设置编码参数
   AVCodecParameters* param = g_fmtCtx->streams[g_avStream->index]->codecpar;
   param->codec_type = AVMEDIA_TYPE_VIDEO;
   param->width = picWidth;
   param->height = picHeight;

   avcodec_parameters_to_context(g_codecCtx, param);
   //参数绑定设置
   g_codecCtx->pix_fmt = AV_PIX_FMT_NV12;
   g_codecCtx->time_base = AVRational{1, g_frameRate};
   g_codecCtx->bit_rate = g_bitRate;
   g_codecCtx->gop_size = g_gopSize;
   g_codecCtx->max_b_frames = 0;

   if (g_codecCtx->codec_id == AV_CODEC_ID_H264) {
       g_codecCtx->qmin = 10;
       g_codecCtx->qmax = 51;
       g_codecCtx->qcompress = (float)0.6;
   }

   if (g_codecCtx->codec_id == AV_CODEC_ID_MPEG1VIDEO)
       g_codecCtx->mb_decision = 2;
       //初始化code
   if (avcodec_open2(g_codecCtx, g_codec, NULL) < 0) {
       ACLLITE_LOG_ERROR("Open encoder failed");
       return ACLLITE_ERROR;
   }
   //g_codecCtx参数传递给codecpar
   avcodec_parameters_from_context(g_avStream->codecpar, g_codecCtx);
   //指定输出数据的形式
   av_dump_format(g_fmtCtx, 0, g_outFile.c_str(), 1);
   //写文件头
   int ret1 = avformat_write_header(g_fmtCtx, NULL);
   if (ret1 != AVSTREAM_INIT_IN_WRITE_HEADER) {
       ACLLITE_LOG_ERROR("Write file header fail");
       return ACLLITE_ERROR;
   }
   g_pkt = av_packet_alloc();



//传输数据初始化
       g_rgbFrame = av_frame_alloc();
       g_yuvFrame = av_frame_alloc();
       g_rgbFrame->width = g_codecCtx->width;
       g_yuvFrame->width = g_codecCtx->width;
       g_rgbFrame->height = g_codecCtx->height;
       g_yuvFrame->height = g_codecCtx->height;
       g_rgbFrame->format = AV_PIX_FMT_BGR24;
       g_yuvFrame->format = g_codecCtx->pix_fmt;

       g_rgbSize = av_image_get_buffer_size(AV_PIX_FMT_BGR24, g_codecCtx->width, g_codecCtx->height, 1);
       g_yuvSize = av_image_get_buffer_size(g_codecCtx->pix_fmt, g_codecCtx->width, g_codecCtx->height, 1);

       g_brgBuf = (uint8_t*)av_malloc(g_rgbSize);
       g_yuvBuf = (uint8_t*)av_malloc(g_yuvSize);


       //内存分配
       int ret2 = av_image_fill_arrays(g_rgbFrame->data, g_rgbFrame->linesize,
           g_brgBuf, AV_PIX_FMT_BGR24,
           g_codecCtx->width, g_codecCtx->height, 1);

       ret2 = av_image_fill_arrays(g_yuvFrame->data, g_yuvFrame->linesize,
           g_yuvBuf, g_codecCtx->pix_fmt,
           g_codecCtx->width, g_codecCtx->height, 1);
       g_imgCtx = sws_getContext(
           g_codecCtx->width, g_codecCtx->height, AV_PIX_FMT_BGR24,
           g_codecCtx->width, g_codecCtx->height, g_codecCtx->pix_fmt,
           SWS_BILINEAR, NULL, NULL, NULL);
     //2.类变量初始化
    AclLiteError ret = aclDev.Init();
    if (ret) {
        ACLLITE_LOG_ERROR("Init resource failed, error %d", ret);
        return ACLLITE_ERROR;
    }

    if (ACLLITE_OK != OpenVideoCapture()) {
        return ACLLITE_ERROR;
    }

    ret = g_dvpp_.Init();
    if (ret) {
        ACLLITE_LOG_ERROR("Dvpp init failed, error %d", ret);
        return ACLLITE_ERROR;
    }

    cap_ = nullptr;

    ret = g_model_.Init();
    if (ret) {
        ACLLITE_LOG_ERROR("Model init failed, error %d", ret);
        return ACLLITE_ERROR;
    }
    //3.创建模型img_info的输入以及数据拷贝操作
    g_runMode_ = g_aclDev_.GetRunMode();

    const float imageInfo[4] = {(float)g_modelInputWidth, (float)g_modelInputHeight,
                            (float)g_modelInputWidth, (float)g_modelInputHeight};

    g_imageInfoSize_ = sizeof(imageInfo);

    g_imageInfoBuf_ = CopyDataToDevice((void *)imageInfo, g_imageInfoSize_,
                                    g_runMode_, MEMORY_DEVICE);

    if (g_imageInfoBuf_ == nullptr) {
    ACLLITE_LOG_ERROR("Copy image info to device failed");
    return ACLLITE_ERROR;
    }
    //4.获取视频源
    cap_ = new AclLiteVideoProc(streamName_);
    //5.视频流解码以及dvpp硬件-resize
    int i =0;
    while(true)
    {
        //6.获取解码图片(在device侧的YUV420图片)(存放在ImageDta结构体中)
//         struct ImageData {
//     acldvppPixelFormat format;
//     uint32_t width = 0;
//     uint32_t height = 0;
//     uint32_t alignWidth = 0;
//     uint32_t alignHeight = 0;
//     uint32_t size = 0;
//     std::shared_ptr data = nullptr;
// };
i++;
        ImageData image;
        ret = cap_->Read(image);
        ImageData resizedImage;
        ret = g_dvpp_.Resize(resizedImage, image, 640, 640);
        //7.创建模型输入进行模型推理
                ret = g_model_.CreateInput(resizedImage.data.get(), resizedImage.size,
                                   g_imageInfoBuf_, g_imageInfoSize_);
        if (ret != ACLLITE_OK) {
        ACLLITE_LOG_ERROR("Create mode input dataset failed, error:%d", ret);
        return ACLLITE_ERROR;
        }
        std::vector inferenceOutput;
        ret = g_model_.Execute(inferenceOutput);
        if (ret != ACLLITE_OK) {
            g_model_.DestroyInput();
            ACLLITE_LOG_ERROR("Execute model inference failed, error: %d", ret);
            return ACLLITE_ERROR;
        }
        g_model_.DestroyInput();
        //8.将YUV图像转换为opencv图像
        ImageData yuvImage;
        ret = CopyImageToLocal(yuvImage, image, g_runMode_);
        if (ret == ACLLITE_ERROR) {
        ACLLITE_LOG_ERROR("Copy image to host failed");
        return ACLLITE_ERROR;
        }

        cv::Mat yuvimg(yuvImage.height * 3 / 2, yuvImage.width, CV_8UC1, yuvImage.data.get());
        cv::Mat origImage;
        cv::cvtColor(yuvimg, origImage, CV_YUV2BGR_NV12);
        //模型后处理(根据目标跟踪需要的输入进行获取xywh)
        float* detectData = (float *)inferenceOutput[0].data.get();
        float* boxNum = (float *)inferenceOutput[1].data.get();
        uint32_t totalBox = boxNum[0];
        //获取(x,y,w,h) 
        std::vector obj;
    float widthScale = (float)(origImage.cols) / 640.0;
    float heightScale = (float)(origImage.rows) / 640.0;
    vector detectResults;
    for (uint32_t i = 0; i < totalBox; i++) {
        box boundBox;
        boundBox.score = float(detectData[totalBox * SCORE + i]);
        boundBox.x = detectData[totalBox * TOPLEFTX + i] * widthScale;
        boundBox.y = detectData[totalBox * TOPLEFTY + i] * heightScale;
        boundBox.w = detectData[totalBox * BOTTOMRIGHTX + i] * widthScale;
        boundBox.h = detectData[totalBox * BOTTOMRIGHTY + i] * heightScale;
        boundBox.classIndex = (uint32_t)detectData[totalBox * LABEL + i];
        detectResults.emplace_back(boundBox);
    }
      for (size_t i = 0; i < detectResults.size(); i++){

              if (res[i].classId != class_id){ continue; }
              obj[i].label = detectResults[i].classIndex;
              obj[i].rect.x = detectResults[i].x;
              obj[i].rect.y = detectResults[i].y;
              obj[i].rect.height = detectResults[i].h;
              obj[i].rect.width = detectResults[i].w;
              obj[i].prob = detectResults[i].score;
    
      }


        std::vector output_stracks = tracker.update(obj);
      for (size_t i = 0; i < output_stracks.size(); i++){
          std::vector tlwh = output_stracks[i].tlwh;
          cv::Scalar __color = tracker.get_color(output_stracks[i].track_id);
          cv::putText(origImage, std::to_string(output_stracks[i].track_id), cv::Point(tlwh[0], tlwh[1] - 10), cv::FONT_ITALIC, 0.75, __color, 2);
          cv::rectangle(origImage, cv::Rect(tlwh[0], tlwh[1], tlwh[2], tlwh[3]), __color, 2);    
      }
    //跟踪完成后写推流
        memcpy(g_brgBuf, origImage.data, g_rgbSize);
    sws_scale(g_imgCtx,
        g_rgbFrame->data,
        g_rgbFrame->linesize,
        0,
        g_codecCtx->height,
        g_yuvFrame->data,
        g_yuvFrame->linesize);
    g_yuvFrame->pts = i;
    if (avcodec_send_frame(g_codecCtx, g_yuvFrame) >= 0) {
        // cout<= 0) {
            cout<<"avcodec_receive_packet"<stream_index = g_avStream->index;
            av_packet_rescale_ts(g_pkt, g_codecCtx->time_base, g_avStream->time_base);
            g_pkt->pos = -1;
            int ret = av_interleaved_write_frame(g_fmtCtx, g_pkt);

            if (ret < 0) {
                ACLLITE_LOG_ERROR("error is: %d", ret);
            }
        }
    }


    }

    av_packet_free(&g_pkt);
   avcodec_close(g_codecCtx);
    if (g_fmtCtx) {
        avio_close(g_fmtCtx->pb);
        avformat_free_context(g_fmtCtx);
    }
     if (cap_ != nullptr) {
            cout << "cap is not open" << endl;
        cap_->Close();
        delete cap_;
    }
    dvpp_.DestroyResource();











    return 0;
} 
  

跟踪器方面的函数,可以搜索开源代码yolov5-bytetrack-main.cpp截取内部跟踪部分,检测部分使用华为ACL编写的推理代码进行检测;

可以加入学习讨论:1076799627

你可能感兴趣的:(YOLO,人工智能,深度学习)