def __init__(self, datasets_root, image_root, mask_root, data_aug_prob, mean, std, seed):
"""Initialize file paths or a list of file names. """
super(Dataset, self).__init__(seed)
self.datasets_root = Path(datasets_root)
self.image_root = self.datasets_root / image_root
self.mask_root = self.datasets_root / mask_root
# Preparing images and labels filename lists.
image_list = util.get_fname_list(self.image_root)
mask_list = util.get_fname_list(self.mask_root, suffix='*.png')
image_names = pd.DataFrame({'image_fname': image_list})
mask_names = pd.DataFrame({'mask_fname': mask_list})
image_names['match_substr'] = [Path(f).stem for f in image_list]
mask_names['match_substr'] = [Path(f).stem for f in mask_list]
self.match_df = image_names.merge(mask_names, on='match_substr', how='inner')
# mean & std
self.mean = eval(mean)
self.std = eval(std)
self.normalize = transforms.Normalize(self.mean, self.std)
self.data_aug_prob = data_aug_prob
self.transform = None
if self.data_aug_prob > 0:
self.transform = transforms.Compose([
util.RandomCrop(crop_size),
util.Jitter_HSV(self.data_aug_prob),
util.RandomHorizontalFlip(self.data_aug_prob),
util.RandomVerticleFlip(self.data_aug_prob),
util.RandomRotate90(self.data_aug_prob)
])
def __init__(self, datasets_root, image_root, mask_root, edge_root, mini_root, direct_root, data_aug_prob, mean, std, seed):
super(XGDataset, self).__init__(seed)
self.datasets_root = Path(datasets_root)
self.image_root = self.datasets_root / image_root
self.mask_root = self.datasets_root / mask_root
self.edge_root = self.datasets_root / edge_root
self.mini_root = self.datasets_root / mini_root
self.direct_root = self.datasets_root / direct_root
# Preparing images and labels filename lists.
image_list = util.get_fname_list(self.image_root)
mask_list = util.get_fname_list(self.mask_root, suffix='*.png')
edge_list = util.get_fname_list(self.edge_root, suffix='*.png')
mini_list = util.get_fname_list(self.mini_root, suffix='*.png')
direct_list = util.get_fname_list(self.direct_root, suffix='*.png')
image_names = pd.DataFrame({'image_fname': image_list})
mask_names = pd.DataFrame({'mask_fname': mask_list})
edge_names = pd.DataFrame({'edge_fname': edge_list})
mini_names = pd.DataFrame({'mini_fname': mini_list})
direct_names = pd.DataFrame({'direct_fname': direct_list})
image_names['match_substr'] = [Path(f).stem for f in image_list]
mask_names['match_substr'] = [Path(f).stem for f in mask_list]
edge_names['match_substr'] = [Path(f).stem for f in edge_list]
mini_names['match_substr'] = [Path(f).stem for f in mini_list]
direct_names['match_substr'] = [Path(f).stem for f in direct_list]
self.match_df = image_names.merge(mask_names, on='match_substr', how='inner').merge(
edge_names, on='match_substr', how='inner').merge(
mini_names, on='match_substr', how='inner').merge(
direct_names, on='match_substr', how='inner')
# mean & std
self.mean = eval(mean)
self.std = eval(std)
self.normalize = transforms.Normalize(self.mean, self.std)
self.data_aug_prob = data_aug_prob
self.transform = None
if self.data_aug_prob > 0:
self.transform = transforms.Compose([
util.RandomCrop2(crop_size),
util.RandomHorizontalFlip(self.data_aug_prob),
util.RandomVerticleFlip(self.data_aug_prob),
util.RandomRotate90(self.data_aug_prob)
])