def _verify_classes(self, classes):
categories = ['bedroom', 'bridge', 'church_outdoor', 'classroom',
'conference_room', 'dining_room', 'kitchen',
'living_room', 'restaurant', 'tower']
dset_opts = ['train', 'val', 'test']
try:
verify_str_arg(classes, "classes", dset_opts)
if classes == 'test':
classes = [classes]
else:
classes = [c + '_' + classes for c in categories]
except ValueError:
if not isinstance(classes, Iterable):
msg = ("Expected type str or Iterable for argument classes, "
"but got type {}.")
raise ValueError(msg.format(type(classes)))
classes = list(classes)
msg_fmtstr = ("Expected type str for elements in argument classes, "
"but got type {}.")
for c in classes:
verify_str_arg(c, custom_msg=msg_fmtstr.format(type(c)))
c_short = c.split('_')
category, dset_opt = '_'.join(c_short[:-1]), c_short[-1]
msg_fmtstr = "Unknown value '{}' for {}. Valid values are {{{}}}."
msg = msg_fmtstr.format(category, "LSUN class", iterable_to_str(categories))
verify_str_arg(category, valid_values=categories, custom_msg=msg)
msg = msg_fmtstr.format(dset_opt, "postfix", iterable_to_str(dset_opts))
verify_str_arg(dset_opt, valid_values=dset_opts, custom_msg=msg)
return classes
def verify_mode_type(self, split, image_mode, image_type):
image_mode = self.verify_mode(image_mode)
if image_mode == "gtFine":
valid_splits = ("train", "test", "val")
valid_types = ("_instanceIds.png", "instance", "_labelIds.png",
"semantic", "_color.png", "color", "_polygons.json",
"polygon")
elif image_mode == "gtCoarse":
valid_splits = ("train", "train_extra", "val")
valid_types = ("_instanceIds.png", "instance", "_labelIds.png",
"semantic", "_color.png", "color", "_polygons.json",
"polygon")
elif image_mode == "leftImg8bit":
valid_splits = ("train", "train_extra", "test", "val")
valid_types = (
"_leftImg8bit.png",
None,
)
for i in range(len(valid_types) // 2):
if image_type == valid_types[i * 2 + 1]:
image_type = valid_types[i * 2]
break
msg = ("Unknown value '{}' for argument split if image_mode is '{}'. "
"Valid values are {{{}}}.")
msg = msg.format(split, image_mode, iterable_to_str(valid_splits))
verify_str_arg(split, "split", valid_splits, msg)
msg = (
"Unknown value '{}' for argument image_type if image_mode is '{}'. "
"Valid values are {{{}}}.")
msg = msg.format(image_type, image_mode, iterable_to_str(valid_types))
verify_str_arg(image_type, "image_type", valid_types, msg)
return image_mode, image_type
def __init__(self, root, split='train', transform=None, target_transform=None, transforms=None):
super(Deepfashion, self).__init__(root, transforms, transform, target_transform)
self.images_dir = os.path.join(self.root, 'img')
self.targets_dir = os.path.join(self.root, 'lbl')
self.split = split
self.images = []
self.targets = []
valid_modes = ("train", "test", "val")
msg = ("Unknown value '{}' for argument split if mode is '{}'. "
"Valid values are {{{}}}.")
msg = msg.format(split, mode, iterable_to_str(valid_modes))
verify_str_arg(split, "split", valid_modes, msg)
if not os.path.isdir(self.images_dir) or not os.path.isdir(self.targets_dir):
image_dir_zip = os.path.join(self.root, '{}'.format('img.zip'))
target_dir_zip = os.path.join(self.root, '{}'.format('lbl.zip'))
if os.path.isfile(image_dir_zip) and os.path.isfile(target_dir_zip):
extract_archive(from_path=image_dir_zip, to_path=self.root)
extract_archive(from_path=target_dir_zip, to_path=self.root)
else:
raise RuntimeError('Dataset not found or incomplete. Please make sure all required folders for the'
' specified "split" and "mode" are inside the "root" directory')
data_list = pd.read_csv(os.path.join(self.root, 'list_eval_partition.txt'), sep='\t', skiprows=1)
data_list = data_list[data_list['evaluation_status'] == self.split]
for image_path in data_list['image_name']:
target_path = 'lbl/' + '/'.join(image_path.split('/')[1:])
self.images.append(image_path)
self.targets.append(target_path)
def __init__(self, root, split='train', mode='fine', target_type='instance',
transform=None, target_transform=None, transforms=None):
super(Cityscapes, self).__init__(root, transforms, transform, target_transform)
self.mode = 'gtFine' if mode == 'fine' else 'gtCoarse'
self.images_dir = os.path.join(self.root, 'leftImg8bit', split)
self.targets_dir = os.path.join(self.root, self.mode, split)
self.target_type = target_type
self.split = split
self.images = []
self.targets = []
verify_str_arg(mode, "mode", ("fine", "coarse"))
if mode == "fine":
valid_modes = ("train", "test", "val")
else:
valid_modes = ("train", "train_extra", "val")
msg = ("Unknown value '{}' for argument split if mode is '{}'. "
"Valid values are {{{}}}.")
msg = msg.format(split, mode, iterable_to_str(valid_modes))
verify_str_arg(split, "split", valid_modes, msg)
if not isinstance(target_type, list):
self.target_type = [target_type]
[verify_str_arg(value, "target_type",
("instance", "semantic", "polygon", "color"))
for value in self.target_type]
if not os.path.isdir(self.images_dir) or not os.path.isdir(self.targets_dir):
if split == 'train_extra':
image_dir_zip = os.path.join(self.root, 'leftImg8bit{}'.format('_trainextra.zip'))
else:
image_dir_zip = os.path.join(self.root, 'leftImg8bit{}'.format('_trainvaltest.zip'))
if self.mode == 'gtFine':
target_dir_zip = os.path.join(self.root, '{}{}'.format(self.mode, '_trainvaltest.zip'))
elif self.mode == 'gtCoarse':
target_dir_zip = os.path.join(self.root, '{}{}'.format(self.mode, '.zip'))
if os.path.isfile(image_dir_zip) and os.path.isfile(target_dir_zip):
extract_archive(from_path=image_dir_zip, to_path=self.root)
extract_archive(from_path=target_dir_zip, to_path=self.root)
else:
raise RuntimeError('Dataset not found or incomplete. Please make sure all required folders for the'
' specified "split" and "mode" are inside the "root" directory')
for city in os.listdir(self.images_dir):
img_dir = os.path.join(self.images_dir, city)
target_dir = os.path.join(self.targets_dir, city)
for file_name in os.listdir(img_dir):
target_types = []
for t in self.target_type:
target_name = file_name
target_types.append(os.path.join(target_dir, target_name))
self.images.append(os.path.join(img_dir, file_name))
self.targets.append(target_types)
def __init__(
self,
root: str,
split: str = "train",
mode: str = "fine",
target_type: Union[List[str], str] = "semantic",
transform: Optional[Callable] = None,
target_transform: Optional[Callable] = None,
transforms: Optional[Callable] = None,
) -> None:
super(Cityscapes, self).__init__(root, transforms, transform,
target_transform)
self.mode = 'gtFine' if mode == 'fine' else 'gtCoarse'
self.images_dir = os.path.join(self.root, 'leftImg8bit', split)
self.targets_dir = os.path.join(self.root, self.mode, split)
self.target_type = target_type
self.split = split
self.images = []
self.targets = []
verify_str_arg(mode, "mode", ("fine", "coarse"))
if mode == "fine":
valid_modes = ("train", "test", "val")
else:
valid_modes = ("train", "train_extra", "val")
msg = ("Unknown value '{}' for argument split if mode is '{}'. "
"Valid values are {{{}}}.")
msg = msg.format(split, mode, iterable_to_str(valid_modes))
verify_str_arg(split, "split", valid_modes, msg)
if not isinstance(target_type, list):
self.target_type = [target_type]
[
verify_str_arg(value, "target_type",
("instance", "semantic", "polygon", "color"))
for value in self.target_type
]
self.images = self._get_files('image', self.split)
self.targets = self._get_files('label', self.split)
def verify_mode_split(self, split, image_mode):
image_mode = self.verify_mode(image_mode)
if image_mode == "gt_color":
valid_splits = ("testv1", "testv2")
elif image_mode == "gt_labelIds":
valid_splits = ("testv1", "testv2")
elif image_mode == "gt_labelTrainIds":
valid_splits = ("testv1", "testv2")
elif image_mode == "RGB":
valid_splits = ("light", "medium", "testv1", "testv2")
msg = ("Unknown value '{}' for argument split if image_mode is '{}'. "
"Valid values are {{{}}}.")
msg = msg.format(split, image_mode, iterable_to_str(valid_splits))
verify_str_arg(split, "split", valid_splits, msg)
return image_mode
# def verify_dataset(self):
if not os.path.isdir(os.path.join(self.root, self.images_dir)):
if self.image_mode == 'gtFine':
image_dir_zip = os.path.join(self.root, '{}_trainvaltest.zip'.format(self.image_mode))
elif self.image_mode == 'gtCoarse':
image_dir_zip = os.path.join(self.root, '{}.zip'.format(self.image_mode))
else:
if split == 'train_extra':
image_dir_zip = os.path.join(self.root, '{}_trainextra.zip'.format(self.image_mode))
else:
image_dir_zip = os.path.join(self.root, '{}_trainvaltest.zip'.format(self.image_mode))
if os.path.isfile(image_dir_zip):
extract_archive(from_path=image_dir_zip, to_path=self.root)
extract_archive(from_path=target_dir_zip, to_path=self.root)
else:
raise RuntimeError('Dataset not found or incomplete. Please make sure all required folders for the'
' specified "split" and "image_mode" are inside the "root" directory')
def __init__(
self,
data_array,
mask_array,
patch_size,
stride,
split="train",
transforms=None,
max_inlines=None,
n_channels=1,
complete_patches_only=True,
val_ratio=0.1,
test_ratio=0.2,
):
"""Initialise Numpy Dataset
Args:
data_array (numpy.Array): a 3D numpy array that contain the seismic info
mask_array (numpy.Array): a 3D numpy array that contains the labels
patch_size (int): the size of the patch in pixels
stride (int): the stride applied when extracting patches
split (str, optional): what split to load, (train, val, test). Defaults to `train`
transforms (albumentations.augmentations.transforms, optional): albumentation transforms to apply to patches. Defaults to None
exclude_files (list[str], optional): list of files to exclude. Defaults to None
max_inlines (int, optional): maximum number of inlines to load. Defaults to None
n_channels (int, optional): number of channels that the output should contain. Defaults to 3
complete_patches_only (bool, optional): whether to load incomplete patches that are padded to patch_size. Defaults to True
val_ratio (float): ratio to use for validation. Defaults to 0.1
test_ratio (float): ratio to use for test. Defaults to 0.2
"""
super(InlinePatchDataset, self).__init__()
self._data_array = data_array
self._slice_mask_array = mask_array
self._split = split
self._max_inlines = max_inlines
self._n_channels = n_channels
self._complete_patches_only = complete_patches_only
self._patch_size = patch_size
self._stride = stride
self._image_array = []
self._mask_array = []
self._ids = []
self._patch_locations = []
self.transforms = transforms
valid_modes = ("train", "test", "val")
msg = "Unknown value '{}' for argument split. " "Valid values are {{{}}}."
msg = msg.format(split, iterable_to_str(valid_modes))
verify_str_arg(split, "split", valid_modes, msg)
# Set the patch and stride for the patch extractor
_extract_patches_from = _extract_patches(patch_size, stride, self._complete_patches_only)
num_partitions = 5
indexes = self._data_array.shape[0]
num_elements = math.ceil(indexes / num_partitions)
train_indexes_list = []
test_indexes_list = []
val_indexes_list = []
for partition in partition_all(num_elements, range(indexes)): # Partition files into N partitions
train_indexes, val_indexes, test_indexes = _split_train_val_test(partition, val_ratio, test_ratio)
train_indexes_list.extend(train_indexes)
test_indexes_list.extend(test_indexes)
val_indexes_list.extend(val_indexes)
if split == "train":
indexes = train_indexes_list
elif split == "val":
indexes = val_indexes_list
elif split == "test":
indexes = test_indexes_list
# Extract patches
for index in indexes:
img_array = self._data_array[index]
mask_array = self._slice_mask_array[index]
self._ids.append(index)
image_generator, mask_generator, patch_locations = _extract_patches_from(img_array, mask_array)
self._patch_locations.extend(patch_locations)
self._image_array.extend(image_generator)
self._mask_array.extend(mask_generator)
assert len(self._image_array) == len(self._patch_locations), "The shape is not the same"
assert len(self._patch_locations) % len(self._ids) == 0, "Something is wrong with the patches"
self._patches_per_image = int(len(self._patch_locations) / len(self._ids))
self._classes, self._class_counts = _get_classes_and_counts(self._mask_array)
def __init__(self,
root,
split='train',
mode='fine',
target_type='semantic',
transform=None,
target_transform=None,
transforms=None,
cache_image_classes=True,
use_train_labels=True,
return_indices=None):
super(CityscapesDataset, self).__init__(root, transforms, transform,
target_transform)
self.mode = 'gtFine' if mode == 'fine' else 'gtCoarse'
self.splits = split if isinstance(split, (list, tuple)) else [split]
self.images_dirs = [
os.path.join(self.root, 'leftImg8bit', split)
for split in self.splits
]
self.targets_dirs = [
os.path.join(self.root, self.mode, split) for split in self.splits
]
self.target_type = target_type
self.images = []
self.targets = []
# Verification
verify_str_arg(mode, "mode", ("fine", "coarse"))
if mode == "fine":
valid_modes = ("train", "test", "val")
else:
valid_modes = ("train", "train_extra", "val")
msg = ("Unknown value '{}' for argument split if mode is '{}'. "
"Valid values are {{{}}}.")
for split in self.splits:
verify_str_arg(
split, "split", valid_modes,
msg.format(split, mode, iterable_to_str(valid_modes)))
if not isinstance(target_type, list):
self.target_type = [target_type]
[
verify_str_arg(value, "target_type",
("instance", "semantic", "polygon", "color"))
for value in self.target_type
]
for i in range(len(self.splits)):
if not os.path.isdir(self.images_dirs[i]) or not os.path.isdir(
self.targets_dirs[i]):
if split == 'train_extra':
image_dir_zip = os.path.join(
self.root, 'leftImg8bit{}'.format('_trainextra.zip'))
else:
image_dir_zip = os.path.join(
self.root, 'leftImg8bit{}'.format('_trainvaltest.zip'))
if self.mode == 'gtFine':
target_dir_zip = os.path.join(
self.root, '{}{}'.format(self.mode,
'_trainvaltest.zip'))
elif self.mode == 'gtCoarse':
target_dir_zip = os.path.join(
self.root, '{}{}'.format(self.mode, '.zip'))
if os.path.isfile(image_dir_zip) and os.path.isfile(
target_dir_zip):
extract_archive(from_path=image_dir_zip, to_path=self.root)
extract_archive(from_path=target_dir_zip,
to_path=self.root)
else:
raise RuntimeError(
'Dataset not found or incomplete. Please make sure all required folders for the'
' specified "split" and "mode" are inside the "root" directory'
)
# Parse image directories
for i in range(len(self.splits)):
for city in os.listdir(self.images_dirs[i]):
img_dir = os.path.join(self.images_dirs[i], city)
target_dir = os.path.join(self.targets_dirs[i], city)
for file_name in os.listdir(img_dir):
target_types = []
for t in self.target_type:
target_name = '{}_{}'.format(
file_name.split('_leftImg8bit')[0],
self._get_target_suffix(self.mode, t))
target_types.append(
os.path.join(target_dir, target_name))
self.images.append(os.path.join(img_dir, file_name))
self.targets.append(target_types)
# Handle the case for using only the train labels
self.classes = [c for c in CityscapesDataset.classes if not c.ignore_in_eval] \
if use_train_labels else CityscapesDataset.classes
self.use_train_labels = use_train_labels
# Add additional necessary arguments
self.weights = np.ones(len(self.images))
# Calculate classes per image
self.image_classes = None
if 'semantic' in self.target_type and 'test' not in self.splits:
cache_file = os.path.join(root,
f'{"_".join(sorted(self.splits))}.npy'
) if cache_image_classes else None
type_index = np.where(
np.array(self.target_type) == 'semantic')[0][0]
masks = [p[type_index] for p in self.targets]
self.image_classes = self.calc_classes_per_image(masks, cache_file)
self.weights = calc_weights_from_image_classes(self.image_classes)
self.return_indices = self.splits[
0] == 'test' if return_indices is None else return_indices
def __init__(
self,
root,
split='train',
mode='fine',
target_type='semantic',
is_transform=True,
img_size=161, # 512
augmentations=None,
):
super(CityscapesLoader, self).__init__()
self.root = root
self.mode = 'gtFine' if mode == 'fine' else 'gtCoarse'
self.images_dir = os.path.join(self.root, 'leftImg8bit', split)
self.targets_dir = os.path.join(self.root, self.mode, split)
self.split = split
self.images = []
self.targets = []
self.augmentations = augmentations
self.is_transform = is_transform
self.img_size = img_size if isinstance(img_size, tuple) else (img_size,
img_size)
self.tf = transforms.Compose([
#ToTensor(): Converts a PIL Image or numpy.ndarray (H x W x C)
# in the range [0, 255] to a torch.FloatTensor of
# shape (C x H x W) in the range [0.0, 1.0] if the
# PIL Image belongs to one of the modes (L, LA, P, I,
# F, RGB, YCbCr, RGBA, CMYK, 1) or if the numpy.ndarray
# has dtype = np.uint8
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
])
verify_str_arg(mode, "mode", ("fine", "coarse"))
if mode == "fine":
valid_modes = ("train", "test", "val")
else:
valid_modes = ("train", "train_extra", "val")
msg = ("Unknown value '{}' for argument split if mode is '{}'. "
"Valid values are {{{}}}.")
msg = msg.format(split, mode, iterable_to_str(valid_modes))
verify_str_arg(split, "split", valid_modes, msg)
#if not isinstance(target_type, list):
# self.target_type = [target_type]
#else:
# self.target_type = target_type
#[verify_str_arg(value, "target_type",
# ("instance", "semantic", "polygon", "color")) for value in self.target_type]
self.target_type = target_type
for city in os.listdir(self.images_dir):
img_dir = os.path.join(self.images_dir, city)
target_dir = os.path.join(self.targets_dir, city)
for file_name in os.listdir(img_dir):
target_name = '{}_{}'.format(
file_name.split('_leftImg8bit')[0],
self._get_target_suffix(self.mode, self.target_type))
self.targets.append(os.path.join(target_dir, target_name))
self.images.append(os.path.join(img_dir, file_name))
def __init__(self,
root,
split='train',
mode='fine',
target_type='semantic',
transform=None,
target_transform=None,
transforms=None,
num_samples=None,
return_image_name=False):
super(Cityscapes, self).__init__(root, transforms, transform,
target_transform)
self.mode = 'gtFine' if mode == 'fine' else 'gtCoarse'
self.images_dir = os.path.join(self.root, 'leftImg8bit', split)
self.targets_dir = os.path.join(self.root, self.mode, split)
self.target_type = target_type
self.split = split
self.images = []
self.targets = []
self.id_to_trainid = {
-1: self.ignore_label,
0: self.ignore_label,
1: self.ignore_label,
2: self.ignore_label,
3: self.ignore_label,
4: self.ignore_label,
5: self.ignore_label,
6: self.ignore_label,
7: 0,
8: 1,
9: self.ignore_label,
10: self.ignore_label,
11: 2,
12: 3,
13: 4,
14: self.ignore_label,
15: self.ignore_label,
16: self.ignore_label,
17: 5,
18: self.ignore_label,
19: 6,
20: 7,
21: 8,
22: 9,
23: 10,
24: 11,
25: 12,
26: 13,
27: 14,
28: 15,
29: self.ignore_label,
30: self.ignore_label,
31: 16,
32: 17,
33: 18
}
verify_str_arg(mode, "mode", ("fine", "coarse"))
self.rt_img_name = return_image_name
if mode == "fine":
valid_modes = ("train", "test", "val")
else:
valid_modes = ("train", "train_extra", "val")
msg = ("Unknown value '{}' for argument split if mode is '{}'. "
"Valid values are {{{}}}.")
msg = msg.format(split, mode, iterable_to_str(valid_modes))
verify_str_arg(split, "split", valid_modes, msg)
if not isinstance(target_type, list):
self.target_type = [target_type]
[
verify_str_arg(value, "target_type",
("instance", "semantic", "polygon", "color"))
for value in self.target_type
]
if not os.path.isdir(self.images_dir) or not os.path.isdir(
self.targets_dir):
if split == 'train_extra':
image_dir_zip = os.path.join(
self.root, 'leftImg8bit{}'.format('_trainextra.zip'))
else:
image_dir_zip = os.path.join(
self.root, 'leftImg8bit{}'.format('_trainvaltest.zip'))
if self.mode == 'gtFine':
target_dir_zip = os.path.join(
self.root, '{}{}'.format(self.mode, '_trainvaltest.zip'))
elif self.mode == 'gtCoarse':
target_dir_zip = os.path.join(self.root,
'{}{}'.format(self.mode, '.zip'))
if os.path.isfile(image_dir_zip) and os.path.isfile(
target_dir_zip):
extract_archive(from_path=image_dir_zip, to_path=self.root)
extract_archive(from_path=target_dir_zip, to_path=self.root)
else:
raise RuntimeError(
'Dataset not found or incomplete. Please make sure all required folders for the'
' specified "split" and "mode" are inside the "root" directory'
)
for city in os.listdir(self.images_dir):
img_dir = os.path.join(self.images_dir, city)
target_dir = os.path.join(self.targets_dir, city)
for file_name in os.listdir(img_dir):
target_types = []
for t in self.target_type:
target_name = '{}_{}'.format(
file_name.split('_leftImg8bit')[0],
self._get_target_suffix(self.mode, t))
target_types.append(os.path.join(target_dir, target_name))
self.images.append(os.path.join(img_dir, file_name))
self.targets.append(target_types)
# Limit numbers of Dataset
if num_samples is not None and num_samples < len(self.images):
idx = np.random.choice(len(self.images), num_samples)
self.images = [self.images[i] for i in idx]
self.targets = [self.targets[i] for i in idx]
def __init__(
self,
root,
patch_size,
stride,
split="train",
transforms=None,
exclude_files=None,
max_inlines=None,
n_channels=1,
complete_patches_only=True,
):
"""Initialise Penobscot Dataset
Args:
root (str): root directory to load data from
patch_size (int): the size of the patch in pixels
stride (int): the stride applied when extracting patches
split (str, optional): what split to load, (train, val, test). Defaults to `train`
transforms (albumentations.augmentations.transforms, optional): albumentation transforms to apply to patches. Defaults to None
exclude_files (list[str], optional): list of files to exclude. Defaults to None
max_inlines (int, optional): maximum number of inlines to load. Defaults to None
n_channels (int, optional): number of channels that the output should contain. Defaults to 3
complete_patches_only (bool, optional): whether to load incomplete patches that are padded to patch_size. Defaults to True
"""
super(PenobscotInlinePatchDataset,
self).__init__(root, transforms=transforms)
self._image_dir = os.path.join(self.root, "inlines", split)
self._mask_dir = os.path.join(self.root, "masks")
self._split = split
self._exclude_files = exclude_files
self._max_inlines = max_inlines
self._n_channels = n_channels
self._complete_patches_only = complete_patches_only
self._patch_size = patch_size
self._stride = stride
self._image_array = []
self._mask_array = []
self._file_ids = []
self._patch_locations = []
valid_modes = ("train", "test", "val")
msg = "Unknown value '{}' for argument split. " "Valid values are {{{}}}."
msg = msg.format(split, iterable_to_str(valid_modes))
verify_str_arg(split, "split", valid_modes, msg)
if not os.path.exists(self._image_dir):
raise DataNotSplitException(
f"Directory {self._image_dir} does not exist. The dataset has not been \
appropriately split into train, val and test.")
# Get the number of inlines that make up dataset
images_iter, self._max_inlines = pipe(
os.path.join(self._image_dir, "*.tiff"),
glob.iglob,
_filter_files(self._exclude_files),
_limit_inlines(self._max_inlines),
)
# Set the patch and stride for the patch extractor
_extract_patches_from = _extract_patches(patch_size, stride,
self._complete_patches_only)
# Extract patches
for image_path, mask_path in _generate_images_and_masks(
images_iter, self._mask_dir):
img_array = self._open_image(image_path)
mask_array = self._open_mask(mask_path)
self._file_ids.append(_extract_filename(image_path))
image_generator, mask_generator, patch_locations = _extract_patches_from(
img_array, mask_array)
self._patch_locations.extend(patch_locations)
self._image_array.extend(image_generator)
self._mask_array.extend(mask_generator)
assert len(self._image_array) == len(
self._patch_locations), "The shape is not the same"
assert len(self._patch_locations) % len(
self._file_ids) == 0, "Something is wrong with the patches"
self._patches_per_image = int(
len(self._patch_locations) / len(self._file_ids))
# Combine classes 2 and 3
self._mask_array = _combine_classes(self._mask_array)
self._classes, self._class_counts = _get_classes_and_counts(
self._mask_array)
def __init__(self,
root,
split='train',
mode='fine',
target_types=['depth', 'semantic'],
transforms=None):
""" Args:
root (string): Root directory of dataset where directory ``leftImg8bit``
and ``gtFine`` or ``gtCoarse`` are located.
split (string, optional): The image split to use, ``train``, ``test`` or ``val`` if mode="gtFine"
otherwise ``train``, ``train_extra`` or ``val``
mode (string, optional): The quality mode to use, ``gtFine`` or ``gtCoarse``
target_type (string or list, optional): Type of target to use, ``instance``, ``semantic``, ``polygon``, ``depth``
or ``color``. Can also be a list to output a tuple with all specified target types.
transforms (callable, optional): A function/transform that takes input sample and its target as entry
and returns a transformed version.
sequence (tuple, optional): Tuple of (begin, end) denoting the range of the sequence. For example with (-3, 3) the output sequence will include timerange [-3, -2, -1, 0, 1, 2]
"""
super().__init__()
if isinstance(root, torch._six.string_classes):
root = os.path.expanduser(root)
self.root = root
self.transforms = transforms
assert mode in ['fine', 'coarse']
self.mode = 'gtFine' if mode == 'fine' else 'gtCoarse'
self.images_dir = os.path.join(self.root, 'leftImg8bit', split)
self.targets_dir = os.path.join(self.root, self.mode, split)
self.disparity_dir = os.path.join(self.root, 'disparity', split)
self.camera_dir = os.path.join(self.root, 'camera', split)
self.target_types = target_types
self.split = split
self.images = []
self.targets = []
verify_str_arg(mode, "mode", ("fine", "coarse"))
if mode == "fine":
valid_modes = ("train", "test", "val", "train_val")
else:
valid_modes = ("train", "train_extra", "val")
msg = ("Unknown value '{}' for argument split with mode '{}'"
"Valid values are {{{}}}.")
msg = msg.format(split, mode, iterable_to_str(valid_modes))
verify_str_arg(split, "split", valid_modes, msg)
# [verify_str_arg(value, "target_types", ('semantic', 'depth', 'camera'))
# for value in self.target_types]
if 'depth' in self.target_types and 'camera' not in self.target_types: # calculating depth requires camera intrinsics
self.target_types.append('camera')
if not os.path.isdir(self.images_dir) or not os.path.isdir(
self.targets_dir):
raise RuntimeError(
'Dataset not found or incomplete. Please make sure all required folders for the'
' specified "split" and "mode" are inside the "root" directory'
)
for city in os.listdir(self.images_dir):
img_dir = os.path.join(self.images_dir, city)
target_dir = os.path.join(self.targets_dir, city)
disparity_dir = os.path.join(self.disparity_dir, city)
camera_dir = os.path.join(self.camera_dir, city)
for file_name in os.listdir(img_dir):
curr_target = {}
for t in self.target_types:
if t == "depth":
disparity_name = '{}_{}'.format(
file_name.split('_leftImg8bit')[0],
'disparity.png')
curr_target[t] = os.path.join(disparity_dir,
disparity_name)
elif t == 'camera':
camera_name = f"{file_name.split('_leftImg8bit')[0]}_camera.json"
curr_target['camera'] = os.path.join(
camera_dir, camera_name)
else:
target_name = '{}_{}'.format(
file_name.split('_leftImg8bit')[0],
self._get_target_suffix(self.mode, t))
curr_target[t] = os.path.join(target_dir, target_name)
self.images.append(os.path.join(img_dir, file_name))
self.targets.append(curr_target)
