prs-eth/rollingdepth

1# Copyright 2024 Bingxin Ke, ETH Zurich. All rights reserved.2# Last modified: 2024-11-283#4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License at7#8#     http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# distributed under the License is distributed on an "AS IS" BASIS,12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.13# See the License for the specific language governing permissions and14# limitations under the License.15# ---------------------------------------------------------------------------------16# If you find this code useful, we kindly ask you to cite our paper in your work.17# Please find bibtex at: https://github.com/prs-eth/RollingDepth#-citation18# More information about the method can be found at https://rollingdepth.github.io19# ---------------------------------------------------------------------------------20import logging21from os import PathLike22from typing import Optional23 24import av25import numpy as np26from tqdm import tqdm27 28 29def get_video_fps(video_path: PathLike) -> float:30    # Open the video file31    container = av.open(video_path)32 33    # Get the video stream34    video_stream = container.streams.video[0]35 36    # Calculate FPS from the stream's time base and average frame rate37    fps = float(video_stream.average_rate)  # type: ignore38 39    # Close the container40    container.close()41 42    return fps43 44 45def write_video_from_numpy(46    frames: np.ndarray,  # shape [n h w 3]47    output_path: PathLike,48    fps: int = 30,49    codec: Optional[str] = None,  # Let PyAV choose default codec50    crf: int = 23,51    preset: str = "medium",52    verbose: bool = False,53) -> None:54    if len(frames.shape) != 4 or frames.shape[-1] != 3:55        raise ValueError(f"Expected shape [n, height, width, 3], got {frames.shape}")56    if frames.dtype != np.uint8:57        raise ValueError(f"Expected dtype uint8, got {frames.dtype}")58 59    n_frames, height, width, _ = frames.shape60 61    # Try to determine codec from output format if not specified62    if codec is None:63        codecs_to_try = ["libx264", "h264", "mpeg4", "mjpeg"]64    else:65        codecs_to_try = [codec]66 67    # Try available codecs68    for try_codec in codecs_to_try:69        try:70            container = av.open(output_path, mode="w")71            stream = container.add_stream(try_codec, rate=fps)72            if verbose:73                logging.info(f"Using codec: {try_codec}")74            break75        except av.codec.codec.UnknownCodecError:  # type: ignore76            if try_codec == codecs_to_try[-1]:  # Last codec in list77                raise ValueError(78                    f"No working codec found. Tried: {codecs_to_try}. "79                    "Please install ffmpeg with necessary codecs."80                )81            continue82 83    stream.width = width  # type: ignore84    stream.height = height  # type: ignore85    stream.pix_fmt = "yuv420p"  # type: ignore86 87    # Only set these options for x264-compatible codecs88    if try_codec in ["libx264", "h264"]:  # type: ignore89        stream.options = {"crf": str(crf), "preset": preset}  # type: ignore90 91    # Create a single VideoFrame object and reuse it92    video_frame = av.VideoFrame(width, height, "rgb24")93 94    frames_iterable = range(n_frames)95    if verbose:96        frames_iterable = tqdm(frames_iterable, desc="Writing video", total=n_frames)97 98    try:99        for frame_idx in frames_iterable:100            # Get view of current frame101            current_frame = frames[frame_idx]102 103            # Update frame data in-place104            video_frame.to_ndarray()[:] = current_frame105 106            packet = stream.encode(video_frame)  # type: ignore107            container.mux(packet)  # type: ignore108 109        # Flush the stream110        packet = stream.encode(None)  # type: ignore111        container.mux(packet)  # type: ignore112    finally:113        container.close()  # type: ignore114