|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
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] |
|
|
