Commit | Line | Data |
---|---|---|
2ba45a60 DM |
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 "parser.h" | |
26 | #include "hevc.h" | |
27 | #include "golomb.h" | |
28 | ||
29 | #define START_CODE 0x000001 ///< start_code_prefix_one_3bytes | |
30 | ||
31 | typedef struct HEVCParseContext { | |
32 | HEVCContext h; | |
33 | ParseContext pc; | |
34 | } HEVCParseContext; | |
35 | ||
36 | /** | |
37 | * Find the end of the current frame in the bitstream. | |
38 | * @return the position of the first byte of the next frame, or END_NOT_FOUND | |
39 | */ | |
40 | static int hevc_find_frame_end(AVCodecParserContext *s, const uint8_t *buf, | |
41 | int buf_size) | |
42 | { | |
43 | int i; | |
44 | ParseContext *pc = &((HEVCParseContext *)s->priv_data)->pc; | |
45 | ||
46 | for (i = 0; i < buf_size; i++) { | |
47 | int nut; | |
48 | ||
49 | pc->state64 = (pc->state64 << 8) | buf[i]; | |
50 | ||
51 | if (((pc->state64 >> 3 * 8) & 0xFFFFFF) != START_CODE) | |
52 | continue; | |
53 | ||
54 | nut = (pc->state64 >> 2 * 8 + 1) & 0x3F; | |
55 | // Beginning of access unit | |
56 | if ((nut >= NAL_VPS && nut <= NAL_AUD) || nut == NAL_SEI_PREFIX || | |
57 | (nut >= 41 && nut <= 44) || (nut >= 48 && nut <= 55)) { | |
58 | if (pc->frame_start_found) { | |
59 | pc->frame_start_found = 0; | |
60 | return i - 5; | |
61 | } | |
62 | } else if (nut <= NAL_RASL_R || | |
63 | (nut >= NAL_BLA_W_LP && nut <= NAL_CRA_NUT)) { | |
64 | int first_slice_segment_in_pic_flag = buf[i] >> 7; | |
65 | if (first_slice_segment_in_pic_flag) { | |
66 | if (!pc->frame_start_found) { | |
67 | pc->frame_start_found = 1; | |
68 | } else { // First slice of next frame found | |
69 | pc->frame_start_found = 0; | |
70 | return i - 5; | |
71 | } | |
72 | } | |
73 | } | |
74 | } | |
75 | ||
76 | return END_NOT_FOUND; | |
77 | } | |
78 | ||
79 | /** | |
80 | * Parse NAL units of found picture and decode some basic information. | |
81 | * | |
82 | * @param s parser context. | |
83 | * @param avctx codec context. | |
84 | * @param buf buffer with field/frame data. | |
85 | * @param buf_size size of the buffer. | |
86 | */ | |
87 | static inline int parse_nal_units(AVCodecParserContext *s, AVCodecContext *avctx, | |
88 | const uint8_t *buf, int buf_size) | |
89 | { | |
90 | HEVCContext *h = &((HEVCParseContext *)s->priv_data)->h; | |
91 | GetBitContext *gb = &h->HEVClc->gb; | |
92 | SliceHeader *sh = &h->sh; | |
93 | const uint8_t *buf_end = buf + buf_size; | |
94 | int state = -1, i; | |
95 | HEVCNAL *nal; | |
96 | ||
97 | /* set some sane default values */ | |
98 | s->pict_type = AV_PICTURE_TYPE_I; | |
99 | s->key_frame = 0; | |
100 | s->picture_structure = AV_PICTURE_STRUCTURE_UNKNOWN; | |
101 | ||
102 | h->avctx = avctx; | |
103 | ||
104 | if (!buf_size) | |
105 | return 0; | |
106 | ||
107 | if (h->nals_allocated < 1) { | |
108 | HEVCNAL *tmp = av_realloc_array(h->nals, 1, sizeof(*tmp)); | |
109 | if (!tmp) | |
110 | return AVERROR(ENOMEM); | |
111 | h->nals = tmp; | |
112 | memset(h->nals, 0, sizeof(*tmp)); | |
113 | h->nals_allocated = 1; | |
114 | } | |
115 | ||
116 | nal = &h->nals[0]; | |
117 | ||
118 | for (;;) { | |
119 | int src_length, consumed; | |
120 | buf = avpriv_find_start_code(buf, buf_end, &state); | |
121 | if (--buf + 2 >= buf_end) | |
122 | break; | |
123 | src_length = buf_end - buf; | |
124 | ||
125 | h->nal_unit_type = (*buf >> 1) & 0x3f; | |
126 | h->temporal_id = (*(buf + 1) & 0x07) - 1; | |
127 | if (h->nal_unit_type <= NAL_CRA_NUT) { | |
128 | // Do not walk the whole buffer just to decode slice segment header | |
129 | if (src_length > 20) | |
130 | src_length = 20; | |
131 | } | |
132 | ||
133 | consumed = ff_hevc_extract_rbsp(h, buf, src_length, nal); | |
134 | if (consumed < 0) | |
135 | return consumed; | |
136 | ||
137 | init_get_bits8(gb, nal->data + 2, nal->size); | |
138 | switch (h->nal_unit_type) { | |
139 | case NAL_VPS: | |
140 | ff_hevc_decode_nal_vps(h); | |
141 | break; | |
142 | case NAL_SPS: | |
143 | ff_hevc_decode_nal_sps(h); | |
144 | break; | |
145 | case NAL_PPS: | |
146 | ff_hevc_decode_nal_pps(h); | |
147 | break; | |
148 | case NAL_SEI_PREFIX: | |
149 | case NAL_SEI_SUFFIX: | |
150 | ff_hevc_decode_nal_sei(h); | |
151 | break; | |
152 | case NAL_TRAIL_N: | |
153 | case NAL_TRAIL_R: | |
154 | case NAL_TSA_N: | |
155 | case NAL_TSA_R: | |
156 | case NAL_STSA_N: | |
157 | case NAL_STSA_R: | |
158 | case NAL_RADL_N: | |
159 | case NAL_RADL_R: | |
160 | case NAL_RASL_N: | |
161 | case NAL_RASL_R: | |
162 | case NAL_BLA_W_LP: | |
163 | case NAL_BLA_W_RADL: | |
164 | case NAL_BLA_N_LP: | |
165 | case NAL_IDR_W_RADL: | |
166 | case NAL_IDR_N_LP: | |
167 | case NAL_CRA_NUT: | |
168 | sh->first_slice_in_pic_flag = get_bits1(gb); | |
169 | s->picture_structure = h->picture_struct; | |
170 | s->field_order = h->picture_struct; | |
171 | ||
172 | if (IS_IRAP(h)) { | |
173 | s->key_frame = 1; | |
174 | sh->no_output_of_prior_pics_flag = get_bits1(gb); | |
175 | } | |
176 | ||
177 | sh->pps_id = get_ue_golomb(gb); | |
178 | if (sh->pps_id >= MAX_PPS_COUNT || !h->pps_list[sh->pps_id]) { | |
179 | av_log(h->avctx, AV_LOG_ERROR, "PPS id out of range: %d\n", sh->pps_id); | |
180 | return AVERROR_INVALIDDATA; | |
181 | } | |
182 | h->pps = (HEVCPPS*)h->pps_list[sh->pps_id]->data; | |
183 | ||
184 | if (h->pps->sps_id >= MAX_SPS_COUNT || !h->sps_list[h->pps->sps_id]) { | |
185 | av_log(h->avctx, AV_LOG_ERROR, "SPS id out of range: %d\n", h->pps->sps_id); | |
186 | return AVERROR_INVALIDDATA; | |
187 | } | |
188 | if (h->sps != (HEVCSPS*)h->sps_list[h->pps->sps_id]->data) { | |
189 | h->sps = (HEVCSPS*)h->sps_list[h->pps->sps_id]->data; | |
190 | h->vps = (HEVCVPS*)h->vps_list[h->sps->vps_id]->data; | |
191 | } | |
192 | ||
193 | if (!sh->first_slice_in_pic_flag) { | |
194 | int slice_address_length; | |
195 | ||
196 | if (h->pps->dependent_slice_segments_enabled_flag) | |
197 | sh->dependent_slice_segment_flag = get_bits1(gb); | |
198 | else | |
199 | sh->dependent_slice_segment_flag = 0; | |
200 | ||
201 | slice_address_length = av_ceil_log2_c(h->sps->ctb_width * | |
202 | h->sps->ctb_height); | |
203 | sh->slice_segment_addr = get_bits(gb, slice_address_length); | |
204 | if (sh->slice_segment_addr >= h->sps->ctb_width * h->sps->ctb_height) { | |
205 | av_log(h->avctx, AV_LOG_ERROR, "Invalid slice segment address: %u.\n", | |
206 | sh->slice_segment_addr); | |
207 | return AVERROR_INVALIDDATA; | |
208 | } | |
209 | } else | |
210 | sh->dependent_slice_segment_flag = 0; | |
211 | ||
212 | if (sh->dependent_slice_segment_flag) | |
213 | break; | |
214 | ||
215 | for (i = 0; i < h->pps->num_extra_slice_header_bits; i++) | |
216 | skip_bits(gb, 1); // slice_reserved_undetermined_flag[] | |
217 | ||
218 | sh->slice_type = get_ue_golomb(gb); | |
219 | if (!(sh->slice_type == I_SLICE || sh->slice_type == P_SLICE || | |
220 | sh->slice_type == B_SLICE)) { | |
221 | av_log(h->avctx, AV_LOG_ERROR, "Unknown slice type: %d.\n", | |
222 | sh->slice_type); | |
223 | return AVERROR_INVALIDDATA; | |
224 | } | |
225 | s->pict_type = sh->slice_type == B_SLICE ? AV_PICTURE_TYPE_B : | |
226 | sh->slice_type == P_SLICE ? AV_PICTURE_TYPE_P : | |
227 | AV_PICTURE_TYPE_I; | |
228 | ||
229 | if (h->pps->output_flag_present_flag) | |
230 | sh->pic_output_flag = get_bits1(gb); | |
231 | ||
232 | if (h->sps->separate_colour_plane_flag) | |
233 | sh->colour_plane_id = get_bits(gb, 2); | |
234 | ||
235 | if (!IS_IDR(h)) { | |
236 | sh->pic_order_cnt_lsb = get_bits(gb, h->sps->log2_max_poc_lsb); | |
237 | s->output_picture_number = h->poc = ff_hevc_compute_poc(h, sh->pic_order_cnt_lsb); | |
238 | } else | |
239 | s->output_picture_number = h->poc = 0; | |
240 | ||
241 | if (h->temporal_id == 0 && | |
242 | h->nal_unit_type != NAL_TRAIL_N && | |
243 | h->nal_unit_type != NAL_TSA_N && | |
244 | h->nal_unit_type != NAL_STSA_N && | |
245 | h->nal_unit_type != NAL_RADL_N && | |
246 | h->nal_unit_type != NAL_RASL_N && | |
247 | h->nal_unit_type != NAL_RADL_R && | |
248 | h->nal_unit_type != NAL_RASL_R) | |
249 | h->pocTid0 = h->poc; | |
250 | ||
251 | return 0; /* no need to evaluate the rest */ | |
252 | } | |
253 | buf += consumed; | |
254 | } | |
255 | /* didn't find a picture! */ | |
256 | av_log(h->avctx, AV_LOG_ERROR, "missing picture in access unit\n"); | |
257 | return -1; | |
258 | } | |
259 | ||
260 | static int hevc_parse(AVCodecParserContext *s, | |
261 | AVCodecContext *avctx, | |
262 | const uint8_t **poutbuf, int *poutbuf_size, | |
263 | const uint8_t *buf, int buf_size) | |
264 | { | |
265 | int next; | |
266 | ParseContext *pc = &((HEVCParseContext *)s->priv_data)->pc; | |
267 | ||
268 | if (s->flags & PARSER_FLAG_COMPLETE_FRAMES) { | |
269 | next = buf_size; | |
270 | } else { | |
271 | next = hevc_find_frame_end(s, buf, buf_size); | |
272 | if (ff_combine_frame(pc, next, &buf, &buf_size) < 0) { | |
273 | *poutbuf = NULL; | |
274 | *poutbuf_size = 0; | |
275 | return buf_size; | |
276 | } | |
277 | } | |
278 | ||
279 | parse_nal_units(s, avctx, buf, buf_size); | |
280 | ||
281 | *poutbuf = buf; | |
282 | *poutbuf_size = buf_size; | |
283 | return next; | |
284 | } | |
285 | ||
286 | // Split after the parameter sets at the beginning of the stream if they exist. | |
287 | static int hevc_split(AVCodecContext *avctx, const uint8_t *buf, int buf_size) | |
288 | { | |
289 | int i; | |
290 | uint32_t state = -1; | |
291 | int has_ps = 0; | |
292 | ||
293 | for (i = 0; i < buf_size; i++) { | |
294 | state = (state << 8) | buf[i]; | |
295 | if (((state >> 8) & 0xFFFFFF) == START_CODE) { | |
296 | int nut = (state >> 1) & 0x3F; | |
297 | if (nut >= NAL_VPS && nut <= NAL_PPS) | |
298 | has_ps = 1; | |
299 | else if (has_ps) | |
300 | return i - 3; | |
301 | else // no parameter set at the beginning of the stream | |
302 | return 0; | |
303 | } | |
304 | } | |
305 | return 0; | |
306 | } | |
307 | ||
308 | static int hevc_init(AVCodecParserContext *s) | |
309 | { | |
310 | HEVCContext *h = &((HEVCParseContext *)s->priv_data)->h; | |
311 | h->HEVClc = av_mallocz(sizeof(HEVCLocalContext)); | |
312 | h->skipped_bytes_pos_size = INT_MAX; | |
313 | ||
314 | return 0; | |
315 | } | |
316 | ||
317 | static void hevc_close(AVCodecParserContext *s) | |
318 | { | |
319 | int i; | |
320 | HEVCContext *h = &((HEVCParseContext *)s->priv_data)->h; | |
321 | ParseContext *pc = &((HEVCParseContext *)s->priv_data)->pc; | |
322 | ||
323 | av_freep(&h->skipped_bytes_pos); | |
324 | av_freep(&h->HEVClc); | |
325 | av_freep(&pc->buffer); | |
326 | ||
327 | for (i = 0; i < FF_ARRAY_ELEMS(h->vps_list); i++) | |
328 | av_buffer_unref(&h->vps_list[i]); | |
329 | for (i = 0; i < FF_ARRAY_ELEMS(h->sps_list); i++) | |
330 | av_buffer_unref(&h->sps_list[i]); | |
331 | for (i = 0; i < FF_ARRAY_ELEMS(h->pps_list); i++) | |
332 | av_buffer_unref(&h->pps_list[i]); | |
333 | ||
334 | av_buffer_unref(&h->current_sps); | |
335 | h->sps = NULL; | |
336 | ||
337 | for (i = 0; i < h->nals_allocated; i++) | |
338 | av_freep(&h->nals[i].rbsp_buffer); | |
339 | av_freep(&h->nals); | |
340 | h->nals_allocated = 0; | |
341 | } | |
342 | ||
343 | AVCodecParser ff_hevc_parser = { | |
344 | .codec_ids = { AV_CODEC_ID_HEVC }, | |
345 | .priv_data_size = sizeof(HEVCParseContext), | |
346 | .parser_init = hevc_init, | |
347 | .parser_parse = hevc_parse, | |
348 | .parser_close = hevc_close, | |
349 | .split = hevc_split, | |
350 | }; |