Fix YUV to RGB Color Matrix Selection

Util/FFmpeg/FFmpegNativeDecoder.cs

@@ -206,6 +206,12 @@ internal sealed unsafe class FFmpegNativeDecoder : IDisposable
     private int _srcWidth;
     private int _srcHeight;
 
+    /// <summary>YUV→RGB matrix id (<c>SWS_CS_*</c>) applied to the context. 0 when no matrix is set (non-RGBA output).</summary>
+    private int _srcSwsColorspace;
+
+    /// <summary>Whether the source range applied to the context is the full range.</summary>
+    private bool _srcFullRange;
+
     // ──────────────────────────────────────────────────────────────
     // Public Properties
     // ──────────────────────────────────────────────────────────────
@@ -296,6 +302,59 @@ internal sealed unsafe class FFmpegNativeDecoder : IDisposable
         }
     }
 
+    /// <summary>
+    /// Color matrix id used to convert the YUV frames of this stream to RGB, selected from the colorspace tag of the stream.
+    /// The ids are shared with the <c>ColorMatrix</c> uniform of ColorConversion.compute:
+    /// 0 = BT.709 (also used when the stream has no usable tag), 1 = BT.601, 2 = BT.2020 (non-constant luminance).
+    /// </summary>
+    public int ColorMatrix
+    {
+        get
+        {
+            lock (_lock)
+            {
+                if (_pCodecContext != null && _pCodecContext->colorspace != AVColorSpace.AVCOL_SPC_UNSPECIFIED)
+                {
+                    return GetColorMatrixId(_pCodecContext->colorspace);
+                }
+                if (_pFormatContext != null && _videoStreamIndex >= 0)
+                {
+                    return GetColorMatrixId(_pFormatContext->streams[_videoStreamIndex]->codecpar->color_space);
+                }
+                return ColorMatrixBT709;
+            }
+        }
+    }
+
+    /// <summary>Color matrix id of BT.709.</summary>
+    // ReSharper disable InconsistentNaming
+    public const int ColorMatrixBT709 = 0;
+
+    /// <summary>Color matrix id of BT.601 (SMPTE 170M / BT.470 BG).</summary>
+    public const int ColorMatrixBT601 = 1;
+
+    /// <summary>Color matrix id of BT.2020 (non-constant luminance).</summary>
+    public const int ColorMatrixBT2020 = 2;
+    // ReSharper restore InconsistentNaming
+
+    /// <summary>
+    /// Converts a colorspace tag to a color matrix id (see <see cref="ColorMatrix"/>).
+    /// Tags that have no matrix of their own (unspecified, RGB, other standards) use BT.709.
+    /// </summary>
+    private static int GetColorMatrixId(AVColorSpace colorspace)
+    {
+        switch (colorspace)
+        {
+            case AVColorSpace.AVCOL_SPC_SMPTE170M:
+            case AVColorSpace.AVCOL_SPC_BT470BG:
+                return ColorMatrixBT601;
+            case AVColorSpace.AVCOL_SPC_BT2020_NCL:
+                return ColorMatrixBT2020;
+            default:
+                return ColorMatrixBT709;
+        }
+    }
+
     /// <summary>
     /// Gets a value indicating whether hardware-accelerated decoding is currently active.
     /// </summary>
@@ -778,12 +837,67 @@ internal sealed unsafe class FFmpegNativeDecoder : IDisposable
     }
 
     /// <summary>
-    /// Lazily initializes or updates the <see cref="SwsContext"/> when source/destination formats or dimensions change.
+    /// Selects the YUV→RGB conversion matrix (libswscale <c>SWS_CS_*</c>) from the colorspace tag of a frame.
+    /// The matrix set is the same as that of the compute shader path (ColorConversion.compute, <c>ColorMatrix</c> uniform):
+    /// BT.709 (also for frames without a usable tag), BT.601 and BT.2020, so both paths give the same colors.
+    /// Reference: https://ffmpeg.org/doxygen/trunk/group__libsws.html (sws_getCoefficients)
+    /// </summary>
+    private static int GetSwsColorspace(AVColorSpace colorspace)
+    {
+        switch (GetColorMatrixId(colorspace))
+        {
+            case ColorMatrixBT601:
+                return ffmpeg.SWS_CS_ITU601;
+            case ColorMatrixBT2020:
+                return ffmpeg.SWS_CS_BT2020;
+            default:
+                return ffmpeg.SWS_CS_ITU709;
+        }
+    }
+
+    /// <summary>
+    /// Returns true when the Y/U/V samples of the frame use the full range (0-255) instead of the limited range (16-235 / 16-240).
+    /// The range is given by the frame's color_range tag or by a deprecated full range pixel format (YUVJ*).
+    /// </summary>
+    private static bool IsFullRange(AVFrame* frame)
+    {
+        if (frame->color_range == AVColorRange.AVCOL_RANGE_JPEG) return true;
+
+        switch ((AVPixelFormat)frame->format)
+        {
+            case AVPixelFormat.AV_PIX_FMT_YUVJ420P:
+            case AVPixelFormat.AV_PIX_FMT_YUVJ422P:
+            case AVPixelFormat.AV_PIX_FMT_YUVJ444P:
+            case AVPixelFormat.AV_PIX_FMT_YUVJ440P:
+                return true;
+            default:
+                return false;
+        }
+    }
+
+    /// <summary>
+    /// Lazily initializes or updates the <see cref="SwsContext"/> when source/destination formats, dimensions,
+    /// or (for RGBA output) the colorspace / range of the source frame change.
+    /// For RGBA output the YUV→RGB matrix and the source range are taken from <paramref name="colorFrame"/>.
     /// Reference: https://ffmpeg.org/doxygen/trunk/group__libsws.html
     /// </summary>
-    private void InitOrUpdateSwsContext(AVPixelFormat srcFormat, int srcWidth, int srcHeight, AVPixelFormat dstFormat)
+    /// <param name="srcFormat">Pixel format of the source frame.</param>
+    /// <param name="srcWidth">Width of the source frame in pixels.</param>
+    /// <param name="srcHeight">Height of the source frame in pixels.</param>
+    /// <param name="dstFormat">Pixel format of the output.</param>
+    /// <param name="colorFrame">Source frame whose colorspace / range tags are applied. Required for RGBA output.</param>
+    private void InitOrUpdateSwsContext(AVPixelFormat srcFormat, int srcWidth, int srcHeight, AVPixelFormat dstFormat, AVFrame* colorFrame = null)
     {
-        if (_pSwsContext != null && _srcPixFmt == srcFormat && _srcWidth == srcWidth && _srcHeight == srcHeight && _dstPixFmt == dstFormat)
+        var swsColorspace = 0;
+        var fullRange = false;
+        if (colorFrame != null && dstFormat == AVPixelFormat.AV_PIX_FMT_RGBA)
+        {
+            swsColorspace = GetSwsColorspace(colorFrame->colorspace);
+            fullRange = IsFullRange(colorFrame);
+        }
+
+        if (_pSwsContext != null && _srcPixFmt == srcFormat && _srcWidth == srcWidth && _srcHeight == srcHeight && _dstPixFmt == dstFormat
+            && _srcSwsColorspace == swsColorspace && _srcFullRange == fullRange)
         {
             return;
         }
@@ -798,6 +912,8 @@ internal sealed unsafe class FFmpegNativeDecoder : IDisposable
         _srcWidth = srcWidth;
         _srcHeight = srcHeight;
         _dstPixFmt = dstFormat;
+        _srcSwsColorspace = swsColorspace;
+        _srcFullRange = fullRange;
 
         _pSwsContext = ffmpeg.sws_getContext(
             srcWidth, srcHeight, srcFormat,
@@ -809,6 +925,44 @@ internal sealed unsafe class FFmpegNativeDecoder : IDisposable
         {
             throw new Exception($"Failed to initialize SwsContext (src: {srcFormat} {srcWidth}x{srcHeight} -> dst: {dstFormat} {Width}x{Height}).");
         }
+
+        if (swsColorspace != 0)
+        {
+            ApplySwsColorspace(swsColorspace, fullRange);
+        }
+    }
+
+    /// <summary>
+    /// Sets the YUV→RGB matrix and the source range of <see cref="_pSwsContext"/>.
+    /// The destination range, brightness, contrast and saturation keep the values of the context.
+    /// Reference: https://ffmpeg.org/doxygen/trunk/group__libsws.html (sws_setColorspaceDetails)
+    /// </summary>
+    /// <param name="swsColorspace">Matrix id (<c>SWS_CS_*</c>) of the source.</param>
+    /// <param name="fullRange">True when the source uses the full range.</param>
+    private void ApplySwsColorspace(int swsColorspace, bool fullRange)
+    {
+        int* currentInvTable;
+        int* currentTable;
+        int srcRange, dstRange, brightness, contrast, saturation;
+
+        // Returns a negative value when the context cannot change its colorspace details (e.g. YUV to YUV conversion).
+        if (ffmpeg.sws_getColorspaceDetails(_pSwsContext, &currentInvTable, &srcRange, &currentTable, &dstRange,
+                &brightness, &contrast, &saturation) < 0)
+        {
+            return;
+        }
+
+        var coefficients = ffmpeg.sws_getCoefficients(swsColorspace);
+        var invTable = new int_array4();
+        var table = new int_array4();
+        for (uint i = 0; i < 4; i++)
+        {
+            invTable[i] = coefficients[i];
+            table[i] = currentTable[i];
+        }
+
+        ffmpeg.sws_setColorspaceDetails(_pSwsContext, invTable, fullRange ? 1 : 0, table, dstRange,
+            brightness, contrast, saturation);
     }
 
     /// <summary>
@@ -900,7 +1054,7 @@ internal sealed unsafe class FFmpegNativeDecoder : IDisposable
             if (frame == null) return false;
 
             timestamp = GetTimestampAndUpdateIndex(frame);
-            InitOrUpdateSwsContext((AVPixelFormat)frame->format, frame->width, frame->height, AVPixelFormat.AV_PIX_FMT_RGBA);
+            InitOrUpdateSwsContext((AVPixelFormat)frame->format, frame->width, frame->height, AVPixelFormat.AV_PIX_FMT_RGBA, frame);
 
             ClearConvertedFramePointers();
             if (vflip)