From b0f96cdee999d1aa1592c1952569c6a69383bcc7 Mon Sep 17 00:00:00 2001 From: imSp4rky Date: Fri, 30 Oct 2026 17:13:14 -0600 Subject: [PATCH] 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 | 111 +++++++++++++++++++++++++++ unshackle/core/tracks/tracks.py | 4 +- unshackle/core/tracks/video.py | 99 ++++++++++++++++++++---- 3 files changed, 200 insertions(+), 14 deletions(-) create mode 100644 tests/tracks/test_vp9_colour_args.py diff --git a/tests/tracks/test_vp9_colour_args.py b/tests/tracks/test_vp9_colour_args.py new file mode 100644 index 00000000..a40fb2c6 --- /dev/null +++ b/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" diff --git a/unshackle/core/tracks/tracks.py b/unshackle/core/tracks/tracks.py index dd2d9cf1..12a3f006 100644 --- a/unshackle/core/tracks/tracks.py +++ b/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", diff --git a/unshackle/core/tracks/video.py b/unshackle/core/tracks/video.py index 1b9441de..a6765584 100644 --- a/unshackle/core/tracks/video.py +++ b/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",)