def get_fields():
fields = {}
input_field = config.get_inputs_field('train', cfg)
fields['inputs'] = input_field
fields['idx'] = data.IndexField()
fields['viewid'] = data.ViewIdField()
return fields
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 == 'img_seq':
if mode == 'train' and 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(),
])
if mode == 'train':
random_view = True
else:
random_view = False
inputs_field = data.ImageSubseqField(cfg['data']['img_seq_folder'],
transform,
random_view=random_view)
elif input_type == 'pcl_seq':
connected_samples = cfg['data']['input_pointcloud_corresponding']
transform = transforms.Compose([
data.SubsamplePointcloudSeq(cfg['data']['input_pointcloud_n'],
connected_samples=connected_samples),
data.PointcloudNoise(cfg['data']['input_pointcloud_noise'])
])
inputs_field = data.PointCloudSubseqField(
cfg['data']['pointcloud_seq_folder'],
transform,
seq_len=cfg['data']['length_sequence'])
elif input_type == 'end_pointclouds':
transform = data.SubsamplePointcloudSeq(
cfg['data']['input_pointcloud_n'],
connected_samples=cfg['data']['input_pointcloud_corresponding'])
inputs_field = data.PointCloudSubseqField(
cfg['data']['pointcloud_seq_folder'],
only_end_points=True,
seq_len=cfg['data']['length_sequence'],
transform=transform)
elif input_type == 'idx':
inputs_field = data.IndexField()
else:
raise ValueError('Invalid input type (%s)' % input_type)
return inputs_field
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']
with_transforms = cfg['data']['with_transforms']
if input_type is None:
inputs_field = None
elif input_type == 'img':
if mode == 'train' and 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(),
])
with_camera = cfg['data']['img_with_camera']
if mode == 'train':
random_view = True
else:
random_view = False
inputs_field = data.ImagesField(
cfg['data']['img_folder'], transform,
with_camera=with_camera, random_view=random_view
)
elif input_type == 'pointcloud':
transform = transforms.Compose([
data.SubsamplePointcloud(cfg['data']['pointcloud_n']),
data.PointcloudNoise(cfg['data']['pointcloud_noise'])
])
with_transforms = cfg['data']['with_transforms']
inputs_field = data.PointCloudField(
cfg['data']['pointcloud_file'], transform,
with_transforms=with_transforms
)
elif input_type == 'voxels':
inputs_field = data.VoxelsField(
cfg['data']['voxels_file']
)
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 == 'Humans':
# 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.HumansDataset(
dataset_folder,
fields,
split=split,
categories=categories,
length_sequence=cfg['data']['length_sequence'],
n_files_per_sequence=cfg['data']['n_files_per_sequence'],
offset_sequence=cfg['data']['offset_sequence'],
ex_folder_name=cfg['data']['pointcloud_seq_folder'],
n_intervals=cfg['data']['n_intervals'])
else:
raise ValueError('Invalid dataset "%s"' % cfg['data']['dataset'])
return dataset
def get_dataset(dataset_root, mode, input_type='img', absolute_depth=True):
if input_type == 'img':
inputs_field = get_img_inputs_field(mode)
elif input_type == 'img_with_depth':
inputs_field = get_img_with_depth_input_field(mode, absolute_depth=absolute_depth)
elif input_type == 'depth_pointcloud':
inputs_field = get_depth_pointcloud_field(mode)
fields = {}
fields['inputs'] = inputs_field
fields['idx'] = data.IndexField()
fields['category'] = data.CategoryField()
dataset = data.Shapes3dDataset(
dataset_root, fields,
split=mode,
categories=None,
)
return dataset
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_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']
with_transforms = cfg['data']['with_transforms']
if input_type is None:
inputs_field = None
elif input_type == 'img':
if mode == 'train' and 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(),
])
with_camera = cfg['data']['img_with_camera']
if mode == 'train':
random_view = True
else:
random_view = False
if 'img_extension' in cfg['data']:
inputs_field = data.ImagesField(
cfg['data']['img_folder'],
transform,
extension=cfg['data']['img_extension'],
with_camera=with_camera,
random_view=random_view)
else:
inputs_field = data.ImagesField(cfg['data']['img_folder'],
transform,
with_camera=with_camera,
random_view=random_view)
elif input_type == 'img_with_depth':
# data augment not supported
transform = transforms.Compose([
transforms.Resize((cfg['data']['img_size'])),
transforms.ToTensor(),
])
with_camera = cfg['data']['img_with_camera']
if mode == 'train':
random_view = True
else:
random_view = False
data_params = {
'with_camera': with_camera,
'random_view': random_view,
}
if 'absolute_depth' in cfg['data']:
data_params['absolute_depth'] = cfg['data']['absolute_depth']
if 'with_minmax' in cfg['data']:
data_params['with_minmax'] = cfg['data']['with_minmax']
if 'img_extension' in cfg['data']:
data_params['extension'] = cfg['data']['img_extension']
inputs_field = data.ImagesWithDepthField('img', 'depth', 'mask',
transform, **data_params)
elif input_type == 'depth_pred':
# data augment not supported
transform = transforms.Compose([
transforms.Resize((cfg['data']['img_size'])),
transforms.ToTensor(),
])
with_camera = cfg['data']['img_with_camera']
if mode == 'train':
random_view = True
else:
random_view = False
data_params = {
'with_camera': with_camera,
'random_view': random_view,
}
if 'absolute_depth' in cfg['data']:
data_params['absolute_depth'] = cfg['data']['absolute_depth']
if 'with_minmax' in cfg['data']:
data_params['with_minmax'] = cfg['data']['with_minmax']
if 'pred_with_img' in cfg['model']:
data_params['with_img'] = cfg['model']['pred_with_img']
if 'img_extension' in cfg['data']:
data_params['extension'] = cfg['data']['img_extension']
inputs_field = data.DepthPredictedField('img', 'depth', 'mask',
cfg['data']['depth_pred_root'],
'depth_pred', transform,
**data_params)
elif input_type == 'depth_pointcloud':
t_lst = []
if 'depth_pointcloud_n' in cfg['data'] and cfg['data'][
'depth_pointcloud_n'] is not None:
t_lst.append(
data.SubsampleDepthPointcloud(
cfg['data']['depth_pointcloud_n']))
if 'depth_pointcloud_noise' in cfg['data'] and cfg['data'][
'depth_pointcloud_noise'] is not None:
t_lst.append(
data.PointcloudNoise(cfg['data']['depth_pointcloud_noise']))
transform = transforms.Compose(t_lst)
if mode == 'train':
random_view = True
else:
random_view = False
data_params = {
'random_view': random_view,
'with_camera': True,
'img_folder_name': 'img'
}
if 'view_penalty' in cfg['training'] and cfg['training'][
'view_penalty']:
data_params['with_mask'] = True
data_params['mask_folder_name'] = 'mask'
data_params['mask_flow_folder_name'] = 'mask_flow'
data_params['extension'] = 'png'
img_transform = transforms.Compose([
transforms.Resize((cfg['data']['img_size'])),
transforms.ToTensor(),
])
data_params['img_transform'] = img_transform
data_params['with_depth_pred'] = True
data_params['depth_pred_folder_name'] = 'depth_pred'
inputs_field = data.DepthPointCloudField(
cfg['data']['depth_pointcloud_root'],
cfg['data']['depth_pointcloud_folder'], transform, **data_params)
elif input_type == 'multi_img':
if mode == 'train' and 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(),
])
with_camera = cfg['data']['img_with_camera']
if mode == 'train':
random_view = True
else:
random_view = False
inputs_field = data.ImagesField(cfg['data']['img_folder'],
transform,
with_camera=with_camera,
random_view=random_view)
elif input_type == 'pointcloud':
transform = transforms.Compose([
data.SubsamplePointcloud(cfg['data']['pointcloud_n']),
data.PointcloudNoise(cfg['data']['pointcloud_noise'])
])
with_transforms = cfg['data']['with_transforms']
inputs_field = data.PointCloudField(cfg['data']['pointcloud_file'],
transform,
with_transforms=with_transforms)
elif input_type == 'voxels':
inputs_field = data.VoxelsField(cfg['data']['voxels_file'])
elif input_type == 'idx':
inputs_field = data.IndexField()
else:
raise ValueError('Invalid input type (%s)' % input_type)
return inputs_field
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'],
)
pointcloud_field = data.PointCloudField(cfg['data']['pointcloud_chamfer_file'])
fields = {
'points_iou': points_field,
'pointcloud_chamfer': pointcloud_field,
'idx': data.IndexField(),
}
print('Test split: ', cfg['data']['test_split'])
if args.da:
dataset_folder = cfg['data']['uda_path']
categories_to_use = cfg['data']['uda_classes']
else:
dataset_folder = cfg['data']['path']
categories_to_use = cfg['data']['classes']
dataset = data.Shapes3dDataset(dataset_folder,
fields,
cfg['data']['test_split'],
categories=categories_to_use)
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
if not args.eval_input:
out_file = os.path.join(generation_dir, 'eval_meshes_full.pkl')
out_file_tmp = os.path.join(generation_dir, 'eval_meshes_full_tmp.pkl')
out_file_class = os.path.join(generation_dir, 'eval_meshes.csv')
out_file_class_tmp = os.path.join(generation_dir, 'eval_meshes_tmp.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
fields = {
'pointcloud_chamfer':
data.PointCloudSubseqField(cfg['data']['pointcloud_seq_folder'],
seq_len=cfg['data']['length_sequence']),
'idx':
data.IndexField(),
}
if cfg['test']['eval_mesh_iou']:
fields['points'] = data.PointsSubseqField(
cfg['data']['points_iou_seq_folder'],
all_steps=True,
seq_len=cfg['data']['length_sequence'])
print('Test split: ', cfg['data']['test_split'])
dataset = data.HumansDataset(
dataset_folder,
fields,
split=cfg['data']['test_split'],
categories=cfg['data']['classes'],
length_sequence=cfg['data']['length_sequence'],
def get_inputs_field(mode, cfg, use_target_domain = False):
''' Returns the inputs fields.
Args:
mode (str): the mode which is used
cfg (dict): config dictionary
use_target_domain (bool): whether to use the target_domain dataset
'''
input_type = cfg['data']['input_type']
with_transforms = cfg['data']['with_transforms']
if input_type is None:
inputs_field = None
elif input_type == 'img':
if mode == 'train' and 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(),
])
with_camera = cfg['data']['img_with_camera']
if mode == 'train':
random_view = True
else:
random_view = False
if use_target_domain:
img_folder_name = cfg['data']['uda_img_folder']
filename_pattern= cfg['data']['uda_bg_configure']
else:
img_folder_name = cfg['data']['img_folder']
filename_pattern = cfg['data']['img_filename_pattern']
inputs_field = data.ImagesField(
img_folder_name, transform,
with_camera=with_camera, random_view=random_view, filename_pattern=filename_pattern, extensions=['jpg', 'jpeg', 'png']
)
elif input_type == 'pointcloud':
transform = transforms.Compose([
data.SubsamplePointcloud(cfg['data']['pointcloud_n']),
data.PointcloudNoise(cfg['data']['pointcloud_noise'])
])
with_transforms = cfg['data']['with_transforms']
inputs_field = data.PointCloudField(
cfg['data']['pointcloud_file'], transform,
with_transforms=with_transforms
)
elif input_type == 'voxels':
inputs_field = data.VoxelsField(
cfg['data']['voxels_file']
)
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, 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_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"]
with_transforms = cfg["data"]["with_transforms"]
if input_type is None:
inputs_field = None
elif input_type == "img":
transform = None
if mode == "train" and cfg["data"]["img_augment"]:
# resize_op = transforms.RandomResizedCrop(cfg["data"]["img_size"],(0.75, 1.0), (1.0, 1.0))
def preprocess(image):
# image = tf.image.crop_and_resize(
# image, crop_size=cfg["data"]["img_size"]) # CHECK
image = tf.image.resize(
image, [cfg["data"]["img_size"], cfg["data"]["img_size"]])
image /= 255.0
return image
transform = preprocess
else:
def preprocess(image):
# image = image[tf.newaxis, ...]
image = tf.image.resize(
image, [cfg["data"]["img_size"], cfg["data"]["img_size"]])
image /= 255.0
return image
transform = preprocess
# transform = transforms.Compose([
# resize_op,
# transforms.ToTensor(),
# ])
with_camera = cfg["data"]["img_with_camera"]
if mode == "train":
random_view = True
else:
random_view = False
inputs_field = data.ImagesField(
cfg["data"]["img_folder"],
transform,
with_camera=with_camera,
random_view=random_view,
)
elif input_type == "pointcloud":
def preprocess(points):
output = data.SubsamplePointcloud(
cfg["data"]["pointcloud_n"])(points)
output = data.PointcloudNoise(
cfg["data"]["pointcloud_noise"])(output)
return output
transform = preprocess
with_transforms = cfg["data"]["with_transforms"]
inputs_field = data.PointCloudField(cfg["data"]["pointcloud_file"],
transform,
with_transforms=with_transforms)
elif input_type == "voxels":
inputs_field = data.VoxelsField(cfg["data"]["voxels_file"])
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,
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