rtsp_tensorrt/pipeline/input/video_reader.cpp

#include "video_reader.hpp"
#include <filesystem>
#include <iostream>

namespace pipeline {

VideoReader::VideoReader(const Config& config)
    : config_(config)
    , stop_flag_(false)
    , is_opened_(false)
    , current_frame_(0)
    , total_frames_(0)
    , original_fps_(0.0f)
    , current_fps_(0.0f)
{
    last_frame_time_ = std::chrono::steady_clock::now();
}

VideoReader::~VideoReader() {
    try {
        if (is_opened_) {
            close();
        }
    } catch (const std::exception& e) {
        std::cerr << "Error in VideoReader destructor: " << e.what() << std::endl;
    } catch (...) {
        std::cerr << "Unknown error in VideoReader destructor" << std::endl;
    }
}

bool VideoReader::open(const std::string& filepath) {
    // 如果已经打开，先关闭
    if (is_opened_) {
        close();
    }

    // 打开视频文件
    if (!cap_.open(filepath)) {
        std::cerr << "Error: Cannot open video file: " << filepath << std::endl;
        return false;
    }

    // 获取视频信息
    total_frames_ = static_cast<int>(cap_.get(cv::CAP_PROP_FRAME_COUNT));
    original_fps_ = cap_.get(cv::CAP_PROP_FPS);
    current_fps_ = original_fps_;
    filepath_ = filepath;

    // 尝试读取第一帧
    cv::Mat test_frame;
    if (!cap_.read(test_frame)) {
        std::cerr << "Error: Cannot read frames from video" << std::endl;
        cap_.release();
        return false;
    }

    // 重置视频位置
    if (!cap_.set(cv::CAP_PROP_POS_FRAMES, 0)) {
        std::cerr << "Error: Cannot reset video position" << std::endl;
        cap_.release();
        return false;
    }

    // 视频打开成功后再创建队列和启动线程
    is_opened_ = true;
    current_frame_ = 0;
    last_frame_time_ = std::chrono::steady_clock::now();

    // 如果配置为异步读取，创建队列并启动读取线程
    if (config_.async_reading) {
        try {
            frame_queue_ = std::make_unique<FrameQueue>(config_.buffer_size);
            stop_flag_ = false;
            read_thread_ = std::make_unique<std::thread>(&VideoReader::readLoop, this);
        } catch (const std::exception& e) {
            std::cerr << "Error: Failed to start async reading: " << e.what() << std::endl;
            close();
            return false;
        }
    }

    return true;
}

void VideoReader::close() {
    if (!is_opened_) {
        return;
    }

    try {
        // 停止异步读取线程
        if (config_.async_reading && read_thread_) {
            // 1. 设置停止标志
            stop_flag_ = true;

            // 2. 关闭队列,防止新的帧被添加
            if (frame_queue_) {
                frame_queue_->close();
            }

            // 3. 等待线程完全停止
            if (read_thread_->joinable()) {
                read_thread_->join();
            }

            // 4. 清理线程
            read_thread_.reset();

            // 5. 清理队列
            if (frame_queue_) {
                frame_queue_->clear();
                frame_queue_.reset();
            }
        }

        // 6. 释放视频资源
        cap_.release();
    } catch (const std::exception& e) {
        std::cerr << "Error in close(): " << e.what() << std::endl;
    } catch (...) {
        std::cerr << "Unknown error in close()" << std::endl;
    }

    // 7. 重置状态
    is_opened_ = false;
    current_frame_ = 0;
    total_frames_ = 0;
    current_fps_ = 0.0f;
    original_fps_ = 0.0f;
    filepath_.clear();
    stop_flag_ = false;  // 重置停止标志
}

bool VideoReader::checkVideoFormat() {
    // 检查视频是否成功打开
    if (!cap_.isOpened()) {
        std::cerr << "Error: Invalid video capture object" << std::endl;
        return false;
    }

    // 检查基本视频属性
    double fps = cap_.get(cv::CAP_PROP_FPS);
    if (fps <= 0) {
        std::cerr << "Error: Invalid video FPS: " << fps << std::endl;
        return false;
    }

    int width = static_cast<int>(cap_.get(cv::CAP_PROP_FRAME_WIDTH));
    int height = static_cast<int>(cap_.get(cv::CAP_PROP_FRAME_HEIGHT));
    if (width <= 0 || height <= 0) {
        std::cerr << "Error: Invalid video dimensions: " << width << "x" << height << std::endl;
        return false;
    }

    int codec = static_cast<int>(cap_.get(cv::CAP_PROP_FOURCC));
    if (codec == 0) {
        std::cerr << "Warning: Unknown video codec" << std::endl;
    }

    // 尝试读取第一帧
    cv::Mat test_frame;
    if (!cap_.read(test_frame)) {
        std::cerr << "Error: Cannot read frames from video" << std::endl;
        return false;
    }
    
    // 如果成功读取第一帧，需要重置视频位置
    cap_.set(cv::CAP_PROP_POS_FRAMES, 0);
    
    return true;
}

void VideoReader::updateFrameRate() {
    static auto last_fps_update = std::chrono::steady_clock::now();
    static int frame_count = 0;

    frame_count++;
    auto now = std::chrono::steady_clock::now();
    auto duration = std::chrono::duration_cast<std::chrono::microseconds>(
        now - last_fps_update).count();

    // 每秒更新一次帧率
    if (duration >= 1000000) {  // 1秒 = 1000000微秒
        current_fps_ = static_cast<float>(frame_count * 1000000) / duration;
        frame_count = 0;
        last_fps_update = now;
    }
}

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

    auto now = std::chrono::steady_clock::now();
    auto frame_duration = std::chrono::microseconds(static_cast<long long>(1000000.0f / config_.target_fps));
    auto elapsed = std::chrono::duration_cast<std::chrono::microseconds>(now - last_frame_time_);

    if (elapsed < frame_duration) {
        std::this_thread::sleep_for(frame_duration - elapsed);
    }

    last_frame_time_ = std::chrono::steady_clock::now();
}

bool VideoReader::read(cv::Mat& frame) {
    if (!is_opened_) {
        std::cerr << "Error: Video is not opened" << std::endl;
        return false;
    }

    // 异步读取模式
    if (config_.async_reading) {
        if (!frame_queue_) {
            std::cerr << "Error: Frame queue is not initialized" << std::endl;
            return false;
        }
        return frame_queue_->pop(frame, std::chrono::milliseconds(config_.frame_timeout_ms));
    }

    // 同步读取模式
    return readNextFrame(frame);
}

bool VideoReader::readNextFrame(cv::Mat& frame) {
    // 检查是否到达视频末尾
    if (current_frame_ >= total_frames_) {
        if (config_.loop_playback) {
            restart();
            // 在循环时添加延迟，避免过快读取
            std::this_thread::sleep_for(std::chrono::milliseconds(100));
        } else {
            return false;
        }
    }

    // 控制帧率
    auto now = std::chrono::steady_clock::now();
    auto frame_duration = std::chrono::microseconds(
        static_cast<long long>(1000000.0f / config_.target_fps));
    auto elapsed = std::chrono::duration_cast<std::chrono::microseconds>(
        now - last_frame_time_);
    if (elapsed < frame_duration) {
        std::this_thread::sleep_for(frame_duration - elapsed);
    }
    last_frame_time_ = std::chrono::steady_clock::now();

    // 读取帧
    if (!cap_.read(frame)) {
        if (config_.loop_playback) {
            restart();
            return cap_.read(frame);
        }
        return false;
    }

    // 更新计数器和帧率
    current_frame_++;
    updateFrameRate();

    return true;
}

void VideoReader::readLoop() {
    if (!cap_.isOpened() || !frame_queue_) {
        std::cerr << "Video capture or frame queue is not initialized" << std::endl;
        return;
    }

    std::cout << "Starting read loop for video: " << filepath_ << std::endl;

    cv::Mat frame;
    const int PUSH_TIMEOUT_MS = 1000;
    const int MAX_RETRIES = 10;
    int consecutive_errors = 0;
    const int MAX_CONSECUTIVE_ERRORS = 5;

    auto last_frame_time = std::chrono::steady_clock::now();
    auto frame_duration = std::chrono::microseconds(
        static_cast<long long>(1000000.0f / config_.target_fps));

    try {
        while (!stop_flag_) {
            if (!cap_.read(frame)) {
                consecutive_errors++;
                if (consecutive_errors >= MAX_CONSECUTIVE_ERRORS) {
                    std::cerr << "Too many consecutive read errors, attempting recovery..." << std::endl;
                    
                    // 尝试重新打开视频文件
                    cap_.release();
                    if (!cap_.open(filepath_)) {
                        std::cerr << "Failed to recover video stream" << std::endl;
                        break;
                    }
                    
                    // 重置位置
                    if (!cap_.set(cv::CAP_PROP_POS_FRAMES, 0)) {
                        std::cerr << "Failed to reset video position during recovery" << std::endl;
                        break;
                    }
                    
                    consecutive_errors = 0;
                    continue;  // 重试读取
                }

                if (config_.loop_playback) {
                    if (!cap_.set(cv::CAP_PROP_POS_FRAMES, 0)) {
                        std::cerr << "Failed to reset video position for loop playback" << std::endl;
                        break;
                    }
                    std::this_thread::sleep_for(std::chrono::milliseconds(100));
                    continue;
                } else {
                    break;
                }
            }

            consecutive_errors = 0;  // 重置错误计数

            if (frame.empty()) {
                continue;
            }

            // 控制帧率
            auto now = std::chrono::steady_clock::now();
            auto elapsed = std::chrono::duration_cast<std::chrono::microseconds>(now - last_frame_time);
            if (elapsed < frame_duration) {
                std::this_thread::sleep_for(frame_duration - elapsed);
            }
            last_frame_time = std::chrono::steady_clock::now();

            // 检查frame_queue_是否有效
            if (!frame_queue_) {
                std::cerr << "Frame queue is not available" << std::endl;
                break;
            }

            // 尝试推送帧到队列
            int retries = 0;
            bool pushed = false;
            while (!stop_flag_ && retries < MAX_RETRIES) {
                if (frame_queue_->push(frame, std::chrono::milliseconds(PUSH_TIMEOUT_MS))) {
                    pushed = true;
                    break;
                }
                retries++;
                std::this_thread::sleep_for(std::chrono::milliseconds(10));
            }

            if (!pushed && !stop_flag_) {
                std::cerr << "Failed to push frame to queue after " << MAX_RETRIES << " retries" << std::endl;
            }
        }
    } catch (const std::exception& e) {
        std::cerr << "Error in read loop: " << e.what() << std::endl;
    } catch (...) {
        std::cerr << "Unknown error in read loop" << std::endl;
    }

    // 确保在退出时关闭队列
    if (frame_queue_) {
        frame_queue_->close();
    }

    std::cout << "Exiting read loop for video: " << filepath_ << std::endl;
}

bool VideoReader::hasFrames() const {
    if (!config_.async_reading) {
        return is_opened_ && (current_frame_ < total_frames_ || config_.loop_playback);
    }
    return frame_queue_ && !frame_queue_->empty();
}

size_t VideoReader::getQueueSize() const {
    if (!config_.async_reading || !frame_queue_) {
        return 0;
    }
    return frame_queue_->size();
}

size_t VideoReader::getQueueCapacity() const {
    if (!config_.async_reading) {
        return 0;
    }
    return frame_queue_->capacity();
}

bool VideoReader::seek(int frame_number) {
    if (!is_opened_) {
        std::cerr << "Error: Video is not opened" << std::endl;
        return false;
    }

    if (frame_number < 0 || frame_number >= total_frames_) {
        std::cerr << "Error: Invalid frame number: " << frame_number << std::endl;
        return false;
    }

    // 如果在异步模式下，需要先停止读取线程
    bool was_async = false;
    if (config_.async_reading && read_thread_) {
        was_async = true;
        stop_flag_ = true;
        if (frame_queue_) {
            frame_queue_->close();
            frame_queue_->clear();
        }
        if (read_thread_->joinable()) {
            read_thread_->join();
        }
        read_thread_.reset();
        frame_queue_.reset();  // 重置队列
    }

    // 设置视频位置
    if (!cap_.set(cv::CAP_PROP_POS_FRAMES, frame_number)) {
        std::cerr << "Error: Failed to seek to frame " << frame_number << std::endl;
        // 如果seek失败，尝试重新打开视频文件
        cap_.release();
        if (!cap_.open(filepath_)) {
            std::cerr << "Error: Failed to reopen video file after seek failure" << std::endl;
            return false;
        }
        if (!cap_.set(cv::CAP_PROP_POS_FRAMES, frame_number)) {
            std::cerr << "Error: Failed to seek after reopening video" << std::endl;
            return false;
        }
    }

    current_frame_ = frame_number;
    last_frame_time_ = std::chrono::steady_clock::now();

    // 如果之前是异步模式，重新启动读取线程
    if (was_async) {
        try {
            frame_queue_ = std::make_unique<FrameQueue>(config_.buffer_size);
            stop_flag_ = false;
            read_thread_ = std::make_unique<std::thread>(&VideoReader::readLoop, this);
        } catch (const std::exception& e) {
            std::cerr << "Error: Failed to restart async reading: " << e.what() << std::endl;
            return false;
        }
    }

    return true;
}

void VideoReader::restart() {
    if (!is_opened_) {
        return;
    }

    // 如果在异步模式下，需要先停止读取线程
    bool was_async = false;
    if (config_.async_reading && read_thread_) {
        was_async = true;
        stop_flag_ = true;
        if (frame_queue_) {
            frame_queue_->close();
            frame_queue_->clear();
        }
        if (read_thread_->joinable()) {
            read_thread_->join();
        }
        read_thread_.reset();
        frame_queue_.reset();  // 重置队列
    }

    // 尝试重置视频位置
    if (!cap_.set(cv::CAP_PROP_POS_FRAMES, 0)) {
        // 如果重置失败，重新打开视频文件
        cap_.release();
        if (!cap_.open(filepath_)) {
            std::cerr << "Error: Failed to reopen video file during restart" << std::endl;
            return;
        }
    }

    current_frame_ = 0;
    last_frame_time_ = std::chrono::steady_clock::now();

    // 如果之前是异步模式，重新启动读取线程
    if (was_async) {
        try {
            frame_queue_ = std::make_unique<FrameQueue>(config_.buffer_size);
            stop_flag_ = false;
            read_thread_ = std::make_unique<std::thread>(&VideoReader::readLoop, this);
        } catch (const std::exception& e) {
            std::cerr << "Error: Failed to restart async reading: " << e.what() << std::endl;
            return;
        }
    }
}

} // namespace pipeline