remux_scratch.cpp

#include <condition_variable>
#include <cstdio>
#include <mutex>
#include <queue>
#include <thread>

extern "C"
{
#include <libavcodec/avcodec.h>
#include <libavformat/avformat.h>
#include <libavutil/timestamp.h>
#include <libswscale/swscale.h>
}

#define QUEUE_SIZE 256

#define OUTPUT_FILENAME "output.mp4"
#define OUTPUT_FORMAT_NAME "mp4"
#define OUTPUT_FORMAT_MIME_TYPE "video/mp4"

#define OUTPUT_BITRATE 2400 * 1000
#define OUTPUT_CODEC_ID AV_CODEC_ID_H264
#define OUTPUT_CODEC_PROFILE 0
#define OUTPUT_DURATION 300
#define OUTPUT_FRAME_HEIGHT 1080
#define OUTPUT_FRAME_WIDTH 1920
#define OUTPUT_PIXEL_FORMAT AV_PIX_FMT_YUV420P

#define OUTPUT_TIMESCALE 90000

#define ASSERT_ERROR(ret, fail_cond, step)                                                                             \
    if (fail_cond)                                                                                                     \
    {                                                                                                                  \
        const int error_str_size = 64;                                                                                 \
        char error_str[error_str_size];                                                                                \
        av_make_error_string(error_str, error_str_size, ret);                                                          \
        fprintf(stderr, "%s: %s (%d)\n", step, error_str, ret);                                                        \
        return -1;                                                                                                     \
    }                                                                                                                  \
    else                                                                                                               \
        printf("%s: ok\n", step);

int open_input(const char *url, AVCodecContext **dec, int *stream_idx, AVFormatContext **fmt)
{
    int ret;

    *fmt = nullptr;
    ret = avformat_open_input(fmt, url, nullptr, nullptr);
    ASSERT_ERROR(ret, ret != 0, "[open_input] avformat_open_input")

    ret = avformat_find_stream_info(*fmt, nullptr);
    ASSERT_ERROR(ret, ret < 0, "[open_input] avformat_find_stream_info")

    AVCodec *codec;
    ret = av_find_best_stream(*fmt, AVMEDIA_TYPE_VIDEO, -1, -1, &codec, 0);
    ASSERT_ERROR(ret, ret < 0 || ret == AVERROR_STREAM_NOT_FOUND, "[open_input] av_find_best_stream")
    *stream_idx = ret;
    auto stream = (*fmt)->streams[*stream_idx];

    *dec = avcodec_alloc_context3(codec);
    avcodec_parameters_to_context(*dec, stream->codecpar);
    ret = avcodec_open2(*dec, codec, nullptr);
    ASSERT_ERROR(ret, ret != 0, "[open_input] avcodec_open2")

    return 0;
}

int close_input(AVCodecContext *dec, AVFormatContext *fmt)
{
    avcodec_close(dec);
    avformat_close_input(&fmt);
    return 0;
}

int open_output(const char *url, AVRational *time_base, AVCodecContext **enc, AVStream **stream, AVFormatContext **fmt)
{
    int ret;

    auto format = av_guess_format(OUTPUT_FORMAT_NAME, url, OUTPUT_FORMAT_MIME_TYPE);
    if (format == nullptr)
    {
        fprintf(stderr, "Format %s (%s) for %s is not supported", OUTPUT_FORMAT_NAME, OUTPUT_FORMAT_MIME_TYPE, url);
        return -1;
    }

    ret = avformat_alloc_output_context2(fmt, format, OUTPUT_FORMAT_NAME, url);
    ASSERT_ERROR(ret, ret < 0, "[open_output] avformat_alloc_output_context2")

    ret = avio_open(&(*fmt)->pb, url, AVIO_FLAG_WRITE);
    ASSERT_ERROR(ret, ret < 0, "[open_output] avio_open")
    printf("[open_output] avio_open: ok\n");

    auto codec = avcodec_find_encoder(OUTPUT_CODEC_ID);
    if (codec == nullptr)
    {
        fprintf(stderr, "Codec Id %d is not supported", OUTPUT_CODEC_ID);
        return -1;
    }

    *stream = avformat_new_stream(*fmt, codec);
    (*stream)->codecpar = avcodec_parameters_alloc();
    (*stream)->r_frame_rate = {OUTPUT_TIMESCALE, OUTPUT_DURATION};
    (*stream)->start_time = 0;
    (*stream)->time_base = *time_base = {1, OUTPUT_TIMESCALE};

    auto codec_par = (*stream)->codecpar;
    codec_par->bit_rate = OUTPUT_BITRATE;
    codec_par->codec_id = OUTPUT_CODEC_ID;
    codec_par->codec_type = AVMEDIA_TYPE_VIDEO;
    codec_par->format = OUTPUT_PIXEL_FORMAT;
    codec_par->height = OUTPUT_FRAME_HEIGHT;
    codec_par->profile = OUTPUT_CODEC_PROFILE;
    codec_par->width = OUTPUT_FRAME_WIDTH;

    ret = avformat_write_header(*fmt, nullptr);
    ASSERT_ERROR(ret, ret < 0, "[open_output] avformat_write_header")
    printf("[open_output] avformat_write_header: ok\n");

    *enc = avcodec_alloc_context3(codec);
    avcodec_parameters_to_context(*enc, codec_par);
    (*enc)->time_base = {1, OUTPUT_TIMESCALE};
    ret = avcodec_open2(*enc, codec, nullptr);
    ASSERT_ERROR(ret, ret != 0, "[open_output] avcodec_open2")
    printf("[open_output] avcodec_open2: ok\n");

    return 0;
}

int close_output(AVCodecContext *enc, AVFormatContext *fmt)
{
    av_interleaved_write_frame(fmt, nullptr);
    av_write_trailer(fmt);
    avformat_flush(fmt);
    avformat_free_context(fmt);

    avio_flush(fmt->pb);
    avio_closep(&fmt->pb);

    return 0;
}

void print_packet(AVPacket *pkt, AVRational time_base, const char *tag)
{
    char pts[64], dts[64];
    av_ts_make_time_string(pts, pkt->pts, &time_base);
    av_ts_make_time_string(dts, pkt->dts, &time_base);
    printf("%s: pts=%s\tdts=%s\n", tag, pts, dts);
}

#define ASSERT_LAV(condition, ret, step)                                                                               \
    if (!(condition))                                                                                                  \
    {                                                                                                                  \
        char error_str[AV_ERROR_MAX_STRING_SIZE];                                                                      \
        av_make_error_string(error_str, AV_ERROR_MAX_STRING_SIZE, ret);                                                \
        fprintf(stderr, "%s: %s (%d)\n", step, error_str, ret);                                                        \
        return ret;                                                                                                    \
    }                                                                                                                  \
    else                                                                                                               \
        printf("%s: ok\n", step);

inline int transform(AVFrame *out_frame, AVFrame *in_frame, SwsContext *&scaler)
{
    if (scaler == nullptr)
        scaler = sws_getContext(in_frame->width, in_frame->height, static_cast<AVPixelFormat>(in_frame->format),
                                out_frame->width, out_frame->height, static_cast<AVPixelFormat>(out_frame->format),
                                SWS_BILINEAR, nullptr, nullptr, nullptr);
    auto ret = sws_scale(scaler, in_frame->data, in_frame->linesize, 0, in_frame->height, out_frame->data,
                         out_frame->linesize);
    ASSERT_LAV(ret, ret == out_frame->height, "transcode# sws_scale")
    return 0;
}

struct FrameQueue
{
    std::condition_variable cv_reader;
    std::condition_variable cv_writer;
    std::mutex mutex;
    std::queue<AVFrame *> q;
};

AVFrame *dequeue_frame(FrameQueue *fq)
{
    while (true)
    {
        std::unique_lock lock(fq->mutex);

        if (fq->q.empty())
            continue;

        auto ret = fq->q.front();
        fq->q.pop();
        fq->cv_reader.notify_all();
        return ret;
    }
}

void enqueue_frame(AVFrame *frame, FrameQueue *fq)
{
    std::unique_lock lock(fq->mutex);
    if (fq->q.size() > QUEUE_SIZE)
        fq->cv_reader.wait(lock);

    fq->q.push(frame);
    fq->cv_writer.notify_one();

    printf("queue size: %lu\n", fq->q.size());
}

int write_frames(FrameQueue &fq, AVRational out_time_base, AVRational in_time_base, AVCodecContext *enc,
                 AVFormatContext *fmt)
{
    int ret;
    SwsContext *scaler;

    auto dump_packet = [&enc, &fmt, out_time_base]() {
        int ret;

        auto pkt = av_packet_alloc();
        ret = avcodec_receive_packet(enc, pkt);
        ASSERT_LAV(ret == 0, ret, "avcodec_receive_packet")

        print_packet(pkt, out_time_base, "output");

        ret = av_interleaved_write_frame(fmt, pkt);
        ASSERT_LAV(ret == 0, ret, "av_interleaved_write_frame")

        return ret;
    };

    auto out_frame = av_frame_alloc();
    out_frame->format = OUTPUT_PIXEL_FORMAT;
    out_frame->height = OUTPUT_FRAME_HEIGHT;
    out_frame->width = OUTPUT_FRAME_WIDTH;
    av_frame_get_buffer(out_frame, 32);

    AVFrame *in_frame;
    auto pkt = av_packet_alloc();
    do
    {
        in_frame = dequeue_frame(&fq);

        if (in_frame != nullptr)
        {
            transform(out_frame, in_frame, scaler);
            out_frame->pts = av_rescale_q(in_frame->pts, in_time_base, out_time_base);
            av_frame_free(&in_frame);

            ret = avcodec_send_frame(enc, out_frame);
        }
        else
            ret = avcodec_send_frame(enc, nullptr);

        ASSERT_LAV(ret == 0, ret, "avcodec_send_frame")

        while (0 == (ret = avcodec_receive_packet(enc, pkt)))
        {
            print_packet(pkt, out_time_base, "output");

            ret = av_interleaved_write_frame(fmt, pkt);
            ASSERT_LAV(ret == 0, ret, "av_interleaved_write_frame")

            pkt = av_packet_alloc();
        }

        if (ret == AVERROR_EOF)
        {
            printf("encoder eof\n");
            break;
        }

        ASSERT_LAV(ret == AVERROR(EAGAIN), ret, "avcodec_receive_packet")

    } while (nullptr != in_frame);

    return 0;
}

int read_frames(FrameQueue &fq, int stream_idx, AVCodecContext *dec, AVFormatContext *fmt)
{
    int ret;

    auto pkt = av_packet_alloc();
    while (0 == (ret = av_read_frame(fmt, pkt)))
    {
        if (pkt->stream_index != stream_idx)
            continue;

        print_packet(pkt, dec->time_base, "input");

        while (AVERROR(EAGAIN) == (ret = avcodec_send_packet(dec, pkt)))
        {
            auto frame = av_frame_alloc();
            ret = avcodec_receive_frame(dec, frame);
            ASSERT_LAV(ret == 0, ret, "avcodec_receive_frame")

            enqueue_frame(frame, &fq);
        }
        ASSERT_LAV(ret == 0, ret, "avcodec_send_packet")
    }
    ASSERT_LAV(ret == AVERROR_EOF, ret, "av_read_frame")

    enqueue_frame(nullptr, &fq);
    return 0;
}

const auto RESCALE_ROUNDING = static_cast<AVRounding>(AV_ROUND_NEAR_INF | AV_ROUND_PASS_MINMAX);

int main(int arg_count, char *args[])
{
    if (arg_count < 2)
    {
        fprintf(stderr, "No input url specified\n");
        return -1;
    }
    auto url = args[1];

    int ret;

    /* initialize input */

    AVCodecContext *dec;
    int in_stream_idx;
    AVFormatContext *in_fmt;

    if (open_input(url, &dec, &in_stream_idx, &in_fmt) != 0)
        return -1;

    // av_dump_format(in_fmt, in_stream_idx, url, 0);

    auto in_time_base = in_fmt->streams[in_stream_idx]->time_base;

    /* initialize output */

    AVCodecContext *enc;
    AVStream *out_stream;
    AVFormatContext *out_fmt;
    AVRational out_time_base;

    if (open_output(OUTPUT_FILENAME, &out_time_base, &enc, &out_stream, &out_fmt) != 0)
        return -1;
    printf("open_output: ok\n");

    // av_dump_format(out_fmt, out_stream->index, OUTPUT_FILENAME, 1);

    /* remux frames */

    FrameQueue fq;

    std::thread reader([&]() {
        auto r = read_frames(fq, in_stream_idx, dec, in_fmt);
        if (r == 0)
            printf("read_frames: ok\n");
        else
            enqueue_frame(nullptr, &fq);
    });
    reader.join();

    std::thread writer([&]() {
        auto r = write_frames(fq, out_time_base, in_time_base, enc, out_fmt);
        ASSERT_LAV(r == 0, r, "write_frames")
        return 0;
    });
    writer.join();

    /* close output */

    if (close_output(enc, out_fmt) != 0)
        return -1;
    printf("close_output: ok\n");

    /* close input */

    if (close_input(dec, in_fmt) != 0)
        return -1;
    printf("close_input: ok\n");

    return 0;
}