mvdata 0.9.0__py3-none-any.whl

mvdata/codec/encode.py

@@ -0,0 +1,313 @@
+"""NVENC encoding pipeline: RGB → elementary stream → MP4."""
+
+from __future__ import annotations
+
+import io
+import shutil
+import subprocess
+import tempfile
+from fractions import Fraction
+from pathlib import Path
+from typing import Iterable, List, Tuple
+
+import numpy as np
+
+from ._imports import normalize_codec_name, try_import_av
+from .probe import (
+    nvenc_encode_compatibility_issue,
+    nvenc_encode_compatibility_issue_for_rgb,
+    safe_get_encoder_caps,
+)
+
+
+def rgb_hwc_to_nvenc_bgra_hwc4(rgb: np.ndarray) -> np.ndarray:
+    """HWC uint8 RGB → contiguous H×W×4 BGRA for ``CreateEncoder(..., 'ARGB', True)``."""
+    if rgb.dtype != np.uint8 or rgb.ndim != 3 or rgb.shape[2] != 3:
+        raise ValueError("rgb must be uint8 HWC with 3 channels")
+    r, g, b = rgb[:, :, 0], rgb[:, :, 1], rgb[:, :, 2]
+    a = np.full(rgb.shape[:2], 255, dtype=np.uint8)
+    return np.stack([b, g, r, a], axis=-1)
+
+
+def rgb_hwc_to_yuv444_planar(rgb: np.ndarray, *, full_range: bool = True) -> np.ndarray:
+    """HWC uint8 RGB → planar YUV444 (H*3, W) uint8 for ``CreateEncoder(..., 'YUV444', True)``."""
+    if rgb.dtype != np.uint8 or rgb.ndim != 3 or rgb.shape[2] != 3:
+        raise ValueError("rgb must be uint8 HWC with 3 channels")
+    r = rgb[:, :, 0].astype(np.float32)
+    g = rgb[:, :, 1].astype(np.float32)
+    b = rgb[:, :, 2].astype(np.float32)
+    if full_range:
+        y = 0.299 * r + 0.587 * g + 0.114 * b
+        cb = -0.168736 * r - 0.331264 * g + 0.5 * b + 128.0
+        cr = 0.5 * r - 0.418688 * g - 0.081312 * b + 128.0
+    else:
+        y = 16.0 + 65.481 * r / 255.0 + 128.553 * g / 255.0 + 24.966 * b / 255.0
+        cb = 128.0 - 37.797 * r / 255.0 - 74.203 * g / 255.0 + 112.0 * b / 255.0
+        cr = 128.0 + 112.0 * r / 255.0 - 93.786 * g / 255.0 - 18.214 * b / 255.0
+    clip = lambda plane: np.clip(plane + 0.5, 0, 255).astype(np.uint8)
+    return np.concatenate([clip(y), clip(cb), clip(cr)], axis=0)
+
+
+def pad_rgb_for_nvenc(rgb: np.ndarray) -> Tuple[np.ndarray, Tuple[int, int, int, int]]:
+    """Return the original frame plus its full-frame crop (portrait is unsupported)."""
+    h, w, c = rgb.shape
+    if c != 3:
+        raise ValueError("expected RGB HWC")
+    if w < h:
+        raise ValueError(f"portrait frames ({w}x{h}) are not encoded via NVENC without padding")
+    return rgb, (0, h, 0, w)
+
+
+def assert_yuv444_encode_supported(nvc, gpu_id: int, codec: str) -> None:
+    codec = normalize_codec_name(codec)
+    caps = safe_get_encoder_caps(nvc, gpu_id, codec)
+    if caps is None:
+        raise RuntimeError(f"GPU reports no encoder support for codec={codec!r}.")
+    if not caps.get("support_yuv444_encode"):
+        raise RuntimeError(
+            f"GPU reports no 4:4:4 encode for codec={codec!r} (support_yuv444_encode=0). "
+            "Use chroma 420 or a different codec/GPU."
+        )
+
+
+def best_quality_extra_kwargs(nvc, gpu_id: int, codec: str) -> dict:
+    """Only enable features advertised by the GPU so the dict is safe to splat into CreateEncoder."""
+    caps = nvc.GetEncoderCaps(gpu_id, codec)
+    kw: dict = {}
+    if caps.get("support_lookahead"):
+        kw["lookahead"] = 32
+    return kw
+
+
+def suggest_nvenc_vbr_bitrate_bps(
+    width: int, height: int, fps: int, *, chroma_subsampling: str = "444"
+) -> int:
+    ref_px = 1920 * 1080
+    ref_fps = 30
+    ref_bps = 80_000_000
+    px = max(1, int(width) * int(height))
+    t = max(1, int(fps))
+    bps = int(ref_bps * ((px / ref_px) ** 0.88) * (t / ref_fps))
+    c = chroma_subsampling.strip()
+    if c == "444":
+        bps = int(bps * 1.5)
+    elif c != "420":
+        raise ValueError(f"chroma_subsampling must be 420 or 444, got {chroma_subsampling!r}")
+    return max(4_000_000, min(bps, 500_000_000))
+
+
+def resolve_nvenc_bitrate_bps(
+    bitrate_arg: int, width: int, height: int, fps: int, *, chroma_subsampling: str
+) -> int:
+    if bitrate_arg > 0:
+        return int(bitrate_arg)
+    return suggest_nvenc_vbr_bitrate_bps(width, height, fps, chroma_subsampling=chroma_subsampling)
+
+
+def create_nvenc_encoder(
+    nvc,
+    width: int,
+    height: int,
+    gpu_id: int,
+    fps: int,
+    bitrate: int,
+    *,
+    codec: str = "h264",
+    full_range: bool = True,
+    chroma_subsampling: str = "444",
+    preset: str = "P7",
+):
+    codec = normalize_codec_name(codec)
+    issue = nvenc_encode_compatibility_issue(
+        nvc, width, height, gpu_id=gpu_id, codec=codec, chroma_subsampling=chroma_subsampling
+    )
+    if issue is not None:
+        raise ValueError(issue)
+
+    surface_fmt = "YUV444" if chroma_subsampling == "444" else "ARGB"
+    kwargs = dict(
+        gpu_id=gpu_id,
+        codec=codec,
+        fps=fps,
+        preset=preset,
+        tuning_info="high_quality",
+        rc="vbr",
+        bitrate=bitrate,
+        maxbitrate=max(bitrate + 2_000_000, int(bitrate * 1.25)),
+        bf=0,
+        **best_quality_extra_kwargs(nvc, gpu_id, codec),
+    )
+    if full_range:
+        kwargs["colorrange"] = nvc.ColorRange.JPEG
+        kwargs["videoFullRangeFlag"] = 1
+    return nvc.CreateEncoder(width, height, surface_fmt, True, **kwargs)
+
+
+def encode_rgb_iter_to_elementary_stream(
+    nvc,
+    frames: Iterable[np.ndarray],
+    gpu_id: int,
+    fps: int,
+    bitrate: int,
+    *,
+    codec: str = "h264",
+    full_range: bool = True,
+    chroma_subsampling: str = "444",
+    preset: str = "P7",
+) -> bytes:
+    """Encode an iterable of RGB uint8 HWC frames to an H.264/HEVC elementary stream."""
+    chunks: List[bytes] = []
+    enc = None
+    ref_h = ref_w = 0
+    use_yuv444 = chroma_subsampling == "444"
+
+    for rgb in frames:
+        issue = nvenc_encode_compatibility_issue_for_rgb(
+            nvc, rgb, gpu_id=gpu_id, codec=codec, chroma_subsampling=chroma_subsampling
+        )
+        if issue is not None:
+            raise ValueError(issue)
+        if enc is None:
+            ref_h, ref_w = rgb.shape[0], rgb.shape[1]
+            enc = create_nvenc_encoder(
+                nvc, ref_w, ref_h, gpu_id, fps, bitrate,
+                codec=codec, full_range=full_range,
+                chroma_subsampling=chroma_subsampling, preset=preset,
+            )
+        elif rgb.shape[0] != ref_h or rgb.shape[1] != ref_w:
+            raise ValueError("All frames must share the same resolution within a stream")
+
+        plane = np.ascontiguousarray(
+            rgb_hwc_to_yuv444_planar(rgb, full_range=full_range)
+            if use_yuv444
+            else rgb_hwc_to_nvenc_bgra_hwc4(rgb)
+        )
+        try:
+            out = enc.Encode(plane)
+        except TypeError:
+            out = enc.Encode(plane.reshape(-1))
+        if out:
+            chunks.append(bytes(out))
+
+    if enc is None:
+        return b""
+    tail = enc.EndEncode()
+    if tail:
+        chunks.append(bytes(tail))
+    return b"".join(chunks)
+
+
+def encode_rgb_frames_to_elementary_stream(
+    nvc,
+    frames: List[np.ndarray],
+    gpu_id: int,
+    fps: int,
+    bitrate: int,
+    *,
+    codec: str = "h264",
+    full_range: bool = True,
+    chroma_subsampling: str = "444",
+    preset: str = "P7",
+) -> bytes:
+    return encode_rgb_iter_to_elementary_stream(
+        nvc, iter(frames), gpu_id, fps, bitrate,
+        codec=codec, full_range=full_range,
+        chroma_subsampling=chroma_subsampling, preset=preset,
+    )
+
+
+def ffmpeg_executable() -> str | None:
+    try:
+        import imageio_ffmpeg  # type: ignore
+
+        return imageio_ffmpeg.get_ffmpeg_exe()
+    except Exception:
+        return shutil.which("ffmpeg")
+
+
+def mux_elementary_to_mp4(
+    es_bytes: bytes, mp4_path: Path, fps: float, *, demux_format: str = "h264"
+) -> None:
+    ffmpeg = ffmpeg_executable()
+    if ffmpeg:
+        suffix = ".hevc" if demux_format in ("hevc", "h265") else ".h264"
+        with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as f:
+            f.write(es_bytes)
+            tmp_es = Path(f.name)
+        try:
+            subprocess.run(
+                [ffmpeg, "-y", "-loglevel", "error", "-f", demux_format,
+                 "-r", str(fps), "-i", str(tmp_es), "-c:v", "copy", str(mp4_path)],
+                check=True,
+            )
+        finally:
+            tmp_es.unlink(missing_ok=True)
+        return
+
+    av = try_import_av()
+    if av is None:
+        raise RuntimeError(
+            "Need ffmpeg in PATH, or install PyAV (`pip install av`) to mux elementary video to MP4."
+        )
+
+    bio = io.BytesIO(es_bytes)
+    with av.open(bio, format=demux_format) as inp:
+        if not inp.streams.video:
+            raise RuntimeError("No video stream in elementary bitstream")
+        vin = inp.streams.video[0]
+        with av.open(str(mp4_path), mode="w") as out:
+            vout = out.add_stream_from_template(vin, opaque=True)
+            tb_den = 90_000
+            vout.time_base = Fraction(1, tb_den)
+            inc = max(1, int(round(tb_den / float(fps))))
+            for i, packet in enumerate(inp.demux(vin)):
+                if getattr(packet, "size", 0) == 0:
+                    continue
+                if packet.pts is None:
+                    packet.pts = i * inc
+                    packet.dts = packet.pts
+                packet.stream = vout
+                out.mux(packet)
+
+
+def encode_rgb_iter_to_mp4(
+    nvc,
+    frames: Iterable[np.ndarray],
+    mp4_path: Path,
+    *,
+    gpu_id: int,
+    fps: int,
+    bitrate: int,
+    codec: str = "h264",
+    full_range: bool = True,
+    chroma_subsampling: str = "444",
+    preset: str = "P7",
+) -> None:
+    codec = normalize_codec_name(codec)
+    es = encode_rgb_iter_to_elementary_stream(
+        nvc, frames, gpu_id, fps, bitrate,
+        codec=codec, full_range=full_range,
+        chroma_subsampling=chroma_subsampling, preset=preset,
+    )
+    demux = "hevc" if codec in ("hevc", "h265") else "h264"
+    mux_elementary_to_mp4(es, mp4_path, float(fps), demux_format=demux)
+
+
+def encode_rgb_sequence_to_mp4(
+    nvc,
+    frames: List[np.ndarray],
+    mp4_path: Path,
+    *,
+    gpu_id: int,
+    fps: int,
+    bitrate: int,
+    codec: str = "h264",
+    full_range: bool = True,
+    chroma_subsampling: str = "444",
+    preset: str = "P7",
+) -> None:
+    encode_rgb_iter_to_mp4(
+        nvc, iter(frames), mp4_path,
+        gpu_id=gpu_id, fps=fps, bitrate=bitrate, codec=codec,
+        full_range=full_range, chroma_subsampling=chroma_subsampling, preset=preset,
+    )

mvdata/codec/frames.py

@@ -0,0 +1,199 @@
+"""Conversion helpers that turn decoded frames into canonical HWC RGB arrays.
+
+The public surface here is deliberately small:
+
+``_decoded_frame_to_rgb_numpy``
+    Takes a decoder output (numpy, dlpack-capable, or torch/cupy tensor) and
+    returns a contiguous HWC numpy array owned by the caller.
+``_decoded_frame_to_rgb_cupy``
+    Same, but returns a fresh ``cupy.ndarray`` that stays on the GPU. Uses
+    DLPack to view decoder device memory, then copies into a new buffer so the
+    decoder is free to reuse its pool slot. Accepts an optional raw CUDA stream
+    pointer to order the copy with user work.
+``numpy_to_cupy_rgb``
+    H2D helper used for streams that can only be produced on the CPU (pyavif,
+    PyAV CPU fallback).
+``_pyav_frame_to_rgb``
+    PyAV → HWC RGB with correct handling of full-range yuv444 and high bit
+    depths.
+"""
+
+from __future__ import annotations
+
+from contextlib import ExitStack
+from typing import Any
+
+import numpy as np
+
+from ._imports import try_import_cupy, try_import_torch
+from .probe import infer_video_bit_depth_from_frame
+
+
+def _cupy_or_raise():
+    cp = try_import_cupy()
+    if cp is None:
+        raise RuntimeError(
+            "device='cuda' requires CuPy. Install it with `pip install mvdata[cuda]` "
+            "(or `pip install cupy-cuda12x`)."
+        )
+    return cp
+
+
+def _device_and_stream_ctx(
+    cp, *, gpu_id: int | None = None, cuda_stream: int | None = None
+):
+    stack = ExitStack()
+    if gpu_id is not None:
+        stack.enter_context(cp.cuda.Device(int(gpu_id)))
+    if cuda_stream is not None and int(cuda_stream) != 0:
+        stack.enter_context(cp.cuda.ExternalStream(int(cuda_stream)))
+    return stack
+
+
+def _target_dtype_for(arr_dtype, *, bit_depth: int | None):
+    if arr_dtype == np.uint8:
+        return np.uint8
+    if arr_dtype == np.uint16:
+        return np.uint16
+    if np.issubdtype(arr_dtype, np.integer):
+        if (bit_depth or 8) > 8 or arr_dtype.itemsize > 1:
+            return np.uint16
+        return np.uint8
+    raise TypeError(f"Cannot convert decoded frame to HWC RGB: dtype={arr_dtype}")
+
+
+def _ensure_hwc(arr):
+    """Transpose 3×H×W to H×W×3 when needed. Returns (arr, needs_copy)."""
+    if arr.ndim == 3 and arr.shape[0] == 3 and arr.shape[-1] != 3:
+        return arr.transpose(1, 2, 0), True
+    if not (arr.ndim == 3 and arr.shape[-1] == 3):
+        raise TypeError(f"Cannot convert decoded frame to HWC RGB: shape={arr.shape}")
+    return arr, False
+
+
+def _normalize_rgb_numpy_output(
+    arr: np.ndarray, *, bit_depth: int | None = None, copy_result: bool = False
+) -> np.ndarray:
+    arr = np.asarray(arr)
+    arr, transposed = _ensure_hwc(arr)
+    copy_result = copy_result or transposed
+
+    target = _target_dtype_for(arr.dtype, bit_depth=bit_depth)
+    if arr.dtype != target:
+        arr = arr.astype(target, copy=False)
+        copy_result = True
+
+    if copy_result:
+        return np.array(arr, copy=True, order="C")
+    if not arr.flags.c_contiguous:
+        return np.ascontiguousarray(arr)
+    return arr
+
+
+def _decoded_frame_to_rgb_numpy(frame, torch_mod, *, bit_depth: int | None = None):
+    """Return a contiguous HWC RGB numpy array owned by the caller."""
+    if hasattr(frame, "__dlpack__"):
+        try:
+            arr = np.from_dlpack(frame)
+            return _normalize_rgb_numpy_output(arr, bit_depth=bit_depth, copy_result=True)
+        except Exception:
+            pass
+        if torch_mod is not None:
+            try:
+                arr = torch_mod.from_dlpack(frame).detach().cpu().numpy()
+                return _normalize_rgb_numpy_output(arr, bit_depth=bit_depth)
+            except Exception:
+                pass
+        cp = try_import_cupy()
+        if cp is not None:
+            arr = cp.asnumpy(cp.from_dlpack(frame))
+            return _normalize_rgb_numpy_output(arr, bit_depth=bit_depth)
+    # Host-memory decode outputs may alias decoder-managed buffers that get
+    # reused on the next decode call, so take ownership here.
+    return _normalize_rgb_numpy_output(frame, bit_depth=bit_depth, copy_result=True)
+
+
+def _decoded_frame_to_rgb_cupy(
+    frame: Any,
+    *,
+    bit_depth: int | None = None,
+    gpu_id: int | None = None,
+    cuda_stream: int | None = None,
+):
+    """Produce a fresh contiguous HWC RGB ``cupy.ndarray`` from a decoder output.
+
+    Resulting array is always a new device allocation: the caller owns it and
+    the decoder is free to overwrite its pool slot immediately after return.
+    The GPU→GPU copy (or H2D copy, if ``frame`` is a host array) is issued on
+    ``cuda_stream`` when provided, otherwise on CuPy's current stream.
+    """
+    cp = _cupy_or_raise()
+
+    def _as_cupy(src):
+        if hasattr(src, "__dlpack__"):
+            return cp.from_dlpack(src)
+        return cp.asarray(np.ascontiguousarray(np.asarray(src)))
+
+    with _device_and_stream_ctx(cp, gpu_id=gpu_id, cuda_stream=cuda_stream):
+        src = _as_cupy(frame)
+        src, _ = _ensure_hwc(src)
+        target = _target_dtype_for(src.dtype, bit_depth=bit_depth)
+        out = cp.empty(src.shape, dtype=target, order="C")
+        if src.dtype == target:
+            cp.copyto(out, src)
+        else:
+            out[...] = src.astype(target, copy=False)
+    return out
+
+
+def numpy_to_cupy_rgb(
+    arr: np.ndarray,
+    *,
+    gpu_id: int | None = None,
+    cuda_stream: int | None = None,
+):
+    """Copy a host HWC RGB array to a fresh ``cupy.ndarray`` on ``cuda_stream``."""
+    cp = _cupy_or_raise()
+    arr = np.ascontiguousarray(arr)
+    with _device_and_stream_ctx(cp, gpu_id=gpu_id, cuda_stream=cuda_stream):
+        out = cp.empty(arr.shape, dtype=arr.dtype)
+        out.set(arr)
+    return out
+
+
+def _pyav_frame_to_rgb(frame, *, bit_depth: int | None = None) -> np.ndarray:
+    """Convert a PyAV VideoFrame to an HWC RGB numpy array.
+
+    Full-range yuv444p streams are converted manually to dodge swscale's
+    limited-range default when VUI is unspecified.
+    """
+    pf = str(frame.format.name)
+    depth = int(bit_depth) if bit_depth is not None else infer_video_bit_depth_from_frame(frame)
+    cr = getattr(frame, "color_range", None)
+    cr_text = str(cr) if cr is not None else ""
+    is_full = cr is not None and cr_text in ("2", "JPEG", "jpeg", "pc")
+
+    if pf.startswith("yuv444p") and (
+        is_full or cr is None or cr_text in ("0", "UNSPECIFIED", "unspecified")
+    ):
+        yuv = frame.to_ndarray(format=pf)
+        y = yuv[0].astype(np.float32)
+        max_value = float((1 << depth) - 1)
+        center = (max_value + 1.0) / 2.0
+        cb = yuv[1].astype(np.float32) - center
+        crv = yuv[2].astype(np.float32) - center
+
+        if depth > 8:
+            max_out = 65535.0
+            target_dtype = np.uint16
+        else:
+            max_out = 255.0
+            target_dtype = np.uint8
+        scale = max_out / max_value
+
+        r = np.clip((y + 1.402 * crv) * scale + 0.5, 0, max_out).astype(target_dtype)
+        g = np.clip((y - 0.344136 * cb - 0.714136 * crv) * scale + 0.5, 0, max_out).astype(target_dtype)
+        b = np.clip((y + 1.772 * cb) * scale + 0.5, 0, max_out).astype(target_dtype)
+        return np.stack([r, g, b], axis=-1)
+
+    return frame.to_ndarray(format="rgb48le" if depth > 8 else "rgb24")