fix(tracks): keep HDR colour metadata when muxing VP9

mkvmerge drops the MP4 vpcC/SmDm/CoLL boxes of VP9, so pass the colour codes and mastering display data from an ffprobe of the source.

tests/tracks/test_vp9_colour_args.py

@@ -0,0 +1,111 @@
+import json
+import subprocess
+from pathlib import Path
+from types import SimpleNamespace
+from typing import Any, Callable, Optional
+
+import pytest
+
+from unshackle.core import binaries
+from unshackle.core.tracks import Tracks, Video
+
+MASTERING = {
+    "side_data_type": "Mastering display metadata",
+    "red_x": "46400/65536",
+    "red_y": "19137/65536",
+    "green_x": "11142/65536",
+    "green_y": "52233/65536",
+    "blue_x": "8586/65536",
+    "blue_y": "3015/65536",
+    "white_point_x": "20494/65536",
+    "white_point_y": "21562/65536",
+    "min_luminance": "2/16384",
+    "max_luminance": "256001/256",
+}
+NO_LIGHT_LEVEL = {"side_data_type": "Content light level metadata", "max_content": 0, "max_average": 0}
+
+
+def make_video(tmp_path: Path, codec: Video.Codec, range_: Video.Range) -> Video:
+    video = Video(url="https://example.test/v.m3u8", language="en", codec=codec, range_=range_)
+    video.path = tmp_path / "v.mp4"
+    video.path.write_bytes(b"\x00" * 16)
+    return video
+
+
+def probe_returns(monkeypatch: pytest.MonkeyPatch, stream: dict[str, Any]) -> None:
+    monkeypatch.setattr(binaries, "FFProbe", Path("ffprobe"))
+    stdout = json.dumps({"streams": [stream]})
+    monkeypatch.setattr(subprocess, "run", lambda *a, **k: SimpleNamespace(stdout=stdout, returncode=0))
+
+
+def pairs(args: list[str]) -> dict[str, str]:
+    return dict(zip(args[::2], args[1::2]))
+
+
+def test_hdr10_vp9_gets_cicp_and_mastering_display(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
+    probe_returns(monkeypatch, {"color_transfer": "smpte2084", "side_data_list": [MASTERING, NO_LIGHT_LEVEL]})
+
+    args = pairs(make_video(tmp_path, Video.Codec.VP9, Video.Range.HDR10).vp9_colour_args(0))
+
+    assert args["--colour-primaries"] == "0:9"
+    assert args["--colour-transfer-characteristics"] == "0:16"
+    assert args["--colour-matrix-coefficients"] == "0:9"
+    assert args["--chromaticity-coordinates"] == "0:0.70801,0.29201,0.17001,0.79701,0.13101,0.04601"
+    assert args["--white-colour-coordinates"] == "0:0.31271,0.32901"
+    assert args["--max-luminance"] == "0:1000.003906"
+    assert args["--min-luminance"] == "0:0.000122"
+    assert "--max-content-light" not in args  # a zero CoLL means unknown
+
+
+def test_partial_mastering_display_keeps_the_cicp_args(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
+    partial = {"side_data_type": "Mastering display metadata", "min_luminance": "1/10000", "max_luminance": "1000/1"}
+    probe_returns(monkeypatch, {"color_transfer": "smpte2084", "side_data_list": [partial]})
+
+    args = pairs(make_video(tmp_path, Video.Codec.VP9, Video.Range.HDR10).vp9_colour_args(1))
+
+    assert args["--colour-transfer-characteristics"] == "1:16"
+    assert "--max-luminance" not in args
+
+
+@pytest.mark.parametrize(
+    ("transfer", "range_", "expected"),
+    [
+        ("smpte2084", Video.Range.SDR, "0:16"),  # an unlabelled manifest still gets the file's HDR
+        ("arib-std-b67", Video.Range.SDR, "0:18"),
+        ("bt709", Video.Range.HDR10, None),  # a wrong HDR label never stamps PQ on SDR
+        ("unknown", Video.Range.HLG, "0:18"),
+        ("unknown", Video.Range.SDR, None),
+    ],
+)
+def test_the_file_transfer_wins_over_the_range(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch, transfer: str, range_: Video.Range, expected: Optional[str]
+) -> None:
+    probe_returns(monkeypatch, {"color_transfer": transfer})
+
+    args = pairs(make_video(tmp_path, Video.Codec.VP9, range_).vp9_colour_args(0))
+
+    assert args.get("--colour-transfer-characteristics") == expected
+
+
+def test_no_ffprobe_falls_back_to_the_range(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(binaries, "FFProbe", None)
+
+    args = pairs(make_video(tmp_path, Video.Codec.VP9, Video.Range.HDR10).vp9_colour_args(0))
+
+    assert args == {
+        "--colour-primaries": "0:9",
+        "--colour-transfer-characteristics": "0:16",
+        "--colour-matrix-coefficients": "0:9",
+    }
+
+
+def test_mux_passes_the_vp9_args_and_keeps_the_hevc_hlg_flag(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch, mux_args: Callable[..., list[str]]
+) -> None:
+    monkeypatch.setattr(binaries, "FFProbe", None)
+
+    vp9 = mux_args(Tracks(make_video(tmp_path, Video.Codec.VP9, Video.Range.HDR10)), {"video": 0})
+    hevc = mux_args(Tracks(make_video(tmp_path, Video.Codec.HEVC, Video.Range.HLG)), {"video": 0})
+
+    assert vp9[vp9.index("--colour-transfer-characteristics") + 1] == "0:16"
+    assert hevc[hevc.index("--color-transfer-characteristics") + 1] == "0:18"

unshackle/core/tracks/tracks.py

@@ -899,7 +899,9 @@ class Tracks:
                     ]
                 )
 
-            if hasattr(vt, "range") and vt.range == Video.Range.HLG:
+            if vt.codec == Video.Codec.VP9:
+                video_args.extend(vt.vp9_colour_args(v))
+            elif hasattr(vt, "range") and vt.range == Video.Range.HLG:
                 video_args.extend(
                     [
                         "--color-transfer-characteristics",

unshackle/core/tracks/video.py

@@ -1,11 +1,13 @@
 from __future__ import annotations
 
+import json
 import logging
 import math
 import re
 import subprocess
 import time
 from enum import Enum
+from fractions import Fraction
 from pathlib import Path
 from typing import Any, Optional, Union
 
@@ -483,22 +485,14 @@ class Video(Track):
         if self.range in (Video.Range.SDR, Video.Range.DV, Video.Range.HYBRID):
             return None
 
-        vui = {
-            Video.Range.HDR10: (9, 16, 9),
-            Video.Range.HDR10P: (9, 16, 9),
-            Video.Range.HLG: (9, 18, 9),
-        }.get(self.range)
-        if not vui:
+        transfer_name = HDR_TRANSFER.get(self.range)
+        if not transfer_name:
             return None
 
         # Skip the rewrite when the bitstream VUI already matches the target, which avoids a full
         # bitstream pass on services that already ship correct colour metadata
-        expected = {
-            Video.Range.HDR10: ("bt2020", "smpte2084", "bt2020nc"),
-            Video.Range.HDR10P: ("bt2020", "smpte2084", "bt2020nc"),
-            Video.Range.HLG: ("bt2020", "arib-std-b67", "bt2020nc"),
-        }.get(self.range)
-        if expected and binaries.FFProbe:
+        expected = ("bt2020", transfer_name, "bt2020nc")
+        if binaries.FFProbe:
             probe = subprocess.run(
                 [
                     binaries.FFProbe,
@@ -522,7 +516,7 @@ class Video(Track):
             if current == expected:
                 return None
 
-        primaries, transfer, matrix = vui
+        primaries, transfer, matrix = TRANSFER_CICP[transfer_name]
         filter_key = {Video.Codec.AVC: "h264_metadata", Video.Codec.HEVC: "hevc_metadata"}[self.codec]
         return (
             f"{filter_key}=colour_primaries={primaries}"
@@ -590,6 +584,78 @@ class Video(Track):
         )
         return True
 
+    def vp9_colour_args(self, track_id: int) -> list[str]:
+        """Return mkvmerge colour args for an HDR VP9 track.
+
+        mkvmerge drops the MP4 vpcC/SmDm/CoLL colour boxes of VP9, which leaves an HDR10 file tagged as SDR.
+        The transfer the file states wins; ``self.range`` applies only when the file states none.
+        """
+        stream = self._probe_colour()
+        transfer = stream.get("color_transfer", "unknown")
+        codes = TRANSFER_CICP.get(transfer if transfer != "unknown" else HDR_TRANSFER.get(self.range, ""))
+        if not codes:
+            return []
+
+        opts: dict[str, object] = dict(
+            zip(("--colour-primaries", "--colour-transfer-characteristics", "--colour-matrix-coefficients"), codes)
+        )
+        for data in stream.get("side_data_list", []):
+            # ffprobe leaves out the fields a source does not carry, so a partial entry adds nothing
+            try:
+                if data.get("side_data_type") == "Mastering display metadata":
+                    xy = ("red_x", "red_y", "green_x", "green_y", "blue_x", "blue_y", "white_point_x", "white_point_y")
+                    coords = [f"{float(Fraction(data[key])):.5f}" for key in xy]
+                    opts.update(
+                        {
+                            "--chromaticity-coordinates": ",".join(coords[:6]),
+                            "--white-colour-coordinates": ",".join(coords[6:]),
+                            "--max-luminance": f"{float(Fraction(data['max_luminance'])):.6f}",
+                            "--min-luminance": f"{float(Fraction(data['min_luminance'])):.6f}",
+                        }
+                    )
+                elif data.get("side_data_type") == "Content light level metadata" and data.get("max_content"):
+                    opts.update(
+                        {"--max-content-light": int(data["max_content"]), "--max-frame-light": int(data["max_average"])}
+                    )
+            except (KeyError, TypeError, ValueError, ZeroDivisionError):
+                continue
+        return [arg for flag, value in opts.items() for arg in (flag, f"{track_id}:{value}")]
+
+    def _probe_colour(self) -> dict[str, Any]:
+        """Return the FFprobe colour transfer and side data of the first video track, or {} when it cannot probe."""
+        if not binaries.FFProbe or not self.path or not self.path.exists():
+            return {}
+        probe_start = time.monotonic()
+        probe = subprocess.run(
+            [
+                binaries.FFProbe,
+                "-v",
+                "error",
+                # without it ffprobe walks every moof of a fragmented MP4 to read the header
+                "-seekable",
+                "0",
+                "-select_streams",
+                "v:0",
+                "-show_entries",
+                "stream=color_transfer:stream_side_data",
+                "-of",
+                "json",
+                str(self.path),
+            ],
+            capture_output=True,
+            text=True,
+            encoding="utf-8",
+            errors="replace",
+        )
+        log_tool_run(
+            "ffprobe colour", "ffprobe", probe.returncode, duration_ms=round((time.monotonic() - probe_start) * 1000, 1)
+        )
+        try:
+            stream: dict[str, Any] = json.loads(probe.stdout)["streams"][0]
+        except (ValueError, KeyError, IndexError, TypeError):
+            return {}
+        return stream
+
     def ccextractor(
         self, track_id: Any, out_path: Union[Path, str], language: Language, original: bool = False
     ) -> Optional[Subtitle]:
@@ -764,4 +830,11 @@ class Video(Track):
         return True
 
 
+HDR_TRANSFER = {
+    Video.Range.HDR10: "smpte2084",
+    Video.Range.HDR10P: "smpte2084",
+    Video.Range.HLG: "arib-std-b67",
+}
+TRANSFER_CICP = {"smpte2084": (9, 16, 9), "arib-std-b67": (9, 18, 9)}
+
 __all__ = ("Video",)