Spaces:
Running
on
Zero
Running
on
Zero
File size: 2,689 Bytes
a3a2e41 |
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 |
import tempfile
from typing import Optional
import numpy as np
from moviepy.editor import ImageSequenceClip, AudioFileClip
from scipy.io import wavfile
def save_video(
output_path: str,
video_numpy: np.ndarray,
audio_numpy: Optional[np.ndarray] = None,
sample_rate: int = 16000,
fps: int = 24,
) -> str:
"""
Combine a sequence of video frames with an optional audio track and save as an MP4.
Args:
output_path (str): Path to the output MP4 file.
video_numpy (np.ndarray): Numpy array of frames. Shape (C, F, H, W).
Values can be in range [-1, 1] or [0, 255].
audio_numpy (Optional[np.ndarray]): 1D or 2D numpy array of audio samples, range [-1, 1].
sample_rate (int): Sample rate of the audio in Hz. Defaults to 16000.
fps (int): Frames per second for the video. Defaults to 24.
Returns:
str: Path to the saved MP4 file.
"""
# Validate inputs
assert isinstance(video_numpy, np.ndarray), "video_numpy must be a numpy array"
assert video_numpy.ndim == 4, "video_numpy must have shape (C, F, H, W)"
assert video_numpy.shape[0] in {1, 3}, "video_numpy must have 1 or 3 channels"
if audio_numpy is not None:
assert isinstance(audio_numpy, np.ndarray), "audio_numpy must be a numpy array"
assert np.abs(audio_numpy).max() <= 1.0, "audio_numpy values must be in range [-1, 1]"
# Reorder dimensions: (C, F, H, W) → (F, H, W, C)
video_numpy = video_numpy.transpose(1, 2, 3, 0)
# Normalize frames if values are in [-1, 1]
if video_numpy.max() <= 1.0:
video_numpy = np.clip(video_numpy, -1, 1)
video_numpy = ((video_numpy + 1) / 2 * 255).astype(np.uint8)
else:
video_numpy = video_numpy.astype(np.uint8)
# Convert numpy array to a list of frames
frames = list(video_numpy)
# Create video clip
clip = ImageSequenceClip(frames, fps=fps)
# Add audio if provided
if audio_numpy is not None:
with tempfile.NamedTemporaryFile(suffix=".wav") as temp_audio_file:
wavfile.write(
temp_audio_file.name,
sample_rate,
(audio_numpy * 32767).astype(np.int16),
)
audio_clip = AudioFileClip(temp_audio_file.name)
final_clip = clip.set_audio(audio_clip)
else:
final_clip = clip
# Write final video to disk
final_clip.write_videofile(
output_path, codec="libx264", audio_codec="aac", fps=fps, verbose=False, logger=None
)
final_clip.close()
return output_path |