Source code for torchvideo.datasets.video_dataset

from pathlib import Path
from typing import Union, Optional, Tuple, List, Any  # noqa

import torch.utils.data

from torchvideo.samplers import FrameSampler, _default_sampler
from .label_sets import LabelSet
from .types import Label, Transform


[docs]class VideoDataset(torch.utils.data.Dataset): """Abstract base class that all ``VideoDatasets`` inherit from. If you are implementing your own ``VideoDataset``, you should inherit from this class. """ def __init__( self, root_path: Union[str, Path], label_set: Optional[LabelSet] = None, sampler: FrameSampler = _default_sampler(), transform: Optional[Transform] = None, ) -> None: """ Args: root_path: Path to dataset on disk. label_set: Optional label set for labelling examples. sampler: Optional sampler for drawing frames from each video. transform: Optional transform over the list of frames. """ self.root_path = Path(root_path) self.label_set = label_set self.sampler = sampler self.transform = transform self.labels = None # type: Optional[List[Any]] """The labels corresponding to the examples in the dataset. To get the label for example at index ``i`` you simple call ``dataset.labels[i]``, although this will be returned by ``__getitem__`` if this field is not None.""" """The unique ID of each video (usually a path is possible)""" @property def video_ids(self): raise NotImplementedError()
[docs] def __getitem__( self, index: int ) -> Union[torch.Tensor, Tuple[torch.Tensor, Label]]: # pragma: no cover """Load an example by index Args: index: index of the example within the dataset. Returns: Example transformed by ``transform`` if one was passed during instantiation, otherwise the example is converted to a tensor without any transformations applied to it. Additionally, if a label set is present, the method return a tuple: ``(video_tensor, label)`` """ raise NotImplementedError()
[docs] def __len__(self) -> int: # pragma: no cover """Total number of examples in the dataset""" raise NotImplementedError()