115 lines
3.8 KiB
Python
115 lines
3.8 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import subprocess
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
|
|
def extract_hdr10plus_metadata(
|
|
*,
|
|
ffmpeg: Path,
|
|
hdr10plus_tool: Path,
|
|
source: Path,
|
|
output: Path,
|
|
) -> None:
|
|
"""Extract HDR10+ JSON without creating a full intermediate video file."""
|
|
output.parent.mkdir(parents=True, exist_ok=True)
|
|
ffmpeg_process = subprocess.Popen(
|
|
[
|
|
str(ffmpeg),
|
|
"-v", "error", "-i", str(source), "-map", "0:v:0",
|
|
"-c", "copy", "-bsf:v", "hevc_mp4toannexb", "-f", "hevc", "-",
|
|
],
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.PIPE,
|
|
)
|
|
assert ffmpeg_process.stdout is not None
|
|
try:
|
|
extracted = subprocess.run(
|
|
[str(hdr10plus_tool), "extract", "-o", str(output), "-"],
|
|
stdin=ffmpeg_process.stdout,
|
|
capture_output=True,
|
|
text=True,
|
|
)
|
|
finally:
|
|
ffmpeg_process.stdout.close()
|
|
ffmpeg_stderr = ffmpeg_process.stderr.read().decode(errors="replace")
|
|
ffmpeg_process.stderr.close()
|
|
ffmpeg_returncode = ffmpeg_process.wait()
|
|
if ffmpeg_returncode != 0:
|
|
raise RuntimeError(
|
|
f"FFmpeg could not extract the HEVC stream from {source}:\n{ffmpeg_stderr.strip()}"
|
|
)
|
|
if extracted.returncode != 0:
|
|
raise RuntimeError(
|
|
f"hdr10plus_tool could not extract metadata from {source}:\n"
|
|
f"{extracted.stderr.strip() or extracted.stdout.strip()}"
|
|
)
|
|
if not output.is_file():
|
|
raise RuntimeError(f"hdr10plus_tool did not create metadata JSON for {source}")
|
|
|
|
|
|
def select_hdr10plus_frames(
|
|
*,
|
|
source_json: Path,
|
|
source_frames: list[int],
|
|
output_json: Path,
|
|
) -> None:
|
|
"""Keep metadata for the encoded source frames and rebuild scene indices."""
|
|
payload = json.loads(source_json.read_text(encoding="utf-8"))
|
|
entries = payload.get("SceneInfo")
|
|
if not isinstance(entries, list) or not entries:
|
|
raise RuntimeError(f"HDR10+ JSON has no SceneInfo entries: {source_json}")
|
|
if not source_frames:
|
|
raise RuntimeError("Cannot create HDR10+ metadata for an empty output")
|
|
if min(source_frames) < 0 or max(source_frames) >= len(entries):
|
|
raise RuntimeError(
|
|
"HDR10+ metadata frame count does not cover the validated source frames "
|
|
f"({len(entries)} metadata entries, highest source frame {max(source_frames)})"
|
|
)
|
|
|
|
selected = [dict(entries[index]) for index in source_frames]
|
|
_reindex_scenes(selected)
|
|
payload["SceneInfo"] = selected
|
|
payload["SceneInfoSummary"] = _scene_summary(selected)
|
|
output_json.write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8")
|
|
|
|
|
|
def _reindex_scenes(entries: list[dict[str, Any]]) -> None:
|
|
scene_id = 0
|
|
scene_frame_index = 0
|
|
previous: dict[str, Any] | None = None
|
|
for sequence_frame_index, entry in enumerate(entries):
|
|
if previous is not None and _scene_payload(entry) != _scene_payload(previous):
|
|
scene_id += 1
|
|
scene_frame_index = 0
|
|
entry["SceneFrameIndex"] = scene_frame_index
|
|
entry["SceneId"] = scene_id
|
|
entry["SequenceFrameIndex"] = sequence_frame_index
|
|
scene_frame_index += 1
|
|
previous = entry
|
|
|
|
|
|
def _scene_payload(entry: dict[str, Any]) -> dict[str, Any]:
|
|
return {
|
|
key: value
|
|
for key, value in entry.items()
|
|
if key not in {"SceneFrameIndex", "SceneId", "SequenceFrameIndex"}
|
|
}
|
|
|
|
|
|
def _scene_summary(entries: list[dict[str, Any]]) -> dict[str, list[int]]:
|
|
starts = [
|
|
index
|
|
for index, entry in enumerate(entries)
|
|
if entry.get("SceneFrameIndex") == 0
|
|
]
|
|
return {
|
|
"SceneFirstFrameIndex": starts,
|
|
"SceneFrameNumbers": [
|
|
next_start - start
|
|
for start, next_start in zip(starts, [*starts[1:], len(entries)])
|
|
],
|
|
}
|