media-batch-tools/tools/video_codec_constraints.py

from __future__ import annotations

from dataclasses import dataclass
from math import sqrt


@dataclass(frozen=True)
class VideoConstraintSpec:
    width: int
    height: int
    bit_depth: int
    chroma: str
    fps: float


_CHROMA_RANK = {"420": 0, "422": 1, "444": 2}

_X264_PROFILES = (
    ("baseline", 8, "420"),
    ("main", 8, "420"),
    ("high", 8, "420"),
    ("high10", 10, "420"),
    ("high422", 10, "422"),
    ("high444", 10, "444"),
)

_X265_PROFILES = (
    ("main", 8, "420"),
    ("main10", 10, "420"),
    ("main12", 12, "420"),
    ("main422-10", 10, "422"),
    ("main422-12", 12, "422"),
    ("main444-8", 8, "444"),
    ("main444-10", 10, "444"),
    ("main444-12", 12, "444"),
)

# (name, maximum luma samples per picture, maximum luma samples per second)
_X265_LEVELS = (
    ("1.0", 36_864, 552_960),
    ("2.0", 122_880, 3_686_400),
    ("2.1", 245_760, 7_372_800),
    ("3.0", 552_960, 16_588_800),
    ("3.1", 983_040, 33_177_600),
    ("4.0", 2_228_224, 66_846_720),
    ("4.1", 2_228_224, 133_693_440),
    ("5.0", 8_912_896, 267_386_880),
    ("5.1", 8_912_896, 534_773_760),
    ("5.2", 8_912_896, 1_069_547_520),
    ("6.0", 35_651_584, 1_069_547_520),
    ("6.1", 35_651_584, 2_139_095_040),
    ("6.2", 35_651_584, 4_278_190_080),
)

# (name, maximum macroblocks per frame, maximum macroblocks per second)
_X264_LEVELS = (
    ("1.0", 99, 1_485),
    ("1.1", 396, 3_000),
    ("1.2", 396, 6_000),
    ("1.3", 396, 11_880),
    ("2.0", 396, 11_880),
    ("2.1", 792, 19_800),
    ("2.2", 1_620, 20_250),
    ("3.0", 1_620, 40_500),
    ("3.1", 3_600, 108_000),
    ("3.2", 5_120, 216_000),
    ("4.0", 8_192, 245_760),
    ("4.1", 8_192, 245_760),
    ("4.2", 8_704, 522_240),
    ("5.0", 22_080, 589_824),
    ("5.1", 36_864, 983_040),
    ("5.2", 36_864, 2_073_600),
    ("6.0", 139_264, 4_177_920),
    ("6.1", 139_264, 8_355_840),
    ("6.2", 139_264, 16_711_680),
)


def profile_names(codec: str) -> list[str]:
    return [name for name, _, _ in _profiles(codec)]


def profile_supports(codec: str, profile: str, spec: VideoConstraintSpec) -> bool:
    _, depth, chroma = _profile(codec, profile)
    return (
        depth >= spec.bit_depth
        and _CHROMA_RANK[chroma] >= _CHROMA_RANK[spec.chroma]
    )


def minimum_profile(codec: str, spec: VideoConstraintSpec) -> str:
    for name in profile_names(codec):
        if profile_supports(codec, name, spec):
            return name
    raise ValueError(
        f"{codec} has no profile for "
        f"{spec.bit_depth}-bit 4:{spec.chroma[1]}:{spec.chroma[2]}"
    )


def resolve_profile(codec: str, choice: str | None, spec: VideoConstraintSpec) -> str | None:
    if choice is None or choice == "none":
        return None
    if choice == "minimum":
        return minimum_profile(codec, spec)

    if codec == "libx264":
        names = profile_names(codec)
        start = names.index(choice)
        for name in names[start:]:
            if profile_supports(codec, name, spec):
                return name
        return minimum_profile(codec, spec)

    _, requested_depth, requested_chroma = _profile(codec, choice)
    required_depth = max(requested_depth, spec.bit_depth)
    required_chroma = max(
        _CHROMA_RANK[requested_chroma],
        _CHROMA_RANK[spec.chroma],
    )
    candidates = [
        (name, depth, chroma)
        for name, depth, chroma in _profiles(codec)
        if depth >= required_depth and _CHROMA_RANK[chroma] >= required_chroma
    ]
    if not candidates:
        return minimum_profile(codec, spec)
    return min(
        candidates,
        key=lambda item: (item[1], _CHROMA_RANK[item[2]]),
    )[0]


def available_profiles(codec: str, specs: list[VideoConstraintSpec]) -> list[str]:
    native_capabilities = {
        _profile(codec, minimum_profile(codec, spec))[1:]
        for spec in specs
    }
    return [
        name
        for name, depth, chroma in _profiles(codec)
        if (depth, chroma) in native_capabilities
    ]


def level_names(codec: str) -> list[str]:
    return [name for name, _, _ in _levels(codec)]


def level_supports(codec: str, level: str, spec: VideoConstraintSpec) -> bool:
    _, max_frame, max_rate = _level(codec, level)
    if codec == "libx264":
        width = (spec.width + 15) // 16
        height = (spec.height + 15) // 16
    else:
        width = spec.width
        height = spec.height
    frame_size = width * height
    max_dimension = sqrt(max_frame * 8)
    return (
        frame_size <= max_frame
        and width <= max_dimension
        and height <= max_dimension
        and frame_size * spec.fps <= max_rate
    )


def minimum_level(codec: str, spec: VideoConstraintSpec) -> str:
    for name in level_names(codec):
        if level_supports(codec, name, spec):
            return name
    raise ValueError(
        f"{codec} has no supported level for {spec.width}x{spec.height} at {spec.fps:.3f} fps"
    )


def resolve_level(codec: str, choice: str | None, spec: VideoConstraintSpec) -> str | None:
    if choice is None or choice == "none":
        return None
    minimum = minimum_level(codec, spec)
    if choice == "minimum":
        return minimum
    names = level_names(codec)
    return names[max(names.index(choice), names.index(minimum))]


def available_levels(codec: str, specs: list[VideoConstraintSpec]) -> list[str]:
    return [
        name
        for name in level_names(codec)
        if any(level_supports(codec, name, spec) for spec in specs)
    ]


def _profiles(codec: str):
    if codec == "libx264":
        return _X264_PROFILES
    if codec == "libx265":
        return _X265_PROFILES
    raise ValueError(f"Unsupported codec: {codec}")


def _profile(codec: str, name: str):
    for profile in _profiles(codec):
        if profile[0] == name:
            return profile
    raise ValueError(f"Unknown {codec} profile: {name}")


def _levels(codec: str):
    if codec == "libx264":
        return _X264_LEVELS
    if codec == "libx265":
        return _X265_LEVELS
    raise ValueError(f"Unsupported codec: {codec}")


def _level(codec: str, name: str):
    for level in _levels(codec):
        if level[0] == name:
            return level
    raise ValueError(f"Unknown {codec} level: {name}")