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)