main_stream.cpp

/**
 * @file main_stream.cpp
 * @brief Entry point for the multi-threaded real-time streaming pipeline (RTSP, local cameras, synthetic feeds).
 * 
 * Part of the TensorRT Inference Pipeline for RF-DETR.
 * 
 * @author Loay Wael (loaywael@github.com)
 * @date 2026
 * 
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 * 
 *     http://www.apache.org/licenses/LICENSE-2.0
 * 
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "core/interfaces.hpp"
#include "engine/trt.hpp"
#include "utils/safe_queue.hpp"
#include "utils/pipeline.hpp"
#include "utils/annotations.hpp"

#include <iostream>
#include <thread>
#include <mutex>
#include <chrono>
#include <memory>
#include <iomanip>
#include <algorithm>

/**
 * @brief Capture Thread function: Continuous decoding from RTSP, local camera, or synthetic stream.
 */
void captureThreadFunc(const std::string& str_rtsp_url, trt::SafeQueue<trt::Frame>& queue_frames) {
    cv::VideoCapture video_cap;
    bool b_is_synthetic = false;
    
    if (str_rtsp_url == "synthetic" || str_rtsp_url.empty()) {
        std::cout << "[Capture] Starting synthetic test video stream generator...\n";
        b_is_synthetic = true;
    } else {
        // If the URL consists purely of digits, treat it as a local camera device index
        bool b_is_device_index = !str_rtsp_url.empty() && std::all_of(str_rtsp_url.begin(), str_rtsp_url.end(), ::isdigit);
        if (b_is_device_index) {
            int si32_device_id = std::stoi(str_rtsp_url);
            std::cout << "[Capture] Connecting to Local Camera: Device " << si32_device_id << " ...\n";
            video_cap.open(si32_device_id, cv::CAP_ANY);
        } else {
            std::cout << "[Capture] Connecting to RTSP Stream: " << str_rtsp_url << " ...\n";
            video_cap.open(str_rtsp_url, cv::CAP_FFMPEG);
        }
        
        if (!video_cap.isOpened()) {
            std::cerr << "[Capture] Warning: Failed to connect to capture source. Falling back to synthetic source!\n";
            b_is_synthetic = true;
        }
    }

    int64_t si64_frame_count = 0;
    int si32_synth_width = 1280;
    int si32_synth_height = 720;
    int si32_shape_x = 100;
    int si32_shape_y = 150;
    int si32_dx = 5;
    int si32_dy = 3;

    while (trt::g_b_running) {
        trt::Frame frame_obj;
        frame_obj.si64_frame_id = ++si64_frame_count;
        frame_obj.si64_timestamp_ms = std::chrono::duration_cast<std::chrono::milliseconds>(
            std::chrono::steady_clock::now().time_since_epoch()
        ).count();
        frame_obj.str_stream_id = b_is_synthetic ? "SyntheticStream" : str_rtsp_url;

        if (b_is_synthetic) {
            // Generate synthetic frame (moving circles/text simulating camera input)
            cv::Mat mat_canvas = cv::Mat::zeros(cv::Size(si32_synth_width, si32_synth_height), CV_8UC3);
            
            // Draw grid background
            for (int si32_i = 0; si32_i < si32_synth_width; si32_i += 80) {
                cv::line(mat_canvas, cv::Point(si32_i, 0), cv::Point(si32_i, si32_synth_height), cv::Scalar(40, 40, 40), 1);
            }
            for (int si32_j = 0; si32_j < si32_synth_height; si32_j += 80) {
                cv::line(mat_canvas, cv::Point(0, si32_j), cv::Point(si32_synth_width, si32_j), cv::Scalar(40, 40, 40), 1);
            }

            // Move and draw dynamic shapes simulating detections
            si32_shape_x += si32_dx;
            si32_shape_y += si32_dy;
            if (si32_shape_x < 50 || si32_shape_x > si32_synth_width - 50) si32_dx = -si32_dx;
            if (si32_shape_y < 50 || si32_shape_y > si32_synth_height - 50) si32_dy = -si32_dy;

            // Simulated Class 0: Moving Target
            cv::circle(mat_canvas, cv::Point(si32_shape_x, si32_shape_y), 40, cv::Scalar(0, 0, 255), -1); // BGR Red
            cv::putText(mat_canvas, "TARGET [Class 0]", cv::Point(si32_shape_x - 60, si32_shape_y - 50), 
                        cv::FONT_HERSHEY_SIMPLEX, 0.6, cv::Scalar(0, 0, 255), 2);

            // Simulated Class 1: Stationary Target
            cv::rectangle(mat_canvas, cv::Rect(400, 300, 100, 100), cv::Scalar(255, 0, 0), -1); // BGR Blue
            cv::putText(mat_canvas, "BOX [Class 1]", cv::Point(390, 280), 
                        cv::FONT_HERSHEY_SIMPLEX, 0.6, cv::Scalar(255, 0, 0), 2);

            // Pipeline info overlay
            cv::putText(mat_canvas, "RTSP EMULATOR (Synthetic Stream Active)", cv::Point(30, 40),
                        cv::FONT_HERSHEY_SIMPLEX, 0.8, cv::Scalar(0, 255, 0), 2);
            cv::putText(mat_canvas, "Frame: " + std::to_string(frame_obj.si64_frame_id), cv::Point(30, 80),
                        cv::FONT_HERSHEY_SIMPLEX, 0.6, cv::Scalar(200, 200, 200), 1);

            frame_obj.mat_image = mat_canvas;
            
            // Limit capture loop speed to ~30 FPS for simulation
            std::this_thread::sleep_for(std::chrono::milliseconds(33));
        } else {
            // Read frame from camera stream
            if (!video_cap.read(frame_obj.mat_image)) {
                std::cerr << "[Capture] Stream disconnected or EOF reached. Retrying connection...\n";
                // For live camera devices or RTSP reconnect
                video_cap.open(str_rtsp_url); 
                std::this_thread::sleep_for(std::chrono::seconds(2));
                continue;
            }
        }

        // Push onto safe queue
        queue_frames.push(frame_obj);
    }
    
    if (video_cap.isOpened()) {
        video_cap.release();
    }
    std::cout << "[Capture] Capture thread terminated successfully.\n";
}

/**
 * @brief Inference Consumer Thread function: Runs TensorRT detection.
 */
void inferenceThreadFunc(std::shared_ptr<trt::TrtEngine> shptr_engine, trt::SafeQueue<trt::Frame>& queue_frames) {
    cudaStream_t cuda_stream;
    cudaStreamCreate(&cuda_stream);

    double fp64_total_latency_ms = 0.0;
    int si32_processed_frames = 0;
    auto time_start = std::chrono::steady_clock::now();

    while (trt::g_b_running || !queue_frames.empty()) {
        trt::Frame frame_obj;
        if (!queue_frames.pop(frame_obj, 100)) {
            continue; // Wait for frames
        }

        auto time_start_infer = std::chrono::steady_clock::now();
        
        // Execute inference
        std::vector<trt::Detection> vdet_detections = shptr_engine->runInference(frame_obj, cuda_stream);
        
        auto time_end_infer = std::chrono::steady_clock::now();
        double fp64_latency_ms = std::chrono::duration<double, std::ratio<1, 1000>>(time_end_infer - time_start_infer).count();
        
        fp64_total_latency_ms += fp64_latency_ms;
        si32_processed_frames++;

        // Print diagnostic statistics every 30 frames
        if (si32_processed_frames % 30 == 0) {
            auto time_current = std::chrono::steady_clock::now();
            double fp64_elapsed_s = std::chrono::duration<double>(time_current - time_start).count();
            double fp64_current_fps = si32_processed_frames / fp64_elapsed_s;
            double fp64_avg_latency_ms = fp64_total_latency_ms / si32_processed_frames;
            
            std::cout << "[Inference] FPS: " << std::fixed << std::setprecision(1) << fp64_current_fps
                      << " | Avg Latency: " << std::setprecision(2) << fp64_avg_latency_ms << " ms"
                      << " | Active Queue Size: " << queue_frames.size() << "\n";
        }

        // Overlay detections in console output
        for (const auto& det_obj : vdet_detections) {
            std::cout << "  -> Frame #" << frame_obj.si64_frame_id 
                      << " | Detected: Class " << det_obj.si32_class_id 
                      << " | Conf: " << std::fixed << std::setprecision(2) << det_obj.fp32_confidence
                      << " | Box: [" << det_obj.rect2f_bbox.x << ", " << det_obj.rect2f_bbox.y 
                      << ", " << det_obj.rect2f_bbox.width << ", " << det_obj.rect2f_bbox.height << "]\n";
        }
    }

    cudaStreamDestroy(cuda_stream);
    std::cout << "[Inference] Inference thread terminated successfully.\n";
}

int main(int argc, char* argv[]) {
    // Register signal handlers for cleanup
    trt::setupSignalHandlers();

    std::cout << "========================================================================\n";
    std::cout << "        TensorRT 10.X C++ Live Streaming Pipeline (Multi-Threaded)       \n";
    std::cout << "========================================================================\n";

    std::string str_engine_path = "model.engine";
    std::string str_rtsp_url = "synthetic"; // Default is synthetic video pipeline

    if (argc > 1) {
        str_engine_path = argv[1];
    }
    if (argc > 2) {
        str_rtsp_url = argv[2];
    }

    float fp32_conf_thresh = 0.3f;
    float fp32_iou_thresh = 0.5f;

    auto shptr_engine = trt::setupEngine(str_engine_path, fp32_conf_thresh, fp32_iou_thresh);
    if (!shptr_engine) {
        return 0; // Exit gracefully if setup failed (to verify compilation)
    }

    trt::SafeQueue<trt::Frame> queue_frames;

    std::cout << "[System] Starting capture and inference threads...\n";
    
    // Spawn threads
    std::thread thread_capture(captureThreadFunc, str_rtsp_url, std::ref(queue_frames));
    std::thread thread_inference(inferenceThreadFunc, shptr_engine, std::ref(queue_frames));

    // Wait for shutdown trigger
    while (trt::g_b_running) {
        std::this_thread::sleep_for(std::chrono::milliseconds(500));
    }

    // Join threads on shutdown
    if (thread_capture.joinable()) thread_capture.join();
    if (thread_inference.joinable()) thread_inference.join();

    std::cout << "[System] Streaming pipeline shut down safely and cleanly.\n";
    std::cout << "========================================================================\n";

    return 0;
}