init commit of samurai

lib/test/evaluation/tc128dataset.py

@@ -0,0 +1,44 @@
+import numpy as np
+from lib.test.evaluation.data import Sequence, BaseDataset, SequenceList
+import os
+import glob
+import six
+
+
+class TC128Dataset(BaseDataset):
+    """
+    TC-128 Dataset
+    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')))
+        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]
+        # valid frame range for each sequence
+        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)
+        # load valid frame range
+        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)]
+
+        # load annotations
+        anno = np.loadtxt(self.anno_files[index], delimiter=',')
+        assert len(img_files) == len(anno)
+        assert anno.shape[1] == 4
+
+        # return img_files, anno
+        return Sequence(sequence_name, img_files, 'tc128', anno.reshape(-1, 4))
+
+    def __len__(self):
+        return len(self.seq_names)