Spaces:
Sleeping
Sleeping
File size: 2,777 Bytes
d4b77ac |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 |
# --------------------------------------------------------
# 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]
|