Spaces:
Running
on
Zero
Running
on
Zero
File size: 2,196 Bytes
a3a3ae4 |
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 |
import torch
import cv2
import numpy as np
import imageio
import torchvision
import math
from einops import rearrange
from typing import List
from PIL import Image, ImageDraw, ImageFont
def tensor2vid(video: torch.Tensor, mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) -> List[np.ndarray]:
mean = torch.tensor(mean, device=video.device).reshape(1, -1, 1, 1, 1) # ncfhw
std = torch.tensor(std, device=video.device).reshape(1, -1, 1, 1, 1) # ncfhw
video = video.mul_(std).add_(mean) # unnormalize back to [0,1]
video.clamp_(0, 1)
images = rearrange(video, 'i c f h w -> (i f) h w c')
images = images.unbind(dim=0)
images = [(image.cpu().numpy() * 255).astype('uint8') for image in images] # f h w c
return images
def export_to_video(video_frames: List[np.ndarray], output_video_path: str = None, save_to_gif=False, use_cv2=True, fps=8) -> str:
h, w, c = video_frames[0].shape
if save_to_gif:
image_lst = []
if output_video_path.endswith('mp4'):
output_video_path = output_video_path[:-3] + 'gif'
for i in range(len(video_frames)):
image_lst.append(video_frames[i])
imageio.mimsave(output_video_path, image_lst, fps=fps)
return output_video_path
else:
if use_cv2:
fourcc = cv2.VideoWriter_fourcc(*"mp4v")
video_writer = cv2.VideoWriter(output_video_path, fourcc, fps=fps, frameSize=(w, h))
for i in range(len(video_frames)):
img = cv2.cvtColor(video_frames[i], cv2.COLOR_RGB2BGR)
video_writer.write(img)
video_writer.release()
else:
duration = math.ceil(len(video_frames) / fps)
append_num = duration * fps - len(video_frames)
for k in range(append_num): video_frames.append(video_frames[-1])
video_stack = np.stack(video_frames, axis=0)
video_tensor = torch.from_numpy(video_stack)
# torchvision.io.write_video(output_video_path, video_tensor, fps=fps, options={"crf": "17"})
torchvision.io.write_video(output_video_path, video_tensor, fps=fps, options={"crf": "17"})
return output_video_path |