|
|
import torch |
|
|
from einops import rearrange |
|
|
from moviepy.editor import VideoFileClip |
|
|
import numpy as np |
|
|
from moviepy.editor import (VideoFileClip, |
|
|
AudioFileClip, |
|
|
concatenate_videoclips, |
|
|
CompositeAudioClip, |
|
|
ImageSequenceClip, |
|
|
) |
|
|
from moviepy.audio.AudioClip import AudioArrayClip |
|
|
def read_frames(video_path:str): |
|
|
video_clip = VideoFileClip(video_path) |
|
|
audio_clip = video_clip.audio |
|
|
frames = np.array(list(video_clip.iter_frames())) |
|
|
video_fps = video_clip.fps |
|
|
return { |
|
|
"video": frames, |
|
|
"fps": video_fps, |
|
|
"audio_clip": audio_clip |
|
|
} |
|
|
def audio2array(audio_clip,start_frame,end_frame,video_fps,sampling_rate,audio_len): |
|
|
audio_fps = audio_clip.fps |
|
|
start_time = start_frame / video_fps |
|
|
end_time = end_frame / video_fps |
|
|
sampled_audio = audio_clip.subclip(start_time, end_time) |
|
|
sampled_audio = sampled_audio.set_fps(sampling_rate) |
|
|
audio_chunks = list(sampled_audio.iter_chunks(fps=sampling_rate,chunksize=50000)) |
|
|
audio = np.concatenate(audio_chunks, axis=0) |
|
|
if len(audio.shape) == 2: |
|
|
audio = np.mean(audio,axis=1) |
|
|
audio = np.interp(np.arange(0,audio_len),np.arange(0,len(audio)),audio) |
|
|
return { |
|
|
"audio": audio, |
|
|
"fps": audio_fps |
|
|
} |
|
|
def read_video(video_path,sampling_rate:int = 16000,frame_idx=None): |
|
|
video_clip = VideoFileClip(video_path) |
|
|
audio_clip = video_clip.audio |
|
|
|
|
|
frames = np.array(list(video_clip.iter_frames())) |
|
|
video_fps = video_clip.fps |
|
|
audio_fps = audio_clip.fps |
|
|
if frame_idx is not None: |
|
|
start_time = frame_idx[0] / video_fps |
|
|
end_time = frame_idx[-1] / video_fps |
|
|
else: |
|
|
start_time = 0 |
|
|
end_time = len(frames) / video_fps |
|
|
sampled_audio = audio_clip.subclip(start_time, end_time) |
|
|
sampled_audio = sampled_audio.set_fps(sampling_rate) |
|
|
audio_chunks = list(sampled_audio.iter_chunks(fps=sampling_rate,chunksize=50000)) |
|
|
audio = np.concatenate(audio_chunks, axis=0) |
|
|
if len(audio.shape) == 2: |
|
|
audio = np.mean(audio,axis=1) |
|
|
return { |
|
|
"video": frames, |
|
|
"audio": audio, |
|
|
"video_fps" : video_fps, |
|
|
"audio_fps": audio_fps |
|
|
} |
|
|
def write_video(output_path, frames, audio, video_fps, audio_fps): |
|
|
video_clip = ImageSequenceClip(list(frames), fps=video_fps) |
|
|
audio_clip = AudioArrayClip(audio,audio_fps) |
|
|
video_clip = video_clip.set_audio(audio_clip) |
|
|
video_clip.write_videofile(output_path, codec='libx264', audio_codec='aac') |
|
|
def tensor2frames(x:torch.Tensor): |
|
|
if len(x.shape) == 5: |
|
|
x = rearrange(x,'b t c h w -> (b t) c h w') |
|
|
frames = [] |
|
|
for image in x: |
|
|
image = image.permute(1, 2, 0) |
|
|
image_np = ((image / 2.0 + 0.5).clamp(0, 1) * 255).to(dtype=torch.uint8).cpu().contiguous().numpy() |
|
|
frames.append(image_np) |
|
|
return frames |
|
|
if __name__ == "__main__": |
|
|
video_path = "/mnt/spaceai-data/tts/team/digital_avatar_group/fenghe/datasets/celebv-hq/-3Bl8i34Z7Q_0.mp4" |
|
|
d = read_frames(video_path) |
|
|
print(audio2array(d["audio_clip"],0,16,24,16000,10000)) |
|
|
|
|
|
|
