diff --git a/data/test_car2.mp4 b/data/test_car2.mp4 new file mode 100644 index 0000000..54b7112 Binary files /dev/null and b/data/test_car2.mp4 differ diff --git a/docker/Dockerfile b/docker/Dockerfile index 9db46f7..4f65d94 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -5,6 +5,7 @@ RUN apt-get update && \ # 基础工具 sudo \ build-essential \ + libmysqlclient-dev \ cmake \ git \ gdb \ diff --git a/src/main.cpp b/src/main.cpp index 57e8be3..e407dcf 100644 --- a/src/main.cpp +++ b/src/main.cpp @@ -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"; diff --git a/src/videoService/video_pipeline.cpp b/src/videoService/video_pipeline.cpp index e72f132..e25b9da 100644 --- a/src/videoService/video_pipeline.cpp +++ b/src/videoService/video_pipeline.cpp @@ -2,7 +2,7 @@ #include // for std::max, std::min #include - +const int YoloDetector::NPU_CORE_CNT; VideoPipeline::VideoPipeline() : running_(false), next_track_id_(0) { detector_ = std::make_unique(); // 模型路径 @@ -39,6 +39,37 @@ void VideoPipeline::Stop() { processingThread_.join(); spdlog::info("VideoPipeline Stopped."); } +void VideoPipeline::inferenceWorker() { + while (running_) { + FrameData data; + { + std::unique_lock 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 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; + 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 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 results = detector_->detect(frame); - auto t3 = Clock::now(); - // 2. [修改] 更新追踪器和平滑分数 - updateTracker(results); - auto t4 = Clock::now(); - // 3. [修改] 准备绘图数据 - std::vector 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 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(t2 - t1).count(); - double ms_infer = std::chrono::duration_cast(t3 - t2).count(); - double ms_track = std::chrono::duration_cast(t4 - t3).count(); - double ms_draw = std::chrono::duration_cast(t5 - t4).count(); - double ms_write = std::chrono::duration_cast(t6 - t5).count(); - double ms_total = std::chrono::duration_cast(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(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 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(); } diff --git a/src/videoService/video_pipeline.hpp b/src/videoService/video_pipeline.hpp index a733b34..bb8a9f7 100644 --- a/src/videoService/video_pipeline.hpp +++ b/src/videoService/video_pipeline.hpp @@ -1,9 +1,12 @@ #pragma once #include -#include // [新增] +#include +#include #include +#include #include +#include #include #include #include @@ -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 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& trackedObjects); - - // [新增] 核心：更新追踪和分数逻辑 void updateTracker(const std::vector& detections); - - // [新增] 辅助函数：计算 IoU (交并比) float computeIOU(const cv::Rect& box1, const cv::Rect& box2); private: std::atomic running_; - std::thread processingThread_; + std::thread processingThread_; // 主处理线程 std::unique_ptr detector_; - // [新增] 追踪列表，Key是ID + // 追踪相关 std::map tracks_; int next_track_id_ = 0; + + // === [新增] 多线程并行处理相关 === + std::vector worker_threads_; // 3个工作线程 + + // 输入队列 (待检测) + std::queue input_queue_; + std::mutex input_mtx_; + std::condition_variable input_cv_; + const size_t MAX_INPUT_QUEUE_SIZE = 5; // 限制缓冲大小，防止内存爆炸 + + // 输出缓冲区 (已检测，等待排序) + // 使用 map 自动按 frame_id 排序，解决乱序问题 + std::map output_buffer_; + std::mutex output_mtx_; + std::condition_variable output_cv_; }; diff --git a/src/yoloDetector/yolo_detector.cpp b/src/yoloDetector/yolo_detector.cpp index 45979e6..5a48a68 100644 --- a/src/yoloDetector/yolo_detector.cpp +++ b/src/yoloDetector/yolo_detector.cpp @@ -1,21 +1,26 @@ #include "yolo_detector.hpp" #include -#include // for memset +#include #include #include #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 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 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 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 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 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 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 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& 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& 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 keep_boxes; for (size_t i = 0; i < boxes.size(); ++i) { if (!is_suppressed[i]) { diff --git a/src/yoloDetector/yolo_detector.hpp b/src/yoloDetector/yolo_detector.hpp index 8d66193..87e9dff 100644 --- a/src/yoloDetector/yolo_detector.hpp +++ b/src/yoloDetector/yolo_detector.hpp @@ -1,6 +1,9 @@ #pragma once +#include +#include #include +#include #include #include @@ -23,24 +26,37 @@ public: ~YoloDetector(); int init(const std::string& modelPath); + // detect 现在是线程安全的，支持多线程并发调用 std::vector detect(const cv::Mat& inputImage); private: unsigned char* load_model(const char* filename, int* model_size); - - // NMS 函数声明 void nms_boxes(std::vector& 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 ctx_pool_; + + // 资源管理：空闲 context 的索引队列 + std::queue 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 CLASS_NAMES = {"fuel_car", "new_energy_car"}; - // 这里的阈值如果 postprocess 内部用了，这里可以作为 fallback 或二次筛选 const float obj_thresh_ = 0.25f; const float nms_thresh_ = 0.45f; };