args = parser.parse_args() update_config(cfg, args) if not os.path.exists(os.path.join(cfg.TEST.PATH, 'SyncedPoses.txt')): logger.info("First run on this captured data, start the pre-processing...") process_data(cfg.TEST.PATH) else: logger.info("Found SyncedPoses.txt, skipping data pre-processing...") logger.info("Running NeuralRecon...") transform = [transforms.ResizeImage((640, 480)), transforms.ToTensor(), transforms.RandomTransformSpace( cfg.MODEL.N_VOX, cfg.MODEL.VOXEL_SIZE, random_rotation=False, random_translation=False, paddingXY=0, paddingZ=0, max_epoch=cfg.TRAIN.EPOCHS), transforms.IntrinsicsPoseToProjection(cfg.TEST.N_VIEWS, 4)] transforms = transforms.Compose(transform) ARKitDataset = find_dataset_def(cfg.DATASET) test_dataset = ARKitDataset(cfg.TEST.PATH, "test", transforms, cfg.TEST.N_VIEWS, len(cfg.MODEL.THRESHOLDS) - 1) data_loader = DataLoader(test_dataset, cfg.BATCH_SIZE, shuffle=False, num_workers=cfg.TEST.N_WORKERS, drop_last=False) # model logger.info("Initializing the model on GPU...") model = NeuralRecon(cfg).cuda().eval() model = torch.nn.DataParallel(model, device_ids=[0]) # use the latest checkpoint file saved_models = [fn for fn in os.listdir(cfg.LOGDIR) if fn.endswith(".ckpt")] saved_models = sorted(saved_models, key=lambda x: int(x.split('_')[-1].split('.')[0])) loadckpt = os.path.join(cfg.LOGDIR, saved_models[-1]) logger.info("Resuming from " + str(loadckpt))
random_translation = False paddingXY = 0 paddingZ = 0 transform = [] transform += [ transforms.ResizeImage((640, 480)), transforms.ToTensor(), transforms.RandomTransformSpace(cfg.MODEL.N_VOX, cfg.MODEL.VOXEL_SIZE, random_rotation, random_translation, paddingXY, paddingZ, max_epoch=cfg.TRAIN.EPOCHS), transforms.IntrinsicsPoseToProjection(n_views, 4), ] transforms = transforms.Compose(transform) # dataset, dataloader MVSDataset = find_dataset_def(cfg.DATASET) train_dataset = MVSDataset(cfg.TRAIN.PATH, "train", transforms, cfg.TRAIN.N_VIEWS, len(cfg.MODEL.THRESHOLDS) - 1) test_dataset = MVSDataset(cfg.TEST.PATH, "test", transforms, cfg.TEST.N_VIEWS, len(cfg.MODEL.THRESHOLDS) - 1) if cfg.DISTRIBUTED: train_sampler = DistributedSampler(train_dataset, shuffle=False) TrainImgLoader = torch.utils.data.DataLoader(