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
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'], tracktor['output_subdir']) sacred_config = osp.join(output_dir, 'sacred_config.yaml') if not osp.exists(output_dir): os.makedirs(output_dir) with open(sacred_config, 'w') as outfile: yaml.dump(_config, outfile, default_flow_style=False) ########################## # Initialize the modules # ########################## # object detection _log.info("Initializing object detector.") obj_detect = FRCNN_FPN(num_classes=2).to(device) obj_detect.load_state_dict( torch.load(_config['tracktor']['obj_detect_model'], map_location=lambda storage, loc: storage)) obj_detect.eval() # reid reid_network = resnet50(pretrained=False, **reid['cnn']).to(device) reid_network.load_state_dict( torch.load(tracktor['reid_weights'], map_location=lambda storage, loc: storage)) reid_network.eval() # tracktor if 'oracle' in tracktor: tracker = OracleTracker(obj_detect, reid_network, tracktor['tracker'], tracktor['oracle']) else: tracker = Tracker(obj_detect, reid_network, tracktor['tracker']) time_total = 0 num_frames = 0 mot_accums = [] dataset = Datasets(tracktor['dataset']) for seq in dataset: tracker.reset() start = time.time() _log.info(f"Tracking: {seq}") data_loader = DataLoader(seq, batch_size=1, shuffle=False) for i, frame in enumerate(tqdm(data_loader)): if len(seq) * tracktor['frame_split'][0] <= i <= len( seq) * tracktor['frame_split'][1]: tracker.step(frame, i) num_frames += 1 results = tracker.get_results() time_total += time.time() - start _log.info(f"Tracks found: {len(results)}") _log.info(f"Runtime for {seq}: {time.time() - start :.1f} s.") if tracktor['interpolate']: results = interpolate(results) if seq.no_gt: _log.info(f"No GT data for evaluation available.") else: mot_accums.append(get_mot_accum(results, seq)) _log.info(f"Writing predictions to: {output_dir}") seq.write_results(results, output_dir) if tracktor['write_images']: plot_sequence(results, seq, osp.join(output_dir, tracktor['dataset'], str(seq))) _log.info( f"Tracking runtime for all sequences (without evaluation or image writing): " f"{time_total:.1f} s ({num_frames / time_total:.1f} Hz)") if mot_accums: summary = evaluate_mot_accums(mot_accums, [str(s) for s in dataset if not s.no_gt], generate_overall=True) summary.to_pickle( "output/finetuning_results/results_{}_{}_{}_{}_{}.pkl".format( tracktor['output_subdir'], tracktor['tracker']['finetuning']['max_displacement'], tracktor['tracker']['finetuning']['batch_size'], tracktor['tracker']['finetuning']['learning_rate'], tracktor['tracker']['finetuning']['iterations']))
def main(tracktor, reid, _config, _log, _run): sacred.commands.print_config(_run) # set all seeds torch.manual_seed(tracktor['seed']) torch.cuda.manual_seed(tracktor['seed']) np.random.seed(tracktor['seed']) torch.backends.cudnn.deterministic = True output_dir = osp.join(get_output_dir(tracktor['module_name']), tracktor['name']) sacred_config = osp.join(output_dir, 'sacred_config.yaml') if not osp.exists(output_dir): os.makedirs(output_dir) with open(sacred_config, 'w') as outfile: yaml.dump(_config, outfile, default_flow_style=False) ########################## # Initialize the modules # ########################## # object detection _log.info("Initializing object detector.") obj_detect = FRCNN_FPN(num_classes=2) obj_detect.load_state_dict( torch.load(_config['tracktor']['obj_detect_model'], map_location=lambda storage, loc: storage)) obj_detect.eval() obj_detect.cuda() # reid reid_network = resnet50(pretrained=False, **reid['cnn']) reid_network.load_state_dict( torch.load(tracktor['reid_weights'], map_location=lambda storage, loc: storage)) reid_network.eval() reid_network.cuda() # motion network motion_network = None if tracktor['tracker']['motion_model_enabled'] and not tracktor['motion'][ 'use_cva_model']: motion_network = eval( tracktor['motion']['model'])(**tracktor['motion']['model_args']) motion_network.load_state_dict( torch.load(tracktor['motion']['network_weights'])['model']) motion_network.eval().cuda() # tracktor if 'oracle' in tracktor: tracker = OracleTracker(obj_detect, reid_network, tracktor['tracker'], tracktor['oracle']) else: tracker = Tracker(obj_detect, reid_network, motion_network, tracktor['tracker'], tracktor['motion'], 2) time_total = 0 num_frames = 0 mot_accums = [] dataset = Datasets(tracktor['dataset']) for seq in dataset: tracker.reset() _log.info(f"Tracking: {seq}") data_loader = DataLoader(seq, batch_size=1, shuffle=False) start = time.time() all_mm_times = [] all_warp_times = [] for i, frame in enumerate(tqdm(data_loader)): if len(seq) * tracktor['frame_split'][0] <= i <= len( seq) * tracktor['frame_split'][1]: with torch.no_grad(): mm_time, warp_time = tracker.step(frame) if mm_time is not None: all_mm_times.append(mm_time) if warp_time is not None: all_warp_times.append(warp_time) num_frames += 1 results = tracker.get_results() time_total += time.time() - start _log.info(f"Tracks found: {len(results)}") _log.info(f"Runtime for {seq}: {time.time() - start :.1f} s.") _log.info( f"Average FPS for {seq}: {len(data_loader) / (time.time() - start) :.3f}" ) _log.info( f"Average MM time for {seq}: {float(np.array(all_mm_times).mean()) :.3f} s" ) if all_warp_times: _log.info( f"Average warp time for {seq}: {float(np.array(all_warp_times).mean()) :.3f} s" ) if tracktor['interpolate']: results = interpolate(results) if 'semi_online' in tracktor and tracktor['semi_online']: for i, track in results.items(): for frame in sorted(track, reverse=True): if track[frame][5] == 0: break del track[frame] if tracktor['write_images']: plot_sequence(results, seq, osp.join(output_dir, tracktor['dataset'], str(seq)), tracktor['tracker']['plot_mm']) if seq.no_gt: _log.info(f"No GT data for evaluation available.") else: mot_accums.append(get_mot_accum(results, seq)) _log.info(f"Writing predictions to: {output_dir}") seq.write_results(results, output_dir) _log.info( f"Tracking runtime for all sequences (without evaluation or image writing): " f"{time_total:.2f} s for {num_frames} frames ({num_frames / time_total:.2f} Hz)" ) if mot_accums: evaluate_mot_accums(mot_accums, [str(s) for s in dataset if not s.no_gt], generate_overall=True)
def main(module_name, name, seed, obj_detect_models, reid_models, tracker, oracle, dataset, load_results, frame_range, interpolate, write_images, _config, _log, _run): sacred.commands.print_config(_run) # set all seeds torch.manual_seed(seed) torch.cuda.manual_seed(seed) np.random.seed(seed) torch.backends.cudnn.deterministic = True output_dir = osp.join(get_output_dir(module_name), name) sacred_config = osp.join(output_dir, 'sacred_config.yaml') if not osp.exists(output_dir): os.makedirs(output_dir) with open(sacred_config, 'w') as outfile: yaml.dump(copy.deepcopy(_config), outfile, default_flow_style=False) ########################## # Initialize the modules # ########################## # object detection _log.info("Initializing object detector(s).") obj_detects = [] for obj_detect_model in obj_detect_models: obj_detect = FRCNN_FPN(num_classes=2) obj_detect.load_state_dict( torch.load(obj_detect_model, map_location=lambda storage, loc: storage)) obj_detects.append(obj_detect) obj_detect.eval() if torch.cuda.is_available(): obj_detect.cuda() # reid _log.info("Initializing reID network(s).") reid_networks = [] for reid_model in reid_models: reid_cfg = os.path.join(os.path.dirname(reid_model), 'sacred_config.yaml') reid_cfg = yaml.safe_load(open(reid_cfg)) reid_network = ReIDNetwork_resnet50(pretrained=False, **reid_cfg['model_args']) reid_network.load_state_dict( torch.load(reid_model, map_location=lambda storage, loc: storage)) reid_network.eval() if torch.cuda.is_available(): reid_network.cuda() reid_networks.append(reid_network) # tracktor if oracle is not None: tracker = OracleTracker(obj_detect, reid_network, tracker, oracle) else: tracker = Tracker(obj_detect, reid_network, tracker) time_total = 0 num_frames = 0 mot_accums = [] dataset = Datasets(dataset) for seq, obj_detect, reid_network in zip(dataset, obj_detects, reid_networks): tracker.obj_detect = obj_detect tracker.reid_network = reid_network tracker.reset() _log.info(f"Tracking: {seq}") start_frame = int(frame_range['start'] * len(seq)) end_frame = int(frame_range['end'] * len(seq)) seq_loader = DataLoader( torch.utils.data.Subset(seq, range(start_frame, end_frame))) num_frames += len(seq_loader) results = {} if load_results: results = seq.load_results(output_dir) if not results: start = time.time() for frame_data in tqdm(seq_loader): with torch.no_grad(): tracker.step(frame_data) results = tracker.get_results() time_total += time.time() - start _log.info(f"Tracks found: {len(results)}") _log.info(f"Runtime for {seq}: {time.time() - start :.2f} s.") if interpolate: results = interpolate_tracks(results) _log.info(f"Writing predictions to: {output_dir}") seq.write_results(results, output_dir) if seq.no_gt: _log.info("No GT data for evaluation available.") else: mot_accums.append(get_mot_accum(results, seq_loader)) if write_images: plot_sequence(results, seq, osp.join(output_dir, str(dataset), str(seq)), write_images) if time_total: _log.info( f"Tracking runtime for all sequences (without evaluation or image writing): " f"{time_total:.2f} s for {num_frames} frames ({num_frames / time_total:.2f} Hz)" ) if mot_accums: _log.info("Evaluation:") evaluate_mot_accums(mot_accums, [str(s) for s in dataset if not s.no_gt], generate_overall=True)
def main(tracktor, reid, _config, _log, _run): 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)
tracker.step(frame) num_frames += 1 results = tracker.get_results() time_total += time.time() - start print(f"Tracks found: {len(results)}") print(f"Runtime for {dataset}: {time.time() - start :.2f} s.") if tracktor['interpolate']: results = interpolate(results) if dataset.no_gt: print(f"No GT data for evaluation available.") else: mot_accums.append(get_mot_accum(results, dataset)) print(f"Writing predictions to: {output_dir}") dataset.write_results(results, output_dir) if tracktor['write_images']: plot_sequence(results, dataset, osp.join(output_dir)) print( 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 # ########################## _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 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)