class Tracktor:
def __init__(self, tracktor_config):
self.tracktor_config = tracktor_config
with open(tracktor_config, 'r') as f:
self.tracktor = yaml.load(f)['tracktor']
self.reid = self.tracktor['reid']
# Set up seed
torch.manual_seed(self.tracktor['seed'])
torch.cuda.manual_seed(self.tracktor['seed'])
np.random.seed(self.tracktor['seed'])
torch.backends.cudnn.deterministic = True
# Output directory
self.output_dir = osp.join(get_output_dir(self.tracktor['module_name']), self.tracktor['name'])
if not osp.exists(self.output_dir):
os.makedirs(self.output_dir)
# object detection
self.obj_detect = FRCNN_FPN(num_classes=2)
self.obj_detect.load_state_dict(torch.load(self.tracktor['obj_detect_model'],
map_location=lambda storage, loc: storage))
self.obj_detect.eval()
self.obj_detect.cuda()
# reid
self.reid_network = resnet50(pretrained=False, **self.reid['cnn'])
self.reid_network.load_state_dict(torch.load(self.tracktor['reid_weights'],
map_location=lambda storage, loc: storage))
self.reid_network.eval()
self.reid_network.cuda()
self.tracker = Tracker(self.obj_detect, self.reid_network, self.tracktor['tracker'])
self.transforms = ToTensor()
self.tracker.reset()
def run(self, image):
image = Image.fromarray(image[..., ::-1])
image = self.transforms(image)[None, ...]
blob = {"dets" : torch.Tensor([]), "img" : image}
self.tracker.step(blob)
def get_results(self):
return self.tracker.get_results()
def write_predictions(self, output_dir=None):
if output_dir is None:
output_dir = self.output_dir
results = self.get_results()
if self.tracktor['interpolate']:
results = interpolate(results)
print(f"Writing predictions to: {output_dir}")
write_results(results, output_dir)
def tracker_obj(base_dir):
tracktor = yaml.safe_load(
open(f'{base_dir}/experiments/cfgs/tracktor.yaml').read())['tracktor']
reid = yaml.safe_load(
open(f"{base_dir}/{tracktor['reid_config']}"))['reid']
# set all seeds
output_dir = osp.join(get_output_dir(tracktor['module_name']),
tracktor['name'])
##########################
# Initialize the modules #
##########################
# object detection
obj_detect = FRCNN_FPN(num_classes=2)
obj_detect.load_state_dict(
torch.load(f"{base_dir}/{tracktor['obj_detect_model']}",
map_location=lambda storage, loc: storage))
obj_detect.eval()
obj_detect.cuda()
# reid
reid_network = resnet50(pretrained=False, **reid['cnn'])
reid_network.load_state_dict(
torch.load(f"{base_dir}/{tracktor['reid_weights']}",
map_location=lambda storage, loc: storage))
reid_network.eval()
reid_network.cuda()
# tracktor
if 'oracle' in tracktor:
tracker = OracleTracker(obj_detect, reid_network, tracktor['tracker'],
tracktor['oracle'])
else:
tracker = Tracker(obj_detect, reid_network, tracktor['tracker'])
return tracker
def main(module_name, name, seed, obj_detect_models, reid_models, tracker,
oracle, dataset, load_results, frame_range, interpolate, write_images,
_config, _log, _run):
sacred.commands.print_config(_run)
# set all seeds
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
np.random.seed(seed)
torch.backends.cudnn.deterministic = True
output_dir = osp.join(get_output_dir(module_name), name)
sacred_config = osp.join(output_dir, 'sacred_config.yaml')
if not osp.exists(output_dir):
os.makedirs(output_dir)
with open(sacred_config, 'w') as outfile:
yaml.dump(copy.deepcopy(_config), outfile, default_flow_style=False)
##########################
# Initialize the modules #
##########################
# object detection
_log.info("Initializing object detector(s).")
obj_detects = []
for obj_detect_model in obj_detect_models:
obj_detect = FRCNN_FPN(num_classes=2)
obj_detect.load_state_dict(
torch.load(obj_detect_model,
map_location=lambda storage, loc: storage))
obj_detects.append(obj_detect)
obj_detect.eval()
if torch.cuda.is_available():
obj_detect.cuda()
# reid
_log.info("Initializing reID network(s).")
reid_networks = []
for reid_model in reid_models:
reid_cfg = os.path.join(os.path.dirname(reid_model),
'sacred_config.yaml')
reid_cfg = yaml.safe_load(open(reid_cfg))
reid_network = ReIDNetwork_resnet50(pretrained=False,
**reid_cfg['model_args'])
reid_network.load_state_dict(
torch.load(reid_model, map_location=lambda storage, loc: storage))
reid_network.eval()
if torch.cuda.is_available():
reid_network.cuda()
reid_networks.append(reid_network)
# tracktor
if oracle is not None:
tracker = OracleTracker(obj_detect, reid_network, tracker, oracle)
else:
tracker = Tracker(obj_detect, reid_network, tracker)
time_total = 0
num_frames = 0
mot_accums = []
dataset = Datasets(dataset)
for seq, obj_detect, reid_network in zip(dataset, obj_detects,
reid_networks):
tracker.obj_detect = obj_detect
tracker.reid_network = reid_network
tracker.reset()
_log.info(f"Tracking: {seq}")
start_frame = int(frame_range['start'] * len(seq))
end_frame = int(frame_range['end'] * len(seq))
seq_loader = DataLoader(
torch.utils.data.Subset(seq, range(start_frame, end_frame)))
num_frames += len(seq_loader)
results = {}
if load_results:
results = seq.load_results(output_dir)
if not results:
start = time.time()
for frame_data in tqdm(seq_loader):
with torch.no_grad():
tracker.step(frame_data)
results = tracker.get_results()
time_total += time.time() - start
_log.info(f"Tracks found: {len(results)}")
_log.info(f"Runtime for {seq}: {time.time() - start :.2f} s.")
if interpolate:
results = interpolate_tracks(results)
_log.info(f"Writing predictions to: {output_dir}")
seq.write_results(results, output_dir)
if seq.no_gt:
_log.info("No GT data for evaluation available.")
else:
mot_accums.append(get_mot_accum(results, seq_loader))
if write_images:
plot_sequence(results, seq,
osp.join(output_dir, str(dataset), str(seq)),
write_images)
if time_total:
_log.info(
f"Tracking runtime for all sequences (without evaluation or image writing): "
f"{time_total:.2f} s for {num_frames} frames ({num_frames / time_total:.2f} Hz)"
)
if mot_accums:
_log.info("Evaluation:")
evaluate_mot_accums(mot_accums,
[str(s) for s in dataset if not s.no_gt],
generate_overall=True)
def main(tracktor, reid, _config, _log, _run):
sacred.commands.print_config(_run)
# set all seeds
torch.manual_seed(tracktor['seed'])
torch.cuda.manual_seed(tracktor['seed'])
np.random.seed(tracktor['seed'])
torch.backends.cudnn.deterministic = True
output_dir = osp.join(get_output_dir(tracktor['module_name']),
tracktor['name'])
sacred_config = osp.join(output_dir, 'sacred_config.yaml')
if not osp.exists(output_dir):
os.makedirs(output_dir)
with open(sacred_config, 'w') as outfile:
yaml.dump(_config, outfile, default_flow_style=False)
##########################
# Initialize the modules #
##########################
# object detection
_log.info("Initializing object detector.")
obj_detect = FRCNN_FPN(num_classes=2)
obj_detect.load_state_dict(
torch.load(_config['tracktor']['obj_detect_model'],
map_location=lambda storage, loc: storage))
obj_detect.eval()
obj_detect.cuda()
# reid
reid_network = resnet50(pretrained=False, **reid['cnn'])
reid_network.load_state_dict(
torch.load(tracktor['reid_weights'],
map_location=lambda storage, loc: storage))
reid_network.eval()
reid_network.cuda()
# motion network
motion_network = None
if tracktor['tracker']['motion_model_enabled'] and not tracktor['motion'][
'use_cva_model']:
motion_network = eval(
tracktor['motion']['model'])(**tracktor['motion']['model_args'])
motion_network.load_state_dict(
torch.load(tracktor['motion']['network_weights'])['model'])
motion_network.eval().cuda()
# tracktor
if 'oracle' in tracktor:
tracker = OracleTracker(obj_detect, reid_network, tracktor['tracker'],
tracktor['oracle'])
else:
tracker = Tracker(obj_detect, reid_network, motion_network,
tracktor['tracker'], tracktor['motion'], 2)
time_total = 0
num_frames = 0
mot_accums = []
dataset = Datasets(tracktor['dataset'])
for seq in dataset:
tracker.reset()
_log.info(f"Tracking: {seq}")
data_loader = DataLoader(seq, batch_size=1, shuffle=False)
start = time.time()
all_mm_times = []
all_warp_times = []
for i, frame in enumerate(tqdm(data_loader)):
if len(seq) * tracktor['frame_split'][0] <= i <= len(
seq) * tracktor['frame_split'][1]:
with torch.no_grad():
mm_time, warp_time = tracker.step(frame)
if mm_time is not None:
all_mm_times.append(mm_time)
if warp_time is not None:
all_warp_times.append(warp_time)
num_frames += 1
results = tracker.get_results()
time_total += time.time() - start
_log.info(f"Tracks found: {len(results)}")
_log.info(f"Runtime for {seq}: {time.time() - start :.1f} s.")
_log.info(
f"Average FPS for {seq}: {len(data_loader) / (time.time() - start) :.3f}"
)
_log.info(
f"Average MM time for {seq}: {float(np.array(all_mm_times).mean()) :.3f} s"
)
if all_warp_times:
_log.info(
f"Average warp time for {seq}: {float(np.array(all_warp_times).mean()) :.3f} s"
)
if tracktor['interpolate']:
results = interpolate(results)
if 'semi_online' in tracktor and tracktor['semi_online']:
for i, track in results.items():
for frame in sorted(track, reverse=True):
if track[frame][5] == 0:
break
del track[frame]
if tracktor['write_images']:
plot_sequence(results, seq,
osp.join(output_dir, tracktor['dataset'], str(seq)),
tracktor['tracker']['plot_mm'])
if seq.no_gt:
_log.info(f"No GT data for evaluation available.")
else:
mot_accums.append(get_mot_accum(results, seq))
_log.info(f"Writing predictions to: {output_dir}")
seq.write_results(results, output_dir)
_log.info(
f"Tracking runtime for all sequences (without evaluation or image writing): "
f"{time_total:.2f} s for {num_frames} frames ({num_frames / time_total:.2f} Hz)"
)
if mot_accums:
evaluate_mot_accums(mot_accums,
[str(s) for s in dataset if not s.no_gt],
generate_overall=True)
def main(tracktor, reid, _config, _log, _run):
target = Target()
targetpath = target.Folder()
targetname = target.TargetName()
vottpath = target.GetVottPath()
vottfile = target.GetVottContent()
dictid, timelist = target.GetTagTime(vottfile)
print(f"{len(timelist)} frames were tagged")
timedict = target.ExtractByTimeList(timelist)
bbdict = target.GetbbWithTime(vottfile)
sacred.commands.print_config(_run)
# set all seeds
torch.manual_seed(tracktor['seed'])
torch.cuda.manual_seed(tracktor['seed'])
np.random.seed(tracktor['seed'])
torch.backends.cudnn.deterministic = True
output_dir = osp.join(get_output_dir(tracktor['module_name']),
tracktor['name'])
sacred_config = osp.join(output_dir, 'sacred_config.yaml')
if not osp.exists(output_dir):
os.makedirs(output_dir)
with open(sacred_config, 'w') as outfile:
yaml.dump(_config, outfile, default_flow_style=False)
##########################
# Initialize the modules #
##########################
# object detection
_log.info("Initializing object detector.")
obj_detect = FRCNN_FPN(num_classes=2)
obj_detect.load_state_dict(
torch.load(_config['tracktor']['obj_detect_model'],
map_location=lambda storage, loc: storage))
obj_detect.eval()
obj_detect.cuda()
# reid
reid_network = resnet50(pretrained=False, **reid['cnn'])
reid_network.load_state_dict(
torch.load(tracktor['reid_weights'],
map_location=lambda storage, loc: storage))
reid_network.eval()
reid_network.cuda()
# tracktor
print("Tracktor初始化完成")
tracker = Tracker(obj_detect, reid_network, tracktor['tracker'])
time_total = 0
num_frames = 0
mot_accums = []
dataset = Datasets(tracktor['dataset'])
for seq in dataset:
tracker.reset()
start = time.time()
_log.info(f"Tracking: {seq}")
data_loader = DataLoader(seq, batch_size=1, shuffle=False)
print(f"{seq}加載完成, tracking開始")
for i, frame in enumerate(tqdm(data_loader)):
if len(seq) * tracktor['frame_split'][0] <= i <= len(
seq) * tracktor['frame_split'][1]:
id = tracker.step(frame, bbdict[timedict["%06d" % num_frames]])
target.WriteID2asset(id, dictid[timedict["%06d" % num_frames]])
num_frames += 1
results = tracker.get_results()
ids = list(results.keys())
target.WriteID2vott(ids, vottfile=vottfile)
time_total += time.time() - start
_log.info(f"Tracks found: {len(results)}")
_log.info(f"Runtime for {seq}: {time.time() - start :.1f} s.")
target.CleanImg()
if tracktor['interpolate']:
results = interpolate(results)
if seq.no_gt:
_log.info(f"No GT data for evaluation available.")
else:
mot_accums.append(get_mot_accum(results, seq))
_log.info(f"Writing predictions to: {output_dir}")
seq.write_results(results, output_dir)
if tracktor['write_images']:
plot_sequence(results, seq,
osp.join(output_dir, tracktor['dataset'], str(seq)))
if tracktor['write_videos']:
plot_sequence_video(
results, seq,
osp.join(output_dir, tracktor['dataset'], str(seq)))
_log.info(
f"Tracking runtime for all sequences (without evaluation or image writing): "
f"{time_total:.1f} s ({num_frames / time_total:.1f} Hz)")
if mot_accums:
evaluate_mot_accums(mot_accums,
[str(s) for s in dataset if not s.no_gt],
generate_overall=True)
def __init__(
self,
ckpt_path,
frcnn_weights_path,
reid_weights_path,
tracking_cfg_path,
preprocessing_cfg_path,
use_gt,
pre_cnn,
pre_track,
):
self.name = "MPNTracker"
self.use_gt = use_gt
if not self.use_gt:
with open(tracking_cfg_path) as config_file:
config = yaml.load(config_file)
with open(preprocessing_cfg_path) as config_file:
pre_config = yaml.load(config_file)
frcnn_prepr_params = pre_config["frcnn_prepr_params"]
tracktor_params = pre_config["tracktor_params"]
CustomMOTNeuralSolver.reid_weights_path = reid_weights_path
# preprocessor
self.pre_track = pre_track
if self.pre_track == "Tracktor" or self.pre_track == "FRCNN":
obj_detect = FRCNN_FPN(num_classes=2)
obj_detect.load_state_dict(
torch.load(
frcnn_weights_path,
map_location=lambda storage, loc: storage,
))
obj_detect.eval()
obj_detect.cuda()
if self.pre_track == "Tracktor":
self.prepr_params = tracktor_params
make_deterministic(self.prepr_params["seed"])
self.preprocessor = Tracker(obj_detect, None,
self.prepr_params["tracker"])
elif self.pre_track == "FRCNN":
self.prepr_params = frcnn_prepr_params
make_deterministic(self.prepr_params["seed"])
self.preprocessor = FRCNNPreprocessor(
obj_detect, self.prepr_params)
self.transforms = ToTensor()
# Load model from checkpoint and update config entries that may vary from the ones used in training
self.model = CustomMOTNeuralSolver.load_from_checkpoint(
checkpoint_path=ckpt_path)
self.model.cnn_model.eval()
self.model.hparams.update({
"eval_params": config["eval_params"],
"data_splits": config["data_splits"],
})
self.model.hparams["dataset_params"][
"precomputed_embeddings"] = False
self.model.hparams["dataset_params"]["img_batch_size"] = 2500
self.pre_cnn = pre_cnn
if self.pre_cnn:
self.extract_transforms = Compose((
Resize(self.model.hparams["dataset_params"]["img_size"]),
ToTensor(),
Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]),
))
def main(tracktor, reid, _config, _log, _run):
sacred.commands.print_config(_run)
# set all seeds
torch.manual_seed(tracktor['seed'])
torch.cuda.manual_seed(tracktor['seed'])
np.random.seed(tracktor['seed'])
torch.backends.cudnn.deterministic = True
output_dir = osp.join(get_output_dir(tracktor['module_name']),
tracktor['name'])
sacred_config = osp.join(output_dir, 'sacred_config.yaml')
if not osp.exists(output_dir):
os.makedirs(output_dir)
with open(sacred_config, 'w') as outfile:
yaml.dump(_config, outfile, default_flow_style=False)
##########################
# Initialize the modules #
##########################
# object detection
_log.info("Initializing object detector.")
obj_detect = FRCNN_FPN(num_classes=2)
obj_detect.load_state_dict(
torch.load(_config['tracktor']['obj_detect_model'],
map_location=lambda storage, loc: storage))
obj_detect.eval()
obj_detect.cuda()
# reid
reid_network = resnet50(pretrained=False, **reid['cnn'])
reid_network.load_state_dict(
torch.load(tracktor['reid_weights'],
map_location=lambda storage, loc: storage))
reid_network.eval()
reid_network.cuda()
# tracktor
if 'oracle' in tracktor:
tracker = OracleTracker(obj_detect, reid_network, tracktor['tracker'],
tracktor['oracle'])
else:
tracker = Tracker(obj_detect, reid_network, tracktor['tracker'])
time_total = 0
num_frames = 0
mot_accums = []
dataset = Datasets(tracktor['dataset'])
for seq in dataset:
tracker.reset()
start = time.time()
_log.info(f"Tracking: {seq}")
data_loader = DataLoader(seq, batch_size=1, shuffle=False)
for i, frame in enumerate(tqdm(data_loader)):
if len(seq) * tracktor['frame_split'][0] <= i <= len(
seq) * tracktor['frame_split'][1]:
tracker.step(frame)
num_frames += 1
results = tracker.get_results()
time_total += time.time() - start
_log.info(f"Tracks found: {len(results)}")
_log.info(f"Runtime for {seq}: {time.time() - start :.1f} s.")
if tracktor['interpolate']:
results = interpolate(results)
if seq.no_gt:
_log.info(f"No GT data for evaluation available.")
else:
mot_accums.append(get_mot_accum(results, seq))
_log.info(f"Writing predictions to: {output_dir}")
seq.write_results(results, output_dir)
if tracktor['write_images']:
plot_sequence(results, seq,
osp.join(output_dir, tracktor['dataset'], str(seq)))
img_array = []
dir = osp.join(output_dir, tracktor['dataset'], str(seq), "*.jpg")
files = glob.glob(dir)
sorted_files = natsorted(files)
for filename in sorted_files:
img = cv2.imread(filename)
height, width, layers = img.shape
size = (width, height)
img_array.append(img)
out = cv2.VideoWriter(
osp.join(output_dir, tracktor['dataset'],
str(seq), "result_video.avi"),
cv2.VideoWriter_fourcc(*'DIVX'), 10, size)
for i in range(len(img_array)):
out.write(img_array[i])
out.release()
_log.info(
f"Tracking runtime for all sequences (without evaluation or image writing): "
f"{time_total:.1f} s ({num_frames / time_total:.1f} Hz)")
if mot_accums:
evaluate_mot_accums(mot_accums,
[str(s) for s in dataset if not s.no_gt],
generate_overall=True)
def main(tracktor, reid, _config, _log, _run):
sacred.commands.print_config(_run)
# set all seeds
torch.manual_seed(tracktor['seed'])
torch.cuda.manual_seed(tracktor['seed'])
np.random.seed(tracktor['seed'])
torch.backends.cudnn.deterministic = True
output_dir = osp.join(get_output_dir(tracktor['module_name']),
tracktor['name'], tracktor['output_subdir'])
sacred_config = osp.join(output_dir, 'sacred_config.yaml')
if not osp.exists(output_dir):
os.makedirs(output_dir)
with open(sacred_config, 'w') as outfile:
yaml.dump(_config, outfile, default_flow_style=False)
##########################
# Initialize the modules #
##########################
# object detection
_log.info("Initializing object detector.")
obj_detect = FRCNN_FPN(num_classes=2).to(device)
obj_detect.load_state_dict(
torch.load(_config['tracktor']['obj_detect_model'],
map_location=lambda storage, loc: storage))
obj_detect.eval()
# reid
reid_network = resnet50(pretrained=False, **reid['cnn']).to(device)
reid_network.load_state_dict(
torch.load(tracktor['reid_weights'],
map_location=lambda storage, loc: storage))
reid_network.eval()
# tracktor
if 'oracle' in tracktor:
tracker = OracleTracker(obj_detect, reid_network, tracktor['tracker'],
tracktor['oracle'])
else:
tracker = Tracker(obj_detect, reid_network, tracktor['tracker'])
time_total = 0
num_frames = 0
mot_accums = []
dataset = Datasets(tracktor['dataset'])
for seq in dataset:
tracker.reset()
start = time.time()
_log.info(f"Tracking: {seq}")
data_loader = DataLoader(seq, batch_size=1, shuffle=False)
for i, frame in enumerate(tqdm(data_loader)):
if len(seq) * tracktor['frame_split'][0] <= i <= len(
seq) * tracktor['frame_split'][1]:
tracker.step(frame, i)
num_frames += 1
results = tracker.get_results()
time_total += time.time() - start
_log.info(f"Tracks found: {len(results)}")
_log.info(f"Runtime for {seq}: {time.time() - start :.1f} s.")
if tracktor['interpolate']:
results = interpolate(results)
if seq.no_gt:
_log.info(f"No GT data for evaluation available.")
else:
mot_accums.append(get_mot_accum(results, seq))
_log.info(f"Writing predictions to: {output_dir}")
seq.write_results(results, output_dir)
if tracktor['write_images']:
plot_sequence(results, seq,
osp.join(output_dir, tracktor['dataset'], str(seq)))
_log.info(
f"Tracking runtime for all sequences (without evaluation or image writing): "
f"{time_total:.1f} s ({num_frames / time_total:.1f} Hz)")
if mot_accums:
summary = evaluate_mot_accums(mot_accums,
[str(s) for s in dataset if not s.no_gt],
generate_overall=True)
summary.to_pickle(
"output/finetuning_results/results_{}_{}_{}_{}_{}.pkl".format(
tracktor['output_subdir'],
tracktor['tracker']['finetuning']['max_displacement'],
tracktor['tracker']['finetuning']['batch_size'],
tracktor['tracker']['finetuning']['learning_rate'],
tracktor['tracker']['finetuning']['iterations']))
def test_motion_model(val_loader, tracker_config, motion_model):
obj_detect = FRCNN_FPN(num_classes=2)
obj_detect.load_state_dict(
torch.load(tracker_config['tracktor']['obj_detect_model'],
map_location=lambda storage, loc: storage))
obj_detect.eval()
obj_detect.cuda()
reid_network = resnet50(pretrained=False, output_dim=128)
reid_network.load_state_dict(
torch.load(tracker_config['tracktor']['reid_weights'],
map_location=lambda storage, loc: storage))
reid_network.eval()
reid_network.cuda()
pred_loss_func = nn.SmoothL1Loss()
loss_iters = []
low_vis_loss_sum = 0.0
low_vis_num = 0
high_vis_loss_sum = 0.0
high_vis_num = 0
total_iters = len(val_loader)
n_iters = 0
with torch.no_grad():
for data in val_loader:
n_iters += 1
early_reid = get_batch_mean_early_reid(reid_network,
data['early_reid_patches'])
curr_reid = reid_network(data['curr_reid_patch'].cuda())
conv_features, repr_features = get_features(
obj_detect, data['curr_img'], data['curr_gt_app'])
prev_loc = data['prev_gt_warped'].cuda()
curr_loc = data['curr_gt_warped'].cuda()
label_loc = data['label_gt'].cuda()
curr_vis = data['curr_vis'].cuda()
pred_loc_wh, vis = motion_model(early_reid, curr_reid,
conv_features, repr_features,
prev_loc, curr_loc)
label_loc_wh = two_p_to_wh(label_loc)
pred_loss = pred_loss_func(pred_loc_wh, label_loc_wh)
loss_iters.append(pred_loss.item())
low_vis_ind = curr_vis < 0.3
if low_vis_ind.any():
low_vis_pred_loss = pred_loss_func(pred_loc_wh[low_vis_ind],
label_loc_wh[low_vis_ind])
low_vis_loss_sum += (low_vis_pred_loss *
torch.sum(low_vis_ind)).item()
low_vis_num += torch.sum(low_vis_ind).item()
high_vis_ind = curr_vis > 0.7
if high_vis_ind.any():
high_vis_pred_loss = pred_loss_func(pred_loc_wh[high_vis_ind],
label_loc_wh[high_vis_ind])
high_vis_loss_sum += (high_vis_pred_loss *
torch.sum(high_vis_ind)).item()
high_vis_num += torch.sum(high_vis_ind).item()
if n_iters % 50 == 0:
print('Iter %5d/%5d finished.' % (n_iters, total_iters),
flush=True)
mean_loss = np.mean(loss_iters)
mean_low_vis_loss = low_vis_loss_sum / low_vis_num
mean_high_vis_loss = high_vis_loss_sum / high_vis_num
print('All finished! Loss %.6f, low vis loss %.6f, high vis loss %.6f.' %
(mean_loss, mean_low_vis_loss, mean_high_vis_loss))
def main(tracktor, reid, _config, _log, _run):
sacred.commands.print_config(_run)
# set all seeds
torch.manual_seed(tracktor['seed'])
torch.cuda.manual_seed(tracktor['seed'])
np.random.seed(tracktor['seed'])
torch.backends.cudnn.deterministic = True
output_dir = osp.join(get_output_dir(tracktor['module_name']), tracktor['name'])
sacred_config = osp.join(output_dir, 'sacred_config.yaml')
if not osp.exists(output_dir):
os.makedirs(output_dir)
with open(sacred_config, 'w') as outfile:
yaml.dump(_config, outfile, default_flow_style=False)
##########################
# Initialize the modules #
##########################
# object detection
_log.info("Initializing object detector.")
obj_detect = FRCNN_FPN(num_classes=2)
obj_detect.load_state_dict(torch.load(_config['tracktor']['obj_detect_model'],
map_location=lambda storage, loc: storage))
obj_detect.eval()
obj_detect.cuda()
# reid
reid_network = resnet50(pretrained=False, **reid['cnn'])
reid_network.load_state_dict(torch.load(tracktor['reid_weights'],
map_location=lambda storage, loc: storage))
reid_network.eval()
reid_network.cuda()
# tracktor
if 'oracle' in tracktor:
tracker = OracleTracker(obj_detect, reid_network, tracktor['tracker'], tracktor['oracle'])
else:
tracker = Tracker(obj_detect, reid_network, tracktor['tracker'])
time_total = 0
num_frames = 0
mot_accums = []
# Data transform
normalize_mean=[0.485, 0.456, 0.406]
normalize_std=[0.229, 0.224, 0.225]
# dataset = Datasets(tracktor['dataset'])
transforms = ToTensor()
# transforms = Compose([ToTensor(), Normalize(normalize_mean,
# normalize_std)])
tracker.reset()
# tracker.public_detections=False
start = time.time()
_log.info(f"Tracking: video")
# Load video and annotations
cap = cv2.VideoCapture("/home/yc3390/camera_detection_demo/data/prid2011_videos/test_b_1min_1min.mp4")
with open("/home/yc3390/camera_detection_demo/data/prid2011_videos/anno_b.pkl", 'rb') as f:
gts = pk.load(f)
det_file = "/data/yc3390/tracktor_output/output/tracktor/MOT17/Tracktor++/Video-result_ReID.txt"
# with open("/data/yc3390/tracktor_output/output/tracktor/MOT17/Tracktor++/Video-result_ReID.pkl", 'rb') as f:
# dts = pk.load(f)
# for dt in dts:
# if len(dt['boxes'][0]):
# for i in range(len(dt['boxes'])):
# dt['boxes'][i][-1] = -1
offset = 25 * 60
dets = {}
for i in range(1, offset+1):
dets[i] = []
assert osp.exists(det_file)
with open(det_file, "r") as inf:
reader = csv.reader(inf, delimiter=',')
for row in reader:
x1 = float(row[2]) - 1
y1 = float(row[3]) - 1
# This -1 accounts for the width (width of 1 x1=x2)
x2 = x1 + float(row[4]) - 1
y2 = y1 + float(row[5]) - 1
score = float(row[6])
bb = np.array([x1,y1,x2,y2], dtype=np.float32)
dets[int(float(row[0]))].append(bb)
frame_count = offset
while True:
ret, image = cap.read()
if not ret:
break
# BGR to RGB
image = Image.fromarray(image[..., ::-1])
image = transforms(image)[None, ...]
# Detection
# if frame_count in gts.keys():
# frames =
blob = {"dets" : torch.Tensor([dets[i]]), "img" : image}
tracker.step(blob)
frame_count += 1
print("Finished ", frame_count, output_dir, image.shape)
results = tracker.get_results()
time_total += time.time() - start
_log.info(f"Tracks found: {len(results)}")
_log.info(f"Runtime for video: {time.time() - start :.1f} s.")
if tracktor['interpolate']:
results = interpolate(results)
if True:
_log.info(f"No GT data for evaluation available.")
else:
mot_accums.append(get_mot_accum(results, seq))
_log.info(f"Writing predictions to: {output_dir}")
write_results(results, output_dir)
box_head_features = obj_detect.roi_heads.box_head(box_features)
return box_features.cpu(), box_head_features.cpu()
with open('experiments/cfgs/tracktor.yaml', 'r') as f:
tracker_config = yaml.safe_load(f)
val_set = MOT17Vis('val', 0.8, 0.0, val_bbox_jitter=True)
val_loader = DataLoader(val_set, batch_size=1, shuffle=True, num_workers=2)
obj_detect = FRCNN_FPN(num_classes=2)
obj_detect.load_state_dict(
torch.load(tracker_config['tracktor']['obj_detect_model'],
map_location=lambda storage, loc: storage))
obj_detect.eval()
obj_detect.cuda()
vis_model = VisEst(conv_only=False)
vis_model.load_state_dict(
torch.load(
'output/tracktor/motion/vis_valjit_flip_fullfeat_l21e-4/vis_model_epoch_94.pth'
))
vis_model.eval()
vis_model.cuda()
for data in val_loader:
conv_features, repr_features = get_features(obj_detect, data['img'],
data['gt'])
conv_features = conv_features.cuda()
repr_features = repr_features.cuda()
def main(dataset_names, prepr_w_tracktor, frcnn_prepr_params, tracktor_params,
frcnn_weights, _config, _log, _run):
sacred.commands.print_config(_run)
if prepr_w_tracktor:
prepr_params = tracktor_params
else:
prepr_params = frcnn_prepr_params
make_deterministic(prepr_params['seed'])
MOV_CAMERA_DICT = {**MOT15_MOV_CAMERA_DICT, **MOT17_MOV_CAMERA_DICT}
# object detection
_log.info("Initializing object detector.")
obj_detect = FRCNN_FPN(num_classes=2)
obj_detect.load_state_dict(
torch.load(osp.join(OUTPUT_PATH, frcnn_weights),
map_location=lambda storage, loc: storage))
obj_detect.eval()
obj_detect.cuda()
if prepr_w_tracktor:
preprocessor = Tracker(obj_detect, None, prepr_params['tracker'])
else:
preprocessor = FRCNNPreprocessor(obj_detect, prepr_params)
_log.info(
f"Starting preprocessing of datasets {dataset_names} with {'Tracktor' if prepr_w_tracktor else 'FRCNN'} \n"
)
for dataset_name in dataset_names:
dataset = Datasets(dataset_name)
_log.info(
f"Preprocessing {len(dataset)} sequences from dataset {dataset_name} \n"
)
time_total = 0
num_frames = 0
for seq in dataset:
preprocessor.reset()
start = time.time()
_log.info(f"Preprocessing : {seq}")
if prepr_w_tracktor:
preprocessor.do_align = tracktor_params['tracker'][
'do_align'] and (MOV_CAMERA_DICT[str(seq)])
data_loader = DataLoader(seq,
batch_size=1,
shuffle=False,
num_workers=8,
pin_memory=True)
for i, frame in enumerate(tqdm(data_loader)):
with torch.no_grad():
preprocessor.step(frame)
num_frames += 1
time_total += time.time() - start
_log.info(f"Runtime for {seq}: {time.time() - start :.1f} s.")
output_file_path = osp.join(seq.seq_path, 'det',
prepr_params['det_file_name'])
if prepr_w_tracktor:
results = preprocessor.get_results()
seq.write_results(results, output_file_path)
else:
_log.info(f"Writing predictions in: {output_file_path}")
preprocessor.save_results(output_file_path)
_log.info(
f"Tracking runtime for all sequences (without evaluation or image writing): "
f"{time_total:.1f} s ({num_frames / time_total:.1f} Hz)")
def main():
args = parse_args()
logging.basicConfig(level=logging.INFO,format='%(asctime)s %(name)-12s %(levelname)-8s %(message)s')
logging.root.setLevel(logging.INFO)
is_cuda = torch.cuda.is_available()
if is_cuda:
LOG.info('-' * 50)
LOG.info('Enabling CUDA')
LOG.info('-' * 50)
device = torch.device('cuda' if is_cuda else 'cpu')
# sacred.commands.print_config(_run)
# set all seeds
torch.manual_seed(tracktor['seed'])
np.random.seed(tracktor['seed'])
if is_cuda:
torch.cuda.manual_seed(tracktor['seed'])
torch.backends.cudnn.deterministic = True
output_dir = osp.join(get_output_dir(tracktor['module_name']), tracktor['name'])
if not osp.exists(output_dir):
os.makedirs(output_dir)
##########################
# Initialize the modules #
##########################
# object detection
LOG.info("Initializing object detector.")
obj_detect = FRCNN_FPN(num_classes=2)
obj_detect_state_dict = torch.load(args.detection_path, map_location=device)
obj_detect.load_state_dict(obj_detect_state_dict)
obj_detect.eval()
if is_cuda:
obj_detect.cuda()
# LOG.info('Load detection model...')
# obj_detect = load_object_detection_driver(args.detection_path)
# LOG.info('Done.')
# reid
LOG.info("Initializing reidentification network.")
reid_network = resnet50(pretrained=False, output_dim=128)
reid_network.load_state_dict(torch.load(args.reid_path, map_location=device))
reid_network.eval()
if is_cuda:
reid_network.cuda()
# tracktor
if 'oracle' in tracktor:
tracker = OracleTracker(obj_detect, reid_network, tracktor['tracker'], tracktor['oracle'])
else:
tracker = Tracker(obj_detect, reid_network, tracktor['tracker'])
time_total = 0
tracker.reset()
start = time.time()
vc = cv2.VideoCapture(args.source)
frame_count = vc.get(cv2.CAP_PROP_FRAME_COUNT)
video_output = args.output
fourcc = cv2.VideoWriter_fourcc(*'avc1')
width = int(vc.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(vc.get(cv2.CAP_PROP_FRAME_HEIGHT))
fps = vc.get(cv2.CAP_PROP_FPS)
each_frame = args.every_frame
writer = cv2.VideoWriter(video_output, fourcc, fps / each_frame, frameSize=(width, height))
LOG.info(f"Tracking: {args.source}")
frame_id = 0
frame_num = 0
results = {}
try:
while True:
frame_num += 1
if frame_num % each_frame == 0:
ret, frame = vc.read()
if not ret:
break
else:
vc.grab()
continue
frame_id += 1
if frame_id % 50 == 0:
LOG.info(f'Processing frame {frame_id}')
if frame_count * tracktor['frame_split'][0] <= frame_id <= frame_count * tracktor['frame_split'][1]:
rgb_frame = frame[:, :, ::-1]
torch_frame = F.to_tensor(rgb_frame.copy())
torch_frame = torch_frame.expand([1, *torch_frame.shape])
if is_cuda:
torch_frame = torch_frame.cuda()
torch_blob = {
'img': torch_frame
}
tracker.step(torch_blob, frame)
# __import__('ipdb').set_trace()
results = tracker.results
output = draw_boxes(frame, frame_id - 1, results=results)
writer.write(output)
except KeyboardInterrupt:
LOG.info('Stopping.')
writer.release()
time_total += time.time() - start
LOG.info(f"Tracks found: {len(results)}")
LOG.info(f"Runtime for {args.source}: {time.time() - start :.1f} s.")
if tracktor['interpolate']:
results = utils.interpolate(results)
# if tracktor['write_images']:
# utils.plot_sequence(results, seq, osp.join(output_dir, args.source))
LOG.info(f"Tracking runtime for all sequences (without evaluation or image writing): "
f"{time_total:.1f} s ({frame_id / time_total:.1f} Hz)")
def main(tracktor, reid, _config, _log, _run):
sacred.commands.print_config(_run)
# set all seeds
torch.manual_seed(tracktor['seed'])
torch.cuda.manual_seed(tracktor['seed'])
np.random.seed(tracktor['seed'])
torch.backends.cudnn.deterministic = True
output_dir = osp.join(get_output_dir(tracktor['module_name']),
tracktor['name'])
sacred_config = osp.join(output_dir, 'sacred_config.yaml')
if not osp.exists(output_dir):
os.makedirs(output_dir)
with open(sacred_config, 'w') as outfile:
yaml.dump(_config, outfile, default_flow_style=False)
##########################
# Initialize the modules #
##########################
_log.info("Initializing object detector.")
# object detection
obj_detect = FRCNN_FPN(num_classes=2, correlation_head=CorrelationHead())
obj_detect_model = torch.load(_config['tracktor']['obj_detect_model'],
map_location=lambda storage, loc: storage)
correlation_weights = torch.load(
_config['tracktor']['correlation_weights'],
map_location=lambda storage, loc: storage)
for k in correlation_weights:
obj_detect_model.update(
{"correlation_head." + k: correlation_weights[k]})
obj_detect.load_state_dict(obj_detect_model)
obj_detect.eval()
obj_detect.cuda()
# reid
reid_network = resnet50(pretrained=False, **reid['cnn'])
reid_network.load_state_dict(
torch.load(tracktor['reid_weights'],
map_location=lambda storage, loc: storage))
reid_network.eval()
reid_network.cuda()
# tracktor
if 'oracle' in tracktor:
tracker = OracleTracker(obj_detect, reid_network, tracktor['tracker'],
tracktor['oracle'])
else:
tracker = Tracker(obj_detect, reid_network, tracktor['tracker'])
time_total = 0
num_frames = 0
mot_accums = []
dataset = Datasets(tracktor['dataset'])
for seq in dataset:
tracker.reset()
start = time.time()
_log.info(f"Tracking: {seq}")
data_loader = DataLoader(seq, batch_size=1, shuffle=False)
for i, frame in enumerate(tqdm(data_loader)):
if len(seq) * tracktor['frame_split'][0] <= i <= len(
seq) * tracktor['frame_split'][1]:
with torch.no_grad():
tracker.step(frame)
num_frames += 1
results = tracker.get_results()
time_total += time.time() - start
_log.info(f"Tracks found: {len(results)}")
_log.info(f"Runtime for {seq}: {time.time() - start :.2f} s.")
if tracktor['interpolate']:
results = interpolate(results)
if seq.no_gt:
_log.info(f"No GT data for evaluation available.")
else:
mot_accums.append(get_mot_accum(results, seq))
_log.info(f"Writing predictions to: {output_dir}")
seq.write_results(results, output_dir)
if tracktor['write_images']:
plot_sequence(results, seq,
osp.join(output_dir, tracktor['dataset'], str(seq)))
score_killed_tracks = tracker.get_score_killed_tracks()
_log.info(f"Score Killed Tracks: ({len(score_killed_tracks)})")
for kill in score_killed_tracks:
_log.info(
f"Track [ {kill['id']:3d} ] killed in frame [ {kill['frame']:3d} ]"
)
nms_killed_tracks = tracker.get_nms_killed_tracks()
_log.info(f"NMS Killed Tracks ({len(nms_killed_tracks)}):")
for kill in nms_killed_tracks:
_log.info(
f"Track [ {kill['id']:3d} ] killed in frame [ {kill['frame']:3d} ]"
)
_log.info(
f"Tracking runtime for all sequences (without evaluation or image writing): "
f"{time_total:.2f} s for {num_frames} frames ({num_frames / time_total:.2f} Hz)"
)
if mot_accums:
evaluate_mot_accums(mot_accums,
[str(s) for s in dataset if not s.no_gt],
generate_overall=True)
def train_main(use_ecc, use_modulator, use_bn, use_residual, use_reid_distance, no_visrepr, vis_loss_ratio, no_vis_loss, motion_noise,
lr, weight_decay, batch_size, output_dir, ex_name):
random.seed(12345)
torch.manual_seed(12345)
torch.cuda.manual_seed(12345)
np.random.seed(12345)
torch.backends.cudnn.deterministic = True
output_dir = osp.join(output_dir, ex_name)
log_file = osp.join(output_dir, 'epoch_log.txt')
if not osp.exists(output_dir):
os.makedirs(output_dir)
with open(log_file, 'w') as f:
f.write('[Experiment name]%s\n\n' % ex_name)
f.write('[Parameters]\n')
f.write('use_ecc=%r\nuse_modulator=%r\nuse_bn=%r\nuse_residual=%r\nuse_reid_distance=%r\nno_visrepr=%r\nvis_loss_ratio=%f\nno_vis_loss=%r\nmotion_noise=%f\nlr=%f\nweight_decay=%f\nbatch_size=%d\n\n' %
(use_ecc, use_modulator, use_bn, use_residual, use_reid_distance, no_visrepr, vis_loss_ratio, no_vis_loss, motion_noise, lr, weight_decay, batch_size))
f.write('[Loss log]\n')
with open('experiments/cfgs/tracktor.yaml', 'r') as f:
tracker_config = yaml.safe_load(f)
#################
# Load Datasets #
#################
train_set = MOT17ClipsWrapper('train', 0.8, 0.0, clip_len=10, motion_noise=motion_noise, train_jitter=True, ecc=True, tracker_cfg=tracker_config)
train_loader = DataLoader(train_set, batch_size=1, shuffle=True, num_workers=1, collate_fn=clips_wrapper_collate)
val_set = MOT17ClipsWrapper('val', 0.8, 0.0, clip_len=10, motion_noise=motion_noise, train_jitter=True, ecc=True, tracker_cfg=tracker_config)
val_loader = DataLoader(val_set, batch_size=1, shuffle=False, num_workers=1, collate_fn=clips_wrapper_collate)
with open(osp.join(cfg.ROOT_DIR, 'output', 'precomputed_ecc_matrices_3.pkl'), 'rb') as f:
ecc_dict = pickle.load(f)
train_set.load_precomputed_ecc_warp_matrices(ecc_dict)
val_set.load_precomputed_ecc_warp_matrices(ecc_dict)
########################
# Initializing Modules #
########################
obj_detect = FRCNN_FPN(num_classes=2)
obj_detect.load_state_dict(torch.load(tracker_config['tracktor']['obj_detect_model'],
map_location=lambda storage, loc: storage))
obj_detect.eval()
obj_detect.cuda()
motion_model = MotionModelReID(use_modulator=use_modulator, use_bn=use_bn, use_residual=use_residual,
use_reid_distance=use_reid_distance, no_visrepr=no_visrepr)
motion_model.train()
motion_model.cuda()
reid_network = resnet50(pretrained=False, output_dim=128)
reid_network.load_state_dict(torch.load(tracker_config['tracktor']['reid_weights'],
map_location=lambda storage, loc: storage))
reid_network.eval()
reid_network.cuda()
optimizer = torch.optim.Adam(motion_model.parameters(), lr=lr, weight_decay=weight_decay)
pred_loss_func = nn.SmoothL1Loss()
vis_loss_func = nn.MSELoss()
#######################
# Training Parameters #
#######################
# usage: historical_reid, curr_reid, roi_pool_output, representation_feature, prev_loc, curr_loc, curr_vis, label_loc
batch_manager = BatchForgerManager([
BatchForgerList(batch_size),
BatchForger(batch_size, (motion_model.reid_dim,)),
BatchForger(batch_size, (motion_model.roi_output_dim, motion_model.pool_size, motion_model.pool_size)),
BatchForger(batch_size, (motion_model.representation_dim,)),
BatchForger(batch_size, (4,)),
BatchForger(batch_size, (4,)),
BatchForger(batch_size, ()),
BatchForger(batch_size, (4,))
])
max_epochs = 100
log_freq = 25
train_pred_loss_epochs = []
train_vis_loss_epochs = []
val_pred_loss_epochs = []
val_vis_loss_epochs = []
lowest_val_loss = 9999999.9
lowest_val_loss_epoch = -1
############
# Training #
############
for epoch in range(max_epochs):
n_iter = 0
train_pred_loss_iters = []
train_vis_loss_iters = []
val_pred_loss_iters = []
val_vis_loss_iters = []
for data in train_loader:
historical_reid = get_batch_reid_features(reid_network, data['imgs'], data['historical'])
curr_reid = get_curr_reid_features(reid_network, data['imgs'], data['curr_frame_offset'], data['curr_gt_app'])
conv_features, repr_features = get_features(obj_detect, data['imgs'], data['curr_frame_offset'], data['curr_gt_app'])
prev_loc = (data['prev_gt_warped'] if use_ecc else data['prev_gt']).cuda()
curr_loc = (data['curr_gt_warped'] if use_ecc else data['curr_gt']).cuda()
curr_vis = data['curr_vis'].cuda()
label_loc = data['label_gt'].cuda()
batch_manager.feed((historical_reid, curr_reid, conv_features, repr_features, prev_loc, curr_loc, curr_vis, label_loc))
while batch_manager.has_one_batch():
n_iter += 1
historical_reid, curr_reid, conv_features, repr_features, prev_loc, curr_loc, curr_vis, label_loc = \
batch_manager.dump()
optimizer.zero_grad()
pred_loc_wh, vis = motion_model(historical_reid, curr_reid, conv_features, repr_features, prev_loc, curr_loc)
label_loc_wh = two_p_to_wh(label_loc)
pred_loss = pred_loss_func(pred_loc_wh, label_loc_wh)
vis_loss = vis_loss_func(vis, curr_vis)
if no_vis_loss:
loss = pred_loss
else:
loss = pred_loss + vis_loss_ratio * vis_loss
loss.backward()
optimizer.step()
train_pred_loss_iters.append(pred_loss.item())
train_vis_loss_iters.append(vis_loss.item())
if n_iter % log_freq == 0:
print('[Train Iter %5d] train pred loss %.6f, vis loss %.6f ...' %
(n_iter, np.mean(train_pred_loss_iters[n_iter-log_freq:n_iter]),
np.mean(train_vis_loss_iters[n_iter-log_freq:n_iter])), flush=True)
mean_train_pred_loss = np.mean(train_pred_loss_iters)
mean_train_vis_loss = np.mean(train_vis_loss_iters)
train_pred_loss_epochs.append(mean_train_pred_loss)
train_vis_loss_epochs.append(mean_train_vis_loss)
print('Train epoch %4d end.' % (epoch + 1))
batch_manager.reset()
motion_model.eval()
with torch.no_grad():
for data in val_loader:
historical_reid = get_batch_reid_features(reid_network, data['imgs'], data['historical'])
curr_reid = get_curr_reid_features(reid_network, data['imgs'], data['curr_frame_offset'], data['curr_gt_app'])
conv_features, repr_features = get_features(obj_detect, data['imgs'], data['curr_frame_offset'], data['curr_gt_app'])
prev_loc = (data['prev_gt_warped'] if use_ecc else data['prev_gt']).cuda()
curr_loc = (data['curr_gt_warped'] if use_ecc else data['curr_gt']).cuda()
curr_vis = data['curr_vis'].cuda()
label_loc = data['label_gt'].cuda()
batch_manager.feed((historical_reid, curr_reid, conv_features, repr_features, prev_loc, curr_loc, curr_vis, label_loc))
while batch_manager.has_one_batch():
historical_reid, curr_reid, conv_features, repr_features, prev_loc, curr_loc, curr_vis, label_loc = \
batch_manager.dump()
pred_loc_wh, vis = motion_model(historical_reid, curr_reid, conv_features, repr_features, prev_loc, curr_loc)
label_loc_wh = two_p_to_wh(label_loc)
pred_loss = pred_loss_func(pred_loc_wh, label_loc_wh)
vis_loss = vis_loss_func(vis, curr_vis)
val_pred_loss_iters.append(pred_loss.item())
val_vis_loss_iters.append(vis_loss.item())
mean_val_pred_loss = np.mean(val_pred_loss_iters)
mean_val_vis_loss = np.mean(val_vis_loss_iters)
val_pred_loss_epochs.append(mean_val_pred_loss)
val_vis_loss_epochs.append(mean_val_vis_loss)
print('[Epoch %4d] train pred loss %.6f, vis loss %.6f; val pred loss %.6f, vis loss %.6f' %
(epoch+1, mean_train_pred_loss, mean_train_vis_loss, mean_val_pred_loss, mean_val_vis_loss), flush=True)
with open(log_file, 'a') as f:
f.write('Epoch %4d: train pred loss %.6f, vis loss %.6f; val pred loss %.6f, vis loss %.6f\n' %
(epoch+1, mean_train_pred_loss, mean_train_vis_loss, mean_val_pred_loss, mean_val_vis_loss))
batch_manager.reset()
motion_model.train()
if mean_val_pred_loss < lowest_val_loss:
lowest_val_loss, lowest_val_loss_epoch = mean_val_pred_loss, epoch + 1
torch.save(motion_model.state_dict(), osp.join(output_dir, 'reid_motion_model_epoch_%d.pth'%(epoch+1)))
def train_main(sgd, lr, weight_decay, batch_size, output_dir, ex_name):
random.seed(12345)
torch.manual_seed(12345)
torch.cuda.manual_seed(12345)
np.random.seed(12345)
torch.backends.cudnn.deterministic = True
output_dir = osp.join(output_dir, ex_name)
log_file = osp.join(output_dir, 'epoch_log.txt')
if not osp.exists(output_dir):
os.makedirs(output_dir)
with open(log_file, 'w') as f:
f.write('[Experiment name]%s\n\n' % ex_name)
f.write('[Parameters]\n')
f.write('lr=%f\nweight_decay=%f\nbatch_size=%d\n\n' %
(lr, weight_decay, batch_size))
f.write('[Loss log]\n')
with open('experiments/cfgs/tracktor.yaml', 'r') as f:
tracker_config = yaml.safe_load(f)
#################
# Load Datasets #
#################
train_set = MOT17SimpleReIDWrapper('train', 0.8, 0.0, 1, train_random_sample=False, ecc=False)
train_loader = DataLoader(train_set, batch_size=batch_size, shuffle=True, num_workers=2, collate_fn=simple_reid_wrapper_collate)
val_set = MOT17SimpleReIDWrapper('val', 0.8, 0.0, 1, train_random_sample=False, ecc=False)
val_loader = DataLoader(val_set, batch_size=batch_size, shuffle=False, num_workers=2, collate_fn=simple_reid_wrapper_collate)
########################
# Initializing Modules #
########################
obj_detect = FRCNN_FPN(num_classes=2)
obj_detect.load_state_dict(torch.load(tracker_config['tracktor']['obj_detect_model'],
map_location=lambda storage, loc: storage))
obj_detect.eval()
obj_detect.cuda()
reid_network = resnet50(pretrained=False, output_dim=128)
reid_network.load_state_dict(torch.load(tracker_config['tracktor']['reid_weights'],
map_location=lambda storage, loc: storage))
reid_network.eval()
reid_network.cuda()
vis_model = VisSimpleReID()
vis_model.train()
vis_model.cuda()
if sgd:
optimizer = torch.optim.SGD(vis_model.parameters(), lr=lr, weight_decay=weight_decay, momentum=0.9)
else:
optimizer = torch.optim.Adam(vis_model.parameters(), lr=lr, weight_decay=weight_decay)
loss_func = nn.MSELoss()
#######################
# Training Parameters #
#######################
max_epochs = 100
log_freq = 25
lowest_val_loss = 9999999.9
lowest_val_loss_epoch = -1
############
# Training #
############
for epoch in range(max_epochs):
n_iter = 0
new_lowest_flag = False
train_loss_iters = []
val_loss_iters = []
for data in train_loader:
with torch.no_grad():
early_reid = get_batch_mean_early_reid(reid_network, data['early_reid_patches'])
curr_reid = reid_network(data['curr_reid_patch'].cuda())
conv_features, repr_features = get_features(obj_detect, data['curr_img'], data['curr_gt_app'])
curr_vis = data['curr_vis'].cuda()
n_iter += 1
optimizer.zero_grad()
vis = vis_model(early_reid, curr_reid, conv_features, repr_features)
loss = loss_func(vis, curr_vis)
loss.backward()
optimizer.step()
train_loss_iters.append(loss.item())
if n_iter % log_freq == 0:
print('[Train Iter %5d] train loss %.6f ...' %
(n_iter, np.mean(train_loss_iters[n_iter-log_freq:n_iter])),
flush=True)
mean_train_loss = np.mean(train_loss_iters)
print('Train epoch %4d end.' % (epoch + 1))
vis_model.eval()
with torch.no_grad():
for data in val_loader:
early_reid = get_batch_mean_early_reid(reid_network, data['early_reid_patches'])
curr_reid = reid_network(data['curr_reid_patch'].cuda())
conv_features, repr_features = get_features(obj_detect, data['curr_img'], data['curr_gt_app'])
curr_vis = data['curr_vis'].cuda()
vis = vis_model(early_reid, curr_reid, conv_features, repr_features)
loss = loss_func(vis, curr_vis)
val_loss_iters.append(loss.item())
mean_val_loss = np.mean(val_loss_iters)
print('[Epoch %4d] train loss %.6f, val loss %.6f' %
(epoch+1, mean_train_loss, mean_val_loss))
vis_model.train()
if mean_val_loss < lowest_val_loss:
lowest_val_loss, lowest_val_loss_epoch = mean_val_loss, epoch + 1
new_lowest_flag = True
torch.save(vis_model.state_dict(), osp.join(output_dir, 'vis_model_epoch_%d.pth'%(epoch+1)))
with open(log_file, 'a') as f:
f.write('[Epoch %4d] train loss %.6f, val loss %.6f %s\n' %
(epoch+1, mean_train_loss, mean_val_loss, '*' if new_lowest_flag else ''))
def test_tracktor_motion(val_loader, tracker_config, bbox_regression=True):
obj_detect = FRCNN_FPN(num_classes=2)
obj_detect.load_state_dict(
torch.load(tracker_config['tracktor']['obj_detect_model'],
map_location=lambda storage, loc: storage))
obj_detect.eval()
obj_detect.cuda()
pred_loss_func = nn.SmoothL1Loss()
loss_iters = []
low_vis_loss_sum = 0.0
low_vis_num = 0
high_vis_loss_sum = 0.0
high_vis_num = 0
total_iters = len(val_loader)
n_iters = 0
print(total_iters)
with torch.no_grad():
for data in val_loader:
n_iters += 1
prev_loc = data['prev_gt_warped'].cuda()
curr_loc = data['curr_gt_warped'].cuda()
label_loc = data['label_gt'].cuda()
curr_vis = data['curr_vis'].cuda()
pred_loc = curr_loc.clone()
last_motion = curr_loc - prev_loc
pred_loc += last_motion
if bbox_regression:
obj_detect.load_image(data['label_img'][0])
pred_loc, _ = obj_detect.predict_boxes(pred_loc)
label_loc_wh = two_p_to_wh(label_loc)
pred_loc_wh = two_p_to_wh(pred_loc)
pred_loss = pred_loss_func(pred_loc_wh, label_loc_wh)
loss_iters.append(pred_loss.item())
low_vis_ind = curr_vis < 0.3
if low_vis_ind.any():
low_vis_pred_loss = pred_loss_func(pred_loc_wh[low_vis_ind],
label_loc_wh[low_vis_ind])
low_vis_loss_sum += (low_vis_pred_loss *
torch.sum(low_vis_ind)).item()
low_vis_num += torch.sum(low_vis_ind).item()
high_vis_ind = curr_vis > 0.7
if high_vis_ind.any():
high_vis_pred_loss = pred_loss_func(pred_loc_wh[high_vis_ind],
label_loc_wh[high_vis_ind])
high_vis_loss_sum += (high_vis_pred_loss *
torch.sum(high_vis_ind)).item()
high_vis_num += torch.sum(high_vis_ind).item()
if n_iters % 500 == 0:
print('Iter %5d/%5d finished.' % (n_iters, total_iters),
flush=True)
mean_loss = np.mean(loss_iters)
mean_low_vis_loss = low_vis_loss_sum / low_vis_num
mean_high_vis_loss = high_vis_loss_sum / high_vis_num
print('All finished! Loss %.6f, low vis loss %.6f, high vis loss %.6f.' %
(mean_loss, mean_low_vis_loss, mean_high_vis_loss))
def train_main(oracle_training, no_visrepr, max_previous_frame, use_ecc,
use_modulator, use_bn, use_residual, vis_loss_ratio,
no_vis_loss, lr, weight_decay, batch_size, output_dir,
pretrain_vis_path, ex_name):
random.seed(12345)
torch.manual_seed(12345)
torch.cuda.manual_seed(12345)
np.random.seed(12345)
torch.backends.cudnn.deterministic = True
output_dir = osp.join(output_dir, ex_name)
log_file = osp.join(output_dir, 'epoch_log.txt')
if not osp.exists(output_dir):
os.makedirs(output_dir)
with open(log_file, 'w') as f:
f.write('[Experiment name]%s\n\n' % ex_name)
f.write('[Parameters]\n')
f.write(
'oracle_training=%r\nno_visrepr=%r\nmax_previous_frame=%d\nuse_ecc=%r\nuse_modulator=%r\nuse_bn=%r\nuse_residual=%r\nvis_loss_ratio=%f\nno_vis_loss=%r\nlr=%f\nweight_decay=%f\nbatch_size=%d\n\n'
% (oracle_training, no_visrepr, max_previous_frame, use_ecc,
use_modulator, use_bn, use_residual, vis_loss_ratio,
no_vis_loss, lr, weight_decay, batch_size))
f.write('[Loss log]\n')
with open('experiments/cfgs/tracktor.yaml', 'r') as f:
tracker_config = yaml.safe_load(f)
#################
# Load Datasets #
#################
train_set = MOT17TracksWrapper('train',
0.8,
0.0,
input_track_len=max_previous_frame + 1,
max_sample_frame=max_previous_frame,
get_data_mode='sample' +
(',ecc' if use_ecc else ''),
tracker_cfg=tracker_config)
train_loader = DataLoader(train_set,
batch_size=batch_size,
shuffle=True,
num_workers=1,
collate_fn=tracks_wrapper_collate)
val_set = MOT17TracksWrapper('val',
0.8,
0.1,
input_track_len=max_previous_frame + 1,
max_sample_frame=max_previous_frame,
get_data_mode='sample' +
(',ecc' if use_ecc else ''),
tracker_cfg=tracker_config)
val_loader = DataLoader(val_set,
batch_size=batch_size,
shuffle=False,
num_workers=1,
collate_fn=tracks_wrapper_collate)
with open(
osp.join(cfg.ROOT_DIR, 'output', 'precomputed_ecc_matrices_3.pkl'),
'rb') as f:
ecc_dict = pickle.load(f)
train_set.load_precomputed_ecc_warp_matrices(ecc_dict)
val_set.load_precomputed_ecc_warp_matrices(ecc_dict)
########################
# Initializing Modules #
########################
obj_detect = FRCNN_FPN(num_classes=2)
obj_detect.load_state_dict(
torch.load(tracker_config['tracktor']['obj_detect_model'],
map_location=lambda storage, loc: storage))
obj_detect.eval()
obj_detect.cuda()
if oracle_training:
motion_model = VisOracleMotionModel(vis_conv_only=False,
use_modulator=use_modulator)
else:
if no_visrepr:
motion_model = MotionModelNoVisRepr(vis_conv_only=False,
use_modulator=use_modulator,
use_bn=use_bn)
else:
motion_model = MotionModelV2(vis_conv_only=False,
use_modulator=use_modulator,
use_bn=use_bn,
use_residual=use_residual)
# motion_model.load_vis_pretrained(pretrain_vis_path)
motion_model.train()
motion_model.cuda()
optimizer = torch.optim.Adam(motion_model.parameters(),
lr=lr,
weight_decay=weight_decay)
pred_loss_func = nn.SmoothL1Loss()
vis_loss_func = nn.MSELoss()
#######################
# Training Parameters #
#######################
max_epochs = 100
log_freq = 25
train_pred_loss_epochs = []
train_vis_loss_epochs = []
val_pred_loss_epochs = []
val_vis_loss_epochs = []
lowest_val_loss = 9999999.9
lowest_val_loss_epoch = -1
############
# Training #
############
for epoch in range(max_epochs):
n_iter = 0
train_pred_loss_iters = []
train_vis_loss_iters = []
val_pred_loss_iters = []
val_vis_loss_iters = []
for data, label in train_loader:
# jitter bboxs for getting roi features
im_w = torch.tensor([img.size()[-1] for img in data['curr_img']],
dtype=data['curr_gt'].dtype)
im_h = torch.tensor([img.size()[-2] for img in data['curr_img']],
dtype=data['curr_gt'].dtype)
jittered_curr_gt = bbox_jitter(data['curr_gt'].clone(), im_w, im_h)
conv_features, repr_features = get_features(
obj_detect, data['curr_img'], jittered_curr_gt)
# for motion calculation, we still use the unjittered bboxs
prev_loc = (data['prev_gt_warped']
if use_ecc else data['prev_gt']).cuda()
curr_loc = (data['curr_gt_warped']
if use_ecc else data['curr_gt']).cuda()
label_loc = label['label_gt'].cuda()
curr_vis = data['curr_vis'].cuda()
n_iter += 1
# TODO the output bbox should be (x,y,w,h)?
optimizer.zero_grad()
if oracle_training:
pred_loc_wh, vis = motion_model(conv_features, repr_features,
prev_loc, curr_loc,
curr_vis.unsqueeze(-1))
else:
pred_loc_wh, vis = motion_model(conv_features, repr_features,
prev_loc, curr_loc)
label_loc_wh = two_p_to_wh(label_loc)
pred_loss = pred_loss_func(pred_loc_wh, label_loc_wh)
vis_loss = vis_loss_func(vis, curr_vis)
if no_vis_loss:
loss = pred_loss
else:
loss = pred_loss + vis_loss_ratio * vis_loss
loss.backward()
optimizer.step()
train_pred_loss_iters.append(pred_loss.item())
train_vis_loss_iters.append(vis_loss.item())
if n_iter % log_freq == 0:
print(
'[Train Iter %5d] train pred loss %.6f, vis loss %.6f ...'
%
(n_iter,
np.mean(train_pred_loss_iters[n_iter - log_freq:n_iter]),
np.mean(train_vis_loss_iters[n_iter - log_freq:n_iter])),
flush=True)
mean_train_pred_loss = np.mean(train_pred_loss_iters)
mean_train_vis_loss = np.mean(train_vis_loss_iters)
train_pred_loss_epochs.append(mean_train_pred_loss)
train_vis_loss_epochs.append(mean_train_vis_loss)
print('Train epoch %4d end.' % (epoch + 1))
motion_model.eval()
with torch.no_grad():
for data, label in val_loader:
# do not jitter for validation
conv_features, repr_features = get_features(
obj_detect, data['curr_img'], data['curr_gt'])
prev_loc = (data['prev_gt_warped']
if use_ecc else data['prev_gt']).cuda()
curr_loc = (data['curr_gt_warped']
if use_ecc else data['curr_gt']).cuda()
label_loc = label['label_gt'].cuda()
curr_vis = data['curr_vis'].cuda()
if oracle_training:
pred_loc_wh, vis = motion_model(conv_features,
repr_features, prev_loc,
curr_loc,
curr_vis.unsqueeze(-1))
else:
pred_loc_wh, vis = motion_model(conv_features,
repr_features, prev_loc,
curr_loc)
label_loc_wh = two_p_to_wh(label_loc)
pred_loss = pred_loss_func(pred_loc_wh, label_loc_wh)
vis_loss = vis_loss_func(vis, curr_vis)
val_pred_loss_iters.append(pred_loss.item())
val_vis_loss_iters.append(vis_loss.item())
mean_val_pred_loss = np.mean(val_pred_loss_iters)
mean_val_vis_loss = np.mean(val_vis_loss_iters)
val_pred_loss_epochs.append(mean_val_pred_loss)
val_vis_loss_epochs.append(mean_val_vis_loss)
print(
'[Epoch %4d] train pred loss %.6f, vis loss %.6f; val pred loss %.6f, vis loss %.6f'
% (epoch + 1, mean_train_pred_loss, mean_train_vis_loss,
mean_val_pred_loss, mean_val_vis_loss))
with open(log_file, 'a') as f:
f.write(
'Epoch %4d: train pred loss %.6f, vis loss %.6f; val pred loss %.6f, vis loss %.6f\n'
% (epoch + 1, mean_train_pred_loss, mean_train_vis_loss,
mean_val_pred_loss, mean_val_vis_loss))
motion_model.train()
if mean_val_pred_loss < lowest_val_loss:
lowest_val_loss, lowest_val_loss_epoch = mean_val_pred_loss, epoch + 1
torch.save(
motion_model.state_dict(),
osp.join(output_dir,
'motion_model_epoch_%d.pth' % (epoch + 1)))
def train_main(v2, use_refine_model, use_ecc, use_modulator, use_bn, use_residual, vis_roi_features, no_visrepr, vis_loss_ratio, no_vis_loss,
modulate_from_vis, max_sample_frame, lr, weight_decay, batch_size, output_dir, ex_name):
random.seed(12345)
torch.manual_seed(12345)
torch.cuda.manual_seed(12345)
np.random.seed(12345)
torch.backends.cudnn.deterministic = True
output_dir = osp.join(output_dir, ex_name)
log_file = osp.join(output_dir, 'epoch_log.txt')
if not osp.exists(output_dir):
os.makedirs(output_dir)
with open(log_file, 'w') as f:
f.write('[Experiment name]%s\n\n' % ex_name)
f.write('[Parameters]\n')
f.write('use_ecc=%r\nuse_modulator=%r\nuse_bn=%r\nuse_residual=%r\nvis_roi_features=%r\nno_visrepr=%r\nvis_loss_ratio=%f\nno_vis_loss=%r\nmodulate_from_vis=%r\nmax_sample_frame=%d\nlr=%f\nweight_decay=%f\nbatch_size=%d\n\n' %
(use_ecc, use_modulator, use_bn, use_residual, vis_roi_features, no_visrepr, vis_loss_ratio, no_vis_loss, modulate_from_vis, max_sample_frame, lr, weight_decay, batch_size))
f.write('[Loss log]\n')
with open('experiments/cfgs/tracktor.yaml', 'r') as f:
tracker_config = yaml.safe_load(f)
#################
# Load Datasets #
#################
train_set = MOT17SimpleReIDWrapper('train', 0.8, 0.0, max_sample_frame, tracker_cfg=tracker_config)
train_loader = DataLoader(train_set, batch_size=batch_size, shuffle=True, num_workers=1, collate_fn=simple_reid_wrapper_collate)
val_set = MOT17SimpleReIDWrapper('val', 0.8, 0.0, max_sample_frame, tracker_cfg=tracker_config)
val_loader = DataLoader(val_set, batch_size=batch_size, shuffle=False, num_workers=1, collate_fn=simple_reid_wrapper_collate)
with open(osp.join(cfg.ROOT_DIR, 'output', 'precomputed_ecc_matrices_3.pkl'), 'rb') as f:
ecc_dict = pickle.load(f)
train_set.load_precomputed_ecc_warp_matrices(ecc_dict)
val_set.load_precomputed_ecc_warp_matrices(ecc_dict)
########################
# Initializing Modules #
########################
obj_detect = FRCNN_FPN(num_classes=2)
obj_detect.load_state_dict(torch.load(tracker_config['tracktor']['obj_detect_model'],
map_location=lambda storage, loc: storage))
obj_detect.eval()
obj_detect.cuda()
if v2:
motion_model = MotionModelSimpleReIDV2(use_modulator=use_modulator, use_bn=use_bn, use_residual=use_residual,
vis_roi_features=vis_roi_features, no_visrepr=no_visrepr, modulate_from_vis=modulate_from_vis)
else:
motion_model = MotionModelSimpleReID(use_modulator=use_modulator, use_bn=use_bn, use_residual=use_residual,
vis_roi_features=vis_roi_features, no_visrepr=no_visrepr, modulate_from_vis=modulate_from_vis)
motion_model.train()
motion_model.cuda()
if use_refine_model:
motion_model = RefineModel(motion_model)
motion_model.train()
motion_model.cuda()
reid_network = resnet50(pretrained=False, output_dim=128)
reid_network.load_state_dict(torch.load(tracker_config['tracktor']['reid_weights'],
map_location=lambda storage, loc: storage))
reid_network.eval()
reid_network.cuda()
optimizer = torch.optim.Adam(motion_model.parameters(), lr=lr, weight_decay=weight_decay)
pred_loss_func = nn.SmoothL1Loss()
vis_loss_func = nn.MSELoss()
#######################
# Training Parameters #
#######################
max_epochs = 100
log_freq = 25
train_pred_loss_epochs = []
train_vis_loss_epochs = []
val_pred_loss_epochs = []
val_vis_loss_epochs = []
lowest_val_loss = 9999999.9
lowest_val_loss_epoch = -1
last_save_epoch = 0
save_epoch_freq = 5
############
# Training #
############
for epoch in range(max_epochs):
n_iter = 0
new_lowest_flag = False
train_pred_loss_iters = []
train_vis_loss_iters = []
val_pred_loss_iters = []
val_vis_loss_iters = []
for data in train_loader:
early_reid = get_batch_mean_early_reid(reid_network, data['early_reid_patches'])
curr_reid = reid_network(data['curr_reid_patch'].cuda())
conv_features, repr_features = get_features(obj_detect, data['curr_img'], data['curr_gt_app'])
prev_loc = (data['prev_gt_warped'] if use_ecc else data['prev_gt']).cuda()
curr_loc = (data['curr_gt_warped'] if use_ecc else data['curr_gt']).cuda()
label_loc = data['label_gt'].cuda()
curr_vis = data['curr_vis'].cuda()
n_iter += 1
optimizer.zero_grad()
if use_refine_model:
pred_loc_wh, vis = motion_model(obj_detect, data['label_img'], conv_features, repr_features, prev_loc, curr_loc,
early_reid=early_reid, curr_reid=curr_reid)
label_loc_wh = two_p_to_wh(label_loc)
pred_loss = weighted_smooth_l1_loss(pred_loc_wh, label_loc_wh, curr_vis)
vis_loss = vis_loss_func(vis, curr_vis)
else:
if v2:
pred_loc_wh, vis = motion_model(early_reid, curr_reid, repr_features, prev_loc, curr_loc)
else:
pred_loc_wh, vis = motion_model(early_reid, curr_reid, conv_features, repr_features, prev_loc, curr_loc)
label_loc_wh = two_p_to_wh(label_loc)
pred_loss = pred_loss_func(pred_loc_wh, label_loc_wh)
vis_loss = vis_loss_func(vis, curr_vis)
if no_vis_loss:
loss = pred_loss
else:
loss = pred_loss + vis_loss_ratio * vis_loss
loss.backward()
optimizer.step()
train_pred_loss_iters.append(pred_loss.item())
train_vis_loss_iters.append(vis_loss.item())
if n_iter % log_freq == 0:
print('[Train Iter %5d] train pred loss %.6f, vis loss %.6f ...' %
(n_iter, np.mean(train_pred_loss_iters[n_iter-log_freq:n_iter]), np.mean(train_vis_loss_iters[n_iter-log_freq:n_iter])),
flush=True)
mean_train_pred_loss = np.mean(train_pred_loss_iters)
mean_train_vis_loss = np.mean(train_vis_loss_iters)
train_pred_loss_epochs.append(mean_train_pred_loss)
train_vis_loss_epochs.append(mean_train_vis_loss)
print('Train epoch %4d end.' % (epoch + 1))
motion_model.eval()
with torch.no_grad():
for data in val_loader:
early_reid = get_batch_mean_early_reid(reid_network, data['early_reid_patches'])
curr_reid = reid_network(data['curr_reid_patch'].cuda())
conv_features, repr_features = get_features(obj_detect, data['curr_img'], data['curr_gt_app'])
prev_loc = (data['prev_gt_warped'] if use_ecc else data['prev_gt']).cuda()
curr_loc = (data['curr_gt_warped'] if use_ecc else data['curr_gt']).cuda()
label_loc = data['label_gt'].cuda()
curr_vis = data['curr_vis'].cuda()
if use_refine_model:
pred_loc_wh, vis = motion_model(obj_detect, data['label_img'], conv_features, repr_features, prev_loc, curr_loc,
early_reid=early_reid, curr_reid=curr_reid)
label_loc_wh = two_p_to_wh(label_loc)
pred_loss = pred_loss_func(pred_loc_wh, label_loc_wh)
vis_loss = vis_loss_func(vis, curr_vis)
else:
if v2:
pred_loc_wh, vis = motion_model(early_reid, curr_reid, repr_features, prev_loc, curr_loc)
else:
pred_loc_wh, vis = motion_model(early_reid, curr_reid, conv_features, repr_features, prev_loc, curr_loc)
label_loc_wh = two_p_to_wh(label_loc)
pred_loss = pred_loss_func(pred_loc_wh, label_loc_wh)
vis_loss = vis_loss_func(vis, curr_vis)
val_pred_loss_iters.append(pred_loss.item())
val_vis_loss_iters.append(vis_loss.item())
mean_val_pred_loss = np.mean(val_pred_loss_iters)
mean_val_vis_loss = np.mean(val_vis_loss_iters)
val_pred_loss_epochs.append(mean_val_pred_loss)
val_vis_loss_epochs.append(mean_val_vis_loss)
print('[Epoch %4d] train pred loss %.6f, vis loss %.6f; val pred loss %.6f, vis loss %.6f' %
(epoch+1, mean_train_pred_loss, mean_train_vis_loss, mean_val_pred_loss, mean_val_vis_loss))
motion_model.train()
if mean_val_pred_loss < lowest_val_loss:
lowest_val_loss, lowest_val_loss_epoch = mean_val_pred_loss, epoch + 1
last_save_epoch = lowest_val_loss_epoch
new_lowest_flag = True
torch.save(motion_model.state_dict(), osp.join(output_dir, 'simple_reid_motion_model_epoch_%d.pth'%(epoch+1)))
elif epoch + 1 - last_save_epoch == save_epoch_freq:
last_save_epoch = epoch + 1
torch.save(motion_model.state_dict(), osp.join(output_dir, 'simple_reid_motion_model_epoch_%d.pth'%(epoch+1)))
with open(log_file, 'a') as f:
f.write('Epoch %4d: train pred loss %.6f, vis loss %.6f; val pred loss %.6f, vis loss %.6f %s\n' %
(epoch+1, mean_train_pred_loss, mean_train_vis_loss, mean_val_pred_loss, mean_val_vis_loss, '*' if new_lowest_flag else ''))
def pretrain_main(conv_only, image_flip, train_jitter, val_jitter, lr,
weight_decay, batch_size, output_dir, ex_name):
random.seed(12345)
torch.manual_seed(12345)
torch.cuda.manual_seed(12345)
np.random.seed(12345)
torch.backends.cudnn.deterministic = True
# output_dir = osp.join(get_output_dir('motion'), 'vis')
output_dir = osp.join(output_dir, ex_name)
log_file = osp.join(output_dir, 'epoch_log.txt')
if not osp.exists(output_dir):
os.makedirs(output_dir)
with open(log_file, 'w') as f:
f.write('[Experiment name]%s\n\n' % ex_name)
f.write('[Parameters]\n')
f.write(
'conv_only=%r\nimage_flip=%r\ntrain_jitter=%r\nval_jitter=%r\nlr=%f\nweight_decay=%f\nbatch_size=%d\n\n'
% (conv_only, image_flip, train_jitter, val_jitter, lr,
weight_decay, batch_size))
f.write('[Loss log]\n')
with open('experiments/cfgs/tracktor.yaml', 'r') as f:
tracker_config = yaml.safe_load(f)
#################
# Load Datasets #
#################
train_set = MOT17Vis('train',
0.8,
0.0,
train_bbox_jitter=train_jitter,
random_image_flip=image_flip)
train_loader = DataLoader(train_set,
batch_size=1,
shuffle=True,
num_workers=2)
val_set = MOT17Vis('val', 0.8, 0.0, val_bbox_jitter=val_jitter)
val_loader = DataLoader(val_set,
batch_size=1,
shuffle=False,
num_workers=2)
########################
# Initializing Modules #
########################
obj_detect = FRCNN_FPN(num_classes=2)
obj_detect.load_state_dict(
torch.load(tracker_config['tracktor']['obj_detect_model'],
map_location=lambda storage, loc: storage))
obj_detect.eval()
obj_detect.cuda()
vis_model = VisEst(conv_only=conv_only)
vis_model.train()
vis_model.cuda()
optimizer = torch.optim.Adam(vis_model.parameters(),
lr=lr,
weight_decay=weight_decay)
loss_func = nn.MSELoss()
#######################
# Training Parameters #
#######################
max_epochs = 100
conv_batch_forger = BatchForger(
batch_size,
(vis_model.output_dim, vis_model.pool_size, vis_model.pool_size))
repr_batch_forger = BatchForger(batch_size,
(vis_model.representation_dim, ))
label_batch_forger = BatchForger(batch_size, (1, ))
log_freq = 100
train_loss_epochs = []
val_loss_epochs = []
lowest_val_loss = 9999999.9
lowest_val_loss_epoch = -1
############
# Training #
############
for epoch in range(max_epochs):
n_iter = 0
train_loss_iters = []
val_loss_iters = []
for data in train_loader:
conv_features, repr_features = get_features(
obj_detect, data['img'], data['gt'])
conv_batch_forger.feed(conv_features)
repr_batch_forger.feed(repr_features)
label_batch_forger.feed(data['vis'].squeeze(0).unsqueeze(-1))
while label_batch_forger.has_one_batch():
n_iter += 1
batch_conv = conv_batch_forger.dump().cuda()
batch_repr = repr_batch_forger.dump().cuda()
batch_label = label_batch_forger.dump().cuda()
optimizer.zero_grad()
pred, _ = vis_model(batch_conv, batch_repr)
loss = loss_func(pred, batch_label)
loss.backward()
optimizer.step()
train_loss_iters.append(loss.item())
if n_iter % log_freq == 0:
print(
'[Train Iter %5d] train loss %.6f ...' %
(n_iter,
np.mean(train_loss_iters[n_iter - log_freq:n_iter])))
mean_train_loss = np.mean(train_loss_iters)
train_loss_epochs.append(mean_train_loss)
print('Train epoch %4d end.' % (epoch + 1))
conv_batch_forger.reset()
repr_batch_forger.reset()
label_batch_forger.reset()
vis_model.eval()
for data in val_loader:
conv_features, repr_features = get_features(
obj_detect, data['img'], data['gt'])
conv_batch_forger.feed(conv_features)
repr_batch_forger.feed(repr_features)
label_batch_forger.feed(data['vis'].squeeze(0).unsqueeze(-1))
while label_batch_forger.has_one_batch():
batch_conv = conv_batch_forger.dump().cuda()
batch_repr = repr_batch_forger.dump().cuda()
batch_label = label_batch_forger.dump().cuda()
pred, _ = vis_model(batch_conv, batch_repr)
loss = loss_func(pred, batch_label)
val_loss_iters.append(loss.item())
mean_val_loss = np.mean(val_loss_iters)
val_loss_epochs.append(mean_val_loss)
print('[Epoch %4d] train loss %.6f, val loss %.6f' %
(epoch + 1, mean_train_loss, mean_val_loss))
with open(log_file, 'a') as f:
f.write('Epoch %4d: train loss %.6f, val loss %.6f\n' %
(epoch + 1, mean_train_loss, mean_val_loss))
conv_batch_forger.reset()
repr_batch_forger.reset()
label_batch_forger.reset()
vis_model.train()
if mean_val_loss < lowest_val_loss:
lowest_val_loss, lowest_val_loss_epoch = mean_val_loss, epoch + 1
torch.save(
vis_model.state_dict(),
osp.join(output_dir, 'vis_model_epoch_%d.pth' % (epoch + 1)))
def main(tracktor, reid, _config, _log, _run):
sacred.commands.print_config(_run)
# set all seeds
torch.manual_seed(tracktor['seed'])
torch.cuda.manual_seed(tracktor['seed'])
np.random.seed(tracktor['seed'])
torch.backends.cudnn.deterministic = True
output_dir = osp.join(get_output_dir(tracktor['module_name']), tracktor['name'])
sacred_config = osp.join(output_dir, 'sacred_config.yaml')
if not osp.exists(output_dir):
os.makedirs(output_dir)
with open(sacred_config, 'w') as outfile:
yaml.dump(_config, outfile, default_flow_style=False)
##########################
# Initialize the modules #
##########################
# object detection
_log.info("Initializing object detector.")
obj_detect = FRCNN_FPN(num_classes=2)
obj_detect.load_state_dict(torch.load(_config['tracktor']['obj_detect_model'],
map_location=lambda storage, loc: storage))
obj_detect.eval()
obj_detect.cuda()
# reid
reid_network = resnet50(pretrained=False, **reid['cnn'])
reid_network.load_state_dict(torch.load(tracktor['reid_weights'],
map_location=lambda storage, loc: storage))
reid_network.eval()
reid_network.cuda()
# neural motion model
vis_model = VisSimpleReID()
motion_model = MotionModelV3(vis_model)
motion_model.load_state_dict(torch.load('output/motion/finetune_motion_model_v3.pth'))
motion_model.eval()
motion_model.cuda()
save_vis_results = False
# tracktor
if 'oracle' in tracktor:
tracker = OracleTracker(obj_detect, reid_network, tracktor['tracker'], tracktor['oracle'])
else:
# tracker = Tracker(obj_detect, reid_network, tracktor['tracker'])
tracker = TrackerNeuralMM(obj_detect, reid_network, motion_model, tracktor['tracker'], save_vis_results=save_vis_results, vis_model=None)
time_total = 0
num_frames = 0
mot_accums = []
dataset = Datasets(tracktor['dataset'], {'use_val_split':True})
for seq in dataset:
tracker.reset()
start = time.time()
_log.info(f"Tracking: {seq}")
data_loader = DataLoader(seq, batch_size=1, shuffle=False)
for i, frame in enumerate(tqdm(data_loader)):
if len(seq) * tracktor['frame_split'][0] <= i <= len(seq) * tracktor['frame_split'][1]:
with torch.no_grad():
tracker.step(frame)
num_frames += 1
results = tracker.get_results()
time_total += time.time() - start
_log.info(f"Tracks found: {len(results)}")
_log.info(f"Runtime for {seq}: {time.time() - start :.1f} s.")
if tracktor['interpolate']:
results = interpolate(results)
if seq.no_gt:
_log.info(f"No GT data for evaluation available.")
else:
mot_accums.append(get_mot_accum(results, seq))
_log.info(f"Writing predictions to: {output_dir}")
seq.write_results(results, output_dir)
if save_vis_results:
vis_results = tracker.get_vis_results()
seq.write_vis_results(vis_results, output_dir)
if tracktor['write_images']:
plot_sequence(results, seq, osp.join(output_dir, tracktor['dataset'], str(seq)))
_log.info(f"Tracking runtime for all sequences (without evaluation or image writing): "
f"{time_total:.1f} s ({num_frames / time_total:.1f} Hz)")
if mot_accums:
evaluate_mot_accums(mot_accums, [str(s) for s in dataset if not s.no_gt], generate_overall=True)
def main(tracktor, reid, _config, _log, _run):
sacred.commands.print_config(_run)
# set all seeds
torch.manual_seed(tracktor['seed'])
torch.cuda.manual_seed(tracktor['seed'])
np.random.seed(tracktor['seed'])
torch.backends.cudnn.deterministic = True
output_dir = osp.join(get_output_dir(tracktor['module_name']), tracktor['name'])
sacred_config = osp.join(output_dir, 'sacred_config.yaml')
if not osp.exists(output_dir):
os.makedirs(output_dir)
with open(sacred_config, 'w') as outfile:
yaml.dump(_config, outfile, default_flow_style=False)
##########################
# Initialize the modules #
##########################
# object detection
_log.info("Initializing object detector.")
use_masks = _config['tracktor']['tracker']['use_masks']
mask_model = Mask_RCNN(num_classes=2)
fast_model = FRCNN_FPN(num_classes=2)
fast_model.load_state_dict(torch.load(_config['tracktor']['fast_rcnn_model'],
map_location=lambda storage, loc: storage))
if(use_masks):
mask_model.load_state_dict(torch.load(_config['tracktor']['mask_rcnn_model'],
map_location=lambda storage, loc: storage)['model_state_dict'])
mask_model.eval()
mask_model.cuda()
fast_model.eval()
fast_model.cuda()
# reid
reid_network = resnet50(pretrained=False, **reid['cnn'])
reid_network.load_state_dict(torch.load(tracktor['reid_weights'],
map_location=lambda storage, loc: storage))
reid_network.eval()
reid_network.cuda()
# tracktor
if 'oracle' in tracktor:
tracker = OracleTracker(fast_model, reid_network, tracktor['tracker'], tracktor['oracle'])
else:
tracker = Tracker(fast_model, reid_network, tracktor['tracker'], mask_model)
time_total = 0
num_frames = 0
mot_accums = []
dataset = Datasets(tracktor['dataset'])
for seq in dataset:
num_frames = 0
tracker.reset()
start = time.time()
_log.info(f"Tracking: {seq}")
data_loader = DataLoader(seq, batch_size=1, shuffle=False)
if tracktor['write_images'] and use_masks:
print("[*] Plotting image to {}".format(osp.join(output_dir, tracktor['dataset'])))
for i, frame in enumerate(tqdm(data_loader)):
if len(seq) * tracktor['frame_split'][0] <= i <= len(seq) * tracktor['frame_split'][1]:
tracker.step(frame)
if tracktor['write_images'] and use_masks:
result = tracker.get_results()
masks = tracker.get_masks()
plot_sequence(result, masks, seq, num_frames, osp.join(output_dir, tracktor['dataset'], str(seq)), plot_masks = True)
num_frames += 1
results = tracker.get_results()
import matplotlib.pyplot as plt
time_total += time.time() - start
_log.info(f"Tracks found: {len(results)}")
_log.info(f"Runtime for {seq}: {time.time() - start :.1f} s.")
if tracktor['interpolate']:
results = interpolate(results)
if seq.no_gt:
_log.info(f"No GT data for evaluation available.")
else:
mot_accums.append(get_mot_accum(results, seq))
_log.info(f"Writing predictions to: {output_dir}")
seq.write_results(results, output_dir)
_log.info(f"Tracking runtime for all sequences (without evaluation or image writing): "
f"{time_total:.1f} s ({num_frames / time_total:.1f} Hz)")
if mot_accums:
evaluate_mot_accums(mot_accums, [str(s) for s in dataset if not s.no_gt], generate_overall=True)
