|
|
|
|
|
from copy import deepcopy
|
|
|
from typing import Any, Callable, List, Optional, Sequence, Union
|
|
|
|
|
|
import numpy as np
|
|
|
from mmengine.dataset import COLLATE_FUNCTIONS, pseudo_collate
|
|
|
|
|
|
from mmaction.registry import DATASETS
|
|
|
from mmaction.utils import ConfigType
|
|
|
from .video_dataset import VideoDataset
|
|
|
|
|
|
|
|
|
def get_type(transform: Union[dict, Callable]) -> str:
|
|
|
"""get the type of the transform."""
|
|
|
if isinstance(transform, dict) and 'type' in transform:
|
|
|
return transform['type']
|
|
|
elif callable(transform):
|
|
|
return transform.__repr__().split('(')[0]
|
|
|
else:
|
|
|
raise TypeError
|
|
|
|
|
|
|
|
|
@DATASETS.register_module()
|
|
|
class RepeatAugDataset(VideoDataset):
|
|
|
"""Video dataset for action recognition use repeat augment.
|
|
|
https://arxiv.org/pdf/1901.09335.pdf.
|
|
|
|
|
|
The dataset loads raw videos and apply specified transforms to return a
|
|
|
dict containing the frame tensors and other information.
|
|
|
|
|
|
The ann_file is a text file with multiple lines, and each line indicates
|
|
|
a sample video with the filepath and label, which are split with a
|
|
|
whitespace. Example of a annotation file:
|
|
|
|
|
|
.. code-block:: txt
|
|
|
|
|
|
some/path/000.mp4 1
|
|
|
some/path/001.mp4 1
|
|
|
some/path/002.mp4 2
|
|
|
some/path/003.mp4 2
|
|
|
some/path/004.mp4 3
|
|
|
some/path/005.mp4 3
|
|
|
|
|
|
|
|
|
Args:
|
|
|
ann_file (str): Path to the annotation file.
|
|
|
pipeline (List[Union[dict, ConfigDict, Callable]]): A sequence of
|
|
|
data transforms.
|
|
|
data_prefix (dict or ConfigDict): Path to a directory where videos
|
|
|
are held. Defaults to ``dict(video='')``.
|
|
|
num_repeats (int): Number of repeat time of one video in a batch.
|
|
|
Defaults to 4.
|
|
|
sample_once (bool): Determines whether use same frame index for
|
|
|
repeat samples. Defaults to False.
|
|
|
multi_class (bool): Determines whether the dataset is a multi-class
|
|
|
dataset. Defaults to False.
|
|
|
num_classes (int, optional): Number of classes of the dataset, used in
|
|
|
multi-class datasets. Defaults to None.
|
|
|
start_index (int): Specify a start index for frames in consideration of
|
|
|
different filename format. However, when taking videos as input,
|
|
|
it should be set to 0, since frames loaded from videos count
|
|
|
from 0. Defaults to 0.
|
|
|
modality (str): Modality of data. Support ``RGB``, ``Flow``.
|
|
|
Defaults to ``RGB``.
|
|
|
test_mode (bool): Store True when building test or validation dataset.
|
|
|
Defaults to False.
|
|
|
"""
|
|
|
|
|
|
def __init__(self,
|
|
|
ann_file: str,
|
|
|
pipeline: List[Union[dict, Callable]],
|
|
|
data_prefix: ConfigType = dict(video=''),
|
|
|
num_repeats: int = 4,
|
|
|
sample_once: bool = False,
|
|
|
multi_class: bool = False,
|
|
|
num_classes: Optional[int] = None,
|
|
|
start_index: int = 0,
|
|
|
modality: str = 'RGB',
|
|
|
**kwargs) -> None:
|
|
|
|
|
|
use_decord = get_type(pipeline[0]) == 'DecordInit' and \
|
|
|
get_type(pipeline[2]) == 'DecordDecode'
|
|
|
|
|
|
assert use_decord, (
|
|
|
'RepeatAugDataset requires decord as the video '
|
|
|
'loading backend, will support more backends in the '
|
|
|
'future')
|
|
|
|
|
|
super().__init__(
|
|
|
ann_file,
|
|
|
pipeline=pipeline,
|
|
|
data_prefix=data_prefix,
|
|
|
multi_class=multi_class,
|
|
|
num_classes=num_classes,
|
|
|
start_index=start_index,
|
|
|
modality=modality,
|
|
|
test_mode=False,
|
|
|
**kwargs)
|
|
|
self.num_repeats = num_repeats
|
|
|
self.sample_once = sample_once
|
|
|
|
|
|
def prepare_data(self, idx) -> List[dict]:
|
|
|
"""Get data processed by ``self.pipeline``.
|
|
|
|
|
|
Reduce the video loading and decompressing.
|
|
|
Args:
|
|
|
idx (int): The index of ``data_info``.
|
|
|
Returns:
|
|
|
List[dict]: A list of length num_repeats.
|
|
|
"""
|
|
|
transforms = self.pipeline.transforms
|
|
|
|
|
|
data_info = self.get_data_info(idx)
|
|
|
data_info = transforms[0](data_info)
|
|
|
|
|
|
frame_inds_list, frame_inds_length = [], [0]
|
|
|
|
|
|
fake_data_info = dict(
|
|
|
total_frames=data_info['total_frames'],
|
|
|
start_index=data_info['start_index'])
|
|
|
|
|
|
if not self.sample_once:
|
|
|
for repeat in range(self.num_repeats):
|
|
|
data_info_ = transforms[1](fake_data_info)
|
|
|
frame_inds = data_info_['frame_inds']
|
|
|
frame_inds_list.append(frame_inds.reshape(-1))
|
|
|
frame_inds_length.append(frame_inds.size +
|
|
|
frame_inds_length[-1])
|
|
|
else:
|
|
|
data_info_ = transforms[1](fake_data_info)
|
|
|
frame_inds = data_info_['frame_inds']
|
|
|
for repeat in range(self.num_repeats):
|
|
|
frame_inds_list.append(frame_inds.reshape(-1))
|
|
|
frame_inds_length.append(frame_inds.size +
|
|
|
frame_inds_length[-1])
|
|
|
|
|
|
for key in data_info_:
|
|
|
data_info[key] = data_info_[key]
|
|
|
|
|
|
data_info['frame_inds'] = np.concatenate(frame_inds_list)
|
|
|
|
|
|
data_info = transforms[2](data_info)
|
|
|
imgs = data_info.pop('imgs')
|
|
|
|
|
|
data_info_list = []
|
|
|
for repeat in range(self.num_repeats):
|
|
|
data_info_ = deepcopy(data_info)
|
|
|
start = frame_inds_length[repeat]
|
|
|
end = frame_inds_length[repeat + 1]
|
|
|
data_info_['imgs'] = imgs[start:end]
|
|
|
for transform in transforms[3:]:
|
|
|
data_info_ = transform(data_info_)
|
|
|
data_info_list.append(data_info_)
|
|
|
del imgs
|
|
|
return data_info_list
|
|
|
|
|
|
|
|
|
@COLLATE_FUNCTIONS.register_module()
|
|
|
def repeat_pseudo_collate(data_batch: Sequence) -> Any:
|
|
|
data_batch = [i for j in data_batch for i in j]
|
|
|
return pseudo_collate(data_batch)
|
|
|
|
