avcodec/hevc_parser: add missing call to ff_hevc_reset_sei()
[ffmpeg.git] / 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 "hevc_ps.h"
28 #include "hevc_sei.h"
29 #include "h2645_parse.h"
30 #include "internal.h"
31 #include "parser.h"
32
33 #define START_CODE 0x000001 ///< start_code_prefix_one_3bytes
34
35 #define IS_IRAP_NAL(nal) (nal->type >= 16 && nal->type <= 23)
36 #define IS_IDR_NAL(nal) (nal->type == HEVC_NAL_IDR_W_RADL || nal->type == HEVC_NAL_IDR_N_LP)
37
38 typedef struct HEVCParserContext {
39     ParseContext pc;
40
41     H2645Packet pkt;
42     HEVCParamSets ps;
43     HEVCSEIContext sei;
44     SliceHeader sh;
45
46     int parsed_extradata;
47
48     int poc;
49     int pocTid0;
50 } HEVCParserContext;
51
52 static int hevc_parse_slice_header(AVCodecParserContext *s, H2645NAL *nal,
53                                    AVCodecContext *avctx)
54 {
55     HEVCParserContext *ctx = s->priv_data;
56     HEVCParamSets *ps = &ctx->ps;
57     HEVCSEIContext *sei = &ctx->sei;
58     SliceHeader *sh = &ctx->sh;
59     GetBitContext *gb = &nal->gb;
60     int i, num = 0, den = 0;
61
62     sh->first_slice_in_pic_flag = get_bits1(gb);
63     s->picture_structure = sei->picture_timing.picture_struct;
64     s->field_order = sei->picture_timing.picture_struct;
65
66     if (IS_IRAP_NAL(nal)) {
67         s->key_frame = 1;
68         sh->no_output_of_prior_pics_flag = get_bits1(gb);
69     }
70
71     sh->pps_id = get_ue_golomb(gb);
72     if (sh->pps_id >= HEVC_MAX_PPS_COUNT || !ps->pps_list[sh->pps_id]) {
73         av_log(avctx, AV_LOG_ERROR, "PPS id out of range: %d\n", sh->pps_id);
74         return AVERROR_INVALIDDATA;
75     }
76     ps->pps = (HEVCPPS*)ps->pps_list[sh->pps_id]->data;
77
78     if (ps->pps->sps_id >= HEVC_MAX_SPS_COUNT || !ps->sps_list[ps->pps->sps_id]) {
79         av_log(avctx, AV_LOG_ERROR, "SPS id out of range: %d\n", ps->pps->sps_id);
80         return AVERROR_INVALIDDATA;
81     }
82     if (ps->sps != (HEVCSPS*)ps->sps_list[ps->pps->sps_id]->data) {
83         ps->sps = (HEVCSPS*)ps->sps_list[ps->pps->sps_id]->data;
84         ps->vps = (HEVCVPS*)ps->vps_list[ps->sps->vps_id]->data;
85     }
86
87     s->coded_width  = ps->sps->width;
88     s->coded_height = ps->sps->height;
89     s->width        = ps->sps->output_width;
90     s->height       = ps->sps->output_height;
91     s->format       = ps->sps->pix_fmt;
92     avctx->profile  = ps->sps->ptl.general_ptl.profile_idc;
93     avctx->level    = ps->sps->ptl.general_ptl.level_idc;
94
95     if (ps->vps->vps_timing_info_present_flag) {
96         num = ps->vps->vps_num_units_in_tick;
97         den = ps->vps->vps_time_scale;
98     } else if (ps->sps->vui.vui_timing_info_present_flag) {
99         num = ps->sps->vui.vui_num_units_in_tick;
100         den = ps->sps->vui.vui_time_scale;
101     }
102
103     if (num != 0 && den != 0)
104         av_reduce(&avctx->framerate.den, &avctx->framerate.num,
105                   num, den, 1 << 30);
106
107     if (!sh->first_slice_in_pic_flag) {
108         int slice_address_length;
109
110         if (ps->pps->dependent_slice_segments_enabled_flag)
111             sh->dependent_slice_segment_flag = get_bits1(gb);
112         else
113             sh->dependent_slice_segment_flag = 0;
114
115         slice_address_length = av_ceil_log2_c(ps->sps->ctb_width *
116                                               ps->sps->ctb_height);
117         sh->slice_segment_addr = get_bitsz(gb, slice_address_length);
118         if (sh->slice_segment_addr >= ps->sps->ctb_width * ps->sps->ctb_height) {
119             av_log(avctx, AV_LOG_ERROR, "Invalid slice segment address: %u.\n",
120                    sh->slice_segment_addr);
121             return AVERROR_INVALIDDATA;
122         }
123     } else
124         sh->dependent_slice_segment_flag = 0;
125
126     if (sh->dependent_slice_segment_flag)
127         return 0; /* break; */
128
129     for (i = 0; i < ps->pps->num_extra_slice_header_bits; i++)
130         skip_bits(gb, 1); // slice_reserved_undetermined_flag[]
131
132     sh->slice_type = get_ue_golomb(gb);
133     if (!(sh->slice_type == HEVC_SLICE_I || sh->slice_type == HEVC_SLICE_P ||
134           sh->slice_type == HEVC_SLICE_B)) {
135         av_log(avctx, AV_LOG_ERROR, "Unknown slice type: %d.\n",
136                sh->slice_type);
137         return AVERROR_INVALIDDATA;
138     }
139     s->pict_type = sh->slice_type == HEVC_SLICE_B ? AV_PICTURE_TYPE_B :
140                    sh->slice_type == HEVC_SLICE_P ? AV_PICTURE_TYPE_P :
141                                                AV_PICTURE_TYPE_I;
142
143     if (ps->pps->output_flag_present_flag)
144         sh->pic_output_flag = get_bits1(gb);
145
146     if (ps->sps->separate_colour_plane_flag)
147         sh->colour_plane_id = get_bits(gb, 2);
148
149     if (!IS_IDR_NAL(nal)) {
150         sh->pic_order_cnt_lsb = get_bits(gb, ps->sps->log2_max_poc_lsb);
151         s->output_picture_number = ctx->poc = ff_hevc_compute_poc(ps->sps, ctx->pocTid0, sh->pic_order_cnt_lsb, nal->type);
152     } else
153         s->output_picture_number = ctx->poc = 0;
154
155     if (nal->temporal_id == 0 &&
156         nal->type != HEVC_NAL_TRAIL_N &&
157         nal->type != HEVC_NAL_TSA_N &&
158         nal->type != HEVC_NAL_STSA_N &&
159         nal->type != HEVC_NAL_RADL_N &&
160         nal->type != HEVC_NAL_RASL_N &&
161         nal->type != HEVC_NAL_RADL_R &&
162         nal->type != HEVC_NAL_RASL_R)
163         ctx->pocTid0 = ctx->poc;
164
165     return 1; /* no need to evaluate the rest */
166 }
167
168 /**
169  * Parse NAL units of found picture and decode some basic information.
170  *
171  * @param s parser context.
172  * @param avctx codec context.
173  * @param buf buffer with field/frame data.
174  * @param buf_size size of the buffer.
175  */
176 static int parse_nal_units(AVCodecParserContext *s, const uint8_t *buf,
177                            int buf_size, AVCodecContext *avctx)
178 {
179     HEVCParserContext *ctx = s->priv_data;
180     HEVCParamSets *ps = &ctx->ps;
181     HEVCSEIContext *sei = &ctx->sei;
182     int is_global = buf == avctx->extradata;
183     int ret, i;
184
185     /* set some sane default values */
186     s->pict_type         = AV_PICTURE_TYPE_I;
187     s->key_frame         = 0;
188     s->picture_structure = AV_PICTURE_STRUCTURE_UNKNOWN;
189
190     ff_hevc_reset_sei(sei);
191
192     ret = ff_h2645_packet_split(&ctx->pkt, buf, buf_size, avctx, 0, 0,
193                                 AV_CODEC_ID_HEVC, 1);
194     if (ret < 0)
195         return ret;
196
197     for (i = 0; i < ctx->pkt.nb_nals; i++) {
198         H2645NAL *nal = &ctx->pkt.nals[i];
199         GetBitContext *gb = &nal->gb;
200
201         switch (nal->type) {
202         case HEVC_NAL_VPS:
203             ff_hevc_decode_nal_vps(gb, avctx, ps);
204             break;
205         case HEVC_NAL_SPS:
206             ff_hevc_decode_nal_sps(gb, avctx, ps, 1);
207             break;
208         case HEVC_NAL_PPS:
209             ff_hevc_decode_nal_pps(gb, avctx, ps);
210             break;
211         case HEVC_NAL_SEI_PREFIX:
212         case HEVC_NAL_SEI_SUFFIX:
213             ff_hevc_decode_nal_sei(gb, avctx, sei, ps, nal->type);
214             break;
215         case HEVC_NAL_TRAIL_N:
216         case HEVC_NAL_TRAIL_R:
217         case HEVC_NAL_TSA_N:
218         case HEVC_NAL_TSA_R:
219         case HEVC_NAL_STSA_N:
220         case HEVC_NAL_STSA_R:
221         case HEVC_NAL_BLA_W_LP:
222         case HEVC_NAL_BLA_W_RADL:
223         case HEVC_NAL_BLA_N_LP:
224         case HEVC_NAL_IDR_W_RADL:
225         case HEVC_NAL_IDR_N_LP:
226         case HEVC_NAL_CRA_NUT:
227         case HEVC_NAL_RADL_N:
228         case HEVC_NAL_RADL_R:
229         case HEVC_NAL_RASL_N:
230         case HEVC_NAL_RASL_R:
231
232             if (is_global) {
233                 av_log(avctx, AV_LOG_ERROR, "Invalid NAL unit: %d\n", nal->type);
234                 return AVERROR_INVALIDDATA;
235             }
236
237             ret = hevc_parse_slice_header(s, nal, avctx);
238             if (ret)
239                 return ret;
240             break;
241         }
242     }
243     /* didn't find a picture! */
244     if (!is_global)
245         av_log(avctx, AV_LOG_ERROR, "missing picture in access unit\n");
246     return -1;
247 }
248
249 /**
250  * Find the end of the current frame in the bitstream.
251  * @return the position of the first byte of the next frame, or END_NOT_FOUND
252  */
253 static int hevc_find_frame_end(AVCodecParserContext *s, const uint8_t *buf,
254                                int buf_size)
255 {
256     HEVCParserContext *ctx = s->priv_data;
257     ParseContext       *pc = &ctx->pc;
258     int i;
259
260     for (i = 0; i < buf_size; i++) {
261         int nut;
262
263         pc->state64 = (pc->state64 << 8) | buf[i];
264
265         if (((pc->state64 >> 3 * 8) & 0xFFFFFF) != START_CODE)
266             continue;
267
268         nut = (pc->state64 >> 2 * 8 + 1) & 0x3F;
269         // Beginning of access unit
270         if ((nut >= HEVC_NAL_VPS && nut <= HEVC_NAL_AUD) || nut == HEVC_NAL_SEI_PREFIX ||
271             (nut >= 41 && nut <= 44) || (nut >= 48 && nut <= 55)) {
272             if (pc->frame_start_found) {
273                 pc->frame_start_found = 0;
274                 return i - 5;
275             }
276         } else if (nut <= HEVC_NAL_RASL_R ||
277                    (nut >= HEVC_NAL_BLA_W_LP && nut <= HEVC_NAL_CRA_NUT)) {
278             int first_slice_segment_in_pic_flag = buf[i] >> 7;
279             if (first_slice_segment_in_pic_flag) {
280                 if (!pc->frame_start_found) {
281                     pc->frame_start_found = 1;
282                 } else { // First slice of next frame found
283                     pc->frame_start_found = 0;
284                     return i - 5;
285                 }
286             }
287         }
288     }
289
290     return END_NOT_FOUND;
291 }
292
293 static int hevc_parse(AVCodecParserContext *s, AVCodecContext *avctx,
294                       const uint8_t **poutbuf, int *poutbuf_size,
295                       const uint8_t *buf, int buf_size)
296 {
297     int next;
298     HEVCParserContext *ctx = s->priv_data;
299     ParseContext *pc = &ctx->pc;
300
301     if (avctx->extradata && !ctx->parsed_extradata) {
302         parse_nal_units(s, avctx->extradata, avctx->extradata_size, avctx);
303         ctx->parsed_extradata = 1;
304     }
305
306     if (s->flags & PARSER_FLAG_COMPLETE_FRAMES) {
307         next = buf_size;
308     } else {
309         next = hevc_find_frame_end(s, buf, buf_size);
310         if (ff_combine_frame(pc, next, &buf, &buf_size) < 0) {
311             *poutbuf      = NULL;
312             *poutbuf_size = 0;
313             return buf_size;
314         }
315     }
316
317     parse_nal_units(s, buf, buf_size, avctx);
318
319     *poutbuf      = buf;
320     *poutbuf_size = buf_size;
321     return next;
322 }
323
324 // Split after the parameter sets at the beginning of the stream if they exist.
325 static int hevc_split(AVCodecContext *avctx, const uint8_t *buf, int buf_size)
326 {
327     const uint8_t *ptr = buf, *end = buf + buf_size;
328     uint32_t state = -1;
329     int has_vps = 0;
330     int has_sps = 0;
331     int has_pps = 0;
332     int nut;
333
334     while (ptr < end) {
335         ptr = avpriv_find_start_code(ptr, end, &state);
336         if ((state >> 8) != START_CODE)
337             break;
338         nut = (state >> 1) & 0x3F;
339         if (nut == HEVC_NAL_VPS)
340             has_vps = 1;
341         else if (nut == HEVC_NAL_SPS)
342             has_sps = 1;
343         else if (nut == HEVC_NAL_PPS)
344             has_pps = 1;
345         else if ((nut != HEVC_NAL_SEI_PREFIX || has_pps) &&
346                   nut != HEVC_NAL_AUD) {
347             if (has_vps && has_sps) {
348                 while (ptr - 4 > buf && ptr[-5] == 0)
349                     ptr--;
350                 return ptr - 4 - buf;
351             }
352         }
353     }
354     return 0;
355 }
356
357 static void hevc_parser_close(AVCodecParserContext *s)
358 {
359     HEVCParserContext *ctx = s->priv_data;
360     int i;
361
362     for (i = 0; i < FF_ARRAY_ELEMS(ctx->ps.vps_list); i++)
363         av_buffer_unref(&ctx->ps.vps_list[i]);
364     for (i = 0; i < FF_ARRAY_ELEMS(ctx->ps.sps_list); i++)
365         av_buffer_unref(&ctx->ps.sps_list[i]);
366     for (i = 0; i < FF_ARRAY_ELEMS(ctx->ps.pps_list); i++)
367         av_buffer_unref(&ctx->ps.pps_list[i]);
368
369     ctx->ps.sps = NULL;
370
371     ff_h2645_packet_uninit(&ctx->pkt);
372     ff_hevc_reset_sei(&ctx->sei);
373
374     av_freep(&ctx->pc.buffer);
375 }
376
377 AVCodecParser ff_hevc_parser = {
378     .codec_ids      = { AV_CODEC_ID_HEVC },
379     .priv_data_size = sizeof(HEVCParserContext),
380     .parser_parse   = hevc_parse,
381     .parser_close   = hevc_parser_close,
382     .split          = hevc_split,
383 };