scene-actor-extraction/src/ffmpeg_decoder.hpp

#pragma once

extern "C" {
#include <libavformat/avformat.h>
#include <libavcodec/avcodec.h>
#include <libavutil/avutil.h>
#include <libavutil/hwcontext.h>
#include <libavutil/imgutils.h>
#include <libavutil/opt.h>
#include <libavutil/pixdesc.h>
#include <libswscale/swscale.h>
}

#include <opencv2/core.hpp>

#include <algorithm>
#include <iostream>
#include <stdexcept>
#include <string>
#include <vector>

// ── FFmpegDecoder ─────────────────────────────────────────────────────────────
// Seek-and-decode video reader backed by FFmpeg.
//
// Hardware decode is selected at runtime via the generic hwaccel API
// (av_hwdevice_ctx_create): the decoder probes the device types the local
// build supports, in priority order CUDA (NVIDIA) → VAAPI (AMD/Intel) →
// software.  This works across GPU vendors without vendor-specific decoder
// names.
//
// Hardware decoders output frames in GPU memory (e.g. AV_PIX_FMT_CUDA,
// AV_PIX_FMT_VAAPI); av_hwframe_transfer_data copies them to a system-memory
// frame (typically NV12), then swscale converts NV12/YUV → BGR24 for the rest
// of the pipeline.
//
// Non-copyable; wrap in unique_ptr if you need to move it.

struct FFmpegDecoder {
    // out_scale in (0,1] downscales decoded frames (applied in the sws_scale
    // colour conversion, so it's nearly free). 1.0 = native resolution.
    explicit FFmpegDecoder(const std::string& path, bool use_hw = true,
                           float out_scale = 1.0f)
        : out_scale_(out_scale > 0.f && out_scale <= 1.f ? out_scale : 1.0f)
    {
        // Quiet FFmpeg's own logging (e.g. the harmless "Could not dynamically
        // load CUDA" emitted while probing hwaccels before VAAPI succeeds).
        av_log_set_level(AV_LOG_ERROR);
        if (avformat_open_input(&fmt_ctx_, path.c_str(), nullptr, nullptr) < 0)
            throw std::runtime_error("[FFmpegDecoder] cannot open: " + path);
        if (avformat_find_stream_info(fmt_ctx_, nullptr) < 0)
            throw std::runtime_error("[FFmpegDecoder] stream info failed");

        stream_idx_ = av_find_best_stream(
            fmt_ctx_, AVMEDIA_TYPE_VIDEO, -1, -1, nullptr, 0);
        if (stream_idx_ < 0)
            throw std::runtime_error("[FFmpegDecoder] no video stream");

        AVStream*  stream = fmt_ctx_->streams[stream_idx_];
        AVCodecID  cid    = stream->codecpar->codec_id;

        // Try hardware backends in priority order; fall back to software.
        if (use_hw)
            try_open_hw(stream, cid);

        if (!hw_active_) {
            const AVCodec* swc = avcodec_find_decoder(cid);
            if (!swc) throw std::runtime_error("[FFmpegDecoder] no software decoder");
            codec_ctx_ = avcodec_alloc_context3(swc);
            avcodec_parameters_to_context(codec_ctx_, stream->codecpar);
            codec_ctx_->thread_count = 0;  // auto: use all cores
            if (avcodec_open2(codec_ctx_, swc, nullptr) < 0)
                throw std::runtime_error("[FFmpegDecoder] cannot open codec");
            std::cerr << "[FFmpegDecoder] " << path
                      << "  codec=" << swc->name << " (CPU)\n";
        }

        frame_    = av_frame_alloc();
        tmp_frame_= av_frame_alloc();
        pkt_      = av_packet_alloc();

        // Decode forward (no seek) when target is within this many pts units.
        // 5 seconds covers typical H.264/H.265 GOP sizes.
        max_forward_pts_ = to_stream_pts(5.0);

        double total = duration_sec();
        double vfps  = fps();
        std::cerr << "[FFmpegDecoder] duration=" << total
                  << "s  fps=" << vfps << "\n";
    }

    ~FFmpegDecoder() {
        if (sws_ctx_) sws_freeContext(sws_ctx_);
        av_frame_free(&frame_);
        av_frame_free(&tmp_frame_);
        av_packet_free(&pkt_);
        avcodec_free_context(&codec_ctx_);
        if (hw_device_ctx_) av_buffer_unref(&hw_device_ctx_);
        avformat_close_input(&fmt_ctx_);
    }

    FFmpegDecoder(const FFmpegDecoder&)            = delete;
    FFmpegDecoder& operator=(const FFmpegDecoder&) = delete;

    double duration_sec() const {
        if (!fmt_ctx_ || fmt_ctx_->duration == AV_NOPTS_VALUE) return 0.0;
        return static_cast<double>(fmt_ctx_->duration) / AV_TIME_BASE;
    }

    double fps() const {
        AVStream* s = fmt_ctx_->streams[stream_idx_];
        if (s->avg_frame_rate.den == 0) return 25.0;
        return av_q2d(s->avg_frame_rate);
    }

    bool        hw_active()   const { return hw_active_; }
    const char* codec_name()  const { return codec_ctx_ ? codec_ctx_->codec->name : "unknown"; }
    // Human-readable backend: "CUDA", "VAAPI", … or "CPU".
    const char* hw_backend()  const {
        return hw_active_ ? av_hwdevice_get_type_name(hw_type_) : "CPU";
    }

    // Decode the frame at target_sec and return it as BGR cv::Mat.
    // Returns an empty Mat at EOF.
    //
    // Smart seek: if the target is within max_forward_sec_ ahead of the last
    // decoded position, decode forward (no seek, no flush).  This is dramatically
    // faster for sequential sampling because avcodec_flush_buffers + re-init on
    // every call is the main bottleneck — especially with GPU decode.
    cv::Mat read_at(double target_sec) {
        AVStream* stream  = fmt_ctx_->streams[stream_idx_];
        int64_t   tgt_pts = to_stream_pts(target_sec);

        // Decide: seek or decode forward?
        bool need_seek = (last_pts_ == AV_NOPTS_VALUE) ||
                         (tgt_pts < last_pts_)          ||
                         (tgt_pts - last_pts_ > max_forward_pts_);

        if (need_seek) {
            if (av_seek_frame(fmt_ctx_, stream_idx_, tgt_pts, AVSEEK_FLAG_BACKWARD) < 0)
                av_seek_frame(fmt_ctx_, -1,
                              static_cast<int64_t>(target_sec * AV_TIME_BASE),
                              AVSEEK_FLAG_BACKWARD);
            avcodec_flush_buffers(codec_ctx_);
            last_pts_ = AV_NOPTS_VALUE;
        }

        // Decode forward until we reach or pass target_pts.
        // Convert to BGR and unref the AVFrame immediately so GPU surfaces
        // are returned to the pool — holding them causes surface exhaustion
        // at higher sample rates.
        cv::Mat out;
        while (out.empty()) {
            int ret = av_read_frame(fmt_ctx_, pkt_);
            if (ret == AVERROR_EOF || ret < 0) break;

            if (pkt_->stream_index != stream_idx_) {
                av_packet_unref(pkt_);
                continue;
            }

            avcodec_send_packet(codec_ctx_, pkt_);
            av_packet_unref(pkt_);

            while (avcodec_receive_frame(codec_ctx_, frame_) == 0) {
                int64_t pts = frame_->best_effort_timestamp;
                if (pts == AV_NOPTS_VALUE) pts = frame_->pts;
                last_pts_ = pts;
                if (pts >= tgt_pts)
                    out = to_bgr(frame_);
                av_frame_unref(frame_);   // release GPU surface immediately
                if (!out.empty()) break;
            }
        }
        return out;
    }

private:
    AVFormatContext* fmt_ctx_   = nullptr;
    AVCodecContext*  codec_ctx_ = nullptr;
    AVBufferRef*     hw_device_ctx_ = nullptr;
    AVFrame*         frame_     = nullptr;
    AVFrame*         tmp_frame_ = nullptr;
    AVPacket*        pkt_       = nullptr;
    SwsContext*      sws_ctx_   = nullptr;
    int              stream_idx_      = -1;
    bool             hw_active_       = false;
    AVHWDeviceType   hw_type_         = AV_HWDEVICE_TYPE_NONE;
    AVPixelFormat    hw_pix_fmt_      = AV_PIX_FMT_NONE;
    int64_t          last_pts_        = AV_NOPTS_VALUE;
    int64_t          max_forward_pts_ = AV_NOPTS_VALUE;  // set after codec opens
    float            out_scale_       = 1.0f;            // decoded-frame downscale (0,1]

    int64_t to_stream_pts(double sec) const {
        AVStream* s = fmt_ctx_->streams[stream_idx_];
        return av_rescale_q(static_cast<int64_t>(sec * AV_TIME_BASE),
                            AV_TIME_BASE_Q, s->time_base);
    }

    // get_format callback: tell the decoder we want the hardware surface
    // format negotiated for this device.  The chosen format is stashed on the
    // codec context's opaque pointer so this static callback can read it.
    static AVPixelFormat get_hw_format(AVCodecContext* ctx,
                                       const AVPixelFormat* fmts) {
        auto want = *static_cast<const AVPixelFormat*>(ctx->opaque);
        for (const AVPixelFormat* p = fmts; *p != AV_PIX_FMT_NONE; ++p)
            if (*p == want) return *p;
        std::cerr << "[FFmpegDecoder] hw surface format unavailable, "
                     "decoder will fall back to software output\n";
        return fmts[0];
    }

    // Probe hardware device types in priority order and open the first that
    // works for this codec.  Detection is fully at runtime: only device types
    // compiled into the local FFmpeg are returned by av_hwdevice_iterate_types,
    // and av_hwdevice_ctx_create only succeeds if a usable device is present.
    void try_open_hw(AVStream* stream, AVCodecID cid) {
        static const AVHWDeviceType kPriority[] = {
            AV_HWDEVICE_TYPE_CUDA,    // NVIDIA
            AV_HWDEVICE_TYPE_VAAPI,   // AMD / Intel (Linux)
        };

        const std::vector<AVHWDeviceType> available = available_hw_types();

        const AVCodec* dec = avcodec_find_decoder(cid);
        if (!dec) return;

        for (AVHWDeviceType type : kPriority) {
            bool present = false;
            for (AVHWDeviceType a : available) present |= (a == type);
            if (!present) continue;

            // Find the hw pixel format this decoder advertises for this device.
            AVPixelFormat pix = hw_pix_fmt_for(dec, type);
            if (pix == AV_PIX_FMT_NONE) continue;

            AVBufferRef* dev_ctx = nullptr;
            if (av_hwdevice_ctx_create(&dev_ctx, type, nullptr, nullptr, 0) < 0)
                continue;  // no usable device of this type on the machine

            codec_ctx_ = avcodec_alloc_context3(dec);
            avcodec_parameters_to_context(codec_ctx_, stream->codecpar);
            codec_ctx_->thread_count   = 1;
            codec_ctx_->hw_device_ctx  = av_buffer_ref(dev_ctx);
            hw_pix_fmt_                 = pix;
            codec_ctx_->opaque         = &hw_pix_fmt_;
            codec_ctx_->get_format     = get_hw_format;

            if (avcodec_open2(codec_ctx_, dec, nullptr) >= 0) {
                hw_active_     = true;
                hw_type_       = type;
                hw_device_ctx_ = dev_ctx;
                std::cerr << "[FFmpegDecoder] codec=" << dec->name
                          << "  hwaccel=" << av_hwdevice_get_type_name(type)
                          << "\n";
                return;
            }

            // This backend failed to open; tear down and try the next.
            avcodec_free_context(&codec_ctx_);
            av_buffer_unref(&dev_ctx);
            hw_pix_fmt_ = AV_PIX_FMT_NONE;
            std::cerr << "[FFmpegDecoder] "
                      << av_hwdevice_get_type_name(type)
                      << " init failed, trying next backend\n";
        }
    }

    static std::vector<AVHWDeviceType> available_hw_types() {
        std::vector<AVHWDeviceType> types;
        AVHWDeviceType t = AV_HWDEVICE_TYPE_NONE;
        while ((t = av_hwdevice_iterate_types(t)) != AV_HWDEVICE_TYPE_NONE)
            types.push_back(t);
        return types;
    }

    // Look up the hw-surface pixel format the decoder exposes for a device type
    // (e.g. AV_PIX_FMT_CUDA for CUDA, AV_PIX_FMT_VAAPI for VAAPI).
    static AVPixelFormat hw_pix_fmt_for(const AVCodec* dec, AVHWDeviceType type) {
        for (int i = 0;; ++i) {
            const AVCodecHWConfig* cfg = avcodec_get_hw_config(dec, i);
            if (!cfg) break;
            if ((cfg->methods & AV_CODEC_HW_CONFIG_METHOD_HW_DEVICE_CTX) &&
                cfg->device_type == type)
                return cfg->pix_fmt;
        }
        return AV_PIX_FMT_NONE;
    }

    cv::Mat to_bgr(AVFrame* src) {
        // Hardware decoders hand back GPU surfaces; transfer to system memory.
        AVFrame* sw = src;
        if (src->format == hw_pix_fmt_ && hw_pix_fmt_ != AV_PIX_FMT_NONE) {
            av_frame_unref(tmp_frame_);
            if (av_hwframe_transfer_data(tmp_frame_, src, 0) < 0) return {};
            av_frame_copy_props(tmp_frame_, src);
            sw = tmp_frame_;
        }

        const int w = sw->width;
        const int h = sw->height;

        // Optional downscale, folded into the colour conversion (near-free).
        // Round to even dimensions for swscale/codec friendliness.
        int out_w = w, out_h = h;
        if (out_scale_ < 1.0f) {
            out_w = std::max(2, (static_cast<int>(w * out_scale_) / 2) * 2);
            out_h = std::max(2, (static_cast<int>(h * out_scale_) / 2) * 2);
        }

        sws_ctx_ = sws_getCachedContext(sws_ctx_,
            w, h, static_cast<AVPixelFormat>(sw->format),
            out_w, out_h, AV_PIX_FMT_BGR24,
            SWS_BILINEAR, nullptr, nullptr, nullptr);
        if (!sws_ctx_) return {};

        cv::Mat out(out_h, out_w, CV_8UC3);
        uint8_t* dst_data[1]     = { out.data };
        int      dst_linesize[1] = { static_cast<int>(out.step) };
        sws_scale(sws_ctx_,
                  sw->data, sw->linesize, 0, h,
                  dst_data, dst_linesize);
        return out;
    }
};