movenc: enable Annex B to MP4 conversion for HEVC tracks.

3 files changed, 158 insertions, 0 deletions

diff --git a/libavformat/hevc.c b/libavformat/hevc.c
index ab4eb10a8b..2c1b4c492d 100644
--- a/libavformat/hevc.c
+++ b/libavformat/hevc.c
@@ -1014,6 +1014,107 @@ static int hvcc_write(AVIOContext *pb, HEVCDecoderConfigurationRecord *hvcc)
     return 0;
 }
 
+int ff_hevc_annexb2mp4(AVIOContext *pb, const uint8_t *buf_in,
+                       int size, int filter_ps, int *ps_count)
+{
+    int num_ps = 0, ret = 0;
+    uint8_t *buf, *end, *start = NULL;
+
+    if (!filter_ps) {
+        ret = ff_avc_parse_nal_units(pb, buf_in, size);
+        goto end;
+    }
+
+    ret = ff_avc_parse_nal_units_buf(buf_in, &start, &size);
+    if (ret < 0)
+        goto end;
+
+    ret = 0;
+    buf = start;
+    end = start + size;
+
+    while (end - buf > 4) {
+        uint32_t len = FFMIN(AV_RB32(buf), end - buf - 4);
+        uint8_t type = (buf[4] >> 1) & 0x3f;
+
+        buf += 4;
+
+        switch (type) {
+        case NAL_VPS:
+        case NAL_SPS:
+        case NAL_PPS:
+            num_ps++;
+            break;
+        default:
+            ret += 4 + len;
+            avio_wb32(pb, len);
+            avio_write(pb, buf, len);
+            break;
+        }
+
+        buf += len;
+    }
+
+end:
+    free(start);
+    if (ps_count)
+        *ps_count = num_ps;
+    return ret;
+}
+
+int ff_hevc_annexb2mp4_buf(const uint8_t *buf_in, uint8_t **buf_out,
+                           int *size, int filter_ps, int *ps_count)
+{
+    AVIOContext *pb;
+    int num_ps = 0, ret = 0;
+    uint8_t *buf, *end, *start = NULL;
+
+    if (!filter_ps) {
+        ret = ff_avc_parse_nal_units_buf(buf_in, buf_out, size);
+        goto end;
+    }
+
+    ret = avio_open_dyn_buf(&pb);
+    if (ret < 0)
+        goto end;
+
+    ret = ff_avc_parse_nal_units_buf(buf_in, &start, size);
+    if (ret < 0)
+        goto end;
+
+    buf = start;
+    end = start + *size;
+
+    while (end - buf > 4) {
+        uint32_t len = FFMIN(AV_RB32(buf), end - buf - 4);
+        uint8_t type = (buf[4] >> 1) & 0x3f;
+
+        buf += 4;
+
+        switch (type) {
+        case NAL_VPS:
+        case NAL_SPS:
+        case NAL_PPS:
+            num_ps++;
+            break;
+        default:
+            avio_wb32(pb, len);
+            avio_write(pb, buf, len);
+            break;
+        }
+
+        buf += len;
+    }
+
+    *size = avio_close_dyn_buf(pb, buf_out);
+
+end:
+    free(start);
+    if (ps_count)
+        *ps_count = num_ps;
+    return ret;
+}
+
 int ff_isom_write_hvcc(AVIOContext *pb, const uint8_t *data,
                        int size, int ps_array_completeness)
 {
diff --git a/libavformat/hevc.h b/libavformat/hevc.h
index 82525ac84b..f394342028 100644
--- a/libavformat/hevc.h
+++ b/libavformat/hevc.h
@@ -30,6 +30,54 @@
 #include "avio.h"
 
 /**
+ * Writes Annex B formatted HEVC NAL units to the provided AVIOContext.
+ *
+ * The NAL units are converted to an MP4-compatible format (start code prefixes
+ * are replaced by 4-byte size fields, as per ISO/IEC 14496-15).
+ *
+ * If filter_ps is non-zero, any HEVC parameter sets found in the input will be
+ * discarded, and *ps_count will be set to the number of discarded PS NAL units.
+ *
+ * @param pb address of the AVIOContext where the data shall be written
+ * @param buf_in address of the buffer holding the input data
+ * @param size size (in bytes) of the input buffer
+ * @param filter_ps whether to write parameter set NAL units to the output (0)
+ *        or to discard them (non-zero)
+ * @param ps_count address of the variable where the number of discarded
+ *        parameter set NAL units shall be written, may be NULL
+ * @return the amount (in bytes) of data written in case of success, a negative
+ *         value corresponding to an AVERROR code in case of failure
+ */
+int ff_hevc_annexb2mp4(AVIOContext *pb, const uint8_t *buf_in,
+                       int size, int filter_ps, int *ps_count);
+
+/**
+ * Writes Annex B formatted HEVC NAL units to a data buffer.
+ *
+ * The NAL units are converted to an MP4-compatible format (start code prefixes
+ * are replaced by 4-byte size fields, as per ISO/IEC 14496-15).
+ *
+ * If filter_ps is non-zero, any HEVC parameter sets found in the input will be
+ * discarded, and *ps_count will be set to the number of discarded PS NAL units.
+ *
+ * On output, *size holds the size (in bytes) of the output data buffer.
+ *
+ * @param buf_in address of the buffer holding the input data
+ * @param size address of the variable holding the size (in bytes) of the input
+ *        buffer (on input) and of the output buffer (on output)
+ * @param buf_out address of the variable holding the address of the output
+ *        buffer
+ * @param filter_ps whether to write parameter set NAL units to the output (0)
+ *        or to discard them (non-zero)
+ * @param ps_count address of the variable where the number of discarded
+ *        parameter set NAL units shall be written, may be NULL
+ * @return 0 in case of success, a negative value corresponding to an AVERROR
+ *         code in case of failure
+ */
+int ff_hevc_annexb2mp4_buf(const uint8_t *buf_in, uint8_t **buf_out,
+                           int *size, int filter_ps, int *ps_count);
+
+/**
  * Writes HEVC extradata (parameter sets, declarative SEI NAL units) to the
  * provided AVIOContext.
  *
diff --git a/libavformat/movenc.c b/libavformat/movenc.c
index 641e1c7f48..6218259664 100644
--- a/libavformat/movenc.c
+++ b/libavformat/movenc.c
@@ -2939,6 +2939,15 @@ int ff_mov_write_packet(AVFormatContext *s, AVPacket *pkt)
         } else {
             size = ff_avc_parse_nal_units(pb, pkt->data, pkt->size);
         }
+    } else if (enc->codec_id == AV_CODEC_ID_HEVC && trk->vos_len > 6 &&
+               (AV_RB24(trk->vos_data) == 1 || AV_RB32(trk->vos_data) == 1)) {
+        /* extradata is Annex B, assume the bitstream is too and convert it */
+        if (trk->hint_track >= 0 && trk->hint_track < mov->nb_streams) {
+            ff_hevc_annexb2mp4_buf(pkt->data, &reformatted_data, &size, 0, NULL);
+            avio_write(pb, reformatted_data, size);
+        } else {
+            size = ff_hevc_annexb2mp4(pb, pkt->data, pkt->size, 0, NULL);
+        }
     } else {
         avio_write(pb, pkt->data, size);
     }