File size: 4,386 Bytes
28e129b | 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 | import logging
import os
import sys
import numpy as np
from torch.utils import data
from torchvision.datasets.video_utils import VideoClips
class VideoIter(data.Dataset):
def __init__(self,
clip_length,
frame_stride,
proc_video=None,
dataset_path=None,
video_transform=None,
use_splits=False,
return_label=False,):
super().__init__()
# video clip properties
self.frames_stride = frame_stride
self.total_clip_length_in_frames = clip_length * frame_stride
self.video_transform = video_transform
# IO
self.dataset_path = dataset_path
self.video_list = self._get_video_list(dataset_path=self.dataset_path, proc_video=proc_video, use_splits=use_splits)
self.return_label = return_label
# data loading
self.video_clips = VideoClips(video_paths=self.video_list,
clip_length_in_frames=self.total_clip_length_in_frames,
frames_between_clips=int(self.total_clip_length_in_frames), # between start frame
)
@property
def video_count(self):
return len(self.video_list)
def getitem_from_raw_video(self, idx):
"""
Load and preprocess a video clip from the dataset.
Args:
idx (int): Index of the video clip to load.
Returns:
tuple: A tuple containing the loaded and preprocessed video clip, clip index, directory, and filename.
"""
video, _, _, _ = self.video_clips.get_clip(idx)
video_idx, clip_idx = self.video_clips.get_clip_location(idx)
video_path = self.video_clips.video_paths[video_idx]
in_clip_frames = list(range(0, self.total_clip_length_in_frames, self.frames_stride))
video = video[in_clip_frames]
if self.video_transform is not None:
video = self.video_transform(video)
dir, file = video_path.split(os.sep)[-2:]
file = file.split('.')[0]
if self.return_label:
label = 0 if "Normal" in video_path else 1
return video, label, clip_idx, dir, file
return video, clip_idx, dir, file
def __len__(self):
return len(self.video_clips)
def __getitem__(self, index):
succ = False
while not succ:
try:
batch = self.getitem_from_raw_video(index)
succ = True
except Exception as e:
index = np.random.choice(range(0, self.__len__()))
trace_back = sys.exc_info()[2]
line = trace_back.tb_lineno
logging.warning(f"VideoIter:: ERROR (line number {line}) !! (Force using another index:\n{index})\n{e}")
return batch
def _get_video_list(self, dataset_path, proc_video, use_splits):
"""
Get a list of video files from the dataset path, considering the specified processing options.
Args:
dataset_path (str): Path to the dataset directory.
proc_video (list or None): List of processed video filenames or None.
use_splits (bool): Whether to use splits for data selection.
Returns:
list: List of video file paths.
"""
assert os.path.exists(dataset_path), "VideoIter:: failed to locate: `{}'".format(dataset_path)
with open("/media/DataDrive/yiling/annotation/recognition/splits/VAD/splits_video_list.txt", "r") as file:
splits_files = [line.strip() for line in file.readlines()]
vid_list = []
for path, subdirs, files in os.walk(dataset_path):
for name in files:
file_path = os.path.join(path, name)
parent_dir = os.path.basename(os.path.dirname(file_path))
relative_path = os.path.join(parent_dir, name)
if use_splits and relative_path not in splits_files:
continue
if ('mp4' not in name and 'avi' not in name) or relative_path.split(".")[0] in proc_video:
continue
vid_list.append(os.path.join(path, name))
logging.info(f"Found {len(vid_list)} unprocessed video files in {dataset_path}")
return vid_list
|