Spaces:
Sleeping
Sleeping
| # -------------------------------------------------------- | |
| # Python Single Object Tracking Evaluation | |
| # Licensed under The MIT License [see LICENSE for details] | |
| # Written by Fangyi Zhang | |
| # @author fangyi.zhang@vipl.ict.ac.cn | |
| # @project https://github.com/StrangerZhang/pysot-toolkit.git | |
| # Revised for SiamMask by foolwood | |
| # -------------------------------------------------------- | |
| import os | |
| import cv2 | |
| from glob import glob | |
| class Video(object): | |
| def __init__(self, name, root, video_dir, init_rect, img_names, | |
| gt_rect, attr): | |
| self.name = name | |
| self.video_dir = video_dir | |
| self.init_rect = init_rect | |
| self.gt_traj = gt_rect | |
| self.attr = attr | |
| self.pred_trajs = {} | |
| self.img_names = [os.path.join(root, x) for x in img_names] | |
| self.imgs = None | |
| def load_tracker(self, path, tracker_names=None, store=True): | |
| """ | |
| Args: | |
| path(str): path to result | |
| tracker_name(list): name of tracker | |
| """ | |
| if not tracker_names: | |
| tracker_names = [x.split('/')[-1] for x in glob(path) | |
| if os.path.isdir(x)] | |
| if isinstance(tracker_names, str): | |
| tracker_names = [tracker_names] | |
| for name in tracker_names: | |
| traj_file = os.path.join(path, name, self.name+'.txt') | |
| if os.path.exists(traj_file): | |
| with open(traj_file, 'r') as f : | |
| pred_traj = [list(map(float, x.strip().split(','))) | |
| for x in f.readlines()] | |
| if len(pred_traj) != len(self.gt_traj): | |
| print(name, len(pred_traj), len(self.gt_traj), self.name) | |
| if store: | |
| self.pred_trajs[name] = pred_traj | |
| else: | |
| return pred_traj | |
| else: | |
| print(traj_file) | |
| self.tracker_names = list(self.pred_trajs.keys()) | |
| def load_img(self): | |
| if self.imgs is None: | |
| self.imgs = [cv2.imread(x) | |
| for x in self.img_names] | |
| self.width = self.imgs[0].shape[1] | |
| self.height = self.imgs[0].shape[0] | |
| def free_img(self): | |
| self.imgs = None | |
| def __len__(self): | |
| return len(self.img_names) | |
| def __getitem__(self, idx): | |
| if self.imgs is None: | |
| return cv2.imread(self.img_names[idx]), \ | |
| self.gt_traj[idx] | |
| else: | |
| return self.imgs[idx], self.gt_traj[idx] | |
| def __iter__(self): | |
| for i in range(len(self.img_names)): | |
| if self.imgs is not None: | |
| yield self.imgs[i], self.gt_traj[i] | |
| else: | |
| yield cv2.imread(self.img_names[i]), \ | |
| self.gt_traj[i] | |