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