tracking.py

Go to the documentation of this file.

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]