From 678b6fda19ef1cbf86ac2f64584508621fb36f53 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Gonzalo=20Garramu=C3=B1o?= <ggarra13@gmail.com>
Date: Sun, 10 Mar 2024 13:04:09 -0300
Subject: [PATCH] Fixed FFmpegReader's color pipeline.

Added tags for color_space, colorprimaries and color_trc.
---
 lib/tlCore/Image.cpp         | 10 ++++++----
 lib/tlIO/FFmpegReadVideo.cpp | 38 +++++++++++++++++++++++++++++++++---
 lib/tlIO/FFmpegWrite.cpp     |  3 +--
 3 files changed, 42 insertions(+), 9 deletions(-)

diff --git a/lib/tlCore/Image.cpp b/lib/tlCore/Image.cpp
index f866b74e..6da922a8 100644
--- a/lib/tlCore/Image.cpp
+++ b/lib/tlCore/Image.cpp
@@ -7,6 +7,7 @@
 #include <tlCore/Assert.h>
 #include <tlCore/Error.h>
 #include <tlCore/String.h>
+#include <tlCore/Locale.h>
 
 #include <algorithm>
 #include <array>
@@ -92,7 +93,10 @@ namespace tl
             const std::array<math::Vector4f, static_cast<size_t>(YUVCoefficients::Count)> data =
             {
                 math::Vector4f(1.79274, 2.1124, 0.213242, 0.532913),
-                math::Vector4f(1.67867, 2.14177, 0.187332, 0.650421)
+                // Is BT2020 right?  These are the coeffs I got:
+                // math::Vector4f(1.4746, 2.0184, 1.5958, 0.0);  // chatgpt3.5
+                // math::Vector4f(1.8556, 2.14101, 0.2124, 0.6301), // gemini
+                math::Vector4f(1.67867, 2.14177, 0.187332, 0.650421) // darby
             };
             return data[static_cast<size_t>(value)];
         }
@@ -355,10 +359,8 @@ namespace tl
             }
             if (2 == split.size())
             {
-                std::string savedLocale = std::setlocale(LC_NUMERIC, NULL);
-                std::setlocale(LC_NUMERIC, "C");
+                locale::SetAndRestore saved;
                 out.pixelAspectRatio = std::stof(split[1]);
-                std::setlocale(LC_NUMERIC, savedLocale.c_str());
             }
             split = string::split(split[0], 'x');
             if (split.size() != 2)
diff --git a/lib/tlIO/FFmpegReadVideo.cpp b/lib/tlIO/FFmpegReadVideo.cpp
index 688212b4..8ed77360 100644
--- a/lib/tlIO/FFmpegReadVideo.cpp
+++ b/lib/tlIO/FFmpegReadVideo.cpp
@@ -347,11 +347,12 @@ namespace tl
                     }
                     break;
                 }
-                if (_avCodecContext[_avStream]->color_range != AVCOL_RANGE_JPEG)
+                const auto params = _avCodecParameters[_avStream];
+                if (params->color_range != AVCOL_RANGE_JPEG)
                 {
                     _info.videoLevels = image::VideoLevels::LegalRange;
                 }
-                switch (_avCodecParameters[_avStream]->color_space)
+                switch (params->color_space)
                 {
                 case AVCOL_SPC_BT2020_NCL:
                     _info.yuvCoefficients = image::YUVCoefficients::BT2020;
@@ -456,6 +457,18 @@ namespace tl
                     _tags["Video Codec"] =
                         avcodec_get_name(_avCodecContext[_avStream]->codec_id);
                 }
+                {
+                    _tags["Video Color Primaries"] =
+                        av_color_primaries_name(params->color_primaries);
+                }
+                {
+                    _tags["Video Color TRC"] =
+                        av_color_transfer_name(params->color_trc);
+                }
+                {
+                    _tags["Video Color Space"] =
+                        av_color_space_name(params->color_space);
+                }
                 {
                     std::stringstream ss;
                     ss << _info.videoLevels;
@@ -599,6 +612,25 @@ namespace tl
                     {
                         throw std::runtime_error(string::Format("{0}: Cannot initialize sws context").arg(_fileName));
                     }
+                    
+                    const auto params = _avCodecParameters[_avStream];
+                    int in_full, out_full, brightness, contrast, saturation;
+                    const int *inv_table, *table;
+
+                    sws_getColorspaceDetails(
+                        _swsContext, (int**)&inv_table, &in_full,
+                        (int**)&table, &out_full, &brightness, &contrast,
+                        &saturation);
+                    
+                    inv_table = sws_getCoefficients(params->color_space);
+                    table = sws_getCoefficients(AVCOL_SPC_BT709);
+                
+                    in_full = (params->color_range == AVCOL_RANGE_JPEG);
+                    out_full = (params->color_range == AVCOL_RANGE_JPEG);
+
+                    sws_setColorspaceDetails(
+                        _swsContext, inv_table, in_full, table, out_full,
+                        brightness, contrast, saturation);
                 }
             }
         }
@@ -851,7 +883,7 @@ namespace tl
                     w,
                     h,
                     1);
-                
+                    
                 sws_scale(
                     _swsContext,
                     (uint8_t const* const*)_avFrame->data,
diff --git a/lib/tlIO/FFmpegWrite.cpp b/lib/tlIO/FFmpegWrite.cpp
index 9e14d95b..f2585891 100644
--- a/lib/tlIO/FFmpegWrite.cpp
+++ b/lib/tlIO/FFmpegWrite.cpp
@@ -1421,8 +1421,7 @@ namespace tl
                         (int**)&table, &out_full, &brightness, &contrast,
                         &saturation);
 
-                    inv_table =
-                        parseYUVType("bt709", p.avCodecContext->colorspace);
+                    inv_table = sws_getCoefficients(p.avCodecContext->colorspace);
                     table = parseYUVType("bt709", AVCOL_SPC_UNSPECIFIED);
 
                     // We use the full range, and we set -color_range to 2