generated from guanyuankai/bonus-edge-proxy
feat(video Service): 完成ev score平均分类。
This commit is contained in:
parent
9f6472c219
commit
5a241a3e9e
|
|
@ -1,16 +1,13 @@
|
|||
#include "video_pipeline.hpp"
|
||||
|
||||
#include <algorithm> // for std::max, std::min
|
||||
#include <chrono>
|
||||
|
||||
VideoPipeline::VideoPipeline() : running_(false) {
|
||||
// [新增] 实例化检测器并加载模型
|
||||
VideoPipeline::VideoPipeline() : running_(false), next_track_id_(0) {
|
||||
detector_ = std::make_unique<YoloDetector>();
|
||||
|
||||
// 模型路径写死为 models/vehicle_model.rknn
|
||||
int ret = detector_->init("../models/vehicle_model.rknn");
|
||||
if (ret != 0) {
|
||||
spdlog::error("Failed to initialize YoloDetector with model: models/vehicle_model.rknn");
|
||||
// 注意：这里如果初始化失败，后续 detect 会直接返回空结果，建议检查日志
|
||||
// 模型路径
|
||||
if (detector_->init("../models/vehicle_model.rknn") != 0) {
|
||||
spdlog::error("Failed to initialize YoloDetector");
|
||||
} else {
|
||||
spdlog::info("YoloDetector initialized successfully.");
|
||||
}
|
||||
|
|
@ -24,8 +21,6 @@ void VideoPipeline::Start(const std::string& inputUrl, const std::string& output
|
|||
if (running_)
|
||||
return;
|
||||
running_ = true;
|
||||
spdlog::info("Starting VideoPipeline with Input: {}", inputUrl);
|
||||
|
||||
processingThread_ = std::thread(&VideoPipeline::processLoop, this, inputUrl, outputUrl, false);
|
||||
}
|
||||
|
||||
|
|
@ -33,9 +28,6 @@ void VideoPipeline::StartTest(const std::string& filePath, const std::string& ou
|
|||
if (running_)
|
||||
return;
|
||||
running_ = true;
|
||||
spdlog::info("Starting VideoPipeline (File Test Mode) Input: {}", filePath);
|
||||
|
||||
// true 表示是文件源
|
||||
processingThread_ = std::thread(&VideoPipeline::processLoop, this, filePath, outputUrl, true);
|
||||
}
|
||||
|
||||
|
|
@ -43,33 +35,120 @@ void VideoPipeline::Stop() {
|
|||
if (!running_)
|
||||
return;
|
||||
running_ = false;
|
||||
if (processingThread_.joinable()) {
|
||||
if (processingThread_.joinable())
|
||||
processingThread_.join();
|
||||
}
|
||||
spdlog::info("VideoPipeline Stopped.");
|
||||
}
|
||||
|
||||
// [删除] mockInference 函数已移除
|
||||
// [新增] 计算两个矩形的交并比 (IoU)
|
||||
float VideoPipeline::computeIOU(const cv::Rect& box1, const cv::Rect& box2) {
|
||||
int x1 = std::max(box1.x, box2.x);
|
||||
int y1 = std::max(box1.y, box2.y);
|
||||
int x2 = std::min(box1.x + box1.width, box2.x + box2.width);
|
||||
int y2 = std::min(box1.y + box1.height, box2.y + box2.height);
|
||||
|
||||
void VideoPipeline::drawOverlay(cv::Mat& frame, const std::vector<DetectionResult>& results) {
|
||||
for (const auto& res : results) {
|
||||
// 根据不同类别使用不同颜色 (示例: 0为红色, 1为绿色)
|
||||
cv::Scalar color = (res.class_id == 0) ? cv::Scalar(0, 0, 255) : cv::Scalar(0, 255, 0);
|
||||
if (x1 >= x2 || y1 >= y2)
|
||||
return 0.0f;
|
||||
|
||||
cv::rectangle(frame, cv::Rect(res.x, res.y, res.width, res.height), color, 2);
|
||||
float intersection = (float)((x2 - x1) * (y2 - y1));
|
||||
float area1 = (float)(box1.width * box1.height);
|
||||
float area2 = (float)(box2.width * box2.height);
|
||||
return intersection / (area1 + area2 - intersection);
|
||||
}
|
||||
|
||||
std::string text = res.label + " " + std::to_string(res.confidence).substr(0, 4);
|
||||
|
||||
// 简单的防止文字跑出边界的处理
|
||||
int y_pos = res.y - 5;
|
||||
if (y_pos < 10)
|
||||
y_pos = res.y + 15;
|
||||
|
||||
cv::putText(frame, text, cv::Point(res.x, y_pos), cv::FONT_HERSHEY_SIMPLEX, 0.6, color, 2);
|
||||
// [新增] 核心逻辑：追踪与平滑
|
||||
void VideoPipeline::updateTracker(const std::vector<DetectionResult>& detections) {
|
||||
// 1. 标记所有现有轨迹为丢失 (missing_frames + 1)
|
||||
for (auto& pair : tracks_) {
|
||||
pair.second.missing_frames++;
|
||||
}
|
||||
|
||||
// 添加水印
|
||||
cv::putText(frame, "RK3588 YOLOv8 Live", cv::Point(20, 50), cv::FONT_HERSHEY_SIMPLEX, 1.0,
|
||||
// 2. 匹配当前帧检测结果与现有轨迹
|
||||
for (const auto& det : detections) {
|
||||
cv::Rect detBox(det.x, det.y, det.width, det.height);
|
||||
|
||||
int best_match_id = -1;
|
||||
float max_iou = 0.0f;
|
||||
|
||||
// 寻找 IoU 最大的匹配
|
||||
for (auto& pair : tracks_) {
|
||||
float iou = computeIOU(detBox, pair.second.box);
|
||||
if (iou > 0.3f && iou > max_iou) { // 阈值 0.3 可根据需要调整
|
||||
max_iou = iou;
|
||||
best_match_id = pair.first;
|
||||
}
|
||||
}
|
||||
|
||||
// 假设: class_id == 1 是新能源(Green), class_id == 0 是油车(Fuel)
|
||||
// 请根据你模型的实际定义修改这里的 ID
|
||||
float current_is_ev = (det.class_id == 1) ? 1.0f : 0.0f;
|
||||
|
||||
if (best_match_id != -1) {
|
||||
// === 匹配成功：更新平滑分数 ===
|
||||
TrackedVehicle& track = tracks_[best_match_id];
|
||||
track.box = detBox;
|
||||
track.missing_frames = 0; // 重置丢失计数
|
||||
|
||||
// [核心算法] 指数移动平均 (EMA)
|
||||
// alpha = 0.1 表示新结果占10%权重，历史占90%，数值越小越平滑，反应越慢
|
||||
float alpha = 0.15f;
|
||||
track.ev_score = track.ev_score * (1.0f - alpha) + current_is_ev * alpha;
|
||||
|
||||
} else {
|
||||
// === 未匹配：创建新轨迹 ===
|
||||
TrackedVehicle newTrack;
|
||||
newTrack.id = next_track_id_++;
|
||||
newTrack.box = detBox;
|
||||
newTrack.missing_frames = 0;
|
||||
newTrack.ev_score = current_is_ev; // 初始分数为当前检测结果
|
||||
newTrack.last_class_id = det.class_id;
|
||||
|
||||
tracks_[newTrack.id] = newTrack;
|
||||
}
|
||||
}
|
||||
|
||||
// 3. 移除长时间丢失的轨迹 & 更新显示状态
|
||||
for (auto it = tracks_.begin(); it != tracks_.end();) {
|
||||
if (it->second.missing_frames > 10) { // 超过10帧未检测到则移除
|
||||
it = tracks_.erase(it);
|
||||
} else {
|
||||
// 根据平滑后的分数决定最终类别
|
||||
// 阈值 0.5: 分数 > 0.5 判定为新能源
|
||||
int final_class = (it->second.ev_score > 0.5f) ? 1 : 0;
|
||||
|
||||
it->second.last_class_id = final_class;
|
||||
it->second.label = (final_class == 1) ? "Green" : "Fuel"; // 显示文本
|
||||
|
||||
++it;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// [修改] 绘制函数使用 TrackedVehicle
|
||||
void VideoPipeline::drawOverlay(cv::Mat& frame, const std::vector<TrackedVehicle>& trackedObjects) {
|
||||
for (const auto& trk : trackedObjects) {
|
||||
// 如果丢失了几帧但还在内存里，用虚线或者灰色表示（可选），这里保持原样
|
||||
if (trk.missing_frames > 0)
|
||||
continue; // 暂时只画当前帧存在的
|
||||
|
||||
// 颜色：新能源用绿色，油车用红色
|
||||
cv::Scalar color = (trk.last_class_id == 1) ? cv::Scalar(0, 255, 0) : cv::Scalar(0, 0, 255);
|
||||
|
||||
cv::rectangle(frame, trk.box, color, 2);
|
||||
|
||||
// 显示 类别 + 平滑后的分数
|
||||
// 例如: Green 0.85
|
||||
std::string text = fmt::format("{} {:.2f}", trk.label, trk.ev_score);
|
||||
|
||||
int y_pos = trk.box.y - 5;
|
||||
if (y_pos < 10)
|
||||
y_pos = trk.box.y + 15;
|
||||
|
||||
cv::putText(frame, text, cv::Point(trk.box.x, y_pos), cv::FONT_HERSHEY_SIMPLEX, 0.6, color,
|
||||
2);
|
||||
}
|
||||
|
||||
cv::putText(frame, "RK3588 YOLOv8 Smooth", cv::Point(20, 50), cv::FONT_HERSHEY_SIMPLEX, 1.0,
|
||||
cv::Scalar(0, 255, 255), 2);
|
||||
}
|
||||
|
||||
|
|
@ -83,18 +162,13 @@ void VideoPipeline::processLoop(std::string inputUrl, std::string outputUrl, boo
|
|||
return;
|
||||
}
|
||||
|
||||
// 降低一点目标 FPS，因为推理需要时间。RK3588 NPU 很快，但 25/30 FPS 比较稳妥
|
||||
const double TARGET_FPS = 30.0;
|
||||
const double TARGET_FPS = 60.0;
|
||||
const double FRAME_DURATION_MS = 1000.0 / TARGET_FPS;
|
||||
|
||||
int width = cap.get(cv::CAP_PROP_FRAME_WIDTH);
|
||||
int height = cap.get(cv::CAP_PROP_FRAME_HEIGHT);
|
||||
|
||||
spdlog::info("Source: {}x{} | Mode: {}", width, height,
|
||||
isFileSource ? "FILE LOOP" : "LIVE STREAM");
|
||||
|
||||
// GStreamer pipeline string (保持原样...)
|
||||
std::stringstream pipeline;
|
||||
// 注意：推流分辨率这里保持和输入一致，或者可以 resize
|
||||
pipeline << "appsrc ! "
|
||||
<< "videoconvert ! "
|
||||
<< "video/x-raw,format=NV12,width=" << width << ",height=" << height
|
||||
|
|
@ -106,70 +180,72 @@ 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);
|
||||
|
||||
if (!writer.isOpened()) {
|
||||
spdlog::error("Failed to initialize VideoWriter.");
|
||||
}
|
||||
|
||||
cv::Mat frame;
|
||||
long frame_count = 0;
|
||||
using Clock = std::chrono::high_resolution_clock;
|
||||
using Ms = std::chrono::duration<double, std::milli>;
|
||||
|
||||
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) {
|
||||
spdlog::info("End of file reached, looping...");
|
||||
cap.set(cv::CAP_PROP_POS_FRAMES, 0);
|
||||
continue;
|
||||
} else {
|
||||
spdlog::warn("Frame read failed. Reconnecting...");
|
||||
std::this_thread::sleep_for(std::chrono::seconds(1));
|
||||
cap.release();
|
||||
cap.open(inputUrl);
|
||||
continue;
|
||||
}
|
||||
}
|
||||
|
||||
if (frame.empty())
|
||||
continue;
|
||||
|
||||
// === 1. 真实 RKNN 推理 ===
|
||||
// detect 内部已经包含了预处理、NPU推理、后处理和坐标还原
|
||||
auto t2 = Clock::now();
|
||||
// 1. 推理
|
||||
std::vector<DetectionResult> results = detector_->detect(frame);
|
||||
|
||||
frame_count++;
|
||||
// spdlog::info("frame count {}", frame_count);
|
||||
if (!results.empty()) {
|
||||
spdlog::info("Frame {}: Detected {} objects", frame_count, results.size());
|
||||
for (const auto& res : results) {
|
||||
spdlog::info(" -> [Class: {} ({})] Conf: {:.2f} Box: [{}, {}, {}x{}]",
|
||||
res.class_id, res.label, res.confidence, res.x, res.y, res.width,
|
||||
res.height);
|
||||
}
|
||||
} else {
|
||||
// 如果连续没有检测到，每隔一秒(20帧)提示一次，确认代码还在跑
|
||||
if (frame_count % 20 == 0) {
|
||||
spdlog::debug("Frame {}: No objects detected.", frame_count);
|
||||
}
|
||||
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);
|
||||
}
|
||||
// === 2. 绘制叠加 ===
|
||||
drawOverlay(frame, results);
|
||||
|
||||
// === 3. 推流 ===
|
||||
if (writer.isOpened()) {
|
||||
// 4. [修改] 绘制
|
||||
drawOverlay(frame, tracks_to_draw);
|
||||
auto t5 = Clock::now();
|
||||
// 5. 推流
|
||||
if (writer.isOpened())
|
||||
writer.write(frame);
|
||||
auto t6 = Clock::now();
|
||||
|
||||
// ----------------- 计算耗时 -----------------
|
||||
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);
|
||||
}
|
||||
|
||||
// 文件模式下控制播放速度，直播流则不需要 wait
|
||||
// FPS控制 (保持原样)
|
||||
if (isFileSource) {
|
||||
auto loop_end = std::chrono::steady_clock::now();
|
||||
std::chrono::duration<double, std::milli> elapsed = loop_end - loop_start;
|
||||
|
||||
double elapsed_ms = elapsed.count();
|
||||
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) {
|
||||
if (wait_ms > 0)
|
||||
std::this_thread::sleep_for(std::chrono::milliseconds((int)wait_ms));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -1,52 +1,53 @@
|
|||
#pragma once
|
||||
|
||||
#include <atomic>
|
||||
#include <memory> // [新增] 用于 unique_ptr
|
||||
#include <map> // [新增]
|
||||
#include <memory>
|
||||
#include <opencv2/opencv.hpp>
|
||||
#include <string>
|
||||
#include <thread>
|
||||
#include <vector>
|
||||
|
||||
#include "spdlog/spdlog.h"
|
||||
#include "yoloDetector/yolo_detector.hpp" // [新增] 包含检测器头文件 (同时也包含了 DetectionResult 定义)
|
||||
#include "yoloDetector/yolo_detector.hpp"
|
||||
|
||||
// [删除] 移除此处定义的 DetectionResult，直接使用 yolo_detector.hpp 中的定义
|
||||
/*
|
||||
struct DetectionResult {
|
||||
int x;
|
||||
int y;
|
||||
int width;
|
||||
int height;
|
||||
std::string label;
|
||||
float confidence;
|
||||
// [新增] 定义追踪车辆的结构体
|
||||
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; // 当前显示的标签
|
||||
};
|
||||
*/
|
||||
|
||||
class VideoPipeline {
|
||||
public:
|
||||
VideoPipeline();
|
||||
~VideoPipeline();
|
||||
|
||||
// 启动视频流处理
|
||||
void Start(const std::string& inputUrl, const std::string& outputUrl);
|
||||
void StartTest(const std::string& filePath, const std::string& outputUrl);
|
||||
|
||||
// 停止处理
|
||||
void Stop();
|
||||
|
||||
private:
|
||||
void processLoop(std::string inputUrl, std::string outputUrl, bool isFileSource);
|
||||
|
||||
// [删除] 移除模拟推理函数
|
||||
// std::vector<DetectionResult> mockInference(const cv::Mat& frame);
|
||||
// [修改] 绘图函数现在接收追踪列表，而不是原始检测结果
|
||||
void drawOverlay(cv::Mat& frame, const std::vector<TrackedVehicle>& trackedObjects);
|
||||
|
||||
// 绘图函数
|
||||
void drawOverlay(cv::Mat& frame, const std::vector<DetectionResult>& results);
|
||||
// [新增] 核心：更新追踪和分数逻辑
|
||||
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_;
|
||||
|
||||
// [新增] YOLO 检测器实例
|
||||
std::unique_ptr<YoloDetector> detector_;
|
||||
|
||||
// [新增] 追踪列表，Key是ID
|
||||
std::map<int, TrackedVehicle> tracks_;
|
||||
int next_track_id_ = 0;
|
||||
};
|
||||
|
|
|
|||
Loading…
Reference in New Issue