Remuxing arbitrary video stream into MKV via FFMPEG

remuxer.cc

#include <atomic>
#include <condition_variable>
#include <deque>
#include <functional>
#include <iostream>
#include <memory>
#include <mutex>
#include <thread>

extern "C" {
#include "libavcodec/avcodec.h"
#include "libavfilter/avfilter.h"
#include "libavfilter/buffersink.h"
#include "libavfilter/buffersrc.h"
#include <libavcodec/avcodec.h>
#include <libavdevice/avdevice.h>
#include <libavformat/avformat.h>
#include <libswscale/swscale.h>
}

// static const char SOURCE_NAME[] = "http://81.83.10.9:8001/mjpg/video.mjpg";
static const char SOURCE_NAME[] = "http://qthttp.apple.com.edgesuite.net/1010qwoeiuryfg/sl.m3u8";
// static const char SOURCE_NAME[] = "https://bitdash-a.akamaihd.net/content/MI201109210084_1/m3u8s/f08e80da-bf1d-4e3d-8899-f0f6155f6efa.m3u8";

using Pkt = std::unique_ptr<AVPacket, void (*)(AVPacket *)>;
std::deque<Pkt> frame_buffer;
std::mutex frame_mtx;
std::condition_variable frame_cv;
std::atomic_bool keep_running{true};

AVCodecParameters *common_codecpar = nullptr;
AVRational time_base{0, 0};
AVRational frame_rate{25, 1};
std::mutex codecpar_mtx;
std::condition_variable codecpar_cv;

void read_frames_from_source(unsigned N)
{
    AVFormatContext *fmt_ctx = avformat_alloc_context();
    int err = avformat_open_input(&fmt_ctx, SOURCE_NAME, nullptr, nullptr);
    if (err < 0) {
        std::cerr << "cannot open input" << std::endl;
        avformat_free_context(fmt_ctx);
        return;
    }

    err = avformat_find_stream_info(fmt_ctx, nullptr);
    if (err < 0) {
        std::cerr << "cannot find stream info" << std::endl;
        avformat_free_context(fmt_ctx);
        return;
    }

    // Simply finding the first video stream, preferrably H.264. Others are ignored below
    int video_stream_id = -1;
    for (unsigned i = 0; i < fmt_ctx->nb_streams; i++) {
        auto *c = fmt_ctx->streams[i]->codecpar;
        if (c->codec_type == AVMEDIA_TYPE_VIDEO) {
            video_stream_id = i;
            if (c->codec_id == AV_CODEC_ID_H264)
                break;
        }
    }

    if (video_stream_id < 0) {
        std::cerr << "failed to find find video stream" << std::endl;
        avformat_free_context(fmt_ctx);
        return;
    }

    std::unique_ptr<AVBSFContext, void (*)(AVBSFContext *)> bsf_ctx{nullptr, nullptr};
    AVCodecID codec_id = fmt_ctx->streams[video_stream_id]->codecpar->codec_id;

    if (codec_id == AV_CODEC_ID_H264 || codec_id == AV_CODEC_ID_HEVC) {
        const char *filtername = codec_id == AV_CODEC_ID_H264 ? "h264_mp4toannexb" : "hevc_mp4toannexb";

        const auto *bsf = av_bsf_get_by_name(filtername);
        if (!bsf) {
            std::cerr << "failed to find bit stream filter" << std::endl;
            return;
        }

        AVBSFContext *bsf_ctx_raw = nullptr;
        err = av_bsf_alloc(bsf, &bsf_ctx_raw);
        if (err < 0) {
            std::cerr << "failed to allocate bit stream filter context" << std::endl;
            return;
        }

        bsf_ctx = std::unique_ptr<AVBSFContext, void (*)(AVBSFContext *)>{bsf_ctx_raw, [](AVBSFContext *p) { av_bsf_free(&p); }};
    }

    {   // Here we have the codec params and can launch the writer
        std::lock_guard<std::mutex> locker(codecpar_mtx);
        common_codecpar = fmt_ctx->streams[video_stream_id]->codecpar;
        time_base = fmt_ctx->streams[video_stream_id]->time_base;
        frame_rate = fmt_ctx->streams[video_stream_id]->avg_frame_rate;

        if (bsf_ctx) {
            err = avcodec_parameters_copy(bsf_ctx->par_in, common_codecpar);
            if (err < 0) {
                std::cerr << "failed to copy parameters to bit stream filter parameters" << std::endl;
                return;
            }

            err = av_bsf_init(bsf_ctx.get());
            if (err < 0) {
                std::cerr << "failed to init bit stream filter" << std::endl;
                return;
            }

            err = avcodec_parameters_copy(common_codecpar, bsf_ctx->par_out);
            if (err < 0) {
                std::cerr << "failed to copy parameters from bit stream filter" << std::endl;
                return;
            }
        }
    }
    codecpar_cv.notify_all();

    unsigned cnt = 0;
    while (++cnt <= N) { // we read some limited number of frames
        Pkt pkt{av_packet_alloc(), [](AVPacket *p) { av_packet_free(&p); }};

        err = av_read_frame(fmt_ctx, pkt.get());
        if (err < 0) {
            std::cerr << "read packet error" << std::endl;
            continue;
        }

        // That's why the cycle above, we write only one video stream here
        if (pkt->stream_index != video_stream_id)
            continue;

        if (bsf_ctx) {
            AVPacket raw_filtered_pkt;
            err = av_bsf_send_packet(bsf_ctx.get(), pkt.get());
            if (err < 0) {
                std::cerr << "failed to send packet to bitstream filter" << std::endl;
                return;
            }

            memset(&raw_filtered_pkt, 0, sizeof(raw_filtered_pkt));
            err = av_bsf_receive_packet(bsf_ctx.get(), &raw_filtered_pkt);
            if (err == AVERROR(EAGAIN))
                continue;
            if (err != 0) {
                std::cerr << "failed to receive packet from bitstream filter" << std::endl;
                return;
            }

            av_packet_unref(pkt.get());
            av_packet_ref(pkt.get(), &raw_filtered_pkt);
            av_packet_unref(&raw_filtered_pkt);
        }

        {
            std::lock_guard<std::mutex> locker(frame_mtx);
            frame_buffer.push_back(std::move(pkt));
        }
        frame_cv.notify_one();
    }

    keep_running.store(false);
    avformat_free_context(fmt_ctx);
}

void write_frames_into_file(std::string filepath)
{
    AVFormatContext *out_ctx = nullptr;

    int err = avformat_alloc_output_context2(&out_ctx, nullptr, "matroska", filepath.c_str());
    if (err < 0) {
        std::cerr << "avformat_alloc_output_context2 failed" << std::endl;
        return;
    }
    out_ctx->strict_std_compliance = FF_COMPLIANCE_UNOFFICIAL;

    AVCodec* codec = avcodec_find_encoder(common_codecpar->codec_id);
    AVStream *video_stream = avformat_new_stream(out_ctx, codec); // the proper way

    int video_stream_id = video_stream->index;

    AVCodecContext *encoder = avcodec_alloc_context3(codec);
    avcodec_parameters_to_context(encoder, common_codecpar);
    encoder->time_base = time_base;
    encoder->framerate = frame_rate;
    if (out_ctx->oformat->flags & AVFMT_GLOBALHEADER)
        encoder->flags |= AV_CODEC_FLAG_GLOBAL_HEADER;

    err = avcodec_open2(encoder, codec, nullptr);
    if (err < 0) {
        std::cerr << "avcodec_open2 failed" << std::endl;
        return;
    }

    err = avcodec_parameters_from_context(video_stream->codecpar, encoder);
    if (err < 0) {
        std::cerr << "avcodec_parameters_from_context failed" << std::endl;
        return;
    }

    if (!(out_ctx->flags & AVFMT_NOFILE)) {
        err =  avio_open(&out_ctx->pb, filepath.c_str(), AVIO_FLAG_WRITE);
        if (err < 0) {
            std::cerr << "avio_open fail" << std::endl;
            return;
        }
    }

    err = avformat_write_header(out_ctx, nullptr);
    if (err < 0) {
        char ffmpeg_err_buf[AV_ERROR_MAX_STRING_SIZE];
        av_make_error_string(&ffmpeg_err_buf[0], AV_ERROR_MAX_STRING_SIZE, err);
        std::cerr << "avformat_write_header failed: " << ffmpeg_err_buf << std::endl;
        return;
    }

    unsigned cnt = 0;
    while (true) {
        std::unique_lock<std::mutex> locker(frame_mtx);
        frame_cv.wait(locker, [&] { return !frame_buffer.empty() || !keep_running; });

        if (!keep_running)
            break;

        Pkt pkt = std::move(frame_buffer.front());
        frame_buffer.pop_front();
        ++cnt;
        locker.unlock();

        pkt->stream_index = video_stream_id; // mandatory
        pkt->pts = cnt * (1000 / frame_rate.num);
        pkt->dts = cnt * (1000 / frame_rate.num);

        err = av_write_frame(out_ctx, pkt.get());
        if (err < 0) {
            std::cerr << "av_write_frame failed " << cnt << std::endl;
        } else if (cnt % 25 == 0) {
            std::cout << cnt << " OK" << std::endl;
        }
    }

    av_write_trailer(out_ctx);
    avcodec_free_context(&encoder);
    avformat_free_context(out_ctx);
}

int main()
{
    std::thread reader(std::bind(&read_frames_from_source, 1000));
    std::thread writer;

    // Writer wont start until reader's got AVCodecParameters
    // In this example it spares us from setting writer's params properly manually

    {   // Waiting for codec params to be set
        std::unique_lock<std::mutex> locker(codecpar_mtx);
        codecpar_cv.wait(locker, [&] { return common_codecpar != nullptr; });
        writer = std::thread(std::bind(&write_frames_into_file, "out.mkv"));
    }

    reader.join();
    keep_running.store(false);
    writer.join();

    return 0;
}