BM1688实时yolov5识别

前言

主要实现的功能是rtsp拉流,然后硬件解码,再做yolo识别,编码推流出去

配置ZLMediaKit

这是一个rtsp推流的服务,具体配置网上一搜都是

源码

首先需要配置sail 一些库,按照sophon的手册搭建

CMakeLists.txt

cmake_minimum_required(VERSION 3.10)
project(yolo_rtsp_pipeline)

SET(EXECUTABLE_OUTPUT_PATH ${PROJECT_SOURCE_DIR})

set(CMAKE_CXX_FLAGS "-O3 -fopenmp")

if (NOT DEFINED TARGET_ARCH)
    set(TARGET_ARCH pcie)
endif()

if (${TARGET_ARCH} STREQUAL "pcie")
    find_package(OpenCV REQUIRED)
    include_directories(${OpenCV_INCLUDE_DIRS})

    find_package(libsophon REQUIRED)
    include_directories(${LIBSOPHON_INCLUDE_DIRS})
    link_directories(${LIBSOPHON_LIB_DIRS})

    # Use sail
    set(SAIL_DIR /opt/sophon/sophon-sail/lib/cmake)
    find_package(SAIL REQUIRED)
    include_directories(${SAIL_INCLUDE_DIRS})
    link_directories(${SAIL_LIB_DIRS})

    set(FFMPEG_DIR /opt/sophon/sophon-ffmpeg-latest/lib/cmake)
    find_package(FFMPEG REQUIRED)
    include_directories(${FFMPEG_INCLUDE_DIRS})
    link_directories(${FFMPEG_LIB_DIRS})

    set(OpenCV_DIR /opt/sophon/sophon-opencv-latest/lib/cmake/)
    find_package(OpenCV REQUIRED)
    include_directories(${OpenCV_INCLUDE_DIRS})

    message(STATUS "SAIL_INCLUDE_DIRS: ${SAIL_INCLUDE_DIRS}")
    message(STATUS "SAIL_LIB_DIRS: ${SAIL_LIB_DIRS}")
    message(STATUS "FFMPEG_INCLUDE_DIRS: ${FFMPEG_INCLUDE_DIRS}")
    message(STATUS "OPENCV_INCLUDE_DIRS: ${OpenCV_INCLUDE_DIRS}")

    aux_source_directory(. SRC_FILES)
    add_executable(yolo_rtsp_pipeline.pcie ${SRC_FILES})
    target_link_libraries(yolo_rtsp_pipeline.pcie ${FFMPEG_LIBS} ${OpenCV_LIBS} sail)

elseif (${TARGET_ARCH} STREQUAL "soc")
    # Cross-compile using buildroot aarch64 toolchain
    if (DEFINED ENV{GCC_COMPILER})
        set(CROSS_PREFIX $ENV{GCC_COMPILER})
    else()
        set(CROSS_PREFIX /home/jiangguo/aarch64-buildroot-linux-gnu_sdk-buildroot/bin/aarch64-buildroot-linux-gnu)
    endif()
    set(CMAKE_C_COMPILER   ${CROSS_PREFIX}-gcc)
    set(CMAKE_ASM_COMPILER ${CROSS_PREFIX}-gcc)
    set(CMAKE_CXX_COMPILER ${CROSS_PREFIX}-g++)
    set(CMAKE_SYSROOT $ENV{SYSROOT_PATH})

    set(BM_LIBS bmlib bmrt bmcv bmjpeg bmvenc bmvd yuv)
    set(OPENCV_LIBS opencv_imgproc opencv_core opencv_highgui opencv_imgcodecs opencv_videoio)
    set(FFMPEG_LIBS avcodec avdevice avfilter avformat avutil swresample swscale)
    set(OPENCV_EXTRA opencv_calib3d opencv_dnn opencv_features2d opencv_flann opencv_ml
        opencv_objdetect opencv_photo opencv_stitching opencv_video)

    set(BUILD_SAIL /home/jiangguo/v2.2_release/sophon-sail/sophon-sail/build_soc/sophon-sail)

    message("CROSS_PREFIX: ${CROSS_PREFIX}")
    message("SOC-SDK: ${SDK}")
    include_directories("${SDK}/include/")
    include_directories("${SDK}/include/opencv4/")
    include_directories("${BUILD_SAIL}/include/")
    include_directories("${BUILD_SAIL}/include/sail/")   # for spdlog internal includes
    link_directories("${SDK}/lib/")
    link_directories("${BUILD_SAIL}/lib/")

    aux_source_directory(. SRC_FILES)
    add_executable(yolo_rtsp_pipeline.soc ${SRC_FILES})
    target_link_libraries(yolo_rtsp_pipeline.soc
        -Wl,--start-group sail ${BM_LIBS} ${OPENCV_LIBS} ${OPENCV_EXTRA} ${FFMPEG_LIBS}
        -Wl,--end-group
        -Wl,-rpath-link,${SDK}/lib pthread dl)

elseif (${TARGET_ARCH} STREQUAL "soc_native")
    # Native compile directly on SOC device (no cross-compiler needed)
    set(BM_LIBS bmlib bmrt bmcv bmjpeg bmvenc bmvd yuv)
    set(OPENCV_LIBS opencv_imgproc opencv_core opencv_highgui opencv_imgcodecs opencv_videoio)
    set(FFMPEG_LIBS avcodec avdevice avfilter avformat avutil swresample swscale)
    set(OPENCV_EXTRA opencv_calib3d opencv_dnn opencv_features2d opencv_flann opencv_ml
        opencv_objdetect opencv_photo opencv_stitching opencv_video)

    # Use system-installed libsophon/sail/ffmpeg/opencv on SOC
    include_directories(/opt/sophon/sophon-sail/include)
    include_directories(/opt/sophon/libsophon-current/include)
    include_directories(/opt/sophon/libsophon-current/include/opencv4)
    include_directories(/opt/sophon/sophon-sail/include/sail)
    include_directories(/opt/sophon/sophon-ffmpeg-latest/include)
    include_directories(/opt/sophon/sophon-opencv-latest/include/opencv4)
    include_directories(/opt/sophon/include)
    link_directories(/opt/sophon/sophon-sail/lib)
    link_directories(/opt/sophon/sophon-ffmpeg-latest/lib)
    link_directories(/opt/sophon/sophon-opencv-latest/lib)
    link_directories(/opt/sophon/libsophon-current/lib)

    aux_source_directory(. SRC_FILES)
    add_executable(yolo_rtsp_pipeline.soc ${SRC_FILES})
    target_link_libraries(yolo_rtsp_pipeline.soc
        -Wl,--start-group sail ${BM_LIBS} ${OPENCV_LIBS} ${OPENCV_EXTRA} ${FFMPEG_LIBS}
        -Wl,--end-group pthread dl)

endif()

yolo_rtsp_pipeline.cpp

/* Copyright 2016-2022 by SOPHGO Technologies Inc. All rights reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

 Unless required by applicable law or agreed to in writing, software
 distributed under the License is distributed on an "AS IS" BASIS,
 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.

==============================================================================*/

/**
 * @brief Full pipeline demo:
 *   RTSP pull → HW decode → YOLOv5 inference → draw bboxes → HW encode → RTSP push
 */

#ifdef _WIN32
#include "getopt_win.h"
#else
#include <getopt.h>
#endif

#define USE_FFMPEG 1
#define USE_OPENCV 1
#define USE_BMCV 1

#include <algorithm>
#ifdef _OPENMP
#include <omp.h>
#endif
#include <cmath>
#include <cstdio>
#include <cstring>
#include <iostream>
#include <numeric>
#include <string>
#include <thread>
#include <vector>

#include <sail/cvwrapper.h>
#include <sail/encoder.h>
#include <sail/engine.h>
#include <sail/tensor.h>

// ---------------------------------------------------------------------------
// COCO 80-class names
// ---------------------------------------------------------------------------
static const std::vector<std::string> kCocoNames = {
    "person",        "bicycle",      "car",           "motorcycle",
    "airplane",      "bus",          "train",         "truck",
    "boat",          "trafficlight", "firehydrant",   "stopsign",
    "parkingmeter",  "bench",        "bird",          "cat",
    "dog",           "horse",        "sheep",         "cow",
    "elephant",      "bear",         "zebra",         "giraffe",
    "backpack",      "umbrella",     "handbag",       "tie",
    "suitcase",      "frisbee",      "skis",          "snowboard",
    "sportsball",    "kite",         "baseballbat",   "baseballglove",
    "skateboard",    "surfboard",    "tennisracket",  "bottle",
    "wineglass",     "cup",          "fork",          "knife",
    "spoon",         "bowl",         "banana",        "apple",
    "sandwich",      "orange",       "broccoli",      "carrot",
    "hotdog",        "pizza",        "donut",         "cake",
    "chair",         "couch",        "pottedplant",   "bed",
    "diningtable",   "toilet",       "tvmonitor",     "laptop",
    "mouse",         "remote",       "keyboard",      "cellphone",
    "microwave",     "oven",         "toaster",       "sink",
    "refrigerator",  "book",         "clock",         "vase",
    "scissors",      "teddybear",    "hairdrier",     "toothbrush"};

// ---------------------------------------------------------------------------
// Detection result structure
// ---------------------------------------------------------------------------
struct Detection {
  float left;
  float top;
  float right;
  float bottom;
  float score;
  int class_id;
};

// ---------------------------------------------------------------------------
// IoU / NMS
// ---------------------------------------------------------------------------
static float box_iou(const Detection &a, const Detection &b) {
  float inter_left = std::max(a.left, b.left);
  float inter_top = std::max(a.top, b.top);
  float inter_right = std::min(a.right, b.right);
  float inter_bottom = std::min(a.bottom, b.bottom);
  float inter_w = std::max(0.0f, inter_right - inter_left);
  float inter_h = std::max(0.0f, inter_bottom - inter_top);
  float inter_area = inter_w * inter_h;
  float area_a = (a.right - a.left) * (a.bottom - a.top);
  float area_b = (b.right - b.left) * (b.bottom - b.top);
  float union_area = area_a + area_b - inter_area;
  return (union_area > 0.0f) ? inter_area / union_area : 0.0f;
}

static void nms(std::vector<Detection> &dets, float nms_threshold) {
  std::sort(dets.begin(), dets.end(),
            [](const Detection &a, const Detection &b) {
              return a.score > b.score;
            });
  std::vector<Detection> keep;
  std::vector<bool> suppressed(dets.size(), false);
  for (size_t i = 0; i < dets.size(); ++i) {
    if (suppressed[i]) continue;
    keep.push_back(dets[i]);
    for (size_t j = i + 1; j < dets.size(); ++j) {
      if (suppressed[j]) continue;
      if (dets[i].class_id == dets[j].class_id &&
          box_iou(dets[i], dets[j]) > nms_threshold) {
        suppressed[j] = true;
      }
    }
  }
  dets = std::move(keep);
}

// ---------------------------------------------------------------------------
// YOLOv5 3-output post-processing
//
//   85 = [tx, ty, tw, th, objectness, class_0 ... class_79]  (raw pre-sigmoid)
// ---------------------------------------------------------------------------
static void yolov5_decode_layer(
    const float *output, int batch_idx, int num_anchors, int grid_h, int grid_w,
    int num_classes, int stride,
    const float anchors[][2],
    int net_w, int net_h, int orig_w, int orig_h,
    int pad_stx, int pad_sty, int pad_w, int pad_h,
    float conf_thres, std::vector<Detection> &detections) {

  const int row_len = 5 + num_classes;
  // scale maps model coords → original coords: orig = model / scale
  // Letterbox: scale = net_w / orig_w (same for x,y because aspect preserved)
  // Stretch:   scale_x = net_w / orig_w, scale_y = net_h / orig_h
  bool is_stretch = (pad_stx == 0 && pad_sty == 0 && pad_w == net_w && pad_h == net_h);
  float scale_x, scale_y;
  if (is_stretch) {
    scale_x = static_cast<float>(net_w) / static_cast<float>(orig_w);
    scale_y = static_cast<float>(net_h) / static_cast<float>(orig_h);
  } else {
    float s = std::min(static_cast<float>(pad_w) / static_cast<float>(orig_w),
                       static_cast<float>(pad_h) / static_cast<float>(orig_h));
    scale_x = scale_y = s;
  }

  for (int a = 0; a < num_anchors; ++a) {
    float anchor_w = anchors[a][0];
    float anchor_h = anchors[a][1];

    for (int gy = 0; gy < grid_h; ++gy) {
      for (int gx = 0; gx < grid_w; ++gx) {
        int offset = ((batch_idx * num_anchors + a) * grid_h + gy) * grid_w + gx;
        const float *row = output + offset * row_len;

        float obj_conf = 1.0f / (1.0f + std::exp(-row[4]));
        if (obj_conf < conf_thres * 0.1f) continue;

        float max_cls_score = 0.0f;
        int best_cls = -1;
        for (int c = 0; c < num_classes; ++c) {
          float score = 1.0f / (1.0f + std::exp(-row[5 + c]));
          if (score > max_cls_score) {
            max_cls_score = score;
            best_cls = c;
          }
        }

        float confidence = obj_conf * max_cls_score;
        if (confidence < conf_thres) continue;

        // YOLOv5 decode: cx = (2*sigmoid(tx)-0.5 + grid_x) * stride
        float tx = 1.0f / (1.0f + std::exp(-row[0]));
        float ty = 1.0f / (1.0f + std::exp(-row[1]));
        float tw = 1.0f / (1.0f + std::exp(-row[2]));
        float th = 1.0f / (1.0f + std::exp(-row[3]));

        float cx = (2.0f * tx - 0.5f + static_cast<float>(gx)) * stride;
        float cy = (2.0f * ty - 0.5f + static_cast<float>(gy)) * stride;
        float bw = (2.0f * tw) * (2.0f * tw) * anchor_w;
        float bh = (2.0f * th) * (2.0f * th) * anchor_h;

        float left  = (cx - bw * 0.5f - pad_stx) / scale_x;
        float top   = (cy - bh * 0.5f - pad_sty) / scale_y;
        float right  = (cx + bw * 0.5f - pad_stx) / scale_x;
        float bottom = (cy + bh * 0.5f - pad_sty) / scale_y;

        left   = std::max(0.0f, left);
        top    = std::max(0.0f, top);
        right  = std::min(static_cast<float>(orig_w), right);
        bottom = std::min(static_cast<float>(orig_h), bottom);

        if (right > left && bottom > top) {
          detections.push_back({left, top, right, bottom, confidence, best_cls});
        }
      }
    }
  }
}

static const float kYolov5Anchors[3][3][2] = {
  {{10.0f, 13.0f}, {16.0f, 30.0f}, {33.0f, 23.0f}},       // stride 8
  {{30.0f, 61.0f}, {62.0f, 45.0f}, {59.0f, 119.0f}},       // stride 16
  {{116.0f, 90.0f}, {156.0f, 198.0f}, {373.0f, 326.0f}}    // stride 32
};
static const int kYolov5Strides[3] = {8, 16, 32};

static void yolov5_3output_postprocess(
    const std::vector<float*> &outputs,
    const std::vector<std::vector<int>> &output_shapes,
    int batch_idx, int num_classes,
    int net_w, int net_h, int orig_w, int orig_h,
    int pad_stx, int pad_sty, int pad_w, int pad_h,
    float conf_thres, float nms_thres,
    std::vector<Detection> &detections) {

  std::vector<Detection> all_dets;
#pragma omp parallel
  {
    std::vector<Detection> local;
#pragma omp for nowait
    for (int layer = 0; layer < 3; ++layer) {
      int grid_h = output_shapes[layer][2];
      int grid_w = output_shapes[layer][3];
      yolov5_decode_layer(outputs[layer], batch_idx, 3, grid_h, grid_w,
                          num_classes, kYolov5Strides[layer],
                          kYolov5Anchors[layer],
                          net_w, net_h, orig_w, orig_h,
                          pad_stx, pad_sty, pad_w, pad_h,
                          conf_thres, local);
    }
#pragma omp critical
    all_dets.insert(all_dets.end(), local.begin(), local.end());
  }
  nms(all_dets, nms_thres);
  detections = std::move(all_dets);
}

// ---------------------------------------------------------------------------
// Simple argument parser
// ---------------------------------------------------------------------------
static std::string get_arg(int argc, char *argv[], const std::string &key,
                           const std::string &default_val) {
  std::string prefix = "--" + key + "=";
  for (int i = 1; i < argc; ++i) {
    if (strncmp(argv[i], prefix.c_str(), prefix.size()) == 0) {
      return std::string(argv[i] + prefix.size());
    }
  }
  return default_val;
}

// ---------------------------------------------------------------------------
// Main
// ---------------------------------------------------------------------------
int main(int argc, char *argv[]) {
  // -------- Defaults ----------------------------------------
  std::string input_url = "rtsp://192.168.2.119/554";
  std::string output_url = "rtsp://127.0.0.1:554/live/test";
  std::string bmodel_path = "yolov5s_v6.1_3output_int8_1b.bmodel";
  int device_id = 0;
  float conf_thres = 0.5f;
  float nms_thres = 0.45f;
  int enc_bitrate = 2000;
  int enc_framerate = 25;
  int enc_gop = 32;
  std::string enc_fmt = "h264_bm";

  // -------- Parse args --------------------------------------
  input_url = get_arg(argc, argv, "input", input_url);
  output_url = get_arg(argc, argv, "output", output_url);
  bmodel_path = get_arg(argc, argv, "bmodel", bmodel_path);
  device_id = std::stoi(get_arg(argc, argv, "dev_id", std::to_string(device_id)));
  conf_thres = std::stof(get_arg(argc, argv, "conf_thres", std::to_string(conf_thres)));
  nms_thres = std::stof(get_arg(argc, argv, "nms_thres", std::to_string(nms_thres)));
  enc_bitrate = std::stoi(get_arg(argc, argv, "enc_bitrate", std::to_string(enc_bitrate)));
  enc_framerate = std::stoi(get_arg(argc, argv, "enc_framerate", std::to_string(enc_framerate)));
  enc_gop = std::stoi(get_arg(argc, argv, "enc_gop", std::to_string(enc_gop)));
  enc_fmt = get_arg(argc, argv, "enc_fmt", enc_fmt);

  fprintf(stderr, "========================================\n");
  fprintf(stderr, "  YOLO RTSP Pipeline Demo\n");
  fprintf(stderr, "========================================\n");
  fprintf(stderr, "Input:        %s\n", input_url.c_str());
  fprintf(stderr, "Output:       %s\n", output_url.c_str());
  fprintf(stderr, "Model:        %s\n", bmodel_path.c_str());
  fprintf(stderr, "Device ID:    %d\n", device_id);
  fprintf(stderr, "Conf thres:   %.2f\n", conf_thres);
  fprintf(stderr, "NMS thres:    %.2f\n", nms_thres);
  fprintf(stderr, "Encode:       %s @ %dkbps, %dfps, GOP=%d\n",
          enc_fmt.c_str(), enc_bitrate, enc_framerate, enc_gop);

  // -------- 1. Init ------------------------------------------
  sail::Handle handle(device_id);
  sail::Bmcv bmcv(handle);

  // -------- 2. Decoder ---------------------------------------
  fprintf(stderr, "Opening decoder: %s\n", input_url.c_str());
  sail::Decoder decoder(input_url, true, device_id);
  if (!decoder.is_opened()) {
    fprintf(stderr, "[ERROR] Failed to open decoder!\n");
    return 1;
  }
  fprintf(stderr, "Decoder OK.\n");

  // -------- 3. Engine (manual preprocess + inference) --------
  fprintf(stderr, "Loading bmodel: %s\n", bmodel_path.c_str());
  sail::Engine engine(device_id);
  if (!engine.load(bmodel_path)) {
    fprintf(stderr, "[ERROR] Failed to load bmodel!\n");
    return 1;
  }

  auto graph_name = engine.get_graph_names().front();
  auto input_name = engine.get_input_names(graph_name).front();
  auto output_names = engine.get_output_names(graph_name);
  auto input_shape = engine.get_input_shape(graph_name, input_name);
  auto input_dtype = engine.get_input_dtype(graph_name, input_name);

  int net_w = input_shape[3];
  int net_h = input_shape[2];
  int num_outputs = output_names.size();
  if (num_outputs != 3) {
    fprintf(stderr, "[ERROR] Expected 3 outputs, got %d\n", num_outputs);
    return 1;
  }

  auto output_dtype = engine.get_output_dtype(graph_name, output_names[0]);
  std::vector<std::vector<int>> output_shapes(num_outputs);
  for (int i = 0; i < num_outputs; ++i) {
    output_shapes[i] = engine.get_output_shape(graph_name, output_names[i]);
    fprintf(stderr, "Output[%d] %s: [%d,%d,%d,%d,%d]\n",
            i, output_names[i].c_str(),
            output_shapes[i][0], output_shapes[i][1],
            output_shapes[i][2], output_shapes[i][3], output_shapes[i][4]);
  }
  int num_classes = output_shapes[0][4] - 5;
  fprintf(stderr, "Model input:  %dx%d (dtype=%d)\n", net_w, net_h,
          static_cast<int>(input_dtype));
  fprintf(stderr, "Model output: %d layers, %d classes\n", num_outputs, num_classes);

  auto in = std::make_shared<sail::Tensor>(handle, input_shape, input_dtype, false, false);
  std::vector<std::shared_ptr<sail::Tensor>> out_tensors;
  std::map<std::string, sail::Tensor *> output_tensors;
  for (int i = 0; i < num_outputs; ++i) {
    out_tensors.push_back(
        std::make_shared<sail::Tensor>(handle, output_shapes[i], output_dtype, true, true));
    output_tensors[output_names[i]] = out_tensors.back().get();
  }
  std::map<std::string, sail::Tensor *> input_tensors = {{input_name, in.get()}};
  engine.set_io_mode(graph_name, sail::SYSO);

  float input_scale = engine.get_input_scale(graph_name, input_name);
  float output_scale = engine.get_output_scale(graph_name, output_names[0]);
  fprintf(stderr, "Input scale:  %.4f\n", input_scale);
  fprintf(stderr, "Output scale: %.4f\n", output_scale);

  // -------- 4. Encoder -----------------------------------------
  auto frame_shape = decoder.get_frame_shape();
  int src_h = frame_shape[2];
  int src_w = frame_shape[3];
  fprintf(stderr, "Source frame: %dx%d\n", src_w, src_h);

  char enc_params_buf[256];
  snprintf(enc_params_buf, sizeof(enc_params_buf),
           "width=%d:height=%d:bitrate=%d:gop=%d:gop_preset=2:framerate=%d",
           src_w, src_h, enc_bitrate, enc_gop, enc_framerate);
  std::string enc_params = enc_params_buf;

  fprintf(stderr, "Opening encoder: %s\n", output_url.c_str());
  fprintf(stderr, "Encoder params: %s\n", enc_params.c_str());
  sail::Encoder encoder(output_url, handle, enc_fmt, "NV12", enc_params, 1, 1);

  // -------- 5. Preprocess image --------------------------------
  auto img_dtype = bmcv.get_bm_image_data_format(input_dtype);
  sail::BMImage preprocessed(handle, net_h, net_w, FORMAT_RGB_PLANAR, img_dtype);
  preprocessed.align();

  // -------- 6. Main loop ---------------------------------------
  fprintf(stderr, "========================================\n");
  fprintf(stderr, "  Pipeline running...\n");
  fprintf(stderr, "========================================\n");

  int frame_count = 0;

  while (true) {
    // 6a. HW decode
    sail::BMImage img;
    int ret = decoder.read(handle, img);
    if (ret != 0) {
      fprintf(stderr, "[WARN] Decode failed (frame %d), reconnect...\n", frame_count);
      decoder.reconnect();
      std::this_thread::sleep_for(std::chrono::milliseconds(100));
      continue;
    }
    int orig_w = img.width();
    int orig_h = img.height();

    // 6b. Preprocess: NV12→RGB→resize→normalize→tensor
    sail::BMImage rgb_img(handle, img.height(), img.width(),
                          FORMAT_RGB_PLANAR, DATA_TYPE_EXT_1N_BYTE);
    rgb_img.align();
    bmcv.convert_format(img, rgb_img);
    sail::BMImage resized = bmcv.vpp_resize(rgb_img, net_w, net_h);

    float alpha = input_scale / 255.0f;
    bmcv.convert_to(resized, preprocessed,
                    std::make_tuple(std::make_pair(alpha, 0.0f),
                                    std::make_pair(alpha, 0.0f),
                                    std::make_pair(alpha, 0.0f)));
    bmcv.bm_image_to_tensor(preprocessed, *in);

    // 6c. TPU inference
    engine.process(graph_name, input_tensors, output_tensors);

    // 6d. Extract output
    std::vector<float*> output_ptrs(num_outputs);
    std::vector<std::vector<float>> output_bufs(num_outputs);
    for (int i = 0; i < num_outputs; ++i) {
      int out_size = std::accumulate(output_shapes[i].begin(),
                                     output_shapes[i].end(), 1,
                                     std::multiplies<int>());
      if (output_dtype == BM_FLOAT32) {
        output_ptrs[i] = reinterpret_cast<float *>(out_tensors[i]->sys_data());
      } else {
        output_bufs[i].resize(out_size);
        output_ptrs[i] = output_bufs[i].data();
        if (output_dtype == BM_INT8) {
          engine.scale_int8_to_fp32(
              reinterpret_cast<int8_t *>(out_tensors[i]->sys_data()),
              output_ptrs[i], output_scale, out_size);
        } else if (output_dtype == BM_UINT8) {
          engine.scale_uint8_to_fp32(
              reinterpret_cast<uint8_t *>(out_tensors[i]->sys_data()),
              output_ptrs[i], output_scale, out_size);
        }
      }
    }

    // 6e. YOLOv5 postprocess (stretch resize, no letterbox)
    std::vector<Detection> detections;
    yolov5_3output_postprocess(output_ptrs, output_shapes, 0, num_classes,
                               net_w, net_h, orig_w, orig_h,
                               0, 0, net_w, net_h,
                               conf_thres, nms_thres, detections);

    // 6f. Draw boxes
    for (const auto &det : detections) {
      int x0 = static_cast<int>(det.left);
      int y0 = static_cast<int>(det.top);
      int w = static_cast<int>(det.right - det.left);
      int h = static_cast<int>(det.bottom - det.top);
      bmcv.rectangle(img, x0, y0, w, h, std::make_tuple(0, 255, 0), 2);
    }

    // 6g. HW encode + push
    ret = encoder.video_write(img);
    if (ret != 0)
      fprintf(stderr, "[WARN] Encode push failed (frame %d)\n", frame_count);

    // 6h. Progress
    frame_count++;
    if (frame_count % 100 == 0)
      fprintf(stderr, "Frame %d | %zu detections\n", frame_count, detections.size());
  }

  decoder.release();
  encoder.release();
  fprintf(stderr, "Pipeline exited. Total frames: %d\n", frame_count);
  return 0;
}

更多推荐