Spaces:
Sleeping
Sleeping
File size: 5,734 Bytes
40b178e |
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 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 |
import logging
import tempfile
from typing import List, Optional, Tuple, Union
import matplotlib
import mediapy
import numpy as np
import PIL.Image
import torch
from decord import VideoReader, cpu
logger = logging.getLogger(__name__)
dataset_res_dict = {
"sintel": [448, 1024],
"scannet": [640, 832],
"KITTI": [384, 1280],
"bonn": [512, 640],
"NYUv2": [448, 640],
}
def read_video_frames(
video_path: str,
process_length: int,
target_fps: int,
max_res: int,
dataset: str = "open",
) -> Tuple[np.ndarray, int]:
"""
Read video frames from a file, resize and downsample them.
Args:
video_path (str): Path to the video file.
process_length (int): Maximum number of frames to process.
target_fps (int): Target FPS for the output.
max_res (int): Maximum resolution (height or width).
dataset (str): Dataset name for resolution settings.
Returns:
Tuple[np.ndarray, int]: A tuple containing the frames (numpy array) and the actual FPS.
"""
if dataset == "open":
logger.info(f"Processing video: {video_path}")
vid = VideoReader(video_path, ctx=cpu(0))
logger.info(
f"Original video shape: {(len(vid), *vid.get_batch([0]).shape[1:])}"
)
original_height, original_width = vid.get_batch([0]).shape[1:3]
height = round(original_height / 64) * 64
width = round(original_width / 64) * 64
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
else:
height = dataset_res_dict[dataset][0]
width = dataset_res_dict[dataset][1]
vid = VideoReader(video_path, ctx=cpu(0), width=width, height=height)
fps = vid.get_avg_fps() if target_fps == -1 else target_fps
stride = round(vid.get_avg_fps() / fps)
stride = max(stride, 1)
frames_idx = list(range(0, len(vid), stride))
logger.info(
f"Downsampled shape: {(len(frames_idx), *vid.get_batch([0]).shape[1:])}, with stride: {stride}"
)
if process_length != -1 and process_length < len(frames_idx):
frames_idx = frames_idx[:process_length]
logger.info(
f"Final processing shape: {(len(frames_idx), *vid.get_batch([0]).shape[1:])}"
)
frames = vid.get_batch(frames_idx).asnumpy().astype("float32") / 255.0
return frames, fps
def save_video(
video_frames: Union[List[np.ndarray], List[PIL.Image.Image], np.ndarray],
output_video_path: Optional[str] = None,
fps: int = 10,
crf: int = 18,
) -> str:
"""
Save video frames to a file.
Args:
video_frames (Union[List[np.ndarray], List[PIL.Image.Image], np.ndarray]): List of frames or numpy array.
output_video_path (Optional[str]): Path to save the video. If None, a temporary file is created.
fps (int): Frames per second.
crf (int): Constant Rate Factor for encoding quality.
Returns:
str: Path to the saved video.
"""
if output_video_path is None:
output_video_path = tempfile.NamedTemporaryFile(suffix=".mp4").name
if isinstance(video_frames, np.ndarray):
# If it's a numpy array, we assume it's already in the correct format or needs simple conversion
if video_frames.dtype != np.uint8:
video_frames = (video_frames * 255).astype(np.uint8)
elif isinstance(video_frames[0], np.ndarray):
video_frames = [(frame * 255).astype(np.uint8) for frame in video_frames]
elif isinstance(video_frames[0], PIL.Image.Image):
video_frames = [np.array(frame) for frame in video_frames]
mediapy.write_video(output_video_path, video_frames, fps=fps, crf=crf)
return output_video_path
class ColorMapper:
"""
A color mapper to map depth values to a certain colormap.
"""
def __init__(self, colormap: str = "inferno"):
"""
Initialize the ColorMapper.
Args:
colormap (str): Name of the colormap to use.
"""
self.colormap = torch.tensor(matplotlib.colormaps[colormap].colors)
def apply(
self,
image: torch.Tensor,
v_min: Optional[float] = None,
v_max: Optional[float] = None,
) -> torch.Tensor:
"""
Apply the colormap to an image.
Args:
image (torch.Tensor): Input image tensor.
v_min (Optional[float]): Minimum value for normalization.
v_max (Optional[float]): Maximum value for normalization.
Returns:
torch.Tensor: Color-mapped image.
"""
if v_min is None:
v_min = image.min()
if v_max is None:
v_max = image.max()
image = (image - v_min) / (v_max - v_min)
image = (image * 255).long()
# Clamp values to be within valid range for indexing
image = torch.clamp(image, 0, 255)
image = self.colormap[image]
return image
def vis_sequence_depth(
depths: np.ndarray, v_min: Optional[float] = None, v_max: Optional[float] = None
) -> np.ndarray:
"""
Visualize a sequence of depth maps.
Args:
depths (np.ndarray): Input depth maps.
v_min (Optional[float]): Minimum value for normalization.
v_max (Optional[float]): Maximum value for normalization.
Returns:
np.ndarray: Visualized depth maps.
"""
visualizer = ColorMapper()
if v_min is None:
v_min = depths.min()
if v_max is None:
v_max = depths.max()
res = visualizer.apply(torch.tensor(depths), v_min=v_min, v_max=v_max).numpy()
return res
|