Skip to content

Instantly share code, notes, and snippets.

@prabindh

prabindh/encode-video.cpp

Created July 28, 2020 03:41
Show Gist options
  • Save prabindh/1587d4b6cde284d30b082df2e4fb7df1 to your computer and use it in GitHub Desktop.
Save prabindh/1587d4b6cde284d30b082df2e4fb7df1 to your computer and use it in GitHub Desktop.
Download ZIP
calling ffmpeg encoder
Raw
encode-video.cpp
#define OUTPUT_W 1920
#define OUTPUT_H 1080
#define OUTPUT_FPS 25
int main()
{
int ret = 0;
bool bUseHW = false;
bool bNv = false;
Encoder* pEnc = new Encoder(bUseHW, bNv, OUTPUT_W, OUTPUT_H, OUTPUT_FPS);
format = AV_PIX_FMT_YUV420P;
size = avpicture_get_size(format, OUTPUT_W, OUTPUT_H);
picture = av_frame_alloc();
picture_buffer = (uint8_t*)(av_malloc(size));
ret = avpicture_fill((AVPicture *)picture, picture_buffer, format,
OUTPUT_W, OUTPUT_H);
picture = av_frame_alloc();
if (!picture) {
fprintf(stderr, "Could not allocate video frame\n");
exit(1);
}
picture->format = format;
picture->width = OUTPUT_W;
picture->height = OUTPUT_H;
ret = av_frame_get_buffer(picture, 0);
if (ret < 0) {
fprintf(stderr, "Could not allocate the video frame data\n");
exit(1);
}
for (int i = 0; i < 600; i++)
{
prepare_dummy_420P(OUTPUT_W, OUTPUT_H, i, picture);
pEnc->addFrame(picture);
}
pEnc->flush();
return ret;
}
@prabindh
Copy link
Author

//
#include "encoder.h"

extern "C" {

static enum AVPixelFormat get_vaapi_format(AVCodecContext*, const enum AVPixelFormat *pix_fmts)
{
    const enum AVPixelFormat *p;
    for (p = pix_fmts; *p != AV_PIX_FMT_NONE; p++) {
        if (*p == AV_PIX_FMT_VAAPI)
            return *p;
    }
    fprintf(stderr, "Failed to get HW surface format.\n");
    return AV_PIX_FMT_NONE;
}
static enum AVPixelFormat get_nv12_format(AVCodecContext*, const enum AVPixelFormat *pix_fmts)
{
    const enum AVPixelFormat *p;
    for (p = pix_fmts; *p != AV_PIX_FMT_NONE; p++) {
        if (*p == AV_PIX_FMT_NV12)
            return *p;
    }
    fprintf(stderr, "Failed to get HW surface format.\n");
    return AV_PIX_FMT_NONE;
}
static enum AVPixelFormat get_YUV420P_format(AVCodecContext*, const enum AVPixelFormat *pix_fmts)
{
    const enum AVPixelFormat *p;
    for (p = pix_fmts; *p != AV_PIX_FMT_NONE; p++) {
        if (*p == AV_PIX_FMT_YUV420P)
            return *p;
    }
    fprintf(stderr, "Failed to get 420P surface format.\n");
    return AV_PIX_FMT_NONE;
}

}

Encoder::Encoder(const bool hwAccel, bool bNv, int w, int h, int fps)
: m_hardwareAcceleration(hwAccel)
{
m_width = w;
m_height = h;
m_fps = fps;
setup(bNv);
}
void Encoder::addFrame(AVFrame* frame)
{
AVFrame* frameToEncode = frame;
if (m_hardwareAcceleration) {
filterFrame(frame, m_hwFrame);
assert(m_hwFrame->format == AV_PIX_FMT_VAAPI);
frameToEncode = m_hwFrame;
}

frameToEncode->pts = m_frameId++;
encodeFrame(frameToEncode);

}
void Encoder::flush()
{
encodeFrame(nullptr);
av_write_trailer(m_muxer);
}

void Encoder::setup(bool bNv)
{
AVOutputFormat * outFmt = av_guess_format("mp4", NULL, NULL);
assert(avformat_alloc_output_context2(&m_muxer, outFmt, nullptr, nullptr) == 0);

assert(m_muxer != nullptr);

setupEncoder(bNv);

m_avStream = avformat_new_stream(m_muxer, nullptr);
assert(m_avStream != nullptr);
m_avStream->id = m_muxer->nb_streams - 1;
m_avStream->time_base = m_encoder->time_base;

// Some formats want stream headers to be separate.
if (m_muxer->oformat->flags & AVFMT_GLOBALHEADER)
    m_encoder->flags |= AV_CODEC_FLAG_GLOBAL_HEADER;

assert(avcodec_parameters_from_context(m_avStream->codecpar, m_encoder) == 0);
assert(avio_open(&m_muxer->pb, m_hardwareAcceleration ? "hardware.mp4" : "software.mp4", AVIO_FLAG_WRITE) == 0);
assert(avformat_write_header(m_muxer, nullptr) == 0);

}

void Encoder::setupEncoder(bool bNv)
{
const char* encoderName = m_hardwareAcceleration ? "h264_vaapi" : "libx264";
AVCodec* videoCodec = nullptr;
AVRational timeBase, frameRate;

if (m_hardwareAcceleration && bNv)
{
    encoderName = "h264_nvenc";
}
videoCodec = avcodec_find_encoder_by_name(encoderName);

m_encoder = avcodec_alloc_context3(videoCodec);

for (int i = 0;; i++) {
    const AVCodecHWConfig *config = avcodec_get_hw_config(videoCodec, i);
    if (!config) {
        fprintf(stderr, "Encoder %s does not support device type %s.\n",
            videoCodec->name, av_hwdevice_get_type_name(AV_HWDEVICE_TYPE_CUDA));
        return;
    }
    if (config->methods & AV_CODEC_HW_CONFIG_METHOD_HW_DEVICE_CTX &&
        config->device_type == AV_HWDEVICE_TYPE_CUDA) {
        m_encoder->pix_fmt = config->pix_fmt;
        break;
    }
}


m_encoder->bit_rate = m_width * m_height * m_fps * 2;
m_encoder->width = m_width;
m_encoder->height = m_height;

m_encoder->gop_size = m_fps;
m_encoder->max_b_frames = 1;

timeBase = {1, m_fps};
frameRate = { m_fps, 1 };
m_encoder->time_base = timeBase;
m_encoder->framerate = frameRate;

m_encoder->gop_size = m_fps;  // have at least 1 I-frame per second
m_encoder->max_b_frames = 1;
m_encoder->pix_fmt = AV_PIX_FMT_YUV420P;

if (m_hardwareAcceleration) {
    if (bNv)
    {
        m_encoder->pix_fmt = AV_PIX_FMT_YUV420P; // AV_PIX_FMT_NV12;
        m_encoder->get_format = get_YUV420P_format; // get_nv12_format;
    }
    else
    {
        m_encoder->pix_fmt = AV_PIX_FMT_VAAPI;
        m_encoder->get_format = get_vaapi_format;
    }

    if (bNv)
    {
        assert(av_hwdevice_ctx_create(&m_device, AV_HWDEVICE_TYPE_CUDA, "", nullptr, 0) == 0);
        const AVHWDeviceContext* deviceCtx = (AVHWDeviceContext*)m_device->data;
        assert(deviceCtx->type == AV_HWDEVICE_TYPE_CUDA);
    }
    else
    {
        assert(av_hwdevice_ctx_create(&m_device, AV_HWDEVICE_TYPE_VAAPI, "/dev/dri/renderD128", nullptr, 0) == 0);
        const AVHWDeviceContext* deviceCtx = (AVHWDeviceContext*)m_device->data;
        assert(deviceCtx->type == AV_HWDEVICE_TYPE_VAAPI);
    }

    initFilters(bNv);

#if 0
m_encoder->hw_device_ctx = nullptr;
AVBufferRef* bufferRef = av_buffersink_get_hw_frames_ctx(m_bufferSink);
m_encoder->hw_frames_ctx = av_buffer_ref(bufferRef);
#else
{
AVBufferRef *hw_frames_ref;
AVHWFramesContext *frames_ctx = NULL;

        assert((hw_frames_ref = av_hwframe_ctx_alloc(m_device)) != nullptr);
        frames_ctx = (AVHWFramesContext *)(hw_frames_ref->data);
        frames_ctx->format = AV_PIX_FMT_CUDA;
        frames_ctx->sw_format = AV_PIX_FMT_NV12;
        frames_ctx->width = m_width;
        frames_ctx->height = m_height;
        frames_ctx->initial_pool_size = 20;
        assert(av_hwframe_ctx_init(hw_frames_ref) == 0);
        m_encoder->hw_frames_ctx = av_buffer_ref(hw_frames_ref);
        assert(m_encoder->hw_frames_ctx != nullptr);

        // Fix warning, cosmetical only: No quality level set; using default (20).
        m_encoder->global_quality = 20;

        m_encoder->hw_device_ctx = av_hwframe_ctx_alloc(m_device);
        //m_encoder->hw_frames_ctx = av_buffer_ref(m_device);           // Fix: Not required, done by av_hwframe_ctx_alloc
        m_hwFrame = av_frame_alloc();
        av_hwframe_get_buffer(m_encoder->hw_frames_ctx, m_hwFrame, 0); // Fix: Must pass hw_frames_ctx, not m_encoder->hw_device_ctx
        assert(m_hwFrame != nullptr);
    }

#endif
}

assert(avcodec_open2(m_encoder, videoCodec, nullptr) == 0);

if (m_hardwareAcceleration) {
    m_encoder->hw_device_ctx = av_hwframe_ctx_alloc(m_device);
    m_hwFrame = av_frame_alloc();
    av_hwframe_get_buffer(m_encoder->hw_device_ctx, m_hwFrame, 0);
}

m_muxer->video_codec_id = videoCodec->id;
m_muxer->video_codec = videoCodec;

}
void Encoder::initFilters(bool bNv)
{
AVFilterInOut* inputs = nullptr;
AVFilterInOut* outputs = nullptr;
m_filterGraph = avfilter_graph_alloc();

if (bNv)
{
    // Move to more optimal input format later
    //format= has to match in all cases ?
    // format=pix_fmts=yuv420p,hwupload_cuda,scale_cuda,hwdownload,format=pix_fmts=nv12
    assert(avfilter_graph_parse2(m_filterGraph, "format=pix_fmts=yuv420p,hwupload_cuda,scale_cuda,hwdownload,format=pix_fmts=yuv420p", &inputs, &outputs) == 0);
}
else
{
    assert(avfilter_graph_parse2(m_filterGraph, "format=nv12,hwupload", &inputs, &outputs) == 0);
}

for (unsigned i = 0; i < m_filterGraph->nb_filters; i++) {
    m_filterGraph->filters[i]->hw_device_ctx = av_buffer_ref(m_device);
    assert(m_filterGraph->filters[i]->hw_device_ctx != nullptr);
}

initInputFilters(inputs, bNv);
initOutputFilters(outputs, bNv);

assert(avfilter_graph_config(m_filterGraph, nullptr) == 0);

}
void Encoder::initInputFilters(AVFilterInOut* inputs, bool bNv)
{
assert(inputs != nullptr);
assert(inputs->next == nullptr);

char args[512];
snprintf(args, sizeof(args),
    "video_size=%dx%d:pix_fmt=%d:time_base=%d/%d:pixel_aspect=%d/%d",
    m_width, m_height, AV_PIX_FMT_YUV420P,
    1, m_fps,
    1, 1);

assert(avfilter_graph_create_filter(&m_bufferSrc, avfilter_get_by_name("buffer"), "in",
    args, nullptr, m_filterGraph) == 0);
assert(avfilter_link(m_bufferSrc, 0, inputs->filter_ctx, inputs->pad_idx) == 0);

}
void Encoder::initOutputFilters(AVFilterInOut* outputs, bool bNv)
{
assert(outputs != nullptr);
assert(outputs->next == nullptr);

assert(avfilter_graph_create_filter(&m_bufferSink, avfilter_get_by_name("buffersink"), "out",
    nullptr, nullptr, m_filterGraph) == 0);
if (bNv)
{
    assert(avfilter_graph_create_filter(&m_formatFilter, avfilter_get_by_name("format"), "format",
        "pix_fmts=yuv420p", nullptr, m_filterGraph) == 0);
}
else
{
    assert(avfilter_graph_create_filter(&m_formatFilter, avfilter_get_by_name("format"), "format",
        "vaapi_vld", nullptr, m_filterGraph) == 0);
}
assert(avfilter_link(outputs->filter_ctx, outputs->pad_idx, m_formatFilter, 0) == 0);
assert(avfilter_link(m_formatFilter, 0, m_bufferSink, 0) == 0);

}
void Encoder::filterFrame(AVFrame* inFrame, AVFrame* outFrame)
{
assert(av_buffersrc_add_frame_flags(m_bufferSrc, inFrame, AV_BUFFERSRC_FLAG_KEEP_REF) == 0);
assert(av_buffersink_get_frame(m_bufferSink, outFrame) == 0);
}
void Encoder::encodeFrame(AVFrame* frame)
{
assert(avcodec_send_frame(m_encoder, frame) == 0);

AVPacket packet;
av_init_packet(&packet);
int ret = 0;
while (ret >= 0) {
    ret = avcodec_receive_packet(m_encoder, &packet);
    if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF) {
        return;  // nothing to write
    }
    assert(ret >= 0);

    av_packet_rescale_ts(&packet, m_encoder->time_base, m_avStream->time_base);
    packet.stream_index = m_avStream->index;
    av_interleaved_write_frame(m_muxer, &packet);
    av_packet_unref(&packet);
}

}

// Results in [h264_nvenc @ 00007FFF05513D40] Mismatching AVCodecContext.pix_fmt and AVHWFramesContext.format

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment