gem/utils/video_io_utils.py

# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
import os
import shutil
from pathlib import Path

import cv2
import ffmpeg
import imageio.v3 as iio
import numpy as np
import torch
from tqdm import tqdm


def get_video_lwh(video_path):
    L, H, W, _ = iio.improps(video_path, plugin="pyav").shape
    return L, W, H


def read_video_np(video_path, start_frame=0, end_frame=-1, scale=1.0):
    """
    Args:
        video_path: str
    Returns:
        frames: np.array, (N, H, W, 3) RGB, uint8
    """
    # If video path not exists, an error will be raised by ffmpegs
    filter_args = []
    should_check_length = False

    # 1. Trim
    if not (start_frame == 0 and end_frame == -1):
        if end_frame == -1:
            filter_args.append(("trim", f"start_frame={start_frame}"))
        else:
            should_check_length = True
            filter_args.append(("trim", f"start_frame={start_frame}:end_frame={end_frame}"))

    # 2. Scale
    if scale != 1.0:
        filter_args.append(("scale", f"iw*{scale}:ih*{scale}"))

    # Excute then check
    frames = iio.imread(video_path, plugin="pyav", filter_sequence=filter_args)
    if should_check_length:
        assert len(frames) == end_frame - start_frame

    return frames


def get_video_reader(video_path):
    return iio.imiter(video_path, plugin="pyav")


def read_images_np(image_paths, verbose=False):
    """
    Args:
        image_paths: list of str
    Returns:
        images: np.array, (N, H, W, 3) RGB, uint8
    """
    if verbose:
        images = [cv2.imread(str(img_path))[..., ::-1] for img_path in tqdm(image_paths)]
    else:
        images = [cv2.imread(str(img_path))[..., ::-1] for img_path in image_paths]
    images = np.stack(images, axis=0)
    return images


def save_video(images, video_path, fps=30, crf=17):
    """
    Args:
        images: (N, H, W, 3) RGB, uint8
        crf: 17 is visually lossless, 23 is default, +6 results in half the bitrate
    0 is lossless, https://trac.ffmpeg.org/wiki/Encode/H.264#crf
    """
    if isinstance(images, torch.Tensor):
        images = images.cpu().numpy().astype(np.uint8)
    elif isinstance(images, list):
        images = np.array(images).astype(np.uint8)

    with iio.imopen(video_path, "w", plugin="pyav") as writer:
        writer.init_video_stream("libx264", fps=fps)
        writer._video_stream.options = {"crf": str(crf)}
        writer.write(images)


def get_writer(video_path, fps=30, crf=17):
    """remember to .close()"""
    writer = iio.imopen(video_path, "w", plugin="pyav")
    writer.init_video_stream("libx264", fps=fps)
    writer._video_stream.options = {"crf": str(crf)}
    return writer


def copy_file(video_path, out_video_path, overwrite=True):
    if not overwrite and Path(out_video_path).exists():
        return
    shutil.copy(video_path, out_video_path)


def concat_videos(cfg, out_video_path: str, in_video_paths=None):
    # if len(in_video_paths) < 2:
    #     raise ValueError("At least two video paths are required for merging.")
    # in_video_paths = [cfg.video1_path, cfg.text1_video_path, cfg.video2_path]
    if in_video_paths is None:
        in_video_paths = [
            cfg.paths.incam_video1,
            cfg.text1_video_path,
            cfg.paths.incam_video2,
        ]

    # Get the size of the first video to use as target size
    probe = ffmpeg.probe(in_video_paths[0])
    video_stream = next(
        (stream for stream in probe["streams"] if stream["codec_type"] == "video"), None
    )
    target_size = (int(video_stream["width"]), int(video_stream["height"]))

    # Resize and pad all videos to match the target size
    temp_paths = [resize_and_pad_video(path, target_size) for path in in_video_paths]

    try:
        # Create inputs from the resized videos
        inputs = [ffmpeg.input(path) for path in temp_paths]
        merged_video = ffmpeg.concat(*inputs)
        output = ffmpeg.output(merged_video, out_video_path)
        ffmpeg.run(output, overwrite_output=True, quiet=True)
    finally:
        # Clean up temporary files
        for path in temp_paths:
            if os.path.exists(path):
                os.unlink(path)


def resize_and_pad_video(video_path, target_size):
    """
    Resize and pad a video to match the target size.

    Args:
        video_path: Path to the input video
        target_size: Tuple of (width, height) for the target size

    Returns:
        Path to the resized and padded temporary video
    """
    import tempfile

    target_width, target_height = target_size

    # Create a temporary file for the resized video
    temp_file = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)
    temp_path = temp_file.name
    temp_file.close()

    # Get video info
    probe = ffmpeg.probe(video_path)
    video_stream = next(
        (stream for stream in probe["streams"] if stream["codec_type"] == "video"), None
    )
    width = int(video_stream["width"])
    height = int(video_stream["height"])

    # Calculate scaling to maintain aspect ratio
    if width / height > target_width / target_height:
        # Width is the limiting factor
        scale_w = target_width
        scale_h = -1  # Maintain aspect ratio
    else:
        # Height is the limiting factor
        scale_w = -1  # Maintain aspect ratio
        scale_h = target_height

    # Resize and pad
    stream = ffmpeg.input(video_path)
    stream = ffmpeg.filter(stream, "scale", scale_w, scale_h)
    stream = ffmpeg.filter(stream, "pad", target_width, target_height, "(ow-iw)/2", "(oh-ih)/2")
    stream = ffmpeg.output(stream, temp_path)
    ffmpeg.run(stream, quiet=True, overwrite_output=True)

    return temp_path


def merge_videos_horizontal(in_video_paths: list, out_video_path: str):
    if len(in_video_paths) < 2:
        raise ValueError("At least two video paths are required for merging.")
    inputs = [ffmpeg.input(path) for path in in_video_paths]
    merged_video = ffmpeg.filter(inputs, "hstack", inputs=len(inputs))
    output = ffmpeg.output(merged_video, out_video_path)
    ffmpeg.run(output, overwrite_output=True, quiet=True)


def merge_videos_grid_2x2(in_video_paths: list, out_video_path: str):
    """Merge exactly 4 videos into a 2x2 grid: [top-left, top-right, bottom-left, bottom-right]."""
    if len(in_video_paths) != 4:
        raise ValueError("Exactly four video paths are required for 2x2 grid merge.")
    inputs = [ffmpeg.input(path) for path in in_video_paths]
    top_row = ffmpeg.filter([inputs[0], inputs[1]], "hstack", inputs=2)
    bottom_row = ffmpeg.filter([inputs[2], inputs[3]], "hstack", inputs=2)
    grid = ffmpeg.filter([top_row, bottom_row], "vstack", inputs=2)
    output = ffmpeg.output(grid, out_video_path)
    ffmpeg.run(output, overwrite_output=True, quiet=True)


def merge_videos_vertical(in_video_paths: list, out_video_path: str):
    if len(in_video_paths) < 2:
        raise ValueError("At least two video paths are required for merging.")
    inputs = [ffmpeg.input(path) for path in in_video_paths]
    merged_video = ffmpeg.filter(inputs, "vstack", inputs=len(inputs))
    output = ffmpeg.output(merged_video, out_video_path)
    ffmpeg.run(output, overwrite_output=True, quiet=True)