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) ])