Spaces:
Running on Zero
Running on Zero
File size: 3,814 Bytes
41ff959 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 | from pathlib import Path
from typing import Union
import imageio.v2 as imageio
import imageio_ffmpeg
import numpy as np
import torch
from einops import rearrange, repeat
from jaxtyping import Float, UInt8
from torch import Tensor
FloatImage = Union[
Float[Tensor, "height width"],
Float[Tensor, "channel height width"],
Float[Tensor, "batch channel height width"],
]
def prep_image(image: FloatImage) -> UInt8[np.ndarray, "height width channel"]:
# Handle batched images.
if image.ndim == 4:
image = rearrange(image, "b c h w -> c h (b w)")
# Handle single-channel images.
if image.ndim == 2:
image = rearrange(image, "h w -> () h w")
# Ensure that there are 3 or 4 channels.
channel, _, _ = image.shape
if channel == 1:
image = repeat(image, "() h w -> c h w", c=3)
assert image.shape[0] in (3, 4)
image = (image.detach().clip(min=0, max=1) * 255).type(torch.uint8)
return rearrange(image, "c h w -> h w c").cpu().numpy()
def save_video(
images: list[FloatImage],
path: Union[Path, str],
fps: int = 10,
) -> None:
"""Save an RGB video whose input frames are in range 0-1."""
# Create the parent directory if it doesn't already exist.
path = Path(path)
path.parent.mkdir(exist_ok=True, parents=True)
frames = [prep_image(image) for image in images]
if len(frames) == 0:
raise ValueError("save_video received an empty frame list.")
reference_height, reference_width = frames[0].shape[:2]
reference_channels = frames[0].shape[2] if frames[0].ndim == 3 else 1
normalized_frames = []
for idx, frame in enumerate(frames):
if frame.ndim != 3:
raise ValueError(f"Frame {idx} has invalid shape {frame.shape}; expected HWC.")
if frame.shape[2] != reference_channels:
raise ValueError(
f"Frame {idx} has channel count {frame.shape[2]} but expected {reference_channels}."
)
if frame.shape[:2] != (reference_height, reference_width):
raise ValueError(
f"Frame {idx} has shape {frame.shape[:2]} but expected {(reference_height, reference_width)}."
)
normalized_frames.append(np.ascontiguousarray(frame))
# yuv420p requires even spatial resolution. Pad the bottom/right edge if needed.
if reference_height % 2 != 0 or reference_width % 2 != 0:
padded_frames = []
pad_height = reference_height % 2
pad_width = reference_width % 2
for frame in normalized_frames:
padded_frames.append(
np.pad(
frame,
((0, pad_height), (0, pad_width), (0, 0)),
mode="edge",
)
)
normalized_frames = padded_frames
reference_height += pad_height
reference_width += pad_width
ffmpeg_path = imageio_ffmpeg.get_ffmpeg_exe()
writer = imageio.get_writer(
str(path),
format="FFMPEG",
mode="I",
fps=fps,
codec="libx264",
pixelformat="yuv420p",
macro_block_size=1,
ffmpeg_log_level="error",
output_params=["-crf", "21"],
)
try:
for frame in normalized_frames:
writer.append_data(frame)
except Exception as exc:
frame_shapes = [frame.shape for frame in normalized_frames[:3]]
if len(normalized_frames) > 3:
frame_shapes.append(("...", len(normalized_frames)))
raise OSError(
f"{exc}\n[save_video] ffmpeg={ffmpeg_path}, "
f"fps={fps}, first_frame_shape={(reference_height, reference_width, reference_channels)}, "
f"sample_frame_shapes={frame_shapes}, output={path}"
) from exc
finally:
writer.close()
|