generated from guanyuankai/bonus-edge-proxy
feat(video Service): 完善车辆识别算法业务。
This commit is contained in:
parent
5a241a3e9e
commit
430a69ad51
Binary file not shown.
|
|
@ -5,6 +5,7 @@ RUN apt-get update && \
|
|||
# 基础工具
|
||||
sudo \
|
||||
build-essential \
|
||||
libmysqlclient-dev \
|
||||
cmake \
|
||||
git \
|
||||
gdb \
|
||||
|
|
|
|||
|
|
@ -75,7 +75,7 @@ void poll_system_metrics(boost::asio::steady_timer& timer, SystemMonitor::System
|
|||
int main(int argc, char* argv[]) {
|
||||
// TODO: [GYK] DEV#1: 将 URL 放入 config.json 中读取
|
||||
// std::string cam_rtsp_input = "rtsp://admin:123456@192.168.1.57:554/stream0";
|
||||
std::string cam_rtsp_input = "../data/test_car.mp4";
|
||||
std::string cam_rtsp_input = "../data/test_car2.mp4";
|
||||
// std::string cam_rtsp_input =
|
||||
// "rtsp://admin:hzx12345@192.168.1.10:554/Streaming/Channels/1901";
|
||||
|
||||
|
|
|
|||
|
|
@ -2,7 +2,7 @@
|
|||
|
||||
#include <algorithm> // for std::max, std::min
|
||||
#include <chrono>
|
||||
|
||||
const int YoloDetector::NPU_CORE_CNT;
|
||||
VideoPipeline::VideoPipeline() : running_(false), next_track_id_(0) {
|
||||
detector_ = std::make_unique<YoloDetector>();
|
||||
// 模型路径
|
||||
|
|
@ -39,6 +39,37 @@ void VideoPipeline::Stop() {
|
|||
processingThread_.join();
|
||||
spdlog::info("VideoPipeline Stopped.");
|
||||
}
|
||||
void VideoPipeline::inferenceWorker() {
|
||||
while (running_) {
|
||||
FrameData data;
|
||||
{
|
||||
std::unique_lock<std::mutex> lock(input_mtx_);
|
||||
// 等待有数据或者停止信号
|
||||
input_cv_.wait(lock, [this] { return !input_queue_.empty() || !running_; });
|
||||
|
||||
if (!running_ && input_queue_.empty())
|
||||
return;
|
||||
|
||||
data = input_queue_.front();
|
||||
input_queue_.pop();
|
||||
// 通知主线程队列有空位了（流控）
|
||||
input_cv_.notify_all();
|
||||
}
|
||||
|
||||
// === 核心：这里是并行执行的 ===
|
||||
// YoloDetector::detect 现在是线程安全的，会自动分配 NPU 核心
|
||||
data.results = detector_->detect(data.original_frame);
|
||||
|
||||
// 将结果放入输出缓冲区
|
||||
{
|
||||
std::lock_guard<std::mutex> lock(output_mtx_);
|
||||
// 使用 move 减少拷贝
|
||||
output_buffer_[data.frame_id] = std::move(data);
|
||||
}
|
||||
// 通知主线程有结果了
|
||||
output_cv_.notify_one();
|
||||
}
|
||||
}
|
||||
|
||||
// [新增] 计算两个矩形的交并比 (IoU)
|
||||
float VideoPipeline::computeIOU(const cv::Rect& box1, const cv::Rect& box2) {
|
||||
|
|
@ -155,19 +186,23 @@ void VideoPipeline::drawOverlay(cv::Mat& frame, const std::vector<TrackedVehicle
|
|||
void VideoPipeline::processLoop(std::string inputUrl, std::string outputUrl, bool isFileSource) {
|
||||
cv::VideoCapture cap;
|
||||
cap.open(inputUrl);
|
||||
|
||||
if (!cap.isOpened()) {
|
||||
spdlog::error("Failed to open input: {}", inputUrl);
|
||||
running_ = false;
|
||||
return;
|
||||
}
|
||||
|
||||
const double TARGET_FPS = 60.0;
|
||||
const double FRAME_DURATION_MS = 1000.0 / TARGET_FPS;
|
||||
// 1. 启动 3 个工作线程 (对应 3 个 NPU 核心)
|
||||
// 你的 NPU 利用率将在这里被填满
|
||||
for (int i = 0; i < 3; ++i) {
|
||||
worker_threads_.emplace_back(&VideoPipeline::inferenceWorker, this);
|
||||
}
|
||||
|
||||
int width = cap.get(cv::CAP_PROP_FRAME_WIDTH);
|
||||
int height = cap.get(cv::CAP_PROP_FRAME_HEIGHT);
|
||||
const double TARGET_FPS = 30.0; // 提升目标 FPS
|
||||
|
||||
// GStreamer pipeline string (保持原样...)
|
||||
// ... GStreamer pipeline string 设置保持不变 ...
|
||||
std::stringstream pipeline;
|
||||
pipeline << "appsrc ! "
|
||||
<< "videoconvert ! "
|
||||
|
|
@ -180,75 +215,92 @@ void VideoPipeline::processLoop(std::string inputUrl, std::string outputUrl, boo
|
|||
cv::VideoWriter writer;
|
||||
writer.open(pipeline.str(), cv::CAP_GSTREAMER, 0, TARGET_FPS, cv::Size(width, height), true);
|
||||
|
||||
cv::Mat frame;
|
||||
long frame_count = 0;
|
||||
using Clock = std::chrono::high_resolution_clock;
|
||||
using Ms = std::chrono::duration<double, std::milli>;
|
||||
long read_frame_idx = 0; // 读取计数
|
||||
long write_frame_idx = 0; // 写入计数
|
||||
|
||||
while (running_) {
|
||||
auto loop_start = std::chrono::steady_clock::now();
|
||||
auto loop_begin_tp = Clock::now();
|
||||
auto t1 = Clock::now();
|
||||
if (!cap.read(frame)) {
|
||||
if (isFileSource) {
|
||||
cap.set(cv::CAP_PROP_POS_FRAMES, 0);
|
||||
continue;
|
||||
} else {
|
||||
std::this_thread::sleep_for(std::chrono::seconds(1));
|
||||
cap.release();
|
||||
cap.open(inputUrl);
|
||||
continue;
|
||||
// === 阶段 A: 读取并分发任务 (生产者) ===
|
||||
// 限制预读数量，防止内存爆满 (例如最多预读 5 帧)
|
||||
{
|
||||
std::unique_lock<std::mutex> lock(input_mtx_);
|
||||
// 如果输入队列满了，等待工作线程处理
|
||||
if (input_queue_.size() < MAX_INPUT_QUEUE_SIZE) {
|
||||
cv::Mat frame;
|
||||
if (cap.read(frame) && !frame.empty()) {
|
||||
FrameData data;
|
||||
data.frame_id = read_frame_idx++;
|
||||
data.original_frame = frame; // 拷贝一份 (必须，因为 cv::Mat 是引用计数)
|
||||
|
||||
input_queue_.push(data);
|
||||
input_cv_.notify_one(); // 唤醒一个工作线程
|
||||
} else {
|
||||
if (isFileSource) {
|
||||
// 文件读完了的处理...
|
||||
running_ = false;
|
||||
} else {
|
||||
// 网络流断线的处理...
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
if (frame.empty())
|
||||
continue;
|
||||
auto t2 = Clock::now();
|
||||
// 1. 推理
|
||||
std::vector<DetectionResult> results = detector_->detect(frame);
|
||||
auto t3 = Clock::now();
|
||||
// 2. [修改] 更新追踪器和平滑分数
|
||||
updateTracker(results);
|
||||
auto t4 = Clock::now();
|
||||
// 3. [修改] 准备绘图数据
|
||||
std::vector<TrackedVehicle> tracks_to_draw;
|
||||
for (const auto& pair : tracks_) {
|
||||
tracks_to_draw.push_back(pair.second);
|
||||
|
||||
// === 阶段 B: 按顺序收集结果并处理 (消费者) ===
|
||||
FrameData current_data;
|
||||
bool has_data = false;
|
||||
|
||||
{
|
||||
std::unique_lock<std::mutex> lock(output_mtx_);
|
||||
// 检查输出缓冲区里是否有我们期待的下一帧 (write_frame_idx)
|
||||
// 因为多线程处理，第5帧可能比第4帧先处理完，必须等待第4帧
|
||||
auto it = output_buffer_.find(write_frame_idx);
|
||||
if (it != output_buffer_.end()) {
|
||||
current_data = std::move(it->second);
|
||||
output_buffer_.erase(it);
|
||||
has_data = true;
|
||||
write_frame_idx++;
|
||||
}
|
||||
}
|
||||
|
||||
// 4. [修改] 绘制
|
||||
drawOverlay(frame, tracks_to_draw);
|
||||
auto t5 = Clock::now();
|
||||
// 5. 推流
|
||||
if (writer.isOpened())
|
||||
writer.write(frame);
|
||||
auto t6 = Clock::now();
|
||||
if (has_data) {
|
||||
// 注意：updateTracker 和 drawOverlay 必须在主线程串行执行
|
||||
// 因为它们依赖 tracks_ 状态，且必须按时间顺序更新
|
||||
|
||||
// ----------------- 计算耗时 -----------------
|
||||
double ms_read = std::chrono::duration_cast<Ms>(t2 - t1).count();
|
||||
double ms_infer = std::chrono::duration_cast<Ms>(t3 - t2).count();
|
||||
double ms_track = std::chrono::duration_cast<Ms>(t4 - t3).count();
|
||||
double ms_draw = std::chrono::duration_cast<Ms>(t5 - t4).count();
|
||||
double ms_write = std::chrono::duration_cast<Ms>(t6 - t5).count();
|
||||
double ms_total = std::chrono::duration_cast<Ms>(t6 - t1).count();
|
||||
frame_count++;
|
||||
if (frame_count % 60 == 0 || ms_total > 18.0) {
|
||||
spdlog::info(
|
||||
"Frame[{}] Cost: Total={:.1f}ms | Read={:.1f} Infer={:.1f} Track={:.1f} "
|
||||
"Draw={:.1f} Write={:.1f}",
|
||||
frame_count, ms_total, ms_read, ms_infer, ms_track, ms_draw, ms_write);
|
||||
}
|
||||
// 1. 追踪 (CPU)
|
||||
updateTracker(current_data.results);
|
||||
|
||||
// FPS控制 (保持原样)
|
||||
if (isFileSource) {
|
||||
auto loop_end = std::chrono::steady_clock::now();
|
||||
double elapsed_ms =
|
||||
std::chrono::duration<double, std::milli>(loop_end - loop_start).count();
|
||||
double wait_ms = FRAME_DURATION_MS - elapsed_ms;
|
||||
if (wait_ms > 0)
|
||||
std::this_thread::sleep_for(std::chrono::milliseconds((int)wait_ms));
|
||||
// 2. 准备绘图数据 (CPU)
|
||||
std::vector<TrackedVehicle> tracks_to_draw;
|
||||
for (const auto& pair : tracks_) {
|
||||
tracks_to_draw.push_back(pair.second);
|
||||
}
|
||||
|
||||
// 3. 绘图 (CPU)
|
||||
drawOverlay(current_data.original_frame, tracks_to_draw);
|
||||
|
||||
// 4. 推流 (IO)
|
||||
if (writer.isOpened()) {
|
||||
writer.write(current_data.original_frame);
|
||||
}
|
||||
|
||||
// 简单的 FPS 打印
|
||||
if (write_frame_idx % 60 == 0) {
|
||||
spdlog::info("Processed Frame ID: {}", write_frame_idx);
|
||||
}
|
||||
} else {
|
||||
// 如果没有等到当前帧，稍微休眠一下避免死循环占满 CPU
|
||||
// 但不要睡太久，否则延迟高
|
||||
std::this_thread::sleep_for(std::chrono::milliseconds(1));
|
||||
}
|
||||
}
|
||||
|
||||
// 清理：通知线程退出并 Join
|
||||
input_cv_.notify_all();
|
||||
for (auto& t : worker_threads_) {
|
||||
if (t.joinable())
|
||||
t.join();
|
||||
}
|
||||
worker_threads_.clear();
|
||||
|
||||
cap.release();
|
||||
writer.release();
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,9 +1,12 @@
|
|||
#pragma once
|
||||
|
||||
#include <atomic>
|
||||
#include <map> // [新增]
|
||||
#include <condition_variable>
|
||||
#include <map>
|
||||
#include <memory>
|
||||
#include <mutex>
|
||||
#include <opencv2/opencv.hpp>
|
||||
#include <queue>
|
||||
#include <string>
|
||||
#include <thread>
|
||||
#include <vector>
|
||||
|
|
@ -11,14 +14,21 @@
|
|||
#include "spdlog/spdlog.h"
|
||||
#include "yoloDetector/yolo_detector.hpp"
|
||||
|
||||
// [新增] 定义追踪车辆的结构体
|
||||
// ... TrackedVehicle 结构体保持不变 ...
|
||||
struct TrackedVehicle {
|
||||
int id; // 唯一ID
|
||||
cv::Rect box; // 当前位置
|
||||
float ev_score; // 新能源分数 (0.0 - 1.0), 越接近1.0越可能是新能源
|
||||
int missing_frames; // 连续丢失帧数 (用于删除消失的车辆)
|
||||
int last_class_id; // 上一次显示的类别ID (避免标签闪烁)
|
||||
std::string label; // 当前显示的标签
|
||||
int id;
|
||||
cv::Rect box;
|
||||
float ev_score;
|
||||
int missing_frames;
|
||||
int last_class_id;
|
||||
std::string label;
|
||||
};
|
||||
|
||||
// [新增] 每一帧的数据包
|
||||
struct FrameData {
|
||||
long frame_id; // 帧序号，用于排序
|
||||
cv::Mat original_frame; // 原始图像
|
||||
std::vector<DetectionResult> results; // 检测结果
|
||||
};
|
||||
|
||||
class VideoPipeline {
|
||||
|
|
@ -33,21 +43,34 @@ public:
|
|||
private:
|
||||
void processLoop(std::string inputUrl, std::string outputUrl, bool isFileSource);
|
||||
|
||||
// [修改] 绘图函数现在接收追踪列表，而不是原始检测结果
|
||||
// [新增] 工作线程函数
|
||||
void inferenceWorker();
|
||||
|
||||
void drawOverlay(cv::Mat& frame, const std::vector<TrackedVehicle>& trackedObjects);
|
||||
|
||||
// [新增] 核心：更新追踪和分数逻辑
|
||||
void updateTracker(const std::vector<DetectionResult>& detections);
|
||||
|
||||
// [新增] 辅助函数：计算 IoU (交并比)
|
||||
float computeIOU(const cv::Rect& box1, const cv::Rect& box2);
|
||||
|
||||
private:
|
||||
std::atomic<bool> running_;
|
||||
std::thread processingThread_;
|
||||
std::thread processingThread_; // 主处理线程
|
||||
std::unique_ptr<YoloDetector> detector_;
|
||||
|
||||
// [新增] 追踪列表，Key是ID
|
||||
// 追踪相关
|
||||
std::map<int, TrackedVehicle> tracks_;
|
||||
int next_track_id_ = 0;
|
||||
|
||||
// === [新增] 多线程并行处理相关 ===
|
||||
std::vector<std::thread> worker_threads_; // 3个工作线程
|
||||
|
||||
// 输入队列 (待检测)
|
||||
std::queue<FrameData> input_queue_;
|
||||
std::mutex input_mtx_;
|
||||
std::condition_variable input_cv_;
|
||||
const size_t MAX_INPUT_QUEUE_SIZE = 5; // 限制缓冲大小，防止内存爆炸
|
||||
|
||||
// 输出缓冲区 (已检测，等待排序)
|
||||
// 使用 map 自动按 frame_id 排序，解决乱序问题
|
||||
std::map<long, FrameData> output_buffer_;
|
||||
std::mutex output_mtx_;
|
||||
std::condition_variable output_cv_;
|
||||
};
|
||||
|
|
|
|||
|
|
@ -1,21 +1,26 @@
|
|||
#include "yolo_detector.hpp"
|
||||
|
||||
#include <algorithm>
|
||||
#include <cstring> // for memset
|
||||
#include <cstring>
|
||||
#include <fstream>
|
||||
#include <iostream>
|
||||
|
||||
#include "spdlog/spdlog.h"
|
||||
|
||||
YoloDetector::YoloDetector()
|
||||
: ctx_(0), is_initialized_(false), model_data_(nullptr), output_attrs_(nullptr) {
|
||||
: is_initialized_(false), model_data_(nullptr), output_attrs_(nullptr) {
|
||||
memset(&io_num_, 0, sizeof(io_num_));
|
||||
}
|
||||
|
||||
YoloDetector::~YoloDetector() {
|
||||
if (ctx_ > 0) {
|
||||
rknn_destroy(ctx_);
|
||||
// 销毁所有 context
|
||||
for (rknn_context ctx : ctx_pool_) {
|
||||
if (ctx > 0) {
|
||||
rknn_destroy(ctx);
|
||||
}
|
||||
}
|
||||
ctx_pool_.clear();
|
||||
|
||||
if (model_data_) {
|
||||
free(model_data_);
|
||||
}
|
||||
|
|
@ -48,72 +53,122 @@ int YoloDetector::init(const std::string& modelPath) {
|
|||
if (!model_data_)
|
||||
return -1;
|
||||
|
||||
// 2. 初始化 RKNN
|
||||
int ret = rknn_init(&ctx_, model_data_, model_data_size_, 0, NULL);
|
||||
if (ret < 0) {
|
||||
spdlog::error("rknn_init failed! ret={}", ret);
|
||||
return -1;
|
||||
// 2. 初始化 NPU Pool
|
||||
ctx_pool_.resize(NPU_CORE_CNT);
|
||||
|
||||
for (int i = 0; i < NPU_CORE_CNT; ++i) {
|
||||
int ret = rknn_init(&ctx_pool_[i], model_data_, model_data_size_, 0, NULL);
|
||||
if (ret < 0) {
|
||||
spdlog::error("rknn_init failed for core {}! ret={}", i, ret);
|
||||
return -1;
|
||||
}
|
||||
|
||||
// ================= [新增] 强制核心绑定 =================
|
||||
// RK3588 有 3 个核心:
|
||||
// i=0 -> Core 0
|
||||
// i=1 -> Core 1
|
||||
// i=2 -> Core 2
|
||||
rknn_core_mask core_mask;
|
||||
switch (i) {
|
||||
case 0:
|
||||
core_mask = RKNN_NPU_CORE_0;
|
||||
break;
|
||||
case 1:
|
||||
core_mask = RKNN_NPU_CORE_1;
|
||||
break;
|
||||
case 2:
|
||||
core_mask = RKNN_NPU_CORE_2;
|
||||
break;
|
||||
default:
|
||||
core_mask = RKNN_NPU_CORE_AUTO;
|
||||
break;
|
||||
}
|
||||
|
||||
ret = rknn_set_core_mask(ctx_pool_[i], core_mask);
|
||||
if (ret < 0) {
|
||||
spdlog::warn("Failed to set core mask for context {}: {}", i, ret);
|
||||
} else {
|
||||
spdlog::info("Context {} bound to NPU Core {}", i, i);
|
||||
}
|
||||
// =======================================================
|
||||
|
||||
free_ctx_indices_.push(i);
|
||||
}
|
||||
|
||||
// 3. 查询模型输入输出信息
|
||||
ret = rknn_query(ctx_, RKNN_QUERY_IN_OUT_NUM, &io_num_, sizeof(io_num_));
|
||||
// 3. 查询模型信息 (只需要查询其中一个 context 即可，所有 context 模型结构相同)
|
||||
rknn_context main_ctx = ctx_pool_[0];
|
||||
int ret = rknn_query(main_ctx, RKNN_QUERY_IN_OUT_NUM, &io_num_, sizeof(io_num_));
|
||||
if (ret < 0)
|
||||
return -1;
|
||||
|
||||
// 获取输入属性
|
||||
rknn_tensor_attr input_attrs[io_num_.n_input];
|
||||
memset(input_attrs, 0, sizeof(input_attrs));
|
||||
for (int i = 0; i < io_num_.n_input; i++) {
|
||||
input_attrs[i].index = i;
|
||||
rknn_query(ctx_, RKNN_QUERY_INPUT_ATTR, &(input_attrs[i]), sizeof(rknn_tensor_attr));
|
||||
rknn_query(main_ctx, RKNN_QUERY_INPUT_ATTR, &(input_attrs[i]), sizeof(rknn_tensor_attr));
|
||||
}
|
||||
|
||||
// [新增] 获取输出属性 (关键：用于后处理中的量化参数 scale/zp)
|
||||
output_attrs_ = (rknn_tensor_attr*)malloc(sizeof(rknn_tensor_attr) * io_num_.n_output);
|
||||
memset(output_attrs_, 0, sizeof(rknn_tensor_attr) * io_num_.n_output);
|
||||
for (int i = 0; i < io_num_.n_output; i++) {
|
||||
output_attrs_[i].index = i;
|
||||
rknn_query(ctx_, RKNN_QUERY_OUTPUT_ATTR, &(output_attrs_[i]), sizeof(rknn_tensor_attr));
|
||||
rknn_query(main_ctx, RKNN_QUERY_OUTPUT_ATTR, &(output_attrs_[i]), sizeof(rknn_tensor_attr));
|
||||
}
|
||||
|
||||
// 假设只有一个输入 (图像)
|
||||
model_width_ = input_attrs[0].dims[2]; // width
|
||||
model_height_ = input_attrs[0].dims[1]; // height
|
||||
model_channel_ = input_attrs[0].dims[0]; // channel
|
||||
// 假设只有一个输入
|
||||
model_width_ = input_attrs[0].dims[2];
|
||||
model_height_ = input_attrs[0].dims[1];
|
||||
model_channel_ = input_attrs[0].dims[0];
|
||||
|
||||
// 兼容部分 NCHW 格式
|
||||
if (model_width_ == 3) {
|
||||
model_width_ = input_attrs[0].dims[1];
|
||||
model_height_ = input_attrs[0].dims[2];
|
||||
}
|
||||
|
||||
spdlog::info("Model Input Shape: {}x{}x{}", model_width_, model_height_, model_channel_);
|
||||
spdlog::info("Model initialized on {} NPU cores. Input Shape: {}x{}x{}", NPU_CORE_CNT,
|
||||
model_width_, model_height_, model_channel_);
|
||||
is_initialized_ = true;
|
||||
return 0;
|
||||
}
|
||||
|
||||
// 资源获取逻辑
|
||||
int YoloDetector::get_context_id() {
|
||||
std::unique_lock<std::mutex> lock(mtx_);
|
||||
// 如果队列为空，等待直到有空闲 context 归还
|
||||
cv_.wait(lock, [this] { return !free_ctx_indices_.empty(); });
|
||||
|
||||
int id = free_ctx_indices_.front();
|
||||
free_ctx_indices_.pop();
|
||||
return id;
|
||||
}
|
||||
|
||||
// 资源释放逻辑
|
||||
void YoloDetector::return_context_id(int id) {
|
||||
{
|
||||
std::lock_guard<std::mutex> lock(mtx_);
|
||||
free_ctx_indices_.push(id);
|
||||
}
|
||||
// 通知等待的线程
|
||||
cv_.notify_one();
|
||||
}
|
||||
|
||||
cv::Mat YoloDetector::letterbox(const cv::Mat& src, int target_width, int target_height,
|
||||
float& ratio, int& dw, int& dh) {
|
||||
// 保持原有逻辑不变
|
||||
int w = src.cols;
|
||||
int h = src.rows;
|
||||
|
||||
float scale_w = (float)target_width / w;
|
||||
float scale_h = (float)target_height / h;
|
||||
ratio = std::min(scale_w, scale_h);
|
||||
|
||||
int new_unpad_w = (int)(w * ratio);
|
||||
int new_unpad_h = (int)(h * ratio);
|
||||
|
||||
dw = (target_width - new_unpad_w) / 2;
|
||||
dh = (target_height - new_unpad_h) / 2;
|
||||
|
||||
cv::Mat resized;
|
||||
cv::resize(src, resized, cv::Size(new_unpad_w, new_unpad_h));
|
||||
|
||||
cv::Mat padded;
|
||||
cv::copyMakeBorder(resized, padded, dh, target_height - new_unpad_h - dh, dw,
|
||||
target_width - new_unpad_w - dw, cv::BORDER_CONSTANT, cv::Scalar(0, 0, 0));
|
||||
|
||||
return padded;
|
||||
}
|
||||
|
||||
|
|
@ -122,12 +177,20 @@ std::vector<DetectionResult> YoloDetector::detect(const cv::Mat& inputImage) {
|
|||
if (!is_initialized_ || inputImage.empty())
|
||||
return results;
|
||||
|
||||
// === 0. 获取 NPU 资源 ===
|
||||
// 这一步会阻塞，直到有空闲的 NPU
|
||||
int ctx_id = get_context_id();
|
||||
rknn_context current_ctx = ctx_pool_[ctx_id];
|
||||
|
||||
// 使用 try-catch 或者是 RAII 确保资源一定会归还，这里手动保证
|
||||
// 建议在生产环境中封装一个 ContextGuard 类
|
||||
|
||||
// === 1. 预处理 ===
|
||||
// 注意：OpenCV 的操作是 CPU 上的，多线程调用时互不干扰（只要 inputImage 不同）
|
||||
float ratio;
|
||||
int dw, dh;
|
||||
cv::Mat input_rgb;
|
||||
cv::cvtColor(inputImage, input_rgb, cv::COLOR_BGR2RGB);
|
||||
|
||||
cv::Mat img = letterbox(input_rgb, model_width_, model_height_, ratio, dw, dh);
|
||||
|
||||
// === 2. 设置输入 ===
|
||||
|
|
@ -139,46 +202,40 @@ std::vector<DetectionResult> YoloDetector::detect(const cv::Mat& inputImage) {
|
|||
inputs[0].fmt = RKNN_TENSOR_NHWC;
|
||||
inputs[0].buf = img.data;
|
||||
|
||||
rknn_inputs_set(ctx_, io_num_.n_input, inputs);
|
||||
rknn_inputs_set(current_ctx, io_num_.n_input, inputs);
|
||||
|
||||
// === 3. 执行推理 ===
|
||||
int ret = rknn_run(ctx_, nullptr);
|
||||
int ret = rknn_run(current_ctx, nullptr);
|
||||
if (ret < 0) {
|
||||
spdlog::error("rknn_run failed: {}", ret);
|
||||
spdlog::error("rknn_run failed on core {}: {}", ctx_id, ret);
|
||||
return_context_id(ctx_id); // 错误发生也要归还资源
|
||||
return results;
|
||||
}
|
||||
|
||||
// === 4. 获取输出 ===
|
||||
// 使用 vector 或 new 替代变长数组(VLA)，更符合 C++ 标准
|
||||
// 每个线程都有自己的 outputs vector，不会冲突
|
||||
std::vector<rknn_output> outputs(io_num_.n_output);
|
||||
memset(outputs.data(), 0, outputs.size() * sizeof(rknn_output));
|
||||
|
||||
// YoloV8 输出通常需要 float 处理 (want_float=1)，
|
||||
// 但 postprocess.cc 中的 process_i8 是针对 int8 优化的。
|
||||
// 如果你的模型是 int8 量化且 output_attrs 正常，want_float=0 也可以，
|
||||
// 但 post_process 函数里强转为了 float* 处理 fp32 分支，或者 int8* 处理 i8 分支。
|
||||
// 为了保险起见，如果 postprocess 逻辑是处理原始输出的，这里通常不需要 want_float=1，除非是 fp32
|
||||
// 模型。
|
||||
// **注意**：你的 postprocess.cc 能够处理量化(process_i8)和非量化(process_fp32)数据。
|
||||
// 通常建议让 rknn 直接输出原始数据，由后处理决定如何解析。
|
||||
for (int i = 0; i < io_num_.n_output; i++) {
|
||||
outputs[i].want_float = 0; // 0: 输出原始量化数据(int8), 1: 驱动转为float
|
||||
outputs[i].want_float = 0;
|
||||
outputs[i].index = i;
|
||||
}
|
||||
|
||||
ret = rknn_outputs_get(ctx_, io_num_.n_output, outputs.data(), NULL);
|
||||
ret = rknn_outputs_get(current_ctx, io_num_.n_output, outputs.data(), NULL);
|
||||
if (ret < 0) {
|
||||
spdlog::error("rknn_outputs_get failed: {}", ret);
|
||||
return_context_id(ctx_id);
|
||||
return results;
|
||||
}
|
||||
|
||||
// === 5. 后处理 ===
|
||||
// 构造当前线程专用的 app_ctx
|
||||
rknn_app_context_t app_ctx;
|
||||
app_ctx.io_num = io_num_;
|
||||
app_ctx.output_attrs = output_attrs_; // 现在这个变量已在 init 中初始化
|
||||
app_ctx.output_attrs = output_attrs_; // 只读共享
|
||||
app_ctx.model_width = model_width_;
|
||||
app_ctx.model_height = model_height_;
|
||||
app_ctx.is_quant = true; // 假设模型是量化的
|
||||
app_ctx.is_quant = true;
|
||||
|
||||
letterbox_t letter_box_param;
|
||||
letter_box_param.x_pad = dw;
|
||||
|
|
@ -186,39 +243,8 @@ std::vector<DetectionResult> YoloDetector::detect(const cv::Mat& inputImage) {
|
|||
letter_box_param.scale = ratio;
|
||||
|
||||
object_detect_result_list od_results;
|
||||
// spdlog::info("--- Debugging Raw Output ---");
|
||||
// for (int i = 0; i < io_num_.n_output; ++i) {
|
||||
// int8_t* buffer = (int8_t*)outputs[i].buf;
|
||||
// int size = outputs[i].size;
|
||||
// int32_t zp = output_attrs_[i].zp;
|
||||
// float scale = output_attrs_[i].scale;
|
||||
|
||||
// // 统计这一层输出的最大值、最小值和平均值
|
||||
// float max_val = -9999.0f;
|
||||
// float min_val = 9999.0f;
|
||||
// int8_t raw_max = -128;
|
||||
|
||||
// // 为了性能只采样前 1000 个数值或全部
|
||||
// int check_count = (size > 1000) ? 1000 : size;
|
||||
|
||||
// for (int j = 0; j < check_count; j++) {
|
||||
// // 手动反量化： float = (int - zp) * scale
|
||||
// float val = (buffer[j] - zp) * scale;
|
||||
// if (val > max_val)
|
||||
// max_val = val;
|
||||
// if (val < min_val)
|
||||
// min_val = val;
|
||||
// if (buffer[j] > raw_max)
|
||||
// raw_max = buffer[j];
|
||||
// }
|
||||
|
||||
// spdlog::info("Output[{}]: Size={}, ZP={}, Scale={}", i, size, zp, scale);
|
||||
// spdlog::info(" -> Raw int8 max: {}, Dequantized float range: [{:.4f}, {:.4f}]", raw_max,
|
||||
// min_val, max_val);
|
||||
// }
|
||||
// spdlog::info("----------------------------");
|
||||
// === [调试代码结束] ===
|
||||
// 调用官方/Demo提供的后处理
|
||||
// 执行后处理 (纯 CPU 计算)
|
||||
post_process(&app_ctx, outputs.data(), &letter_box_param, obj_thresh_, nms_thresh_,
|
||||
&od_results);
|
||||
|
||||
|
|
@ -233,20 +259,21 @@ std::vector<DetectionResult> YoloDetector::detect(const cv::Mat& inputImage) {
|
|||
res.confidence = det->prop;
|
||||
res.class_id = det->cls_id;
|
||||
res.label = coco_cls_to_name(det->cls_id);
|
||||
|
||||
results.push_back(res); // [修复] 变量名从 final_results 改为 results
|
||||
results.push_back(res);
|
||||
}
|
||||
|
||||
// === 6. 释放 RKNN 输出内存 (必须！) ===
|
||||
rknn_outputs_release(ctx_, io_num_.n_output, outputs.data());
|
||||
// === 6. 释放 RKNN 输出内存 ===
|
||||
rknn_outputs_release(current_ctx, io_num_.n_output, outputs.data());
|
||||
|
||||
// === 7. 归还资源 ===
|
||||
// 非常重要！不归还的话，后续请求会一直死锁
|
||||
return_context_id(ctx_id);
|
||||
|
||||
return results;
|
||||
}
|
||||
|
||||
// [修复] 将 nms_boxes 移出 detect 函数，单独实现
|
||||
// 注意：post_process 内部通常已经包含了 NMS，这里如果是为了二次处理可以保留，
|
||||
// 否则可以直接使用 post_process 的输出。
|
||||
void YoloDetector::nms_boxes(std::vector<DetectionResult>& boxes, float nms_thresh) {
|
||||
// 保持原有逻辑不变
|
||||
std::sort(boxes.begin(), boxes.end(), [](const DetectionResult& a, const DetectionResult& b) {
|
||||
return a.confidence > b.confidence;
|
||||
});
|
||||
|
|
@ -258,27 +285,21 @@ void YoloDetector::nms_boxes(std::vector<DetectionResult>& boxes, float nms_thre
|
|||
for (size_t j = i + 1; j < boxes.size(); ++j) {
|
||||
if (is_suppressed[j])
|
||||
continue;
|
||||
|
||||
// 计算 IoU
|
||||
int xx1 = std::max(boxes[i].x, boxes[j].x);
|
||||
int yy1 = std::max(boxes[i].y, boxes[j].y);
|
||||
int xx2 = std::min(boxes[i].x + boxes[i].width, boxes[j].x + boxes[j].width);
|
||||
int yy2 = std::min(boxes[i].y + boxes[i].height, boxes[j].y + boxes[j].height);
|
||||
|
||||
int w = std::max(0, xx2 - xx1);
|
||||
int h = std::max(0, yy2 - yy1);
|
||||
int inter_area = w * h;
|
||||
int union_area =
|
||||
boxes[i].width * boxes[i].height + boxes[j].width * boxes[j].height - inter_area;
|
||||
|
||||
float iou = (float)inter_area / union_area;
|
||||
if (iou > nms_thresh) {
|
||||
is_suppressed[j] = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 移除被抑制的框
|
||||
std::vector<DetectionResult> keep_boxes;
|
||||
for (size_t i = 0; i < boxes.size(); ++i) {
|
||||
if (!is_suppressed[i]) {
|
||||
|
|
|
|||
|
|
@ -1,6 +1,9 @@
|
|||
#pragma once
|
||||
|
||||
#include <condition_variable>
|
||||
#include <mutex>
|
||||
#include <opencv2/opencv.hpp>
|
||||
#include <queue>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
|
|
@ -23,24 +26,37 @@ public:
|
|||
~YoloDetector();
|
||||
|
||||
int init(const std::string& modelPath);
|
||||
// detect 现在是线程安全的，支持多线程并发调用
|
||||
std::vector<DetectionResult> detect(const cv::Mat& inputImage);
|
||||
|
||||
private:
|
||||
unsigned char* load_model(const char* filename, int* model_size);
|
||||
|
||||
// NMS 函数声明
|
||||
void nms_boxes(std::vector<DetectionResult>& boxes, float nms_thresh);
|
||||
|
||||
cv::Mat letterbox(const cv::Mat& src, int target_width, int target_height, float& ratio,
|
||||
int& dw, int& dh);
|
||||
|
||||
// 获取一个空闲的 NPU 上下文索引
|
||||
int get_context_id();
|
||||
// 归还 NPU 上下文索引
|
||||
void return_context_id(int id);
|
||||
|
||||
private:
|
||||
rknn_context ctx_;
|
||||
// NPU 核心数量
|
||||
static const int NPU_CORE_CNT = 3;
|
||||
|
||||
// 使用 vector 存储多个 context
|
||||
std::vector<rknn_context> ctx_pool_;
|
||||
|
||||
// 资源管理：空闲 context 的索引队列
|
||||
std::queue<int> free_ctx_indices_;
|
||||
std::mutex mtx_;
|
||||
std::condition_variable cv_;
|
||||
|
||||
bool is_initialized_;
|
||||
|
||||
// 模型输入输出参数
|
||||
rknn_input_output_num io_num_; // [新增] 存储输入输出数量
|
||||
rknn_tensor_attr* output_attrs_; // [新增] 存储输出属性(zp, scale等)
|
||||
// 模型输入输出参数 (这些是共享的元数据，不需要复制)
|
||||
rknn_input_output_num io_num_;
|
||||
rknn_tensor_attr* output_attrs_;
|
||||
|
||||
int model_width_;
|
||||
int model_height_;
|
||||
|
|
@ -49,9 +65,7 @@ private:
|
|||
unsigned char* model_data_;
|
||||
int model_data_size_;
|
||||
|
||||
// 配置参数
|
||||
const std::vector<std::string> CLASS_NAMES = {"fuel_car", "new_energy_car"};
|
||||
// 这里的阈值如果 postprocess 内部用了，这里可以作为 fallback 或二次筛选
|
||||
const float obj_thresh_ = 0.25f;
|
||||
const float nms_thresh_ = 0.45f;
|
||||
};
|
||||
|
|
|
|||
Loading…
Reference in New Issue