[FFmpeg-devel] [PATCH 2/3] avformat/mpegtsenc: use an AVParser for h264 streams

[James Almer]

Use it to enable keyframe tagging heuristics to mark as many RAPs as possible.

Signed-off-by: James Almer <jamrial at gmail.com>
---
 configure               |  2 +-
 libavformat/mpegtsenc.c | 45 +++++++++++++++++++++++++++++++++++++----
 2 files changed, 42 insertions(+), 5 deletions(-)

diff --git a/configure b/configure
index 2d2d125fd3..bbf0283a67 100755
--- a/configure
+++ b/configure
@@ -3378,7 +3378,7 @@ mp3_demuxer_select="mpegaudio_parser"
 mp3_muxer_select="mpegaudioheader"
 mp4_muxer_select="mov_muxer"
 mpegts_demuxer_select="iso_media"
-mpegts_muxer_select="ac3_parser adts_muxer latm_muxer h264_mp4toannexb_bsf hevc_mp4toannexb_bsf"
+mpegts_muxer_select="ac3_parser adts_muxer latm_muxer h264_parser h264_mp4toannexb_bsf hevc_mp4toannexb_bsf"
 mpegtsraw_demuxer_select="mpegts_demuxer"
 mxf_muxer_select="golomb pcm_rechunk_bsf"
 mxf_d10_muxer_select="mxf_muxer"
diff --git a/libavformat/mpegtsenc.c b/libavformat/mpegtsenc.c
index 98dac17994..4e15f3da07 100644
--- a/libavformat/mpegtsenc.c
+++ b/libavformat/mpegtsenc.c
@@ -257,6 +257,10 @@ typedef struct MpegTSWriteStream {
     int opus_queued_samples;
     int opus_pending_trim_start;
 
+    /* For H.264 */
+    AVCodecParserContext *parser;
+    AVCodecContext *parser_avctx;
+
     DVBAC3Descriptor *dvb_ac3_desc;
 } MpegTSWriteStream;
 
@@ -1216,6 +1220,21 @@ static int mpegts_init(AVFormatContext *s)
         ts_st->payload_dts     = AV_NOPTS_VALUE;
         ts_st->cc              = 15;
         ts_st->discontinuity   = ts->flags & MPEGTS_FLAG_DISCONT;
+        if (st->codecpar->codec_id == AV_CODEC_ID_H264) {
+            ts_st->parser = av_parser_init(st->codecpar->codec_id);
+            if (!ts_st->parser)
+                return AVERROR(ENOMEM);
+            ts_st->parser_avctx = avcodec_alloc_context3(NULL);
+            if (!ts_st->parser_avctx)
+                return AVERROR(ENOMEM);
+            ret = avcodec_parameters_to_context(ts_st->parser_avctx, st->codecpar);
+            if (ret < 0)
+                return ret;
+            // We only want to parse frame headers
+            ts_st->parser->flags |= PARSER_FLAG_COMPLETE_FRAMES;
+            // And we want keyframe tagging heuristics
+            ts_st->parser->flags |= PARSER_FLAG_USE_KEYFRAME_HEURISTICS;
+        }
         if (st->codecpar->codec_id == AV_CODEC_ID_AAC &&
             st->codecpar->extradata_size > 0) {
             AVStream *ast;
@@ -1832,6 +1851,7 @@ static int mpegts_write_packet_internal(AVFormatContext *s, AVPacket *pkt)
     const int64_t delay = av_rescale(s->max_delay, 90000, AV_TIME_BASE) * 2;
     const int64_t max_audio_delay = av_rescale(s->max_delay, 90000, AV_TIME_BASE) / 2;
     int64_t dts = pkt->dts, pts = pkt->pts;
+    int flags = pkt->flags;
     int opus_samples = 0;
     size_t side_data_size;
     uint8_t *side_data = NULL;
@@ -1864,11 +1884,18 @@ static int mpegts_write_packet_internal(AVFormatContext *s, AVPacket *pkt)
     if (st->codecpar->codec_id == AV_CODEC_ID_H264) {
         const uint8_t *p = buf, *buf_end = p + size;
         uint32_t state = -1;
-        int extradd = (pkt->flags & AV_PKT_FLAG_KEY) ? st->codecpar->extradata_size : 0;
+        int extradd;
         int ret = ff_check_h264_startcode(s, st, pkt);
         if (ret < 0)
             return ret;
 
+        av_parser_parse2(ts_st->parser, ts_st->parser_avctx,
+                         &buf, &size, pkt->data, pkt->size,
+                         pkt->pts, pkt->dts, pkt->pos);
+        if (ts_st->parser->key_frame)
+            flags |= AV_PKT_FLAG_KEY;
+
+        extradd = (flags & AV_PKT_FLAG_KEY) ? st->codecpar->extradata_size : 0;
         if (extradd && AV_RB24(st->codecpar->extradata) > 1)
             extradd = 0;
 
@@ -2106,7 +2133,7 @@ static int mpegts_write_packet_internal(AVFormatContext *s, AVPacket *pkt)
         av_assert0(!ts_st->payload_size);
         // for video and subtitle, write a single pes packet
         mpegts_write_pes(s, st, buf, size, pts, dts,
-                         pkt->flags & AV_PKT_FLAG_KEY, stream_id);
+                         flags & AV_PKT_FLAG_KEY, stream_id);
         ts_st->opus_queued_samples = 0;
         av_free(data);
         return 0;
@@ -2115,7 +2142,7 @@ static int mpegts_write_packet_internal(AVFormatContext *s, AVPacket *pkt)
     if (!ts_st->payload_size) {
         ts_st->payload_pts   = pts;
         ts_st->payload_dts   = dts;
-        ts_st->payload_flags = pkt->flags;
+        ts_st->payload_flags = flags;
     }
 
     memcpy(ts_st->payload + ts_st->payload_size, buf, size);
@@ -2184,6 +2211,9 @@ static void mpegts_deinit(AVFormatContext *s)
         if (ts_st) {
             av_freep(&ts_st->dvb_ac3_desc);
             av_freep(&ts_st->payload);
+            avcodec_free_context(&ts_st->parser_avctx);
+            av_parser_close(ts_st->parser);
+            ts_st->parser = NULL;
             if (ts_st->amux) {
                 avformat_free_context(ts_st->amux);
                 ts_st->amux = NULL;
@@ -2207,8 +2237,15 @@ static int mpegts_check_bitstream(struct AVFormatContext *s, const AVPacket *pkt
         if (pkt->size >= 5 && AV_RB32(pkt->data) != 0x0000001 &&
                              (AV_RB24(pkt->data) != 0x000001 ||
                               (st->codecpar->extradata_size > 0 &&
-                               st->codecpar->extradata[0] == 1)))
+                               st->codecpar->extradata[0] == 1))) {
             ret = ff_stream_add_bitstream_filter(st, "h264_mp4toannexb", NULL);
+            if (ret == 1) {
+                MpegTSWriteStream *ts_st = st->priv_data;
+                int err = avcodec_parameters_to_context(ts_st->parser_avctx, st->internal->bsfc->par_out);
+                if (err < 0)
+                    return err;
+            }
+        }
     } else if (st->codecpar->codec_id == AV_CODEC_ID_HEVC) {
         if (pkt->size >= 5 && AV_RB32(pkt->data) != 0x0000001 &&
                              (AV_RB24(pkt->data) != 0x000001 ||
-- 
2.32.0