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