| 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 |
