[FFmpeg-devel] [PATCH] avformat: add H264 and HEVC support in IVF muxer

[alx.sukhanov at gmail.com]

From: Alex Sukhanov <asukhanov at google.com>

---
 libavformat/ivfenc.c | 50 +++++++++++++++++++++++++++++++++-----------
 1 file changed, 38 insertions(+), 12 deletions(-)

diff --git a/libavformat/ivfenc.c b/libavformat/ivfenc.c
index 66441a2a43..6410828533 100644
--- a/libavformat/ivfenc.c
+++ b/libavformat/ivfenc.c
@@ -36,19 +36,29 @@ static int ivf_write_header(AVFormatContext *s)
         return AVERROR(EINVAL);
     }
     par = s->streams[0]->codecpar;
-    if (par->codec_type != AVMEDIA_TYPE_VIDEO ||
-        !(par->codec_id == AV_CODEC_ID_AV1 ||
-          par->codec_id == AV_CODEC_ID_VP8 ||
-          par->codec_id == AV_CODEC_ID_VP9)) {
-        av_log(s, AV_LOG_ERROR, "Currently only VP8, VP9 and AV1 are supported!\n");
-        return AVERROR(EINVAL);
-    }
     avio_write(pb, "DKIF", 4);
     avio_wl16(pb, 0); // version
     avio_wl16(pb, 32); // header length
-    avio_wl32(pb,
-              par->codec_id == AV_CODEC_ID_VP9 ? AV_RL32("VP90") :
-              par->codec_id == AV_CODEC_ID_VP8 ? AV_RL32("VP80") : AV_RL32("AV01"));
+    switch (par->codec_id) {
+      case AV_CODEC_ID_AV1:
+        avio_wl32(pb, AV_RL32("AV01"));
+        break;
+      case AV_CODEC_ID_H264:
+        avio_wl32(pb, AV_RL32("H264"));
+        break;
+      case AV_CODEC_ID_HEVC:
+        avio_wl32(pb, AV_RL32("HEVC"));
+        break;
+      case AV_CODEC_ID_VP8:
+        avio_wl32(pb, AV_RL32("VP80"));
+        break;
+      case AV_CODEC_ID_VP9:
+        avio_wl32(pb, AV_RL32("VP90"));
+        break;
+      default:
+        av_log(s, AV_LOG_ERROR, "Currently only AV1, H264, HEVC, VP8 and VP9 and AV1 are supported!\n");
+        return AVERROR(EINVAL);
+    }
     avio_wl16(pb, par->width);
     avio_wl16(pb, par->height);
     avio_wl32(pb, s->streams[0]->time_base.den);
@@ -95,16 +105,32 @@ static int ivf_check_bitstream(struct AVFormatContext *s, const AVPacket *pkt)
     int ret = 1;
     AVStream *st = s->streams[pkt->stream_index];
 
-    if (st->codecpar->codec_id == AV_CODEC_ID_VP9)
+    if (st->codecpar->codec_id == AV_CODEC_ID_H264) {
+        if (pkt->size >= 5 && AV_RB32(pkt->data) != 0x0000001 &&
+                             (AV_RB24(pkt->data) != 0x000001 ||
+                              (st->codecpar->extradata_size > 0 &&
+                               st->codecpar->extradata[0] == 1)))
+            ret = ff_stream_add_bitstream_filter(st, "h264_mp4toannexb", NULL);
+    } else if (st->codecpar->codec_id == AV_CODEC_ID_HEVC) {
+        if (pkt->size >= 5 && AV_RB32(pkt->data) != 0x0000001 &&
+                             (AV_RB24(pkt->data) != 0x000001 ||
+                              (st->codecpar->extradata_size > 0 &&
+                               st->codecpar->extradata[0] == 1)))
+            ret = ff_stream_add_bitstream_filter(st, "hevc_mp4toannexb", NULL);
+    } else if (st->codecpar->codec_id == AV_CODEC_ID_VP9) {
         ret = ff_stream_add_bitstream_filter(st, "vp9_superframe", NULL);
+    }
 
     return ret;
 }
 
 static const AVCodecTag codec_ivf_tags[] = {
+    { AV_CODEC_ID_AV1,  MKTAG('A', 'V', '0', '1') },
+    { AV_CODEC_ID_H264, MKTAG('H', '2', '6', '4') },
+    { AV_CODEC_ID_HEVC, MKTAG('H', 'E', 'V', 'C') },
     { AV_CODEC_ID_VP8,  MKTAG('V', 'P', '8', '0') },
     { AV_CODEC_ID_VP9,  MKTAG('V', 'P', '9', '0') },
-    { AV_CODEC_ID_AV1,  MKTAG('A', 'V', '0', '1') },
+
     { AV_CODEC_ID_NONE, 0 }
 };
 
-- 
2.19.0.605.g01d371f741-goog