def get_inputs_field(mode, cfg): ''' Returns the inputs fields. Args: mode (str): the mode which is used cfg (dict): config dictionary ''' input_type = cfg['data']['input_type'] if input_type is None: inputs_field = None elif input_type == 'pointcloud': transform = transforms.Compose([ data.SubsamplePointcloud(cfg['data']['pointcloud_n']), data.PointcloudNoise(cfg['data']['pointcloud_noise']) ]) inputs_field = data.PointCloudField( cfg['data']['pointcloud_file'], transform=transform, multi_files=cfg['data']['multi_files']) elif input_type == 'idx': inputs_field = data.IndexField() else: raise ValueError('Invalid input type (%s)' % input_type) return inputs_field
def get_dataset(mode, cfg, return_idx=False, return_category=False): ''' Returns the dataset. Args: model (nn.Module): the model which is used cfg (dict): config dictionary return_idx (bool): whether to include an ID field ''' method = cfg['method'] dataset_type = cfg['data']['dataset'] dataset_folder = cfg['data']['path'] categories = cfg['data']['classes'] # Get split splits = { 'train': cfg['data']['train_split'], 'val': cfg['data']['val_split'], 'test': cfg['data']['test_split'], } split = splits[mode] # Create dataset if dataset_type == 'Shapes3D': # Dataset fields # Method specific fields (usually correspond to output) fields = method_dict[method].config.get_data_fields(mode, cfg) # Input fields inputs_field = get_inputs_field(mode, cfg) if inputs_field is not None: fields['inputs'] = inputs_field if return_idx: fields['idx'] = data.IndexField() if return_category: fields['category'] = data.CategoryField() dataset = data.Shapes3dDataset(dataset_folder, fields, split=split, categories=categories) else: raise ValueError('Invalid dataset "%s"' % cfg['data']['dataset']) return dataset
out_file = os.path.join(generation_dir, 'eval_meshes_full.pkl') out_file_class = os.path.join(generation_dir, 'eval_meshes.csv') else: out_file = os.path.join(generation_dir, 'eval_input_full.pkl') out_file_class = os.path.join(generation_dir, 'eval_input.csv') # Dataset points_field = data.PointsField(cfg['data']['points_iou_file'], unpackbits=cfg['data']['points_unpackbits'], multi_files=cfg['data']['multi_files']) pointcloud_field = data.PointCloudField(cfg['data']['pointcloud_chamfer_file'], multi_files=cfg['data']['multi_files']) fields = { 'points_iou': points_field, 'pointcloud_chamfer': pointcloud_field, 'idx': data.IndexField(), } print('Test split: ', cfg['data']['test_split']) dataset_folder = cfg['data']['path'] dataset = data.Shapes3dDataset(dataset_folder, fields, cfg['data']['test_split'], categories=cfg['data']['classes'], cfg=cfg) # Evaluator evaluator = MeshEvaluator(n_points=100000) # Loader