def parameters(yaml_name: str):
params = TrackerParams()
prj_dir = env_settings().prj_dir
save_dir = env_settings().save_dir
# update default config from yaml file
yaml_file = os.path.join(prj_dir,
'experiments/stark_s/%s.yaml' % yaml_name)
update_config_from_file(yaml_file)
params.cfg = cfg
print("test config: ", cfg)
# template and search region
params.template_factor = cfg.TEST.TEMPLATE_FACTOR
params.template_size = cfg.TEST.TEMPLATE_SIZE
params.search_factor = cfg.TEST.SEARCH_FACTOR
params.search_size = cfg.TEST.SEARCH_SIZE
# Network checkpoint path
params.checkpoint = os.path.join(
save_dir, "checkpoints/train/stark_s/%s/STARKS_ep%04d.pth.tar" %
(yaml_name, cfg.TEST.EPOCH))
# whether to save boxes from all queries
params.save_all_boxes = False
return params
def __init__(self,
name: str,
parameter_name: str,
dataset_name: str,
run_id: int = None,
display_name: str = None,
result_only=False):
assert run_id is None or isinstance(run_id, int)
self.name = name
self.parameter_name = parameter_name
self.dataset_name = dataset_name
self.run_id = run_id
self.display_name = display_name
env = env_settings()
if self.run_id is None:
self.results_dir = '{}/{}/{}'.format(env.results_path, self.name,
self.parameter_name)
else:
self.results_dir = '{}/{}/{}_{:03d}'.format(
env.results_path, self.name, self.parameter_name, self.run_id)
if result_only:
self.results_dir = '{}/{}/{}'.format(env.results_path, "LaSOT",
self.name)
tracker_module_abspath = os.path.abspath(
os.path.join(os.path.dirname(__file__), '..', 'tracker',
'%s.py' % self.name))
if os.path.isfile(tracker_module_abspath):
tracker_module = importlib.import_module(
'lib.test.tracker.{}'.format(self.name))
self.tracker_class = tracker_module.get_tracker_class()
else:
self.tracker_class = None
def check_and_load_precomputed_results(trackers,
dataset,
report_name,
force_evaluation=False,
**kwargs):
# Load data
settings = env_settings()
# Load pre-computed results
result_plot_path = os.path.join(settings.result_plot_path, report_name)
eval_data_path = os.path.join(result_plot_path, 'eval_data.pkl')
if os.path.isfile(eval_data_path) and not force_evaluation:
with open(eval_data_path, 'rb') as fh:
eval_data = pickle.load(fh)
else:
# print('Pre-computed evaluation data not found. Computing results!')
eval_data = extract_results(trackers, dataset, report_name, **kwargs)
if not check_eval_data_is_valid(eval_data, trackers, dataset):
# print('Pre-computed evaluation data invalid. Re-computing results!')
eval_data = extract_results(trackers, dataset, report_name, **kwargs)
# pass
else:
# Update display names
tracker_names = [{
'name': t.name,
'param': t.parameter_name,
'run_id': t.run_id,
'disp_name': t.display_name
} for t in trackers]
eval_data['trackers'] = tracker_names
with open(eval_data_path, 'wb') as fh:
pickle.dump(eval_data, fh)
return eval_data
def unpack_tracking_results(download_path, output_path=None):
"""
Unpacks zipped benchmark results. The directory 'download_path' should have the following structure
- root
- tracker1
- param1.zip
- param2.zip
.
.
- tracker2
- param1.zip
- param2.zip
.
.
args:
download_path - Path to the directory where the zipped results are stored
output_path - Path to the directory where the results will be unpacked. Set to env_settings().results_path
by default
"""
if output_path is None:
output_path = env_settings().results_path
if not os.path.exists(output_path):
os.makedirs(output_path)
trackers = os.listdir(download_path)
for t in trackers:
runfiles = os.listdir(os.path.join(download_path, t))
for r in runfiles:
save_path = os.path.join(output_path, t)
if not os.path.exists(save_path):
os.makedirs(save_path)
shutil.unpack_archive(os.path.join(download_path, t, r),
os.path.join(save_path, r[:-4]), 'zip')
def transform_got10k(tracker_name, cfg_name):
env = env_settings()
result_dir = env.results_path
src_dir = os.path.join(result_dir,
"%s/%s/got10k/" % (tracker_name, cfg_name))
dest_dir = os.path.join(result_dir,
"%s/%s/got10k_submit/" % (tracker_name, cfg_name))
if not os.path.exists(dest_dir):
os.makedirs(dest_dir)
items = os.listdir(src_dir)
for item in items:
if "all" in item:
continue
src_path = os.path.join(src_dir, item)
if "time" not in item:
seq_name = item.replace(".txt", '')
seq_dir = os.path.join(dest_dir, seq_name)
if not os.path.exists(seq_dir):
os.makedirs(seq_dir)
new_item = item.replace(".txt", '_001.txt')
dest_path = os.path.join(seq_dir, new_item)
bbox_arr = np.loadtxt(src_path, dtype=np.int, delimiter='\t')
np.savetxt(dest_path, bbox_arr, fmt='%d', delimiter=',')
else:
seq_name = item.replace("_time.txt", '')
seq_dir = os.path.join(dest_dir, seq_name)
if not os.path.exists(seq_dir):
os.makedirs(seq_dir)
dest_path = os.path.join(seq_dir, item)
os.system("cp %s %s" % (src_path, dest_path))
# make zip archive
shutil.make_archive(src_dir, "zip", src_dir)
shutil.make_archive(dest_dir, "zip", dest_dir)
# Remove the original files
shutil.rmtree(src_dir)
shutil.rmtree(dest_dir)
def transform_trackingnet(tracker_name, cfg_name):
env = env_settings()
result_dir = env.results_path
src_dir = os.path.join(result_dir,
"%s/%s/trackingnet/" % (tracker_name, cfg_name))
dest_dir = os.path.join(
result_dir, "%s/%s/trackingnet_submit/" % (tracker_name, cfg_name))
if not os.path.exists(dest_dir):
os.makedirs(dest_dir)
items = os.listdir(src_dir)
for item in items:
if "all" in item:
continue
if "time" not in item:
src_path = os.path.join(src_dir, item)
dest_path = os.path.join(dest_dir, item)
bbox_arr = np.loadtxt(src_path, dtype=np.int, delimiter='\t')
np.savetxt(dest_path, bbox_arr, fmt='%d', delimiter=',')
# make zip archive
shutil.make_archive(src_dir, "zip", src_dir)
shutil.make_archive(dest_dir, "zip", dest_dir)
# Remove the original files
shutil.rmtree(src_dir)
shutil.rmtree(dest_dir)
def plot_got_success(trackers, report_name):
""" Plot success plot for GOT-10k dataset using the json reports.
Save the json reports from http://got-10k.aitestunion.com/leaderboard in the directory set to
env_settings.got_reports_path
The tracker name in the experiment file should be set to the name of the report file for that tracker,
e.g. DiMP50_report_2019_09_02_15_44_25 if the report is name DiMP50_report_2019_09_02_15_44_25.json
args:
trackers - List of trackers to evaluate
report_name - Name of the folder in env_settings.perm_mat_path where the computed results and plots are saved
"""
# Load data
settings = env_settings()
plot_draw_styles = get_plot_draw_styles()
result_plot_path = os.path.join(settings.result_plot_path, report_name)
auc_curve = torch.zeros((len(trackers), 101))
scores = torch.zeros(len(trackers))
# Load results
tracker_names = []
for trk_id, trk in enumerate(trackers):
json_path = '{}/{}.json'.format(settings.got_reports_path, trk.name)
if os.path.isfile(json_path):
with open(json_path, 'r') as f:
eval_data = json.load(f)
else:
raise Exception('Report not found {}'.format(json_path))
if len(eval_data.keys()) > 1:
raise Exception
# First field is the tracker name. Index it out
eval_data = eval_data[list(eval_data.keys())[0]]
if 'succ_curve' in eval_data.keys():
curve = eval_data['succ_curve']
ao = eval_data['ao']
elif 'overall' in eval_data.keys(
) and 'succ_curve' in eval_data['overall'].keys():
curve = eval_data['overall']['succ_curve']
ao = eval_data['overall']['ao']
else:
raise Exception('Invalid JSON file {}'.format(json_path))
auc_curve[trk_id, :] = torch.tensor(curve) * 100.0
scores[trk_id] = ao * 100.0
tracker_names.append({
'name': trk.name,
'param': trk.parameter_name,
'run_id': trk.run_id,
'disp_name': trk.display_name
})
threshold_set_overlap = torch.arange(0.0, 1.01, 0.01, dtype=torch.float64)
success_plot_opts = {
'plot_type': 'success',
'legend_loc': 'lower left',
'xlabel': 'Overlap threshold',
'ylabel': 'Overlap Precision [%]',
'xlim': (0, 1.0),
'ylim': (0, 100),
'title': 'Success plot'
}
plot_draw_save(auc_curve, threshold_set_overlap, scores, tracker_names,
plot_draw_styles, result_plot_path, success_plot_opts)
plt.show()
def plot_results(trackers,
dataset,
report_name,
merge_results=False,
plot_types=('success'),
force_evaluation=False,
**kwargs):
"""
Plot results for the given trackers
args:
trackers - List of trackers to evaluate
dataset - List of sequences to evaluate
report_name - Name of the folder in env_settings.perm_mat_path where the computed results and plots are saved
merge_results - If True, multiple random runs for a non-deterministic trackers are averaged
plot_types - List of scores to display. Can contain 'success',
'prec' (precision), and 'norm_prec' (normalized precision)
"""
# Load data
settings = env_settings()
plot_draw_styles = get_plot_draw_styles()
# Load pre-computed results
result_plot_path = os.path.join(settings.result_plot_path, report_name)
eval_data = check_and_load_precomputed_results(trackers, dataset,
report_name,
force_evaluation, **kwargs)
# Merge results from multiple runs
if merge_results:
eval_data = merge_multiple_runs(eval_data)
tracker_names = eval_data['trackers']
valid_sequence = torch.tensor(eval_data['valid_sequence'],
dtype=torch.bool)
print('\nPlotting results over {} / {} sequences'.format(
valid_sequence.long().sum().item(), valid_sequence.shape[0]))
print('\nGenerating plots for: {}'.format(report_name))
# ******************************** Success Plot **************************************
if 'success' in plot_types:
ave_success_rate_plot_overlap = torch.tensor(
eval_data['ave_success_rate_plot_overlap'])
# Index out valid sequences
auc_curve, auc = get_auc_curve(ave_success_rate_plot_overlap,
valid_sequence)
threshold_set_overlap = torch.tensor(
eval_data['threshold_set_overlap'])
success_plot_opts = {
'plot_type': 'success',
'legend_loc': 'lower left',
'xlabel': 'Overlap threshold',
'ylabel': 'Overlap Precision [%]',
'xlim': (0, 1.0),
'ylim': (0, 88),
'title': 'Success'
}
plot_draw_save(auc_curve, threshold_set_overlap, auc, tracker_names,
plot_draw_styles, result_plot_path, success_plot_opts)
# ******************************** Precision Plot **************************************
if 'prec' in plot_types:
ave_success_rate_plot_center = torch.tensor(
eval_data['ave_success_rate_plot_center'])
# Index out valid sequences
prec_curve, prec_score = get_prec_curve(ave_success_rate_plot_center,
valid_sequence)
threshold_set_center = torch.tensor(eval_data['threshold_set_center'])
precision_plot_opts = {
'plot_type': 'precision',
'legend_loc': 'lower right',
'xlabel': 'Location error threshold [pixels]',
'ylabel': 'Distance Precision [%]',
'xlim': (0, 50),
'ylim': (0, 100),
'title': 'Precision plot'
}
plot_draw_save(prec_curve, threshold_set_center, prec_score,
tracker_names, plot_draw_styles, result_plot_path,
precision_plot_opts)
# ******************************** Norm Precision Plot **************************************
if 'norm_prec' in plot_types:
ave_success_rate_plot_center_norm = torch.tensor(
eval_data['ave_success_rate_plot_center_norm'])
# Index out valid sequences
prec_curve, prec_score = get_prec_curve(
ave_success_rate_plot_center_norm, valid_sequence)
threshold_set_center_norm = torch.tensor(
eval_data['threshold_set_center_norm'])
norm_precision_plot_opts = {
'plot_type': 'norm_precision',
'legend_loc': 'lower right',
'xlabel': 'Location error threshold',
'ylabel': 'Distance Precision [%]',
'xlim': (0, 0.5),
'ylim': (0, 85),
'title': 'Normalized Precision'
}
plot_draw_save(prec_curve, threshold_set_center_norm, prec_score,
tracker_names, plot_draw_styles, result_plot_path,
norm_precision_plot_opts)
plt.show()
if __name__ == "__main__":
load_checkpoint = True
save_name = "backbone_bottleneck_pe.onnx"
"""update cfg"""
args = parse_args()
yaml_fname = 'experiments/%s/%s.yaml' % (args.script, args.config)
update_config_from_file(yaml_fname)
'''set some values'''
bs = 1
z_sz = cfg.TEST.TEMPLATE_SIZE
# build the stark model
model = build_stark_lightning_x_trt(cfg, phase='test')
# load checkpoint
if load_checkpoint:
save_dir = env_settings().save_dir
checkpoint_name = os.path.join(
save_dir,
"checkpoints/train/%s/%s/STARKLightningXtrt_ep0500.pth.tar" %
(args.script, args.config))
model.load_state_dict(torch.load(checkpoint_name,
map_location='cpu')['net'],
strict=True)
# transfer to test mode
model = repvgg_model_convert(model)
model.eval()
""" rebuild the inference-time model """
backbone = model.backbone
bottleneck = model.bottleneck
position_embed = model.pos_emb_z0
torch_model = Backbone_Bottleneck_PE(backbone, bottleneck, position_embed)
def __init__(self):
self.env_settings = env_settings()