Vehicle_Road_Counter/src/videoService/video_pipeline.cpp

#include "video_pipeline.hpp"

#include <chrono>

VideoPipeline::VideoPipeline() : running_(false) {
	// [新增] 实例化检测器并加载模型
	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 会直接返回空结果，建议检查日志
	} else {
		spdlog::info("YoloDetector initialized successfully.");
	}
}

VideoPipeline::~VideoPipeline() {
	Stop();
}

void VideoPipeline::Start(const std::string& inputUrl, const std::string& outputUrl) {
	if (running_)
		return;
	running_ = true;
	spdlog::info("Starting VideoPipeline with Input: {}", inputUrl);

	processingThread_ = std::thread(&VideoPipeline::processLoop, this, inputUrl, outputUrl, false);
}

void VideoPipeline::StartTest(const std::string& filePath, const std::string& outputUrl) {
	if (running_)
		return;
	running_ = true;
	spdlog::info("Starting VideoPipeline (File Test Mode) Input: {}", filePath);

	// true 表示是文件源
	processingThread_ = std::thread(&VideoPipeline::processLoop, this, filePath, outputUrl, true);
}

void VideoPipeline::Stop() {
	if (!running_)
		return;
	running_ = false;
	if (processingThread_.joinable()) {
		processingThread_.join();
	}
	spdlog::info("VideoPipeline Stopped.");
}

// [删除] mockInference 函数已移除

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);

		cv::rectangle(frame, cv::Rect(res.x, res.y, res.width, res.height), color, 2);

		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);
	}

	// 添加水印
	cv::putText(frame, "RK3588 YOLOv8 Live", cv::Point(20, 50), cv::FONT_HERSHEY_SIMPLEX, 1.0,
				cv::Scalar(0, 255, 255), 2);
}

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;
	}

	// 降低一点目标 FPS，因为推理需要时间。RK3588 NPU 很快，但 25/30 FPS 比较稳妥
	const double TARGET_FPS = 30.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");

	std::stringstream pipeline;
	// 注意：推流分辨率这里保持和输入一致，或者可以 resize
	pipeline << "appsrc ! "
			 << "videoconvert ! "
			 << "video/x-raw,format=NV12,width=" << width << ",height=" << height
			 << ",framerate=" << (int)TARGET_FPS << "/1 ! "
			 << "mpph264enc ! "
			 << "h264parse ! "
			 << "rtspclientsink location=" << outputUrl << " protocols=tcp";

	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;
	while (running_) {
		auto loop_start = std::chrono::steady_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推理、后处理和坐标还原
		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);
			}
		}
		// === 2. 绘制叠加 ===
		drawOverlay(frame, results);

		// === 3. 推流 ===
		if (writer.isOpened()) {
			writer.write(frame);
		}

		// 文件模式下控制播放速度，直播流则不需要 wait
		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 wait_ms = FRAME_DURATION_MS - elapsed_ms;

			if (wait_ms > 0) {
				std::this_thread::sleep_for(std::chrono::milliseconds((int)wait_ms));
			}
		}
	}

	cap.release();
	writer.release();
}