Ticket #7472: main.cpp

extern "C" {
#include <libavutil/opt.h>
#include <libavcodec/avcodec.h>
#include <libavutil/channel_layout.h>
#include <libavutil/common.h>
#include <libavutil/imgutils.h>
#include <libavutil/mathematics.h>
#include <libavutil/samplefmt.h>

#include <libavformat/avformat.h>
#include <libavcodec/avcodec.h>
#include <libavutil/imgutils.h>
#include <libswscale/swscale.h>
}

#include <string>

#include <opencv2/opencv.hpp>

using namespace std;
using namespace cv;



int write_frame_(AVFormatContext *fmt_ctx, const AVRational *time_base, AVStream *st, AVPacket *pkt)
{
    /* rescale output packet timestamp values from codec to stream timebase */
    pkt->pts = av_rescale_q_rnd(pkt->pts, *time_base, st->time_base, AVRounding(AV_ROUND_NEAR_INF|AV_ROUND_PASS_MINMAX));
    pkt->dts = av_rescale_q_rnd(pkt->dts, *time_base, st->time_base, AVRounding(AV_ROUND_NEAR_INF|AV_ROUND_PASS_MINMAX));
    pkt->duration = av_rescale_q(pkt->duration, *time_base, st->time_base);
    pkt->stream_index = st->index;

    /* Write the compressed frame to the media file. */
    return av_interleaved_write_frame(fmt_ctx, pkt);
}


void initialize_avformat_context(AVFormatContext **fctx, const char *format_name)
{
    int ret = avformat_alloc_output_context2(fctx, nullptr, format_name, nullptr);
    if (ret < 0)
    {
        std::cout << "Could not allocate output format context!" << std::endl;
        exit(1);
    }
}


int initialize_io_context(AVFormatContext *fctx, const char *output)
{
    int ret = 0;
    if (!(fctx->oformat->flags & AVFMT_NOFILE))   {
        ret = avio_open2(&fctx->pb, output, AVIO_FLAG_WRITE, nullptr, nullptr);
        if (ret < 0)  {
            std::cout << "Could not open output IO context!" << std::endl;
            return ret;
        }
    }
    return ret;
}


void set_codec_params(AVFormatContext *fctx, AVCodecContext *codec_ctx, double width, double height, int fps, int bitrate)
{
    const AVRational dst_fps = {fps, 1};

    codec_ctx->codec_tag = 0;
    codec_ctx->codec_id = AV_CODEC_ID_H264;
    codec_ctx->codec_type = AVMEDIA_TYPE_VIDEO;
    codec_ctx->width = width;
    codec_ctx->height = height;
    codec_ctx->gop_size = 12;
    codec_ctx->pix_fmt = AV_PIX_FMT_YUV420P;
    codec_ctx->framerate = dst_fps;
    codec_ctx->time_base = av_inv_q(dst_fps);
    codec_ctx->bit_rate = bitrate;
    if (fctx->oformat->flags & AVFMT_GLOBALHEADER)
    {
        codec_ctx->flags |= AV_CODEC_FLAG_GLOBAL_HEADER;
    }
}


void initialize_codec_stream(AVStream *stream, AVCodecContext *codec_ctx, AVCodec *codec, std::string codec_profile)
{
    int ret = avcodec_parameters_from_context(stream->codecpar, codec_ctx);
    if (ret < 0)
    {
        std::cout << "Could not initialize stream codec parameters!" << std::endl;
        exit(1);
    }

    AVDictionary *codec_options = nullptr;
    av_dict_set(&codec_options, "profile", codec_profile.c_str(), 0);
    av_dict_set(&codec_options, "preset", "superfast", 0);
    av_dict_set(&codec_options, "tune", "zerolatency", 0);

    // open video encoder
    ret = avcodec_open2(codec_ctx, codec, &codec_options);
    if (ret < 0)
    {
        std::cout << "Could not open video encoder!" << std::endl;
        exit(1);
    }
}


SwsContext *initialize_sample_scaler(AVCodecContext *codec_ctx, int width, int height)
{
    SwsContext *swsctx = sws_getContext(width, height, AV_PIX_FMT_BGR24, width, height, codec_ctx->pix_fmt, SWS_BICUBIC, nullptr, nullptr, nullptr);
    if (!swsctx)
    {
        std::cout << "Could not initialize sample scaler!" << std::endl;
        exit(1);
    }

    return swsctx;
}


AVFrame *set_frame_buffer(AVCodecContext *codec_ctx, int width, int height, std::vector<uint8_t> &framebuf)
{
    AVFrame *frame = av_frame_alloc();

    av_image_fill_arrays(frame->data, frame->linesize, framebuf.data(), codec_ctx->pix_fmt, width, height, 1);

    frame->width = width;
    frame->height = height;
    frame->format = static_cast<int>(codec_ctx->pix_fmt);

    return frame;
}


void write_frame(AVCodecContext *codec_ctx, AVFormatContext *fmt_ctx, AVFrame *frame)
{
    AVPacket pkt = {0};
    av_init_packet(&pkt);

    int ret = avcodec_send_frame(codec_ctx, frame);
    if (ret < 0)
    {
        std::cout << "Error sending frame to codec context!" << std::endl;
        exit(1);
    }

    ret = avcodec_receive_packet(codec_ctx, &pkt);
    if (ret < 0)
    {
        std::cout << "Error receiving packet from codec context!" << std::endl;
        exit(1);
    }

    av_interleaved_write_frame(fmt_ctx, &pkt);
    av_packet_unref(&pkt);
}


class Streamer
{
    AVFormatContext *ofmt_ctx;
    AVCodec *out_codec;
    AVStream *out_stream;
    AVCodecContext *out_codec_ctx;
    AVFrame *frame;
    SwsContext *swsctx;
    std::vector<uint8_t> framebuf;
    bool init_failure;

public:
    Streamer();
    ~Streamer();

    void write_video_frame(cv::Mat &image, int play);
    void put_frame(cv::Mat &image);

    void stream_init(int width, int height, int fps, int bitrate, const std::string &codec_profile, const std::string &server);


};



void Streamer::put_frame(cv::Mat &image)
{
    if (init_failure) {
        return;
    }

    const int stride[] = {static_cast<int>(image.step[0])};
    sws_scale(swsctx, &image.data, stride, 0, image.rows, frame->data, frame->linesize);

    frame->pts += av_rescale_q(1, out_codec_ctx->time_base, out_stream->time_base);

    write_frame(out_codec_ctx, ofmt_ctx, frame);
}


void Streamer::stream_init(int width, int height, int fps, int bitrate, const std::string &codec_profile, const std::string &server)
{

    av_log_set_level(AV_LOG_DEBUG);

    printf("init network\n");
    avformat_network_init();

    int ret;


    initialize_avformat_context(&ofmt_ctx, "flv");

    int init = 1;
    init = initialize_io_context(ofmt_ctx, server.c_str());
    if (init != 0) { // failure to initialize
        init_failure = true;
        exit(1);
        return;
    }


    out_codec = avcodec_find_encoder(AV_CODEC_ID_H264);
    out_stream = avformat_new_stream(ofmt_ctx, out_codec);
    out_codec_ctx = avcodec_alloc_context3(out_codec);

    set_codec_params(ofmt_ctx, out_codec_ctx, width, height, fps, bitrate);
    initialize_codec_stream(out_stream, out_codec_ctx, out_codec, codec_profile);

    out_stream->codecpar->extradata = out_codec_ctx->extradata;
    out_stream->codecpar->extradata_size = out_codec_ctx->extradata_size;

    av_dump_format(ofmt_ctx, 0, server.c_str(), 1);

    swsctx = initialize_sample_scaler(out_codec_ctx, width, height);
    framebuf.resize(av_image_get_buffer_size(out_codec_ctx->pix_fmt, width, height, 1));
    frame = set_frame_buffer(out_codec_ctx, width, height, framebuf);

    ret = avformat_write_header(ofmt_ctx, nullptr);
    if (ret < 0)
    {
        std::cout << "Could not write header!" << std::endl;
        exit(1);
    }
}



Streamer::Streamer()
{
    ofmt_ctx = nullptr;
    out_codec = nullptr;
    out_stream = nullptr;
    out_codec_ctx = nullptr;
    frame = nullptr;
    swsctx = nullptr;
    init_failure = false;

#if LIBAVCODEC_VERSION_INT < AV_VERSION_INT(58, 9, 100)
    av_register_all();
#endif
}


Streamer::~Streamer()
{
    if(ofmt_ctx) av_write_trailer(ofmt_ctx);
    if(frame) {
        av_frame_free(&frame);
        frame = nullptr;
    }


    if(ofmt_ctx) {
        avio_close(ofmt_ctx->pb);
        avformat_free_context(ofmt_ctx);
        ofmt_ctx = nullptr;
    }

    if(out_codec_ctx) {
        avcodec_close(out_codec_ctx); //double free
        out_codec_ctx = nullptr;
    }

}


int main(int argc, char *argv[])
{
    std::string video_fname;
    video_fname = std::string(argv[1]);

    Streamer streamer;

    VideoCapture video_capture(video_fname, CAP_FFMPEG);
    if(!video_capture.isOpened()) {
        CV_Error(cv::Error::StsBadArg, "could not open video path/file \n");
        return 1;
    }

    int cap_frame_width = video_capture.get(CAP_PROP_FRAME_WIDTH);
    int cap_frame_height = video_capture.get(CAP_PROP_FRAME_HEIGHT);

    int cap_fps = video_capture.get(CAP_PROP_FPS);
    printf("video info w = %d, h = %d, fps = %d\n", cap_frame_width, cap_frame_height, cap_fps);

    int bitrate = 1000000;

    streamer.stream_init(cap_frame_width, cap_frame_height, cap_fps, bitrate,
                         "high444", "rtmp://localhost/live/mystream");

    Mat frame;
    bool ok = video_capture.read(frame);
    while(ok) {
        streamer.put_frame(frame);
        ok = video_capture.read(frame);
    }
    video_capture.release();

    return 0;
}