From b2dd35958908eb9a859b99697fad1fa7a809c2b7 Mon Sep 17 00:00:00 2001
From: Hans-Kristian Arntzen <post@themaister.net>
Date: Thu, 21 Nov 2024 12:02:00 +0100
Subject: [PATCH] Add most basic way of piping through PQ video encoding.

No metadata implemented yet.
---
 assets/shaders/util/rgb_to_yuv.comp | 11 +++++------
 video/ffmpeg_encode.cpp             | 30 +++++++++++++++++++++++++++--
 video/ffmpeg_encode.hpp             |  1 +
 3 files changed, 34 insertions(+), 8 deletions(-)

diff --git a/assets/shaders/util/rgb_to_yuv.comp b/assets/shaders/util/rgb_to_yuv.comp
index 7a9cd2a1..606c41fc 100644
--- a/assets/shaders/util/rgb_to_yuv.comp
+++ b/assets/shaders/util/rgb_to_yuv.comp
@@ -17,11 +17,10 @@ layout(push_constant) uniform Registers
     float dither_strength;
 } registers;
 
-// BT709. Row-major.
-const mediump mat3 rgb_to_yuv = mat3(
-    vec3(0.2126, 0.7152, 0.0722),
-    vec3(-0.114572, -0.385428, 0.5),
-    vec3(0.5, -0.454153, -0.0458471));
+layout(set = 0, binding = 3) uniform Transform
+{
+	mat3 rgb_to_yuv;
+};
 
 #define D(x) ((x) - 0.5)
 const mediump float dither[] = float[](
@@ -46,4 +45,4 @@ void main()
         imageStore(uLuma, ivec2(coord), ycbcr.xxxx);
         imageStore(uChroma, ivec2(coord), ycbcr.yzyz);
     }
-}
\ No newline at end of file
+}
diff --git a/video/ffmpeg_encode.cpp b/video/ffmpeg_encode.cpp
index 4ac01359..6e768168 100644
--- a/video/ffmpeg_encode.cpp
+++ b/video/ffmpeg_encode.cpp
@@ -1014,8 +1014,18 @@ bool VideoEncoder::Impl::init_video_codec_av(const AVCodec *codec)
 	}
 
 	video.av_ctx->color_range = AVCOL_RANGE_MPEG;
-	video.av_ctx->colorspace = AVCOL_SPC_BT709;
-	video.av_ctx->color_primaries = AVCOL_PRI_BT709;
+
+	if (options.hdr10)
+	{
+		video.av_ctx->colorspace = AVCOL_SPC_BT2020_NCL;
+		video.av_ctx->color_primaries = AVCOL_PRI_BT2020;
+		video.av_ctx->color_trc = AVCOL_TRC_SMPTE2084; // PQ
+	}
+	else
+	{
+		video.av_ctx->colorspace = AVCOL_SPC_BT709;
+		video.av_ctx->color_primaries = AVCOL_PRI_BT709;
+	}
 
 	switch (options.siting)
 	{
@@ -1736,6 +1746,22 @@ void VideoEncoder::process_rgb(Vulkan::CommandBuffer &cmd, YCbCrPipeline &pipeli
 	push.dither_strength = pipeline.constants.dither_strength;
 	cmd.push_constants(&push, 0, sizeof(push));
 
+	auto *color_transform = cmd.allocate_typed_constant_data<vec4>(0, 3, 3);
+
+	// Minimal crude implementation to get something HDR10 working.
+	if (impl->options.hdr10)
+	{
+		color_transform[0] = vec4(0.2627f, 0.678f, 0.0593f, 0.0f);
+		color_transform[1] = vec4(-0.13963f, -0.36037f, 0.5f, 0.0f);
+		color_transform[2] = vec4(0.5f, -0.459786f, -0.0402143f, 0.0f);
+	}
+	else
+	{
+		color_transform[0] = vec4(0.2126f, 0.7152f, 0.0722f, 0.0f);
+		color_transform[1] = vec4(-0.114572f, -0.385428f, 0.5f, 0.0f);
+		color_transform[2] = vec4(0.5f, -0.454153f, -0.0458471f, 0.0f);
+	}
+
 	auto start_yuv_ts = cmd.write_timestamp(VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT);
 	cmd.dispatch(pipeline.constants.luma_dispatch[0], pipeline.constants.luma_dispatch[1], 1);
 	auto end_yuv_ts = cmd.write_timestamp(VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT);
diff --git a/video/ffmpeg_encode.hpp b/video/ffmpeg_encode.hpp
index 17769078..d5632815 100644
--- a/video/ffmpeg_encode.hpp
+++ b/video/ffmpeg_encode.hpp
@@ -82,6 +82,7 @@ class VideoEncoder
 		bool realtime = false;
 		const char *encoder = "libx264";
 		bool low_latency = false;
+		bool hdr10 = false;
 
 		struct
 		{