Spaces:
Running
Running
File size: 3,823 Bytes
55ac26c bddb8a1 55ac26c bddb8a1 55ac26c bddb8a1 891bd26 bddb8a1 891bd26 55ac26c bddb8a1 55ac26c bddb8a1 55ac26c bddb8a1 55ac26c bddb8a1 55ac26c bddb8a1 55ac26c bddb8a1 55ac26c bddb8a1 55ac26c fecb864 bddb8a1 fecb864 bddb8a1 55ac26c bddb8a1 |
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 |
# This file is originally from DepthCrafter/depthcrafter/utils.py at main · Tencent/DepthCrafter
# SPDX-License-Identifier: MIT License license
#
# This file may have been modified by ByteDance Ltd. and/or its affiliates on [date of modification]
# Original file is released under [ MIT License license], with the full license text available at [https://github.com/Tencent/DepthCrafter?tab=License-1-ov-file].
import numpy as np
import matplotlib.cm as cm
import imageio
try:
from decord import VideoReader, cpu
DECORD_AVAILABLE = True
except:
import cv2
DECORD_AVAILABLE = False
def ensure_even(value):
return value if value % 2 == 0 else value + 1
def read_video_frames(video_path, process_length, target_fps=-1, max_res=-1):
if DECORD_AVAILABLE:
vid = VideoReader(video_path, ctx=cpu(0))
original_height, original_width = vid.get_batch([0]).shape[1:3]
height = original_height
width = original_width
print(f'==> original video size: {original_height} x {original_width}')
if max_res > 0 and max(height, width) > max_res:
scale = max_res / max(original_height, original_width)
height = ensure_even(round(original_height * scale))
width = ensure_even(round(original_width * scale))
print(f'==> downsample video size: {height} x {width}')
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))
if process_length != -1 and process_length < len(frames_idx):
frames_idx = frames_idx[:process_length]
frames = vid.get_batch(frames_idx).asnumpy()
else:
cap = cv2.VideoCapture(video_path)
original_fps = cap.get(cv2.CAP_PROP_FPS)
original_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
original_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
if max_res > 0 and max(original_height, original_width) > max_res:
scale = max_res / max(original_height, original_width)
height = round(original_height * scale)
width = round(original_width * scale)
fps = original_fps if target_fps < 0 else target_fps
stride = max(round(original_fps / fps), 1)
frames = []
frame_count = 0
while cap.isOpened():
ret, frame = cap.read()
if not ret or (process_length > 0 and frame_count >= process_length):
break
if frame_count % stride == 0:
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) # Convert BGR to RGB
if max_res > 0 and max(original_height, original_width) > max_res:
frame = cv2.resize(frame, (width, height)) # Resize frame
frames.append(frame)
frame_count += 1
cap.release()
frames = np.stack(frames, axis=0)
return frames, fps
def save_video(frames, output_video_path, fps=10, is_depths=False, grayscale=False):
writer = imageio.get_writer(output_video_path, fps=fps, macro_block_size=1, codec='libx264', ffmpeg_params=['-crf', '18'])
if is_depths:
colormap = np.array(cm.get_cmap("inferno").colors)
d_min, d_max = frames.min(), frames.max()
for i in range(frames.shape[0]):
depth = frames[i]
depth_norm = ((depth - d_min) / (d_max - d_min) * 255).astype(np.uint8)
depth_vis = (colormap[depth_norm] * 255).astype(np.uint8) if not grayscale else depth_norm
writer.append_data(depth_vis)
else:
for i in range(frames.shape[0]):
writer.append_data(frames[i])
writer.close()
|