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 json | |
| import numpy as np | |
| from glob import glob | |
| from tqdm import tqdm | |
| from .dataset import Dataset | |
| from .video import Video | |
| class VOTVideo(Video): | |
| """ | |
| Args: | |
| name: video name | |
| root: dataset root | |
| video_dir: video directory | |
| init_rect: init rectangle | |
| img_names: image names | |
| gt_rect: groundtruth rectangle | |
| camera_motion: camera motion tag | |
| illum_change: illum change tag | |
| motion_change: motion change tag | |
| size_change: size change | |
| occlusion: occlusion | |
| """ | |
| def __init__(self, name, root, video_dir, init_rect, img_names, gt_rect, | |
| camera_motion, illum_change, motion_change, size_change, occlusion, width, height): | |
| super(VOTVideo, self).__init__(name, root, video_dir, init_rect, img_names, gt_rect, None) | |
| self.tags= {'all': [1] * len(gt_rect)} | |
| self.tags['camera_motion'] = camera_motion | |
| self.tags['illum_change'] = illum_change | |
| self.tags['motion_change'] = motion_change | |
| self.tags['size_change'] = size_change | |
| self.tags['occlusion'] = occlusion | |
| self.width = width | |
| self.height = height | |
| # empty tag | |
| all_tag = [v for k, v in self.tags.items() if len(v) > 0 ] | |
| self.tags['empty'] = np.all(1 - np.array(all_tag), axis=1).astype(np.int32).tolist() | |
| self.tag_names = list(self.tags.keys()) | |
| def select_tag(self, tag, start=0, end=0): | |
| if tag == 'empty': | |
| return self.tags[tag] | |
| return self.tags[tag][start:end] | |
| 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_files = glob(os.path.join(path, name, 'baseline', self.name, '*0*.txt')) | |
| if len(traj_files) == 15: | |
| traj_files = traj_files | |
| else: | |
| traj_files = traj_files[0:1] | |
| pred_traj = [] | |
| for traj_file in traj_files: | |
| with open(traj_file, 'r') as f: | |
| traj = [list(map(float, x.strip().split(','))) | |
| for x in f.readlines()] | |
| pred_traj.append(traj) | |
| if store: | |
| self.pred_trajs[name] = pred_traj | |
| else: | |
| return pred_traj | |
| class VOTDataset(Dataset): | |
| """ | |
| Args: | |
| name: dataset name, should be 'VOT2018', 'VOT2016' | |
| dataset_root: dataset root | |
| load_img: wether to load all imgs | |
| """ | |
| def __init__(self, name, dataset_root): | |
| super(VOTDataset, self).__init__(name, dataset_root) | |
| try: | |
| with open(os.path.join(dataset_root, name+'.json'), 'r') as f: | |
| meta_data = json.load(f) | |
| except: | |
| download_str = '# download json file for eval toolkit\n'+\ | |
| 'cd $SiamMask/data\n'+\ | |
| 'wget http://www.robots.ox.ac.uk/~qwang/VOT2016.json\n'+\ | |
| 'wget http://www.robots.ox.ac.uk/~qwang/VOT2018.json' | |
| print(download_str) | |
| exit() | |
| # load videos | |
| pbar = tqdm(meta_data.keys(), desc='loading '+name, ncols=100) | |
| self.videos = {} | |
| for video in pbar: | |
| pbar.set_postfix_str(video) | |
| self.videos[video] = VOTVideo(video, | |
| dataset_root, | |
| meta_data[video]['video_dir'], | |
| meta_data[video]['init_rect'], | |
| meta_data[video]['img_names'], | |
| meta_data[video]['gt_rect'], | |
| meta_data[video]['camera_motion'], | |
| meta_data[video]['illum_change'], | |
| meta_data[video]['motion_change'], | |
| meta_data[video]['size_change'], | |
| meta_data[video]['occlusion'], | |
| meta_data[video]['width'], | |
| meta_data[video]['height']) | |
| self.tags = ['all', 'camera_motion', 'illum_change', 'motion_change', | |
| 'size_change', 'occlusion', 'empty'] | |