def main(args): model = Pointnet2StructurePointNet( num_structure_points=args.num_structure_points, input_channels=0, use_xyz=True) model.cuda() checkpoint_util.load_checkpoint(model_3d=model, filename=args.model_fname) if os.path.exists(args.output_dir) is False: os.makedirs(args.output_dir) fnames = glob.glob(os.path.join(args.data_dir, '*.off')) for fname in fnames: fname = os.path.basename(fname) pts = point_cloud_utils.read_points_off( os.path.join(args.data_dir, fname)) batch_pts = torch.from_numpy(pts)[None, :, :].cuda() structure_points = model(batch_pts) structure_points = structure_points[0].cpu().detach().numpy() outfname = os.path.join(args.output_dir, fname[:-4] + '_stpts.off') point_cloud_utils.write_points_off( outfname, structure_points, COLOR_LIST[:structure_points.shape[0], :]) print('output saved to {0}'.format(args.output_dir))
def main(args): test_set = bhcp_dataloader.bhcp_dataloader( args.data_dir, args.category, is_pts_aligned=args.test_on_aligned) test_loader = DataLoader( test_set, batch_size=1, shuffle=False, num_workers=1, pin_memory=False, ) model = Pointnet2StructurePointNet( num_structure_points=args.num_structure_points, input_channels=0, use_xyz=True) model.cuda() checkpoint_util.load_checkpoint(model_3d=model, filename=args.model_fname) dis_ratios, dis_thresholds = bhcp_evaluate_one_epoch( model, test_loader, args.test_on_aligned) print('distance threshold|radio:') for i in range(dis_ratios.shape[0]): print(' {0}|{1}'.format(dis_thresholds[i], ("%.6f" % dis_ratios[i])), end='') print('\n')
def main(args): model = Pointnet2StructurePointNet( num_structure_points=args.num_structure_points, input_channels=0, use_xyz=True) model.cuda() checkpoint_util.load_checkpoint(model_3d=model, filename=args.model_fname) src_shape_fname = args.src_shape_fname query_pts_fname = args.query_pts_fname tgt_shape_fname = args.tgt_shape_fname out_corres_pts_fname = args.out_corres_pts_fname src_pts = point_cloud_utils.read_points_off(src_shape_fname) tgt_pts = point_cloud_utils.read_points_off(tgt_shape_fname) query_pts, query_pts_colors = point_cloud_utils.read_points_off( query_pts_fname, read_color=True) batch_src_pts = torch.from_numpy(src_pts)[None, :, :].cuda() batch_tgt_pts = torch.from_numpy(tgt_pts)[None, :, :].cuda() query_pts = torch.from_numpy(query_pts)[:, :] if args.test_on_aligned is False: batch_src_pts, src_rot_mats, _ = d_utils.AddPCATransformsToBatchPoints( batch_src_pts, num_of_trans=1) batch_src_pts = batch_src_pts.squeeze(dim=0) src_rot_mat = src_rot_mats.squeeze(dim=0)[0] batch_tgt_pts, tgt_rot_mats, _ = d_utils.AddPCATransformsToBatchPoints( batch_tgt_pts, num_of_trans=1) batch_tgt_pts = batch_tgt_pts.squeeze(dim=0) tgt_rot_mat = tgt_rot_mats.squeeze(dim=0)[0] structure_points_src = model(batch_src_pts)[0] structure_points_tgt = model(batch_tgt_pts)[0] structure_points_src = torch.mm(src_rot_mat.transpose(0, 1), structure_points_src.transpose( 0, 1)).transpose(0, 1) structure_points_tgt = torch.mm(tgt_rot_mat.transpose(0, 1), structure_points_tgt.transpose( 0, 1)).transpose(0, 1) corres_pts_in_tgt = compute_correspondence(structure_points_src, query_pts, structure_points_tgt) point_cloud_utils.write_points_off(out_corres_pts_fname, corres_pts_in_tgt, query_pts_colors) print('output saved to {0}'.format(out_corres_pts_fname))
def main(args): model = Pointnet2StructurePointNet( num_structure_points=args.num_structure_points, input_channels=0, use_xyz=True) model.cuda() checkpoint_util.load_checkpoint(model_3d=model, filename=args.model_fname) model.eval() if os.path.exists(args.output_dir) is False: os.makedirs(args.output_dir) fnames = glob.glob(os.path.join(args.data_dir, '*.off')) for fname in fnames: fname = os.path.basename(fname) pts = point_cloud_utils.read_points_off( os.path.join(args.data_dir, fname)) batch_pts = torch.from_numpy(pts)[None, :, :].cuda() if args.test_on_aligned is False: batch_pts, rot_mats, _ = d_utils.AddPCATransformsToBatchPoints( batch_pts, num_of_trans=1) batch_pts = batch_pts.squeeze(dim=0).contiguous() rot_mats = rot_mats.squeeze(dim=0) structure_points = model(batch_pts) if args.test_on_aligned is False: inv_rot_mats = rot_mats.transpose(1, 2) structure_points = torch.matmul(inv_rot_mats, structure_points.transpose( 1, 2)).transpose(1, 2) structure_points = structure_points[0].cpu().detach().numpy() outfname = os.path.join(args.output_dir, fname[:-4] + '_stpts.off') point_cloud_utils.write_points_off( outfname, structure_points, COLOR_LIST[:structure_points.shape[0], :]) print('output saved to {0}'.format(args.output_dir))
def train(cmd_args): if os.path.exists(cmd_args.log_dir) is not True: os.makedirs(cmd_args.log_dir) checkpoints_dir = os.path.join(cmd_args.log_dir, "checkpoints") if os.path.exists(checkpoints_dir) is not True: os.mkdir(checkpoints_dir) lr_clip = 1e-5 bnm_clip = 1e-2 trans = transforms.Compose([ d_utils.PointcloudToTensor(), #d_utils.PointcloudRandomPermutation(), d_utils.PointcloudJitter() ]) train_set = bhcp_dataloader.bhcp_dataloader(cmd_args.data_dir, cmd_args.category, transforms=trans, is_pts_aligned=False) train_loader = DataLoader( train_set, batch_size=cmd_args.batch_size, shuffle=True, num_workers=5, pin_memory=False, ) model = Pointnet2StructurePointNet( num_structure_points=cmd_args.num_structure_points, input_channels=0, use_xyz=True) model.cuda() optimizer = optim.Adam(model.parameters(), lr=cmd_args.lr, weight_decay=cmd_args.weight_decay) lr_lbmd = lambda it: max( cmd_args.lr_decay** (int(it * cmd_args.batch_size / cmd_args.decay_step)), lr_clip / cmd_args.lr, ) bn_lbmd = lambda it: max( cmd_args.bn_momentum * cmd_args.bnm_decay** (int(it * cmd_args.batch_size / cmd_args.decay_step)), bnm_clip, ) iters = -1 criterions = {'ComputeLoss3d': ComputeLoss3d()} # load status from checkpoint log_open_mode = 'w' start_epoch = 0 if cmd_args.checkpoint is not None: fname = os.path.join(checkpoints_dir, cmd_args.checkpoint) start_epoch, iters = checkpoint_util.load_checkpoint( model_3d=model, optimizer=optimizer, filename=fname) start_epoch += 1 log_open_mode = 'a' log = LogUtils(os.path.join(cmd_args.log_dir, 'logfile'), log_open_mode) log.write('train unsupervised structure points for bhcp\n') log.write_args(cmd_args) lr_scheduler = lr_sched.LambdaLR(optimizer, lr_lambda=lr_lbmd, last_epoch=iters) bnm_scheduler = pt_utils.BNMomentumScheduler(model, bn_lambda=bn_lbmd, last_epoch=iters) iters = max(iters, 0) for epoch_i in range(start_epoch, cmd_args.max_epochs): log.write('\nepoch: {0} ###########################'.format(epoch_i)) train_loss, train_cd_loss, train_consistent_loss, iters = train_one_epoch( model, optimizer, train_loader, lr_scheduler, bnm_scheduler, iters, criterions, num_of_trans=cmd_args.num_of_transform, num_inpts=cmd_args.num_inpts) log.write('\nave_train_loss:{0}, cd_loss:{1}, consis_loss:{2}'.format( ("%.8f" % train_loss), ("%.8f" % train_cd_loss), ("%.8f" % train_consistent_loss))) if cmd_args.checkpoint_save_step != -1 and ( epoch_i + 1) % cmd_args.checkpoint_save_step is 0: fname = os.path.join(checkpoints_dir, 'checkpoint_{}'.format(epoch_i)) checkpoint_util.save_checkpoint(filename=fname, model_3d=model, optimizer=optimizer, iters=iters, epoch=epoch_i) fname = os.path.join(checkpoints_dir, 'model') checkpoint_util.save_checkpoint(filename=fname, model_3d=model, optimizer=optimizer, iters=iters, epoch=epoch_i)
# Load the model input_channel = 1 model = ftp_psp3(input_channel).cuda() model = nn.DataParallel(model) # For using multiple GPUs # Define the optimizer optimizer = optim.Adam(model.parameters(), lr=opt.lr) #Load status from checkpoint log_open_mode = 'w' start_epoch = 0 if opt.checkpoint is not None: fname = os.path.join(checkpoints_dir, opt.checkpoint) start_epoch, iters = checkpoint_util.load_checkpoint(model_3d=model, optimizer=optimizer, filename=fname) start_epoch += 1 log_open_mode = 'a' log = LogUtils(os.path.join(opt.log_dir, 'logfile'), log_open_mode) log.write('Supervised learning for phase map enhancement - Training\n') log.write_args(opt) lr_scheduler = lr_scd.StepLR(optimizer, step_size=opt.decay_step, gamma=opt.lr_decay) iters = max(iters, 0) reg = 1e-7 # Train the network on the training dataset for epoch_num in range(start_epoch, opt.num_epochs): trainData = iter(trainLoader)
if os.path.exists(opt.log_dir) is not True: os.makedirs(opt.log_dir) checkpoints_dir = os.path.join(opt.log_dir, "checkpoints") if os.path.exists(checkpoints_dir) is not True: os.mkdir(checkpoints_dir) # Load the model input_channel = 1 model = ftp_psp5(input_channel).cuda() # model = nn.DataParallel(model) # For using multiple GPUs #Load status from checkpoint log_open_mode = 'w' checkpoint_util.load_checkpoint(model_3d=model, filename=opt.checkpoint) log = LogUtils(os.path.join(opt.log_dir, 'logfile'), log_open_mode) log.write('Supervised learning for phase map enhancement - Testing\n') log.write_args(opt) if opt.checkpoint is None: print('Checkpoint is missing! Load the checkpoint to start the testing') sys.exit() # Test the network using the trained model testData = iter(testLoader) ave_loss = 0 count = 0 for data in iter(testLoader): inp_PM, gt_PM, mask_PM, filename_PM = next(testData)