File size: 2,377 Bytes
b20701a | 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 | '''
Utility functions for handling data
'''
import torch
import cv2
import numpy as np
import unicodedata
from pathlib import Path
def get_sample_frames(video_path, num_frames=5):
cap = cv2.VideoCapture(video_path)
if not cap.isOpened():
return None
total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
if total_frames == 0:
cap.release()
return None
indices = np.linspace(0, total_frames - 1, num_frames, dtype=int)
frames = []
for idx in indices:
cap.set(cv2.CAP_PROP_POS_FRAMES, idx)
ret, frame = cap.read()
if ret:
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
frames.append(frame)
cap.release()
return frames
def read_video(video_path):
cap = cv2.VideoCapture(video_path)
frames = []
while True:
ret, frame = cap.read()
if not ret:
break
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
frames.append(frame)
cap.release()
if len(frames) == 0:
raise ValueError(f"Could not read any frames from {video_path}")
frames = torch.from_numpy(np.stack(frames, axis=0))
return frames
def get_all_path(root, labeled=True):
root = Path(root)
all_path = []
if labeled:
for cls in root.iterdir():
if not cls.is_dir():
continue
for video_path in cls.iterdir():
if not video_path.is_file():
continue
all_path.append(video_path)
else:
for video_path in root.iterdir():
if not video_path.is_file():
continue
all_path.append(video_path)
return all_path
def get_video_metadata(video_path):
cap = cv2.VideoCapture(video_path)
if not cap.isOpened():
return None
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
fps = cap.get(cv2.CAP_PROP_FPS)
cap.release()
return height, width, frame_count, fps
def nfc_normalize(s):
return unicodedata.normalize("NFC", s) |