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     H2645Packet   *pkt = &ctx->pkt;
 196     const uint8_t *buf_end = buf + buf_size;
 197     int state = -1, i;
 198     H2645NAL *nal;
 199     int is_global = buf == avctx->extradata;
 200
 201     if (!h->HEVClc)
 202         h->HEVClc = av_mallocz(sizeof(HEVCLocalContext));
 203     if (!h->HEVClc)
 204         return AVERROR(ENOMEM);
 205
 206     gb = &h->HEVClc->gb;
 207
 208     /* set some sane default values */
 209     s->pict_type         = AV_PICTURE_TYPE_I;
 210     s->key_frame         = 0;
 211     s->picture_structure = AV_PICTURE_STRUCTURE_UNKNOWN;
 212
 213     h->avctx = avctx;
 214
 215     ff_hevc_reset_sei(h);
 216
 217     if (!buf_size)
 218         return 0;
 219
 220     if (pkt->nals_allocated < 1) {
 221         H2645NAL *tmp = av_realloc_array(pkt->nals, 1, sizeof(*tmp));
 222         if (!tmp)
 223             return AVERROR(ENOMEM);
 224         pkt->nals = tmp;
 225         memset(pkt->nals, 0, sizeof(*tmp));
 226         pkt->nals_allocated = 1;
 227     }
 228
 229     nal = &pkt->nals[0];
 230
 231     for (;;) {
 232         int src_length, consumed;
 233         int ret;
 234         int num = 0, den = 0;
 235         buf = avpriv_find_start_code(buf, buf_end, &state);
 236         if (--buf + 2 >= buf_end)
 237             break;
 238         src_length = buf_end - buf;
 239
 240         h->nal_unit_type = (*buf >> 1) & 0x3f;
 241         h->temporal_id   = (*(buf + 1) & 0x07) - 1;
 242         if (h->nal_unit_type <= HEVC_NAL_CRA_NUT) {
 243             // Do not walk the whole buffer just to decode slice segment header
 244             if (src_length > 20)
 245                 src_length = 20;
 246         }
 247
 248         consumed = ff_h2645_extract_rbsp(buf, src_length, nal, 1);
 249         if (consumed < 0)
 250             return consumed;
 251
 252         ret = init_get_bits8(gb, nal->data + 2, nal->size);
 253         if (ret < 0)
 254             return ret;
 255
 256         switch (h->nal_unit_type) {
 257         case HEVC_NAL_VPS:
 258             ff_hevc_decode_nal_vps(gb, avctx, ps);
 259             break;
 260         case HEVC_NAL_SPS:
 261             ff_hevc_decode_nal_sps(gb, avctx, ps, 1);
 262             break;
 263         case HEVC_NAL_PPS:
 264             ff_hevc_decode_nal_pps(gb, avctx, ps);
 265             break;
 266         case HEVC_NAL_SEI_PREFIX:
 267         case HEVC_NAL_SEI_SUFFIX:
 268             ff_hevc_decode_nal_sei(h);
 269             break;
 270         case HEVC_NAL_TRAIL_N:
 271         case HEVC_NAL_TRAIL_R:
 272         case HEVC_NAL_TSA_N:
 273         case HEVC_NAL_TSA_R:
 274         case HEVC_NAL_STSA_N:
 275         case HEVC_NAL_STSA_R:
 276         case HEVC_NAL_RADL_N:
 277         case HEVC_NAL_RADL_R:
 278         case HEVC_NAL_RASL_N:
 279         case HEVC_NAL_RASL_R:
 280         case HEVC_NAL_BLA_W_LP:
 281         case HEVC_NAL_BLA_W_RADL:
 282         case HEVC_NAL_BLA_N_LP:
 283         case HEVC_NAL_IDR_W_RADL:
 284         case HEVC_NAL_IDR_N_LP:
 285         case HEVC_NAL_CRA_NUT:
 286
 287             if (is_global) {
 288                 av_log(avctx, AV_LOG_ERROR, "Invalid NAL unit: %d\n", h->nal_unit_type);
 289                 return AVERROR_INVALIDDATA;
 290             }
 291
 292             sh->first_slice_in_pic_flag = get_bits1(gb);
 293             s->picture_structure = h->picture_struct;
 294             s->field_order = h->picture_struct;
 295
 296             if (IS_IRAP(h)) {
 297                 s->key_frame = 1;
 298                 sh->no_output_of_prior_pics_flag = get_bits1(gb);
 299             }
 300
 301             sh->pps_id = get_ue_golomb(gb);
 302             if (sh->pps_id >= HEVC_MAX_PPS_COUNT || !ps->pps_list[sh->pps_id]) {
 303                 av_log(avctx, AV_LOG_ERROR, "PPS id out of range: %d\n", sh->pps_id);
 304                 return AVERROR_INVALIDDATA;
 305             }
 306             ps->pps = (HEVCPPS*)ps->pps_list[sh->pps_id]->data;
 307
 308             if (ps->pps->sps_id >= HEVC_MAX_SPS_COUNT || !ps->sps_list[ps->pps->sps_id]) {
 309                 av_log(avctx, AV_LOG_ERROR, "SPS id out of range: %d\n", ps->pps->sps_id);
 310                 return AVERROR_INVALIDDATA;
 311             }
 312             if (ps->sps != (HEVCSPS*)ps->sps_list[ps->pps->sps_id]->data) {
 313                 ps->sps = (HEVCSPS*)ps->sps_list[ps->pps->sps_id]->data;
 314                 ps->vps = (HEVCVPS*)ps->vps_list[ps->sps->vps_id]->data;
 315             }
 316
 317             s->coded_width  = ps->sps->width;
 318             s->coded_height = ps->sps->height;
 319             s->width        = ps->sps->output_width;
 320             s->height       = ps->sps->output_height;
 321             s->format       = ps->sps->pix_fmt;
 322             avctx->profile  = ps->sps->ptl.general_ptl.profile_idc;
 323             avctx->level    = ps->sps->ptl.general_ptl.level_idc;
 324
 325             if (ps->vps->vps_timing_info_present_flag) {
 326                 num = ps->vps->vps_num_units_in_tick;
 327                 den = ps->vps->vps_time_scale;
 328             } else if (ps->sps->vui.vui_timing_info_present_flag) {
 329                 num = ps->sps->vui.vui_num_units_in_tick;
 330                 den = ps->sps->vui.vui_time_scale;
 331             }
 332
 333             if (num != 0 && den != 0)
 334                 av_reduce(&avctx->framerate.den, &avctx->framerate.num,
 335                           num, den, 1 << 30);
 336
 337             if (!sh->first_slice_in_pic_flag) {
 338                 int slice_address_length;
 339
 340                 if (ps->pps->dependent_slice_segments_enabled_flag)
 341                     sh->dependent_slice_segment_flag = get_bits1(gb);
 342                 else
 343                     sh->dependent_slice_segment_flag = 0;
 344
 345                 slice_address_length = av_ceil_log2_c(ps->sps->ctb_width *
 346                                                       ps->sps->ctb_height);
 347                 sh->slice_segment_addr = get_bitsz(gb, slice_address_length);
 348                 if (sh->slice_segment_addr >= ps->sps->ctb_width * ps->sps->ctb_height) {
 349                     av_log(avctx, AV_LOG_ERROR, "Invalid slice segment address: %u.\n",
 350                            sh->slice_segment_addr);
 351                     return AVERROR_INVALIDDATA;
 352                 }
 353             } else
 354                 sh->dependent_slice_segment_flag = 0;
 355
 356             if (sh->dependent_slice_segment_flag)
 357                 break;
 358
 359             for (i = 0; i < ps->pps->num_extra_slice_header_bits; i++)
 360                 skip_bits(gb, 1); // slice_reserved_undetermined_flag[]
 361
 362             sh->slice_type = get_ue_golomb(gb);
 363             if (!(sh->slice_type == HEVC_SLICE_I || sh->slice_type == HEVC_SLICE_P ||
 364                   sh->slice_type == HEVC_SLICE_B)) {
 365                 av_log(avctx, AV_LOG_ERROR, "Unknown slice type: %d.\n",
 366                        sh->slice_type);
 367                 return AVERROR_INVALIDDATA;
 368             }
 369             s->pict_type = sh->slice_type == HEVC_SLICE_B ? AV_PICTURE_TYPE_B :
 370                            sh->slice_type == HEVC_SLICE_P ? AV_PICTURE_TYPE_P :
 371                                                        AV_PICTURE_TYPE_I;
 372
 373             if (ps->pps->output_flag_present_flag)
 374                 sh->pic_output_flag = get_bits1(gb);
 375
 376             if (ps->sps->separate_colour_plane_flag)
 377                 sh->colour_plane_id = get_bits(gb, 2);
 378
 379             if (!IS_IDR(h)) {
 380                 sh->pic_order_cnt_lsb = get_bits(gb, ps->sps->log2_max_poc_lsb);
 381                 s->output_picture_number = h->poc = ff_hevc_compute_poc(h, sh->pic_order_cnt_lsb);
 382             } else
 383                 s->output_picture_number = h->poc = 0;
 384
 385             if (h->temporal_id == 0 &&
 386                 h->nal_unit_type != HEVC_NAL_TRAIL_N &&
 387                 h->nal_unit_type != HEVC_NAL_TSA_N &&
 388                 h->nal_unit_type != HEVC_NAL_STSA_N &&
 389                 h->nal_unit_type != HEVC_NAL_RADL_N &&
 390                 h->nal_unit_type != HEVC_NAL_RASL_N &&
 391                 h->nal_unit_type != HEVC_NAL_RADL_R &&
 392                 h->nal_unit_type != HEVC_NAL_RASL_R)
 393                 h->pocTid0 = h->poc;
 394
 395             return 0; /* no need to evaluate the rest */
 396         }
 397         buf += consumed;
 398     }
 399     /* didn't find a picture! */
 400     if (!is_global)
 401         av_log(h->avctx, AV_LOG_ERROR, "missing picture in access unit\n");
 402     return -1;
 403 }
 404 #endif
 405
 406 static int hevc_parse(AVCodecParserContext *s,
 407                       AVCodecContext *avctx,
 408                       const uint8_t **poutbuf, int *poutbuf_size,
 409                       const uint8_t *buf, int buf_size)
 410 {
 411     int next;
 412     HEVCParserContext *ctx = s->priv_data;
 413     ParseContext *pc = &ctx->pc;
 414
 415     if (avctx->extradata && !ctx->parsed_extradata) {
 416         parse_nal_units(s, avctx->extradata, avctx->extradata_size, avctx);
 417         ctx->parsed_extradata = 1;
 418     }
 419
 420     if (s->flags & PARSER_FLAG_COMPLETE_FRAMES) {
 421         next = buf_size;
 422     } else {
 423         next = hevc_find_frame_end(s, buf, buf_size);
 424         if (ff_combine_frame(pc, next, &buf, &buf_size) < 0) {
 425             *poutbuf      = NULL;
 426             *poutbuf_size = 0;
 427             return buf_size;
 428         }
 429     }
 430
 431     parse_nal_units(s, buf, buf_size, avctx);
 432
 433     *poutbuf      = buf;
 434     *poutbuf_size = buf_size;
 435     return next;
 436 }
 437
 438 // Split after the parameter sets at the beginning of the stream if they exist.
 439 static int hevc_split(AVCodecContext *avctx, const uint8_t *buf, int buf_size)
 440 {
 441     const uint8_t *ptr = buf, *end = buf + buf_size;
 442     uint32_t state = -1;
 443     int has_vps = 0;
 444     int has_sps = 0;
 445     int has_pps = 0;
 446     int nut;
 447
 448     while (ptr < end) {
 449         ptr = avpriv_find_start_code(ptr, end, &state);
 450         if ((state >> 8) != START_CODE)
 451             break;
 452         nut = (state >> 1) & 0x3F;
 453         if (nut == HEVC_NAL_VPS)
 454             has_vps = 1;
 455         else if (nut == HEVC_NAL_SPS)
 456             has_sps = 1;
 457         else if (nut == HEVC_NAL_PPS)
 458             has_pps = 1;
 459         else if ((nut != HEVC_NAL_SEI_PREFIX || has_pps) &&
 460                   nut != HEVC_NAL_AUD) {
 461             if (has_vps && has_sps) {
 462                 while (ptr - 4 > buf && ptr[-5] == 0)
 463                     ptr--;
 464                 return ptr - 4 - buf;
 465             }
 466         }
 467     }
 468     return 0;
 469 }
 470
 471 static void hevc_parser_close(AVCodecParserContext *s)
 472 {
 473     HEVCParserContext *ctx = s->priv_data;
 474     int i;
 475
 476 #if ADVANCED_PARSER
 477     HEVCContext  *h  = &ctx->h;
 478
 479     for (i = 0; i < FF_ARRAY_ELEMS(h->ps.vps_list); i++)
 480         av_buffer_unref(&h->ps.vps_list[i]);
 481     for (i = 0; i < FF_ARRAY_ELEMS(h->ps.sps_list); i++)
 482         av_buffer_unref(&h->ps.sps_list[i]);
 483     for (i = 0; i < FF_ARRAY_ELEMS(h->ps.pps_list); i++)
 484         av_buffer_unref(&h->ps.pps_list[i]);
 485
 486     h->ps.sps = NULL;
 487
 488     av_freep(&h->HEVClc);
 489 #endif
 490
 491     for (i = 0; i < FF_ARRAY_ELEMS(ctx->ps.vps_list); i++)
 492         av_buffer_unref(&ctx->ps.vps_list[i]);
 493     for (i = 0; i < FF_ARRAY_ELEMS(ctx->ps.sps_list); i++)
 494         av_buffer_unref(&ctx->ps.sps_list[i]);
 495     for (i = 0; i < FF_ARRAY_ELEMS(ctx->ps.pps_list); i++)
 496         av_buffer_unref(&ctx->ps.pps_list[i]);
 497
 498     ctx->ps.sps = NULL;
 499
 500     ff_h2645_packet_uninit(&ctx->pkt);
 501
 502     av_freep(&ctx->pc.buffer);
 503 }
 504
 505 AVCodecParser ff_hevc_parser = {
 506     .codec_ids      = { AV_CODEC_ID_HEVC },
 507     .priv_data_size = sizeof(HEVCParserContext),
 508     .parser_parse   = hevc_parse,
 509     .parser_close   = hevc_parser_close,
 510     .split          = hevc_split,
 511 };