File size: 6,534 Bytes
aae3ba1 |
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 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 |
# Utils for video processing, frame manipulation, and visualization.
import cv2
import numpy as np
from typing import List
import imageio
def rotate_frame(frame: np.ndarray) -> np.ndarray:
"""Rotate a frame 90 degrees clockwise."""
return cv2.rotate(frame, cv2.ROTATE_90_CLOCKWISE)
def center_crop_image(frame: np.ndarray, crop_percent: float = 1.0) -> np.ndarray:
"""
Center crop an image to a specified percentage of its original size.
Args:
frame (np.ndarray): Input image.
crop_percent (float): Percentage of original size to crop to (0 < crop_percent <= 1).
Returns:
np.ndarray: Cropped image.
"""
if crop_percent == 1.0:
return frame
# Get original dimensions
original_height, original_width = frame.shape[:2]
# Calculate new dimensions
new_width = int(original_width * crop_percent)
new_height = int(original_height * crop_percent)
# Calculate top-left corner of the cropping box
start_x = (original_width - new_width) // 2
start_y = (original_height - new_height) // 2
# Perform the crop
cropped_frame = frame[start_y:start_y + new_height, start_x:start_x + new_width]
return cropped_frame
def read_video_frames(
cap,
start_frame: int = None,
end_frame: int = None,
interval: int = 1,
rotate: bool = False,
crop_percent: float = 1.0
) -> List[np.ndarray]:
"""
Read frames from a video capture object with optional rotation and cropping.
Args:
cap: OpenCV VideoCapture object.
start_frame (int): Starting frame index.
end_frame (int): Ending frame index.
interval (int): Frame interval for sampling.
rotate (bool): Whether to rotate frames 90 degrees clockwise.
crop_percent (float): Center crop percentage.
Returns:
List[np.ndarray]: List of frames.
"""
frame_count = 0
frame_list = []
# If a start frame is specified, move the video pointer
if start_frame is not None:
cap.set(cv2.CAP_PROP_POS_FRAMES, start_frame)
frame_count += start_frame
# If no end frame is specified, read until the end of the video
if end_frame is None:
end_frame = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
# Read frames in a loop
while frame_count < end_frame:
ret, frame = cap.read()
if not ret:
break
# Process frames at the specified interval
if frame_count % interval == 0:
# Optionally rotate the frame
if rotate:
frame = rotate_frame(frame)
# Optionally center crop the frame
frame = center_crop_image(frame, crop_percent=crop_percent)
frame_list.append(frame)
frame_count += 1
return frame_list
def save_to_video(frames: List[np.ndarray], output_path: str, fps: int = 30):
"""
Save a list of frames to a video file.
Args:
frames (List[np.ndarray]): List of frames.
output_path (str): Output video path.
fps (int): Frames per second.
"""
imageio.mimsave(output_path, frames, fps=fps, codec='libx264')
def resize_frames_to_long_side(frames: List[np.ndarray], target_long_side: int) -> List[np.ndarray]:
"""
Resize frames so the longer side matches the target size, preserving aspect ratio.
Args:
frames (List[np.ndarray]): List of frames.
target_long_side (int): Desired length of the longer side.
Returns:
List[np.ndarray]: Resized frames.
"""
# If target_long_side is None, return original frames without resizing
if target_long_side is None:
return frames
resized_frames = []
# Loop over each frame
for frame in frames:
height, width = frame.shape[:2]
# Determine the scaling factor and new dimensions
if width > height:
scale_factor = target_long_side / width
else:
scale_factor = target_long_side / height
new_width = int(width * scale_factor)
new_height = int(height * scale_factor)
# Resize the frame while maintaining the aspect ratio
resized_frame = cv2.resize(frame, (new_width, new_height), interpolation=cv2.INTER_AREA)
resized_frames.append(resized_frame)
return resized_frames
def sample_frames_evenly(video_frames: List[np.ndarray], num_frames: int) -> List[np.ndarray]:
"""
Sample frames evenly from a video sequence.
Args:
video_frames (List[np.ndarray]): List of frames.
num_frames (int): Number of frames to sample.
Returns:
List[np.ndarray]: Sampled frames.
"""
total = len(video_frames)
if num_frames >= total:
return video_frames.copy()
indices = np.linspace(0, total - 1, num=num_frames, dtype=int)
return [video_frames[i] for i in indices]
def wrap_text(text: str, max_width: int, font, font_scale: float) -> List[str]:
"""
Wraps text to fit within a given width.
Args:
text (str): The text to wrap.
max_width (int): The maximum width in pixels.
font: The font to use.
font_scale (float): The scale of the font.
Returns:
List of lines of wrapped text.
"""
words = text.split(' ')
lines = []
current_line = ''
for word in words:
test_line = current_line + word + ' '
# Measure the width of the line
size = cv2.getTextSize(test_line, font, font_scale, 1)[0]
if size[0] > max_width:
lines.append(current_line.strip())
current_line = word + ' '
else:
current_line = test_line
if current_line:
lines.append(current_line.strip())
return lines
def add_overlay_text(frame: np.ndarray, caption: str) -> np.ndarray:
"""
Add overlay text to a frame.
Args:
frame (np.ndarray): Input image.
caption (str): Text to overlay.
Returns:
np.ndarray: Frame with text.
"""
w = frame.shape[1] # Width of the frame
y0, dy = 30, 20 # Starting Y position and line height
for i, line in enumerate(wrap_text(caption, w, cv2.FONT_HERSHEY_SIMPLEX, 1.0)):
y = y0 + i * dy
cv2.putText(frame, line, (10, y), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 255), 2)
return frame |