ffmpeg encode 2

ffmpeg-encode2.cpp

// https://stackoverflow.com/questions/62279768/avcodec-encode-video2-the-memory-usage-is-very-high

extern "C"
{
#ifndef __STDC_CONSTANT_MACROS
#define __STDC_CONSTANT_MACROS
#endif
#include <libavcodec/avcodec.h>
#include "libavformat/avformat.h"
#include "libswscale/swscale.h"
#include "libavutil/avutil.h"
#include "libavutil/imgutils.h"
#include "libavutil/opt.h"
#include <libavutil/rational.h>
//#include "libavutil/avassert.h"


//#include "libavutil/attributes.h"
//#include "libavutil/avassert.h"
//#include "libavutil/frame.h"
//#include "libavutil/imgutils.h"
//#include "internal1.h"
//#include "libavutil/samplefmt.h"
}
#include <vld.h>
//#include "avcodec.h"
//#include "frame_thread_encoder.h"
//#include "internal.h"

// Globals
AVCodec* m_pCodec = NULL;
AVStream *m_pStream = NULL;
AVOutputFormat* m_pFormat = NULL;
AVFormatContext* m_pFormatContext = NULL;
AVCodecContext* m_pCodecContext = NULL;
AVFrame* m_pFrame = NULL;
int m_frameIndex;

// Output format
AVPixelFormat m_pixType = AV_PIX_FMT_NV12;
// Use for mpeg4
//AVPixelFormat m_pixType = AV_PIX_FMT_YUV420P;

// Output frame rate
int m_frameRate = 30;
// Output image dimensions
int m_imageWidth = 256;
int m_imageHeight = 1537;
// Number of frames to export
int m_frameCount = 20000;
// Output file name
const char* m_fileName = "test.h264";
// Output file type
const char* m_fileType = "H264";
// Codec name used to encode
const char* m_encoderName = "h264_qsv";
// use for mpeg4
//const char* m_encoderName = "mpeg4";
// Target bit rate
int m_targetBitRate = 400000;

void addVideoStream()
{
    m_pStream = avformat_new_stream(m_pFormatContext, m_pCodec);
    m_pStream->id = m_pFormatContext->nb_streams - 1;
    m_pStream->time_base = m_pCodecContext->time_base;
    m_pStream->codec->pix_fmt = m_pixType;
    m_pStream->codec->flags = m_pCodecContext->flags;
    m_pStream->codec->width = m_pCodecContext->width;
    m_pStream->codec->height = m_pCodecContext->height;
    m_pStream->codec->time_base = m_pCodecContext->time_base;
    m_pStream->codec->bit_rate = m_pCodecContext->bit_rate;
}

AVFrame* allocatePicture(enum AVPixelFormat pix_fmt, int width, int height)
{
    AVFrame *frame;

    frame = av_frame_alloc();

    if (!frame)
    {
        return NULL;
    }

    frame->format = pix_fmt;
    frame->width = width;
    frame->height = height;

    int checkImage = av_image_alloc(frame->data, frame->linesize, width, height, pix_fmt, 32);

    if (checkImage < 0)
    {
        return NULL;
    }

    return frame;
}

bool initialize()
{
    AVRational frameRate;
    frameRate.den = m_frameRate;
    frameRate.num = 1;

    av_register_all();

    m_pCodec = avcodec_find_encoder_by_name(m_encoderName);

    if (!m_pCodec)
    {
        return false;
    }

    m_pCodecContext = avcodec_alloc_context3(m_pCodec);
    m_pCodecContext->width = m_imageWidth;
    m_pCodecContext->height = m_imageHeight;
    m_pCodecContext->time_base = frameRate;
    m_pCodecContext->gop_size = 0;
    m_pCodecContext->pix_fmt = m_pixType;
    m_pCodecContext->codec_id = m_pCodec->id;
    m_pCodecContext->bit_rate = m_targetBitRate;

    av_opt_set(m_pCodecContext->priv_data, "+CBR", "", 0);

    return true;
}

bool startExport()
{
    m_frameIndex = 0;
    char fakeFileName[512];
    int checkAllocContext = avformat_alloc_output_context2(&m_pFormatContext, NULL, m_fileType, fakeFileName);

    if (checkAllocContext < 0)
    {
        return false;
    }

    if (!m_pFormatContext)
    {
        return false;
    }

    m_pFormat = m_pFormatContext->oformat;

    if (m_pFormat->video_codec != AV_CODEC_ID_NONE)
    {
        addVideoStream();

        int checkOpen = avcodec_open2(m_pCodecContext, m_pCodec, NULL);

        if (checkOpen < 0)
        {
            return false;
        }

        m_pFrame = allocatePicture(m_pCodecContext->pix_fmt, m_pCodecContext->width, m_pCodecContext->height);
        if (!m_pFrame)
        {
            return false;
        }
        m_pFrame->pts = 0;
    }

    int checkOpen = avio_open(&m_pFormatContext->pb, m_fileName, AVIO_FLAG_WRITE);
    if (checkOpen < 0)
    {
        return false;
    }

    av_dict_set(&(m_pFormatContext->metadata), "title", "QS Test", 0);

    int checkHeader = avformat_write_header(m_pFormatContext, NULL);
    if (checkHeader < 0)
    {
        return false;
    }

    return true;
}

int processFrame(AVPacket& avPacket)
{
    avPacket.stream_index = 0;
    avPacket.pts = av_rescale_q(m_pFrame->pts, m_pStream->codec->time_base, m_pStream->time_base);
    avPacket.dts = av_rescale_q(m_pFrame->pts, m_pStream->codec->time_base, m_pStream->time_base);
    m_pFrame->pts++;

    int retVal = av_interleaved_write_frame(m_pFormatContext, &avPacket);
    return retVal;
}

bool exportFrame()
{
    int success = 1;
    int result = 0;

    AVPacket avPacket;

    av_init_packet(&avPacket);
    avPacket.data = NULL;
    avPacket.size = 0;

    AVPacket* avPacket1 = new AVPacket();

    av_init_packet(avPacket1);
    avPacket1->data = NULL;
    avPacket1->size = 0;

    fflush(stdout);

    std::cout << "Before avcodec_encode_video2 for frame: " << m_frameIndex << std::endl;
    success = avcodec_encode_video2(m_pCodecContext, &avPacket, m_pFrame, &result);
    std::cout << "After avcodec_encode_video2 for frame: " << m_frameIndex << std::endl;

    if (result)
    {
        success = processFrame(avPacket);
    }

    av_packet_unref(&avPacket);

    av_free_packet(&avPacket);
    //av_frame_free(&m_pFrame);

    m_frameIndex++;
    return (success == 0);
}

void endExport()
{
    int result = 0;
    int success = 0;

    if (m_pFrame)
    {
        while (success == 0)
        {
            AVPacket avPacket;
            av_init_packet(&avPacket);
            avPacket.data = NULL;
            avPacket.size = 0;

            fflush(stdout);
            success = avcodec_encode_video2(m_pCodecContext, &avPacket, NULL, &result);

            if (result)
            {
                success = processFrame(avPacket);
            }
            av_packet_unref(&avPacket);

            if (!result)
            {
                break;
            }
        }
    }

    if (m_pFormatContext)
    {
        av_write_trailer(m_pFormatContext);

        if (m_pFrame)
        {
            av_frame_free(&m_pFrame);
        }

        avio_closep(&m_pFormatContext->pb);
        avformat_free_context(m_pFormatContext);
        m_pFormatContext = NULL;
    }
}

void cleanup()
{
    if (m_pFrame || m_pCodecContext)
    {
        if (m_pFrame)
        {
            av_frame_free(&m_pFrame);
        }

        if (m_pCodecContext)
        {
            avcodec_close(m_pCodecContext);
            av_free(m_pCodecContext);
        }
    }
}

int main()
{
    bool success = true;
    if (initialize())
    {
        if (startExport())
        {
            for (int loop = 0; loop < m_frameCount; loop++)
            {
                if (!exportFrame())
                {
                    std::cout << "Failed to export frame\n";
                    success = false;
                    break;
                }
            }
            endExport();
        }
        else
        {
            std::cout << "Failed to start export\n";
            success = false;
        }

        cleanup();
    }
    else
    {
        std::cout << "Failed to initialize export\n";
        success = false;
    }

    if (success)
    {
        std::cout << "Successfully exported file\n";
    }
    return 1;
}