FFmpeg coverage


Directory: ../../../ffmpeg/
File: src/libavcodec/dts2pts_bsf.c
Date: 2022-12-09 07:38:14
Exec Total Coverage
Lines: 0 273 0.0%
Functions: 0 14 0.0%
Branches: 0 169 0.0%

Line Branch Exec Source
1 /*
2 * Copyright (c) 2022 James Almer
3 *
4 * This file is part of FFmpeg.
5 *
6 * FFmpeg is free software; you can redistribute it and/or
7 * modify it under the terms of the GNU Lesser General Public
8 * License as published by the Free Software Foundation; either
9 * version 2.1 of the License, or (at your option) any later version.
10 *
11 * FFmpeg is distributed in the hope that it will be useful,
12 * but WITHOUT ANY WARRANTY; without even the implied warranty of
13 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
14 * Lesser General Public License for more details.
15 *
16 * You should have received a copy of the GNU Lesser General Public
17 * License along with FFmpeg; if not, write to the Free Software
18 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
19 */
20
21 /**
22 * @file
23 * Derive PTS by reordering DTS from supported streams
24 */
25
26 #include "libavutil/avassert.h"
27 #include "libavutil/fifo.h"
28 #include "libavutil/tree.h"
29
30 #include "bsf.h"
31 #include "bsf_internal.h"
32 #include "cbs.h"
33 #include "cbs_h264.h"
34 #include "h264_parse.h"
35 #include "h264_ps.h"
36
37 typedef struct DTS2PTSNode {
38 int64_t dts;
39 int64_t duration;
40 int poc;
41 int gop;
42 } DTS2PTSNode;
43
44 typedef struct DTS2PTSFrame {
45 AVPacket *pkt;
46 int poc;
47 int poc_diff;
48 int gop;
49 } DTS2PTSFrame;
50
51 typedef struct DTS2PTSH264Context {
52 H264POCContext poc;
53 SPS sps;
54 int poc_diff;
55 int last_poc;
56 int highest_poc;
57 int picture_structure;
58 } DTS2PTSH264Context;
59
60 typedef struct DTS2PTSContext {
61 struct AVTreeNode *root;
62 AVFifo *fifo;
63
64 // Codec specific function pointers and constants
65 int (*init)(AVBSFContext *ctx);
66 int (*filter)(AVBSFContext *ctx);
67 void (*flush)(AVBSFContext *ctx);
68 size_t fifo_size;
69
70 CodedBitstreamContext *cbc;
71 CodedBitstreamFragment au;
72
73 union {
74 DTS2PTSH264Context h264;
75 } u;
76
77 int nb_frame;
78 int gop;
79 int eof;
80 } DTS2PTSContext;
81
82 // AVTreeNode callbacks
83 static int cmp_insert(const void *key, const void *node)
84 {
85 int ret = ((const DTS2PTSNode *)key)->poc - ((const DTS2PTSNode *)node)->poc;
86 if (!ret)
87 ret = ((const DTS2PTSNode *)key)->gop - ((const DTS2PTSNode *)node)->gop;
88 return ret;
89 }
90
91 static int cmp_find(const void *key, const void *node)
92 {
93 int ret = ((const DTS2PTSFrame *)key)->poc - ((const DTS2PTSNode *) node)->poc;
94 if (!ret)
95 ret = ((const DTS2PTSFrame *)key)->gop - ((const DTS2PTSNode *) node)->gop;
96 return ret;
97 }
98
99 static int dec_poc(void *opaque, void *elem)
100 {
101 DTS2PTSNode *node = elem;
102 int dec = *(int *)opaque;
103 node->poc -= dec;
104 return 0;
105 }
106
107 static int free_node(void *opaque, void *elem)
108 {
109 DTS2PTSNode *node = elem;
110 av_free(node);
111 return 0;
112 }
113
114 // Shared functions
115 static int alloc_and_insert_node(AVBSFContext *ctx, int64_t ts, int64_t duration,
116 int poc, int poc_diff, int gop)
117 {
118 DTS2PTSContext *s = ctx->priv_data;
119 for (int i = 0; i < poc_diff; i++) {
120 struct AVTreeNode *node = av_tree_node_alloc();
121 DTS2PTSNode *poc_node, *ret;
122 if (!node)
123 return AVERROR(ENOMEM);
124 poc_node = av_malloc(sizeof(*poc_node));
125 if (!poc_node) {
126 av_free(node);
127 return AVERROR(ENOMEM);
128 }
129 if (i && ts != AV_NOPTS_VALUE)
130 ts += duration / poc_diff;
131 *poc_node = (DTS2PTSNode) { ts, duration, poc++, gop };
132 ret = av_tree_insert(&s->root, poc_node, cmp_insert, &node);
133 if (ret && ret != poc_node) {
134 *ret = *poc_node;
135 av_free(poc_node);
136 av_free(node);
137 }
138 }
139 return 0;
140 }
141
142 // H.264
143 static const CodedBitstreamUnitType h264_decompose_unit_types[] = {
144 H264_NAL_SPS,
145 H264_NAL_PPS,
146 H264_NAL_IDR_SLICE,
147 H264_NAL_SLICE,
148 };
149
150 static int h264_init(AVBSFContext *ctx)
151 {
152 DTS2PTSContext *s = ctx->priv_data;
153 DTS2PTSH264Context *h264 = &s->u.h264;
154
155 s->cbc->decompose_unit_types = h264_decompose_unit_types;
156 s->cbc->nb_decompose_unit_types = FF_ARRAY_ELEMS(h264_decompose_unit_types);
157
158 s->nb_frame = -(ctx->par_in->video_delay << 1);
159 h264->last_poc = h264->highest_poc = INT_MIN;
160
161 return 0;
162 }
163
164 static int get_mmco_reset(const H264RawSliceHeader *header)
165 {
166 if (header->nal_unit_header.nal_ref_idc == 0 ||
167 !header->adaptive_ref_pic_marking_mode_flag)
168 return 0;
169
170 for (int i = 0; i < H264_MAX_MMCO_COUNT; i++) {
171 if (header->mmco[i].memory_management_control_operation == 0)
172 return 0;
173 else if (header->mmco[i].memory_management_control_operation == 5)
174 return 1;
175 }
176
177 return 0;
178 }
179
180 static int h264_queue_frame(AVBSFContext *ctx, AVPacket *pkt, int poc, int *queued)
181 {
182 DTS2PTSContext *s = ctx->priv_data;
183 DTS2PTSH264Context *h264 = &s->u.h264;
184 DTS2PTSFrame frame;
185 int poc_diff, ret;
186
187 poc_diff = (h264->picture_structure == 3) + 1;
188 if (h264->sps.frame_mbs_only_flag && h264->poc_diff)
189 poc_diff = FFMIN(poc_diff, h264->poc_diff);
190 if (poc < 0) {
191 av_tree_enumerate(s->root, &poc_diff, NULL, dec_poc);
192 s->nb_frame -= poc_diff;
193 }
194 // Check if there was a POC reset (Like an IDR slice)
195 if (s->nb_frame > h264->highest_poc) {
196 s->nb_frame = 0;
197 s->gop = (s->gop + 1) % s->fifo_size;
198 h264->highest_poc = h264->last_poc;
199 }
200
201 ret = alloc_and_insert_node(ctx, pkt->dts, pkt->duration, s->nb_frame, poc_diff, s->gop);
202 if (ret < 0)
203 return ret;
204 av_log(ctx, AV_LOG_DEBUG, "Queueing frame with POC %d, GOP %d, dts %"PRId64"\n",
205 poc, s->gop, pkt->dts);
206 s->nb_frame += poc_diff;
207
208 // Add frame to output FIFO only once
209 if (*queued)
210 return 0;
211
212 frame = (DTS2PTSFrame) { pkt, poc, poc_diff, s->gop };
213 ret = av_fifo_write(s->fifo, &frame, 1);
214 av_assert2(ret >= 0);
215 *queued = 1;
216
217 return 0;
218 }
219
220 static int h264_filter(AVBSFContext *ctx)
221 {
222 DTS2PTSContext *s = ctx->priv_data;
223 DTS2PTSH264Context *h264 = &s->u.h264;
224 CodedBitstreamFragment *au = &s->au;
225 AVPacket *in;
226 int output_picture_number = INT_MIN;
227 int field_poc[2];
228 int queued = 0, ret;
229
230 ret = ff_bsf_get_packet(ctx, &in);
231 if (ret < 0)
232 return ret;
233
234 ret = ff_cbs_read_packet(s->cbc, au, in);
235 if (ret < 0) {
236 av_log(ctx, AV_LOG_WARNING, "Failed to parse access unit.\n");
237 goto fail;
238 }
239
240 for (int i = 0; i < au->nb_units; i++) {
241 CodedBitstreamUnit *unit = &au->units[i];
242
243 switch (unit->type) {
244 case H264_NAL_IDR_SLICE:
245 h264->poc.prev_frame_num = 0;
246 h264->poc.prev_frame_num_offset = 0;
247 h264->poc.prev_poc_msb =
248 h264->poc.prev_poc_lsb = 0;
249 // fall-through
250 case H264_NAL_SLICE: {
251 const H264RawSlice *slice = unit->content;
252 const H264RawSliceHeader *header = &slice->header;
253 const CodedBitstreamH264Context *cbs_h264 = s->cbc->priv_data;
254 const H264RawSPS *sps = cbs_h264->active_sps;
255 int got_reset;
256
257 if (!sps) {
258 av_log(ctx, AV_LOG_ERROR, "No active SPS for a slice\n");
259 goto fail;
260 }
261 // Initialize the SPS struct with the fields ff_h264_init_poc() cares about
262 h264->sps.frame_mbs_only_flag = sps->frame_mbs_only_flag;
263 h264->sps.log2_max_frame_num = sps->log2_max_frame_num_minus4 + 4;
264 h264->sps.poc_type = sps->pic_order_cnt_type;
265 h264->sps.log2_max_poc_lsb = sps->log2_max_pic_order_cnt_lsb_minus4 + 4;
266 h264->sps.offset_for_non_ref_pic = sps->offset_for_non_ref_pic;
267 h264->sps.offset_for_top_to_bottom_field = sps->offset_for_top_to_bottom_field;
268 h264->sps.poc_cycle_length = sps->num_ref_frames_in_pic_order_cnt_cycle;
269 for (int i = 0; i < h264->sps.poc_cycle_length; i++)
270 h264->sps.offset_for_ref_frame[i] = sps->offset_for_ref_frame[i];
271
272 h264->picture_structure = sps->frame_mbs_only_flag ? 3 :
273 (header->field_pic_flag ?
274 header->field_pic_flag + header->bottom_field_flag : 3);
275
276 h264->poc.frame_num = header->frame_num;
277 h264->poc.poc_lsb = header->pic_order_cnt_lsb;
278 h264->poc.delta_poc_bottom = header->delta_pic_order_cnt_bottom;
279 h264->poc.delta_poc[0] = header->delta_pic_order_cnt[0];
280 h264->poc.delta_poc[1] = header->delta_pic_order_cnt[1];
281
282 field_poc[0] = field_poc[1] = INT_MAX;
283 ret = ff_h264_init_poc(field_poc, &output_picture_number, &h264->sps,
284 &h264->poc, h264->picture_structure,
285 header->nal_unit_header.nal_ref_idc);
286 if (ret < 0) {
287 av_log(ctx, AV_LOG_ERROR, "ff_h264_init_poc() failure\n");
288 goto fail;
289 }
290
291 got_reset = get_mmco_reset(header);
292 h264->poc.prev_frame_num = got_reset ? 0 : h264->poc.frame_num;
293 h264->poc.prev_frame_num_offset = got_reset ? 0 : h264->poc.frame_num_offset;
294 if (header->nal_unit_header.nal_ref_idc != 0) {
295 h264->poc.prev_poc_msb = got_reset ? 0 : h264->poc.poc_msb;
296 if (got_reset)
297 h264->poc.prev_poc_lsb = h264->picture_structure == 2 ? 0 : field_poc[0];
298 else
299 h264->poc.prev_poc_lsb = h264->poc.poc_lsb;
300 }
301
302 if (output_picture_number != h264->last_poc) {
303 if (h264->last_poc != INT_MIN) {
304 int64_t diff = FFABS(h264->last_poc - (int64_t)output_picture_number);
305
306 if ((output_picture_number < 0) && !h264->last_poc)
307 h264->poc_diff = 0;
308 else if (FFABS((int64_t)output_picture_number) < h264->poc_diff) {
309 diff = FFABS(output_picture_number);
310 h264->poc_diff = 0;
311 }
312 if ((!h264->poc_diff || (h264->poc_diff > diff)) && diff <= INT_MAX) {
313 h264->poc_diff = diff;
314 if (h264->poc_diff == 1 && h264->sps.frame_mbs_only_flag) {
315 av_tree_enumerate(s->root, &h264->poc_diff, NULL, dec_poc);
316 s->nb_frame -= 2;
317 }
318 }
319 }
320 h264->last_poc = output_picture_number;
321 h264->highest_poc = FFMAX(h264->highest_poc, output_picture_number);
322
323 ret = h264_queue_frame(ctx, in, output_picture_number, &queued);
324 if (ret < 0)
325 goto fail;
326 }
327 break;
328 }
329 default:
330 break;
331 }
332 }
333
334 if (output_picture_number == INT_MIN) {
335 av_log(ctx, AV_LOG_ERROR, "No slices in access unit\n");
336 ret = AVERROR_INVALIDDATA;
337 goto fail;
338 }
339
340 ret = 0;
341 fail:
342 ff_cbs_fragment_reset(au);
343 if (!queued)
344 av_packet_free(&in);
345
346 return ret;
347 }
348
349 static void h264_flush(AVBSFContext *ctx)
350 {
351 DTS2PTSContext *s = ctx->priv_data;
352 DTS2PTSH264Context *h264 = &s->u.h264;
353
354 memset(&h264->sps, 0, sizeof(h264->sps));
355 memset(&h264->poc, 0, sizeof(h264->poc));
356 s->nb_frame = -(ctx->par_in->video_delay << 1);
357 h264->last_poc = h264->highest_poc = INT_MIN;
358 }
359
360 // Core functions
361 static const struct {
362 enum AVCodecID id;
363 int (*init)(AVBSFContext *ctx);
364 int (*filter)(AVBSFContext *ctx);
365 void (*flush)(AVBSFContext *ctx);
366 size_t fifo_size;
367 } func_tab[] = {
368 { AV_CODEC_ID_H264, h264_init, h264_filter, h264_flush, H264_MAX_DPB_FRAMES * 2 * 2 },
369 };
370
371 static int dts2pts_init(AVBSFContext *ctx)
372 {
373 DTS2PTSContext *s = ctx->priv_data;
374 CodedBitstreamFragment *au = &s->au;
375 int i, ret;
376
377 for (i = 0; i < FF_ARRAY_ELEMS(func_tab); i++) {
378 if (func_tab[i].id == ctx->par_in->codec_id) {
379 s->init = func_tab[i].init;
380 s->filter = func_tab[i].filter;
381 s->flush = func_tab[i].flush;
382 s->fifo_size = func_tab[i].fifo_size;
383 break;
384 }
385 }
386 if (i == FF_ARRAY_ELEMS(func_tab))
387 return AVERROR_BUG;
388 av_assert0(s->filter && s->fifo_size);
389
390 s->fifo = av_fifo_alloc2(s->fifo_size, sizeof(DTS2PTSFrame), 0);
391 if (!s->fifo)
392 return AVERROR(ENOMEM);
393
394 ret = ff_cbs_init(&s->cbc, ctx->par_in->codec_id, ctx);
395 if (ret < 0)
396 return ret;
397
398 if (s->init) {
399 ret = s->init(ctx);
400 if (ret < 0)
401 return ret;
402 }
403
404 if (!ctx->par_in->extradata_size)
405 return 0;
406
407 ret = ff_cbs_read_extradata(s->cbc, au, ctx->par_in);
408 if (ret < 0)
409 av_log(ctx, AV_LOG_WARNING, "Failed to parse extradata.\n");
410
411 ff_cbs_fragment_reset(au);
412
413 return 0;
414 }
415
416 static int dts2pts_filter(AVBSFContext *ctx, AVPacket *out)
417 {
418 DTS2PTSContext *s = ctx->priv_data;
419 DTS2PTSNode *poc_node = NULL, *next[2] = { NULL, NULL };
420 DTS2PTSFrame frame;
421 int ret;
422
423 // Fill up the FIFO and POC tree
424 while (!s->eof && av_fifo_can_write(s->fifo)) {
425 ret = s->filter(ctx);
426 if (ret < 0) {
427 if (ret != AVERROR_EOF)
428 return ret;
429 s->eof = 1;
430 }
431 }
432
433 if (!av_fifo_can_read(s->fifo))
434 return AVERROR_EOF;
435
436 // Fetch a packet from the FIFO
437 ret = av_fifo_read(s->fifo, &frame, 1);
438 av_assert2(ret >= 0);
439 av_packet_move_ref(out, frame.pkt);
440 av_packet_free(&frame.pkt);
441
442 // Search the timestamp for the requested POC and set PTS
443 poc_node = av_tree_find(s->root, &frame, cmp_find, (void **)next);
444 if (!poc_node) {
445 poc_node = next[1];
446 if (!poc_node || poc_node->poc != frame.poc)
447 poc_node = next[0];
448 }
449 if (poc_node && poc_node->poc == frame.poc) {
450 out->pts = poc_node->dts;
451 if (!s->eof) {
452 // Remove the found entry from the tree
453 DTS2PTSFrame dup = (DTS2PTSFrame) { NULL, frame.poc + 1, frame.poc_diff, frame.gop };
454 for (; dup.poc_diff > 0; dup.poc++, dup.poc_diff--) {
455 struct AVTreeNode *node = NULL;
456 if (!poc_node || poc_node->dts != out->pts)
457 continue;
458 av_tree_insert(&s->root, poc_node, cmp_insert, &node);
459 av_free(poc_node);
460 av_free(node);
461 poc_node = av_tree_find(s->root, &dup, cmp_find, NULL);
462 }
463 }
464 } else {
465 DTS2PTSFrame dup = (DTS2PTSFrame) { NULL, frame.poc - 1, frame.poc_diff, frame.gop };
466 if (s->eof && (poc_node = av_tree_find(s->root, &dup, cmp_find, NULL)) && poc_node->poc == dup.poc) {
467 out->pts = poc_node->dts;
468 if (out->pts != AV_NOPTS_VALUE)
469 out->pts += poc_node->duration;
470 ret = alloc_and_insert_node(ctx, out->pts, out->duration,
471 frame.poc, frame.poc_diff, frame.gop);
472 if (ret < 0) {
473 av_packet_unref(out);
474 return ret;
475 }
476 if (!ret)
477 av_log(ctx, AV_LOG_DEBUG, "Queueing frame for POC %d, GOP %d, dts %"PRId64", "
478 "generated from POC %d, GOP %d, dts %"PRId64", duration %"PRId64"\n",
479 frame.poc, frame.gop, out->pts,
480 poc_node->poc, poc_node->gop, poc_node->dts, poc_node->duration);
481 } else
482 av_log(ctx, AV_LOG_WARNING, "No timestamp for POC %d in tree\n", frame.poc);
483 }
484 av_log(ctx, AV_LOG_DEBUG, "Returning frame for POC %d, GOP %d, dts %"PRId64", pts %"PRId64"\n",
485 frame.poc, frame.gop, out->dts, out->pts);
486
487 return 0;
488 }
489
490 static void dts2pts_flush(AVBSFContext *ctx)
491 {
492 DTS2PTSContext *s = ctx->priv_data;
493 DTS2PTSFrame frame;
494
495 if (s->flush)
496 s->flush(ctx);
497 s->eof = 0;
498 s->gop = 0;
499
500 while (s->fifo && av_fifo_read(s->fifo, &frame, 1) >= 0)
501 av_packet_free(&frame.pkt);
502
503 av_tree_enumerate(s->root, NULL, NULL, free_node);
504 av_tree_destroy(s->root);
505 s->root = NULL;
506
507 ff_cbs_fragment_reset(&s->au);
508 if (s->cbc)
509 ff_cbs_flush(s->cbc);
510 }
511
512 static void dts2pts_close(AVBSFContext *ctx)
513 {
514 DTS2PTSContext *s = ctx->priv_data;
515
516 dts2pts_flush(ctx);
517
518 av_fifo_freep2(&s->fifo);
519 ff_cbs_fragment_free(&s->au);
520 ff_cbs_close(&s->cbc);
521 }
522
523 static const enum AVCodecID dts2pts_codec_ids[] = {
524 AV_CODEC_ID_H264,
525 AV_CODEC_ID_NONE,
526 };
527
528 const FFBitStreamFilter ff_dts2pts_bsf = {
529 .p.name = "dts2pts",
530 .p.codec_ids = dts2pts_codec_ids,
531 .priv_data_size = sizeof(DTS2PTSContext),
532 .init = dts2pts_init,
533 .flush = dts2pts_flush,
534 .close = dts2pts_close,
535 .filter = dts2pts_filter,
536 };
537