FFMEPG FrameIOContext

FrameIO.cc

#include "FrameIO.h"
#include <stdexcept>

FrameIOContext::FrameIOContext(FILE* in)
    : m_file{ in }
    , m_inputCtx{ avformat_alloc_context() }
    , m_outputCtx{ nullptr }
    , m_ioCtx{ avio_alloc_context(
          reinterpret_cast<unsigned char*>(av_malloc(BufferSize)), BufferSize,
          0, m_file, &FrameIOContext::mem_read, nullptr, &FrameIOContext::mem_seek) }
{
    if (!m_inputCtx)
        throw std::runtime_error("could not allocate input context");
    if (!m_ioCtx)
        throw std::runtime_error("coult not allocate io context");
    m_inputCtx->pb = m_ioCtx;
}

FrameIOContext::~FrameIOContext()
{
    if (m_outputCtx) {
        av_write_trailer(m_outputCtx);
        if (m_outputCtx->pb && !(m_outputCtx->oformat->flags & AVFMT_NOFILE))
            avio_closep(&m_outputCtx->pb);
        avformat_free_context(m_outputCtx);
        m_outputCtx = nullptr;
    }
    avformat_close_input(&m_inputCtx);
    if (m_ioCtx) {
        av_freep(&m_ioCtx->buffer);
        av_freep(&m_ioCtx);
    }
}

static inline const char* getShortName(const std::string& uri)
{
    if (uri.compare(0, 7, "rtsp://") == 0)
        return "rtsp";
    else if (uri.compare(0, 7, "rtmp://") == 0)
        return "flv";
    return nullptr;
}

bool FrameIOContext::connect(const std::string& uri)
{
    if (m_outputCtx) {
        throw std::runtime_error("output already established");
        return false;
    }

    avformat_alloc_output_context2(&m_outputCtx, nullptr, ::getShortName(uri), uri.c_str());
    if (!m_outputCtx)
        return false;

    for (unsigned i = 0; i < m_inputCtx->nb_streams; ++i) {
        auto istream = m_inputCtx->streams[i];
        auto ostream = avformat_new_stream(m_outputCtx, istream->codec->codec);
        if (!ostream) {
            av_log(m_outputCtx, AV_LOG_ERROR, "could not create output stream");
            goto fail;
        }
        if (avcodec_copy_context(ostream->codec, istream->codec) < 0) {
            av_log(m_outputCtx, AV_LOG_ERROR, "could not copy codec context");
            goto fail;
        }
        ostream->codec->codec_tag = 0;
        ostream->codec->time_base.num = 0;
        ostream->codec->time_base.den = 1;
    }

    if (!(m_outputCtx->oformat->flags & AVFMT_NOFILE)) {
        if (avio_open(&m_outputCtx->pb, m_outputCtx->filename, AVIO_FLAG_WRITE) < 0) {
            av_log(m_outputCtx, AV_LOG_ERROR, "could not open output context io");
            goto fail;
        }
    }
    if (avformat_write_header(m_outputCtx, nullptr) < 0) {
        av_log(m_outputCtx, AV_LOG_ERROR, "could not write output context header");
        goto fail;
    }
    av_dump_format(m_outputCtx, 0, m_outputCtx->filename, true);
    return true;
fail:
    avformat_free_context(m_outputCtx);
    m_outputCtx = nullptr;
    return false;
}

int FrameIOContext::open()
{
    auto r = avformat_open_input(&m_inputCtx, nullptr, nullptr, nullptr);
    if (r < 0) {
        av_log(m_inputCtx, AV_LOG_ERROR, "could not open input context\n");
        return r;
    }
    r = avformat_find_stream_info(m_inputCtx, nullptr);
    if (r < 0) {
        av_log(m_inputCtx, AV_LOG_ERROR, "could not find stream info\n");
        return r;
    }
    av_dump_format(m_inputCtx, 0, nullptr, false);
    return 0;
}

bool FrameIOContext::isVideo(AVPacket& pkt, int isOutput)
{
    auto index = pkt.stream_index;
    auto stream = isOutput ? m_outputCtx->streams[index] : m_inputCtx->streams[index];
    return stream->codec->codec_type == AVMEDIA_TYPE_VIDEO;
}

int64_t FrameIOContext::rescalePts(AVPacket& pkt, AVRational& timeBase)
{
    auto index = pkt.stream_index;
    return av_rescale_q(pkt.pts, m_inputCtx->streams[index]->time_base, timeBase);
}

int FrameIOContext::write(AVPacket& pkt)
{
    auto index = pkt.stream_index;
    auto istream = m_inputCtx->streams[index];
    auto ostream = m_outputCtx->streams[index];
    pkt.pts = av_rescale_q_rnd(pkt.pts, istream->time_base, ostream->time_base,
        (AVRounding)(AV_ROUND_NEAR_INF | AV_ROUND_PASS_MINMAX));
    pkt.dts = av_rescale_q_rnd(pkt.dts, istream->time_base, ostream->time_base,
        (AVRounding)(AV_ROUND_NEAR_INF | AV_ROUND_PASS_MINMAX));
    pkt.duration = av_rescale_q(pkt.duration, istream->time_base, ostream->time_base);
    pkt.pos = -1;
    return av_interleaved_write_frame(m_outputCtx, &pkt);
}

FrameIOContext& FrameIOContext::operator>>(AVPacket& pkt)
{
    if (av_read_frame(m_inputCtx, &pkt) < 0)
        throw std::runtime_error("av_read_frame failed");
    return *this;
}

int FrameIOContext::mem_read(void* opaque, uint8_t* buf, int buf_size)
{
    auto file = reinterpret_cast<FILE*>(opaque);
    auto len = ::fread(buf, 1, buf_size, file);
    return len == 0 ? AVERROR_EOF : len;
}

int64_t FrameIOContext::mem_seek(void* opaque, int64_t pos, int whence)
{
    auto file = reinterpret_cast<FILE*>(opaque);
    if (::fseek(file, (long)pos, whence) != 0)
        return -1;
    return ::ftell(file);
}

FrameIO.h

#ifndef FrameIO_h
#define FrameIO_h

#include <string>
extern "C" {
#include <libavformat/avformat.h>
}

class FrameIOContext {
public:
    FrameIOContext(FILE* in);
    virtual ~FrameIOContext();

    int open();
    int write(AVPacket&);
    bool connect(const std::string&);
    bool isVideo(AVPacket&, int isOutput);
    int64_t rescalePts(AVPacket&, AVRational&);
    FrameIOContext& operator>>(AVPacket&);

private:
    FILE* m_file;
    AVFormatContext* m_inputCtx;
    AVFormatContext* m_outputCtx;
    AVIOContext* m_ioCtx;

private:
    const static uint64_t BufferSize = 32 * 1024;
    static int mem_read(void* opaque, uint8_t* buf, int buf_size);
    static int64_t mem_seek(void* opaque, int64_t pos, int whence);
};

#endif

main.cc

// clang++ -std=c++11 FrameIO.cc main.cc -lavformat -lavcodec -lavutil
#include "FrameIO.h"
#include <iostream>
#include <unistd.h>

extern "C" {
#include <libavutil/time.h>
}

int main(int argc, char const* argv[])
{
    av_register_all();
    avformat_network_init();
    av_log_set_level(AV_LOG_DEBUG);

    auto file = argc < 2 ? stdin : fopen(argv[1], "r");
    try {
        FrameIOContext ctx{ file };
        ctx.open();
        AVPacket pkt;
        if (argc > 2) {
            if (!ctx.connect(argv[2])) {
                std::cerr << "could not set output to " << argv[2] << std::endl;
                return 1;
            }
        }
        AVRational time_base_q = { 1, AV_TIME_BASE };
        auto start_time = av_gettime();
        while (true) {
            try {
                ctx >> pkt;
            } catch (const std::exception&) {
                break;
            }
            if (argc > 2) {
                if (ctx.isVideo(pkt, 0)) {
                    auto sleepTime = ctx.rescalePts(pkt, time_base_q) - (av_gettime() - start_time);
                    if (sleepTime > 0)
                        av_usleep(sleepTime);
                }
                ctx.write(pkt);
            } else
                std::cout << "pts: " << pkt.pts << " dts: " << pkt.dts << " duration: " << pkt.duration << std::endl;
            av_free_packet(&pkt);
        }
    } catch (const std::exception& e) {
        std::cerr << "caught " << e.what() << std::endl;
    }

    fclose(file);
    return 0;
}