From 4e18bb6e50da17cb005b05a89f571573d34aa156 Mon Sep 17 00:00:00 2001
From: derrod <dennis@obsproject.com>
Date: Tue, 26 Mar 2024 17:23:16 +0100
Subject: [PATCH] obs-outputs: Set AV1 packet priority

---
 plugins/obs-outputs/rtmp-av1.c | 51 ++++++++++++++++++++++++++++++++--
 1 file changed, 49 insertions(+), 2 deletions(-)

diff --git a/plugins/obs-outputs/rtmp-av1.c b/plugins/obs-outputs/rtmp-av1.c
index f8a66d997..0aca20e06 100644
--- a/plugins/obs-outputs/rtmp-av1.c
+++ b/plugins/obs-outputs/rtmp-av1.c
@@ -33,11 +33,19 @@
 #define AV1_OBU_TILE_GROUP 4
 #define AV1_OBU_TILE_LIST 8
 #define AV1_OBU_FRAME 6
+#define AV1_OBU_FRAME_HEADER 3
 
 #define FF_PROFILE_AV1_MAIN 0
 #define FF_PROFILE_AV1_HIGH 1
 #define FF_PROFILE_AV1_PROFESSIONAL 2
 
+enum frame_type {
+	AV1_KEY_FRAME,
+	AV1_INTER_FRAME,
+	AV1_INTRA_FRAME,
+	AV1_SWITCH_FRAME,
+};
+
 typedef struct AV1SequenceParameters {
 	uint8_t profile;
 	uint8_t level;
@@ -536,11 +544,45 @@ size_t obs_parse_av1_header(uint8_t **header, const uint8_t *data, size_t size)
 	return output.bytes.num;
 }
 
+static void compute_av1_keyframe_priority(const uint8_t *buf, bool *is_keyframe,
+					  int *priority)
+{
+	/* Skip if the packet already has a priority, e.g., assigned by the
+	 * encoder implementation (currently QSV/AMF). */
+	if (*priority)
+		return;
+
+	Av1GetBitContext gb;
+
+	init_get_bits8(&gb, buf, 1);
+
+	// show_existing_frame
+	if (get_bit1(&gb))
+		return;
+
+	enum frame_type type = get_bits(&gb, 2);
+	bool show_frame = get_bit1(&gb);
+
+	switch (type) {
+	case AV1_KEY_FRAME:
+		*is_keyframe = true;
+		*priority = 3;
+		break;
+	case AV1_INTER_FRAME:
+		*priority = show_frame ? 1 : 2;
+		break;
+	case AV1_INTRA_FRAME:
+		*priority = 3;
+		break;
+	case AV1_SWITCH_FRAME:
+		*priority = 2;
+		break;
+	}
+}
+
 static void serialize_av1_data(struct serializer *s, const uint8_t *data,
 			       size_t size, bool *is_keyframe, int *priority)
 {
-	(void)is_keyframe;
-	(void)priority;
 	uint8_t *buf = (uint8_t *)data;
 	uint8_t *end = (uint8_t *)data + size;
 
@@ -559,6 +601,11 @@ static void serialize_av1_data(struct serializer *s, const uint8_t *data,
 		case AV1_OBU_REDUNDANT_FRAME_HEADER:
 		case AV1_OBU_TILE_LIST:
 			break;
+		case AV1_OBU_FRAME:
+		case AV1_OBU_FRAME_HEADER:
+			compute_av1_keyframe_priority(buf + start_pos,
+						      is_keyframe, priority);
+			/* Falls through. */
 		default:
 			s_write(s, buf, len);
 			size += len;