cmvr-es/cmvr-es/common/media/ffmpeg/video_frame_encoder.cpp

// VideoFrameEncoder.cpp
#include "common/media/ffmpeg/video_frame_encoder.h"
#include "common/base/logging/logger.h"
#include <chrono>

namespace ffmpeg {

VideoFrameEncoder::VideoFrameEncoder() {
    converted_frame_ = av_frame_alloc();
    start_time_ = av_gettime(); // 获取初始时间
}

VideoFrameEncoder::~VideoFrameEncoder() {
    cleanup();
    if (converted_frame_) av_frame_free(&converted_frame_);
}

int64_t VideoFrameEncoder::get_current_timestamp() {
    int64_t current_time = av_gettime();
    int64_t elapsed = current_time - start_time_;

    // 将微秒转换为时间基单位
    return av_rescale_q(elapsed, {1, 1000000}, time_base_);
}

int VideoFrameEncoder::initialize(const Config& config) {
    config_ = config;

    // 查找编码器
    //const AVCodec* encoder = avcodec_find_encoder(AV_CODEC_ID_H264);
    const AVCodec* encoder = avcodec_find_encoder_by_name(config.codec.c_str());
    if (!encoder) {
        CMVR_LOG(ERROR) << "找不到H.264编码器";
        return -1;
    }

    // 分配编码器上下文
    encoder_ctx_ = avcodec_alloc_context3(encoder);
    if (!encoder_ctx_) {
        CMVR_LOG(ERROR) << "分配编码器上下文失败";
        return -1;
    }

    // 配置编码器参数
    encoder_ctx_->width = config_.width;
    encoder_ctx_->height = config_.height;
    encoder_ctx_->pix_fmt = config_.pixel_format;

    // 正确设置时间基
    encoder_ctx_->time_base = av_make_q(1, config_.fps);
    encoder_ctx_->framerate = av_make_q(config_.fps, 1);

    // 设置编码器的时间基为帧率倒数
    time_base_ = encoder_ctx_->time_base;

    encoder_ctx_->bit_rate = config_.bitrate;
    encoder_ctx_->gop_size = config_.gop_size;
    encoder_ctx_->max_b_frames = config_.max_b_frames;

    // 设置编码器选项
    if (encoder->id == AV_CODEC_ID_H264) {
        av_opt_set(encoder_ctx_->priv_data, "preset", config_.preset.c_str(), 0);
        av_opt_set(encoder_ctx_->priv_data, "tune", config_.tune.c_str(), 0);
        av_opt_set(encoder_ctx_->priv_data, "profile", config_.profile.c_str(), 0);
        av_opt_set(encoder_ctx_->priv_data, "x264-params", "log-level=error", 0);
    }

    // 打开编码器
    int ret = avcodec_open2(encoder_ctx_, encoder, nullptr);
    if (ret < 0) {
        char err_buf[1024];
        av_strerror(ret, err_buf, sizeof(err_buf));
        CMVR_LOG(ERROR) << "打开编码器失败: " << err_buf;
        return ret;
    }

    // 准备转换后的帧
    converted_frame_->width = config_.width;
    converted_frame_->height = config_.height;
    converted_frame_->format = config_.pixel_format;

    ret = av_frame_get_buffer(converted_frame_, 0);
    if (ret < 0) {
        CMVR_LOG(ERROR) << "分配帧缓冲区失败";
        return ret;
    }

    return 0;
}

int VideoFrameEncoder::init_sws_context(AVFrame* frame) {
    if (!frame) {
        CMVR_LOG(ERROR) << "输入帧为空";
        return -1;
    }

    // 如果格式相同，不需要转换
    if (frame->format == config_.pixel_format &&
        frame->width == config_.width &&
        frame->height == config_.height) {
        sws_ctx_ = nullptr;
        return 0;
    }

    // 创建转换上下文
    sws_ctx_ = sws_getContext(
        frame->width, frame->height, static_cast<AVPixelFormat>(frame->format),
        config_.width, config_.height, config_.pixel_format,
        SWS_BILINEAR, nullptr, nullptr, nullptr
    );

    if (!sws_ctx_) {
        CMVR_LOG(ERROR) << "创建SWS上下文失败";
        return -1;
    }

    return 0;
}

int VideoFrameEncoder::encode_frame(AVFrame* frame) {
    if (!encoder_ctx_ || !frame) {
        CMVR_LOG(ERROR) << "编码器未初始化或输入帧为空";
        return -1;
    }

    // 初始化格式转换（如果需要）
    if (!sws_ctx_) {
        int ret = init_sws_context(frame);
        if (ret < 0) {
            return ret;
        }
    }

    AVFrame* frame_to_encode = frame;

    // 如果需要格式转换
    if (sws_ctx_) {
        sws_scale(
            sws_ctx_,
            frame->data, frame->linesize, 0, frame->height,
            converted_frame_->data, converted_frame_->linesize
        );
        frame_to_encode = converted_frame_;
    }

    // 设置正确的PTS
    frame_to_encode->pts = frame_index_;

    // 发送帧到编码器
    int ret = avcodec_send_frame(encoder_ctx_, frame_to_encode);
    if (ret < 0) {
        char err_buf[1024];
        av_strerror(ret, err_buf, sizeof(err_buf));
        CMVR_LOG(ERROR) << "发送帧到编码器失败: " << err_buf;
        return ret;
    }

    // 接收编码后的数据包
    AVPacket* packet = av_packet_alloc();
    while (ret >= 0) {
        ret = avcodec_receive_packet(encoder_ctx_, packet);
        if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF) {
            break;
        } else if (ret < 0) {
            CMVR_LOG(ERROR) << "接收编码数据包失败";
            av_packet_free(&packet);
            return ret;
        }

        // 通过回调发送数据包，传递正确的PTS/DTS
        if (packet_callback_) {
            packet_callback_(packet, packet->pts, packet->dts);
        }

        av_packet_unref(packet);
    }

    av_packet_free(&packet);

    // 递增帧索引
    frame_index_++;

    return 0;
}

int VideoFrameEncoder::flush() {
    if (!encoder_ctx_) {
        return 0;
    }

    // 发送空帧刷新编码器
    int ret = avcodec_send_frame(encoder_ctx_, nullptr);
    if (ret < 0) {
        CMVR_LOG(ERROR) << "发送刷新帧失败";
        return ret;
    }

    // 接收所有剩余数据包
    AVPacket* packet = av_packet_alloc();
    while (true) {
        ret = avcodec_receive_packet(encoder_ctx_, packet);
        if (ret == AVERROR_EOF) {
            break;
        } else if (ret < 0) {
            CMVR_LOG(ERROR) << "接收刷新数据包失败";
            av_packet_free(&packet);
            return ret;
        }

        // 通过回调发送数据包
        if (packet_callback_) {
            packet_callback_(packet, packet->pts, packet->dts);
        }

        av_packet_unref(packet);
    }

    av_packet_free(&packet);
    return 0;
}

void VideoFrameEncoder::send_packet(AVPacket* packet) {
    if (packet_callback_) {
        packet_callback_(packet, packet->pts, packet->dts);
    }
}

void VideoFrameEncoder::cleanup() {
    if (sws_ctx_) {
        sws_freeContext(sws_ctx_);
        sws_ctx_ = nullptr;
    }

    if (encoder_ctx_) {
        avcodec_close(encoder_ctx_);
        avcodec_free_context(&encoder_ctx_);
    }
}

} // namespace ffmpeg