|
|
import os |
|
|
import cv2 |
|
|
import torch |
|
|
import numpy as np |
|
|
import imageio |
|
|
import torchvision |
|
|
from einops import rearrange |
|
|
|
|
|
|
|
|
def save_videos_grid(videos: torch.Tensor, path: str, rescale=False, n_rows=6, fps=8, quality=8): |
|
|
videos = rearrange(videos, "b c t h w -> t b c h w") |
|
|
outputs = [] |
|
|
for x in videos: |
|
|
x = torchvision.utils.make_grid(x, nrow=n_rows) |
|
|
x = x.transpose(0, 1).transpose(1, 2).squeeze(-1) |
|
|
if rescale: |
|
|
x = (x + 1.0) / 2.0 |
|
|
x = torch.clamp(x,0,1) |
|
|
x = (x * 255).numpy().astype(np.uint8) |
|
|
outputs.append(x) |
|
|
|
|
|
os.makedirs(os.path.dirname(path), exist_ok=True) |
|
|
imageio.mimsave(path, outputs, fps=fps, quality=quality) |
|
|
|
|
|
def pad_image(crop_img, size, color=(255, 255, 255), resize_ratio=1): |
|
|
crop_h, crop_w = crop_img.shape[:2] |
|
|
target_w, target_h = size |
|
|
scale_h, scale_w = target_h / crop_h, target_w / crop_w |
|
|
if scale_w > scale_h: |
|
|
resize_h = int(target_h*resize_ratio) |
|
|
resize_w = int(crop_w / crop_h * resize_h) |
|
|
else: |
|
|
resize_w = int(target_w*resize_ratio) |
|
|
resize_h = int(crop_h / crop_w * resize_w) |
|
|
crop_img = cv2.resize(crop_img, (resize_w, resize_h)) |
|
|
pad_left = (target_w - resize_w) // 2 |
|
|
pad_top = (target_h - resize_h) // 2 |
|
|
pad_right = target_w - resize_w - pad_left |
|
|
pad_bottom = target_h - resize_h - pad_top |
|
|
crop_img = cv2.copyMakeBorder(crop_img, pad_top, pad_bottom, pad_left, pad_right, cv2.BORDER_CONSTANT, value=color) |
|
|
return crop_img |
