Spaces:
Running
on
Zero
Running
on
Zero
File size: 1,815 Bytes
861fa04 |
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 |
import numpy as np
import cv2
import matplotlib.pyplot as plt
import torch
def resize_max_res(video_rgb, max_res, interpolation=cv2.INTER_LINEAR):
"""
Resize the video to the max resolution while keeping the aspect ratio.
Args:
video_rgb: (T, H, W, 3), RGB video, uint8
max_res: int, max resolution
Returns:
video_rgb: (T, H_new, W_new, 3), resized RGB video, uint8
"""
original_height = video_rgb.shape[1]
original_width = video_rgb.shape[2]
# round the height and width to the nearest multiple of 64
height = round(original_height / 64) * 64
width = round(original_width / 64) * 64
# resize the video if the height or width is larger than max_res
if max(height, width) > max_res:
scale = max_res / max(original_height, original_width)
height = round(original_height * scale / 64) * 64
width = round(original_width * scale / 64) * 64
frames = []
for i in range(video_rgb.shape[0]):
frames.append(cv2.resize(video_rgb[i], (width, height), interpolation=interpolation))
frames = np.array(frames)
return frames
def colorize_video_depth(depth_video, colormap="Spectral"):
"""
Colorize the depth video using the specified colormap.
depth_video: (T, H, W), depth video, [0, 1]
return:
colored_depth_video: (T, H, W, 3), colored depth video, dtype=uint8
"""
if isinstance(depth_video, torch.Tensor):
depth_video = depth_video.cpu().numpy()
T, H, W = depth_video.shape
colored_depth_video = []
for i in range(T):
colored_depth = plt.get_cmap(colormap)(depth_video[i], bytes=True)[...,:3]
colored_depth_video.append(colored_depth)
colored_depth_video = np.stack(colored_depth_video, axis=0)
return colored_depth_video |