DSatishchandra commited on
Commit
ea1a9c8
·
verified ·
1 Parent(s): 6512082

Create utils.py

Browse files
Files changed (1) hide show
  1. utils.py +24 -0
utils.py ADDED
@@ -0,0 +1,24 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import cv2
2
+ import torch
3
+ import numpy as np
4
+ from torchvision import transforms
5
+
6
+ def extract_frames(video_path, num_frames=71, resize=(224, 224)):
7
+ cap = cv2.VideoCapture(video_path)
8
+ total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
9
+ interval = max(total // num_frames, 1)
10
+
11
+ frames = []
12
+ for i in range(num_frames):
13
+ cap.set(cv2.CAP_PROP_POS_FRAMES, i * interval)
14
+ ret, frame = cap.read()
15
+ if not ret:
16
+ break
17
+ frame = cv2.resize(frame, resize)
18
+ frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
19
+ frames.append(frame / 255.0)
20
+ cap.release()
21
+
22
+ frames = np.stack(frames, axis=0)
23
+ frames = torch.tensor(frames, dtype=torch.float32).permute(0, 3, 1, 2) # (T, C, H, W)
24
+ return frames.unsqueeze(0) # (1, T, C, H, W)