From 6953ace80adba07fcbba5cfc99d865829a7b0f19 Mon Sep 17 00:00:00 2001
From: WyattBlue <wyattblue@auto-editor.com>
Date: Tue, 5 Mar 2024 19:07:48 -0500
Subject: [PATCH] Get color range info from PyAV

---
 auto_editor/ffwrapper.py        | 16 ++++++++--------
 auto_editor/formats/fcp11.py    | 10 +++++-----
 auto_editor/output.py           | 19 +++++++++++++------
 auto_editor/subcommands/info.py | 28 +++++++++++++++++++++-------
 pyproject.toml                  |  2 +-
 5 files changed, 48 insertions(+), 27 deletions(-)

diff --git a/auto_editor/ffwrapper.py b/auto_editor/ffwrapper.py
index 1f65ad80e..edb2b4750 100644
--- a/auto_editor/ffwrapper.py
+++ b/auto_editor/ffwrapper.py
@@ -135,10 +135,10 @@ class VideoStream:
     sar: Fraction
     time_base: Fraction | None
     pix_fmt: str | None
-    color_range: str | None
-    color_space: str | None
-    color_primaries: str | None
-    color_transfer: str | None
+    color_range: int
+    color_space: int
+    color_primaries: int
+    color_transfer: int
     bitrate: int
     lang: str | None
 
@@ -259,10 +259,10 @@ def initFileInfo(path: str, ffmpeg: FFmpeg, log: Log) -> FileInfo:
                 sar,
                 v.time_base,
                 v.codec_context.pix_fmt,
-                c_range,
-                c_space,
-                c_primary,
-                c_transfer,
+                v.color_range,
+                v.colorspace,
+                v.color_primaries,
+                v.color_trc,
                 0 if v.bit_rate is None else v.bit_rate,
                 v.language,
             ),
diff --git a/auto_editor/formats/fcp11.py b/auto_editor/formats/fcp11.py
index 053ac57e8..1dec587b9 100644
--- a/auto_editor/formats/fcp11.py
+++ b/auto_editor/formats/fcp11.py
@@ -32,13 +32,13 @@ def get_colorspace(src: FileInfo) -> str:
     s = src.videos[0]
     if s.pix_fmt == "rgb24":
         return "sRGB IEC61966-2.1"
-    if s.color_space == "smpte170m":
-        return "6-1-6 (Rec. 601 NTSC)"
-    if s.color_space == "bt470bg":
+    if s.color_space == 5:  # "bt470bg"
         return "5-1-6 (Rec. 601 PAL)"
-    if s.color_primaries == "bt2020":
+    if s.color_space == 6:  # "smpte170m"
+        return "6-1-6 (Rec. 601 NTSC)"
+    if s.color_primaries == 9:  # "bt2020"
         # See: https://video.stackexchange.com/questions/22059/how-to-identify-hdr-video
-        if s.color_transfer in ("arib-std-b67", "smpte2084"):
+        if s.color_transfer in (16, 18):  # "smpte2084" "arib-std-b67"
             return "9-18-9 (Rec. 2020 HLG)"
         return "9-1-9 (Rec. 2020)"
 
diff --git a/auto_editor/output.py b/auto_editor/output.py
index 4cb65c7e3..1a4f91d3a 100644
--- a/auto_editor/output.py
+++ b/auto_editor/output.py
@@ -182,12 +182,19 @@ def mux_quality_media(
         cmd += _ffset("-c:a", args.audio_codec) + _ffset("-b:a", args.audio_bitrate)
 
     if same_container and v_tracks > 0:
-        cmd += (
-            _ffset("-color_range", src.videos[0].color_range)
-            + _ffset("-colorspace", src.videos[0].color_space)
-            + _ffset("-color_primaries", src.videos[0].color_primaries)
-            + _ffset("-color_trc", src.videos[0].color_transfer)
-        )
+        color_range = src.videos[0].color_range
+        colorspace = src.videos[0].color_space
+        color_prim = src.videos[0].color_primaries
+        color_trc = src.videos[0].color_transfer
+
+        if color_range == 1 or color_range == 2:
+            cmd.extend(["-color_range", f"{color_range}"])
+        if colorspace in (0, 1) or (colorspace >= 3 and colorspace < 16):
+            cmd.extend(["-colorspace", f"{colorspace}"])
+        if color_prim in (0, 1) or (color_prim >= 4 and color_prim < 17):
+            cmd.extend(["-color_primaries", f"{color_prim}"])
+        if color_trc == 1 or (color_trc >= 4 and color_trc < 22):
+            cmd.extend(["-color_trc", f"{color_trc}"])
 
     if args.extras is not None:
         cmd.extend(args.extras.split(" "))
diff --git a/auto_editor/subcommands/info.py b/auto_editor/subcommands/info.py
index 27ccff31a..184078226 100644
--- a/auto_editor/subcommands/info.py
+++ b/auto_editor/subcommands/info.py
@@ -42,10 +42,10 @@ class VideoJson(TypedDict):
     pixel_aspect_ratio: str
     duration: float
     pix_fmt: str | None
-    color_range: str | None
-    color_space: str | None
-    color_primaries: str | None
-    color_transfer: str | None
+    color_range: int
+    color_space: int
+    color_primaries: int
+    color_transfer: int
     timebase: str
     bitrate: int
     lang: str | None
@@ -177,13 +177,27 @@ def stream_to_text(text: str, label: str, streams: list[dict[str, Any]]) -> str:
             text += f"   - track {s}:\n"
             for key, value in stream.items():
                 if not is_null(key, value):
-                    key = key.replace("_", " ")
                     if isinstance(value, list):
                         sep = "x" if key == "resolution" else ":"
-
                         value = sep.join(f"{x}" for x in value)
 
-                    text += f"     - {key}: {value}\n"
+                    if key in (
+                        "color_range",
+                        "color_space",
+                        "color_transfer",
+                        "color_primaries",
+                    ):
+                        if key == "color_range":
+                            if value == 1:
+                                text += "     - color range: 1 (tv)\n"
+                            elif value == 2:
+                                text += "     - color range: 2 (pc)\n"
+                        elif value == 1:
+                            text += f"     - {key.replace('_', ' ')}: 1 (bt709)\n"
+                        elif value != 2:
+                            text += f"     - {key.replace('_', ' ')}: {value}\n"
+                    else:
+                        text += f"     - {key.replace('_', ' ')}: {value}\n"
         return text
 
     text = ""
diff --git a/pyproject.toml b/pyproject.toml
index a6cae538b..d033f6f77 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -10,7 +10,7 @@ authors = [{ name = "WyattBlue", email = "wyattblue@auto-editor.com" }]
 requires-python = ">=3.10"
 dependencies = [
     "numpy>=1.22.0",
-    "pyav==12.0.2",
+    "pyav==12.0.4",
     "ae-ffmpeg==1.1.*",
 ]
 keywords = [