Spaces:
No application file
No application file
File size: 5,126 Bytes
152e3c7 |
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 155 156 |
""" "This module contains a video loader."""
import logging
import os
import sys
from typing import List, Tuple, Union
import numpy as np
from torch import Tensor
from torch.utils import data
from torchvision.datasets.video_utils import VideoClips
class VideoIter(data.Dataset):
"""This class implements a loader for videos."""
def __init__(
self,
clip_length,
frame_stride,
dataset_path=None,
video_transform=None,
return_label=False,
) -> None:
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)
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=self.total_clip_length_in_frames,
)
@property
def video_count(self) -> int:
"""Retrieve the number of the videos in the dataset."""
return len(self.video_list)
def getitem_from_raw_video(
self, idx: int
) -> Union[Tuple[Tensor, int, str, str], Tuple[Tensor, int, int, str, str]]:
"""Fetch a sample from the dataset.
Args:
idx (int): Index of the sample the retrieve.
Returns:
Tuple[Tensor, int, str, str]: Video clip, clip idx in the video, directory name, and file
Tuple[Tensor, int, int, str, str]: Video clip, label, clip idx in the video, directory name, and file
"""
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) -> int:
return len(self.video_clips)
def __getitem__(self, index: int):
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]
if trace_back is not None:
line = str(trace_back.tb_lineno)
else:
line = "no-line"
# pylint: disable=line-too-long
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: str) -> List[str]:
"""Fetche all videos in a directory and sub-directories.
Args:
dataset_path (str): A string that represents the directory of the dataset.
Raises:
FileNotFoundError: The directory could not be found in the provided path.
Returns:
List[str]
"""
if not os.path.exists(dataset_path):
raise FileNotFoundError(f"VideoIter:: failed to locate: `{dataset_path}'")
vid_list = []
for path, _, files in os.walk(dataset_path):
for name in files:
if "mp4" not in name:
continue
vid_list.append(os.path.join(path, name))
logging.info(f"Found {len(vid_list)} video files in {dataset_path}")
return vid_list
class SingleVideoIter(VideoIter):
"""Loader for a single video."""
def __init__(
self,
clip_length,
frame_stride,
video_path,
video_transform=None,
return_label=False,
) -> None:
super().__init__(
clip_length, frame_stride, video_path, video_transform, return_label
)
def _get_video_list(self, dataset_path: str) -> List[str]:
return [dataset_path]
def __getitem__(self, idx: int) -> Tensor:
video, _, _, _ = self.video_clips.get_clip(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)
return video
|