blob: 4cd0ca11376265b2dc97fabf731c4cb06ca60185
1 | /* |
2 | * This file is part of FFmpeg. |
3 | * |
4 | * FFmpeg is free software; you can redistribute it and/or |
5 | * modify it under the terms of the GNU Lesser General Public |
6 | * License as published by the Free Software Foundation; either |
7 | * version 2.1 of the License, or (at your option) any later version. |
8 | * |
9 | * FFmpeg is distributed in the hope that it will be useful, |
10 | * but WITHOUT ANY WARRANTY; without even the implied warranty of |
11 | * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU |
12 | * Lesser General Public License for more details. |
13 | * |
14 | * You should have received a copy of the GNU Lesser General Public |
15 | * License along with FFmpeg; if not, write to the Free Software |
16 | * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA |
17 | */ |
18 | |
19 | #include <stdint.h> |
20 | |
21 | #include "libavutil/common.h" |
22 | #include "libavutil/intreadwrite.h" |
23 | #include "libavutil/log.h" |
24 | #include "libavutil/opt.h" |
25 | |
26 | #include "avcodec.h" |
27 | #include "bsf.h" |
28 | #include "h2645_parse.h" |
29 | #include "h264.h" |
30 | #include "hevc.h" |
31 | #include "vc1_common.h" |
32 | |
33 | typedef struct ExtractExtradataContext { |
34 | const AVClass *class; |
35 | |
36 | int (*extract)(AVBSFContext *ctx, AVPacket *pkt, |
37 | uint8_t **data, int *size); |
38 | |
39 | /* AVOptions */ |
40 | int remove; |
41 | } ExtractExtradataContext; |
42 | |
43 | static int val_in_array(const int *arr, int len, int val) |
44 | { |
45 | int i; |
46 | for (i = 0; i < len; i++) |
47 | if (arr[i] == val) |
48 | return 1; |
49 | return 0; |
50 | } |
51 | |
52 | static int extract_extradata_h2645(AVBSFContext *ctx, AVPacket *pkt, |
53 | uint8_t **data, int *size) |
54 | { |
55 | static const int extradata_nal_types_hevc[] = { |
56 | HEVC_NAL_VPS, HEVC_NAL_SPS, HEVC_NAL_PPS, |
57 | }; |
58 | static const int extradata_nal_types_h264[] = { |
59 | H264_NAL_SPS, H264_NAL_PPS, |
60 | }; |
61 | |
62 | ExtractExtradataContext *s = ctx->priv_data; |
63 | |
64 | H2645Packet h2645_pkt = { 0 }; |
65 | int extradata_size = 0; |
66 | const int *extradata_nal_types; |
67 | int nb_extradata_nal_types; |
68 | int i, has_sps = 0, has_vps = 0, ret = 0; |
69 | |
70 | if (ctx->par_in->codec_id == AV_CODEC_ID_HEVC) { |
71 | extradata_nal_types = extradata_nal_types_hevc; |
72 | nb_extradata_nal_types = FF_ARRAY_ELEMS(extradata_nal_types_hevc); |
73 | } else { |
74 | extradata_nal_types = extradata_nal_types_h264; |
75 | nb_extradata_nal_types = FF_ARRAY_ELEMS(extradata_nal_types_h264); |
76 | } |
77 | |
78 | ret = ff_h2645_packet_split(&h2645_pkt, pkt->data, pkt->size, |
79 | ctx, 0, 0, ctx->par_in->codec_id, 1); |
80 | if (ret < 0) |
81 | return ret; |
82 | |
83 | for (i = 0; i < h2645_pkt.nb_nals; i++) { |
84 | H2645NAL *nal = &h2645_pkt.nals[i]; |
85 | if (val_in_array(extradata_nal_types, nb_extradata_nal_types, nal->type)) { |
86 | extradata_size += nal->raw_size + 3; |
87 | if (ctx->par_in->codec_id == AV_CODEC_ID_HEVC) { |
88 | if (nal->type == HEVC_NAL_SPS) has_sps = 1; |
89 | if (nal->type == HEVC_NAL_VPS) has_vps = 1; |
90 | } else { |
91 | if (nal->type == H264_NAL_SPS) has_sps = 1; |
92 | } |
93 | } |
94 | } |
95 | |
96 | if (extradata_size && |
97 | ((ctx->par_in->codec_id == AV_CODEC_ID_HEVC && has_sps && has_vps) || |
98 | (ctx->par_in->codec_id == AV_CODEC_ID_H264 && has_sps))) { |
99 | AVBufferRef *filtered_buf; |
100 | uint8_t *extradata, *filtered_data; |
101 | |
102 | if (s->remove) { |
103 | filtered_buf = av_buffer_alloc(pkt->size + AV_INPUT_BUFFER_PADDING_SIZE); |
104 | if (!filtered_buf) |
105 | goto fail; |
106 | filtered_data = filtered_buf->data; |
107 | } |
108 | |
109 | extradata = av_malloc(extradata_size + AV_INPUT_BUFFER_PADDING_SIZE); |
110 | if (!extradata) { |
111 | av_buffer_unref(&filtered_buf); |
112 | goto fail; |
113 | } |
114 | |
115 | *data = extradata; |
116 | *size = extradata_size; |
117 | |
118 | for (i = 0; i < h2645_pkt.nb_nals; i++) { |
119 | H2645NAL *nal = &h2645_pkt.nals[i]; |
120 | if (val_in_array(extradata_nal_types, nb_extradata_nal_types, |
121 | nal->type)) { |
122 | AV_WB24(extradata, 1); // startcode |
123 | memcpy(extradata + 3, nal->raw_data, nal->raw_size); |
124 | extradata += 3 + nal->raw_size; |
125 | } else if (s->remove) { |
126 | AV_WB24(filtered_data, 1); // startcode |
127 | memcpy(filtered_data + 3, nal->raw_data, nal->raw_size); |
128 | filtered_data += 3 + nal->raw_size; |
129 | } |
130 | } |
131 | |
132 | if (s->remove) { |
133 | av_buffer_unref(&pkt->buf); |
134 | pkt->buf = filtered_buf; |
135 | pkt->data = filtered_buf->data; |
136 | pkt->size = filtered_data - filtered_buf->data; |
137 | } |
138 | } |
139 | |
140 | fail: |
141 | ff_h2645_packet_uninit(&h2645_pkt); |
142 | return ret; |
143 | } |
144 | |
145 | static int extract_extradata_vc1(AVBSFContext *ctx, AVPacket *pkt, |
146 | uint8_t **data, int *size) |
147 | { |
148 | ExtractExtradataContext *s = ctx->priv_data; |
149 | const uint8_t *ptr = pkt->data, *end = pkt->data + pkt->size; |
150 | uint32_t state = UINT32_MAX; |
151 | int has_extradata = 0, extradata_size = 0; |
152 | |
153 | while (ptr < end) { |
154 | ptr = avpriv_find_start_code(ptr, end, &state); |
155 | if (state == VC1_CODE_SEQHDR || state == VC1_CODE_ENTRYPOINT) { |
156 | has_extradata = 1; |
157 | } else if (has_extradata && IS_MARKER(state)) { |
158 | extradata_size = ptr - 4 - pkt->data; |
159 | break; |
160 | } |
161 | } |
162 | |
163 | if (extradata_size) { |
164 | *data = av_malloc(extradata_size + AV_INPUT_BUFFER_PADDING_SIZE); |
165 | if (!*data) |
166 | return AVERROR(ENOMEM); |
167 | |
168 | memcpy(*data, pkt->data, extradata_size); |
169 | *size = extradata_size; |
170 | |
171 | if (s->remove) { |
172 | pkt->data += extradata_size; |
173 | pkt->size -= extradata_size; |
174 | } |
175 | } |
176 | |
177 | return 0; |
178 | } |
179 | |
180 | static int extract_extradata_mpeg12(AVBSFContext *ctx, AVPacket *pkt, |
181 | uint8_t **data, int *size) |
182 | { |
183 | ExtractExtradataContext *s = ctx->priv_data; |
184 | uint32_t state = UINT32_MAX; |
185 | int i, found = 0; |
186 | |
187 | for (i = 0; i < pkt->size; i++) { |
188 | state = (state << 8) | pkt->data[i]; |
189 | if (state == 0x1B3) |
190 | found = 1; |
191 | else if (found && state != 0x1B5 && state < 0x200 && state >= 0x100) { |
192 | if (i > 3) { |
193 | *size = i - 3; |
194 | *data = av_malloc(*size + AV_INPUT_BUFFER_PADDING_SIZE); |
195 | if (!*data) |
196 | return AVERROR(ENOMEM); |
197 | |
198 | memcpy(*data, pkt->data, *size); |
199 | |
200 | if (s->remove) { |
201 | pkt->data += *size; |
202 | pkt->size -= *size; |
203 | } |
204 | } |
205 | break; |
206 | } |
207 | } |
208 | return 0; |
209 | } |
210 | |
211 | static int extract_extradata_mpeg4(AVBSFContext *ctx, AVPacket *pkt, |
212 | uint8_t **data, int *size) |
213 | { |
214 | ExtractExtradataContext *s = ctx->priv_data; |
215 | const uint8_t *ptr = pkt->data, *end = pkt->data + pkt->size; |
216 | uint32_t state = UINT32_MAX; |
217 | |
218 | while (ptr < end) { |
219 | ptr = avpriv_find_start_code(ptr, end, &state); |
220 | if (state == 0x1B3 || state == 0x1B6) { |
221 | if (ptr - pkt->data > 4) { |
222 | *size = ptr - 4 - pkt->data; |
223 | *data = av_malloc(*size + AV_INPUT_BUFFER_PADDING_SIZE); |
224 | if (!*data) |
225 | return AVERROR(ENOMEM); |
226 | |
227 | memcpy(*data, pkt->data, *size); |
228 | |
229 | if (s->remove) { |
230 | pkt->data += *size; |
231 | pkt->size -= *size; |
232 | } |
233 | } |
234 | break; |
235 | } |
236 | } |
237 | return 0; |
238 | } |
239 | |
240 | static const struct { |
241 | enum AVCodecID id; |
242 | int (*extract)(AVBSFContext *ctx, AVPacket *pkt, |
243 | uint8_t **data, int *size); |
244 | } extract_tab[] = { |
245 | { AV_CODEC_ID_CAVS, extract_extradata_mpeg4 }, |
246 | { AV_CODEC_ID_H264, extract_extradata_h2645 }, |
247 | { AV_CODEC_ID_HEVC, extract_extradata_h2645 }, |
248 | { AV_CODEC_ID_MPEG1VIDEO, extract_extradata_mpeg12 }, |
249 | { AV_CODEC_ID_MPEG2VIDEO, extract_extradata_mpeg12 }, |
250 | { AV_CODEC_ID_MPEG4, extract_extradata_mpeg4 }, |
251 | { AV_CODEC_ID_VC1, extract_extradata_vc1 }, |
252 | }; |
253 | |
254 | static int extract_extradata_init(AVBSFContext *ctx) |
255 | { |
256 | ExtractExtradataContext *s = ctx->priv_data; |
257 | int i; |
258 | |
259 | for (i = 0; i < FF_ARRAY_ELEMS(extract_tab); i++) { |
260 | if (extract_tab[i].id == ctx->par_in->codec_id) { |
261 | s->extract = extract_tab[i].extract; |
262 | break; |
263 | } |
264 | } |
265 | if (!s->extract) |
266 | return AVERROR_BUG; |
267 | |
268 | return 0; |
269 | } |
270 | |
271 | static int extract_extradata_filter(AVBSFContext *ctx, AVPacket *out) |
272 | { |
273 | ExtractExtradataContext *s = ctx->priv_data; |
274 | AVPacket *in; |
275 | uint8_t *extradata = NULL; |
276 | int extradata_size; |
277 | int ret = 0; |
278 | |
279 | ret = ff_bsf_get_packet(ctx, &in); |
280 | if (ret < 0) |
281 | return ret; |
282 | |
283 | ret = s->extract(ctx, in, &extradata, &extradata_size); |
284 | if (ret < 0) |
285 | goto fail; |
286 | |
287 | if (extradata) { |
288 | ret = av_packet_add_side_data(in, AV_PKT_DATA_NEW_EXTRADATA, |
289 | extradata, extradata_size); |
290 | if (ret < 0) { |
291 | av_freep(&extradata); |
292 | goto fail; |
293 | } |
294 | } |
295 | |
296 | av_packet_move_ref(out, in); |
297 | |
298 | fail: |
299 | av_packet_free(&in); |
300 | return ret; |
301 | } |
302 | |
303 | static const enum AVCodecID codec_ids[] = { |
304 | AV_CODEC_ID_CAVS, |
305 | AV_CODEC_ID_H264, |
306 | AV_CODEC_ID_HEVC, |
307 | AV_CODEC_ID_MPEG1VIDEO, |
308 | AV_CODEC_ID_MPEG2VIDEO, |
309 | AV_CODEC_ID_MPEG4, |
310 | AV_CODEC_ID_VC1, |
311 | AV_CODEC_ID_NONE, |
312 | }; |
313 | |
314 | #define OFFSET(x) offsetof(ExtractExtradataContext, x) |
315 | static const AVOption options[] = { |
316 | { "remove", "remove the extradata from the bitstream", OFFSET(remove), AV_OPT_TYPE_INT, |
317 | { .i64 = 0 }, 0, 1 }, |
318 | { NULL }, |
319 | }; |
320 | |
321 | static const AVClass extract_extradata_class = { |
322 | .class_name = "extract_extradata", |
323 | .item_name = av_default_item_name, |
324 | .option = options, |
325 | .version = LIBAVUTIL_VERSION_INT, |
326 | }; |
327 | |
328 | const AVBitStreamFilter ff_extract_extradata_bsf = { |
329 | .name = "extract_extradata", |
330 | .codec_ids = codec_ids, |
331 | .priv_data_size = sizeof(ExtractExtradataContext), |
332 | .priv_class = &extract_extradata_class, |
333 | .init = extract_extradata_init, |
334 | .filter = extract_extradata_filter, |
335 | }; |
336 |