def prepare_data(self, data_dict, want_reverse=False):
input_keys, target_keys = self.input_keys(), self.target_keys()
if self.conditions.use_ground_truth:
input_keys += target_keys
Log.info_once('Input keys: {}'.format(input_keys))
Log.info_once('Target keys: {}'.format(target_keys))
inputs = [data_dict[k] for k in input_keys]
batch_size = len(inputs[0])
targets = [data_dict[k] for k in target_keys]
sequences = [
self._prepare_sequence(inputs, force_list=True),
self._prepare_sequence(targets, force_list=False)
]
if want_reverse:
rev_data_dict = self._reverse_data_dict(data_dict)
sequences.extend([
self._prepare_sequence([rev_data_dict[k] for k in input_keys],
force_list=True),
self._prepare_sequence([rev_data_dict[k] for k in target_keys],
force_list=False)
])
return sequences, batch_size
def __list_dirs(self, root_dir, dataset):
img_list = list()
label_list = list()
offset_h_list = list()
offset_w_list = list()
name_list = list()
image_dir = os.path.join(root_dir, dataset, 'image')
label_dir = os.path.join(root_dir, dataset, 'label')
offset_h_dir = None
offset_w_dir = None
subdir = os.environ.get('offset_dir')
if subdir is not None:
Log.info_once('Using offset dir: {}'.format(subdir))
offset_h_dir = os.path.join(root_dir, dataset, subdir, 'h')
offset_w_dir = os.path.join(root_dir, dataset, subdir, 'w')
else:
offset_type = self.configer.get('data', 'offset_type')
assert (offset_type is not None)
offset_h_dir = os.path.join(root_dir, dataset, offset_type, 'h')
offset_w_dir = os.path.join(root_dir, dataset, offset_type, 'w')
img_extension = os.listdir(image_dir)[0].split('.')[-1]
for file_name in os.listdir(label_dir):
image_name = '.'.join(file_name.split('.')[:-1])
img_path = os.path.join(image_dir,
'{}.{}'.format(image_name, img_extension))
label_path = os.path.join(label_dir, file_name)
offset_h_path = os.path.join(offset_h_dir,
self._replace_ext(file_name, 'mat'))
offset_w_path = os.path.join(offset_w_dir,
self._replace_ext(file_name, 'mat'))
if not os.path.exists(label_path) or not os.path.exists(img_path):
Log.error('Label Path: {} not exists.'.format(label_path))
continue
img_list.append(img_path)
label_list.append(label_path)
offset_h_list.append(offset_h_path)
offset_w_list.append(offset_w_path)
name_list.append(image_name)
if dataset == 'train' and self.configer.get('data', 'include_val'):
image_dir = os.path.join(root_dir, 'val/image')
label_dir = os.path.join(root_dir, 'val/label')
subdir = os.environ.get('offset_dir')
if subdir is not None:
Log.info_once('Using offset dir: {}'.format(subdir))
offset_h_dir = os.path.join(root_dir, 'val', subdir, 'h')
offset_w_dir = os.path.join(root_dir, 'val', subdir, 'w')
else:
offset_type = self.configer.get('data', 'offset_type')
assert (offset_type is not None)
offset_h_dir = os.path.join(root_dir, 'val', offset_type, 'h')
offset_w_dir = os.path.join(root_dir, 'val', offset_type, 'w')
for file_name in os.listdir(label_dir):
image_name = '.'.join(file_name.split('.')[:-1])
img_path = os.path.join(
image_dir, '{}.{}'.format(image_name, img_extension))
label_path = os.path.join(label_dir, file_name)
offset_h_path = os.path.join(
offset_h_dir, self._replace_ext(file_name, 'mat'))
offset_w_path = os.path.join(
offset_w_dir, self._replace_ext(file_name, 'mat'))
if not os.path.exists(label_path) or not os.path.exists(
img_path):
Log.error('Label Path: {} not exists.'.format(label_path))
continue
img_list.append(img_path)
label_list.append(label_path)
offset_h_list.append(offset_h_path)
offset_w_list.append(offset_w_path)
name_list.append(image_name)
return img_list, label_list, offset_h_list, offset_w_list, name_list
def __list_dirs(self, root_dir, dataset):
if os.environ.get('use_cityscapes_style'):
if 'GTA5_small' in root_dir:
root_dir = root_dir.replace('GTA5_small', 'GTA5_Cityscapes')
else:
root_dir = root_dir.replace('GTA5', 'GTA5_Cityscapes')
Log.info_once(
'Using Cityscapes style, switch to {}'.format(root_dir))
else:
Log.info_once('Using default root dir: {}'.format(root_dir))
img_list = list()
label_list = list()
offset_list = list()
name_list = list()
image_subdir = os.environ.get('image_subdir', 'image')
label_subdir = os.environ.get('label_dir', 'label')
Log.info_once('Using label dir: {}'.format(label_subdir))
offset_subdir = os.environ.get('offset_dir', 'dt_offset')
Log.info_once(
'Using distance transform based offset: {}'.format(offset_subdir))
image_dir = os.path.join(root_dir, dataset, image_subdir)
label_dir = os.path.join(root_dir, dataset, label_subdir)
offset_dir = os.path.join(root_dir, dataset, offset_subdir)
img_extension = os.listdir(image_dir)[0].split('.')[-1]
file_list_txt = os.environ.get('use_file_list')
if file_list_txt is None:
Log.info_once('Using file list: all')
files = sorted(os.listdir(label_dir))
else:
Log.info_once('Using file list: {}'.format(file_list_txt))
with open(
os.path.join(root_dir, dataset, 'file_list',
file_list_txt)) as f:
files = [x.strip() for x in f]
if os.environ.get('chunk'):
n, i = map(int, os.environ.get('chunk').split('_'))
step = len(files) // n + 4
files = files[step * i:step * (i + 1)]
for file_name in files:
image_name = '.'.join(file_name.split('.')[:-1])
img_path = os.path.join(image_dir,
'{}.{}'.format(image_name, img_extension))
label_path = os.path.join(label_dir, file_name)
offset_path = os.path.join(offset_dir,
self._replace_ext(file_name, 'mat'))
if not os.path.exists(label_path) or not os.path.exists(img_path):
Log.error('Label Path: {} not exists.'.format(label_path))
continue
img_list.append(img_path)
label_list.append(label_path)
offset_list.append(offset_path)
name_list.append(image_name)
if dataset == 'train' and self.configer.get('data', 'include_val'):
Log.info_once('Include val set for training ...')
image_dir = os.path.join(root_dir, 'val', image_subdir)
label_dir = os.path.join(root_dir, 'val', label_subdir)
offset_dir = os.path.join(root_dir, 'val', offset_subdir)
if file_list_txt is None:
files = sorted(os.listdir(label_dir))
else:
with open(
os.path.join(root_dir, 'val', 'file_list',
file_list_txt)) as f:
files = [x.strip() for x in f]
for file_name in files:
image_name = '.'.join(file_name.split('.')[:-1])
img_path = os.path.join(
image_dir, '{}.{}'.format(image_name, img_extension))
label_path = os.path.join(label_dir, file_name)
offset_path = os.path.join(offset_dir,
self._replace_ext(file_name, 'mat'))
if not os.path.exists(label_path) or not os.path.exists(
img_path):
Log.error('Label Path: {} not exists.'.format(label_path))
continue
img_list.append(img_path)
label_list.append(label_path)
offset_list.append(offset_path)
name_list.append(image_name)
return img_list, label_list, offset_list, name_list
