def load_files(self, field_name=None, field_id=None, field_ext=None):
"""
Load all files meet the given filters, each one above can be left blank
:param field_name: name of the field
:param field_id: name of the id
:param field_ext: name of the field extension
:return:
"""
if field_name is None:
field_name = self.field_name
if field_id is None:
field_id = self.field_id
field_ext = ersa_utils.str2list(field_ext, d_type=str)
files = []
for fe in field_ext:
if fe in self.rgb_ext:
select_file_ext = self.file_ext[self.rgb_ext.index(fe)]
else:
select_file_ext = self.file_ext[-1]
if type(field_id) is not str:
field_id = str(field_id)
file = ersa_utils.rotate_list(
self.get_file_selection(field_name, field_id, fe,
select_file_ext))[0]
files.append(file)
files = ersa_utils.rotate_list(files)
if len(files) == 1:
if len(files[0]) == 1:
# only one file been requested
return files[0][0]
return files
def process(self, **kwargs):
"""
Extract the patches
:param kwargs:
file_list: list of lists of the files, can be generated by using collectionMaker.load_files()
file_exts: extensions of the new files
:return:
"""
assert len(kwargs['file_exts']) == len(kwargs['file_list'][0])
grid_list = None
if self.tile_size is not None:
grid_list = make_grid(self.tile_size + 2*self.pad, self.patch_size, self.overlap)
pbar = tqdm(kwargs['file_list'])
record_file = open(os.path.join(self.path, 'file_list.txt'), 'w')
for files in pbar:
pbar.set_description('Extracting {}'.format(os.path.basename(files[0])))
patch_list = []
for f, ext in zip(files, kwargs['file_exts']):
patch_list_ext = []
img = ersa_utils.load_file(f)
if self.tile_size is None:
grid_list = make_grid(np.array(img.shape[:2])+2*self.pad, self.patch_size, self.overlap)
# extract images
for patch, y, x in patch_block(img, self.pad, grid_list, self.patch_size, return_coord=True):
patch_name = '{}_y{}x{}.{}'.format(os.path.basename(f).split('.')[0], int(y), int(x), ext)
patch_name = os.path.join(self.path, patch_name)
ersa_utils.save_file(patch_name, patch.astype(np.uint8))
patch_list_ext.append(patch_name)
patch_list.append(patch_list_ext)
patch_list = ersa_utils.rotate_list(patch_list)
for items in patch_list:
record_file.write('{}\n'.format(' '.join(items)))
record_file.close()
def extract_patches(img_pair,
patch_size,
pad,
overlap,
patch_dir,
prefix,
file_exts=('jpg', 'png'),
file_suffix=('rgb', 'gt')):
grid_list = patchExtractor.make_grid(tile_size + 2 * pad, patch_size,
overlap)
record_file = open(os.path.join(patch_dir, 'file_list.txt'), 'a+')
patch_list = []
for suffix_cnt, (img, ext) in enumerate(zip(img_pair, file_exts)):
patch_list_ext = []
# extract images
for patch, y, x in patchExtractor.patch_block(img,
pad,
grid_list,
patch_size,
return_coord=True):
patch_name = '{}_y{}x{}.{}'.format(
prefix + '_{}'.format(file_suffix[suffix_cnt]), int(y), int(x),
ext)
patch_name = os.path.join(patch_dir, patch_name)
ersa_utils.save_file(patch_name, patch)
patch_list_ext.append(patch_name)
patch_list.append(patch_list_ext)
patch_list = ersa_utils.rotate_list(patch_list)
for items in patch_list:
record_file.write('{}\n'.format(' '.join(items)))
record_file.close()
def get_file_selection(self, field_name, field_id, field_ext, file_ext):
"""
Get list of lists of files selected by given field names, field ids, field extensions and file extensions
:param field_name: name of the fields (e.g., city names)
:param field_id: id of the fields (e.g., tile numbers)
:param field_ext: extension of the fields (e.g., RGB)
:param file_ext: file extension (e.g., tif)
:return: list of lists, where each row is file names of same place with different files
"""
field_name = ersa_utils.str2list(field_name, d_type=str)
field_id = ersa_utils.str2list(field_id, d_type=str)
field_ext = ersa_utils.str2list(field_ext, d_type=str)
file_ext = ersa_utils.str2list(file_ext, d_type=str)
if len(file_ext) == 1:
file_ext = [file_ext[0] for _ in range(len(field_ext))]
file_selection = []
for field, file in zip(field_ext, file_ext):
regexp = self.make_regexp(field_name, field_id, field, file)
file_selection.append(self.get_files(regexp, full_path=True))
file_selection = ersa_utils.rotate_list(file_selection)
return file_selection
def make_collection(self):
"""
Make meta data of the collection, including tile dimension, ground truth and rgb files list
means of all channels in rgb files
:return:
"""
# collect files selection
gt_regexp = self.make_regexp(self.field_name, self.field_id,
self.gt_ext, self.file_ext[-1])
gt_files = self.get_files(gt_regexp, full_path=True)
# rgb data can have multiple channels and be stored in multiple files
rgb_files = []
for cnt, ext in enumerate(self.rgb_ext):
rgb_regexp = self.make_regexp(self.field_name, self.field_id, ext,
self.file_ext[cnt])
rgb_files.append(self.get_files(rgb_regexp, full_path=True))
rgb_files = ersa_utils.rotate_list(rgb_files)
# make meta_data
tile_dim = ersa_utils.load_file(rgb_files[0][0]).shape[:2]
channel_mean = get_channel_mean(self.raw_data_path, rgb_files)
meta_data = {
'raw_data_path': self.raw_data_path,
'field_name': self.field_name,
'field_id': self.field_id,
'rgb_ext': self.rgb_ext,
'gt_ext': self.gt_ext,
'file_ext': self.file_ext,
'clc_name': self.clc_name,
'tile_dim': tile_dim,
'gt_files': gt_files,
'rgb_files': rgb_files,
'chan_mean': channel_mean,
'files': self.files
}
ersa_utils.save_file(os.path.join(self.clc_dir, 'meta.pkl'), meta_data)
print(files)
code_list = ['02', '6', '46', '6', '0', '8']
save_dir = r'/media/ei-edl01/data/uab_datasets/aemo_comb/data/Original_Tiles'
for i in range(6):
test_file = files[i]
print('processing {}'.format(test_file[0]))
rgb = ersa_utils.load_file(test_file[0])
gt = ersa_utils.load_file(test_file[1])
rgb_new, gt_new = makeup_aemo_img(rgb, gt, code_list[i])
save_name = os.path.join(save_dir, os.path.basename(test_file[0]))
ersa_utils.save_file(save_name, rgb_new)
save_name = os.path.join(save_dir, os.path.basename(test_file[1]))
ersa_utils.save_file(save_name, (gt_new/255).astype(np.uint8))'''
data_dir = r'/media/ei-edl01/data/uab_datasets/aemo_comb/data/Original_Tiles'
rgb_files = sorted(glob(os.path.join(data_dir, '*rgb.tif')))
gt_files = sorted(glob(os.path.join(data_dir, '*comb.tif')))
files = [rgb_files, gt_files]
files = ersa_utils.rotate_list(files)
from visualize import visualize_utils
for i in range(6):
test_file = files[i]
print('processing {}'.format(test_file[0]))
rgb = ersa_utils.load_file(test_file[0])
gt = ersa_utils.load_file(test_file[1])
print(np.unique(gt))
visualize_utils.compare_two_figure(rgb, gt)