| import numpy as np |
| from lib.test.evaluation.data import Sequence, BaseDataset, SequenceList |
| import os |
| import glob |
| import six |
|
|
|
|
| class TC128CEDataset(BaseDataset): |
| """ |
| TC-128 Dataset (78 newly added sequences) |
| modified from the implementation in got10k-toolkit (https://github.com/got-10k/toolkit) |
| """ |
| def __init__(self): |
| super().__init__() |
| self.base_path = self.env_settings.tc128_path |
| self.anno_files = sorted(glob.glob( |
| os.path.join(self.base_path, '*/*_gt.txt'))) |
| """filter the newly added sequences (_ce)""" |
| self.anno_files = [s for s in self.anno_files if "_ce" in s] |
| self.seq_dirs = [os.path.dirname(f) for f in self.anno_files] |
| self.seq_names = [os.path.basename(d) for d in self.seq_dirs] |
| |
| self.range_files = [glob.glob(os.path.join(d, '*_frames.txt'))[0] for d in self.seq_dirs] |
|
|
| def get_sequence_list(self): |
| return SequenceList([self._construct_sequence(s) for s in self.seq_names]) |
|
|
| def _construct_sequence(self, sequence_name): |
| if isinstance(sequence_name, six.string_types): |
| if not sequence_name in self.seq_names: |
| raise Exception('Sequence {} not found.'.format(sequence_name)) |
| index = self.seq_names.index(sequence_name) |
| |
| frames = np.loadtxt(self.range_files[index], dtype=int, delimiter=',') |
| img_files = [os.path.join(self.seq_dirs[index], 'img/%04d.jpg' % f) for f in range(frames[0], frames[1] + 1)] |
|
|
| |
| anno = np.loadtxt(self.anno_files[index], delimiter=',') |
| assert len(img_files) == len(anno) |
| assert anno.shape[1] == 4 |
|
|
| |
| return Sequence(sequence_name, img_files, 'tc128', anno.reshape(-1, 4)) |
|
|
| def __len__(self): |
| return len(self.seq_names) |
|
|
