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,
)
elif dataset_type == 'kitti':
dataset = data.KittiDataset(dataset_folder,
img_size=cfg['data']['img_size'],
return_idx=return_idx)
elif dataset_type == 'online_products':
dataset = data.OnlineProductDataset(
dataset_folder,
img_size=cfg['data']['img_size'],
classes=cfg['data']['classes'],
max_number_imgs=cfg['generation']['max_number_imgs'],
return_idx=return_idx,
return_category=return_category)
elif dataset_type == 'images':
dataset = data.ImageDataset(
dataset_folder,
img_size=cfg['data']['img_size'],
return_idx=return_idx,
)
else:
raise ValueError('Invalid dataset "%s"' % cfg['data']['dataset'])
return dataset
def get_dataset(mode, cfg, return_idx=False, return_category=False, use_target_domain = 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
use_target_domain (bool): whether to use the target_domain dataset
'''
method = cfg['method']
dataset_type = cfg['data']['dataset']
if use_target_domain:
#dataset_type = cfg['data']['uda_dataset']
dataset_folder = cfg['data']['uda_path']
categories = cfg['data']['uda_classes']
else:
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, use_target_domain)
if inputs_field is not None:
fields['inputs'] = inputs_field
# adding field for UDA input when training
if mode == 'train' and cfg['training']['uda_type'] is not None:
# Also data-augment target domain imgs?
if cfg['data']['img_augment']:
resize_op = transforms.RandomResizedCrop(
cfg['data']['img_size'], (0.75, 1.), (1., 1.))
else:
resize_op = transforms.Resize((cfg['data']['img_size']))
transform = transforms.Compose([
resize_op, transforms.ToTensor(),
])
# random_view=True enables randomness
fields['inputs_target_domain'] = data.ImagesField(
#cfg['data']['uda_path_train'], transform=transform, random_view=True, image_based_hier=True
cfg['data']['uda_path_train'], transform=transform, random_view=True, extensions=['jpg', 'jpeg', 'png'], image_based_hier=True
)
if return_idx:
fields['idx'] = data.IndexField()
if return_category:
fields['category'] = data.CategoryField()
dataset = data.Shapes3dDataset(
dataset_folder, fields,
split=split,
categories=categories,
)
elif dataset_type == 'kitti':
dataset = data.KittiDataset(
dataset_folder, img_size=cfg['data']['img_size'],
return_idx=return_idx
)
elif dataset_type == 'online_products':
dataset = data.OnlineProductDataset(
dataset_folder, img_size=cfg['data']['img_size'],
classes=cfg['data']['classes'],
max_number_imgs=cfg['generation']['max_number_imgs'],
return_idx=return_idx, return_category=return_category
)
elif dataset_type == 'images':
dataset = data.ImageDataset(
dataset_folder, img_size=cfg['data']['img_size'],
return_idx=return_idx,
)
else:
raise ValueError('Invalid dataset "%s"' % cfg['data']['dataset'])
return dataset
def get_dataset(cfg, mode='train', return_idx=False, return_category=False,
**kwargs):
''' Returns a dataset instance.
Args:
cfg (dict): config dictionary
mode (string): which mode is used (train / val /test / render)
return_idx (bool): whether to return model index
return_category (bool): whether to return model category
'''
# Get fields with cfg
method = cfg['method']
input_type = cfg['data']['input_type']
dataset_name = cfg['data']['dataset_name']
dataset_folder = cfg['data']['path']
categories = cfg['data']['classes']
cache_fields = cfg['data']['cache_fields']
n_views = cfg['data']['n_views']
split_model_for_images = cfg['data']['split_model_for_images']
splits = {
'train': cfg['data']['train_split'],
'val': cfg['data']['val_split'],
'test': cfg['data']['test_split'],
'render': cfg['data']['test_split'],
}
split = splits[mode]
fields = method_dict[method].config.get_data_fields(cfg, mode=mode)
if input_type == 'idx':
input_field = data.IndexField()
fields['inputs'] = input_field
elif input_type == 'image':
random_view = True if \
(mode == 'train' or dataset_name == 'NMR') else False
resize_img_transform = data.ResizeImage(cfg['data']['img_size_input'])
fields['inputs'] = data.ImagesField(
cfg['data']['img_folder_input'],
transform=resize_img_transform,
with_mask=False, with_camera=False,
extension=cfg['data']['img_extension_input'],
n_views=cfg['data']['n_views_input'], random_view=random_view)
else:
input_field = None
if return_idx:
fields['idx'] = data.IndexField()
if return_category:
fields['category'] = data.CategoryField()
manager = Manager()
shared_dict = manager.dict()
if ((dataset_name == 'Shapes3D') or
(dataset_name == 'DTU') or
(dataset_name == 'NMR')):
dataset = data.Shapes3dDataset(
dataset_folder, fields, split=split,
categories=categories,
shared_dict=shared_dict,
n_views=n_views, cache_fields=cache_fields,
split_model_for_images=split_model_for_images)
elif dataset_name == 'images':
dataset = data.ImageDataset(
dataset_folder, return_idx=True
)
else:
raise ValueError('Invalid dataset_name!')
return dataset
def get_dataset(mode,
cfg,
batch_size,
shuffle,
repeat_count,
epoch,
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,
batch_size=batch_size,
shuffle=shuffle,
repeat_count=repeat_count,
epoch=epoch,
categories=categories,
)
elif dataset_type == "kitti":
dataset = data.KittiDataset(dataset_folder,
batch_size=batch_size,
shuffle=shuffle,
img_size=cfg["data"]["img_size"],
return_idx=return_idx)
elif dataset_type == "online_products":
dataset = data.OnlineProductDataset(
dataset_folder,
batch_size=batch_size,
shuffle=shuffle,
img_size=cfg["data"]["img_size"],
classes=cfg["data"]["classes"],
max_number_imgs=cfg["generation"]["max_number_imgs"],
return_idx=return_idx,
return_category=return_category,
)
elif dataset_type == "images":
dataset = data.ImageDataset(
dataset_folder,
batch_size=batch_size,
shuffle=shuffle,
img_size=cfg["data"]["img_size"],
return_idx=return_idx,
)
else:
raise ValueError('Invalid dataset "%s"' % cfg["data"]["dataset"])
return dataset
