| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
|
|
| from abc import ABC, abstractmethod |
| from typing import Any, List, Optional |
|
|
| from src.models.data.datafield import DataField |
|
|
|
|
| class BaseDataset(ABC): |
| """ |
| Base class for all datasets. |
| Note that this is not directly wrap-able by a dataloader. It is meant to be |
| subclassed / included in another dataset class. |
| """ |
|
|
| def __init__(self): |
| pass |
|
|
| @abstractmethod |
| def available_data_fields(self) -> List[DataField]: |
| """ |
| Return a list of available data fields in the dataset. |
| """ |
| pass |
|
|
| @abstractmethod |
| def num_videos(self) -> int: |
| """ |
| Returns: |
| Number of videos in the dataset. |
| """ |
| pass |
|
|
| @abstractmethod |
| def num_views(self, video_idx: int) -> int: |
| """ |
| Args: |
| video_idx: Index of the video. |
| |
| Returns: |
| Number of views in the video. |
| """ |
| pass |
|
|
| @abstractmethod |
| def num_frames(self, video_idx: int, view_idx: int = 0) -> int: |
| """ |
| Args: |
| video_idx: Index of the video. |
| view_idx: Index of the view. |
| |
| Returns: |
| Number of frames in the given view. |
| """ |
| pass |
|
|
| def read_video_metadata(self, video_idx: int) -> dict[str, Any]: |
| """ |
| Read metadata of the video. |
| |
| Args: |
| video_idx: Index of the video. |
| |
| Returns: |
| A dictionary containing metadata of the video. |
| """ |
| return {} |
|
|
| def read_view_metadata(self, video_idx: int, view_idx: int) -> dict[str, Any]: |
| """ |
| Read metadata of the view. |
| |
| Args: |
| video_idx: Index of the video. |
| view_idx: Index of the view. |
| |
| Returns: |
| A dictionary containing metadata of the view. |
| """ |
| return {} |
|
|
| def read( |
| self, |
| video_idx: int, |
| frame_idxs: Optional[List[int]] = None, |
| view_idxs: Optional[List[int]] = None, |
| data_fields: Optional[List[DataField]] = None, |
| ) -> dict[DataField, Any]: |
| """ |
| Read data from the dataset. |
| Args: |
| video_idx: Index of the video. |
| frame_idxs: List of frame indices. |
| view_idxs: List of view indices. |
| data_fields: List of data fields to read. If None, read all data fields. |
| |
| Example: |
| if frame_idxs is None, view_idxs is None, read all frames from all views. |
| if frame_idxs is not None, view_idxs is None, read frames from the first view. |
| if frame_idxs is None, view_idxs is not None, read all frames from the specified views. |
| if frame_idxs is not None, view_idxs is not None, read frames from the specified views. |
| |
| Returns: |
| A dictionary mapping data fields to their values. |
| """ |
|
|
| if data_fields is None: |
| data_fields = self.available_data_fields() |
|
|
| if frame_idxs is None: |
| |
| if view_idxs is None: |
| view_iterator = range(self.num_views(video_idx)) |
| else: |
| view_iterator = view_idxs |
|
|
| new_frame_idxs, new_view_idxs = [], [] |
| for view_idx in view_iterator: |
| num_frames = self.num_frames(video_idx, view_idx) |
| new_frame_idxs.extend(list(range(num_frames))) |
| new_view_idxs.extend([view_idx] * num_frames) |
| frame_idxs, view_idxs = new_frame_idxs, new_view_idxs |
|
|
| elif view_idxs is None: |
| |
| view_idxs = [0] * len(frame_idxs) |
|
|
| else: |
| |
| assert len(frame_idxs) == len(view_idxs), ( |
| "Frame and view indices must match." |
| ) |
|
|
| return self._read_data( |
| video_idx=video_idx, |
| frame_idxs=frame_idxs, |
| view_idxs=view_idxs, |
| data_fields=data_fields, |
| ) |
|
|
| @abstractmethod |
| def _read_data( |
| self, |
| video_idx: int, |
| frame_idxs: List[int], |
| view_idxs: List[int], |
| data_fields: List[DataField], |
| ) -> dict[DataField, Any]: |
| pass |
|
|
