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 __init__(self, obj_detect_weights, correlation_args, batch_norm, conv_channels, n_box_channels, roi_output_size, avg_box_features, hidden_size, input_length, n_layers, dropout, correlation_only, use_env_features, fixed_env, correlation_last_only, feature_level=1): super().__init__(correlation_args, batch_norm, conv_channels, n_box_channels, roi_output_size, avg_box_features, hidden_size, input_length, n_layers, dropout, correlation_only, use_env_features, fixed_env, correlation_last_only) self.feature_level = feature_level self.frcnn: nn.Module = FRCNN_FPN(num_classes=2) self.frcnn.load_state_dict(torch.load(obj_detect_weights))
def __init__(self, tracker_config, output_dim=256, pool_size=7, representation_dim=1024, motion_repr_dim=512, vis_conv_only=True, use_modulator=True, use_bn=False): super().__init__(output_dim, pool_size, representation_dim, motion_repr_dim, vis_conv_only, use_modulator, use_bn) 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)) self.transform = obj_detect.transform self.backbone = obj_detect.backbone self.box_roi_pool = obj_detect.roi_heads.box_roi_pool self.box_head = obj_detect.roi_heads.box_head
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
box_features = obj_detect.roi_heads.box_roi_pool( obj_detect.features, gts, obj_detect.preprocessed_images.image_sizes) box_head_features = obj_detect.roi_heads.box_head(box_features) return box_features.cpu(), box_head_features.cpu() with open('experiments/cfgs/tracktor.yaml', 'r') as f: tracker_config = yaml.safe_load(f) val_set = MOT17Vis('val', 0.8, 0.0, val_bbox_jitter=True) val_loader = DataLoader(val_set, batch_size=1, shuffle=True, num_workers=2) obj_detect = FRCNN_FPN(num_classes=2) obj_detect.load_state_dict( torch.load(tracker_config['tracktor']['obj_detect_model'], map_location=lambda storage, loc: storage)) obj_detect.eval() obj_detect.cuda() vis_model = VisEst(conv_only=False) vis_model.load_state_dict( torch.load( 'output/tracktor/motion/vis_valjit_flip_fullfeat_l21e-4/vis_model_epoch_94.pth' )) vis_model.eval() vis_model.cuda() for data in val_loader:
def main(tracktor, reid, _config, _log, _run): target = Target() targetpath = target.Folder() targetname = target.TargetName() vottpath = target.GetVottPath() vottfile = target.GetVottContent() dictid, timelist = target.GetTagTime(vottfile) print(f"{len(timelist)} frames were tagged") timedict = target.ExtractByTimeList(timelist) bbdict = target.GetbbWithTime(vottfile) sacred.commands.print_config(_run) # set all seeds torch.manual_seed(tracktor['seed']) torch.cuda.manual_seed(tracktor['seed']) np.random.seed(tracktor['seed']) torch.backends.cudnn.deterministic = True output_dir = osp.join(get_output_dir(tracktor['module_name']), tracktor['name']) sacred_config = osp.join(output_dir, 'sacred_config.yaml') if not osp.exists(output_dir): os.makedirs(output_dir) with open(sacred_config, 'w') as outfile: yaml.dump(_config, outfile, default_flow_style=False) ########################## # Initialize the modules # ########################## # object detection _log.info("Initializing object detector.") obj_detect = FRCNN_FPN(num_classes=2) obj_detect.load_state_dict( torch.load(_config['tracktor']['obj_detect_model'], map_location=lambda storage, loc: storage)) obj_detect.eval() obj_detect.cuda() # reid reid_network = resnet50(pretrained=False, **reid['cnn']) reid_network.load_state_dict( torch.load(tracktor['reid_weights'], map_location=lambda storage, loc: storage)) reid_network.eval() reid_network.cuda() # tracktor print("Tracktor初始化完成") tracker = Tracker(obj_detect, reid_network, tracktor['tracker']) time_total = 0 num_frames = 0 mot_accums = [] dataset = Datasets(tracktor['dataset']) for seq in dataset: tracker.reset() start = time.time() _log.info(f"Tracking: {seq}") data_loader = DataLoader(seq, batch_size=1, shuffle=False) print(f"{seq}加載完成, tracking開始") for i, frame in enumerate(tqdm(data_loader)): if len(seq) * tracktor['frame_split'][0] <= i <= len( seq) * tracktor['frame_split'][1]: id = tracker.step(frame, bbdict[timedict["%06d" % num_frames]]) target.WriteID2asset(id, dictid[timedict["%06d" % num_frames]]) num_frames += 1 results = tracker.get_results() ids = list(results.keys()) target.WriteID2vott(ids, vottfile=vottfile) time_total += time.time() - start _log.info(f"Tracks found: {len(results)}") _log.info(f"Runtime for {seq}: {time.time() - start :.1f} s.") target.CleanImg() if tracktor['interpolate']: results = interpolate(results) if seq.no_gt: _log.info(f"No GT data for evaluation available.") else: mot_accums.append(get_mot_accum(results, seq)) _log.info(f"Writing predictions to: {output_dir}") seq.write_results(results, output_dir) if tracktor['write_images']: plot_sequence(results, seq, osp.join(output_dir, tracktor['dataset'], str(seq))) if tracktor['write_videos']: plot_sequence_video( results, seq, osp.join(output_dir, tracktor['dataset'], str(seq))) _log.info( f"Tracking runtime for all sequences (without evaluation or image writing): " f"{time_total:.1f} s ({num_frames / time_total:.1f} Hz)") if mot_accums: evaluate_mot_accums(mot_accums, [str(s) for s in dataset if not s.no_gt], generate_overall=True)
def __init__( self, ckpt_path, frcnn_weights_path, reid_weights_path, tracking_cfg_path, preprocessing_cfg_path, use_gt, pre_cnn, pre_track, ): self.name = "MPNTracker" self.use_gt = use_gt if not self.use_gt: with open(tracking_cfg_path) as config_file: config = yaml.load(config_file) with open(preprocessing_cfg_path) as config_file: pre_config = yaml.load(config_file) frcnn_prepr_params = pre_config["frcnn_prepr_params"] tracktor_params = pre_config["tracktor_params"] CustomMOTNeuralSolver.reid_weights_path = reid_weights_path # preprocessor self.pre_track = pre_track if self.pre_track == "Tracktor" or self.pre_track == "FRCNN": obj_detect = FRCNN_FPN(num_classes=2) obj_detect.load_state_dict( torch.load( frcnn_weights_path, map_location=lambda storage, loc: storage, )) obj_detect.eval() obj_detect.cuda() if self.pre_track == "Tracktor": self.prepr_params = tracktor_params make_deterministic(self.prepr_params["seed"]) self.preprocessor = Tracker(obj_detect, None, self.prepr_params["tracker"]) elif self.pre_track == "FRCNN": self.prepr_params = frcnn_prepr_params make_deterministic(self.prepr_params["seed"]) self.preprocessor = FRCNNPreprocessor( obj_detect, self.prepr_params) self.transforms = ToTensor() # Load model from checkpoint and update config entries that may vary from the ones used in training self.model = CustomMOTNeuralSolver.load_from_checkpoint( checkpoint_path=ckpt_path) self.model.cnn_model.eval() self.model.hparams.update({ "eval_params": config["eval_params"], "data_splits": config["data_splits"], }) self.model.hparams["dataset_params"][ "precomputed_embeddings"] = False self.model.hparams["dataset_params"]["img_batch_size"] = 2500 self.pre_cnn = pre_cnn if self.pre_cnn: self.extract_transforms = Compose(( Resize(self.model.hparams["dataset_params"]["img_size"]), ToTensor(), Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ))
def main(tracktor, reid, _config, _log, _run): sacred.commands.print_config(_run) # set all seeds torch.manual_seed(tracktor['seed']) torch.cuda.manual_seed(tracktor['seed']) np.random.seed(tracktor['seed']) torch.backends.cudnn.deterministic = True output_dir = osp.join(get_output_dir(tracktor['module_name']), tracktor['name']) sacred_config = osp.join(output_dir, 'sacred_config.yaml') if not osp.exists(output_dir): os.makedirs(output_dir) with open(sacred_config, 'w') as outfile: yaml.dump(_config, outfile, default_flow_style=False) ########################## # Initialize the modules # ########################## # object detection _log.info("Initializing object detector.") obj_detect = FRCNN_FPN(num_classes=2) obj_detect.load_state_dict( torch.load(_config['tracktor']['obj_detect_model'], map_location=lambda storage, loc: storage)) obj_detect.eval() obj_detect.cuda() # reid reid_network = resnet50(pretrained=False, **reid['cnn']) reid_network.load_state_dict( torch.load(tracktor['reid_weights'], map_location=lambda storage, loc: storage)) reid_network.eval() reid_network.cuda() # tracktor if 'oracle' in tracktor: tracker = OracleTracker(obj_detect, reid_network, tracktor['tracker'], tracktor['oracle']) else: tracker = Tracker(obj_detect, reid_network, tracktor['tracker']) time_total = 0 num_frames = 0 mot_accums = [] dataset = Datasets(tracktor['dataset']) for seq in dataset: tracker.reset() start = time.time() _log.info(f"Tracking: {seq}") data_loader = DataLoader(seq, batch_size=1, shuffle=False) for i, frame in enumerate(tqdm(data_loader)): if len(seq) * tracktor['frame_split'][0] <= i <= len( seq) * tracktor['frame_split'][1]: tracker.step(frame) num_frames += 1 results = tracker.get_results() time_total += time.time() - start _log.info(f"Tracks found: {len(results)}") _log.info(f"Runtime for {seq}: {time.time() - start :.1f} s.") if tracktor['interpolate']: results = interpolate(results) if seq.no_gt: _log.info(f"No GT data for evaluation available.") else: mot_accums.append(get_mot_accum(results, seq)) _log.info(f"Writing predictions to: {output_dir}") seq.write_results(results, output_dir) if tracktor['write_images']: plot_sequence(results, seq, osp.join(output_dir, tracktor['dataset'], str(seq))) img_array = [] dir = osp.join(output_dir, tracktor['dataset'], str(seq), "*.jpg") files = glob.glob(dir) sorted_files = natsorted(files) for filename in sorted_files: img = cv2.imread(filename) height, width, layers = img.shape size = (width, height) img_array.append(img) out = cv2.VideoWriter( osp.join(output_dir, tracktor['dataset'], str(seq), "result_video.avi"), cv2.VideoWriter_fourcc(*'DIVX'), 10, size) for i in range(len(img_array)): out.write(img_array[i]) out.release() _log.info( f"Tracking runtime for all sequences (without evaluation or image writing): " f"{time_total:.1f} s ({num_frames / time_total:.1f} Hz)") if mot_accums: evaluate_mot_accums(mot_accums, [str(s) for s in dataset if not s.no_gt], generate_overall=True)
def main(): 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'], tracktor['output_subdir']) sacred_config = osp.join(output_dir, 'sacred_config.yaml') if not osp.exists(output_dir): os.makedirs(output_dir) with open(sacred_config, 'w') as outfile: yaml.dump(_config, outfile, default_flow_style=False) ########################## # Initialize the modules # ########################## # object detection _log.info("Initializing object detector.") obj_detect = FRCNN_FPN(num_classes=2).to(device) obj_detect.load_state_dict( torch.load(_config['tracktor']['obj_detect_model'], map_location=lambda storage, loc: storage)) obj_detect.eval() # reid reid_network = resnet50(pretrained=False, **reid['cnn']).to(device) reid_network.load_state_dict( torch.load(tracktor['reid_weights'], map_location=lambda storage, loc: storage)) reid_network.eval() # tracktor if 'oracle' in tracktor: tracker = OracleTracker(obj_detect, reid_network, tracktor['tracker'], tracktor['oracle']) else: tracker = Tracker(obj_detect, reid_network, tracktor['tracker']) time_total = 0 num_frames = 0 mot_accums = [] dataset = Datasets(tracktor['dataset']) for seq in dataset: tracker.reset() start = time.time() _log.info(f"Tracking: {seq}") data_loader = DataLoader(seq, batch_size=1, shuffle=False) for i, frame in enumerate(tqdm(data_loader)): if len(seq) * tracktor['frame_split'][0] <= i <= len( seq) * tracktor['frame_split'][1]: tracker.step(frame, i) num_frames += 1 results = tracker.get_results() time_total += time.time() - start _log.info(f"Tracks found: {len(results)}") _log.info(f"Runtime for {seq}: {time.time() - start :.1f} s.") if tracktor['interpolate']: results = interpolate(results) if seq.no_gt: _log.info(f"No GT data for evaluation available.") else: mot_accums.append(get_mot_accum(results, seq)) _log.info(f"Writing predictions to: {output_dir}") seq.write_results(results, output_dir) if tracktor['write_images']: plot_sequence(results, seq, osp.join(output_dir, tracktor['dataset'], str(seq))) _log.info( f"Tracking runtime for all sequences (without evaluation or image writing): " f"{time_total:.1f} s ({num_frames / time_total:.1f} Hz)") if mot_accums: summary = evaluate_mot_accums(mot_accums, [str(s) for s in dataset if not s.no_gt], generate_overall=True) summary.to_pickle( "output/finetuning_results/results_{}_{}_{}_{}_{}.pkl".format( tracktor['output_subdir'], tracktor['tracker']['finetuning']['max_displacement'], tracktor['tracker']['finetuning']['batch_size'], tracktor['tracker']['finetuning']['learning_rate'], tracktor['tracker']['finetuning']['iterations']))
def test_motion_model(val_loader, tracker_config, motion_model): obj_detect = FRCNN_FPN(num_classes=2) obj_detect.load_state_dict( torch.load(tracker_config['tracktor']['obj_detect_model'], map_location=lambda storage, loc: storage)) obj_detect.eval() obj_detect.cuda() reid_network = resnet50(pretrained=False, output_dim=128) reid_network.load_state_dict( torch.load(tracker_config['tracktor']['reid_weights'], map_location=lambda storage, loc: storage)) reid_network.eval() reid_network.cuda() pred_loss_func = nn.SmoothL1Loss() loss_iters = [] low_vis_loss_sum = 0.0 low_vis_num = 0 high_vis_loss_sum = 0.0 high_vis_num = 0 total_iters = len(val_loader) n_iters = 0 with torch.no_grad(): for data in val_loader: n_iters += 1 early_reid = get_batch_mean_early_reid(reid_network, data['early_reid_patches']) curr_reid = reid_network(data['curr_reid_patch'].cuda()) conv_features, repr_features = get_features( obj_detect, data['curr_img'], data['curr_gt_app']) prev_loc = data['prev_gt_warped'].cuda() curr_loc = data['curr_gt_warped'].cuda() label_loc = data['label_gt'].cuda() curr_vis = data['curr_vis'].cuda() pred_loc_wh, vis = motion_model(early_reid, curr_reid, conv_features, repr_features, prev_loc, curr_loc) label_loc_wh = two_p_to_wh(label_loc) pred_loss = pred_loss_func(pred_loc_wh, label_loc_wh) loss_iters.append(pred_loss.item()) low_vis_ind = curr_vis < 0.3 if low_vis_ind.any(): low_vis_pred_loss = pred_loss_func(pred_loc_wh[low_vis_ind], label_loc_wh[low_vis_ind]) low_vis_loss_sum += (low_vis_pred_loss * torch.sum(low_vis_ind)).item() low_vis_num += torch.sum(low_vis_ind).item() high_vis_ind = curr_vis > 0.7 if high_vis_ind.any(): high_vis_pred_loss = pred_loss_func(pred_loc_wh[high_vis_ind], label_loc_wh[high_vis_ind]) high_vis_loss_sum += (high_vis_pred_loss * torch.sum(high_vis_ind)).item() high_vis_num += torch.sum(high_vis_ind).item() if n_iters % 50 == 0: print('Iter %5d/%5d finished.' % (n_iters, total_iters), flush=True) mean_loss = np.mean(loss_iters) mean_low_vis_loss = low_vis_loss_sum / low_vis_num mean_high_vis_loss = high_vis_loss_sum / high_vis_num print('All finished! Loss %.6f, low vis loss %.6f, high vis loss %.6f.' % (mean_loss, mean_low_vis_loss, mean_high_vis_loss))
def 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 main(dataset_names, prepr_w_tracktor, frcnn_prepr_params, tracktor_params, frcnn_weights, _config, _log, _run): sacred.commands.print_config(_run) if prepr_w_tracktor: prepr_params = tracktor_params else: prepr_params = frcnn_prepr_params make_deterministic(prepr_params['seed']) MOV_CAMERA_DICT = {**MOT15_MOV_CAMERA_DICT, **MOT17_MOV_CAMERA_DICT} # object detection _log.info("Initializing object detector.") obj_detect = FRCNN_FPN(num_classes=2) obj_detect.load_state_dict( torch.load(osp.join(OUTPUT_PATH, frcnn_weights), map_location=lambda storage, loc: storage)) obj_detect.eval() obj_detect.cuda() if prepr_w_tracktor: preprocessor = Tracker(obj_detect, None, prepr_params['tracker']) else: preprocessor = FRCNNPreprocessor(obj_detect, prepr_params) _log.info( f"Starting preprocessing of datasets {dataset_names} with {'Tracktor' if prepr_w_tracktor else 'FRCNN'} \n" ) for dataset_name in dataset_names: dataset = Datasets(dataset_name) _log.info( f"Preprocessing {len(dataset)} sequences from dataset {dataset_name} \n" ) time_total = 0 num_frames = 0 for seq in dataset: preprocessor.reset() start = time.time() _log.info(f"Preprocessing : {seq}") if prepr_w_tracktor: preprocessor.do_align = tracktor_params['tracker'][ 'do_align'] and (MOV_CAMERA_DICT[str(seq)]) data_loader = DataLoader(seq, batch_size=1, shuffle=False, num_workers=8, pin_memory=True) for i, frame in enumerate(tqdm(data_loader)): with torch.no_grad(): preprocessor.step(frame) num_frames += 1 time_total += time.time() - start _log.info(f"Runtime for {seq}: {time.time() - start :.1f} s.") output_file_path = osp.join(seq.seq_path, 'det', prepr_params['det_file_name']) if prepr_w_tracktor: results = preprocessor.get_results() seq.write_results(results, output_file_path) else: _log.info(f"Writing predictions in: {output_file_path}") preprocessor.save_results(output_file_path) _log.info( f"Tracking runtime for all sequences (without evaluation or image writing): " f"{time_total:.1f} s ({num_frames / time_total:.1f} Hz)")
if verbose: print(pairs_stored) # # For developing, only do the first couple. # if i == 101: # break # if seq == 'MOT17-02': # break h5.close() print(100 * '#') print("Finished creating dataset {}".format(filename)) # Load model print('Loading model...') obj_detect = FRCNN_FPN(num_classes=2, correlation_head=None) obj_detect.load_state_dict( torch.load( "output/faster_rcnn_fpn_training_mot_17/model_epoch_27_original.model", map_location=lambda storage, loc: storage)) obj_detect.eval() obj_detect.cuda() print('Model loaded!') # Hardcoded loader for MOT17 # sequences = ['MOT20-01', 'MOT20-02', 'MOT20-03', 'MOT20-05' ,'MOT17-02','MOT17-04', 'MOT17-05', 'MOT17-09', 'MOT17-10', 'MOT17-11', 'MOT17-13'] sequences = [ 'MOT17-13' ] #['MOT20-01','MOT20-02', 'MOT20-03', 'MOT20-05','MOT17-02','MOT17-04', 'MOT17-05', 'MOT17-09', 'MOT17-10', 'MOT17-11', 'MOT17-13'] # Hardcoded parameters
'MOT17-02', 'MOT17-04', 'MOT17-05', 'MOT17-09', 'MOT17-10', 'MOT17-11', 'MOT17-13' ] FEATURE_LEVEL = 3 DO_FLIP = False FP16 = True rcnn_transform: GeneralizedRCNNTransform = GeneralizedRCNNTransform( MIN_SIZE, MAX_SIZE, IMAGE_MEAN, IMAGE_STD) if DO_FLIP: transform = Compose([RandomHorizontalFlip(1.), ToTensor()]) else: transform = ToTensor() if FEATURE_LEVEL != -1: obj_detect: nn.Module = FRCNN_FPN(num_classes=2) obj_detect.load_state_dict(torch.load(WEIGHTS_PATH)) obj_detect.eval().cuda() # inspired by torchvision.models.detection.transform.GeneralizedRCNNTransform def batch_images(images, max_size, size_divisible=32): stride = size_divisible max_size = list(max_size) max_size[1] = int(math.ceil(float(max_size[1]) / stride) * stride) max_size[2] = int(math.ceil(float(max_size[2]) / stride) * stride) max_size = tuple(max_size) batch_shape = (len(images), ) + max_size batched_imgs = images[0].new(*batch_shape).zero_() for img, pad_img in zip(images, batched_imgs):
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(module_name, name, seed, obj_detect_models, reid_models, tracker, oracle, dataset, load_results, frame_range, interpolate, write_images, _config, _log, _run): sacred.commands.print_config(_run) # set all seeds torch.manual_seed(seed) torch.cuda.manual_seed(seed) np.random.seed(seed) torch.backends.cudnn.deterministic = True output_dir = osp.join(get_output_dir(module_name), name) sacred_config = osp.join(output_dir, 'sacred_config.yaml') if not osp.exists(output_dir): os.makedirs(output_dir) with open(sacred_config, 'w') as outfile: yaml.dump(copy.deepcopy(_config), outfile, default_flow_style=False) ########################## # Initialize the modules # ########################## # object detection _log.info("Initializing object detector(s).") obj_detects = [] for obj_detect_model in obj_detect_models: obj_detect = FRCNN_FPN(num_classes=2) obj_detect.load_state_dict( torch.load(obj_detect_model, map_location=lambda storage, loc: storage)) obj_detects.append(obj_detect) obj_detect.eval() if torch.cuda.is_available(): obj_detect.cuda() # reid _log.info("Initializing reID network(s).") reid_networks = [] for reid_model in reid_models: reid_cfg = os.path.join(os.path.dirname(reid_model), 'sacred_config.yaml') reid_cfg = yaml.safe_load(open(reid_cfg)) reid_network = ReIDNetwork_resnet50(pretrained=False, **reid_cfg['model_args']) reid_network.load_state_dict( torch.load(reid_model, map_location=lambda storage, loc: storage)) reid_network.eval() if torch.cuda.is_available(): reid_network.cuda() reid_networks.append(reid_network) # tracktor if oracle is not None: tracker = OracleTracker(obj_detect, reid_network, tracker, oracle) else: tracker = Tracker(obj_detect, reid_network, tracker) time_total = 0 num_frames = 0 mot_accums = [] dataset = Datasets(dataset) for seq, obj_detect, reid_network in zip(dataset, obj_detects, reid_networks): tracker.obj_detect = obj_detect tracker.reid_network = reid_network tracker.reset() _log.info(f"Tracking: {seq}") start_frame = int(frame_range['start'] * len(seq)) end_frame = int(frame_range['end'] * len(seq)) seq_loader = DataLoader( torch.utils.data.Subset(seq, range(start_frame, end_frame))) num_frames += len(seq_loader) results = {} if load_results: results = seq.load_results(output_dir) if not results: start = time.time() for frame_data in tqdm(seq_loader): with torch.no_grad(): tracker.step(frame_data) results = tracker.get_results() time_total += time.time() - start _log.info(f"Tracks found: {len(results)}") _log.info(f"Runtime for {seq}: {time.time() - start :.2f} s.") if interpolate: results = interpolate_tracks(results) _log.info(f"Writing predictions to: {output_dir}") seq.write_results(results, output_dir) if seq.no_gt: _log.info("No GT data for evaluation available.") else: mot_accums.append(get_mot_accum(results, seq_loader)) if write_images: plot_sequence(results, seq, osp.join(output_dir, str(dataset), str(seq)), write_images) if time_total: _log.info( f"Tracking runtime for all sequences (without evaluation or image writing): " f"{time_total:.2f} s for {num_frames} frames ({num_frames / time_total:.2f} Hz)" ) if mot_accums: _log.info("Evaluation:") evaluate_mot_accums(mot_accums, [str(s) for s in dataset if not s.no_gt], generate_overall=True)
def main(tracktor, reid, _config, _log, _run): sacred.commands.print_config(_run) # set all seeds torch.manual_seed(tracktor['seed']) torch.cuda.manual_seed(tracktor['seed']) np.random.seed(tracktor['seed']) torch.backends.cudnn.deterministic = True output_dir = osp.join(get_output_dir(tracktor['module_name']), tracktor['name']) sacred_config = osp.join(output_dir, 'sacred_config.yaml') if not osp.exists(output_dir): os.makedirs(output_dir) with open(sacred_config, 'w') as outfile: yaml.dump(_config, outfile, default_flow_style=False) ########################## # Initialize the modules # ########################## # object detection _log.info("Initializing object detector.") obj_detect = FRCNN_FPN(num_classes=2) obj_detect.load_state_dict(torch.load(_config['tracktor']['obj_detect_model'], map_location=lambda storage, loc: storage)) obj_detect.eval() obj_detect.cuda() # reid reid_network = resnet50(pretrained=False, **reid['cnn']) reid_network.load_state_dict(torch.load(tracktor['reid_weights'], map_location=lambda storage, loc: storage)) reid_network.eval() reid_network.cuda() # 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 test_tracktor_motion(val_loader, tracker_config, bbox_regression=True): obj_detect = FRCNN_FPN(num_classes=2) obj_detect.load_state_dict( torch.load(tracker_config['tracktor']['obj_detect_model'], map_location=lambda storage, loc: storage)) obj_detect.eval() obj_detect.cuda() pred_loss_func = nn.SmoothL1Loss() loss_iters = [] low_vis_loss_sum = 0.0 low_vis_num = 0 high_vis_loss_sum = 0.0 high_vis_num = 0 total_iters = len(val_loader) n_iters = 0 print(total_iters) with torch.no_grad(): for data in val_loader: n_iters += 1 prev_loc = data['prev_gt_warped'].cuda() curr_loc = data['curr_gt_warped'].cuda() label_loc = data['label_gt'].cuda() curr_vis = data['curr_vis'].cuda() pred_loc = curr_loc.clone() last_motion = curr_loc - prev_loc pred_loc += last_motion if bbox_regression: obj_detect.load_image(data['label_img'][0]) pred_loc, _ = obj_detect.predict_boxes(pred_loc) label_loc_wh = two_p_to_wh(label_loc) pred_loc_wh = two_p_to_wh(pred_loc) pred_loss = pred_loss_func(pred_loc_wh, label_loc_wh) loss_iters.append(pred_loss.item()) low_vis_ind = curr_vis < 0.3 if low_vis_ind.any(): low_vis_pred_loss = pred_loss_func(pred_loc_wh[low_vis_ind], label_loc_wh[low_vis_ind]) low_vis_loss_sum += (low_vis_pred_loss * torch.sum(low_vis_ind)).item() low_vis_num += torch.sum(low_vis_ind).item() high_vis_ind = curr_vis > 0.7 if high_vis_ind.any(): high_vis_pred_loss = pred_loss_func(pred_loc_wh[high_vis_ind], label_loc_wh[high_vis_ind]) high_vis_loss_sum += (high_vis_pred_loss * torch.sum(high_vis_ind)).item() high_vis_num += torch.sum(high_vis_ind).item() if n_iters % 500 == 0: print('Iter %5d/%5d finished.' % (n_iters, total_iters), flush=True) mean_loss = np.mean(loss_iters) mean_low_vis_loss = low_vis_loss_sum / low_vis_num mean_high_vis_loss = high_vis_loss_sum / high_vis_num print('All finished! Loss %.6f, low vis loss %.6f, high vis loss %.6f.' % (mean_loss, mean_low_vis_loss, mean_high_vis_loss))
def train_main(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 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(oracle_training, no_visrepr, max_previous_frame, use_ecc, use_modulator, use_bn, use_residual, vis_loss_ratio, no_vis_loss, lr, weight_decay, batch_size, output_dir, pretrain_vis_path, ex_name): random.seed(12345) torch.manual_seed(12345) torch.cuda.manual_seed(12345) np.random.seed(12345) torch.backends.cudnn.deterministic = True output_dir = osp.join(output_dir, ex_name) log_file = osp.join(output_dir, 'epoch_log.txt') if not osp.exists(output_dir): os.makedirs(output_dir) with open(log_file, 'w') as f: f.write('[Experiment name]%s\n\n' % ex_name) f.write('[Parameters]\n') f.write( 'oracle_training=%r\nno_visrepr=%r\nmax_previous_frame=%d\nuse_ecc=%r\nuse_modulator=%r\nuse_bn=%r\nuse_residual=%r\nvis_loss_ratio=%f\nno_vis_loss=%r\nlr=%f\nweight_decay=%f\nbatch_size=%d\n\n' % (oracle_training, no_visrepr, max_previous_frame, use_ecc, use_modulator, use_bn, use_residual, vis_loss_ratio, no_vis_loss, lr, weight_decay, batch_size)) f.write('[Loss log]\n') with open('experiments/cfgs/tracktor.yaml', 'r') as f: tracker_config = yaml.safe_load(f) ################# # Load Datasets # ################# train_set = MOT17TracksWrapper('train', 0.8, 0.0, input_track_len=max_previous_frame + 1, max_sample_frame=max_previous_frame, get_data_mode='sample' + (',ecc' if use_ecc else ''), tracker_cfg=tracker_config) train_loader = DataLoader(train_set, batch_size=batch_size, shuffle=True, num_workers=1, collate_fn=tracks_wrapper_collate) val_set = MOT17TracksWrapper('val', 0.8, 0.1, input_track_len=max_previous_frame + 1, max_sample_frame=max_previous_frame, get_data_mode='sample' + (',ecc' if use_ecc else ''), tracker_cfg=tracker_config) val_loader = DataLoader(val_set, batch_size=batch_size, shuffle=False, num_workers=1, collate_fn=tracks_wrapper_collate) with open( osp.join(cfg.ROOT_DIR, 'output', 'precomputed_ecc_matrices_3.pkl'), 'rb') as f: ecc_dict = pickle.load(f) train_set.load_precomputed_ecc_warp_matrices(ecc_dict) val_set.load_precomputed_ecc_warp_matrices(ecc_dict) ######################## # Initializing Modules # ######################## obj_detect = FRCNN_FPN(num_classes=2) obj_detect.load_state_dict( torch.load(tracker_config['tracktor']['obj_detect_model'], map_location=lambda storage, loc: storage)) obj_detect.eval() obj_detect.cuda() if oracle_training: motion_model = VisOracleMotionModel(vis_conv_only=False, use_modulator=use_modulator) else: if no_visrepr: motion_model = MotionModelNoVisRepr(vis_conv_only=False, use_modulator=use_modulator, use_bn=use_bn) else: motion_model = MotionModelV2(vis_conv_only=False, use_modulator=use_modulator, use_bn=use_bn, use_residual=use_residual) # motion_model.load_vis_pretrained(pretrain_vis_path) motion_model.train() motion_model.cuda() optimizer = torch.optim.Adam(motion_model.parameters(), lr=lr, weight_decay=weight_decay) pred_loss_func = nn.SmoothL1Loss() vis_loss_func = nn.MSELoss() ####################### # Training Parameters # ####################### max_epochs = 100 log_freq = 25 train_pred_loss_epochs = [] train_vis_loss_epochs = [] val_pred_loss_epochs = [] val_vis_loss_epochs = [] lowest_val_loss = 9999999.9 lowest_val_loss_epoch = -1 ############ # Training # ############ for epoch in range(max_epochs): n_iter = 0 train_pred_loss_iters = [] train_vis_loss_iters = [] val_pred_loss_iters = [] val_vis_loss_iters = [] for data, label in train_loader: # jitter bboxs for getting roi features im_w = torch.tensor([img.size()[-1] for img in data['curr_img']], dtype=data['curr_gt'].dtype) im_h = torch.tensor([img.size()[-2] for img in data['curr_img']], dtype=data['curr_gt'].dtype) jittered_curr_gt = bbox_jitter(data['curr_gt'].clone(), im_w, im_h) conv_features, repr_features = get_features( obj_detect, data['curr_img'], jittered_curr_gt) # for motion calculation, we still use the unjittered bboxs prev_loc = (data['prev_gt_warped'] if use_ecc else data['prev_gt']).cuda() curr_loc = (data['curr_gt_warped'] if use_ecc else data['curr_gt']).cuda() label_loc = label['label_gt'].cuda() curr_vis = data['curr_vis'].cuda() n_iter += 1 # TODO the output bbox should be (x,y,w,h)? optimizer.zero_grad() if oracle_training: pred_loc_wh, vis = motion_model(conv_features, repr_features, prev_loc, curr_loc, curr_vis.unsqueeze(-1)) else: pred_loc_wh, vis = motion_model(conv_features, repr_features, prev_loc, curr_loc) label_loc_wh = two_p_to_wh(label_loc) pred_loss = pred_loss_func(pred_loc_wh, label_loc_wh) vis_loss = vis_loss_func(vis, curr_vis) if no_vis_loss: loss = pred_loss else: loss = pred_loss + vis_loss_ratio * vis_loss loss.backward() optimizer.step() train_pred_loss_iters.append(pred_loss.item()) train_vis_loss_iters.append(vis_loss.item()) if n_iter % log_freq == 0: print( '[Train Iter %5d] train pred loss %.6f, vis loss %.6f ...' % (n_iter, np.mean(train_pred_loss_iters[n_iter - log_freq:n_iter]), np.mean(train_vis_loss_iters[n_iter - log_freq:n_iter])), flush=True) mean_train_pred_loss = np.mean(train_pred_loss_iters) mean_train_vis_loss = np.mean(train_vis_loss_iters) train_pred_loss_epochs.append(mean_train_pred_loss) train_vis_loss_epochs.append(mean_train_vis_loss) print('Train epoch %4d end.' % (epoch + 1)) motion_model.eval() with torch.no_grad(): for data, label in val_loader: # do not jitter for validation conv_features, repr_features = get_features( obj_detect, data['curr_img'], data['curr_gt']) prev_loc = (data['prev_gt_warped'] if use_ecc else data['prev_gt']).cuda() curr_loc = (data['curr_gt_warped'] if use_ecc else data['curr_gt']).cuda() label_loc = label['label_gt'].cuda() curr_vis = data['curr_vis'].cuda() if oracle_training: pred_loc_wh, vis = motion_model(conv_features, repr_features, prev_loc, curr_loc, curr_vis.unsqueeze(-1)) else: pred_loc_wh, vis = motion_model(conv_features, repr_features, prev_loc, curr_loc) label_loc_wh = two_p_to_wh(label_loc) pred_loss = pred_loss_func(pred_loc_wh, label_loc_wh) vis_loss = vis_loss_func(vis, curr_vis) val_pred_loss_iters.append(pred_loss.item()) val_vis_loss_iters.append(vis_loss.item()) mean_val_pred_loss = np.mean(val_pred_loss_iters) mean_val_vis_loss = np.mean(val_vis_loss_iters) val_pred_loss_epochs.append(mean_val_pred_loss) val_vis_loss_epochs.append(mean_val_vis_loss) print( '[Epoch %4d] train pred loss %.6f, vis loss %.6f; val pred loss %.6f, vis loss %.6f' % (epoch + 1, mean_train_pred_loss, mean_train_vis_loss, mean_val_pred_loss, mean_val_vis_loss)) with open(log_file, 'a') as f: f.write( 'Epoch %4d: train pred loss %.6f, vis loss %.6f; val pred loss %.6f, vis loss %.6f\n' % (epoch + 1, mean_train_pred_loss, mean_train_vis_loss, mean_val_pred_loss, mean_val_vis_loss)) motion_model.train() if mean_val_pred_loss < lowest_val_loss: lowest_val_loss, lowest_val_loss_epoch = mean_val_pred_loss, epoch + 1 torch.save( motion_model.state_dict(), osp.join(output_dir, 'motion_model_epoch_%d.pth' % (epoch + 1)))
def 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 pretrain_main(conv_only, image_flip, train_jitter, val_jitter, lr, weight_decay, batch_size, output_dir, ex_name): random.seed(12345) torch.manual_seed(12345) torch.cuda.manual_seed(12345) np.random.seed(12345) torch.backends.cudnn.deterministic = True # output_dir = osp.join(get_output_dir('motion'), 'vis') output_dir = osp.join(output_dir, ex_name) log_file = osp.join(output_dir, 'epoch_log.txt') if not osp.exists(output_dir): os.makedirs(output_dir) with open(log_file, 'w') as f: f.write('[Experiment name]%s\n\n' % ex_name) f.write('[Parameters]\n') f.write( 'conv_only=%r\nimage_flip=%r\ntrain_jitter=%r\nval_jitter=%r\nlr=%f\nweight_decay=%f\nbatch_size=%d\n\n' % (conv_only, image_flip, train_jitter, val_jitter, lr, weight_decay, batch_size)) f.write('[Loss log]\n') with open('experiments/cfgs/tracktor.yaml', 'r') as f: tracker_config = yaml.safe_load(f) ################# # Load Datasets # ################# train_set = MOT17Vis('train', 0.8, 0.0, train_bbox_jitter=train_jitter, random_image_flip=image_flip) train_loader = DataLoader(train_set, batch_size=1, shuffle=True, num_workers=2) val_set = MOT17Vis('val', 0.8, 0.0, val_bbox_jitter=val_jitter) val_loader = DataLoader(val_set, batch_size=1, shuffle=False, num_workers=2) ######################## # Initializing Modules # ######################## obj_detect = FRCNN_FPN(num_classes=2) obj_detect.load_state_dict( torch.load(tracker_config['tracktor']['obj_detect_model'], map_location=lambda storage, loc: storage)) obj_detect.eval() obj_detect.cuda() vis_model = VisEst(conv_only=conv_only) vis_model.train() vis_model.cuda() optimizer = torch.optim.Adam(vis_model.parameters(), lr=lr, weight_decay=weight_decay) loss_func = nn.MSELoss() ####################### # Training Parameters # ####################### max_epochs = 100 conv_batch_forger = BatchForger( batch_size, (vis_model.output_dim, vis_model.pool_size, vis_model.pool_size)) repr_batch_forger = BatchForger(batch_size, (vis_model.representation_dim, )) label_batch_forger = BatchForger(batch_size, (1, )) log_freq = 100 train_loss_epochs = [] val_loss_epochs = [] lowest_val_loss = 9999999.9 lowest_val_loss_epoch = -1 ############ # Training # ############ for epoch in range(max_epochs): n_iter = 0 train_loss_iters = [] val_loss_iters = [] for data in train_loader: conv_features, repr_features = get_features( obj_detect, data['img'], data['gt']) conv_batch_forger.feed(conv_features) repr_batch_forger.feed(repr_features) label_batch_forger.feed(data['vis'].squeeze(0).unsqueeze(-1)) while label_batch_forger.has_one_batch(): n_iter += 1 batch_conv = conv_batch_forger.dump().cuda() batch_repr = repr_batch_forger.dump().cuda() batch_label = label_batch_forger.dump().cuda() optimizer.zero_grad() pred, _ = vis_model(batch_conv, batch_repr) loss = loss_func(pred, batch_label) loss.backward() optimizer.step() train_loss_iters.append(loss.item()) if n_iter % log_freq == 0: print( '[Train Iter %5d] train loss %.6f ...' % (n_iter, np.mean(train_loss_iters[n_iter - log_freq:n_iter]))) mean_train_loss = np.mean(train_loss_iters) train_loss_epochs.append(mean_train_loss) print('Train epoch %4d end.' % (epoch + 1)) conv_batch_forger.reset() repr_batch_forger.reset() label_batch_forger.reset() vis_model.eval() for data in val_loader: conv_features, repr_features = get_features( obj_detect, data['img'], data['gt']) conv_batch_forger.feed(conv_features) repr_batch_forger.feed(repr_features) label_batch_forger.feed(data['vis'].squeeze(0).unsqueeze(-1)) while label_batch_forger.has_one_batch(): batch_conv = conv_batch_forger.dump().cuda() batch_repr = repr_batch_forger.dump().cuda() batch_label = label_batch_forger.dump().cuda() pred, _ = vis_model(batch_conv, batch_repr) loss = loss_func(pred, batch_label) val_loss_iters.append(loss.item()) mean_val_loss = np.mean(val_loss_iters) val_loss_epochs.append(mean_val_loss) print('[Epoch %4d] train loss %.6f, val loss %.6f' % (epoch + 1, mean_train_loss, mean_val_loss)) with open(log_file, 'a') as f: f.write('Epoch %4d: train loss %.6f, val loss %.6f\n' % (epoch + 1, mean_train_loss, mean_val_loss)) conv_batch_forger.reset() repr_batch_forger.reset() label_batch_forger.reset() vis_model.train() if mean_val_loss < lowest_val_loss: lowest_val_loss, lowest_val_loss_epoch = mean_val_loss, epoch + 1 torch.save( vis_model.state_dict(), osp.join(output_dir, 'vis_model_epoch_%d.pth' % (epoch + 1)))
def main(tracktor, reid, _config, _log, _run): sacred.commands.print_config(_run) # set all seeds torch.manual_seed(tracktor['seed']) torch.cuda.manual_seed(tracktor['seed']) np.random.seed(tracktor['seed']) torch.backends.cudnn.deterministic = True output_dir = osp.join(get_output_dir(tracktor['module_name']), tracktor['name']) sacred_config = osp.join(output_dir, 'sacred_config.yaml') if not osp.exists(output_dir): os.makedirs(output_dir) with open(sacred_config, 'w') as outfile: yaml.dump(_config, outfile, default_flow_style=False) ########################## # Initialize the modules # ########################## # object detection _log.info("Initializing object detector.") obj_detect = FRCNN_FPN(num_classes=2) obj_detect.load_state_dict(torch.load(_config['tracktor']['obj_detect_model'], map_location=lambda storage, loc: storage)) obj_detect.eval() obj_detect.cuda() # reid reid_network = resnet50(pretrained=False, **reid['cnn']) reid_network.load_state_dict(torch.load(tracktor['reid_weights'], map_location=lambda storage, loc: storage)) reid_network.eval() reid_network.cuda() # neural motion model vis_model = VisSimpleReID() motion_model = MotionModelV3(vis_model) motion_model.load_state_dict(torch.load('output/motion/finetune_motion_model_v3.pth')) motion_model.eval() motion_model.cuda() save_vis_results = False # tracktor if 'oracle' in tracktor: tracker = OracleTracker(obj_detect, reid_network, tracktor['tracker'], tracktor['oracle']) else: # tracker = Tracker(obj_detect, reid_network, tracktor['tracker']) tracker = TrackerNeuralMM(obj_detect, reid_network, motion_model, tracktor['tracker'], save_vis_results=save_vis_results, vis_model=None) time_total = 0 num_frames = 0 mot_accums = [] dataset = Datasets(tracktor['dataset'], {'use_val_split':True}) for seq in dataset: tracker.reset() start = time.time() _log.info(f"Tracking: {seq}") data_loader = DataLoader(seq, batch_size=1, shuffle=False) for i, frame in enumerate(tqdm(data_loader)): if len(seq) * tracktor['frame_split'][0] <= i <= len(seq) * tracktor['frame_split'][1]: with torch.no_grad(): tracker.step(frame) num_frames += 1 results = tracker.get_results() time_total += time.time() - start _log.info(f"Tracks found: {len(results)}") _log.info(f"Runtime for {seq}: {time.time() - start :.1f} s.") if tracktor['interpolate']: results = interpolate(results) if seq.no_gt: _log.info(f"No GT data for evaluation available.") else: mot_accums.append(get_mot_accum(results, seq)) _log.info(f"Writing predictions to: {output_dir}") seq.write_results(results, output_dir) if save_vis_results: vis_results = tracker.get_vis_results() seq.write_vis_results(vis_results, output_dir) if tracktor['write_images']: plot_sequence(results, seq, osp.join(output_dir, tracktor['dataset'], str(seq))) _log.info(f"Tracking runtime for all sequences (without evaluation or image writing): " f"{time_total:.1f} s ({num_frames / time_total:.1f} Hz)") if mot_accums: evaluate_mot_accums(mot_accums, [str(s) for s in dataset if not s.no_gt], generate_overall=True)
def main(tracktor, reid, _config, _log, _run): sacred.commands.print_config(_run) # set all seeds torch.manual_seed(tracktor['seed']) torch.cuda.manual_seed(tracktor['seed']) np.random.seed(tracktor['seed']) torch.backends.cudnn.deterministic = True output_dir = osp.join(get_output_dir(tracktor['module_name']), tracktor['name']) sacred_config = osp.join(output_dir, 'sacred_config.yaml') if not osp.exists(output_dir): os.makedirs(output_dir) with open(sacred_config, 'w') as outfile: yaml.dump(_config, outfile, default_flow_style=False) ########################## # Initialize the modules # ########################## # object detection _log.info("Initializing object detector.") use_masks = _config['tracktor']['tracker']['use_masks'] mask_model = Mask_RCNN(num_classes=2) fast_model = FRCNN_FPN(num_classes=2) fast_model.load_state_dict(torch.load(_config['tracktor']['fast_rcnn_model'], map_location=lambda storage, loc: storage)) if(use_masks): mask_model.load_state_dict(torch.load(_config['tracktor']['mask_rcnn_model'], map_location=lambda storage, loc: storage)['model_state_dict']) mask_model.eval() mask_model.cuda() fast_model.eval() fast_model.cuda() # reid reid_network = resnet50(pretrained=False, **reid['cnn']) reid_network.load_state_dict(torch.load(tracktor['reid_weights'], map_location=lambda storage, loc: storage)) reid_network.eval() reid_network.cuda() # tracktor if 'oracle' in tracktor: tracker = OracleTracker(fast_model, reid_network, tracktor['tracker'], tracktor['oracle']) else: tracker = Tracker(fast_model, reid_network, tracktor['tracker'], mask_model) time_total = 0 num_frames = 0 mot_accums = [] dataset = Datasets(tracktor['dataset']) for seq in dataset: num_frames = 0 tracker.reset() start = time.time() _log.info(f"Tracking: {seq}") data_loader = DataLoader(seq, batch_size=1, shuffle=False) if tracktor['write_images'] and use_masks: print("[*] Plotting image to {}".format(osp.join(output_dir, tracktor['dataset']))) for i, frame in enumerate(tqdm(data_loader)): if len(seq) * tracktor['frame_split'][0] <= i <= len(seq) * tracktor['frame_split'][1]: tracker.step(frame) if tracktor['write_images'] and use_masks: result = tracker.get_results() masks = tracker.get_masks() plot_sequence(result, masks, seq, num_frames, osp.join(output_dir, tracktor['dataset'], str(seq)), plot_masks = True) num_frames += 1 results = tracker.get_results() import matplotlib.pyplot as plt time_total += time.time() - start _log.info(f"Tracks found: {len(results)}") _log.info(f"Runtime for {seq}: {time.time() - start :.1f} s.") if tracktor['interpolate']: results = interpolate(results) if seq.no_gt: _log.info(f"No GT data for evaluation available.") else: mot_accums.append(get_mot_accum(results, seq)) _log.info(f"Writing predictions to: {output_dir}") seq.write_results(results, output_dir) _log.info(f"Tracking runtime for all sequences (without evaluation or image writing): " f"{time_total:.1f} s ({num_frames / time_total:.1f} Hz)") if mot_accums: evaluate_mot_accums(mot_accums, [str(s) for s in dataset if not s.no_gt], generate_overall=True)