Exemplo n.º 1
0
def visu_caller(path, step_ind, relu_idx):

    ##########################
    # Initiate the environment
    ##########################

    # Choose which gpu to use
    GPU_ID = '0'

    # Set GPU visible device
    os.environ['CUDA_VISIBLE_DEVICES'] = GPU_ID

    # Disable warnings
    os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'

    ###########################
    # Load the model parameters
    ###########################

    # Load model parameters
    config = Config()
    config.load(path)

    ##################################
    # Change model parameters for test
    ##################################

    # Change parameters for the test here. For example, you can stop augmenting the input data.

    # No augmentation to avoid random inputs
    config.augment_scale_anisotropic = False
    config.augment_symmetries = [False, False, False]
    config.augment_rotation = 'none'
    config.augment_scale_min = 1.0
    config.augment_scale_max = 1.0
    config.augment_noise = 0.0
    config.augment_occlusion = 'none'
    config.augment_color = 1.0

    config.batch_num = 2
    config.in_radius = 5

    ##############
    # Prepare Data
    ##############

    print()
    print('Dataset Preparation')
    print('*******************')

    # Initiate dataset configuration
    if config.dataset.startswith('ModelNet40'):
        dataset = ModelNet40Dataset(config.input_threads)
    elif config.dataset == 'S3DIS':
        dataset = S3DISDataset(config.input_threads)
        on_val = True
    elif config.dataset == 'Scannet':
        dataset = ScannetDataset(config.input_threads, load_test=True)
    elif config.dataset.startswith('ShapeNetPart'):
        dataset = ShapeNetPartDataset(
            config.dataset.split('_')[1], config.input_threads)
    elif config.dataset == 'NPM3D':
        dataset = NPM3DDataset(config.input_threads, load_test=True)
    elif config.dataset == 'Semantic3D':
        dataset = Semantic3DDataset(config.input_threads)
    else:
        raise ValueError('Unsupported dataset : ' + config.dataset)

    # Create subsample clouds of the models
    dl0 = config.first_subsampling_dl
    dataset.load_subsampled_clouds(dl0)

    # Initiate ERF input pipeleine (only diff is that it is not random)
    dataset.init_ERF_input_pipeline(config)

    ##############
    # Define Model
    ##############

    print('Creating Model')
    print('**************\n')
    t1 = time.time()

    if config.dataset.startswith('ShapeNetPart'):
        model = KernelPointFCNN(dataset.flat_inputs, config)
    elif config.dataset.startswith('S3DIS'):
        model = KernelPointFCNN(dataset.flat_inputs, config)
    elif config.dataset.startswith('Scannet'):
        model = KernelPointFCNN(dataset.flat_inputs, config)
    elif config.dataset.startswith('NPM3D'):
        model = KernelPointFCNN(dataset.flat_inputs, config)
    elif config.dataset.startswith('ModelNet40'):
        model = KernelPointCNN(dataset.flat_inputs, config)
    elif config.dataset.startswith('Semantic3D'):
        model = KernelPointFCNN(dataset.flat_inputs, config)
    else:
        raise ValueError('Unsupported dataset : ' + config.dataset)

    # Find all snapshot in the chosen training folder
    snap_path = os.path.join(path, 'snapshots')
    snap_steps = [
        int(f[:-5].split('-')[-1]) for f in os.listdir(snap_path)
        if f[-5:] == '.meta'
    ]

    # Find which snapshot to restore
    chosen_step = np.sort(snap_steps)[step_ind]
    chosen_snap = os.path.join(path, 'snapshots',
                               'snap-{:d}'.format(chosen_step))

    # Create a tester class
    visualizer = ModelVisualizer(model, restore_snap=chosen_snap)
    t2 = time.time()

    print('\n----------------')
    print('Done in {:.1f} s'.format(t2 - t1))
    print('----------------\n')

    #####################
    # Start visualization
    #####################

    print('Start visualization')
    print('*******************\n')

    visualizer.show_effective_recep_field(model, dataset, relu_idx)
Exemplo n.º 2
0
def test_caller(path, step_ind, on_val, dataset_path, noise, calc_tsne):
    ##########################
    # Initiate the environment
    ##########################

    # Choose which gpu to use
    GPU_ID = '0'

    # Set GPU visible device
    os.environ['CUDA_VISIBLE_DEVICES'] = GPU_ID

    # Disable warnings
    os.environ['TF_CPP_MIN_LOG_LEVEL'] = '0'

    ###########################
    # Load the model parameters
    ###########################

    # Load model parameters
    config = Config()
    config.load(path)

    ##################################
    # Change model parameters for test
    ##################################

    # Change parameters for the test here. For example, you can stop augmenting the input data.

    # config.augment_noise = 0.0001
    # config.augment_color = 1.0

    # Adjust batch num if only a single model is to be completed
    if on_val:
        val_data_paths = sorted([
            os.path.join(dataset_path, 'val', 'partial',
                         k.rstrip() + '.h5')
            for k in open(os.path.join(dataset_path, 'val.list')).readlines()
        ])
        if int(len(val_data_paths)) == 1:
            config.validation_size = 1
            config.batch_num = 1
    else:
        test_data_paths = sorted([
            os.path.join(dataset_path, 'test', 'partial',
                         k.rstrip() + '.h5')
            for k in open(os.path.join(dataset_path, 'val.list')).readlines()
        ])
        if int(len(test_data_paths)) == 1:
            config.validation_size = 1
            config.batch_num = 1

    # Augmentations
    config.augment_scale_anisotropic = True
    config.augment_symmetries = [False, False, False]
    config.augment_rotation = 'none'
    config.augment_scale_min = 1.0
    config.augment_scale_max = 1.0
    config.augment_noise = noise
    config.augment_occlusion = 'none'

    ##############
    # Prepare Data
    ##############

    print()
    print('Dataset Preparation')
    print('*******************')

    # Initiate dataset configuration
    dl0 = 0  # config.first_subsampling_dl
    if config.dataset.startswith('ShapeNetV1'):
        dataset = ShapeNetV1Dataset()
        # Create subsample clouds of the models
        dataset.load_subsampled_clouds(dl0)
    elif config.dataset.startswith("pc_shapenetCompletionBenchmark2048"):
        dataset = ShapeNetBenchmark2048Dataset(config.batch_num,
                                               config.num_input_points,
                                               dataset_path)
        # Create subsample clouds of the models
        dataset.load_subsampled_clouds(
            dl0)  # TODO: careful dl0 is used here - padding?
    else:
        raise ValueError('Unsupported dataset : ' + config.dataset)

    # Initialize input pipelines
    if on_val:
        dataset.init_input_pipeline(config)
    else:
        dataset.init_test_input_pipeline(config)

    ##############
    # Define Model
    ##############

    print('Creating Model')
    print('**************\n')
    t1 = time.time()

    if config.dataset.startswith('ShapeNetV1') or config.dataset.startswith(
            "pc_shapenetCompletionBenchmark2048"):
        model = KernelPointCompletionNetwork(dataset.flat_inputs, config,
                                             args.double_fold)
    else:
        raise ValueError('Unsupported dataset : ' + config.dataset)

    # Find all snapshot in the chosen training folder
    snap_path = os.path.join(path, 'snapshots')
    snap_steps = [
        int(f[:-5].split('-')[-1]) for f in os.listdir(snap_path)
        if f[-5:] == '.meta'
    ]

    # Find which snapshot to restore
    if step_ind == -1:
        chosen_step = np.sort(snap_steps)[step_ind]
    else:
        chosen_step = step_ind + 1
    chosen_snap = os.path.join(path, 'snapshots',
                               'snap-{:d}'.format(chosen_step))

    # Create a tester class
    tester = ModelTester(model, restore_snap=chosen_snap)
    t2 = time.time()

    print('\n----------------')
    print('Done in {:.1f} s'.format(t2 - t1))
    print('----------------\n')

    ############
    # Start test
    ############

    print('Start Test')
    print('**********\n')

    if config.dataset.startswith('ShapeNetV1') or config.dataset.startswith(
            "pc_shapenetCompletionBenchmark2048"):
        tester.test_completion(model, dataset, on_val, calc_tsne)
    else:
        raise ValueError('Unsupported dataset')