|
|
import os |
|
|
import numpy as np |
|
|
|
|
|
from .base_dataloader import _BaseVideoLoader |
|
|
from ..factory import DATALOADERS |
|
|
|
|
|
@DATALOADERS.register |
|
|
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] |
