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 my_main(_config, reid):
# set all seeds
torch.manual_seed(reid['seed'])
torch.cuda.manual_seed(reid['seed'])
np.random.seed(reid['seed'])
torch.backends.cudnn.deterministic = True
print(_config)
output_dir = 'D:/emartins/reiddd'
tb_dir = 'D:/emartins/reiddd'
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 dataloader #
#########################
print("[*] Initializing Dataloader")
db_train = Datasets(reid['db_train'], reid['dataloader'])
db_train = DataLoader(db_train, batch_size=1, shuffle=True)
if reid['db_val']:
db_val = None
#db_val = DataLoader(db_val, batch_size=1, shuffle=True)
else:
db_val = None
##########################
# Initialize the modules #
##########################
print("[*] Building CNN")
network = resnet50(pretrained=False, **reid['cnn'])
network.train()
network.cuda()
##################
# Begin training #
##################
print("[*] Solving ...")
# build scheduling like in "In Defense of the Triplet Loss for Person Re-Identification"
# from Hermans et al.
lr = reid['solver']['optim_args']['lr']
iters_per_epoch = len(db_train)
# we want to keep lr until iter 15000 and from there to iter 25000 a exponential decay
l = eval("lambda epoch: 1 if epoch*{} < 15000 else 0.001**((epoch*{} - 15000)/(25000-15000))".format(
iters_per_epoch, iters_per_epoch))
#else:
# l = None
max_epochs = 25000 // len(db_train.dataset) + 1 if 25000 % len(db_train.dataset) else 25000 // len(db_train.dataset)
solver = Solver(output_dir, tb_dir, lr_scheduler_lambda=l)
solver.train(network, db_train, db_val, max_epochs, 100, model_args=reid['model_args'])
def build_tracker():
# reid
reid_network = resnet50(pretrained=False)
reid_network.load_state_dict(torch.load('/root/PycharmProjects/mmdet/tracktor/model/ResNet_iter_25245.pth',
map_location=lambda storage, loc: storage))
reid_network.eval()
reid_network.cuda()
f = open("/root/PycharmProjects/mmdet/tracktor/cfg/tracktor.yaml")
tracktor = yaml.load(f)['tracktor']
return Tracker(reid_network=reid_network, tracker_cfg=tracktor['tracker'])
# if __name__ == '__main__':
# a = build_tracker()
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 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 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 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 train_main(use_early_reid, use_reid_distance, mse_loss, 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(
'use_early_reid=%r\nuse_reid_distance=%r\nlr=%f\nweight_decay=%f\nbatch_size=%d\n\n'
%
(use_early_reid, use_reid_distance, 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 = MOT17VisReID('train', 0.8, 0.0)
train_loader = DataLoader(train_set,
batch_size=batch_size,
shuffle=True,
num_workers=2,
collate_fn=simple_reid_wrapper_collate)
val_set = MOT17VisReID('val', 0.8, 0.0)
val_loader = DataLoader(val_set,
batch_size=batch_size,
shuffle=False,
num_workers=2,
collate_fn=simple_reid_wrapper_collate)
########################
# Initializing Modules #
########################
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 = VisModel(use_early_reid=use_early_reid,
use_reid_distance=use_reid_distance)
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
lowest_val_weighed_loss = 9999999.9
lowest_val_weighted_loss_epoch = -1
for epoch in range(max_epochs):
n_iter = 0
new_lowest_flag = False
train_loss_iters = []
val_loss_iters = []
val_weighted_loss_iters = []
for data in train_loader:
curr_patch = data['patches'].cuda()
vis_gt = data['vis'].cuda()
if use_early_reid:
early_reid = get_batch_mean_early_reid(reid_network,
data['early_patches'])
curr_reid = reid_network(curr_patch)
else:
early_reid, curr_reid = None, None
n_iter += 1
optimizer.zero_grad()
vis = vis_model(curr_patch, early_reid, curr_reid)
loss = loss_func(vis, vis_gt) if mse_loss else weighted_mse_loss(
vis, vis_gt, vis_gt)
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:
curr_patch = data['patches'].cuda()
vis_gt = data['vis'].cuda()
if use_early_reid:
early_reid = get_batch_mean_early_reid(
reid_network, data['early_patches'])
curr_reid = reid_network(curr_patch)
else:
early_reid, curr_reid = None, None
vis = vis_model(curr_patch, early_reid, curr_reid)
loss = loss_func(vis, vis_gt)
weighted_loss = weighted_mse_loss(vis, vis_gt, vis_gt)
val_loss_iters.append(loss.item())
val_weighted_loss_iters.append(weighted_loss.item())
mean_val_loss = np.mean(val_loss_iters)
mean_val_weighted_loss = np.mean(val_weighted_loss_iters)
print(
'[Epoch %4d] train %s loss %.6f, val weighted loss %.6f, val mse loss %.6f'
% (epoch + 1, '' if mse_loss else 'weighted', mean_train_loss,
mean_val_weighted_loss, mean_val_loss))
vis_model.train()
if mean_val_weighted_loss < lowest_val_weighed_loss:
lowest_val_weighed_loss, lowest_val_weighted_loss_epoch = mean_val_weighted_loss, epoch + 1
lowest_val_loss = min(mean_val_loss, lowest_val_loss)
new_lowest_flag = True
torch.save(
vis_model.state_dict(),
osp.join(output_dir, 'vis_model_epoch_%d.pth' % (epoch + 1)))
elif mean_val_loss < lowest_val_loss:
lowest_val_loss, lowest_val_loss_epoch = mean_val_loss, epoch + 1
lowest_val_weighed_loss = min(mean_val_weighted_loss,
lowest_val_weighed_loss)
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 %s loss %.6f, val weighted loss %.6f, val mse loss %.6f %s\n'
% (epoch + 1, '' if mse_loss else 'weighted', mean_train_loss,
mean_val_weighted_loss, mean_val_loss,
'*' if new_lowest_flag else ''))
def main():
parser = argparse.ArgumentParser()
## Required parameters
parser.add_argument(
"--model",
default='helmet_tracking/model_final_13000.pth',
type=str, #required=True,
help="path to base D2 model or s3 location")
parser.add_argument(
"--reid_model",
default='helmet_tracking/ResNet_iter_25137.pth',
type=str, #required=True,
help="path to reid model or s3 location")
parser.add_argument(
"--output_dir",
default='/home/model_results',
type=str, #required=True,
help=
"The output directory where the model checkpoints and predictions will be written."
)
parser.add_argument("--bucket",
default='privisaa-bucket-virginia',
type=str)
parser.add_argument(
"--img_paths",
default='nfl-data/live_video',
type=str,
#required=True,
help="path to images or image location in s3")
parser.add_argument(
"--conf_thresh",
default=.5,
type=float, #required=True,
help="base D2 model")
parser.add_argument("--use_mask", default=0, type=int)
# parser.add_argument("--video", default=None, type=str, required=True, help="path to video for tracking job")
parser.add_argument(
"--d2_config",
default=
'/home/detectron2-ResNeSt/configs/quick_schedules/mask_rcnn_R_50_FPN_inference_acc_test.yaml',
type=str,
help="Detectron2 config file")
args = parser.parse_args()
#COCO-Detection/faster_cascade_rcnn_ResNeSt_101_FPN_syncbn_range-scale_1x.yaml'
confidence_threshold = .55
cfg = get_cfg()
config_file = args.d2_config
model_pth = args.model
reid_pth = args.reid_model
try:
os.mkdir('/home/video')
except:
pass
if os.path.exists(args.model) == False:
print('\n')
print('downloading d2 model from s3')
s3.download_file(args.bucket, args.model, '/home/d2_final.pth')
model_pth = '/home/d2_final.pth'
if os.path.exists(args.reid_model) == False:
print('\n')
print('downloading reid model from s3')
s3.download_file(args.bucket, args.reid_model, '/home/reid_final.pth')
reid_pth = '/home/reid_final.pth'
if os.path.exists(args.img_paths) == False:
objs = s3.list_objects(Bucket='privisaa-bucket-virginia',
Prefix=args.img_paths)['Contents']
keys = []
for key in objs:
if key['Size'] > 0:
keys.append(key['Key'])
folders = []
for key in keys:
folders.append(key.split('/')[-2])
folder = list(np.unique(folders))[-1]
print('\n')
print('Loading images from this video: ', folder)
for key in tqdm(keys):
if key.split('/')[-2] == folder:
s3.download_file(args.bucket, key,
f"/home/video/{key.split('/')[-1]}")
cfg.merge_from_file(config_file)
cfg.MODEL.RETINANET.SCORE_THRESH_TEST = args.conf_thresh
cfg.MODEL.WEIGHTS = model_pth #'/home/ubuntu/finetuned-fasterrcnn-cascade-d2-resnest-13000imgs-005lr-1class-tune2/model_final.pth'
#finetuned-fasterrcnn-cascade-d2-resnest-13000imgs-02lr/model_final.pth'
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = args.conf_thresh
cfg.MODEL.PANOPTIC_FPN.COMBINE.INSTANCES_CONFIDENCE_THRESH = args.conf_thresh
cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 256
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1
model = DefaultPredictor(cfg)
#model.model.load_state_dict(torch.load('finetuned-detectron2-maskrcnn_13000imgs_03lr_3rdRun/model_final.pth')['model'])
eval_results = model.model.eval(
) # finetuned-detectron2-maskrcnn_fixed/model_final.pth')['model'])
# 'finetuned-detectron2-maskrcnn_fixed/model_final.pth'
# finetuned-detectron2-maskrcnn_5700imgs_025lr_furthertuned/model_final.pth
if args.use_mask == 0:
obj_detect = MRCNN_FPN(model, num_classes=1)
else:
obj_detect = FRCNN_FPN(model, num_classes=1)
#obj_detect.to(device)
# reid
reid_network = resnet50(pretrained=False,
output_dim=128) # need to insert dictionary here
reid_network.load_state_dict(
torch.load(
reid_pth, #tracktor['reid_weights'],'/home/ubuntu/code/tracking_wo_bnw/output/tracktor/reid/test/ResNet_iter_25137.pth'
map_location=device)) #lambda storage, loc: storage))
reid_network.eval()
#reid_network.cuda()
reid_network.to(device)
tracker = {
# FRCNN score threshold for detections
'detection_person_thresh': 0.5,
# FRCNN score threshold for keeping the track alive
'regression_person_thresh': 0.5,
# NMS threshold for detection
'detection_nms_thresh': 0.3,
# NMS theshold while tracking
'regression_nms_thresh': 0.6,
# motion model settings
'motion_model': {
'enabled': True,
# average velocity over last n_steps steps
'n_steps': 2,
# if true, only model the movement of the bounding box center. If false, width and height are also modeled.
'center_only': True
},
# DPM or DPM_RAW or 0, raw includes the unfiltered (no nms) versions of the provided detections,
# 0 tells the tracker to use private detections (Faster R-CNN)
'public_detections': False,
# How much last appearance features are to keep
'max_features_num': 20,
# Do camera motion compensation
'do_align': False,
# Which warp mode to use (cv2.MOTION_EUCLIDEAN, cv2.MOTION_AFFINE, ...)
'warp_mode': str(cv2.MOTION_EUCLIDEAN),
# maximal number of iterations (original 50)
'number_of_iterations': 500,
# Threshold increment between two iterations (original 0.001)
'termination_eps': 0.00001,
# Use siamese network to do reid
'do_reid': True,
# How much timesteps dead tracks are kept and cosidered for reid
'inactive_patience': 15,
# How similar do image and old track need to be to be considered the same person
'reid_sim_threshold': 20.0,
# How much IoU do track and image need to be considered for matching
'reid_iou_threshold': 0.05
}
tracker = IJP_tracker(obj_detect, reid_network, tracker)
# tracker = Tracker(obj_detect, reid_network, tracker)
def transform_img(i, pth='/home/ubuntu/videos/'):
if (i < 10):
ind = f'0000{i}'
elif (i < 100):
ind = f'000{i}'
else:
ind = f'00{i}'
frame = Image.open(f'{pth}split_frames_{ind}.jpg')
frame_ten = torch.tensor(np.reshape(np.array(frame),
(1, 3, 720, 1280)),
device=device,
dtype=torch.float32)
blob = {'img': frame_ten}
return blob
# cannot make frame size assumption...
# need to make this an argparse argument
pth = '/home/video/'
# vid_name = args.video #'57675_003286_Endzone'
img_paths = glob(f'{pth}/*') # {vid_name}
img_paths.sort()
# print(img_paths)
tracker.reset()
# results = []
# print('Starting tracking!')
# for i in tqdm(range(1,len(img_paths)+1)):
# #tracker.reset()
# blob = transform_img(i ,pth=pth)
# with torch.no_grad():
# tracker.step(blob)
track_dict = {}
print("\n")
print("###########################################")
print("############## BEGIN TRACKING #############")
print("###########################################")
print('\n')
for i in tqdm(range(1, len(img_paths) + 1)):
track_dict[i] = {}
blob = transform_img(i, pth=pth)
with torch.no_grad():
tracker.step(blob)
for tr in tracker.tracks:
track_dict[i][tr.id] = tr.pos[0].detach().cpu().numpy()
iou_dict = {}
iou_dict2 = {}
new_tracks = {}
missing_tracks = {}
for track in tqdm(track_dict):
if track == 1: # set dictionaries with 1st frame
new_tracks[track] = list(track_dict[track].keys())
missing_tracks[track] = set()
iou_dict[track] = {}
iou_dict[track] = {}
for tr in track_dict[track]:
iou_dict[track][tr] = {}
else:
new_tracks[track] = set(list(track_dict[track].keys())) - set(
new_tracks[1]) # - set(list(track_dict[track-1].keys()))
missing_tracks[track] = set(new_tracks[1]) - set(
list(track_dict[track - 1].keys()))
iou_dict[track] = {}
iou_dict2[track] = {}
for tr in track_dict[track]:
iou_dict[track][tr] = {}
iou_dict2[track][tr] = {}
for t in track_dict[track - 1]:
iou = bbox_iou(track_dict[track][tr],
track_dict[track - 1][t])
if iou > 0.001:
iou_dict[track][tr][t] = iou
if track > 2:
for t in track_dict[track - 2]:
iou = bbox_iou(track_dict[track][tr],
track_dict[track - 2][t])
if iou > 0.001:
iou_dict2[track][tr][t] = iou
tracks_to_delete = {}
tracks_to_change = {}
momentum_dict = {}
for track in track_dict:
if track == 1:
for tr in track_dict[track]:
momentum_dict[tr] = 0 # initialize momentum dict
elif len(track_dict[track]
) > 22: # if there are more than 22 annotations
tracks_to_delete[track] = []
for tr in track_dict[track]:
try:
momentum_dict[tr] += 1
except:
momentum_dict[tr] = 0
for ind in iou_dict[track]:
if (iou_dict[track][ind] == {}) & (
ind > 22
): # need to adjust this, right now just looking if there is ANY IoU
tracks_to_delete[track].append(ind)
else: # if less than 22 annotations
tracks_to_change[track] = {}
if new_tracks[track] != set(
): # if there are new tracks, check them
for tr in track_dict[track]: # update momentum dict
try:
momentum_dict[tr] += 1
except:
momentum_dict[tr] = 0
for newt in new_tracks[track]: # cycle through new tracks
# print('For track ',track,"and ID ",newt,iou_dict[track][newt])
if newt > 22: # if there is a new track and it's greater than 22, change it, figure out what to change to
if (missing_tracks[track] != set()) | (
missing_tracks[track - 1] != set()
): # if there are missing tracks, cycle through them and compare
mis_iou = {}
if missing_tracks[track] != set():
for mis in missing_tracks[track]:
try:
# iou = bbox_iou(track_dict[track-1][mis], track_dict[track][newt])
dist = distance.euclidean(
track_dict[track - 1][mis],
track_dict[track][newt])
except:
try:
# iou = bbox_iou(track_dict[track-2][mis], track_dict[track][newt])
dist = distance.euclidean(
track_dict[track - 2][mis],
track_dict[track][newt])
except:
try:
# iou = bbox_iou(track_dict[track-3][mis], track_dict[track][newt])
dist = distance.euclidean(
track_dict[track - 3][mis],
track_dict[track][newt])
except:
pass
mis_iou[mis] = dist
# tracks_to_change[track][newt] =
tracks_to_change[track][newt] = mis_iou
# if missing_tracks[track-1]!=set():
# for mis in missing_tracks[track-1]:
# iou = bbox_iou(track_dict[track-2][mis], track_dict[track][newt])
# mis_iou[mis] = iou
# tracks_to_change[track][newt] = mis_iou
# try:
# if max(tracks_to_change[track][newt].values())>.2:
# for t in tracks_to_change[trc][tr]:
to_ind = np.argmin(
list(tracks_to_change[track]
[newt].values()))
to_id = list(tracks_to_change[track]
[newt].keys())[to_ind]
to_pos = track_dict[track][newt]
del track_dict[track][newt]
track_dict[track][to_id] = to_pos
# need to send results to s3
result = tracker.get_results()
file = 'tracking_results.json'
tracking_dict = {}
for res in result:
tracking_dict[res] = {}
for r in result[res]:
tracking_dict[res][r] = list(result[res][r][0:4])
with open(file, 'w') as f:
json.dump(tracking_dict, f)
# now = str(datetime.datetime.now()).replace(' ','').replace(':','-')
k = f'nfl-data/tracking_results_{folder}.json'
s3.upload_file(Filename=file, Bucket=args.bucket, Key=k)
print(f'Tracking finished and results saved to: s3://{args.bucket}/{k}')
os.makedirs('/home/labeled_frames')
# create labeled images
print('\n')
print("###########################################")
print("############ Generating Video! ############")
print("###########################################")
print('\n')
print('...')
for j, pth in tqdm(enumerate(img_paths)):
fig, ax = plt.subplots(1, figsize=(24, 14))
img = Image.open(pth)
# Display the image
ax.imshow(np.array(img))
# Create a Rectangle patch
label_list = {}
for r in track_dict[j + 1]:
try:
res = track_dict[j + 1][r]
label_list[r] = res
except:
pass
for i, r in enumerate(label_list):
labs = label_list[r]
rect = patches.Rectangle((labs[0], labs[1]),
labs[2] - labs[0],
labs[3] - labs[1],
linewidth=1,
edgecolor='r',
facecolor='none') # 50,100),40,30
ax.add_patch(rect)
plt.text(labs[0] - 10, labs[1] - 10, f'H:{r}', fontdict=None)
plt.savefig(
f"/home/labeled_frames/{pth.split('/')[-1].replace('.jpg','.png')}"
)
plt.close()
# create video of labels
os.system(
'ffmpeg -r 15 -f image2 -s 1280x720 -i /home/labeled_frames/split_frames_%05d.png -vcodec libx264 -crf 25 -pix_fmt yuv420p /home/labeled_frames.mp4'
)
k = f'nfl-data/tracking_results_{folder}.mp4'
s3.upload_file(Filename='/home/labeled_frames.mp4',
Bucket=args.bucket,
Key=k)
print('\n')
print(f'Video uploaded to: s3://{args.bucket}/{k}')
# for launching A2I job set a conditional here
s3_fname = f's3://{args.bucket}/{k}'
workteam = 'arn:aws:sagemaker:us-east-1:209419068016:workteam/private-crowd/ijp-private-workteam'
flowDefinitionName = 'ijp-video-flow-official'
humanTaskUiArn = 'arn:aws:sagemaker:us-east-1:209419068016:human-task-ui/ijp-video-task3'
#'s3://privisaa-bucket-virginia/nfl-data/nfl-frames/nfl-video-frame00001.png'
# create_workflow_definition_response = sagemaker_client.create_flow_definition(
# FlowDefinitionName= flowDefinitionName,
# RoleArn= role,
# HumanLoopConfig= {
# "WorkteamArn": workteam,
# "HumanTaskUiArn": humanTaskUiArn,
# "TaskCount": 1,
# "TaskDescription": "Identify if the labels in the video look correct.",
# "TaskTitle": "Video classification a2i demo"
# },
# OutputConfig={
# "S3OutputPath" : OUTPUT_PATH
# }
# )
# flowDefinitionArn = create_workflow_definition_response['FlowDefinitionArn']
inputContent = {
"initialValue": .2,
"taskObject":
s3_fname # the s3 object will be passed to the worker task UI to render
}
now = str(datetime.datetime.now()).replace(' ', '-').replace(':',
'-').replace(
'.', '-')
response = a2i.start_human_loop(
HumanLoopName=f'ijp-video-{now}',
FlowDefinitionArn=flowDefinitionArn,
HumanLoopInput={"InputContent": json.dumps(inputContent)},
DataAttributes={
'ContentClassifiers': ['FreeOfPersonallyIdentifiableInformation']
})
print(f'Launched A2I loop ijp-video-{now}')
sns.publish(
TopicArn='arn:aws:sns:us-east-1:209419068016:ijp-topic',
Message=
f'Your video inference is done! You can find the output here: s3://{args.bucket}/{k}',
Subject='Video labeling')
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)
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 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 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 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 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']))
if not osp.exists(output_dir):
os.makedirs(output_dir)
##########################
# Initialize the modules #
##########################
args = parse_args()
# object detection
print("Initializing object detector.")
obj_detect = Model(args.config, args.checkpoint, device)
# 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'], device)
time_total = 0
num_frames = 0
def main(tracktor, reid, train, _config, _log, _run):
# instantiate tracktor
sacred.commands.print_config(_run)
# set all seeds
seed_everything(tracktor['seed'])
output_dir = osp.abspath(
osp.join(osp.dirname(__file__), '..', '..', 'output', 'motion',
train['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)
# object detection
_log.info("Initializing object detector.")
obj_detect: nn.Module = FRCNN_FPN(num_classes=2)
obj_detect_state_dict = torch.load(_config['tracktor']['obj_detect_model'])
obj_detect.load_state_dict(obj_detect_state_dict)
obj_detect.eval()
obj_detect.cuda()
# reid
reid_network = resnet50(pretrained=False, **reid['cnn'])
reid_network.load_state_dict(torch.load(tracktor['reid_weights']))
reid_network.eval()
reid_network.cuda()
# tracktor
if 'oracle' in tracktor:
raise ValueError('Oracle tracker not supported.')
else:
tracker = Tracker(obj_detect, reid_network, None, tracktor['tracker'],
tracktor['motion'],
train['dataset_args']['train']['min_length'])
assert train['dataset_args']['train']['min_length'] == train[
'dataset_args']['val']['min_length']
assert train['dataset_args']['train']['episode_length'] == train[
'dataset_args']['val']['episode_length']
assert train['dataset_args']['train']['offset'] == train['dataset_args'][
'val']['offset']
assert train['loss'] in ['mse', 'iou']
config = TrainingConfig(
data=[eval(train['dataset'])] * 2,
data_args=[
train['dataset_args']['train'], train['dataset_args']['val']
],
model=eval(tracktor['motion']['model']),
model_args=tracktor['motion']['model_args'],
batch_size=train['batch_size'],
epochs=train['epochs'],
log_path=output_dir,
primary_metric='miou',
smaller_is_better=False,
save_every=train['save_every'],
num_workers=train['num_workers'],
lr=train['lr'],
patience=train['patience'],
data_mean=train['data_mean'],
data_std=train['data_std'],
shuffle=train['shuffle'],
teacher_forcing=train['teacher_forcing'],
loss=train['loss'],
lmbda=train['lmbda'],
tracktor_val_every=train['tracktor_val_every'],
use_box_coding=tracktor['motion']['use_box_coding'],
box_coding_weights=tracktor['motion']['box_coding_weights'],
)
def validate_tracktor(motion_network, epoch):
# inject current network into tracker
tracker.motion_network = motion_network
time_total = 0
num_frames = 0
mot_accums = []
dataset = Datasets(train['tracktor_val_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 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(epoch),
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)")
metrics = {}
if mot_accums:
summary = evaluate_mot_accums(
mot_accums, [str(s) for s in dataset if not s.no_gt],
generate_overall=True,
return_summary=True,
metrics=train['tracktor_val_metrics'])
metrics = {
m: summary.loc['OVERALL', m]
for m in train['tracktor_val_metrics']
}
return metrics
# re-seed
seed_everything(tracktor['seed'])
trainer = Trainer(config, validate_tracktor)
trainer.train()
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 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)
