|
|
|
|
|
|
|
|
|
|
|
from os.path import join, realpath, dirname, exists, isdir |
|
from os import listdir |
|
import logging |
|
import glob |
|
import numpy as np |
|
import json |
|
from collections import OrderedDict |
|
|
|
|
|
def get_dataset_zoo(): |
|
root = realpath(join(dirname(__file__), '../data')) |
|
zoos = listdir(root) |
|
|
|
def valid(x): |
|
y = join(root, x) |
|
if not isdir(y): return False |
|
|
|
return exists(join(y, 'list.txt')) \ |
|
or exists(join(y, 'train', 'meta.json')) \ |
|
or exists(join(y, 'ImageSets', '2016', 'val.txt')) \ |
|
or exists(join(y, 'ImageSets', '2017', 'test-dev.txt')) |
|
|
|
zoos = list(filter(valid, zoos)) |
|
return zoos |
|
|
|
|
|
dataset_zoo = get_dataset_zoo() |
|
|
|
|
|
def load_dataset(dataset): |
|
info = OrderedDict() |
|
if 'VOT' in dataset: |
|
base_path = join(realpath(dirname(__file__)), '../data', dataset) |
|
if not exists(base_path): |
|
logging.error("Please download test dataset!!!") |
|
exit() |
|
list_path = join(base_path, 'list.txt') |
|
with open(list_path) as f: |
|
videos = [v.strip() for v in f.readlines()] |
|
for video in videos: |
|
video_path = join(base_path, video) |
|
image_path = join(video_path, '*.jpg') |
|
image_files = sorted(glob.glob(image_path)) |
|
if len(image_files) == 0: |
|
image_path = join(video_path, 'color', '*.jpg') |
|
image_files = sorted(glob.glob(image_path)) |
|
gt_path = join(video_path, 'groundtruth.txt') |
|
gt = np.loadtxt(gt_path, delimiter=',').astype(np.float64) |
|
if gt.shape[1] == 4: |
|
gt = np.column_stack((gt[:, 0], gt[:, 1], gt[:, 0], gt[:, 1] + gt[:, 3]-1, |
|
gt[:, 0] + gt[:, 2]-1, gt[:, 1] + gt[:, 3]-1, gt[:, 0] + gt[:, 2]-1, gt[:, 1])) |
|
info[video] = {'image_files': image_files, 'gt': gt, 'name': video} |
|
elif 'DAVIS' in dataset and 'TEST' not in dataset: |
|
base_path = join(realpath(dirname(__file__)), '../data', 'DAVIS') |
|
list_path = join(realpath(dirname(__file__)), '../data', 'DAVIS', 'ImageSets', dataset[-4:], 'val.txt') |
|
with open(list_path) as f: |
|
videos = [v.strip() for v in f.readlines()] |
|
for video in videos: |
|
info[video] = {} |
|
info[video]['anno_files'] = sorted(glob.glob(join(base_path, 'Annotations/480p', video, '*.png'))) |
|
info[video]['image_files'] = sorted(glob.glob(join(base_path, 'JPEGImages/480p', video, '*.jpg'))) |
|
info[video]['name'] = video |
|
elif 'ytb_vos' in dataset: |
|
base_path = join(realpath(dirname(__file__)), '../data', 'ytb_vos', 'valid') |
|
json_path = join(realpath(dirname(__file__)), '../data', 'ytb_vos', 'valid', 'meta.json') |
|
meta = json.load(open(json_path, 'r')) |
|
meta = meta['videos'] |
|
info = dict() |
|
for v in meta.keys(): |
|
objects = meta[v]['objects'] |
|
frames = [] |
|
anno_frames = [] |
|
info[v] = dict() |
|
for obj in objects: |
|
frames += objects[obj]['frames'] |
|
anno_frames += [objects[obj]['frames'][0]] |
|
frames = sorted(np.unique(frames)) |
|
info[v]['anno_files'] = [join(base_path, 'Annotations', v, im_f+'.png') for im_f in frames] |
|
info[v]['anno_init_files'] = [join(base_path, 'Annotations', v, im_f + '.png') for im_f in anno_frames] |
|
info[v]['image_files'] = [join(base_path, 'JPEGImages', v, im_f+'.jpg') for im_f in frames] |
|
info[v]['name'] = v |
|
|
|
info[v]['start_frame'] = dict() |
|
info[v]['end_frame'] = dict() |
|
for obj in objects: |
|
start_file = objects[obj]['frames'][0] |
|
end_file = objects[obj]['frames'][-1] |
|
info[v]['start_frame'][obj] = frames.index(start_file) |
|
info[v]['end_frame'][obj] = frames.index(end_file) |
|
elif 'TEST' in dataset: |
|
base_path = join(realpath(dirname(__file__)), '../data', 'DAVIS2017TEST') |
|
list_path = join(realpath(dirname(__file__)), '../data', 'DAVIS2017TEST', 'ImageSets', '2017', 'test-dev.txt') |
|
with open(list_path) as f: |
|
videos = [v.strip() for v in f.readlines()] |
|
for video in videos: |
|
info[video] = {} |
|
info[video]['anno_files'] = sorted(glob.glob(join(base_path, 'Annotations/480p', video, '*.png'))) |
|
info[video]['image_files'] = sorted(glob.glob(join(base_path, 'JPEGImages/480p', video, '*.jpg'))) |
|
info[video]['name'] = video |
|
else: |
|
logging.error('Not support') |
|
exit() |
|
return info |
|
|
