Spaces:
Sleeping
Sleeping
File size: 4,920 Bytes
eb3afa1 | 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 |
import cv2
import numpy as np
import torch
import os
from typing import List, Tuple, Optional
def get_video_properties(video_path: str) -> dict:
"""
Get video properties.
"""
cap = cv2.VideoCapture(video_path)
if not cap.isOpened():
return {}
props = {
'width': int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)),
'height': int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)),
'fps': cap.get(cv2.CAP_PROP_FPS),
'frame_count': int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
}
cap.release()
return props
def load_video_frames(
video_path: str,
num_frames: int = 30,
frame_size: Tuple[int, int] = (224, 224),
frame_skip: int = 1
) -> Optional[np.ndarray]:
"""
Load frames from a video with robust preprocessing:
1. Center crop to square (min dimension).
2. Resize to frame_size.
3. Sample frames uniformly.
"""
cap = cv2.VideoCapture(video_path)
if not cap.isOpened():
print(f"Error opening video: {video_path}")
return None
total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
if total_frames <= 0:
cap.release()
return None
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
# Calculate crop coordinates for center square crop
min_dim = min(width, height)
start_x = (width - min_dim) // 2
start_y = (height - min_dim) // 2
# Calculate frame indices to sample
# We want 'num_frames' frames.
# Strategy: evenly space them across the video duration we look at.
# But for simplicity and consistency, let's just grab them with a stride,
# or if video is short, grab all and pad.
# Let's try to span as much of the video as possible?
# Or just stick to the requested architecture of sampling segments.
# The prompt asked for "preprocess ... not based on aspect ratio".
# Simple strategy: Read frames with skip, up to num_frames.
# If video is too short, loop/pad?
# Better: Reservoir sampling or Linspace if we want fixed count?
# Let's stick to the user's likely need: Fixed number of frames.
sampled_frames = []
# We'll seek effectively.
# But looping is safer for some codecs.
# Improved sampling: pick indices using linspace if we want to span whole video?
# Or just sequential for temporal consistency (RNN prefer sequences).
# Let's do sequential with skip.
frame_idx = 0
frames_collected = 0
while frames_collected < num_frames:
ret, frame = cap.read()
if not ret:
break
if frame_idx % frame_skip == 0:
# Preprocess Frame
# 1. Center Crop
crop = frame[start_y:start_y+min_dim, start_x:start_x+min_dim]
# 2. Resize
resized = cv2.resize(crop, frame_size)
# 3. Convert BGR to RGB
rgb = cv2.cvtColor(resized, cv2.COLOR_BGR2RGB)
sampled_frames.append(rgb)
frames_collected += 1
frame_idx += 1
cap.release()
# Handle insufficient frames
if len(sampled_frames) == 0:
return None
if len(sampled_frames) < num_frames:
# Pad with last frame or zeros?
# Let's pad with zeros (black frames) or loop?
# Zero padding is safer to avoid motion artifacts.
padding = [np.zeros((frame_size[1], frame_size[0], 3), dtype=np.uint8)] * (num_frames - len(sampled_frames))
sampled_frames.extend(padding)
return np.array(sampled_frames)
def normalize_frames(frames: np.ndarray) -> np.ndarray:
"""
Normalize frames to [0, 1] and then standard ImageNet mean/std.
Frames input: (N, H, W, C) in RGB, uint8 [0,255]
"""
# Convert to float32 [0, 1]
frames_norm = frames.astype(np.float32) / 255.0
# Standard ImageNet mean and std
mean = np.array([0.485, 0.456, 0.406], dtype=np.float32)
std = np.array([0.229, 0.224, 0.225], dtype=np.float32)
# Apply normalization
# frames is (N, H, W, C), mean/std are (3,)
# We allow broadcasting on the last dimension
frames_norm = (frames_norm - mean) / std
return frames_norm
def validate_video(video_path: str) -> bool:
"""
Check if video is valid and openable.
"""
if not os.path.exists(video_path):
return False
try:
cap = cv2.VideoCapture(video_path)
if not cap.isOpened():
return False
# Read one frame to be sure
ret, _ = cap.read()
cap.release()
return ret
except Exception:
return False
|