save_tool.cpp 4.05 KB
extern "C"
{
	#include <libavcodec/avcodec.h> 
	#include <libavdevice/avdevice.h> 
	#include <libavformat/avformat.h> 
	#include <libavfilter/avfilter.h> 
	#include <libavutil/avutil.h> 
    #include <libavutil/pixdesc.h> 
	#include <libswscale/swscale.h>
    #include <libavutil/imgutils.h>
}

#include <chrono>
#include<string>

using namespace std;

bool bFirstFrame = false;
long last_src_pts = 0;
long last_pts = 0;

void update_pts(AVPacket* pkt) {
	if (pkt->pts > 0) {
		if (bFirstFrame) {
			bFirstFrame = false;
			last_src_pts = pkt->pts;
		}
		int64_t pkt_pts = pkt->pts;
		pkt->pts = last_pts + (pkt_pts - last_src_pts);
		last_src_pts = pkt_pts;
		last_pts = pkt->pts;
		pkt->dts = pkt->pts;
	}
	else {
		if (bFirstFrame) {
			bFirstFrame = false;
			last_pts = 0;
		}
		pkt->pts = last_pts + 512;
		last_pts = pkt->pts;
	}
	
}

void *save_tool(const string uri)
{
    AVFormatContext *i_fmt_ctx;
    AVCodecContext* codec_ctx = nullptr;
    AVCodec *decoder = nullptr;
    AVStream *i_video_stream;
    int video_index;

    AVFormatContext *o_fmt_ctx;
    AVStream *o_video_stream;

    bool bStop = false;
    int frame_nums = 0;

    avcodec_register_all();
    av_register_all();
    avformat_network_init();

    /* should set to NULL so that avformat_open_input() allocate a new one */
    i_fmt_ctx = NULL;

    const char *filename = "2.mp4";

    if (avformat_open_input(&i_fmt_ctx, uri.c_str(), NULL, NULL)!=0)
    {
        fprintf(stderr, " = could not open input file\n");
        return nullptr;
    }

    if (avformat_find_stream_info(i_fmt_ctx, NULL)<0)
    {
        fprintf(stderr, " = could not find stream info\n");
        return nullptr;
    }

    video_index = av_find_best_stream(i_fmt_ctx, AVMEDIA_TYPE_VIDEO, -1, -1, &decoder, 0);
	if (video_index < 0) {
		av_log(nullptr, AV_LOG_ERROR, "Cannot find a video stream in the input file ! \n");
		return nullptr;
	}

    codec_ctx = avcodec_alloc_context3(decoder);
    if (!codec_ctx){
        return nullptr;
    }

    i_video_stream = i_fmt_ctx->streams[video_index];
	if(avcodec_parameters_to_context(codec_ctx, i_video_stream->codecpar) < 0){
		return nullptr;
	}

    avformat_alloc_output_context2(&o_fmt_ctx, NULL, NULL, filename);

    o_video_stream = avformat_new_stream(o_fmt_ctx, NULL);
    {
        AVCodecContext *c;
        c = o_video_stream->codec;
        c->bit_rate = 400000;
        c->codec_id = i_video_stream->codec->codec_id;
        c->codec_type = i_video_stream->codec->codec_type;
        c->time_base.num = i_video_stream->time_base.num;
        c->time_base.den = i_video_stream->time_base.den;
        fprintf(stderr, " = time_base.num = %d time_base.den = %d\n", c->time_base.num, c->time_base.den);
        c->width = i_video_stream->codec->width;
        c->height = i_video_stream->codec->height;
        c->pix_fmt = i_video_stream->codec->pix_fmt;
        printf(" = width: %d height: %d pix_fmt: %d\n", c->width, c->height, c->pix_fmt);
        c->flags = i_video_stream->codec->flags;
        c->flags |= AV_CODEC_FLAG_GLOBAL_HEADER;
        c->me_range = i_video_stream->codec->me_range;
        c->max_qdiff = i_video_stream->codec->max_qdiff;

        c->qmin = i_video_stream->codec->qmin;
        c->qmax = i_video_stream->codec->qmax;

        c->qcompress = i_video_stream->codec->qcompress;
    }

    avio_open(&o_fmt_ctx->pb, filename, AVIO_FLAG_WRITE);

    avformat_write_header(o_fmt_ctx, NULL);

    AVPacket* i_pkt = av_packet_alloc();
    av_init_packet(i_pkt);

    while (av_read_frame(i_fmt_ctx, i_pkt) >= 0){

        update_pts(i_pkt);

        static int num = 1;
        printf(" = frame %d\n", num++);
        av_interleaved_write_frame(o_fmt_ctx, i_pkt);

        av_packet_unref(i_pkt);

        if(frame_nums > 750){
            break;
        }
        frame_nums++;
    }

    av_packet_free(&i_pkt);

    avformat_close_input(&i_fmt_ctx);

    av_write_trailer(o_fmt_ctx);

    avcodec_close(o_fmt_ctx->streams[0]->codec);
    av_freep(&o_fmt_ctx->streams[0]->codec);
    av_freep(&o_fmt_ctx->streams[0]);

    avio_close(o_fmt_ctx->pb);
    av_free(o_fmt_ctx);

    printf("end. \n");
}