def __init__(self, n, configs): self.node = n #self.pub = rospy.Publisher('tracking_result', Object_with_id, queue_size=10) self.model = create_model(configs) self.model.print_network() print('\n\n' + '-*=' * 30 + '\n\n') assert os.path.isfile(configs.pretrained_path), "No file at {}".format(configs.pretrained_path) self.model.load_state_dict(torch.load(configs.pretrained_path,map_location='cuda:0')) configs.device = torch.device('cpu' if configs.no_cuda else 'cuda:{}'.format(configs.gpu_idx)) self.model = self.model.to(device=configs.device) self.model.eval() rospy.Subscriber("detect_object", PointCloud2, self.callback) # your cloud topic name rospy.spin()
def __init__(self): self.configs = parse_test_configs() self.configs.distributed = False # For testing host = self.configs.host port = self.configs.port print("> server start.... ") socketer = socket.socket(socket.AF_INET, socket.SOCK_STREAM) # set the port reuesd socketer.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) socketer.bind((host, port)) # define the max connection socketer.listen(10) self.sock = socketer # Model Initial self.model = create_model(self.configs) self.model.print_network() print('\n\n' + '-*=' * 30 + '\n\n') device_string = 'cpu' if self.configs.no_cuda else 'cuda:{}'.format( self.configs.gpu_idx) assert os.path.isfile(self.configs.pretrained_path), "No file at {}".format( self.configs.pretrained_path) # model.load_state_dict(torch.load(configs.pretrained_path)) self.model.load_state_dict(torch.load( self.configs.pretrained_path, map_location=device_string)) # configs.device = torch.device('cpu' if configs.no_cuda else 'cuda:{}'.format(configs.gpu_idx)) self.configs.device = torch.device(device_string) self.model = self.model.to(device=self.configs.device) self.model.eval() # self.test_dataloader = create_test_dataloader(self.configs) # self.test_dataloader_iter = self.test_dataloader._get_iterator() self.test_dataset = create_test_dataset(self.configs) self.batch_idx = 0 self.need_create_window = True
def main_worker(gpu_idx, configs): configs.gpu_idx = gpu_idx if configs.gpu_idx is not None: print("Use GPU: {} for training".format(configs.gpu_idx)) configs.device = torch.device('cuda:{}'.format(configs.gpu_idx)) if configs.distributed: if configs.dist_url == "env://" and configs.rank == -1: configs.rank = int(os.environ["RANK"]) if configs.multiprocessing_distributed: configs.rank = configs.rank * configs.ngpus_per_node + gpu_idx dist.init_process_group(backend=configs.dist_backend, init_method=configs.dist_url, world_size=configs.world_size, rank=configs.rank) configs.is_master_node = (not configs.distributed) or ( configs.distributed and (configs.rank % configs.ngpus_per_node == 0)) # model model = create_model(configs) model = make_data_parallel(model, configs) if configs.is_master_node: num_parameters = get_num_parameters(model) print('number of trained parameters of the model: {}'.format( num_parameters)) if configs.pretrained_path is not None: model = load_pretrained_model(model, configs.pretrained_path, gpu_idx, configs.overwrite_global_2_local) # Load dataset test_loader = create_test_dataloader(configs) test(test_loader, model, configs)
detections = post_processing(detections, configs.num_classes, configs.down_ratio, configs.peak_thresh) t2 = time_synchronized() detections = detections[0] # only first batch # Draw prediction in the image bev_map = (bev_maps.squeeze().permute(1, 2, 0).numpy() * 255).astype(np.uint8) bev_map = cv2.resize(bev_map, (cnf.BEV_WIDTH, cnf.BEV_HEIGHT)) bev_map = draw_predictions(bev_map, detections.copy(), configs.num_classes) bev_map = cv2.rotate(bev_map, cv2.ROTATE_180) cv2.imshow("BEV", bev_map) print('\tDone testing in time: {:.1f}ms, speed {:.2f}FPS'.format((t2 - t1) * 1000,1 / (t2 - t1))) if __name__ == '__main__': stereo_args= parse_config() configs, args = parse_test_configs() model = create_model(configs) print('\n\n' + '-*=' * 30 + '\n\n') assert os.path.isfile(configs.pretrained_path), "No file at {}".format(configs.pretrained_path) model.load_state_dict(torch.load(configs.pretrained_path, map_location='cpu')) print('Loaded weights from {}\n'.format(configs.pretrained_path)) configs.device = torch.device('cpu' if configs.no_cuda else 'cuda:0') model = model.to(device=configs.device) model.eval() # ============================= Sersy Edit ======================================= stereo_model = Stereo_Depth_Estimation(stereo_args,None) dataset_root = os.path.join(configs.dataset_dir, 'testing') print(configs.dataset_dir) KITTI_stereo = KittiDataset(dataset_root, stereo_mode=True) KITTI = KittiDataset(dataset_root, stereo_mode=False)
def main_worker(gpu_idx, configs): configs.gpu_idx = gpu_idx if configs.gpu_idx is not None: print("Use GPU: {} for training".format(configs.gpu_idx)) configs.device = torch.device('cuda:{}'.format(configs.gpu_idx)) if configs.distributed: if configs.dist_url == "env://" and configs.rank == -1: configs.rank = int(os.environ["RANK"]) if configs.multiprocessing_distributed: # For multiprocessing distributed training, rank needs to be the # global rank among all the processes configs.rank = configs.rank * configs.ngpus_per_node + gpu_idx dist.init_process_group(backend=configs.dist_backend, init_method=configs.dist_url, world_size=configs.world_size, rank=configs.rank) configs.is_master_node = (not configs.distributed) or ( configs.distributed and (configs.rank % configs.ngpus_per_node == 0)) if configs.is_master_node: logger = Logger(configs.logs_dir, configs.saved_fn) logger.info('>>> Created a new logger') logger.info('>>> configs: {}'.format(configs)) tb_writer = SummaryWriter( log_dir=os.path.join(configs.logs_dir, 'tensorboard')) else: logger = None tb_writer = None # model model = create_model(configs) # Data Parallel model = make_data_parallel(model, configs) # Freeze model model = freeze_model(model, configs.freeze_modules_list) if configs.is_master_node: num_parameters = get_num_parameters(model) logger.info('number of trained parameters of the model: {}'.format( num_parameters)) optimizer = create_optimizer(configs, model) lr_scheduler = create_lr_scheduler(optimizer, configs) best_val_loss = np.inf earlystop_count = 0 is_best = False # optionally load weight from a checkpoint if configs.pretrained_path is not None: model = load_pretrained_model(model, configs.pretrained_path, gpu_idx, configs.overwrite_global_2_local) if logger is not None: logger.info('loaded pretrained model at {}'.format( configs.pretrained_path)) # optionally resume from a checkpoint if configs.resume_path is not None: checkpoint = resume_model(configs.resume_path, configs.arch, configs.gpu_idx) if hasattr(model, 'module'): model.module.load_state_dict(checkpoint['state_dict']) else: model.load_state_dict(checkpoint['state_dict']) optimizer.load_state_dict(checkpoint['optimizer']) lr_scheduler.load_state_dict(checkpoint['lr_scheduler']) best_val_loss = checkpoint['best_val_loss'] earlystop_count = checkpoint['earlystop_count'] configs.start_epoch = checkpoint['epoch'] + 1 if logger is not None: logger.info(">>> Loading dataset & getting dataloader...") # Create dataloader train_loader, val_loader, train_sampler = create_train_val_dataloader( configs) test_loader = create_test_dataloader(configs) if logger is not None: logger.info('number of batches in train set: {}'.format( len(train_loader))) if val_loader is not None: logger.info('number of batches in val set: {}'.format( len(val_loader))) logger.info('number of batches in test set: {}'.format( len(test_loader))) if configs.evaluate: assert val_loader is not None, "The validation should not be None" val_loss = evaluate_one_epoch(val_loader, model, configs.start_epoch - 1, configs, logger) print('Evaluate, val_loss: {}'.format(val_loss)) return for epoch in range(configs.start_epoch, configs.num_epochs + 1): # Get the current learning rate for param_group in optimizer.param_groups: lr = param_group['lr'] if logger is not None: logger.info('{}'.format('*-' * 40)) logger.info('{} {}/{} {}'.format('=' * 35, epoch, configs.num_epochs, '=' * 35)) logger.info('{}'.format('*-' * 40)) logger.info('>>> Epoch: [{}/{}] learning rate: {:.2e}'.format( epoch, configs.num_epochs, lr)) if configs.distributed: train_sampler.set_epoch(epoch) # train for one epoch train_loss = train_one_epoch(train_loader, model, optimizer, epoch, configs, logger) loss_dict = {'train': train_loss} if not configs.no_val: val_loss = evaluate_one_epoch(val_loader, model, epoch, configs, logger) is_best = val_loss <= best_val_loss best_val_loss = min(val_loss, best_val_loss) loss_dict['val'] = val_loss if not configs.no_test: test_loss = evaluate_one_epoch(test_loader, model, epoch, configs, logger) loss_dict['test'] = test_loss # Write tensorboard if tb_writer is not None: tb_writer.add_scalars('Loss', loss_dict, epoch) # Save checkpoint if configs.is_master_node and (is_best or ( (epoch % configs.checkpoint_freq) == 0)): saved_state = get_saved_state(model, optimizer, lr_scheduler, epoch, configs, best_val_loss, earlystop_count) save_checkpoint(configs.checkpoints_dir, configs.saved_fn, saved_state, is_best, epoch) # Check early stop training if configs.earlystop_patience is not None: earlystop_count = 0 if is_best else (earlystop_count + 1) print_string = ' |||\t earlystop_count: {}'.format(earlystop_count) if configs.earlystop_patience <= earlystop_count: print_string += '\n\t--- Early stopping!!!' break else: print_string += '\n\t--- Continue training..., earlystop_count: {}'.format( earlystop_count) if logger is not None: logger.info(print_string) # Adjust learning rate if configs.lr_type == 'plateau': assert (not configs.no_val ), "Only use plateau when having validation set" lr_scheduler.step(val_loss) else: lr_scheduler.step() if tb_writer is not None: tb_writer.close() if configs.distributed: cleanup()
def main_worker(gpu_idx, configs): configs.gpu_idx = gpu_idx if configs.gpu_idx is not None: print("Use GPU: {} for training".format(configs.gpu_idx)) configs.device = torch.device('cuda:{}'.format(configs.gpu_idx)) if configs.distributed: if configs.dist_url == "env://" and configs.rank == -1: configs.rank = int(os.environ["RANK"]) if configs.multiprocessing_distributed: # For multiprocessing distributed training, rank needs to be the # global rank among all the processes configs.rank = configs.rank * configs.ngpus_per_node + gpu_idx dist.init_process_group(backend=configs.dist_backend, init_method=configs.dist_url, world_size=configs.world_size, rank=configs.rank) configs.is_master_node = (not configs.distributed) or ( configs.distributed and (configs.rank % configs.ngpus_per_node == 0)) if configs.is_master_node: logger = Logger(configs.logs_dir, configs.saved_fn) logger.info('>>> Created a new logger') logger.info('>>> configs: {}'.format(configs)) tb_writer = SummaryWriter( log_dir=os.path.join(configs.logs_dir, 'tensorboard')) else: logger = None tb_writer = None # model model = create_model(configs) # load weight from a checkpoint if configs.pretrained_path is not None: assert os.path.isfile( configs.pretrained_path), "=> no checkpoint found at '{}'".format( configs.pretrained_path) model.load_weights(weightfile=configs.pretrained_path) if logger is not None: logger.info('loaded pretrained model at {}'.format( configs.pretrained_path)) # resume weights of model from a checkpoint if configs.resume_path is not None: assert os.path.isfile( configs.resume_path), "=> no checkpoint found at '{}'".format( configs.resume_path) model.load_weights(weightfile=configs.resume_path) if logger is not None: logger.info('resume training model from checkpoint {}'.format( configs.pretrained_path)) # Data Parallel model = make_data_parallel(model, configs) # Make sure to create optimizer after moving the model to cuda optimizer = create_optimizer(configs, model) lr_scheduler = create_lr_scheduler(optimizer, configs) # resume optimizer, lr_scheduler from a checkpoint if configs.resume_path is not None: utils_path = configs.resume_path.replace('Model_', 'Utils_') assert os.path.isfile( utils_path), "=> no checkpoint found at '{}'".format(utils_path) utils_state_dict = torch.load(utils_path, map_location='cuda:{}'.format( configs.gpu_idx)) optimizer.load_state_dict(utils_state_dict['optimizer']) lr_scheduler.load_state_dict(utils_state_dict['lr_scheduler']) configs.start_epoch = utils_state_dict['epoch'] + 1 if configs.is_master_node: num_parameters = get_num_parameters(model) logger.info('number of trained parameters of the model: {}'.format( num_parameters)) if logger is not None: logger.info(">>> Loading dataset & getting dataloader...") # Create dataloader train_loader, val_loader, train_sampler = create_train_val_dataloader( configs) if logger is not None: logger.info('number of batches in train set: {}'.format( len(train_loader))) if val_loader is not None: logger.info('number of batches in val set: {}'.format( len(val_loader))) if configs.evaluate: assert val_loader is not None, "The validation should not be None" eval_metrics = evaluate_one_epoch(val_loader, model, configs.start_epoch - 1, configs, logger) precision, recall, AP, f1, ap_class = eval_metrics print( 'Evaluate - precision: {}, recall: {}, AP: {}, f1: {}, ap_class: {}' .format(precision, recall, AP, f1, ap_class)) return for epoch in range(configs.start_epoch, configs.num_epochs + 1): if logger is not None: logger.info('{}'.format('*-' * 40)) logger.info('{} {}/{} {}'.format('=' * 35, epoch, configs.num_epochs, '=' * 35)) logger.info('{}'.format('*-' * 40)) logger.info('>>> Epoch: [{}/{}]'.format(epoch, configs.num_epochs)) if configs.distributed: train_sampler.set_epoch(epoch) # train for one epoch train_one_epoch(train_loader, model, optimizer, lr_scheduler, epoch, configs, logger, tb_writer) if not configs.no_val: precision, recall, AP, f1, ap_class = evaluate_one_epoch( val_loader, model, epoch, configs, logger) val_metrics_dict = { 'precision': precision, 'recall': recall, 'AP': AP, 'f1': f1, 'ap_class': ap_class } if tb_writer is not None: tb_writer.add_scalars('Validation', val_metrics_dict, epoch) # Save checkpoint if configs.is_master_node and ((epoch % configs.checkpoint_freq) == 0): model_state_dict, utils_state_dict = get_saved_state( model, optimizer, lr_scheduler, epoch, configs) save_checkpoint(configs.checkpoints_dir, configs.saved_fn, model_state_dict, utils_state_dict, epoch) if tb_writer is not None: tb_writer.close() if configs.distributed: cleanup()
def main_worker(gpu_idx, configs): configs.gpu_idx = gpu_idx configs.device = torch.device('cpu' if configs.gpu_idx is None else 'cuda:{}'.format(configs.gpu_idx)) if configs.distributed: if configs.dist_url == "env://" and configs.rank == -1: configs.rank = int(os.environ["RANK"]) if configs.multiprocessing_distributed: # For multiprocessing distributed training, rank needs to be the # global rank among all the processes configs.rank = configs.rank * configs.ngpus_per_node + gpu_idx dist.init_process_group(backend=configs.dist_backend, init_method=configs.dist_url, world_size=configs.world_size, rank=configs.rank) configs.subdivisions = int(64 / configs.batch_size / configs.ngpus_per_node) else: configs.subdivisions = int(64 / configs.batch_size) configs.is_master_node = (not configs.distributed) or ( configs.distributed and (configs.rank % configs.ngpus_per_node == 0)) if configs.is_master_node: logger = Logger(configs.logs_dir, configs.saved_fn) logger.info('>>> Created a new logger') logger.info('>>> configs: {}'.format(configs)) tb_writer = SummaryWriter(log_dir=os.path.join(configs.logs_dir, 'tensorboard')) else: logger = None tb_writer = None # model model = create_model(configs) # load weight from a checkpoint if configs.pretrained_path is not None: assert os.path.isfile(configs.pretrained_path), "=> no checkpoint found at '{}'".format(configs.pretrained_path) model.load_state_dict(torch.load(configs.pretrained_path, map_location='cpu')) if logger is not None: logger.info('loaded pretrained model at {}'.format(configs.pretrained_path)) # resume weights of model from a checkpoint if configs.resume_path is not None: assert os.path.isfile(configs.resume_path), "=> no checkpoint found at '{}'".format(configs.resume_path) model.load_state_dict(torch.load(configs.resume_path, map_location='cpu')) if logger is not None: logger.info('resume training model from checkpoint {}'.format(configs.resume_path)) # Data Parallel model = make_data_parallel(model, configs) # Make sure to create optimizer after moving the model to cuda optimizer = create_optimizer(configs, model) lr_scheduler = create_lr_scheduler(optimizer, configs) configs.step_lr_in_epoch = False if configs.lr_type in ['multi_step', 'cosin', 'one_cycle'] else True # resume optimizer, lr_scheduler from a checkpoint if configs.resume_path is not None: utils_path = configs.resume_path.replace('Model_', 'Utils_') assert os.path.isfile(utils_path), "=> no checkpoint found at '{}'".format(utils_path) utils_state_dict = torch.load(utils_path, map_location='cuda:{}'.format(configs.gpu_idx)) optimizer.load_state_dict(utils_state_dict['optimizer']) lr_scheduler.load_state_dict(utils_state_dict['lr_scheduler']) configs.start_epoch = utils_state_dict['epoch'] + 1 if configs.is_master_node: num_parameters = get_num_parameters(model) logger.info('number of trained parameters of the model: {}'.format(num_parameters)) if logger is not None: logger.info(">>> Loading dataset & getting dataloader...") # Create dataloader train_dataloader, train_sampler = create_train_dataloader(configs) if logger is not None: logger.info('number of batches in training set: {}'.format(len(train_dataloader))) if configs.evaluate: val_dataloader = create_val_dataloader(configs) val_loss = validate(val_dataloader, model, configs) print('val_loss: {:.4e}'.format(val_loss)) return for epoch in range(configs.start_epoch, configs.num_epochs + 1): if logger is not None: logger.info('{}'.format('*-' * 40)) logger.info('{} {}/{} {}'.format('=' * 35, epoch, configs.num_epochs, '=' * 35)) logger.info('{}'.format('*-' * 40)) logger.info('>>> Epoch: [{}/{}]'.format(epoch, configs.num_epochs)) if configs.distributed: train_sampler.set_epoch(epoch) # train for one epoch train_one_epoch(train_dataloader, model, optimizer, lr_scheduler, epoch, configs, logger, tb_writer) if (not configs.no_val) and (epoch % configs.checkpoint_freq == 0): val_dataloader = create_val_dataloader(configs) print('number of batches in val_dataloader: {}'.format(len(val_dataloader))) val_loss = validate(val_dataloader, model, configs) print('val_loss: {:.4e}'.format(val_loss)) if tb_writer is not None: tb_writer.add_scalar('Val_loss', val_loss, epoch) # Save checkpoint if configs.is_master_node and ((epoch % configs.checkpoint_freq) == 0): model_state_dict, utils_state_dict = get_saved_state(model, optimizer, lr_scheduler, epoch, configs) save_checkpoint(configs.checkpoints_dir, configs.saved_fn, model_state_dict, utils_state_dict, epoch) if not configs.step_lr_in_epoch: lr_scheduler.step() if tb_writer is not None: tb_writer.add_scalar('LR', lr_scheduler.get_lr()[0], epoch) if tb_writer is not None: tb_writer.close() if configs.distributed: cleanup()
return configs if __name__ == '__main__': configs = parse_eval_configs() configs.distributed = False # For evaluation class_names = load_classes(configs.classnames_infor_path) print(configs.iou_thresh) voxel_generator = VoxelGeneratorV2( voxel_size=list(configs.voxel_size), point_cloud_range=list(configs.point_cloud_range), max_num_points=configs.max_number_of_points_per_voxel, max_voxels=20000) model = create_model(configs, voxel_generator) print('\n\n' + '-*=' * 30 + '\n\n') assert os.path.isfile(configs.pretrained_path), "No file at {}".format( configs.pretrained_path) model.load_state_dict( torch.load(configs.pretrained_path, map_location='cpu')) print('Loaded weights from {}\n'.format(configs.pretrained_path)) configs.device = torch.device( 'cpu' if configs.no_cuda else 'cuda:{}'.format(configs.gpu_idx)) model = model.to(device=configs.device) out_cap = None model.eval()
def demo(configs): video_loader = TTNet_Video_Loader(configs.video_path, configs.input_size, configs.num_frames_sequence) result_filename = os.path.join(configs.save_demo_dir, 'results.txt') frame_rate = video_loader.video_fps if configs.save_demo_output: configs.frame_dir = os.path.join(configs.save_demo_dir, 'frame') if not os.path.isdir(configs.frame_dir): os.makedirs(configs.frame_dir) configs.device = torch.device('cuda:{}'.format(configs.gpu_idx)) # model model = create_model(configs) model.cuda() assert configs.pretrained_path is not None, "Need to load the pre-trained model" model = load_pretrained_model(model, configs.pretrained_path, configs.gpu_idx, configs.overwrite_global_2_local) model.eval() middle_idx = int(configs.num_frames_sequence / 2) queue_frames = deque(maxlen=middle_idx + 1) frame_idx = 0 w_original, h_original = 1920, 1080 w_resize, h_resize = 320, 128 w_ratio = w_original / w_resize h_ratio = h_original / h_resize with torch.no_grad(): for count, resized_imgs in video_loader: # take the middle one img = cv2.resize(resized_imgs[3 * middle_idx: 3 * (middle_idx + 1)].transpose(1, 2, 0), (w_original, h_original)) # Expand the first dim resized_imgs = torch.from_numpy(resized_imgs).to(configs.device, non_blocking=True).float().unsqueeze(0) t1 = time_synchronized() pred_ball_global, pred_ball_local, pred_events, pred_seg = model.run_demo(resized_imgs) t2 = time_synchronized() prediction_global, prediction_local, prediction_seg, prediction_events = post_processing( pred_ball_global, pred_ball_local, pred_events, pred_seg, configs.input_size[0], configs.thresh_ball_pos_mask, configs.seg_thresh, configs.event_thresh) prediction_ball_final = [ int(prediction_global[0] * w_ratio + prediction_local[0] - w_resize / 2), int(prediction_global[1] * h_ratio + prediction_local[1] - h_resize / 2) ] # Get infor of the (middle_idx + 1)th frame if len(queue_frames) == middle_idx + 1: frame_pred_infor = queue_frames.popleft() seg_img = frame_pred_infor['seg'].astype(np.uint8) ball_pos = frame_pred_infor['ball'] seg_img = cv2.resize(seg_img, (w_original, h_original)) ploted_img = plot_detection(img, ball_pos, seg_img, prediction_events) ploted_img = cv2.cvtColor(ploted_img, cv2.COLOR_RGB2BGR) if configs.show_image: cv2.imshow('ploted_img', ploted_img) cv2.waitKey(10) if configs.save_demo_output: cv2.imwrite(os.path.join(configs.frame_dir, '{:06d}.jpg'.format(frame_idx)), ploted_img) frame_pred_infor = { 'seg': prediction_seg, 'ball': prediction_ball_final } queue_frames.append(frame_pred_infor) frame_idx += 1 print('Done frame_idx {} - time {:.3f}s'.format(frame_idx, t2 - t1)) if configs.output_format == 'video': output_video_path = os.path.join(configs.save_demo_dir, 'result.mp4') cmd_str = 'ffmpeg -f image2 -i {}/%05d.jpg -b 5000k -c:v mpeg4 {}'.format( os.path.join(configs.frame_dir), output_video_path) os.system(cmd_str)