-
Notifications
You must be signed in to change notification settings - Fork 0
/
video_loader.py
51 lines (36 loc) · 1.3 KB
/
video_loader.py
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
"""VideoFolder Class used in the script entitled test_inet_p.py.
This code has been adapted from https://github.com/hendrycks/robustness """
import cv2
# from skvideo.io import VideoCapture
# import skvideo.io
import torch
import torch.utils.data as data
from torchvision.datasets.folder import DatasetFolder
from PIL import Image
import os
import os.path
import sys
class VideoFolder(DatasetFolder):
def __init__(self, root, transform=None, target_transform=None, loader=None):
super(VideoFolder, self).__init__(
root, loader, '.mp4', transform=transform, target_transform=target_transform)
self.vids = self.samples
def __getitem__(self, index):
"""
Args:
index (int): Index
Returns:
tuple: (sample, target) where target is class_index of the target class.
"""
path, target = self.samples[index]
# cap = VideoCapture(path)
cap = cv2.VideoCapture(path)
frames = []
while True:
# Capture frame-by-frame
ret, frame = cap.read()
if not ret: break
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
frames.append(self.transform(Image.fromarray(frame)).unsqueeze(0))
cap.release()
return torch.cat(frames, 0), target