148 lines
5.7 KiB
Python
148 lines
5.7 KiB
Python
import torch
|
|
import os
|
|
import os.path
|
|
import numpy as np
|
|
import random
|
|
from collections import OrderedDict
|
|
|
|
from lib.train.data import jpeg4py_loader
|
|
from .base_video_dataset import BaseVideoDataset
|
|
from lib.train.admin import env_settings
|
|
import json
|
|
from lib.utils.lmdb_utils import decode_img, decode_str
|
|
|
|
|
|
def list_sequences(root):
|
|
""" Lists all the videos in the input set_ids. Returns a list of tuples (set_id, video_name)
|
|
|
|
args:
|
|
root: Root directory to TrackingNet
|
|
|
|
returns:
|
|
list - list of tuples (set_id, video_name) containing the set_id and video_name for each sequence
|
|
"""
|
|
fname = os.path.join(root, "seq_list.json")
|
|
with open(fname, "r") as f:
|
|
sequence_list = json.loads(f.read())
|
|
return sequence_list
|
|
|
|
|
|
class TrackingNet_lmdb(BaseVideoDataset):
|
|
""" TrackingNet dataset.
|
|
|
|
Publication:
|
|
TrackingNet: A Large-Scale Dataset and Benchmark for Object Tracking in the Wild.
|
|
Matthias Mueller,Adel Bibi, Silvio Giancola, Salman Al-Subaihi and Bernard Ghanem
|
|
ECCV, 2018
|
|
https://ivul.kaust.edu.sa/Documents/Publications/2018/TrackingNet%20A%20Large%20Scale%20Dataset%20and%20Benchmark%20for%20Object%20Tracking%20in%20the%20Wild.pdf
|
|
|
|
Download the dataset using the toolkit https://github.com/SilvioGiancola/TrackingNet-devkit.
|
|
"""
|
|
def __init__(self, root=None, image_loader=jpeg4py_loader, set_ids=None, data_fraction=None):
|
|
"""
|
|
args:
|
|
root - The path to the TrackingNet folder, containing the training sets.
|
|
image_loader (jpeg4py_loader) - The function to read the images. jpeg4py (https://github.com/ajkxyz/jpeg4py)
|
|
is used by default.
|
|
set_ids (None) - List containing the ids of the TrackingNet sets to be used for training. If None, all the
|
|
sets (0 - 11) will be used.
|
|
data_fraction - Fraction of dataset to be used. The complete dataset is used by default
|
|
"""
|
|
root = env_settings().trackingnet_lmdb_dir if root is None else root
|
|
super().__init__('TrackingNet_lmdb', root, image_loader)
|
|
|
|
if set_ids is None:
|
|
set_ids = [i for i in range(12)]
|
|
|
|
self.set_ids = set_ids
|
|
|
|
# Keep a list of all videos. Sequence list is a list of tuples (set_id, video_name) containing the set_id and
|
|
# video_name for each sequence
|
|
self.sequence_list = list_sequences(self.root)
|
|
|
|
if data_fraction is not None:
|
|
self.sequence_list = random.sample(self.sequence_list, int(len(self.sequence_list) * data_fraction))
|
|
|
|
self.seq_to_class_map, self.seq_per_class = self._load_class_info()
|
|
|
|
# we do not have the class_lists for the tracking net
|
|
self.class_list = list(self.seq_per_class.keys())
|
|
self.class_list.sort()
|
|
|
|
def _load_class_info(self):
|
|
ltr_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), '..')
|
|
class_map_path = os.path.join(ltr_path, 'data_specs', 'trackingnet_classmap.txt')
|
|
|
|
with open(class_map_path, 'r') as f:
|
|
seq_to_class_map = {seq_class.split('\t')[0]: seq_class.rstrip().split('\t')[1] for seq_class in f}
|
|
|
|
seq_per_class = {}
|
|
for i, seq in enumerate(self.sequence_list):
|
|
class_name = seq_to_class_map.get(seq[1], 'Unknown')
|
|
if class_name not in seq_per_class:
|
|
seq_per_class[class_name] = [i]
|
|
else:
|
|
seq_per_class[class_name].append(i)
|
|
|
|
return seq_to_class_map, seq_per_class
|
|
|
|
def get_name(self):
|
|
return 'trackingnet_lmdb'
|
|
|
|
def has_class_info(self):
|
|
return True
|
|
|
|
def get_sequences_in_class(self, class_name):
|
|
return self.seq_per_class[class_name]
|
|
|
|
def _read_bb_anno(self, seq_id):
|
|
set_id = self.sequence_list[seq_id][0]
|
|
vid_name = self.sequence_list[seq_id][1]
|
|
gt_str_list = decode_str(os.path.join(self.root, "TRAIN_%d_lmdb" % set_id),
|
|
os.path.join("anno", vid_name + ".txt")).split('\n')[:-1]
|
|
gt_list = [list(map(float, line.split(','))) for line in gt_str_list]
|
|
gt_arr = np.array(gt_list).astype(np.float32)
|
|
return torch.tensor(gt_arr)
|
|
|
|
def get_sequence_info(self, seq_id):
|
|
bbox = self._read_bb_anno(seq_id)
|
|
|
|
valid = (bbox[:, 2] > 0) & (bbox[:, 3] > 0)
|
|
visible = valid.clone().byte()
|
|
return {'bbox': bbox, 'valid': valid, 'visible': visible}
|
|
|
|
def _get_frame(self, seq_id, frame_id):
|
|
set_id = self.sequence_list[seq_id][0]
|
|
vid_name = self.sequence_list[seq_id][1]
|
|
return decode_img(os.path.join(self.root, "TRAIN_%d_lmdb" % set_id),
|
|
os.path.join("frames", vid_name, str(frame_id) + ".jpg"))
|
|
|
|
def _get_class(self, seq_id):
|
|
seq_name = self.sequence_list[seq_id][1]
|
|
return self.seq_to_class_map[seq_name]
|
|
|
|
def get_class_name(self, seq_id):
|
|
obj_class = self._get_class(seq_id)
|
|
|
|
return obj_class
|
|
|
|
def get_frames(self, seq_id, frame_ids, anno=None):
|
|
frame_list = [self._get_frame(seq_id, f) for f in frame_ids]
|
|
|
|
if anno is None:
|
|
anno = self.get_sequence_info(seq_id)
|
|
|
|
anno_frames = {}
|
|
for key, value in anno.items():
|
|
anno_frames[key] = [value[f_id, ...].clone() for f_id in frame_ids]
|
|
|
|
obj_class = self._get_class(seq_id)
|
|
|
|
object_meta = OrderedDict({'object_class_name': obj_class,
|
|
'motion_class': None,
|
|
'major_class': None,
|
|
'root_class': None,
|
|
'motion_adverb': None})
|
|
|
|
return frame_list, anno_frames, object_meta
|