fix(renderer): preserve video colors in annotations

src/pozify/steps/annotated_renderer.py

@@ -1,6 +1,9 @@
 from __future__ import annotations
 
+import json
 from pathlib import Path
+import shutil
+import subprocess
 from typing import Any
 
 import cv2
@@ -24,6 +27,186 @@ SKELETON_EDGES = [
 ]
 
 PREFERRED_VIDEO_CODECS = ("avc1", "H264", "mp4v")
+HDR_TRANSFERS = {"arib-std-b67", "smpte2084"}
+HDR_PRIMARIES = {"bt2020"}
+BT709_COLOR_ARGS = (
+    "-color_primaries",
+    "bt709",
+    "-color_trc",
+    "bt709",
+    "-colorspace",
+    "bt709",
+)
+
+
+def _tool_path(name: str) -> str | None:
+    return shutil.which(name)
+
+
+def _video_color_metadata(video_path: str) -> dict[str, str]:
+    ffprobe = _tool_path("ffprobe")
+    if ffprobe is None:
+        return {}
+
+    command = [
+        ffprobe,
+        "-v",
+        "error",
+        "-select_streams",
+        "v:0",
+        "-show_entries",
+        "stream=color_space,color_transfer,color_primaries,color_range",
+        "-of",
+        "json",
+        video_path,
+    ]
+    try:
+        result = subprocess.run(
+            command,
+            check=True,
+            capture_output=True,
+            text=True,
+            timeout=10,
+        )
+        payload = json.loads(result.stdout)
+    except (subprocess.SubprocessError, json.JSONDecodeError, OSError):
+        return {}
+
+    streams = payload.get("streams")
+    if not isinstance(streams, list) or not streams:
+        return {}
+    stream = streams[0]
+    if not isinstance(stream, dict):
+        return {}
+    metadata: dict[str, str] = {}
+    for key, value in stream.items():
+        if value is None or isinstance(value, (dict, list)):
+            continue
+        normalized_value = str(value).lower()
+        if normalized_value == "unknown":
+            continue
+        metadata[key] = normalized_value
+    return metadata
+
+
+def _needs_sdr_conversion(color_metadata: dict[str, str]) -> bool:
+    transfer = color_metadata.get("color_transfer", "")
+    primaries = color_metadata.get("color_primaries", "")
+    return transfer in HDR_TRANSFERS or primaries in HDR_PRIMARIES
+
+
+def _sdr_filter(color_metadata: dict[str, str]) -> str:
+    transfer = color_metadata.get("color_transfer", "arib-std-b67")
+    primaries = color_metadata.get("color_primaries", "bt2020")
+    matrix = color_metadata.get("color_space", "bt2020nc")
+    if transfer not in HDR_TRANSFERS:
+        transfer = "arib-std-b67"
+    if primaries not in HDR_PRIMARIES:
+        primaries = "bt2020"
+    if matrix not in {"bt2020nc", "bt2020c"}:
+        matrix = "bt2020nc"
+
+    return (
+        f"zscale=transfer=linear:transferin={transfer}:"
+        f"primariesin={primaries}:matrixin={matrix}:npl=100,"
+        "tonemap=tonemap=hable:desat=0,"
+        "zscale=transfer=bt709:primaries=bt709:matrix=bt709:range=tv,"
+        "format=yuv420p"
+    )
+
+
+def _transcode_hdr_to_sdr(
+    input_path: Path,
+    output_path: Path,
+    color_metadata: dict[str, str],
+) -> bool:
+    ffmpeg = _tool_path("ffmpeg")
+    if ffmpeg is None:
+        return False
+
+    command = [
+        ffmpeg,
+        "-y",
+        "-v",
+        "error",
+        "-i",
+        str(input_path),
+        "-vf",
+        _sdr_filter(color_metadata),
+        "-an",
+        "-c:v",
+        "libx264",
+        "-preset",
+        "veryfast",
+        "-crf",
+        "18",
+        "-pix_fmt",
+        "yuv420p",
+        *BT709_COLOR_ARGS,
+        str(output_path),
+    ]
+    try:
+        subprocess.run(command, check=True, capture_output=True, timeout=120)
+    except (subprocess.SubprocessError, OSError):
+        return False
+    return output_path.exists() and output_path.stat().st_size > 0
+
+
+def _encode_bt709_output(
+    raw_video_path: Path,
+    output_path: Path,
+    audio_source_path: Path | None,
+) -> bool:
+    ffmpeg = _tool_path("ffmpeg")
+    if ffmpeg is None:
+        return False
+
+    command = [
+        ffmpeg,
+        "-y",
+        "-v",
+        "error",
+        "-i",
+        str(raw_video_path),
+    ]
+    if audio_source_path is not None:
+        command.extend(["-i", str(audio_source_path)])
+
+    command.extend(
+        [
+            "-map",
+            "0:v:0",
+        ]
+    )
+    if audio_source_path is not None:
+        command.extend(["-map", "1:a?"])
+
+    command.extend(
+        [
+            "-c:v",
+            "libx264",
+            "-preset",
+            "veryfast",
+            "-crf",
+            "18",
+            "-vf",
+            "setparams=color_primaries=bt709:color_trc=bt709:colorspace=bt709,format=yuv420p",
+            "-pix_fmt",
+            "yuv420p",
+            *BT709_COLOR_ARGS,
+        ]
+    )
+    if audio_source_path is not None:
+        command.extend(["-c:a", "aac", "-b:a", "128k", "-shortest"])
+    else:
+        command.append("-an")
+    command.append(str(output_path))
+
+    try:
+        subprocess.run(command, check=True, capture_output=True, timeout=120)
+    except (subprocess.SubprocessError, OSError):
+        return False
+    return output_path.exists() and output_path.stat().st_size > 0
 
 
 def _frame_landmark_points(
@@ -103,7 +286,12 @@ def _draw_overlays(
         )
 
 
-def _open_video_writer(output_path: Path, fps: float, width: int, height: int) -> tuple[cv2.VideoWriter | None, str | None]:
+def _open_video_writer(
+    output_path: Path,
+    fps: float,
+    width: int,
+    height: int,
+) -> tuple[cv2.VideoWriter | None, str | None]:
     for codec in PREFERRED_VIDEO_CODECS:
         writer = cv2.VideoWriter(
             str(output_path),
@@ -127,9 +315,21 @@ def run(
     if not manifest.analysis_allowed or not manifest.video_path:
         return manifest.video_path
 
-    capture = cv2.VideoCapture(manifest.video_path)
+    source_path = Path(manifest.video_path)
+    color_metadata = _video_color_metadata(manifest.video_path)
+    render_input_path = source_path
+    temporary_paths: list[Path] = []
+    if _needs_sdr_conversion(color_metadata):
+        sdr_input_path = run_dir / "renderer_sdr_input.mp4"
+        if _transcode_hdr_to_sdr(source_path, sdr_input_path, color_metadata):
+            render_input_path = sdr_input_path
+            temporary_paths.append(sdr_input_path)
+
+    capture = cv2.VideoCapture(str(render_input_path))
     if not capture.isOpened():
         capture.release()
+        for temporary_path in temporary_paths:
+            temporary_path.unlink(missing_ok=True)
         return manifest.video_path
 
     fps = manifest.fps if manifest.fps > 0 else 30.0
@@ -140,9 +340,12 @@ def run(
         return manifest.video_path
 
     output_path = run_dir / "annotated_video.mp4"
-    writer, _codec = _open_video_writer(output_path, fps, width, height)
+    raw_output_path = run_dir / "annotated_video_raw.mp4" if _tool_path("ffmpeg") else output_path
+    writer, _codec = _open_video_writer(raw_output_path, fps, width, height)
     if writer is None:
         capture.release()
+        for temporary_path in temporary_paths:
+            temporary_path.unlink(missing_ok=True)
         return manifest.video_path
 
     pose_by_frame = {frame.frame_index: frame for frame in pose_sequence.frames}
@@ -180,5 +383,13 @@ def run(
     finally:
         writer.release()
         capture.release()
+        for temporary_path in temporary_paths:
+            temporary_path.unlink(missing_ok=True)
+
+    if raw_output_path != output_path:
+        encoded = _encode_bt709_output(raw_output_path, output_path, source_path)
+        if not encoded:
+            raw_output_path.replace(output_path)
+        raw_output_path.unlink(missing_ok=True)
 
     return str(output_path)

tests/test_annotated_renderer.py

@@ -95,6 +95,26 @@ class AnnotatedRendererTests(unittest.TestCase):
         self.assertIsNotNone(output_path)
         self.assertTrue(Path(str(output_path)).exists())
 
+    def test_hdr_metadata_requires_sdr_conversion(self) -> None:
+        self.assertTrue(
+            annotated_renderer._needs_sdr_conversion(
+                {
+                    "color_space": "bt2020nc",
+                    "color_transfer": "arib-std-b67",
+                    "color_primaries": "bt2020",
+                }
+            )
+        )
+        self.assertFalse(
+            annotated_renderer._needs_sdr_conversion(
+                {
+                    "color_space": "bt709",
+                    "color_transfer": "bt709",
+                    "color_primaries": "bt709",
+                }
+            )
+        )
+
 
 if __name__ == "__main__":
     unittest.main()