ffmpegRecord.h

Go to the documentation of this file.

#ifndef FFMPEGRECORD_H
#define FFMPEGRECORD_H
 
#include <stdexcept>
#include <cstdint>
#include <vector>
#include <filesystem>
#include <string>
 
#include <glm/glm.hpp>
 
extern "C"
{
    #include <libavcodec/avcodec.h>
    #include <libavformat/avformat.h>
    #include <libavutil/avutil.h>
    #include <libavutil/imgutils.h>
    #include <libavutil/time.h>
    #include <libavutil/opt.h>
    #include <libswscale/swscale.h>
}
 
#include <record.h>
 
class FFmpegRecord : public Record
{
 
public:
 
    FFmpegRecord
    (
        std::filesystem::path file,
        glm::ivec2 resolution,
        uint8_t fps,
        std::string preset = "slow",
        std::string profile = "high",
        uint8_t crf = 0,
        uint8_t cq = 0,
        uint8_t qp = 0,
        std::string codecName = "libx264",
        uint32_t max_b_frames = 0,
        uint32_t gop_size = 1,
        uint32_t bitrate = 2000000
    )
    :
      Record(file, resolution, fps), bitrate(bitrate)
    {
        outputFormat = av_guess_format(nullptr, file.extension().string().c_str(), nullptr);
        if (!outputFormat)
        {
            throw std::runtime_error("Failed to create FFmpeg output format");
        }
 
        codec = avcodec_find_encoder_by_name(codecName.c_str());
        if (!codec)
        {
            throw std::runtime_error("Failed to create FFmpeg codec");
        }
 
        outputFormat->video_codec = codec->id;
 
        if (avformat_alloc_output_context2(&outputContext, outputFormat, nullptr, file.string().c_str()))
        {
            throw std::runtime_error("Failed to create FFmpeg output context");
        }
 
        AVStream * stream = avformat_new_stream(outputContext, codec);
        if (!stream)
        {
            throw std::runtime_error("Failed to create FFmpeg stream");
        }
 
        context = avcodec_alloc_context3(codec);
        if (!context)
        {
            throw std::runtime_error("Failed to create FFmpeg codec context");
        }
 
        stream->codecpar->codec_id = outputFormat->video_codec;
        stream->codecpar->codec_type = AVMEDIA_TYPE_VIDEO;
        stream->codecpar->width = resolution.x;
        stream->codecpar->height = resolution.y;
        stream->codecpar->format = AV_PIX_FMT_YUV420P;
        stream->codecpar->color_primaries = AVCOL_PRI_BT709;
        stream->codecpar->color_space = AVCOL_SPC_BT709;
        stream->codecpar->color_trc = AVCOL_TRC_BT709;
        stream->codecpar->bit_rate = bitrate;
        stream->time_base = (AVRational){ 1, fps };
        stream->r_frame_rate = (AVRational){ fps, 1 };
        stream->avg_frame_rate = (AVRational){ fps, 1 };
        avcodec_parameters_to_context(context, stream->codecpar);
        context->time_base = (AVRational){ 1, fps };
        context->max_b_frames = max_b_frames;
        context->gop_size = gop_size;
        context->framerate = (AVRational){ fps, 1 };
        context->pix_fmt = AV_PIX_FMT_YUV420P;
 
        av_opt_set(context->priv_data, "preset", preset.c_str(), 0);
        av_opt_set(context->priv_data, "profile", profile.c_str(), 0);
        av_opt_set_int(context->priv_data, "cq", cq, 0);
        av_opt_set_int(context->priv_data, "qp", qp, 0);
        av_opt_set_int(context->priv_data, "crf", crf, 0);
        avcodec_parameters_from_context(stream->codecpar, context);
    }
 
    ~FFmpegRecord()
    {
        if (videoFrame)
        {
            av_frame_free(&videoFrame);
        }
        if (context)
        {
            avcodec_free_context(&context);
        }
        if (outputContext)
        {
            avformat_free_context(outputContext);
        }
        if (scaleContext)
        {
            sws_freeContext(scaleContext);
        }
    }
 
    void open()
    {
        if (fileOpen) { return; }
        if ((avcodec_open2(context, codec, NULL)) < 0)
        {
            throw std::runtime_error("Failed to open FFmpeg codec");
        }
 
        if (!(outputFormat->flags & AVFMT_NOFILE))
        {
            auto code = avio_open(&outputContext->pb, file.string().c_str(), AVIO_FLAG_WRITE);
            if ((code) < 0)
            {
                throw std::runtime_error("Failed to open video file");
            }
        }
 
        if ((avformat_write_header(outputContext, NULL)) < 0)
        {
            throw std::runtime_error("Failed to write video header");
        }
 
        av_dump_format(outputContext, 0, file.string().c_str(), 1);
        fileOpen = true;
    }
 
    void close()
    {
        if (!fileOpen) { return; }
        fileOpen = false;
        AVPacket * pkt = av_packet_alloc();
        pkt->data = NULL;
        pkt->size = 0;
 
        while(true)
        {
            avcodec_send_frame(context, NULL);
            if (avcodec_receive_packet(context, pkt) == 0)
            {
                av_interleaved_write_frame(outputContext, pkt);
                av_packet_unref(pkt);
            }
            else
            {
                break;
            }
        }
 
        av_write_trailer(outputContext);
        if (!(outputFormat->flags & AVFMT_NOFILE))
        {
            if (avio_close(outputContext->pb) < 0)
            {
                throw std::runtime_error("Failed to close file");
            }
        }
    }
 
private:
 
    AVFrame * videoFrame = nullptr;
    AVCodec * codec = nullptr;
    AVStream * stream = nullptr;
    AVCodecContext * context = nullptr;
    SwsContext * scaleContext = nullptr;
    AVFormatContext * outputContext = nullptr;
    AVOutputFormat * outputFormat = nullptr;
 
    uint32_t bitrate = 2000000;
    uint16_t frameCounter = 0;
 
    void write(const uint8_t * frame)
    {
        if (!videoFrame)
        {
            videoFrame = av_frame_alloc();
            videoFrame->format = AV_PIX_FMT_YUV420P;
            videoFrame->width = context->width;
            videoFrame->height = context->height;
            if ((av_frame_get_buffer(videoFrame, 0)) < 0)
            {
                throw std::runtime_error("Failed to allocate FFmpeg frame");
            }
        }
        if (!scaleContext)
        {
            scaleContext = sws_getContext
            (
                context->width,
                context->height,
                AV_PIX_FMT_RGBA,
                context->width,
                context->height,
                AV_PIX_FMT_YUV420P,
                SWS_BICUBIC,
                0,
                0,
                0
            );
        }
 
        // Fixed, stride should be 4x for RBGA data.
        int stride[1] = {4 * context->width};
 
        sws_scale
        (
            scaleContext,
            (const uint8_t * const *)&frame,
            stride,
            0,
            context->height,
            videoFrame->data,
            videoFrame->linesize
        );
 
        videoFrame->pts = (1.0 / float(fps)) * 15360 * (frameCounter++);
        if ((avcodec_send_frame(context, videoFrame)) < 0)
        {
            throw std::runtime_error("Failed to send frame");
        }
 
        AVPacket * pkt = av_packet_alloc();
        pkt->data = NULL;
        pkt->size = 0;
        pkt->flags |= AV_PKT_FLAG_KEY;
        if (avcodec_receive_packet(context, pkt) == 0)
        {
            av_interleaved_write_frame(outputContext, pkt);
            av_packet_unref(pkt);
        }
    }
};
 
#endif /* FFMPEGRECORD_H */