Permalink
Cannot retrieve contributors at this time
Name already in use
A tag already exists with the provided branch name. Many Git commands accept both tag and branch names, so creating this branch may cause unexpected behavior. Are you sure you want to create this branch?
Aleksandar_Mitev_Dissertation/video_dataset.py
Go to fileThis commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
270 lines (225 sloc)
11.6 KB
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
import os | |
import os.path | |
import numpy as np | |
from PIL import Image | |
from torchvision import transforms | |
import torch | |
from typing import List, Union, Tuple, Any | |
# Code copied from Koot., R., E.(2021.) Efficient Video Fataset Loading and Augmentation, GitHub. https://github.com/RaivoKoot/Video-Dataset-Loading-Pytorch | |
class VideoRecord(object): | |
""" | |
Helper class for class VideoFrameDataset. This class | |
represents a video sample's metadata. | |
Args: | |
root_datapath: the system path to the root folder | |
of the videos. | |
row: A list with four or more elements where 1) The first | |
element is the path to the video sample's frames excluding | |
the root_datapath prefix 2) The second element is the starting frame id of the video | |
3) The third element is the inclusive ending frame id of the video | |
4) The fourth element is the label index. | |
5) any following elements are labels in the case of multi-label classification | |
""" | |
def __init__(self, row, root_datapath): | |
self._data = row | |
self._path = os.path.join(root_datapath, row[0]) | |
@property | |
def path(self) -> str: | |
return self._path | |
@property | |
def num_frames(self) -> int: | |
return self.end_frame - self.start_frame + 1 # +1 because end frame is inclusive | |
@property | |
def start_frame(self) -> int: | |
return int(self._data[1]) | |
@property | |
def end_frame(self) -> int: | |
return int(self._data[2]) | |
@property | |
def label(self) -> Union[int, List[int]]: | |
# just one label_id | |
if len(self._data) == 4: | |
return int(self._data[3]) | |
# sample associated with multiple labels | |
else: | |
return [int(label_id) for label_id in self._data[3:]] | |
class VideoFrameDataset(torch.utils.data.Dataset): | |
r""" | |
A highly efficient and adaptable dataset class for videos. | |
Instead of loading every frame of a video, | |
loads x RGB frames of a video (sparse temporal sampling) and evenly | |
chooses those frames from start to end of the video, returning | |
a list of x PIL images or ``FRAMES x CHANNELS x HEIGHT x WIDTH`` | |
tensors where FRAMES=x if the ``ImglistToTensor()`` | |
transform is used. | |
More specifically, the frame range [START_FRAME, END_FRAME] is divided into NUM_SEGMENTS | |
segments and FRAMES_PER_SEGMENT consecutive frames are taken from each segment. | |
Note: | |
A demonstration of using this class can be seen | |
in ``demo.py`` | |
https://github.com/RaivoKoot/Video-Dataset-Loading-Pytorch | |
Note: | |
This dataset broadly corresponds to the frame sampling technique | |
introduced in ``Temporal Segment Networks`` at ECCV2016 | |
https://arxiv.org/abs/1608.00859. | |
Note: | |
This class relies on receiving video data in a structure where | |
inside a ``ROOT_DATA`` folder, each video lies in its own folder, | |
where each video folder contains the frames of the video as | |
individual files with a naming convention such as | |
img_001.jpg ... img_059.jpg. | |
For enumeration and annotations, this class expects to receive | |
the path to a .txt file where each video sample has a row with four | |
(or more in the case of multi-label, see README on Github) | |
space separated values: | |
``VIDEO_FOLDER_PATH START_FRAME END_FRAME LABEL_INDEX``. | |
``VIDEO_FOLDER_PATH`` is expected to be the path of a video folder | |
excluding the ``ROOT_DATA`` prefix. For example, ``ROOT_DATA`` might | |
be ``home\data\datasetxyz\videos\``, inside of which a ``VIDEO_FOLDER_PATH`` | |
might be ``jumping\0052\`` or ``sample1\`` or ``00053\``. | |
Args: | |
root_path: The root path in which video folders lie. | |
this is ROOT_DATA from the description above. | |
annotationfile_path: The .txt annotation file containing | |
one row per video sample as described above. | |
num_segments: The number of segments the video should | |
be divided into to sample frames from. | |
frames_per_segment: The number of frames that should | |
be loaded per segment. For each segment's | |
frame-range, a random start index or the | |
center is chosen, from which frames_per_segment | |
consecutive frames are loaded. | |
imagefile_template: The image filename template that video frame files | |
have inside of their video folders as described above. | |
transform: Transform pipeline that receives a list of PIL images/frames. | |
test_mode: If True, frames are taken from the center of each | |
segment, instead of a random location in each segment. | |
""" | |
def __init__(self, | |
root_path: str, | |
annotationfile_path: str, | |
num_segments: int = 3, | |
frames_per_segment: int = 1, | |
imagefile_template: str='img_{:05d}.jpg', | |
transform = None, | |
test_mode: bool = False): | |
super(VideoFrameDataset, self).__init__() | |
self.root_path = root_path | |
self.annotationfile_path = annotationfile_path | |
self.num_segments = num_segments | |
self.frames_per_segment = frames_per_segment | |
self.imagefile_template = imagefile_template | |
self.transform = transform | |
self.test_mode = test_mode | |
self._parse_annotationfile() | |
self._sanity_check_samples() | |
def _load_image(self, directory: str, idx: int) -> Image.Image: | |
return Image.open(os.path.join(directory, self.imagefile_template.format(idx))).convert('RGB') | |
def _parse_annotationfile(self): | |
self.video_list = [VideoRecord(x.strip().split(), self.root_path) for x in open(self.annotationfile_path)] | |
def _sanity_check_samples(self): | |
for record in self.video_list: | |
if record.num_frames <= 0 or record.start_frame == record.end_frame: | |
print(f"\nDataset Warning: video {record.path} seems to have zero RGB frames on disk!\n") | |
elif record.num_frames < (self.num_segments * self.frames_per_segment): | |
print(f"\nDataset Warning: video {record.path} has {record.num_frames} frames " | |
f"but the dataloader is set up to load " | |
f"(num_segments={self.num_segments})*(frames_per_segment={self.frames_per_segment})" | |
f"={self.num_segments * self.frames_per_segment} frames. Dataloader will throw an " | |
f"error when trying to load this video.\n") | |
def _get_start_indices(self, record: VideoRecord) -> 'np.ndarray[int]': | |
""" | |
For each segment, choose a start index from where frames | |
are to be loaded from. | |
Args: | |
record: VideoRecord denoting a video sample. | |
Returns: | |
List of indices of where the frames of each | |
segment are to be loaded from. | |
""" | |
# choose start indices that are perfectly evenly spread across the video frames. | |
if self.test_mode: | |
distance_between_indices = (record.num_frames - self.frames_per_segment + 1) / float(self.num_segments) | |
start_indices = np.array([int(distance_between_indices / 2.0 + distance_between_indices * x) | |
for x in range(self.num_segments)]) | |
# randomly sample start indices that are approximately evenly spread across the video frames. | |
else: | |
max_valid_start_index = (record.num_frames - self.frames_per_segment + 1) // self.num_segments | |
start_indices = np.multiply(list(range(self.num_segments)), max_valid_start_index) + \ | |
np.random.randint(max_valid_start_index, size=self.num_segments) | |
return start_indices | |
def __getitem__(self, idx: int) -> Union[ | |
Tuple[List[Image.Image], Union[int, List[int]]], | |
Tuple['torch.Tensor[num_frames, channels, height, width]', Union[int, List[int]]], | |
Tuple[Any, Union[int, List[int]]], | |
]: | |
""" | |
For video with id idx, loads self.NUM_SEGMENTS * self.FRAMES_PER_SEGMENT | |
frames from evenly chosen locations across the video. | |
Args: | |
idx: Video sample index. | |
Returns: | |
A tuple of (video, label). Label is either a single | |
integer or a list of integers in the case of multiple labels. | |
Video is either 1) a list of PIL images if no transform is used | |
2) a batch of shape (NUM_IMAGES x CHANNELS x HEIGHT x WIDTH) in the range [0,1] | |
if the transform "ImglistToTensor" is used | |
3) or anything else if a custom transform is used. | |
""" | |
record: VideoRecord = self.video_list[idx] | |
frame_start_indices: 'np.ndarray[int]' = self._get_start_indices(record) | |
return self._get(record, frame_start_indices) | |
def _get(self, record: VideoRecord, frame_start_indices: 'np.ndarray[int]') -> Union[ | |
Tuple[List[Image.Image], Union[int, List[int]]], | |
Tuple['torch.Tensor[num_frames, channels, height, width]', Union[int, List[int]]], | |
Tuple[Any, Union[int, List[int]]], | |
]: | |
""" | |
Loads the frames of a video at the corresponding | |
indices. | |
Args: | |
record: VideoRecord denoting a video sample. | |
frame_start_indices: Indices from which to load consecutive frames from. | |
Returns: | |
A tuple of (video, label). Label is either a single | |
integer or a list of integers in the case of multiple labels. | |
Video is either 1) a list of PIL images if no transform is used | |
2) a batch of shape (NUM_IMAGES x CHANNELS x HEIGHT x WIDTH) in the range [0,1] | |
if the transform "ImglistToTensor" is used | |
3) or anything else if a custom transform is used. | |
""" | |
frame_start_indices = frame_start_indices + record.start_frame | |
images = list() | |
# from each start_index, load self.frames_per_segment | |
# consecutive frames | |
for start_index in frame_start_indices: | |
frame_index = int(start_index) | |
# load self.frames_per_segment consecutive frames | |
for _ in range(self.frames_per_segment): | |
image = self._load_image(record.path, frame_index) | |
images.append(image) | |
if frame_index < record.end_frame: | |
frame_index += 1 | |
if self.transform is not None: | |
images = self.transform(images) | |
return images, record.label | |
def __len__(self): | |
return len(self.video_list) | |
class ImglistToTensor(torch.nn.Module): | |
""" | |
Converts a list of PIL images in the range [0,255] to a torch.FloatTensor | |
of shape (NUM_IMAGES x CHANNELS x HEIGHT x WIDTH) in the range [0,1]. | |
Can be used as first transform for ``VideoFrameDataset``. | |
""" | |
@staticmethod | |
def forward(img_list: List[Image.Image]) -> 'torch.Tensor[NUM_IMAGES, CHANNELS, HEIGHT, WIDTH]': | |
""" | |
Converts each PIL image in a list to | |
a torch Tensor and stacks them into | |
a single tensor. | |
Args: | |
img_list: list of PIL images. | |
Returns: | |
tensor of size ``NUM_IMAGES x CHANNELS x HEIGHT x WIDTH`` | |
""" | |
return torch.stack([transforms.functional.to_tensor(pic) for pic in img_list]) |