def save_image_file(image, img_type, dataset, path):
"""
Brings the image in the right format to be saved. Therefore, it reverses all load transforms
:param image: PIL image
:param img_type: color, depth, segmentation
:param dataset: e.g. 'cityscapes'
:param path: path of the destination path
"""
depth_transformer = mytransforms.LoadDepth()
segmentation_transformer = mytransforms.LoadSegmentation()
image = np.array(image) # convert from PIL to openCV
if 'color' in img_type:
image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
elif 'depth' in img_type:
image = depth_transformer.inverse({img_type: image})[img_type]
# image = image.astype(np.uint16)
elif 'segmentation' in img_type:
image = segmentation_transformer.inverse({img_type: image})[img_type]
# image = np.array(image, dtype=np.uint8)
cv2.imwrite(path, image)
def _load_sample(self, si):
"""
Loads the specified sample.
:param si: dict of the form {global_index: {"color": path, "depth": path, ...}}
:return: - the global index of the sample in the dataset
- the sample in the standard dataloader format
- a path dictionary that is built similar to the sample dictionary and contains the image path instead of
an image object
"""
set_idx, sample = si
dataset = sample.pop('dataset')
dataset_path = sample.pop('dataset_path')
new_sample = dict()
paths = dict()
for key, content in sample.items():
assert any(key.startswith(s) for s in SEGMENTATION_KEYS)
filepath = os.path.join(dataset_path, content)
filepath = filepath.replace('/', os.sep)
filepath = filepath.replace('\\', os.sep)
image = cv2.imread(filepath, -1)
new_sample[key, 0, 0] = image
paths[key, 0, 0] = content
# By using the transform ConvertSegmentation, the segmentation images will be converted to the train_id format
load_transforms = transforms.Compose([
mytransforms.LoadSegmentation(),
mytransforms.ConvertSegmentation(labels=self.labels, labels_mode=self.labels_mode)
])
new_sample = load_transforms(new_sample)
return set_idx, new_sample, paths
def load_sample(si):
"""
Loads the specified sample.
:param si: dict of the form {global_index: {"color": path, "depth": path, ...}}
:return: - the global index of the sample in the dataset
- the sample in the standard dataloader format
- a path dictionary that is built similar to the sample dictionary and contains the image path instead of
an image object
"""
set_idx, sample = si
dataset = sample.pop('dataset')
dataset_path = sample.pop('dataset_path')
depth_mode = sample.pop('depth_mode')
new_sample = dict()
paths = dict()
for key, content in sample.items():
if any(key.startswith(s) for s in IMAGE_KEYS):
filepath = os.path.join(dataset_path, content)
filepath = filepath.replace('/', os.sep)
filepath = filepath.replace('\\', os.sep)
image = cv2.imread(filepath, -1)
new_sample[key, 0, 0] = image
paths[key, 0, 0] = content
else:
new_sample[key, 0, 0] = content
load_transforms = transforms.Compose([
mytransforms.LoadRGB(),
mytransforms.LoadSegmentation(),
mytransforms.LoadDepth(),
mytransforms.LoadNumerics()
])
new_sample = load_transforms(new_sample)
return set_idx, new_sample, paths
def __init__(self,
dataset,
trainvaltest_split,
video_mode='mono',
stereo_mode='mono',
cluster_mode=None,
simple_mode=False,
labels=None,
labels_mode=None,
data_transforms=None,
scales=None,
keys_to_load=None,
keys_to_video=None,
keys_to_stereo=None,
split=None,
video_frames=None,
disable_const_items=True,
folders_to_load=None,
files_to_load=None,
n_files=None,
output_filenames=False,
flow_validation_mode=True):
"""Initialises the dataset by loading the desired data from the json file
:param dataset: name of the dataset
:param trainvaltest_split: can be train, validation or test
:param video_mode: can be mono or video and defines if only the images or image sequences are to be loaded
:param stereo_mode: can be mono or stereo and defines if the stereo images are to be loaded
:param simple_mode: if True, the Data is read directly from a folder without using a .json file
:param labels: gives the labels as defined in the named tuples style in Cityscapes. Get the labels from
defintions folder
:param labels_mode: can be fromid or fromrgb and defines if the segmentation masks are given as id or color
:param data_transforms: takes the transforms.compose list
:param scales: list of all scales at which the images should be loaded (list of exponents for powers of 2)
:param keys_to_load: defines all keys which should be loaded
:param keys_to_video: defines for which keys the sequences are to be loaded
:param keys_to_stereo: defines for which keys the stereo images are supposed to be loaded
:param split: dataset split that is supposed to be loaded. default is the complete dataset itself
:param video_frames: all frames of the sequence that are supposed to be loaded (list of frame numbers relative
to the main frame, e.g. [0, -2, -1, 1, 2])
:param disable_const_items: removes the constant items like camera calibration from loading procedure
:param folders_to_load: list of folders from which data should be loaded; folders not mentioned are skipped in
the respective set. Only the last folder in a path is considered; filter is case insensitive.
Default: None -> all folders are loaded from dataset
:param files_to_load: list of files that should be loaded; files not mentioned are skipped in the respective
set. File names need not be complete; filter is case insensitive.
Default: None -> all files are loaded from dataset
:param n_files: How many files shall be loaded. Files are selected randomly if there are more files than n_files
Seeded by numpy.random.seed()
"""
super(BaseDataset, self).__init__()
assert isinstance(dataset, str)
assert trainvaltest_split in (
'train', 'validation',
'test'), '''trainvaltest_split must be train,
validation or test'''
assert video_mode in ('mono',
'video'), 'video_mode must be mono or video'
assert stereo_mode in ('mono',
'stereo'), 'stereo_mode must be mono or stereo'
assert isinstance(simple_mode, bool)
if data_transforms is None:
data_transforms = [
mytransforms.CreateScaledImage(),
mytransforms.CreateColoraug(),
mytransforms.ToTensor()
]
if scales is None:
scales = [0]
if keys_to_load is None:
keys_to_load = ['color']
if keys_to_stereo is None and stereo_mode == 'stereo':
keys_to_stereo = ['color']
if keys_to_video is None and video_mode == 'video':
keys_to_video = ['color']
if video_frames is None:
video_frames = [0, -1, 1]
self.dataset = dataset
self.video_mode = video_mode
self.stereo_mode = stereo_mode
self.scales = scales
self.disable_const_items = disable_const_items
self.output_filenames = output_filenames
self.parameters = dps.DatasetParameterset(dataset)
if labels is not None:
self.parameters.labels = labels
if labels_mode is not None:
self.parameters.labels_mode = labels_mode
path_getter = gp.GetPath()
dataset_folder = path_getter.get_data_path()
datasetpath = os.path.join(dataset_folder, self.dataset)
self.datasetpath = datasetpath
if split is None:
splitpath = None
else:
splitpath = os.path.join(dataset_folder,
self.dataset + '_' + split)
if simple_mode is False:
self.data = self.read_json_file(datasetpath, splitpath,
trainvaltest_split, keys_to_load,
keys_to_stereo, keys_to_video,
video_frames, folders_to_load,
files_to_load, n_files)
else:
self.data = self.read_from_folder(datasetpath, keys_to_load,
video_mode, video_frames)
self.load_transforms = transforms.Compose([
mytransforms.LoadRGB(),
mytransforms.LoadSegmentation(),
mytransforms.LoadDepth(),
mytransforms.LoadFlow(validation_mode=flow_validation_mode),
mytransforms.LoadNumerics()
])
# IMPORTANT to create a new list if the same list is passed to multiple datasets. Otherwise, due to the
# mutability of lists, ConvertSegmentation will only be added once. Hence, the labels may be wrong for the 2nd,
# 3rd, ... dataset!
self.data_transforms = list(data_transforms)
# Error if CreateColorAug and CreateScaledImage not in transforms.
if mytransforms.CreateScaledImage not in data_transforms:
raise Exception(
'The transform CreateScaledImage() has to be part of the data_transforms list'
)
if mytransforms.CreateColoraug not in data_transforms:
raise Exception(
'The transform CreateColoraug() has to be part of the data_transforms list'
)
# Error if depth, segmentation or flow keys are given but not the corresponding Convert-Transform
if any([key.startswith('segmentation') for key in keys_to_load]) and \
mytransforms.ConvertSegmentation not in self.data_transforms:
raise Exception(
'When loading segmentation images, please add mytransforms.ConvertSegmentation() to '
'the data_transforms')
if any([key.startswith('depth') for key in keys_to_load]) and \
mytransforms.ConvertDepth not in self.data_transforms:
raise Exception(
'When loading depth images, please add mytransforms.ConvertDepth() to the data_transforms'
)
if any([key.startswith('flow') for key in keys_to_load]) and \
mytransforms.ConvertFlow not in self.data_transforms:
raise Exception(
'When loading flow images, please add mytransforms.ConvertFlow() to the data_transforms'
)
# In the flow validation mode, it is not allowed to use data-altering transforms
if any([key.startswith('flow')
for key in keys_to_load]) and flow_validation_mode:
allowed_transforms = [
mytransforms.CreateScaledImage, mytransforms.CreateColoraug,
mytransforms.ConvertSegmentation, mytransforms.ConvertDepth,
mytransforms.ConvertFlow, mytransforms.RemoveOriginals,
mytransforms.ToTensor, mytransforms.Relabel,
mytransforms.OneHotEncoding, mytransforms.NormalizeZeroMean,
mytransforms.AdjustKeys, mytransforms.RemapKeys,
mytransforms.AddKeyValue
]
for transform in self.data_transforms:
if transform not in allowed_transforms:
raise Exception(
'In flow validation mode, it is not allowed to use data-altering transforms'
)
# Set the correct parameters to the ConvertDepth and ConvertSegmentation transforms
for i, transform in zip(range(len(self.data_transforms)),
self.data_transforms):
if isinstance(transform, mytransforms.ConvertDepth):
transform.set_mode(self.parameters.depth_mode)
elif isinstance(transform, mytransforms.ConvertSegmentation):
transform.set_mode(self.parameters.labels,
self.parameters.labels_mode)
elif isinstance(transform, mytransforms.ConvertFlow):
transform.set_mode(self.parameters.flow_mode,
flow_validation_mode)
self.data_transforms = transforms.Compose(self.data_transforms)