Function main

tutorials/ctracker/test.py:299–332

(args=None)

		run_each_dataset(model_dir, retinanet, root_path, 'test', 'MOT17-{:02d}'.format(seq_num))

def main(args=None):
	parser = argparse.ArgumentParser(description='Simple script for testing a CTracker network.')
	parser.add_argument('--dataset_path', default='/dockerdata/home/jeromepeng/data/MOT/MOT17/', type=str, help='Dataset path, location of the images sequence.')
	parser.add_argument('--model_dir', default='./trained_model/', help='Path to model (.pt) file.')
	parser.add_argument('--model_path', default='./trained_model/model_final.pth', help='Path to model (.pt) file.')
	parser = parser.parse_args(args)

	if not os.path.exists(os.path.join(parser.model_dir, 'results')):
		os.makedirs(os.path.join(parser.model_dir, 'results'))
    
	retinanet = model.resnet50(num_classes=1, pretrained=True)
# 	retinanet_save = torch.load(os.path.join(parser.model_dir, 'model_final.pth'))
	retinanet_save = torch.load(os.path.join(parser.model_path))
    
    # rename moco pre-trained keys
	state_dict = retinanet_save.state_dict()
	for k in list(state_dict.keys()):
        # retain only encoder up to before the embedding layer
		if k.startswith('module.'):
            # remove prefix
			state_dict[k[len("module."):]] = state_dict[k]
        # delete renamed or unused k
		del state_dict[k]
                
	retinanet.load_state_dict(state_dict)
    
	use_gpu = True

	if use_gpu: retinanet = retinanet.cuda()

	retinanet.eval()

	for seq_num in [2, 4, 5, 9, 10, 11, 13]:
		run_each_dataset(parser.model_dir, retinanet, parser.dataset_path, 'train', 'MOT17-{:02d}'.format(seq_num))
# 	for seq_num in [1, 3, 6, 7, 8, 12, 14]:
# 		run_each_dataset(parser.model_dir, retinanet, parser.dataset_path, 'test', 'MOT17-{:02d}'.format(seq_num))