git.sesse.net Git - ffmpeg/blob - libavcodec/hevc_parser.c

   1 /*
   2  * HEVC Annex B format parser
   3  *
   4  * Copyright (C) 2012 - 2013 Guillaume Martres
   5  *
   6  * This file is part of FFmpeg.
   7  *
   8  * FFmpeg is free software; you can redistribute it and/or
   9  * modify it under the terms of the GNU Lesser General Public
  10  * License as published by the Free Software Foundation; either
  11  * version 2.1 of the License, or (at your option) any later version.
  12  *
  13  * FFmpeg is distributed in the hope that it will be useful,
  14  * but WITHOUT ANY WARRANTY; without even the implied warranty of
  15  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  16  * Lesser General Public License for more details.
  17  *
  18  * You should have received a copy of the GNU Lesser General Public
  19  * License along with FFmpeg; if not, write to the Free Software
  20  * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
  21  */
  22
  23 #include "libavutil/common.h"
  24
  25 #include "golomb.h"
  26 #include "hevc.h"
  27 #include "hevcdec.h"
  28 #include "h2645_parse.h"
  29 #include "parser.h"
  30
  31 #define START_CODE 0x000001 ///< start_code_prefix_one_3bytes
  32
  33 #define IS_IRAP_NAL(nal) (nal->type >= 16 && nal->type <= 23)
  34
  35 #define ADVANCED_PARSER CONFIG_HEVC_DECODER
  36
  37 typedef struct HEVCParserContext {
  38     ParseContext pc;
  39
  40     H2645Packet pkt;
  41     HEVCParamSets ps;
  42
  43     int parsed_extradata;
  44
  45 #if ADVANCED_PARSER
  46     HEVCContext h;
  47 #endif
  48 } HEVCParserContext;
  49
  50 #if !ADVANCED_PARSER
  51 static int hevc_parse_slice_header(AVCodecParserContext *s, H2645NAL *nal,
  52                                    AVCodecContext *avctx)
  53 {
  54     HEVCParserContext *ctx = s->priv_data;
  55     GetBitContext *gb = &nal->gb;
  56
  57     HEVCPPS *pps;
  58     HEVCSPS *sps;
  59     unsigned int pps_id;
  60
  61     get_bits1(gb);          // first slice in pic
  62     if (IS_IRAP_NAL(nal))
  63         get_bits1(gb);      // no output of prior pics
  64
  65     pps_id = get_ue_golomb_long(gb);
  66     if (pps_id >= HEVC_MAX_PPS_COUNT || !ctx->ps.pps_list[pps_id]) {
  67         av_log(avctx, AV_LOG_ERROR, "PPS id out of range: %d\n", pps_id);
  68         return AVERROR_INVALIDDATA;
  69     }
  70     pps = (HEVCPPS*)ctx->ps.pps_list[pps_id]->data;
  71     sps = (HEVCSPS*)ctx->ps.sps_list[pps->sps_id]->data;
  72
  73     /* export the stream parameters */
  74     s->coded_width  = sps->width;
  75     s->coded_height = sps->height;
  76     s->width        = sps->output_width;
  77     s->height       = sps->output_height;
  78     s->format       = sps->pix_fmt;
  79     avctx->profile  = sps->ptl.general_ptl.profile_idc;
  80     avctx->level    = sps->ptl.general_ptl.level_idc;
  81
  82     /* ignore the rest for now*/
  83
  84     return 0;
  85 }
  86
  87 static int parse_nal_units(AVCodecParserContext *s, const uint8_t *buf,
  88                            int buf_size, AVCodecContext *avctx)
  89 {
  90     HEVCParserContext *ctx = s->priv_data;
  91     int ret, i;
  92
  93     ret = ff_h2645_packet_split(&ctx->pkt, buf, buf_size, avctx, 0, 0,
  94                                 AV_CODEC_ID_HEVC, 1);
  95     if (ret < 0)
  96         return ret;
  97
  98     for (i = 0; i < ctx->pkt.nb_nals; i++) {
  99         H2645NAL *nal = &ctx->pkt.nals[i];
 100
 101         /* ignore everything except parameter sets and VCL NALUs */
 102         switch (nal->type) {
 103         case HEVC_NAL_VPS: ff_hevc_decode_nal_vps(&nal->gb, avctx, &ctx->ps);    break;
 104         case HEVC_NAL_SPS: ff_hevc_decode_nal_sps(&nal->gb, avctx, &ctx->ps, 1); break;
 105         case HEVC_NAL_PPS: ff_hevc_decode_nal_pps(&nal->gb, avctx, &ctx->ps);    break;
 106         case HEVC_NAL_TRAIL_R:
 107         case HEVC_NAL_TRAIL_N:
 108         case HEVC_NAL_TSA_N:
 109         case HEVC_NAL_TSA_R:
 110         case HEVC_NAL_STSA_N:
 111         case HEVC_NAL_STSA_R:
 112         case HEVC_NAL_BLA_W_LP:
 113         case HEVC_NAL_BLA_W_RADL:
 114         case HEVC_NAL_BLA_N_LP:
 115         case HEVC_NAL_IDR_W_RADL:
 116         case HEVC_NAL_IDR_N_LP:
 117         case HEVC_NAL_CRA_NUT:
 118         case HEVC_NAL_RADL_N:
 119         case HEVC_NAL_RADL_R:
 120         case HEVC_NAL_RASL_N:
 121         case HEVC_NAL_RASL_R:
 122             if (buf == avctx->extradata) {
 123                 av_log(avctx, AV_LOG_ERROR, "Invalid NAL unit: %d\n", nal->type);
 124                 return AVERROR_INVALIDDATA;
 125             }
 126             hevc_parse_slice_header(s, nal, avctx);
 127             break;
 128         }
 129     }
 130
 131     return 0;
 132 }
 133 #endif
 134
 135 /**
 136  * Find the end of the current frame in the bitstream.
 137  * @return the position of the first byte of the next frame, or END_NOT_FOUND
 138  */
 139 static int hevc_find_frame_end(AVCodecParserContext *s, const uint8_t *buf,
 140                                int buf_size)
 141 {
 142     int i;
 143     ParseContext *pc = s->priv_data;
 144
 145     for (i = 0; i < buf_size; i++) {
 146         int nut;
 147
 148         pc->state64 = (pc->state64 << 8) | buf[i];
 149
 150         if (((pc->state64 >> 3 * 8) & 0xFFFFFF) != START_CODE)
 151             continue;
 152
 153         nut = (pc->state64 >> 2 * 8 + 1) & 0x3F;
 154         // Beginning of access unit
 155         if ((nut >= HEVC_NAL_VPS && nut <= HEVC_NAL_AUD) || nut == HEVC_NAL_SEI_PREFIX ||
 156             (nut >= 41 && nut <= 44) || (nut >= 48 && nut <= 55)) {
 157             if (pc->frame_start_found) {
 158                 pc->frame_start_found = 0;
 159                 return i - 5;
 160             }
 161         } else if (nut <= HEVC_NAL_RASL_R ||
 162                    (nut >= HEVC_NAL_BLA_W_LP && nut <= HEVC_NAL_CRA_NUT)) {
 163             int first_slice_segment_in_pic_flag = buf[i] >> 7;
 164             if (first_slice_segment_in_pic_flag) {
 165                 if (!pc->frame_start_found) {
 166                     pc->frame_start_found = 1;
 167                 } else { // First slice of next frame found
 168                     pc->frame_start_found = 0;
 169                     return i - 5;
 170                 }
 171             }
 172         }
 173     }
 174
 175     return END_NOT_FOUND;
 176 }
 177
 178 #if ADVANCED_PARSER
 179 /**
 180  * Parse NAL units of found picture and decode some basic information.
 181  *
 182  * @param s parser context.
 183  * @param avctx codec context.
 184  * @param buf buffer with field/frame data.
 185  * @param buf_size size of the buffer.
 186  */
 187 static inline int parse_nal_units(AVCodecParserContext *s, const uint8_t *buf,
 188                            int buf_size, AVCodecContext *avctx)
 189 {
 190     HEVCParserContext *ctx = s->priv_data;
 191     HEVCContext       *h   = &ctx->h;
 192     GetBitContext      *gb;
 193     SliceHeader        *sh = &h->sh;
 194     HEVCParamSets *ps = &h->ps;
 195     HEVCSEIContext *sei = &h->sei;
 196     H2645Packet   *pkt = &ctx->pkt;
 197     const uint8_t *buf_end = buf + buf_size;
 198     int state = -1, i;
 199     H2645NAL *nal;
 200     int is_global = buf == avctx->extradata;
 201
 202     if (!h->HEVClc)
 203         h->HEVClc = av_mallocz(sizeof(HEVCLocalContext));
 204     if (!h->HEVClc)
 205         return AVERROR(ENOMEM);
 206
 207     gb = &h->HEVClc->gb;
 208
 209     /* set some sane default values */
 210     s->pict_type         = AV_PICTURE_TYPE_I;
 211     s->key_frame         = 0;
 212     s->picture_structure = AV_PICTURE_STRUCTURE_UNKNOWN;
 213
 214     h->avctx = avctx;
 215
 216     ff_hevc_reset_sei(sei);
 217
 218     if (!buf_size)
 219         return 0;
 220
 221     if (pkt->nals_allocated < 1) {
 222         H2645NAL *tmp = av_realloc_array(pkt->nals, 1, sizeof(*tmp));
 223         if (!tmp)
 224             return AVERROR(ENOMEM);
 225         pkt->nals = tmp;
 226         memset(pkt->nals, 0, sizeof(*tmp));
 227         pkt->nals_allocated = 1;
 228     }
 229
 230     nal = &pkt->nals[0];
 231
 232     for (;;) {
 233         int src_length, consumed;
 234         int ret;
 235         int num = 0, den = 0;
 236         buf = avpriv_find_start_code(buf, buf_end, &state);
 237         if (--buf + 2 >= buf_end)
 238             break;
 239         src_length = buf_end - buf;
 240
 241         h->nal_unit_type = (*buf >> 1) & 0x3f;
 242         h->temporal_id   = (*(buf + 1) & 0x07) - 1;
 243         if (h->nal_unit_type <= HEVC_NAL_CRA_NUT) {
 244             // Do not walk the whole buffer just to decode slice segment header
 245             if (src_length > 20)
 246                 src_length = 20;
 247         }
 248
 249         consumed = ff_h2645_extract_rbsp(buf, src_length, nal, 1);
 250         if (consumed < 0)
 251             return consumed;
 252
 253         ret = init_get_bits8(gb, nal->data + 2, nal->size);
 254         if (ret < 0)
 255             return ret;
 256
 257         switch (h->nal_unit_type) {
 258         case HEVC_NAL_VPS:
 259             ff_hevc_decode_nal_vps(gb, avctx, ps);
 260             break;
 261         case HEVC_NAL_SPS:
 262             ff_hevc_decode_nal_sps(gb, avctx, ps, 1);
 263             break;
 264         case HEVC_NAL_PPS:
 265             ff_hevc_decode_nal_pps(gb, avctx, ps);
 266             break;
 267         case HEVC_NAL_SEI_PREFIX:
 268         case HEVC_NAL_SEI_SUFFIX:
 269             ff_hevc_decode_nal_sei(gb, avctx, sei, ps, h->nal_unit_type);
 270             break;
 271         case HEVC_NAL_TRAIL_N:
 272         case HEVC_NAL_TRAIL_R:
 273         case HEVC_NAL_TSA_N:
 274         case HEVC_NAL_TSA_R:
 275         case HEVC_NAL_STSA_N:
 276         case HEVC_NAL_STSA_R:
 277         case HEVC_NAL_RADL_N:
 278         case HEVC_NAL_RADL_R:
 279         case HEVC_NAL_RASL_N:
 280         case HEVC_NAL_RASL_R:
 281         case HEVC_NAL_BLA_W_LP:
 282         case HEVC_NAL_BLA_W_RADL:
 283         case HEVC_NAL_BLA_N_LP:
 284         case HEVC_NAL_IDR_W_RADL:
 285         case HEVC_NAL_IDR_N_LP:
 286         case HEVC_NAL_CRA_NUT:
 287
 288             if (is_global) {
 289                 av_log(avctx, AV_LOG_ERROR, "Invalid NAL unit: %d\n", h->nal_unit_type);
 290                 return AVERROR_INVALIDDATA;
 291             }
 292
 293             sh->first_slice_in_pic_flag = get_bits1(gb);
 294             s->picture_structure = h->sei.picture_timing.picture_struct;
 295             s->field_order = h->sei.picture_timing.picture_struct;
 296
 297             if (IS_IRAP(h)) {
 298                 s->key_frame = 1;
 299                 sh->no_output_of_prior_pics_flag = get_bits1(gb);
 300             }
 301
 302             sh->pps_id = get_ue_golomb(gb);
 303             if (sh->pps_id >= HEVC_MAX_PPS_COUNT || !ps->pps_list[sh->pps_id]) {
 304                 av_log(avctx, AV_LOG_ERROR, "PPS id out of range: %d\n", sh->pps_id);
 305                 return AVERROR_INVALIDDATA;
 306             }
 307             ps->pps = (HEVCPPS*)ps->pps_list[sh->pps_id]->data;
 308
 309             if (ps->pps->sps_id >= HEVC_MAX_SPS_COUNT || !ps->sps_list[ps->pps->sps_id]) {
 310                 av_log(avctx, AV_LOG_ERROR, "SPS id out of range: %d\n", ps->pps->sps_id);
 311                 return AVERROR_INVALIDDATA;
 312             }
 313             if (ps->sps != (HEVCSPS*)ps->sps_list[ps->pps->sps_id]->data) {
 314                 ps->sps = (HEVCSPS*)ps->sps_list[ps->pps->sps_id]->data;
 315                 ps->vps = (HEVCVPS*)ps->vps_list[ps->sps->vps_id]->data;
 316             }
 317
 318             s->coded_width  = ps->sps->width;
 319             s->coded_height = ps->sps->height;
 320             s->width        = ps->sps->output_width;
 321             s->height       = ps->sps->output_height;
 322             s->format       = ps->sps->pix_fmt;
 323             avctx->profile  = ps->sps->ptl.general_ptl.profile_idc;
 324             avctx->level    = ps->sps->ptl.general_ptl.level_idc;
 325
 326             if (ps->vps->vps_timing_info_present_flag) {
 327                 num = ps->vps->vps_num_units_in_tick;
 328                 den = ps->vps->vps_time_scale;
 329             } else if (ps->sps->vui.vui_timing_info_present_flag) {
 330                 num = ps->sps->vui.vui_num_units_in_tick;
 331                 den = ps->sps->vui.vui_time_scale;
 332             }
 333
 334             if (num != 0 && den != 0)
 335                 av_reduce(&avctx->framerate.den, &avctx->framerate.num,
 336                           num, den, 1 << 30);
 337
 338             if (!sh->first_slice_in_pic_flag) {
 339                 int slice_address_length;
 340
 341                 if (ps->pps->dependent_slice_segments_enabled_flag)
 342                     sh->dependent_slice_segment_flag = get_bits1(gb);
 343                 else
 344                     sh->dependent_slice_segment_flag = 0;
 345
 346                 slice_address_length = av_ceil_log2_c(ps->sps->ctb_width *
 347                                                       ps->sps->ctb_height);
 348                 sh->slice_segment_addr = get_bitsz(gb, slice_address_length);
 349                 if (sh->slice_segment_addr >= ps->sps->ctb_width * ps->sps->ctb_height) {
 350                     av_log(avctx, AV_LOG_ERROR, "Invalid slice segment address: %u.\n",
 351                            sh->slice_segment_addr);
 352                     return AVERROR_INVALIDDATA;
 353                 }
 354             } else
 355                 sh->dependent_slice_segment_flag = 0;
 356
 357             if (sh->dependent_slice_segment_flag)
 358                 break;
 359
 360             for (i = 0; i < ps->pps->num_extra_slice_header_bits; i++)
 361                 skip_bits(gb, 1); // slice_reserved_undetermined_flag[]
 362
 363             sh->slice_type = get_ue_golomb(gb);
 364             if (!(sh->slice_type == HEVC_SLICE_I || sh->slice_type == HEVC_SLICE_P ||
 365                   sh->slice_type == HEVC_SLICE_B)) {
 366                 av_log(avctx, AV_LOG_ERROR, "Unknown slice type: %d.\n",
 367                        sh->slice_type);
 368                 return AVERROR_INVALIDDATA;
 369             }
 370             s->pict_type = sh->slice_type == HEVC_SLICE_B ? AV_PICTURE_TYPE_B :
 371                            sh->slice_type == HEVC_SLICE_P ? AV_PICTURE_TYPE_P :
 372                                                        AV_PICTURE_TYPE_I;
 373
 374             if (ps->pps->output_flag_present_flag)
 375                 sh->pic_output_flag = get_bits1(gb);
 376
 377             if (ps->sps->separate_colour_plane_flag)
 378                 sh->colour_plane_id = get_bits(gb, 2);
 379
 380             if (!IS_IDR(h)) {
 381                 sh->pic_order_cnt_lsb = get_bits(gb, ps->sps->log2_max_poc_lsb);
 382                 s->output_picture_number = h->poc = ff_hevc_compute_poc(h, sh->pic_order_cnt_lsb);
 383             } else
 384                 s->output_picture_number = h->poc = 0;
 385
 386             if (h->temporal_id == 0 &&
 387                 h->nal_unit_type != HEVC_NAL_TRAIL_N &&
 388                 h->nal_unit_type != HEVC_NAL_TSA_N &&
 389                 h->nal_unit_type != HEVC_NAL_STSA_N &&
 390                 h->nal_unit_type != HEVC_NAL_RADL_N &&
 391                 h->nal_unit_type != HEVC_NAL_RASL_N &&
 392                 h->nal_unit_type != HEVC_NAL_RADL_R &&
 393                 h->nal_unit_type != HEVC_NAL_RASL_R)
 394                 h->pocTid0 = h->poc;
 395
 396             return 0; /* no need to evaluate the rest */
 397         }
 398         buf += consumed;
 399     }
 400     /* didn't find a picture! */
 401     if (!is_global)
 402         av_log(h->avctx, AV_LOG_ERROR, "missing picture in access unit\n");
 403     return -1;
 404 }
 405 #endif
 406
 407 static int hevc_parse(AVCodecParserContext *s,
 408                       AVCodecContext *avctx,
 409                       const uint8_t **poutbuf, int *poutbuf_size,
 410                       const uint8_t *buf, int buf_size)
 411 {
 412     int next;
 413     HEVCParserContext *ctx = s->priv_data;
 414     ParseContext *pc = &ctx->pc;
 415
 416     if (avctx->extradata && !ctx->parsed_extradata) {
 417         parse_nal_units(s, avctx->extradata, avctx->extradata_size, avctx);
 418         ctx->parsed_extradata = 1;
 419     }
 420
 421     if (s->flags & PARSER_FLAG_COMPLETE_FRAMES) {
 422         next = buf_size;
 423     } else {
 424         next = hevc_find_frame_end(s, buf, buf_size);
 425         if (ff_combine_frame(pc, next, &buf, &buf_size) < 0) {
 426             *poutbuf      = NULL;
 427             *poutbuf_size = 0;
 428             return buf_size;
 429         }
 430     }
 431
 432     parse_nal_units(s, buf, buf_size, avctx);
 433
 434     *poutbuf      = buf;
 435     *poutbuf_size = buf_size;
 436     return next;
 437 }
 438
 439 // Split after the parameter sets at the beginning of the stream if they exist.
 440 static int hevc_split(AVCodecContext *avctx, const uint8_t *buf, int buf_size)
 441 {
 442     const uint8_t *ptr = buf, *end = buf + buf_size;
 443     uint32_t state = -1;
 444     int has_vps = 0;
 445     int has_sps = 0;
 446     int has_pps = 0;
 447     int nut;
 448
 449     while (ptr < end) {
 450         ptr = avpriv_find_start_code(ptr, end, &state);
 451         if ((state >> 8) != START_CODE)
 452             break;
 453         nut = (state >> 1) & 0x3F;
 454         if (nut == HEVC_NAL_VPS)
 455             has_vps = 1;
 456         else if (nut == HEVC_NAL_SPS)
 457             has_sps = 1;
 458         else if (nut == HEVC_NAL_PPS)
 459             has_pps = 1;
 460         else if ((nut != HEVC_NAL_SEI_PREFIX || has_pps) &&
 461                   nut != HEVC_NAL_AUD) {
 462             if (has_vps && has_sps) {
 463                 while (ptr - 4 > buf && ptr[-5] == 0)
 464                     ptr--;
 465                 return ptr - 4 - buf;
 466             }
 467         }
 468     }
 469     return 0;
 470 }
 471
 472 static void hevc_parser_close(AVCodecParserContext *s)
 473 {
 474     HEVCParserContext *ctx = s->priv_data;
 475     int i;
 476
 477 #if ADVANCED_PARSER
 478     HEVCContext  *h  = &ctx->h;
 479
 480     for (i = 0; i < FF_ARRAY_ELEMS(h->ps.vps_list); i++)
 481         av_buffer_unref(&h->ps.vps_list[i]);
 482     for (i = 0; i < FF_ARRAY_ELEMS(h->ps.sps_list); i++)
 483         av_buffer_unref(&h->ps.sps_list[i]);
 484     for (i = 0; i < FF_ARRAY_ELEMS(h->ps.pps_list); i++)
 485         av_buffer_unref(&h->ps.pps_list[i]);
 486
 487     h->ps.sps = NULL;
 488
 489     av_freep(&h->HEVClc);
 490 #endif
 491
 492     for (i = 0; i < FF_ARRAY_ELEMS(ctx->ps.vps_list); i++)
 493         av_buffer_unref(&ctx->ps.vps_list[i]);
 494     for (i = 0; i < FF_ARRAY_ELEMS(ctx->ps.sps_list); i++)
 495         av_buffer_unref(&ctx->ps.sps_list[i]);
 496     for (i = 0; i < FF_ARRAY_ELEMS(ctx->ps.pps_list); i++)
 497         av_buffer_unref(&ctx->ps.pps_list[i]);
 498
 499     ctx->ps.sps = NULL;
 500
 501     ff_h2645_packet_uninit(&ctx->pkt);
 502
 503     av_freep(&ctx->pc.buffer);
 504 }
 505
 506 AVCodecParser ff_hevc_parser = {
 507     .codec_ids      = { AV_CODEC_ID_HEVC },
 508     .priv_data_size = sizeof(HEVCParserContext),
 509     .parser_parse   = hevc_parse,
 510     .parser_close   = hevc_parser_close,
 511     .split          = hevc_split,
 512 };