Spaces:
Running
Running
File size: 2,344 Bytes
99bbd30 |
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 |
from pathlib import Path
from typing import Union
import torch
from torio.io import StreamingMediaDecoder, StreamingMediaEncoder
class VideoJoiner:
def __init__(self, src_root: Union[str, Path], output_root: Union[str, Path], sample_rate: int,
duration_seconds: float):
self.src_root = Path(src_root)
self.output_root = Path(output_root)
self.sample_rate = sample_rate
self.duration_seconds = duration_seconds
self.output_root.mkdir(parents=True, exist_ok=True)
def join(self, video_id: str, output_name: str, audio: torch.Tensor):
video_path = self.src_root / f'{video_id}.mp4'
output_path = self.output_root / f'{output_name}.mp4'
merge_audio_into_video(video_path, output_path, audio, self.sample_rate,
self.duration_seconds)
def merge_audio_into_video(video_path: Union[str, Path], output_path: Union[str, Path],
audio: torch.Tensor, sample_rate: int, duration_seconds: float):
# audio: (num_samples, num_channels=1/2)
frame_rate = 24
# read the video
reader = StreamingMediaDecoder(video_path)
reader.add_basic_video_stream(
frames_per_chunk=int(frame_rate * duration_seconds),
# buffer_chunk_size=1, # does not work with this -- extracted audio would be too short
format="rgb24",
frame_rate=frame_rate,
)
reader.fill_buffer()
video_chunk = reader.pop_chunks()[0]
t, _, h, w = video_chunk.shape
writer = StreamingMediaEncoder(output_path)
writer.add_audio_stream(
sample_rate=sample_rate,
num_channels=audio.shape[-1],
encoder="libmp3lame",
)
writer.add_video_stream(frame_rate=frame_rate,
width=w,
height=h,
format="rgb24",
encoder="libx264",
encoder_format="yuv420p")
with writer.open():
writer.write_audio_chunk(0, audio.float())
writer.write_video_chunk(1, video_chunk)
if __name__ == '__main__':
# Usage example
import sys
audio = torch.randn(16000 * 4, 1)
merge_audio_into_video(sys.argv[1], sys.argv[2], audio, 16000, 4)
|