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