def __init__(self, node_name=''): self.rosparam_(node_name) self.create_logger_(node_name) self.logger.info('ROS parameters loaded successfully') self.logger.info( '********************** Start logging **********************') save_config_to_file(cfg, logger=self.logger) self.logger.info('ROS buffer_len: {}'.format(self.buffer_len)) self.logger.info('ROS lidar_topic: {}'.format(self.lidar_topic)) self.logger.info('ROS debug_flag: {}'.format(self.debug_flag)) self.logger.info( 'ROS depth_threshold: {}'.format(self.depth_threshold)) self.logger.info( 'ROS score_threshold: {}'.format(self.score_threshold)) self.logger.info('ROS model_checkpoint: {}'.format(self.model_checkpoint)) self.model = PointRCNN(num_classes=self.num_class, use_xyz=True, mode='TEST') self.model.cuda() load_checkpoint( self.model, filename=str(self.base_dir / self.model_checkpoint), logger=self.logger) self.model.eval() self.logger.info('Model initialization complete')
def repeat_eval_ckpt(root_result_dir, ckpt_dir): root_result_dir = os.path.join(root_result_dir, 'eval', 'eval_all_' + args.extra_tag) os.makedirs(root_result_dir, exist_ok=True) log_file = os.path.join(root_result_dir, 'log_eval_all_%s.txt' % cfg.TEST.SPLIT) logger = create_logger(log_file) logger.info('**********************Start logging**********************') # save config for key, val in vars(args).items(): logger.info("{:16} {}".format(key, val)) save_config_to_file(cfg, logger=logger) # create dataloader & network test_loader = create_dataloader(logger) model = PointRCNN(num_classes=test_loader.dataset.num_class, use_xyz=True, mode='TEST') model.cuda() # copy important files to backup backup_dir = os.path.join(root_result_dir, 'backup_files') os.makedirs(backup_dir, exist_ok=True) os.system('cp *.py %s/' % backup_dir) os.system('cp ../lib/net/*.py %s/' % backup_dir) os.system('cp ../lib/datasets/kitti_rcnn_dataset.py %s/' % backup_dir) # evaluated ckpt record ckpt_record_file = os.path.join(root_result_dir, 'eval_list_%s.txt' % cfg.TEST.SPLIT) with open(ckpt_record_file, 'a'): pass # tensorboard log tb_log = SummaryWriter(log_dir=os.path.join(root_result_dir, 'tensorboard_%s' % cfg.TEST.SPLIT)) while True: # check whether there is checkpoint which is not evaluated cur_epoch_id, cur_ckpt = get_no_evaluated_ckpt(ckpt_dir, ckpt_record_file) if cur_epoch_id == -1 or int(float(cur_epoch_id)) < args.start_epoch: wait_second = 30 print('Wait %s second for next check: %s' % (wait_second, ckpt_dir)) time.sleep(wait_second) continue # load checkpoint train_utils.load_checkpoint(model, filename=cur_ckpt) # start evaluation cur_result_dir = os.path.join(root_result_dir, 'epoch_%s' % cur_epoch_id, cfg.TEST.SPLIT) tb_dict = eval_one_epoch(model, test_loader, cur_epoch_id, cur_result_dir, logger) step = int(float(cur_epoch_id)) if step == float(cur_epoch_id): for key, val in tb_dict.items(): tb_log.add_scalar(key, val, step) # record this epoch which has been evaluated with open(ckpt_record_file, 'a') as f: print('%s' % cur_epoch_id, file=f) logger.info('Epoch %s has been evaluated' % cur_epoch_id)
def eval_single_ckpt(root_result_dir): root_result_dir = os.path.join(root_result_dir, 'eval') # set epoch_id and output dir num_list = re.findall(r'\d+', args.ckpt) if args.ckpt is not None else [] epoch_id = num_list[-1] if num_list.__len__() > 0 else 'no_number' root_result_dir = os.path.join(root_result_dir, 'epoch_%s' % epoch_id, cfg.TEST.SPLIT) if args.test: root_result_dir = os.path.join(root_result_dir, 'test_mode') if args.extra_tag != 'default': root_result_dir = os.path.join(root_result_dir, args.extra_tag) root_result_dir = os.path.join( root_result_dir, args.depth_ckpt.split('/')[-1].split('.')[0]) os.makedirs(root_result_dir, exist_ok=True) log_file = os.path.join(root_result_dir, 'log_eval_one.txt') logger = create_logger(log_file) logger.info('**********************Start logging**********************') for key, val in vars(args).items(): logger.info("{:16} {}".format(key, val)) save_config_to_file(cfg, logger=logger) # create dataloader & network test_loader = create_dataloader(logger) model = PointRCNN(num_classes=test_loader.dataset.num_class, use_xyz=True, mode='TEST') model.cuda() # Create depth model depth_model = DepthModel(maxdepth=80, maxdisp=192, down=2, pretrain=args.depth_ckpt, save_tag=cfg.TAG, mode='TEST', mgpus=args.mgpus) # copy important files to backup backup_dir = os.path.join(root_result_dir, 'backup_files') os.makedirs(backup_dir, exist_ok=True) os.system('cp *.py %s/' % backup_dir) os.system('cp ../lib/net/*.py %s/' % backup_dir) os.system('cp ../lib/datasets/kitti_rcnn_dataset.py %s/' % backup_dir) # load checkpoint load_ckpt_based_on_args(model, logger) # start evaluation eval_one_epoch(model, depth_model, test_loader, epoch_id, root_result_dir, logger)
def eval_all_ckpt(root_result_dir): root_result_dir = os.path.join('/'.join(args.rcnn_ckpt.split('/')[:-1]), 'all', 'eval') os.makedirs(root_result_dir, exist_ok=True) # set epoch_id and output dir ckpt_dir = '/'.join(args.rcnn_ckpt.split('/')[:-1]) ckpt_list = os.listdir(ckpt_dir) ckpt_list = [x for x in ckpt_list if x[-4:] == '.pth'] ckpt_list.sort() BEST_precision = 0. BEST_iter = None log_file = os.path.join(root_result_dir, 'log_eval_all.txt') logger = create_logger(log_file) logger.info('**********************Start logging**********************') for key, val in vars(args).items(): logger.info("{:16} {}".format(key, val)) save_config_to_file(cfg, logger=logger) # create dataloader & network test_loader = create_dataloader(logger) model = PointRCNN(num_classes=test_loader.dataset.num_class, use_xyz=True, mode='TEST') model.cuda() for ckpt in tqdm.tqdm(reversed(ckpt_list[25:])): args.rcnn_ckpt = os.path.join(ckpt_dir, ckpt) num_list = re.findall( r'\d+', args.rcnn_ckpt) if args.rcnn_ckpt is not None else [] iter_id = num_list[-1] if num_list.__len__() > 0 else 'no_number' cur_root_result_dir = os.path.join(root_result_dir, cfg.TEST.SPLIT) if args.test: cur_root_result_dir = os.path.join(root_result_dir, 'test_mode') if args.extra_tag != 'default': cur_root_result_dir = os.path.join(cur_root_result_dir, args.extra_tag) os.makedirs(cur_root_result_dir, exist_ok=True) # load checkpoint load_ckpt_based_on_args(model, logger) precision, _, _ = eval_one_epoch_joint(model, test_loader, iter_id, cur_root_result_dir, logger) if precision > BEST_precision: BEST_precision = precision BEST_iter = iter_id print('best_precision: %.4f, best_iter: %s,' % (BEST_precision, BEST_iter)) print(args.rcnn_ckpt[-4:])
def eval_single_ckpt(root_result_dir): root_result_dir = os.path.join(root_result_dir, 'eval') # set epoch_id and output dir num_list = re.findall(r'\d+', args.ckpt) if args.ckpt is not None else [] epoch_id = num_list[-1] if num_list.__len__() > 0 else 'no_number' root_result_dir = os.path.join(root_result_dir, 'epoch_%s' % epoch_id, cfg.TEST.SPLIT) # Checks if TEST mode in on or not if args.test: root_result_dir = os.path.join(root_result_dir, 'test_mode') if args.extra_tag != 'default': root_result_dir = os.path.join(root_result_dir, args.extra_tag) # Create root_result_dir if it doesn't exists os.makedirs(root_result_dir, exist_ok=True) # Log File initialize log_file = os.path.join(root_result_dir, 'log_eval_one.txt') logger = create_logger(log_file) logger.info('**********************Start logging**********************') for key, val in vars(args).items(): logger.info("{:16} {}".format(key, val)) save_config_to_file(cfg, logger=logger) # create dataloader & network test_loader = create_dataloader(logger) model = PointRCNN(num_classes=test_loader.dataset.num_class, use_xyz=True, mode='TEST') model.cuda() # copy important files to backup (Can comment out as takes extra space by backing up files) ''' backup_dir = os.path.join(root_result_dir, 'backup_files') os.makedirs(backup_dir, exist_ok=True) os.system('cp *.py %s/' % backup_dir) os.system('cp ../lib/net/*.py %s/' % backup_dir) os.system('cp ../lib/datasets/kitti_rcnn_dataset.py %s/' % backup_dir) ''' # load checkpoint load_ckpt_based_on_args(model, logger) # start evaluation eval_one_epoch(model, test_loader, epoch_id, root_result_dir, logger)
def eval_single_ckpt(root_result_dir): root_result_dir = os.path.join(root_result_dir, 'eval') # set epoch_id and output dir num_list = re.findall(r'\d+', args.ckpt) if args.ckpt is not None else [] epoch_id = num_list[-1] if num_list.__len__() > 0 else 'no_number' root_result_dir = os.path.join(root_result_dir, 'epoch_%s' % epoch_id, cfg.TEST.SPLIT) if args.test: root_result_dir = os.path.join(root_result_dir, 'test_mode') if args.extra_tag != 'default': root_result_dir = os.path.join(root_result_dir, args.extra_tag) os.makedirs(root_result_dir, exist_ok=True) log_file = os.path.join(root_result_dir, 'log_eval_one.txt') logger = create_logger(log_file) logger.info('**********************Start logging**********************') for key, val in vars(args).items(): logger.info("{:16} {}".format(key, val)) save_config_to_file(cfg, logger=logger) # create dataloader & network while True: try: test_loader = create_dataloader(logger) while test_loader.dataset.npoints < 100 or test_loader.dataset.sample_id_list == 0: test_loader = create_dataloader(logger) model = PointRCNN(num_classes=test_loader.dataset.num_class, use_xyz=True, mode='TEST') model.cuda() # copy important files to backup backup_dir = os.path.join(root_result_dir, 'backup_files') os.makedirs(backup_dir, exist_ok=True) os.system('cp *.py %s/' % backup_dir) os.system('cp ../lib/net/*.py %s/' % backup_dir) os.system('cp ../lib/datasets/kitti_rcnn_dataset.py %s/' % backup_dir) # load checkpoint load_ckpt_based_on_args(model, logger) # start evaluation eval_one_epoch(model, test_loader, epoch_id, root_result_dir, logger) except Exception as e: print("No Poincloud found. Waiting for the next frame.") print("Chill time: Sleeping for 2 sec while waiting") print(e) time.sleep(2)
def eval_single_ckpt(root_result_dir): root_result_dir = os.path.join(root_result_dir, 'eval') # set epoch_id and output dir num_list = re.findall(r'\d+', args.ckpt) if args.ckpt is not None else [] epoch_id = num_list[-1] if num_list.__len__() > 0 else 'no_number' root_result_dir = os.path.join(root_result_dir, 'epoch_%s' % epoch_id, cfg.TEST.SPLIT) if args.test: root_result_dir = os.path.join(root_result_dir, 'test_mode') if args.extra_tag != 'default': root_result_dir = os.path.join(root_result_dir, args.extra_tag) os.makedirs(root_result_dir, exist_ok=True) log_file = os.path.join(root_result_dir, 'log_eval_one.txt') logger = create_logger(log_file) logger.info('**********************Start logging**********************') for key, val in vars(args).items(): logger.info("{:16} {}".format(key, val)) save_config_to_file(cfg, logger=logger) # create dataloader & network test_loader = create_dataloader(logger) model = PointRCNN(num_classes=test_loader.dataset.num_class, use_xyz=True, mode='TEST') model.cuda() # copy important files to backup backup_dir = os.path.join(root_result_dir, 'backup_files') os.makedirs(backup_dir, exist_ok=True) os.system('cp *.py %s/' % backup_dir) os.system('cp ../lib/net/*.py %s/' % backup_dir) os.system('cp ../lib/datasets/kitti_rcnn_dataset.py %s/' % backup_dir) # load checkpoint load_ckpt_based_on_args(model, logger) # start evaluation ret_dict = eval_one_epoch(model, test_loader, epoch_id, root_result_dir, logger) json_file = open(root_result_dir + 'evaldict.json', 'w') json_file.write(json.dumps(ret_dict)) json_file.close()
def __init__(self, use_masked=True): super().__init__() np.random.seed(1024) # label path self.label_root = os.path.join(HOME_DIR, 'data/KITTI/object/training/label_2/') # load config config_path = os.path.join(HOME_DIR, 'tools/configs/pg.json') self.config = load_config(config_path) self.npoints = cfg.RPN.NUM_POINTS root_result_dir = os.path.join('../', 'output', 'rcnn', cfg.TAG) ckpt_dir = os.path.join('../', 'output', 'rcnn', cfg.TAG, 'ckpt') # create logger logger = create_logger(os.path.join(OUTPUT_DIR, 'log_pg.txt')) logger.info( '**********************Start logging**********************') for key, val in self.config.items(): logger.info("{:16} {}".format(key, val)) save_config_to_file(cfg, logger=logger) # create PointRCNN dataloader & network self.test_loader = create_dataloader(self.config, logger) self.test_iter = iter(self.test_loader) self.model = PointRCNN(num_classes=self.test_loader.dataset.num_class, use_xyz=True, mode='TEST') self.use_masked = use_masked # load checkpoint load_ckpt_based_on_cfg(self.config, self.model, logger) # If want parallel # self.model = torch.nn.DataParallel(self.model) self.model.cuda() self.model.eval() self.data = None
def eval_single_ckpt(root_result_dir): root_result_dir = os.path.join(root_result_dir, 'eval') # set epoch_id and output dir epoch_id = 'no_number' root_result_dir = os.path.join(root_result_dir, 'epoch_%s' % epoch_id, cfg.TEST.SPLIT) # Checks if TEST mode in on or not if args.test: root_result_dir = os.path.join(root_result_dir, 'test_mode') # Create root_result_dir if it doesn't exists os.makedirs(root_result_dir, exist_ok=True) # Log File initialize log_file = os.path.join(root_result_dir, 'log_eval_one.txt') logger = create_logger(log_file) logger.info('**********************Start logging**********************') for key, val in vars(args).items(): logger.info("{:16} {}".format(key, val)) save_config_to_file(cfg, logger=logger) # Check if its a Single file or not model = PointRCNN(num_classes=2, use_xyz=True, mode='TEST') model.cuda() # load checkpoint load_ckpt_based_on_args(model, logger) if (args.single_file != None): eval_one_epoch_joint_single_file(model, get_lidar(args.single_file), args.single_file, root_result_dir, logger) else: # create dataloader & network test_loader = create_dataloader(logger) # start evaluation eval_one_epoch(model, test_loader, epoch_id, root_result_dir, logger)
def __init__(self): # dnn cfg_from_file(cfg_file) cfg.RCNN.ENABLED = True cfg.RPN.ENABLED = cfg.RPN.FIXED = True cfg.RPN.LOC_XZ_FINE = False self.pc_roi = [[-25, 25], [-3, 2], [-25, 25]] self.down_sample = {'axis': int(0), 'depth': self.pc_roi[0][1] / 2} # [axis,depth] self.mode = 'TEST' with torch.no_grad(): self.model = PointRCNN(num_classes=2, use_xyz=True, mode=self.mode) self.model.cuda() self.model.eval() load_checkpoint(model=self.model, optimizer=None, filename=pointrcnn_weight) # ros self.pc_sub = rospy.Subscriber(pc_topic, PointCloud2, self.pc_cb, queue_size=1, buff_size=2 ** 24) self.pc_pub = rospy.Publisher(pack_name + "/networks_input", PointCloud2, queue_size=1) if is_viz else None self.mk_pub = rospy.Publisher(pack_name + "/networks_output", MarkerArray, queue_size=1) if is_viz else None self.Tr_velo_kitti_cam = np.array([0.0, - 1.0, 0.0, 0.0, 0.0, 0.0, -1.0, 1.5, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0]).reshape(4, 4) if is_tf else np.identity(4)
class ROSPointRCNN(object): def __init__(self): # dnn cfg_from_file(cfg_file) cfg.RCNN.ENABLED = True cfg.RPN.ENABLED = cfg.RPN.FIXED = True cfg.RPN.LOC_XZ_FINE = False self.pc_roi = [[-25, 25], [-3, 2], [-25, 25]] self.down_sample = {'axis': int(0), 'depth': self.pc_roi[0][1] / 2} # [axis,depth] self.mode = 'TEST' with torch.no_grad(): self.model = PointRCNN(num_classes=2, use_xyz=True, mode=self.mode) self.model.cuda() self.model.eval() load_checkpoint(model=self.model, optimizer=None, filename=pointrcnn_weight) # ros self.pc_sub = rospy.Subscriber(pc_topic, PointCloud2, self.pc_cb, queue_size=1, buff_size=2 ** 24) self.pc_pub = rospy.Publisher(pack_name + "/networks_input", PointCloud2, queue_size=1) if is_viz else None self.mk_pub = rospy.Publisher(pack_name + "/networks_output", MarkerArray, queue_size=1) if is_viz else None self.Tr_velo_kitti_cam = np.array([0.0, - 1.0, 0.0, 0.0, 0.0, 0.0, -1.0, 1.5, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0]).reshape(4, 4) if is_tf else np.identity(4) def pc_cb(self, data): pts_input = self.extract_networks_input_from_pc2rosmsg(data) if self.pc_pub is not None: self.pc_pub.publish(numpy2pc2(pts_input, data.header.frame_id)) np.random.seed(666) with torch.no_grad(): # 准备输入数据 MEAN_SIZE = torch.from_numpy(cfg.CLS_MEAN_SIZE[0]).cuda() inputs = torch.from_numpy(pts_input).cuda(non_blocking=True).float() inputs = torch.unsqueeze(inputs, 0) # 模型推理 input_data = {'pts_input': inputs} ret_dict = self.model(input_data) # 分析结果 batch_size = 1 roi_scores_raw = ret_dict['roi_scores_raw'] # (B, M) 提案置信度预测 roi_boxes3d = ret_dict['rois'] # (B, M, 7) 提案框 seg_result = ret_dict['seg_result'].long() # (B, N) 前景点分割 rcnn_cls = ret_dict['rcnn_cls'].view(batch_size, -1, ret_dict['rcnn_cls'].shape[1]) # (B, M, n) bin分类结果 rcnn_reg = ret_dict['rcnn_reg'].view(batch_size, -1, ret_dict['rcnn_reg'].shape[1]) # (B, M, C) res回归结果 # 解算3D BBOx anchor_size = MEAN_SIZE pred_boxes3d = decode_bbox_target(roi_boxes3d.view(-1, 7), rcnn_reg.view(-1, rcnn_reg.shape[-1]), anchor_size=anchor_size, loc_scope=cfg.RCNN.LOC_SCOPE, loc_bin_size=cfg.RCNN.LOC_BIN_SIZE, num_head_bin=cfg.RCNN.NUM_HEAD_BIN, get_xz_fine=True, get_y_by_bin=cfg.RCNN.LOC_Y_BY_BIN, loc_y_scope=cfg.RCNN.LOC_Y_SCOPE, loc_y_bin_size=cfg.RCNN.LOC_Y_BIN_SIZE, get_ry_fine=True).view(batch_size, -1, 7) # cfg.SCORE_THRESH 置信度阈值 if rcnn_cls.shape[2] == 1: batch_raw_scores = rcnn_cls # (B, M, 1) batch_norm_scores = torch.sigmoid(batch_raw_scores) # (B,M,1) batch_pred_classes = (batch_norm_scores > cfg.RCNN.SCORE_THRESH).long() # (B,M,1) else: batch_pred_classes = torch.argmax(rcnn_cls, dim=1).view(-1) batch_raw_scores = rcnn_cls[:, batch_pred_classes] batch_norm_scores = F.softmax(rcnn_cls, dim=1)[:, batch_pred_classes] # scores threshold inds = batch_norm_scores > cfg.RCNN.SCORE_THRESH for batch in range(batch_size): inds_in_each_batch = inds[batch].view(-1) if inds_in_each_batch.sum() == 0: # batch 内没有超过阈值的3dbbox continue pred_boxes3d_in_each_batch = pred_boxes3d[batch, inds_in_each_batch] raw_scores_in_each_batch = batch_raw_scores[batch, inds_in_each_batch] norm_scores_in_each_batch = batch_norm_scores[batch, inds_in_each_batch] # 非极大值抑制 boxes_bev_in_each_batch = kitti_utils.boxes3d_to_bev_torch(pred_boxes3d_in_each_batch) keep_idx = iou3d_utils.nms_gpu(boxes_bev_in_each_batch, raw_scores_in_each_batch, cfg.RCNN.NMS_THRESH).view(-1) pred_boxes3d_in_each_batch = pred_boxes3d_in_each_batch[keep_idx] raw_scores_in_each_batch = raw_scores_in_each_batch[keep_idx] output = {'boxes3d': pred_boxes3d_in_each_batch.cpu().numpy(), 'scores': raw_scores_in_each_batch.cpu().numpy()} self.visualize(output, data.header.frame_id) def visualize(self, result, frame_id): boxes = result['boxes3d'] scores = result['scores'] self.visualize_lidar_plane(boxes, frame_id) print("Number of detections pr msg: ", boxes.shape[0]) def visualize_lidar_plane(self, bbox3d, frame_id): marker_array = MarkerArray() marker = Marker() marker.header.frame_id = frame_id marker.type = marker.LINE_LIST marker.action = marker.ADD marker.header.stamp = rospy.Time.now() # marker scale (scale y and z not used due to being linelist) marker.scale.x = 0.08 # marker color marker.color.a = 1.0 marker.color.r = 1.0 marker.color.g = 1.0 marker.color.b = 0.0 marker.pose.position.x = 0.0 marker.pose.position.y = 0.0 marker.pose.position.z = 0.0 marker.pose.orientation.x = 0.0 marker.pose.orientation.y = 0.0 marker.pose.orientation.z = 0.0 marker.pose.orientation.w = 1.0 marker.points = [] corner_for_box_list = [0, 1, 0, 3, 2, 3, 2, 1, 4, 5, 4, 7, 6, 7, 6, 5, 3, 7, 0, 4, 1, 5, 2, 6] corners3d = kitti_utils.boxes3d_to_corners3d(bbox3d) # (N,8,3) for box_nr in range(corners3d.shape[0]): box3d_pts_3d_velo = corners3d[box_nr] # (8,3) for corner in corner_for_box_list: p = np.array(box3d_pts_3d_velo[corner, 0:4]) transformed_p = transform_point(p, np.linalg.inv(self.Tr_velo_kitti_cam)) p = Point() p.x = transformed_p[0] p.y = transformed_p[1] p.z = transformed_p[2] marker.points.append(p) marker_array.markers.append(marker) id = 0 for m in marker_array.markers: m.id = id id += 1 self.mk_pub.publish(marker_array) marker_array.markers = [] pass def extract_networks_input_from_pc2rosmsg(self, data): random_select = True pts_lidar = ros_numpy.point_cloud2.pointcloud2_to_xyz_array(data) pts_lidar = transform_pointcloud(pts_lidar, self.Tr_velo_kitti_cam) pts_lidar = cropped_roi(pts_lidar, self.pc_roi) pts_lidar = pts_lidar[:, 0:3] # only xyz if cfg.RPN.NUM_POINTS < len(pts_lidar): pts_depth = pts_lidar[:, self.down_sample['axis']] pts_near_flag = pts_depth < self.down_sample['depth'] far_idxs_choice = np.where(pts_near_flag == 0)[0] near_idxs = np.where(pts_near_flag == 1)[0] near_idxs_choice = np.random.choice(near_idxs, cfg.RPN.NUM_POINTS - len(far_idxs_choice), replace=False) choice = np.concatenate((near_idxs_choice, far_idxs_choice), axis=0) \ if len(far_idxs_choice) > 0 else near_idxs_choice np.random.shuffle(choice) pts_input = pts_lidar[choice, :] return pts_input
# copy important files to backup backup_dir = os.path.join(root_result_dir, 'backup_files') os.makedirs(backup_dir, exist_ok=True) os.system('cp *.py %s/' % backup_dir) os.system('cp ../lib/net/*.py %s/' % backup_dir) os.system('cp ../lib/datasets/kitti_rcnn_dataset.py %s/' % backup_dir) # tensorboard log tb_log = SummaryWriter( log_dir=os.path.join(root_result_dir, 'tensorboard')) # create dataloader & network & optimizer train_loader, test_loader = create_dataloader(logger) model = PointRCNN(num_classes=train_loader.dataset.num_class, use_xyz=True, mode='TRAIN') optimizer = create_optimizer(model) if args.mgpus: model = nn.DataParallel(model) model.cuda() # load checkpoint if it is possible start_epoch = it = 0 last_epoch = -1 if args.ckpt is not None: pure_model = model.module if isinstance( model, torch.nn.DataParallel) else model it, start_epoch = train_utils.load_checkpoint(pure_model, optimizer,
os.makedirs(backup_dir, exist_ok=True) os.system('cp *.py %s/' % backup_dir) os.system('cp ../lib/net/*.py %s/' % backup_dir) os.system('cp ../lib/datasets/kitti_rcnn_dataset.py %s/' % backup_dir) os.system('cp ../lib/datasets/kitti_boxplace_dataset.py %s/' % backup_dir) os.system('cp ./train_utils/train_utils.py %s/' % backup_dir) os.system('cp ../lib/utils/loss_utils.py %s/' % backup_dir) # tensorboard log tb_log = SummaryWriter( log_dir=os.path.join(root_result_dir, 'tensorboard')) # create dataloader & network & optimizer train_loader, test_loader = create_dataloader(logger) model = PointRCNN(num_classes=train_loader.dataset.num_class, num_point=cfg.RCNN.NUM_POINTS, use_xyz=True, mode='TRAIN') optimizer = create_optimizer(model) if args.mgpus: model = nn.DataParallel(model) model.cuda() # load checkpoint if it is possible start_iter = it = 0 last_iter = -1 if args.pretrain_ckpt is not None: pure_model = model.module if isinstance( model, torch.nn.DataParallel) else model train_utils.load_part_ckpt(pure_model,
class PointRCNNEnv(): def __init__(self, use_masked=True): super().__init__() np.random.seed(1024) # label path self.label_root = os.path.join(HOME_DIR, 'data/KITTI/object/training/label_2/') # load config config_path = os.path.join(HOME_DIR, 'tools/configs/pg.json') self.config = load_config(config_path) self.npoints = cfg.RPN.NUM_POINTS root_result_dir = os.path.join('../', 'output', 'rcnn', cfg.TAG) ckpt_dir = os.path.join('../', 'output', 'rcnn', cfg.TAG, 'ckpt') # create logger logger = create_logger(os.path.join(OUTPUT_DIR, 'log_pg.txt')) logger.info( '**********************Start logging**********************') for key, val in self.config.items(): logger.info("{:16} {}".format(key, val)) save_config_to_file(cfg, logger=logger) # create PointRCNN dataloader & network self.test_loader = create_dataloader(self.config, logger) self.test_iter = iter(self.test_loader) self.model = PointRCNN(num_classes=self.test_loader.dataset.num_class, use_xyz=True, mode='TEST') self.use_masked = use_masked # load checkpoint load_ckpt_based_on_cfg(self.config, self.model, logger) # If want parallel # self.model = torch.nn.DataParallel(self.model) self.model.cuda() self.model.eval() self.data = None # def _batch_detector(self, batch_pts): # """ Input a single or batch sample of point clouds, output prediction result # """ # with torch.no_grad(): # self.model.eval() # thresh_list = [0.1, 0.3, 0.5, 0.7, 0.9] def reset(self): """ reset env; here it is equivlent to load an image and a bin from the KITTI dataset. Set the image returned as s0 data = {'sample_id': sample_id, 'random_select': self.random_select, 'pts_rect': pts_rect, 'pts_intensity': pts_intensity, 'gt_boxes3d': all_gt_boxes3d, 'npoints': self.npoints, 'image': image} """ # load the data sample at the reset step self.data = next(self.test_iter) RGB_Image = self.data['image'] return RGB_Image def step(self, action): """step [Input the sampled map, output ] """ # TODO: this is where we need to obs_pts = self._get_obs(action) # Here we output masked_pts as the obs rew = self._get_reward(obs_pts) # we set it as 1 step MDP so done is always true done = True info = {} return obs_pts, rew, done, info def _get_reward(self, obs): """step [Input the sampled point cloud, output the detection success] """ batch_mAP = self._eval_data(masked_pts=obs) return batch_mAP def _get_obs(self, scanning_mask): """Here we set next obs as the sampled point cloud """ masked_pts = self._get_pts_from_mask(scanning_mask) return masked_pts def _get_pts_from_mask(self, scanning_mask): """ mask pts from 2d angular map Input: :param scanning_mask: (B, H, W) Return: :param pts: (B, N, 4) """ # load ang_depth_map from dir ang_depth_map = self.data['angle_map'] # expand mask 2d->3d to enable broadcast mask = np.expand_dims(scanning_mask, axis=3) masked_ang_depth_map = [ ang_depth_map[k] * mask[k] for k in range(self.config['batch_size']) ] # masked_pts = masked_ang_depth_map.reshape((self.config['batch_size'], -1, 4)) masked_pts_arr = [ masked_pts[masked_pts[:, :, 0] > 0] for masked_pts in masked_ang_depth_map ] # around ~(15000,4) adjusted_masked_pts = [] for masked_pts in masked_pts_arr: if masked_pts.shape[0] <= self.npoints: padding = np.full((self.npoints - masked_pts.shape[0], 4), -1) adjusted_masked_pts.append( np.concatenate((masked_pts, padding), axis=0)) else: adjusted_masked_pts.append(masked_pts[:self.npoints, :]) masked_pts = np.array(adjusted_masked_pts) return masked_pts def render(self): """Placeholder for the rendering capacity """ raise NotImplementedError def _eval_data(self, masked_pts=None): """eval data with sampled pts """ with torch.no_grad(): MEAN_SIZE = torch.from_numpy(cfg.CLS_MEAN_SIZE[0]).cuda() batch_size = self.config['batch_size'] # get valid point (projected points should be in image) sample_id, pts_rect, pts_intensity, gt_boxes3d, npoints, labels = \ self.data['sample_id'], self.data['pts_rect'], self.data['pts_intensity'], self.data['gt_boxes3d'], self.data['npoints'], self.data['label'] cls_types = [[ labels[k][i].cls_type for i in range(len(labels[k])) ] for k in range(batch_size)] calib = [ self.test_loader.dataset.get_calib(idx) for idx in sample_id ] if self.use_masked: # use masked/sampled pts if True pts_rect = np.array([ c.lidar_to_rect(masked_pts[k][:, 0:3]) for k, c in enumerate(calib) ]) pts_intensity = [ masked_pts[k][:, 3] for k in range(batch_size) ] npoints = masked_pts.shape[0] inputs = torch.from_numpy(pts_rect).cuda( non_blocking=True).float().view(self.config['batch_size'], -1, 3) gt_boxes3d = torch.from_numpy(gt_boxes3d).cuda(non_blocking=True) input_data = {'pts_input': inputs} # model inference ret_dict = self.model(input_data) roi_scores_raw = ret_dict['roi_scores_raw'] # (B, M) roi_boxes3d = ret_dict['rois'] # (B, M, 7) # seg_result = ret_dict['seg_result'].long() # (B, N) rcnn_cls = ret_dict['rcnn_cls'].view(batch_size, -1, ret_dict['rcnn_cls'].shape[1]) rcnn_reg = ret_dict['rcnn_reg'].view( batch_size, -1, ret_dict['rcnn_reg'].shape[1]) # (B, M, C) norm_scores = torch.sigmoid(rcnn_cls) # remove low confidence scores thresh_mask = norm_scores > cfg.RCNN.SCORE_THRESH # bounding box regression anchor_size = MEAN_SIZE pred_boxes3d = decode_bbox_target( roi_boxes3d.view(-1, 7), rcnn_reg.view(-1, rcnn_reg.shape[-1]), anchor_size=anchor_size, loc_scope=cfg.RCNN.LOC_SCOPE, loc_bin_size=cfg.RCNN.LOC_BIN_SIZE, num_head_bin=cfg.RCNN.NUM_HEAD_BIN, get_xz_fine=True, get_y_by_bin=cfg.RCNN.LOC_Y_BY_BIN, loc_y_scope=cfg.RCNN.LOC_Y_SCOPE, loc_y_bin_size=cfg.RCNN.LOC_Y_BIN_SIZE, get_ry_fine=True).view(batch_size, -1, 7) # select boxes (list of tensors) pred_boxes3d_selected = [ pred_boxes3d[k][thresh_mask[k].view(-1)] for k in range(batch_size) ] raw_scores_selected = [ roi_scores_raw[k][thresh_mask[k].view(-1)] for k in range(batch_size) ] norm_scores_selected = [ norm_scores[k][thresh_mask[k].view(-1)] for k in range(batch_size) ] # rotated NMS boxes_bev_selected = [ kitti_utils.boxes3d_to_bev_torch(bboxes) for bboxes in pred_boxes3d_selected ] keep_idx = [ iou3d_utils.nms_gpu(boxes_bev_selected[k], raw_scores_selected[k], cfg.RCNN.NMS_THRESH).view(-1) for k in range(batch_size) ] pred_boxes3d_selected = [ pred_boxes3d_selected[k][keep_idx[k]] for k in range(batch_size) ] scores_selected = [ raw_scores_selected[k][keep_idx[k]] for k in range(batch_size) ] norm_scores_selected = [ norm_scores_selected[k][keep_idx[k]] for k in range(batch_size) ] # want car gt_boxes keep_idx = [[ i for i in range(len(cls_types[k])) if cls_types[k][i] == 'Car' ] for k in range(batch_size)] gt_boxes3d_selected = [ gt_boxes3d[k][keep_idx[k]] for k in range(batch_size) ] # what if no boxes with cars? has_info = [k for k in range(batch_size) if len(keep_idx[k]) > 0] gt_boxes3d_selected = [gt_boxes3d_selected[x] for x in has_info] pred_boxes3d_selected = [ pred_boxes3d_selected[x] for x in has_info ] batch_size = len(has_info) if batch_size == 0: return None # Intersect over union iou3d = [ iou3d_utils.boxes_iou3d_gpu(gt_boxes3d_selected[k], pred_boxes3d_selected[k]) for k in range(batch_size) ] # get the max iou for each ground truth bounding box gt_max_iou = [ torch.max(iou3d[k], dim=0)[0] for k in range(batch_size) ] # get precision at each index (to get auc) precision_vals = [] for k in range(batch_size): batch_iou = gt_max_iou[k] batch_precision = [] num_correct = 0 for i in range(len(batch_iou)): if batch_iou[i] > 0.7: num_correct += 1 batch_precision.append(num_correct / (i + 1)) precision_vals.append(batch_precision) aps = [] for k in range(batch_size): batch_prec = precision_vals[k] ap = 0 for i in range(len(batch_prec)): ap += max(batch_prec[i:]) aps.append(ap) num_gt_boxes = sum([len(gt_max_iou[k]) for k in range(batch_size)]) return sum(aps) / num_gt_boxes
os.mkdir(save_dir) ckpt_file = '/raid/meng/Pointcloud_Detection/PointRCNN4_weak/output/rpn/weaklyRPN0500/410_floss03_8000/ckpt/checkpoint_iter_07620.pth' cfg_from_file( '/raid/meng/Pointcloud_Detection/PointRCNN1.1_weak/tools/cfgs/weaklyRPN.yaml' ) cfg.RPN.SCORE_THRESH = 0.1 PROP_DIST = 0.3 BACKGROUND_ADDING = False BACK_THRESH = 0.3 COSINE_DISTANCE = False COS_THRESH = 0.3 from lib.net.point_rcnn import PointRCNN model = PointRCNN(num_classes=data_loader.dataset.num_class, use_xyz=True, mode='TEST') model.cuda() checkpoint = torch.load(ckpt_file) model_state = checkpoint['model_state'] update_model_state = { key: val for key, val in model_state.items() if key in model.state_dict() } state_dict = model.state_dict() state_dict.update(update_model_state) model.load_state_dict(state_dict) update_keys = update_model_state.keys().__len__() model.eval()
# copy important files to backup backup_dir = os.path.join(root_result_dir, 'backup_files') os.makedirs(backup_dir, exist_ok=True) os.system('cp *.py %s/' % backup_dir) os.system('cp ../lib/net/*.py %s/' % backup_dir) os.system('cp ../lib/datasets/kitti_rcnn_dataset.py %s/' % backup_dir) # tensorboard log tb_log = SummaryWriter( log_dir=os.path.join(root_result_dir, 'tensorboard')) # create dataloader & network & optimizer train_loader, test_loader = create_dataloader(logger) model = PointRCNN(num_classes=train_loader.dataset.num_class, use_xyz=True, mode='TRAIN').to(device) optimizer = create_optimizer(model) # if args.mgpus: # model = nn.DataParallel(model) # model.cuda() model, optimizer = amp.initialize(model, optimizer, opt_level="O2") # load checkpoint if it is possible start_epoch = it = 0 last_epoch = -1 if args.ckpt is not None: pure_model = model.module if isinstance( model, torch.nn.DataParallel) else model