File size: 2,859 Bytes
272e50d | 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 | import os
import numpy as np
import torch
import torchvision
import torch.nn.functional as F
from PIL import Image
from pathlib import Path
import imageio
from einops import rearrange
import torchvision.transforms as transforms
def concat_pil(img_pil_lst):
# horizontally concat PIL images
# NOTE(ZSH): assume all images are of same size
W, H = img_pil_lst[0].size
num_img = len(img_pil_lst)
new_width = num_img * W
new_image = Image.new("RGB", (new_width, H), color=0)
for img_idx, img in enumerate(img_pil_lst):
new_image.paste(img, (W * img_idx, 0))
return new_image
def save_videos_from_pil(pil_images, path, fps=24, crf=23):
save_fmt = Path(path).suffix
os.makedirs(os.path.dirname(path), exist_ok=True)
if save_fmt == ".mp4":
with imageio.get_writer(path, fps=fps) as writer:
for img in pil_images:
img_array = np.array(img) # Convert PIL Image to numpy array
writer.append_data(img_array)
elif save_fmt == ".gif":
pil_images[0].save(
fp=path,
format="GIF",
append_images=pil_images[1:],
save_all=True,
duration=(1 / fps * 1000),
loop=0,
)
else:
raise ValueError("Unsupported file type. Use .mp4 or .gif.")
def save_videos_grid(videos: torch.Tensor, path: str, rescale=False, n_rows=6, fps=24):
videos = rearrange(videos, "b c t h w -> t b c h w")
height, width = videos.shape[-2:]
outputs = []
for x in videos:
x = torchvision.utils.make_grid(x, nrow=n_rows) # (c h w)
x = x.transpose(0, 1).transpose(1, 2).squeeze(-1) # (h w c)
if rescale:
x = (x + 1.0) / 2.0 # -1,1 -> 0,1
x = (x * 255).numpy().astype(np.uint8)
x = Image.fromarray(x)
outputs.append(x)
os.makedirs(os.path.dirname(path), exist_ok=True)
save_videos_from_pil(outputs, path, fps)
def resize_tensor_frames(video_tensor, new_size):
B, C, video_length, H, W = video_tensor.shape
# Reshape video tensor to combine batch and frame dimensions: (B*F, C, H, W)
video_tensor_reshaped = video_tensor.reshape(-1, C, H, W)
# Resize using interpolate
resized_frames = F.interpolate(
video_tensor_reshaped, size=new_size, mode="bilinear", align_corners=False
)
resized_video = resized_frames.reshape(B, C, video_length, new_size[0], new_size[1])
return resized_video
def pil_list_to_tensor(image_list, size=None):
to_tensor = transforms.ToTensor()
if size is not None:
tensor_list = [to_tensor(img.resize(size[::-1])) for img in image_list]
else:
tensor_list = [to_tensor(img) for img in image_list]
stacked_tensor = torch.stack(tensor_list, dim=0)
tensor = stacked_tensor.permute(1, 0, 2, 3)
return tensor
|