From 746c8d378591fe1f4611adef4feb766300989278 Mon Sep 17 00:00:00 2001 From: David de la Iglesia Castro Date: Thu, 24 Sep 2020 18:48:16 +0200 Subject: [PATCH] Use img_prefix and seg_prefix for loading (#153) * Use img_prefix and seg_prefix for loading * flake8 * Fix split --- mmseg/datasets/custom.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/mmseg/datasets/custom.py b/mmseg/datasets/custom.py index 91d7b0b5eb..f055faee28 100644 --- a/mmseg/datasets/custom.py +++ b/mmseg/datasets/custom.py @@ -131,19 +131,16 @@ def load_annotations(self, img_dir, img_suffix, ann_dir, seg_map_suffix, with open(split) as f: for line in f: img_name = line.strip() - img_file = osp.join(img_dir, img_name + img_suffix) - img_info = dict(filename=img_file) + img_info = dict(filename=img_name + img_suffix) if ann_dir is not None: - seg_map = osp.join(ann_dir, img_name + seg_map_suffix) + seg_map = img_name + seg_map_suffix img_info['ann'] = dict(seg_map=seg_map) img_infos.append(img_info) else: for img in mmcv.scandir(img_dir, img_suffix, recursive=True): - img_file = osp.join(img_dir, img) - img_info = dict(filename=img_file) + img_info = dict(filename=img) if ann_dir is not None: - seg_map = osp.join(ann_dir, - img.replace(img_suffix, seg_map_suffix)) + seg_map = img.replace(img_suffix, seg_map_suffix) img_info['ann'] = dict(seg_map=seg_map) img_infos.append(img_info) @@ -165,6 +162,8 @@ def get_ann_info(self, idx): def pre_pipeline(self, results): """Prepare results dict for pipeline.""" results['seg_fields'] = [] + results['img_prefix'] = self.img_dir + results['seg_prefix'] = self.ann_dir if self.custom_classes: results['label_map'] = self.label_map @@ -225,8 +224,9 @@ def get_gt_seg_maps(self): """Get ground truth segmentation maps for evaluation.""" gt_seg_maps = [] for img_info in self.img_infos: + seg_map = osp.join(self.ann_dir, img_info['ann']['seg_map']) gt_seg_map = mmcv.imread( - img_info['ann']['seg_map'], flag='unchanged', backend='pillow') + seg_map, flag='unchanged', backend='pillow') # modify if custom classes if self.label_map is not None: for old_id, new_id in self.label_map.items():