def main(): args = options() torch.backends.cudnn.deterministic = True testset = RegistrationData(ModelNet40Data(train=False, num_points=args.num_points, unseen=args.unseen), partial_source=args.partial_source, noise=args.noise, outliers=args.outliers, additional_params={'use_masknet': True}) test_loader = DataLoader(testset, batch_size=args.test_batch_size, shuffle=False, drop_last=False, num_workers=args.workers) if not torch.cuda.is_available(): args.device = 'cpu' args.device = torch.device(args.device) # Load Pretrained MaskNet. model = MaskNet() if args.pretrained: assert os.path.isfile(args.pretrained) model.load_state_dict(torch.load(args.pretrained, map_location='cpu')) model = model.to(args.device) test(args, model, test_loader)
def main(): args = options() torch.backends.cudnn.deterministic = True if args.user_data: template = np.random.randn(1, 100, 3) # Define/Read template point cloud. [Shape: BATCH x No. of Points x 3] source = np.random.randn(1, 75, 3) # Define/Read source point cloud. [Shape: BATCH x No. of Points x 3] mask = np.zeros((1, 100, 1)) # Define/Read mask for point cloud. [Not mandatory in testing] igt = np.zeros((1, 4, 4)) # Define/Read igt transformation. [Not mandatory during testing] testset = UserData(template=template, source=source, mask=None, igt=None) elif args.any_data: # Read Stanford bunny's point cloud. bunny_path = os.path.join('learning3d/data/bunny/reconstruction/bun_zipper.ply') if not os.path.exists(bunny_path): print("Please download bunny dataset from http://graphics.stanford.edu/data/3Dscanrep/") print("Add the extracted folder in learning3d/data/") data = o3d.io.read_point_cloud(bunny_path) points = np.array(data.points) idx = np.arange(points.shape[0]) np.random.shuffle(idx) points = points[idx[:args.num_points]] testset = AnyData(pc=points, mask=True) else: testset = RegistrationData(ModelNet40Data(train=False, num_points=args.num_points, unseen=args.unseen), partial_source=args.partial_source, noise=args.noise, outliers=args.outliers) test_loader = DataLoader(testset, batch_size=args.test_batch_size, shuffle=False, drop_last=False, num_workers=args.workers) if not torch.cuda.is_available(): args.device = 'cpu' args.device = torch.device(args.device) # Load Pretrained MaskNet. model = MaskNet() if args.pretrained: assert os.path.isfile(args.pretrained) model.load_state_dict(torch.load(args.pretrained, map_location='cpu')) model = model.to(args.device) test(args, model, test_loader)
def main(): args = options() torch.backends.cudnn.deterministic = True torch.manual_seed(args.seed) torch.cuda.manual_seed_all(args.seed) np.random.seed(args.seed) boardio = SummaryWriter(log_dir='checkpoints/' + args.exp_name) _init_(args) textio = IOStream('checkpoints/' + args.exp_name + '/run.log') textio.cprint(str(args)) trainset = RegistrationData('PointNetLK', ModelNet40Data(train=True)) testset = RegistrationData('PointNetLK', ModelNet40Data(train=False)) train_loader = DataLoader(trainset, batch_size=args.batch_size, shuffle=True, drop_last=True, num_workers=args.workers) test_loader = DataLoader(testset, batch_size=8, shuffle=False, drop_last=False, num_workers=args.workers) if not torch.cuda.is_available(): args.device = 'cpu' args.device = torch.device(args.device) # Create PointNet Model. ptnet = PointNet(emb_dims=args.emb_dims, use_bn=True) if args.transfer_ptnet_weights and os.path.isfile( args.transfer_ptnet_weights): ptnet.load_state_dict( torch.load(args.transfer_ptnet_weights, map_location='cpu')) if args.fine_tune_pointnet == 'tune': pass elif args.fine_tune_pointnet == 'fixed': for param in ptnet.parameters(): param.requires_grad_(False) model = PointNetLK(feature_model=ptnet) model = model.to(args.device) checkpoint = None if args.resume: assert os.path.isfile(args.resume) checkpoint = torch.load(args.resume) args.start_epoch = checkpoint['epoch'] model.load_state_dict(checkpoint['model']) if args.pretrained: assert os.path.isfile(args.pretrained) model.load_state_dict(torch.load(args.pretrained, map_location='cpu')) model.to(args.device) if args.eval: test(args, model, test_loader, textio) else: train(args, model, train_loader, test_loader, boardio, textio, checkpoint)
from learning3d.data_utils import ModelNet40Data, ClassificationData, RegistrationData, FlowData modelnet40 = ModelNet40Data(train=True, num_points=1024, download=True) classification_data = ClassificationData(data_class=ModelNet40Data()) registration_data = RegistrationData(algorithm='PointNetLK', data_class=ModelNet40Data(), partial_source=False, partial_template=False, noise=False) flow_data = FlowData() print(flow_data)