Merge pull request #47 from soreau/fix-46

Update to avcodec_send_frame API
ammen99 · Nov 23, 2019 · 0a9b090 · 0a9b090
2 parents 7cb37c4 + 519ba54
commit 0a9b090
Show file tree

Hide file tree

Showing 2 changed files with 49 additions and 27 deletions.
diff --git a/src/frame-writer.cpp b/src/frame-writer.cpp
@@ -189,14 +189,15 @@ void FrameWriter::init_video_stream()
     if (fmtCtx->oformat->flags & AVFMT_GLOBALHEADER)
         videoCodecCtx->flags |= AV_CODEC_FLAG_GLOBAL_HEADER;
 
-    int err;
-    if ((err = avcodec_open2(videoCodecCtx, codec, &options)) < 0)
+    int ret;
+    char err[256];
+    if ((ret = avcodec_open2(videoCodecCtx, codec, &options)) < 0)
     {
-        std::cerr << "avcodec_open2 failed: " << averr(err) << std::endl;
+        av_strerror(ret, err, 256);
+        std::cerr << "avcodec_open2 failed: " << err << std::endl;
         std::exit(-1);
     }
     av_dict_free(&options);
-    videoStream->time_base = (AVRational){ 1, FPS };
 }
 
 static uint64_t get_codec_channel_layout(AVCodec *codec)
@@ -298,9 +299,13 @@ void FrameWriter::init_codecs()
     }
 
     AVDictionary *dummy = NULL;
-    if (avformat_write_header(fmtCtx, &dummy) != 0)
+    char err[256];
+    int ret;
+    if ((ret = avformat_write_header(fmtCtx, &dummy)) != 0)
     {
         std::cerr << "Failed to write file header" << std::endl;
+        av_strerror(ret, err, 256);
+        std::cerr << err << std::endl;
         std::exit(-1);
     }
     av_dict_free(&dummy);
@@ -419,6 +424,35 @@ void FrameWriter::convert_pixels_to_yuv(const uint8_t *pixels,
     }
 }
 
+void FrameWriter::encode(AVCodecContext *enc_ctx, AVFrame *frame, AVPacket *pkt)
+{
+    int ret;
+
+    /* send the frame to the encoder */
+    ret = avcodec_send_frame(enc_ctx, frame);
+    if (ret < 0)
+    {
+        fprintf(stderr, "error sending a frame for encoding\n");
+        return;
+    }
+
+    while (ret >= 0)
+    {
+        ret = avcodec_receive_packet(enc_ctx, pkt);
+        if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF)
+        {
+            return;
+        }
+        if (ret < 0)
+        {
+            fprintf(stderr, "error during encoding\n");
+            return;
+        }
+
+        finish_frame(enc_ctx, *pkt);
+    }
+}
+
 void FrameWriter::add_frame(const uint8_t* pixels, int64_t usec, bool y_invert)
 {
     /* Calculate data after y-inversion */
@@ -474,15 +508,11 @@ void FrameWriter::add_frame(const uint8_t* pixels, int64_t usec, bool y_invert)
     pkt.data = NULL;
     pkt.size = 0;
 
-    int got_output;
-    avcodec_encode_video2(videoCodecCtx, &pkt, *output_frame, &got_output);
+    encode(videoCodecCtx, *output_frame, &pkt);
 
     /* Restore frame buffer, so that it can be properly freed in the end */
     if (saved_buf0)
         encoder_frame->buf[0] = saved_buf0;
-
-    if (got_output)
-      finish_frame(pkt, true);
 }
 
 #define SRC_RATE 1e6
@@ -509,10 +539,7 @@ void FrameWriter::send_audio_pkt(AVFrame *frame)
     pkt.data = NULL;
     pkt.size = 0;
 
-    int got_output;
-    avcodec_encode_audio2(audioCodecCtx, &pkt, frame, &got_output);
-    if (got_output)
-      finish_frame(pkt, false);
+    encode(audioCodecCtx, frame, &pkt);
 }
 
 size_t FrameWriter::get_audio_buffer_size()
@@ -547,11 +574,11 @@ void FrameWriter::add_audio(const void* buffer)
     av_frame_free(&outputf);
 }
 
-void FrameWriter::finish_frame(AVPacket& pkt, bool is_video)
+void FrameWriter::finish_frame(AVCodecContext *enc_ctx, AVPacket& pkt)
 {
     static std::mutex fmt_mutex, pending_mutex;
 
-    if (is_video)
+    if (enc_ctx == videoCodecCtx)
     {
         av_packet_rescale_ts(&pkt, (AVRational){ 1, 1000000 }, videoStream->time_base);
         pkt.stream_index = videoStream->index;
@@ -584,18 +611,11 @@ FrameWriter::~FrameWriter()
     AVPacket pkt;
     av_init_packet(&pkt);
 
-    for (int got_output = 1; got_output;)
-    {
-        avcodec_encode_video2(videoCodecCtx, &pkt, NULL, &got_output);
-        if (got_output)
-            finish_frame(pkt, true);
-    }
+    encode(videoCodecCtx, NULL, &pkt);
 
-    for (int got_output = 1; got_output && params.enable_audio;)
+    if (params.enable_audio)
     {
-        avcodec_encode_audio2(audioCodecCtx, &pkt, NULL, &got_output);
-        if (got_output)
-            finish_frame(pkt, false);
+        encode(audioCodecCtx, NULL, &pkt);
     }
 
     // Writing the end of the file.

diff --git a/src/frame-writer.hpp b/src/frame-writer.hpp
@@ -100,14 +100,16 @@ class FrameWriter
     void convert_pixels_to_yuv(const uint8_t *pixels,
         const uint8_t *formatted_pixels, int stride[]);
 
+    void encode(AVCodecContext *enc_ctx, AVFrame *frame, AVPacket *pkt);
+
     SwrContext *swrCtx;
     AVStream *audioStream;
     AVCodecContext *audioCodecCtx;
     void init_swr();
     void init_audio_stream();
     void send_audio_pkt(AVFrame *frame);
 
-    void finish_frame(AVPacket& pkt, bool isVideo);
+    void finish_frame(AVCodecContext *enc_ctx, AVPacket& pkt);
 
 public :
     FrameWriter(const FrameWriterParams& params);