File size: 1,656 Bytes
f474836 |
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 |
import os
import cv2
import torch
import einops
import torchvision
def resize_and_center_crop(image, target_width, target_height, interpolation=cv2.INTER_AREA):
original_height, original_width = image.shape[:2]
k = max(target_height / original_height, target_width / original_width)
new_width = int(round(original_width * k))
new_height = int(round(original_height * k))
resized_image = cv2.resize(image, (new_width, new_height), interpolation=interpolation)
x_start = (new_width - target_width) // 2
y_start = (new_height - target_height) // 2
cropped_image = resized_image[y_start:y_start + target_height, x_start:x_start + target_width]
return cropped_image
def save_bcthw_as_mp4(x, output_filename, fps=10):
b, c, t, h, w = x.shape
per_row = b
for p in [6, 5, 4, 3, 2]:
if b % p == 0:
per_row = p
break
os.makedirs(os.path.dirname(os.path.abspath(os.path.realpath(output_filename))), exist_ok=True)
x = torch.clamp(x.float(), -1., 1.) * 127.5 + 127.5
x = x.detach().cpu().to(torch.uint8)
x = einops.rearrange(x, '(m n) c t h w -> t (m h) (n w) c', n=per_row)
torchvision.io.write_video(output_filename, x, fps=fps, video_codec='h264', options={'crf': '1'})
return x
def save_bcthw_as_png(x, output_filename):
os.makedirs(os.path.dirname(os.path.abspath(os.path.realpath(output_filename))), exist_ok=True)
x = torch.clamp(x.float(), -1., 1.) * 127.5 + 127.5
x = x.detach().cpu().to(torch.uint8)
x = einops.rearrange(x, 'b c t h w -> c (b h) (t w)')
torchvision.io.write_png(x, output_filename)
return output_filename
|