rtsp_tensorrt/pipeline/input/rtsp_reader.cpp

#include "rtsp_reader.hpp"
#include <opencv2/core.hpp>
#include <opencv2/imgproc.hpp>

namespace pipeline {

RtspReader::RtspReader(const Config& config)
    : config_(config)
    , last_frame_time_(std::chrono::steady_clock::now())
    , last_retry_time_(std::chrono::steady_clock::now()) {
}

RtspReader::~RtspReader() {
    disconnect();
}

bool RtspReader::connect(const std::string& url) {
    // 如果已经连接，先断开
    if (is_connected_) {
        disconnect();
    }

    // 保存URL
    url_ = url;
    retry_count_ = 0;
    last_retry_time_ = std::chrono::steady_clock::now();

    // 如果是模拟源
    if (url_.substr(0, 7) == "mock://") {
        is_connected_ = true;
        mock_frame_ = cv::Mat(480, 640, CV_8UC3, cv::Scalar(0, 255, 0));
        cv::putText(mock_frame_, "Test Frame", cv::Point(50, 50),
                   cv::FONT_HERSHEY_SIMPLEX, 1.0, cv::Scalar(255, 255, 255), 2);
        return true;
    }

    // 设置RTSP流参数并尝试连接
    setRtspParams();
    return tryReconnect();
}

void RtspReader::disconnect() {
    if (is_connected_) {
        if (url_.substr(0, 7) != "mock://") {
            cap_.release();
        }
        is_connected_ = false;
        retry_count_ = 0;  // 重置重试次数
        last_retry_time_ = std::chrono::steady_clock::now();  // 重置重试时间
    }
}

bool RtspReader::read(cv::Mat& frame) {
    if (!is_connected_) {
        if (!tryReconnect()) {
            return false;
        }
    }

    // 控制帧率
    controlFrameRate();

    // 如果是模拟源
    if (url_.substr(0, 7) == "mock://") {
        if (!mock_frame_.empty()) {
            frame = mock_frame_.clone();
            updateFrameRate();
            return true;
        }
        return false;
    }

    // 读取一帧（带超时）
    bool success = false;
    auto start_time = std::chrono::steady_clock::now();
    while (std::chrono::duration_cast<std::chrono::milliseconds>(
            std::chrono::steady_clock::now() - start_time).count() < config_.frame_timeout_ms) {
        if (cap_.read(frame)) {
            success = true;
            break;
        }
        std::this_thread::sleep_for(std::chrono::milliseconds(1));
    }

    // 读取失败，可能需要重连
    if (!success || frame.empty()) {
        is_connected_ = false;
        return false;
    }

    // 更新帧率统计
    updateFrameRate();
    return true;
}

bool RtspReader::tryReconnect() {
    // 如果是模拟源
    if (url_.substr(0, 7) == "mock://") {
        if (mock_frame_.empty()) {
            mock_frame_ = cv::Mat(480, 640, CV_8UC3, cv::Scalar(0, 255, 0));
            cv::putText(mock_frame_, "Test Frame", cv::Point(50, 50),
                       cv::FONT_HERSHEY_SIMPLEX, 1.0, cv::Scalar(255, 255, 255), 2);
        }
        is_connected_ = true;
        retry_count_ = 0;
        return true;
    }

    // 检查重试次数和间隔
    auto now = std::chrono::steady_clock::now();
    if (retry_count_ >= config_.max_retry_count) {
        return false;
    }

    auto elapsed = std::chrono::duration_cast<std::chrono::milliseconds>(
        now - last_retry_time_).count();
    if (elapsed < config_.retry_interval_ms && retry_count_ > 0) {
        return false;
    }

    // 尝试重新打开RTSP流
    if (!cap_.open(url_, cv::CAP_FFMPEG)) {
        retry_count_++;
        last_retry_time_ = now;
        return false;
    }

    // 检查连接是否成功
    if (!cap_.isOpened()) {
        retry_count_++;
        last_retry_time_ = now;
        return false;
    }

    // 连接成功
    is_connected_ = true;
    retry_count_ = 0;
    return true;
}

void RtspReader::updateFrameRate() {
    auto now = std::chrono::steady_clock::now();
    float delta = std::chrono::duration_cast<std::chrono::milliseconds>(
        now - last_frame_time_).count() / 1000.0f;
    current_fps_ = 1.0f / delta;
    last_frame_time_ = now;
}

void RtspReader::controlFrameRate() {
    if (config_.target_fps <= 0) return;

    auto now = std::chrono::steady_clock::now();
    float target_delta = 1.0f / config_.target_fps;
    float actual_delta = std::chrono::duration_cast<std::chrono::milliseconds>(
        now - last_frame_time_).count() / 1000.0f;

    if (actual_delta < target_delta) {
        auto sleep_time = static_cast<int>(
            (target_delta - actual_delta) * 1000);
        std::this_thread::sleep_for(std::chrono::milliseconds(sleep_time));
    }
}

void RtspReader::setRtspParams() {
    // 设置缓冲区大小
    cap_.set(cv::CAP_PROP_BUFFERSIZE, config_.buffer_size);
    
    // 设置RTSP相关参数
    // 使用MJPEG编码器，可以降低延迟
    cap_.set(cv::CAP_PROP_FOURCC, cv::VideoWriter::fourcc('M','J','P','G'));
}

} // namespace pipeline