FFmpeg coverage

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		✗	const DTS2PTSFrame * key1 = key;
94		✗	const DTS2PTSNode *node1 = node;
95		✗	int ret = FFDIFFSIGN(key1->poc, node1->poc);
96		✗	if (!ret)
97		✗	ret = key1->gop - node1->gop;
98		✗	return ret;
99			}
100
101		✗	static int dec_poc(void *opaque, void *elem)
102			{
103		✗	DTS2PTSNode *node = elem;
104		✗	int dec = *(int *)opaque;
105		✗	node->poc -= dec;
106		✗	return 0;
107			}
108
109		✗	static int free_node(void *opaque, void *elem)
110			{
111		✗	DTS2PTSNode *node = elem;
112		✗	av_free(node);
113		✗	return 0;
114			}
115
116			// Shared functions
117		✗	static int alloc_and_insert_node(AVBSFContext *ctx, int64_t ts, int64_t duration,
118			int poc, int poc_diff, int gop)
119			{
120		✗	DTS2PTSContext *s = ctx->priv_data;
121		✗	for (int i = 0; i < poc_diff; i++) {
122		✗	struct AVTreeNode *node = av_tree_node_alloc();
123			DTS2PTSNode *poc_node, *ret;
124		✗	if (!node)
125		✗	return AVERROR(ENOMEM);
126		✗	poc_node = av_malloc(sizeof(*poc_node));
127		✗	if (!poc_node) {
128		✗	av_free(node);
129		✗	return AVERROR(ENOMEM);
130			}
131		✗	if (i && ts != AV_NOPTS_VALUE)
132		✗	ts += duration / poc_diff;
133		✗	*poc_node = (DTS2PTSNode) { ts, duration, poc++, gop };
134		✗	ret = av_tree_insert(&s->root, poc_node, cmp_insert, &node);
135		✗	if (ret && ret != poc_node) {
136		✗	*ret = *poc_node;
137		✗	av_free(poc_node);
138		✗	av_free(node);
139			}
140			}
141		✗	return 0;
142			}
143
144			// H.264
145			static const CodedBitstreamUnitType h264_decompose_unit_types[] = {
146			H264_NAL_SPS,
147			H264_NAL_PPS,
148			H264_NAL_IDR_SLICE,
149			H264_NAL_SLICE,
150			};
151
152		✗	static int h264_init(AVBSFContext *ctx)
153			{
154		✗	DTS2PTSContext *s = ctx->priv_data;
155		✗	DTS2PTSH264Context *h264 = &s->u.h264;
156
157		✗	s->cbc->decompose_unit_types = h264_decompose_unit_types;
158		✗	s->cbc->nb_decompose_unit_types = FF_ARRAY_ELEMS(h264_decompose_unit_types);
159
160		✗	s->nb_frame = -(ctx->par_in->video_delay << 1);
161		✗	h264->last_poc = h264->highest_poc = INT_MIN;
162
163		✗	return 0;
164			}
165
166		✗	static int get_mmco_reset(const H264RawSliceHeader *header)
167			{
168		✗	if (header->nal_unit_header.nal_ref_idc == 0 ||
169		✗	!header->adaptive_ref_pic_marking_mode_flag)
170		✗	return 0;
171
172		✗	for (int i = 0; i < H264_MAX_MMCO_COUNT; i++) {
173		✗	if (header->mmco[i].memory_management_control_operation == 0)
174		✗	return 0;
175		✗	else if (header->mmco[i].memory_management_control_operation == 5)
176		✗	return 1;
177			}
178
179		✗	return 0;
180			}
181
182		✗	static int h264_queue_frame(AVBSFContext *ctx, AVPacket *pkt, int poc, int *queued)
183			{
184		✗	DTS2PTSContext *s = ctx->priv_data;
185		✗	DTS2PTSH264Context *h264 = &s->u.h264;
186			DTS2PTSFrame frame;
187			int poc_diff, ret;
188
189		✗	poc_diff = (h264->picture_structure == 3) + 1;
190		✗	if (h264->sps.frame_mbs_only_flag && h264->poc_diff)
191		✗	poc_diff = FFMIN(poc_diff, h264->poc_diff);
192		✗	if (poc < 0) {
193		✗	av_tree_enumerate(s->root, &poc_diff, NULL, dec_poc);
194		✗	s->nb_frame -= poc_diff;
195			}
196			// Check if there was a POC reset (Like an IDR slice)
197		✗	if (s->nb_frame > h264->highest_poc) {
198		✗	s->nb_frame = 0;
199		✗	s->gop = (s->gop + 1) % s->fifo_size;
200		✗	h264->highest_poc = h264->last_poc;
201			}
202
203		✗	ret = alloc_and_insert_node(ctx, pkt->dts, pkt->duration, s->nb_frame, poc_diff, s->gop);
204		✗	if (ret < 0)
205		✗	return ret;
206		✗	av_log(ctx, AV_LOG_DEBUG, "Queueing frame with POC %d, GOP %d, dts %"PRId64"\n",
207			poc, s->gop, pkt->dts);
208		✗	s->nb_frame += poc_diff;
209
210			// Add frame to output FIFO only once
211		✗	if (*queued)
212		✗	return 0;
213
214		✗	frame = (DTS2PTSFrame) { pkt, poc, poc_diff, s->gop };
215		✗	ret = av_fifo_write(s->fifo, &frame, 1);
216			av_assert2(ret >= 0);
217		✗	*queued = 1;
218
219		✗	return 0;
220			}
221
222		✗	static int h264_filter(AVBSFContext *ctx)
223			{
224		✗	DTS2PTSContext *s = ctx->priv_data;
225		✗	DTS2PTSH264Context *h264 = &s->u.h264;
226		✗	CodedBitstreamFragment *au = &s->au;
227			AVPacket *in;
228		✗	int output_picture_number = INT_MIN;
229			int field_poc[2];
230		✗	int queued = 0, ret;
231
232		✗	ret = ff_bsf_get_packet(ctx, &in);
233		✗	if (ret < 0)
234		✗	return ret;
235
236		✗	ret = ff_cbs_read_packet(s->cbc, au, in);
237		✗	if (ret < 0) {
238		✗	av_log(ctx, AV_LOG_WARNING, "Failed to parse access unit.\n");
239		✗	goto fail;
240			}
241
242		✗	for (int i = 0; i < au->nb_units; i++) {
243		✗	CodedBitstreamUnit *unit = &au->units[i];
244
245		✗	switch (unit->type) {
246		✗	case H264_NAL_IDR_SLICE:
247		✗	h264->poc.prev_frame_num = 0;
248		✗	h264->poc.prev_frame_num_offset = 0;
249		✗	h264->poc.prev_poc_msb =
250		✗	h264->poc.prev_poc_lsb = 0;
251			// fall-through
252		✗	case H264_NAL_SLICE: {
253		✗	const H264RawSlice *slice = unit->content;
254		✗	const H264RawSliceHeader *header = &slice->header;
255		✗	const CodedBitstreamH264Context *cbs_h264 = s->cbc->priv_data;
256		✗	const H264RawSPS *sps = cbs_h264->active_sps;
257			int got_reset;
258
259		✗	if (!sps) {
260		✗	av_log(ctx, AV_LOG_ERROR, "No active SPS for a slice\n");
261		✗	goto fail;
262			}
263			// Initialize the SPS struct with the fields ff_h264_init_poc() cares about
264		✗	h264->sps.frame_mbs_only_flag = sps->frame_mbs_only_flag;
265		✗	h264->sps.log2_max_frame_num = sps->log2_max_frame_num_minus4 + 4;
266		✗	h264->sps.poc_type = sps->pic_order_cnt_type;
267		✗	h264->sps.log2_max_poc_lsb = sps->log2_max_pic_order_cnt_lsb_minus4 + 4;
268		✗	h264->sps.offset_for_non_ref_pic = sps->offset_for_non_ref_pic;
269		✗	h264->sps.offset_for_top_to_bottom_field = sps->offset_for_top_to_bottom_field;
270		✗	h264->sps.poc_cycle_length = sps->num_ref_frames_in_pic_order_cnt_cycle;
271		✗	for (int i = 0; i < h264->sps.poc_cycle_length; i++)
272		✗	h264->sps.offset_for_ref_frame[i] = sps->offset_for_ref_frame[i];
273
274		✗	h264->picture_structure = sps->frame_mbs_only_flag ? 3 :
275		✗	(header->field_pic_flag ?
276		✗	header->field_pic_flag + header->bottom_field_flag : 3);
277
278		✗	h264->poc.frame_num = header->frame_num;
279		✗	h264->poc.poc_lsb = header->pic_order_cnt_lsb;
280		✗	h264->poc.delta_poc_bottom = header->delta_pic_order_cnt_bottom;
281		✗	h264->poc.delta_poc[0] = header->delta_pic_order_cnt[0];
282		✗	h264->poc.delta_poc[1] = header->delta_pic_order_cnt[1];
283
284		✗	field_poc[0] = field_poc[1] = INT_MAX;
285		✗	ret = ff_h264_init_poc(field_poc, &output_picture_number, &h264->sps,
286			&h264->poc, h264->picture_structure,
287		✗	header->nal_unit_header.nal_ref_idc);
288		✗	if (ret < 0) {
289		✗	av_log(ctx, AV_LOG_ERROR, "ff_h264_init_poc() failure\n");
290		✗	goto fail;
291			}
292
293		✗	got_reset = get_mmco_reset(header);
294		✗	h264->poc.prev_frame_num = got_reset ? 0 : h264->poc.frame_num;
295		✗	h264->poc.prev_frame_num_offset = got_reset ? 0 : h264->poc.frame_num_offset;
296		✗	if (header->nal_unit_header.nal_ref_idc != 0) {
297		✗	h264->poc.prev_poc_msb = got_reset ? 0 : h264->poc.poc_msb;
298		✗	if (got_reset)
299		✗	h264->poc.prev_poc_lsb = h264->picture_structure == 2 ? 0 : field_poc[0];
300			else
301		✗	h264->poc.prev_poc_lsb = h264->poc.poc_lsb;
302			}
303
304		✗	if (output_picture_number != h264->last_poc) {
305		✗	if (h264->last_poc != INT_MIN) {
306		✗	int64_t diff = FFABS(h264->last_poc - (int64_t)output_picture_number);
307
308		✗	if ((output_picture_number < 0) && !h264->last_poc)
309		✗	h264->poc_diff = 0;
310		✗	else if (FFABS((int64_t)output_picture_number) < h264->poc_diff) {
311		✗	diff = FFABS(output_picture_number);
312		✗	h264->poc_diff = 0;
313			}
314		✗	if ((!h264->poc_diff || (h264->poc_diff > diff)) && diff <= INT_MAX) {
315		✗	h264->poc_diff = diff;
316		✗	if (h264->poc_diff == 1 && h264->sps.frame_mbs_only_flag) {
317		✗	av_tree_enumerate(s->root, &h264->poc_diff, NULL, dec_poc);
318		✗	s->nb_frame -= 2;
319			}
320			}
321			}
322		✗	h264->last_poc = output_picture_number;
323		✗	h264->highest_poc = FFMAX(h264->highest_poc, output_picture_number);
324
325		✗	ret = h264_queue_frame(ctx, in, output_picture_number, &queued);
326		✗	if (ret < 0)
327		✗	goto fail;
328			}
329		✗	break;
330			}
331		✗	default:
332		✗	break;
333			}
334			}
335
336		✗	if (output_picture_number == INT_MIN) {
337		✗	av_log(ctx, AV_LOG_ERROR, "No slices in access unit\n");
338		✗	ret = AVERROR_INVALIDDATA;
339		✗	goto fail;
340			}
341
342		✗	ret = 0;
343		✗	fail:
344		✗	ff_cbs_fragment_reset(au);
345		✗	if (!queued)
346		✗	av_packet_free(&in);
347
348		✗	return ret;
349			}
350
351		✗	static void h264_flush(AVBSFContext *ctx)
352			{
353		✗	DTS2PTSContext *s = ctx->priv_data;
354		✗	DTS2PTSH264Context *h264 = &s->u.h264;
355
356		✗	memset(&h264->sps, 0, sizeof(h264->sps));
357		✗	memset(&h264->poc, 0, sizeof(h264->poc));
358		✗	s->nb_frame = -(ctx->par_in->video_delay << 1);
359		✗	h264->last_poc = h264->highest_poc = INT_MIN;
360			}
361
362			// Core functions
363			static const struct {
364			enum AVCodecID id;
365			int (*init)(AVBSFContext *ctx);
366			int (*filter)(AVBSFContext *ctx);
367			void (*flush)(AVBSFContext *ctx);
368			size_t fifo_size;
369			} func_tab[] = {
370			{ AV_CODEC_ID_H264, h264_init, h264_filter, h264_flush, H264_MAX_DPB_FRAMES * 2 * 2 },
371			};
372
373		✗	static int dts2pts_init(AVBSFContext *ctx)
374			{
375		✗	DTS2PTSContext *s = ctx->priv_data;
376		✗	CodedBitstreamFragment *au = &s->au;
377			int i, ret;
378
379		✗	for (i = 0; i < FF_ARRAY_ELEMS(func_tab); i++) {
380		✗	if (func_tab[i].id == ctx->par_in->codec_id) {
381		✗	s->init = func_tab[i].init;
382		✗	s->filter = func_tab[i].filter;
383		✗	s->flush = func_tab[i].flush;
384		✗	s->fifo_size = func_tab[i].fifo_size;
385		✗	break;
386			}
387			}
388		✗	if (i == FF_ARRAY_ELEMS(func_tab))
389		✗	return AVERROR_BUG;
390		✗	av_assert0(s->filter && s->fifo_size);
391
392		✗	s->fifo = av_fifo_alloc2(s->fifo_size, sizeof(DTS2PTSFrame), 0);
393		✗	if (!s->fifo)
394		✗	return AVERROR(ENOMEM);
395
396		✗	ret = ff_cbs_init(&s->cbc, ctx->par_in->codec_id, ctx);
397		✗	if (ret < 0)
398		✗	return ret;
399
400		✗	if (s->init) {
401		✗	ret = s->init(ctx);
402		✗	if (ret < 0)
403		✗	return ret;
404			}
405
406		✗	if (!ctx->par_in->extradata_size)
407		✗	return 0;
408
409		✗	ret = ff_cbs_read_extradata(s->cbc, au, ctx->par_in);
410		✗	if (ret < 0)
411		✗	av_log(ctx, AV_LOG_WARNING, "Failed to parse extradata.\n");
412
413		✗	ff_cbs_fragment_reset(au);
414
415		✗	return 0;
416			}
417
418		✗	static int dts2pts_filter(AVBSFContext *ctx, AVPacket *out)
419			{
420		✗	DTS2PTSContext *s = ctx->priv_data;
421		✗	DTS2PTSNode *poc_node = NULL, *next[2] = { NULL, NULL };
422			DTS2PTSFrame frame;
423			int ret;
424
425			// Fill up the FIFO and POC tree
426		✗	while (!s->eof && av_fifo_can_write(s->fifo)) {
427		✗	ret = s->filter(ctx);
428		✗	if (ret < 0) {
429		✗	if (ret != AVERROR_EOF)
430		✗	return ret;
431		✗	s->eof = 1;
432			}
433			}
434
435		✗	if (!av_fifo_can_read(s->fifo))
436		✗	return AVERROR_EOF;
437
438			// Fetch a packet from the FIFO
439		✗	ret = av_fifo_read(s->fifo, &frame, 1);
440			av_assert2(ret >= 0);
441		✗	av_packet_move_ref(out, frame.pkt);
442		✗	av_packet_free(&frame.pkt);
443
444			// Search the timestamp for the requested POC and set PTS
445		✗	poc_node = av_tree_find(s->root, &frame, cmp_find, (void **)next);
446		✗	if (!poc_node) {
447		✗	poc_node = next[1];
448		✗	if (!poc_node || poc_node->poc != frame.poc)
449		✗	poc_node = next[0];
450			}
451		✗	if (poc_node && poc_node->poc == frame.poc) {
452		✗	out->pts = poc_node->dts;
453		✗	if (!s->eof) {
454			// Remove the found entry from the tree
455		✗	DTS2PTSFrame dup = (DTS2PTSFrame) { NULL, frame.poc + 1, frame.poc_diff, frame.gop };
456		✗	for (; dup.poc_diff > 0; dup.poc++, dup.poc_diff--) {
457		✗	struct AVTreeNode *node = NULL;
458		✗	if (!poc_node || poc_node->dts != out->pts)
459		✗	continue;
460		✗	av_tree_insert(&s->root, poc_node, cmp_insert, &node);
461		✗	av_free(poc_node);
462		✗	av_free(node);
463		✗	poc_node = av_tree_find(s->root, &dup, cmp_find, NULL);
464			}
465			}
466		✗	} else if (s->eof && frame.poc > INT_MIN) {
467		✗	DTS2PTSFrame dup = (DTS2PTSFrame) { NULL, frame.poc - 1, frame.poc_diff, frame.gop };
468		✗	poc_node = av_tree_find(s->root, &dup, cmp_find, NULL);
469		✗	if (poc_node && poc_node->poc == dup.poc) {
470		✗	out->pts = poc_node->dts;
471		✗	if (out->pts != AV_NOPTS_VALUE)
472		✗	out->pts += poc_node->duration;
473		✗	ret = alloc_and_insert_node(ctx, out->pts, out->duration,
474			frame.poc, frame.poc_diff, frame.gop);
475		✗	if (ret < 0) {
476		✗	av_packet_unref(out);
477		✗	return ret;
478			}
479		✗	if (!ret)
480		✗	av_log(ctx, AV_LOG_DEBUG, "Queueing frame for POC %d, GOP %d, dts %"PRId64", "
481			"generated from POC %d, GOP %d, dts %"PRId64", duration %"PRId64"\n",
482			frame.poc, frame.gop, out->pts,
483			poc_node->poc, poc_node->gop, poc_node->dts, poc_node->duration);
484			} else
485		✗	av_log(ctx, AV_LOG_WARNING, "No timestamp for POC %d in tree\n", frame.poc);
486			} else
487		✗	av_log(ctx, AV_LOG_WARNING, "No timestamp for POC %d in tree\n", frame.poc);
488		✗	av_log(ctx, AV_LOG_DEBUG, "Returning frame for POC %d, GOP %d, dts %"PRId64", pts %"PRId64"\n",
489			frame.poc, frame.gop, out->dts, out->pts);
490
491		✗	return 0;
492			}
493
494		✗	static void dts2pts_flush(AVBSFContext *ctx)
495			{
496		✗	DTS2PTSContext *s = ctx->priv_data;
497			DTS2PTSFrame frame;
498
499		✗	if (s->flush)
500		✗	s->flush(ctx);
501		✗	s->eof = 0;
502		✗	s->gop = 0;
503
504		✗	while (s->fifo && av_fifo_read(s->fifo, &frame, 1) >= 0)
505		✗	av_packet_free(&frame.pkt);
506
507		✗	av_tree_enumerate(s->root, NULL, NULL, free_node);
508		✗	av_tree_destroy(s->root);
509		✗	s->root = NULL;
510
511		✗	ff_cbs_fragment_reset(&s->au);
512		✗	if (s->cbc)
513		✗	ff_cbs_flush(s->cbc);
514			}
515
516		✗	static void dts2pts_close(AVBSFContext *ctx)
517			{
518		✗	DTS2PTSContext *s = ctx->priv_data;
519
520		✗	dts2pts_flush(ctx);
521
522		✗	av_fifo_freep2(&s->fifo);
523		✗	ff_cbs_fragment_free(&s->au);
524		✗	ff_cbs_close(&s->cbc);
525			}
526
527			static const enum AVCodecID dts2pts_codec_ids[] = {
528			AV_CODEC_ID_H264,
529			AV_CODEC_ID_NONE,
530			};
531
532			const FFBitStreamFilter ff_dts2pts_bsf = {
533			.p.name = "dts2pts",
534			.p.codec_ids = dts2pts_codec_ids,
535			.priv_data_size = sizeof(DTS2PTSContext),
536			.init = dts2pts_init,
537			.flush = dts2pts_flush,
538			.close = dts2pts_close,
539			.filter = dts2pts_filter,
540			};
541