scene-actor-extraction/src/nodes/frame_source_node.hpp

#pragma once
#include "types.hpp"
#include "config.hpp"
#include "ffmpeg_decoder.hpp"

#include <opencv2/imgproc.hpp>
#include <chrono>
#include <iostream>
#include <memory>
#include <stdexcept>
#include <string>
#include <thread>

// ── FrameSourceFunc ───────────────────────────────────────────────────────────
// KPN source node: reads a movie file and emits one Frame per sample interval.
//
// Decode backend: FFmpeg hwaccel (CUDA/VAAPI, runtime-detected) when
// available, CPU otherwise.
//
// Sampling strategy: seek to the next target timestamp rather than decoding
// every frame, which is fast even for 1-FPS sampling of a 2-hour film.
//
// EOF handling: when the movie ends, emits a Frame with eof=true, then sleeps
// 500 ms between subsequent calls until the KPN network stops the thread.

struct FrameSourceFunc {
    static constexpr std::string_view label() { return "frame_source"; }

    explicit FrameSourceFunc(const Config& cfg)
        : decoder_(std::make_unique<FFmpegDecoder>(cfg.movie_path))
    {
        sample_interval_sec_ = 1.0 / cfg.sample_fps;
        next_pos_sec_        = cfg.start_sec;
        end_sec_             = cfg.end_sec;
        cut_threshold_       = cfg.cut_threshold;
        max_decode_fps_      = cfg.max_decode_fps;

        double total_s = decoder_->duration_sec();
        double span_s  = (end_sec_ > 0 ? std::min(end_sec_, total_s) : total_s)
                         - cfg.start_sec;
        int n_frames = static_cast<int>(span_s * cfg.sample_fps);
        std::cerr << "[frame_source] decoder=" << decoder_->codec_name()
                  << " (" << decoder_->hw_backend() << ")"
                  << "  video_fps=" << decoder_->fps()
                  << "  start=" << cfg.start_sec << "s"
                  << (end_sec_ > 0 ? "  end=" + std::to_string(end_sec_) + "s" : "")
                  << "  sample_fps=" << cfg.sample_fps
                  << "  frames_to_emit=" << n_frames << "\n";
    }

    Frame operator()() {
        if (hit_eof_) {
            std::this_thread::sleep_for(std::chrono::milliseconds(500));
            return Frame{{}, 0.0, -1, /*eof=*/true};
        }

        // Wall-clock rate cap. KPN source nodes resubmit immediately on push
        // overflow, with no backpressure; without this cap we'd decode-and-drop
        // in a tight loop whenever downstream stalls. The cap also protects
        // ORT-only deployments where the pipeline can't keep up at decode speed.
        if (max_decode_fps_ > 0.f) {
            const auto now = std::chrono::steady_clock::now();
            if (!rate_started_) {
                rate_started_ = true;
                next_decode_at_ = now;
            }
            if (now < next_decode_at_)
                std::this_thread::sleep_until(next_decode_at_);
            const auto period = std::chrono::nanoseconds(
                static_cast<int64_t>(1e9f / max_decode_fps_));
            // Anchor the next slot off the slot we just consumed, not off
            // wall-clock now() — keeps the average rate stable. If we fell
            // behind by more than one period, snap forward to avoid building
            // up an unbounded sleep debt.
            next_decode_at_ += period;
            if (next_decode_at_ < now)
                next_decode_at_ = now + period;
        }

        auto t0 = std::chrono::steady_clock::now();
        cv::Mat img = decoder_->read_at(next_pos_sec_);
        auto t1 = std::chrono::steady_clock::now();
        double decode_ms = std::chrono::duration<double, std::milli>(t1 - t0).count();
        decode_ms_acc_  += decode_ms;
        ++decode_count_;
        if (decode_count_ % 10 == 0) {
            double avg_ms  = decode_ms_acc_ / 10.0;
            double avg_fps = avg_ms > 0.0 ? 1000.0 / avg_ms : 0.0;
            std::cerr << "[frame_source] decode  avg=" << avg_ms << "ms"
                      << "  fps=" << avg_fps << "\n";
            decode_ms_acc_ = 0.0;
        }

        if (img.empty()) {
            hit_eof_ = true;
            std::cerr << "[frame_source] EOF at t=" << next_pos_sec_ << "s\n";
            return Frame{{}, next_pos_sec_, frame_idx_++, /*eof=*/true};
        }

        // Cut detection: compare grayscale histogram to previous frame
        bool is_cut = false;
        cv::Mat gray;
        cv::cvtColor(img, gray, cv::COLOR_BGR2GRAY);
        cv::Mat hist;
        const int   bins    = 64;
        const float range[] = {0.f, 256.f};
        const float* ranges = range;
        cv::calcHist(&gray, 1, nullptr, cv::Mat(), hist, 1, &bins, &ranges);
        cv::normalize(hist, hist, 1.0, 0.0, cv::NORM_L1);

        if (prev_hist_valid_) {
            double corr = cv::compareHist(prev_hist_, hist, cv::HISTCMP_CORREL);
            is_cut = (corr < cut_threshold_);
            if (is_cut)
                std::cerr << "[frame_source] cut at t=" << next_pos_sec_
                          << "s  hist_corr=" << corr << "\n";
        }
        prev_hist_       = hist;
        prev_hist_valid_ = true;

        Frame f{img, next_pos_sec_, frame_idx_++, /*eof=*/false, is_cut};
        next_pos_sec_ += sample_interval_sec_;
        if (end_sec_ > 0 && next_pos_sec_ > end_sec_) {
            hit_eof_ = true;
            std::cerr << "[frame_source] reached end_sec=" << end_sec_ << "s\n";
        }
        return f;
    }

private:
    std::unique_ptr<FFmpegDecoder> decoder_;
    double           sample_interval_sec_{1.0};
    double           next_pos_sec_{0.0};
    double           end_sec_{-1.0};
    float            cut_threshold_{0.70f};
    float            max_decode_fps_{0.f};
    std::chrono::steady_clock::time_point next_decode_at_{};
    bool             rate_started_{false};
    int64_t          frame_idx_{0};
    bool             hit_eof_{false};
    cv::Mat          prev_hist_;
    bool             prev_hist_valid_{false};
    double           decode_ms_acc_{0.0};
    int              decode_count_{0};
};