hevc_parser: parse and export some stream parameters

Particularly those that will be needed by the QSV decoder.
More can be added later as necessary.

1 files changed, 123 insertions, 5 deletions

diff --git a/libavcodec/hevc_parser.c b/libavcodec/hevc_parser.c
index ac2c6f52d9..030163e7e0 100644
--- a/libavcodec/hevc_parser.c
+++ b/libavcodec/hevc_parser.c
@@ -22,11 +22,99 @@
 
 #include "libavutil/common.h"
 
-#include "parser.h"
+#include "golomb.h"
 #include "hevc.h"
+#include "parser.h"
 
 #define START_CODE 0x000001 ///< start_code_prefix_one_3bytes
 
+#define IS_IRAP_NAL(nal) (nal->type >= 16 && nal->type <= 23)
+
+typedef struct HEVCParserContext {
+    ParseContext pc;
+
+    HEVCPacket pkt;
+    HEVCParamSets ps;
+
+    int parsed_extradata;
+} HEVCParserContext;
+
+static int hevc_parse_slice_header(AVCodecParserContext *s, HEVCNAL *nal,
+                                   AVCodecContext *avctx)
+{
+    HEVCParserContext *ctx = s->priv_data;
+    GetBitContext *gb = &nal->gb;
+
+    HEVCPPS *pps;
+    HEVCSPS *sps;
+    unsigned int pps_id;
+
+    get_bits1(gb);          // first slice in pic
+    if (IS_IRAP_NAL(nal))
+        get_bits1(gb);      // no output of prior pics
+
+    pps_id = get_ue_golomb_long(gb);
+    if (pps_id >= MAX_PPS_COUNT || !ctx->ps.pps_list[pps_id]) {
+        av_log(avctx, AV_LOG_ERROR, "PPS id out of range: %d\n", pps_id);
+        return AVERROR_INVALIDDATA;
+    }
+    pps = (HEVCPPS*)ctx->ps.pps_list[pps_id]->data;
+    sps = (HEVCSPS*)ctx->ps.sps_list[pps->sps_id]->data;
+
+    /* export the stream parameters */
+    s->coded_width  = sps->width;
+    s->coded_height = sps->height;
+    s->width        = sps->output_width;
+    s->height       = sps->output_height;
+    s->format       = sps->pix_fmt;
+    avctx->profile  = sps->ptl.general_ptl.profile_idc;
+    avctx->level    = sps->ptl.general_ptl.level_idc;
+
+    /* ignore the rest for now*/
+
+    return 0;
+}
+
+static int parse_nal_units(AVCodecParserContext *s, const uint8_t *buf,
+                           int buf_size, AVCodecContext *avctx)
+{
+    HEVCParserContext *ctx = s->priv_data;
+    int ret, i;
+
+    ret = ff_hevc_split_packet(&ctx->pkt, buf, buf_size, avctx, 0, 0);
+    if (ret < 0)
+        return ret;
+
+    for (i = 0; i < ctx->pkt.nb_nals; i++) {
+        HEVCNAL *nal = &ctx->pkt.nals[i];
+
+        /* ignore everything except parameter sets and VCL NALUs */
+        switch (nal->type) {
+        case NAL_VPS: ff_hevc_decode_nal_vps(&nal->gb, avctx, &ctx->ps);    break;
+        case NAL_SPS: ff_hevc_decode_nal_sps(&nal->gb, avctx, &ctx->ps, 1); break;
+        case NAL_PPS: ff_hevc_decode_nal_pps(&nal->gb, avctx, &ctx->ps);    break;
+        case NAL_TRAIL_R:
+        case NAL_TRAIL_N:
+        case NAL_TSA_N:
+        case NAL_TSA_R:
+        case NAL_STSA_N:
+        case NAL_STSA_R:
+        case NAL_BLA_W_LP:
+        case NAL_BLA_W_RADL:
+        case NAL_BLA_N_LP:
+        case NAL_IDR_W_RADL:
+        case NAL_IDR_N_LP:
+        case NAL_CRA_NUT:
+        case NAL_RADL_N:
+        case NAL_RADL_R:
+        case NAL_RASL_N:
+        case NAL_RASL_R: hevc_parse_slice_header(s, nal, avctx); break;
+        }
+    }
+
+    return 0;
+}
+
 /**
  * Find the end of the current frame in the bitstream.
  * @return the position of the first byte of the next frame, or END_NOT_FOUND
@@ -34,8 +122,9 @@
 static int hevc_find_frame_end(AVCodecParserContext *s, const uint8_t *buf,
                                int buf_size)
 {
+    HEVCParserContext *ctx = s->priv_data;
+    ParseContext       *pc = &ctx->pc;
     int i;
-    ParseContext *pc = s->priv_data;
 
     for (i = 0; i < buf_size; i++) {
         int nut;
@@ -76,7 +165,14 @@ static int hevc_parse(AVCodecParserContext *s, AVCodecContext *avctx,
                       const uint8_t *buf, int buf_size)
 {
     int next;
-    ParseContext *pc = s->priv_data;
+
+    HEVCParserContext *ctx = s->priv_data;
+    ParseContext *pc = &ctx->pc;
+
+    if (avctx->extradata && !ctx->parsed_extradata) {
+        parse_nal_units(s, avctx->extradata, avctx->extradata_size, avctx);
+        ctx->parsed_extradata = 1;
+    }
 
     if (s->flags & PARSER_FLAG_COMPLETE_FRAMES) {
         next = buf_size;
@@ -89,6 +185,8 @@ static int hevc_parse(AVCodecParserContext *s, AVCodecContext *avctx,
         }
     }
 
+    parse_nal_units(s, buf, buf_size, avctx);
+
     *poutbuf      = buf;
     *poutbuf_size = buf_size;
     return next;
@@ -116,10 +214,30 @@ static int hevc_split(AVCodecContext *avctx, const uint8_t *buf, int buf_size)
     return 0;
 }
 
+static void hevc_parser_close(AVCodecParserContext *s)
+{
+    HEVCParserContext *ctx = s->priv_data;
+    int i;
+
+    for (i = 0; i < FF_ARRAY_ELEMS(ctx->ps.vps_list); i++)
+        av_buffer_unref(&ctx->ps.vps_list[i]);
+    for (i = 0; i < FF_ARRAY_ELEMS(ctx->ps.sps_list); i++)
+        av_buffer_unref(&ctx->ps.sps_list[i]);
+    for (i = 0; i < FF_ARRAY_ELEMS(ctx->ps.pps_list); i++)
+        av_buffer_unref(&ctx->ps.pps_list[i]);
+
+    for (i = 0; i < ctx->pkt.nals_allocated; i++)
+        av_freep(&ctx->pkt.nals[i].rbsp_buffer);
+    av_freep(&ctx->pkt.nals);
+    ctx->pkt.nals_allocated = 0;
+
+    av_freep(&ctx->pc.buffer);
+}
+
 AVCodecParser ff_hevc_parser = {
     .codec_ids      = { AV_CODEC_ID_HEVC },
-    .priv_data_size = sizeof(ParseContext),
+    .priv_data_size = sizeof(HEVCParserContext),
     .parser_parse   = hevc_parse,
-    .parser_close   = ff_parse_close,
+    .parser_close   = hevc_parser_close,
     .split          = hevc_split,
 };