| import os | |
| import numpy as np | |
| from .base_dataloader import _BaseVideoLoader | |
| from ..factory import DATALOADERS | |
| class TrackingVideoLoader(_BaseVideoLoader): | |
| def __init__(self, **kwargs): | |
| super().__init__(**kwargs) | |
| self._first_frames = dict() | |
| for filename in self._files: | |
| stream = self._streams[filename] | |
| self._first_frames[filename] = next(stream) | |
| self._rois = self._load_roi() | |
| def _load_roi(self): | |
| rois = dict.fromkeys(self._files, None) | |
| for filename in self._files: | |
| rois[filename] = np.loadtxt(os.path.join(self._path, '{}.txt'.format(filename[:-4])), dtype=np.int32, ndmin=2) | |
| return rois | |
| def __getitem__(self, idx): | |
| filename = self._files[idx] | |
| return filename, self._streams[filename], self._first_frames[filename], self._rois[filename] |