diff --git a/mmseg/datasets/__init__.py b/mmseg/datasets/__init__.py index dd4705c3e4..4f248dc16b 100644 --- a/mmseg/datasets/__init__.py +++ b/mmseg/datasets/__init__.py @@ -1,13 +1,18 @@ from .ade import ADE20KDataset from .builder import DATASETS, PIPELINES, build_dataloader, build_dataset +from .chase_db1 import ChaseDB1Dataset from .cityscapes import CityscapesDataset from .custom import CustomDataset from .dataset_wrappers import ConcatDataset, RepeatDataset +from .drive import DRIVEDataset +from .hrf import HRFDataset from .pascal_context import PascalContextDataset +from .stare import STAREDataset from .voc import PascalVOCDataset __all__ = [ 'CustomDataset', 'build_dataloader', 'ConcatDataset', 'RepeatDataset', 'DATASETS', 'build_dataset', 'PIPELINES', 'CityscapesDataset', - 'PascalVOCDataset', 'ADE20KDataset', 'PascalContextDataset' + 'PascalVOCDataset', 'ADE20KDataset', 'PascalContextDataset', + 'ChaseDB1Dataset', 'DRIVEDataset', 'HRFDataset', 'STAREDataset' ] diff --git a/mmseg/datasets/chase_db1.py b/mmseg/datasets/chase_db1.py new file mode 100644 index 0000000000..79d544f202 --- /dev/null +++ b/mmseg/datasets/chase_db1.py @@ -0,0 +1,27 @@ +import os.path as osp + +from .builder import DATASETS +from .custom import CustomDataset + + +@DATASETS.register_module() +class ChaseDB1Dataset(CustomDataset): + """Chase_db1 dataset. + + In segmentation map annotation for Chase_db1, 0 stands for background, + which is included in 2 categories. ``reduce_zero_label`` is fixed to False. + The ``img_suffix`` is fixed to '.jpg' and ``seg_map_suffix`` is fixed to + '_1stHO.jpg'. + """ + + CLASSES = ('background', 'vessel') + + PALETTE = [[120, 120, 120], [6, 230, 230]] + + def __init__(self, **kwargs): + super(ChaseDB1Dataset, self).__init__( + img_suffix='.jpg', + seg_map_suffix='_1stHO.jpg', + reduce_zero_label=False, + **kwargs) + assert osp.exists(self.img_dir) diff --git a/mmseg/datasets/drive.py b/mmseg/datasets/drive.py new file mode 100644 index 0000000000..177ca5f691 --- /dev/null +++ b/mmseg/datasets/drive.py @@ -0,0 +1,27 @@ +import os.path as osp + +from .builder import DATASETS +from .custom import CustomDataset + + +@DATASETS.register_module() +class DRIVEDataset(CustomDataset): + """DRIVE dataset. + + In segmentation map annotation for DRIVE, 0 stands for background, which is + included in 2 categories. ``reduce_zero_label`` is fixed to False. The + ``img_suffix`` is fixed to '.jpg' and ``seg_map_suffix`` is fixed to + '_manual1.jpg'. + """ + + CLASSES = ('background', 'vessel') + + PALETTE = [[120, 120, 120], [6, 230, 230]] + + def __init__(self, **kwargs): + super(DRIVEDataset, self).__init__( + img_suffix='.jpg', + seg_map_suffix='_manual1.jpg', + reduce_zero_label=False, + **kwargs) + assert osp.exists(self.img_dir) diff --git a/mmseg/datasets/hrf.py b/mmseg/datasets/hrf.py new file mode 100644 index 0000000000..ff24417831 --- /dev/null +++ b/mmseg/datasets/hrf.py @@ -0,0 +1,27 @@ +import os.path as osp + +from .builder import DATASETS +from .custom import CustomDataset + + +@DATASETS.register_module() +class HRFDataset(CustomDataset): + """HRF dataset. + + In segmentation map annotation for HRF, 0 stands for background, which is + included in 2 categories. ``reduce_zero_label`` is fixed to False. The + ``img_suffix`` is fixed to '.jpg' and ``seg_map_suffix`` is fixed to + '.jpg'. + """ + + CLASSES = ('background', 'vessel') + + PALETTE = [[120, 120, 120], [6, 230, 230]] + + def __init__(self, **kwargs): + super(HRFDataset, self).__init__( + img_suffix='.jpg', + seg_map_suffix='.jpg', + reduce_zero_label=False, + **kwargs) + assert osp.exists(self.img_dir) diff --git a/mmseg/datasets/stare.py b/mmseg/datasets/stare.py new file mode 100644 index 0000000000..97e987a39c --- /dev/null +++ b/mmseg/datasets/stare.py @@ -0,0 +1,27 @@ +import os.path as osp + +from .builder import DATASETS +from .custom import CustomDataset + + +@DATASETS.register_module() +class STAREDataset(CustomDataset): + """STARE dataset. + + In segmentation map annotation for STARE, 0 stands for background, which is + included in 2 categories. ``reduce_zero_label`` is fixed to False. The + ``img_suffix`` is fixed to '.jpg' and ``seg_map_suffix`` is fixed to + '.ah.jpg'. + """ + + CLASSES = ('background', 'vessel') + + PALETTE = [[120, 120, 120], [6, 230, 230]] + + def __init__(self, **kwargs): + super(STAREDataset, self).__init__( + img_suffix='.jpg', + seg_map_suffix='.ah.jpg', + reduce_zero_label=False, + **kwargs) + assert osp.exists(self.img_dir) diff --git a/tools/convert_datasets/drive.py b/tools/convert_datasets/drive.py index 9da30fb1b4..c25a1d2059 100644 --- a/tools/convert_datasets/drive.py +++ b/tools/convert_datasets/drive.py @@ -50,8 +50,10 @@ def main(): img = mmcv.imread(osp.join(now_dir, img_name)) mmcv.imwrite( img, - osp.join(out_dir, 'images', 'training', - osp.splitext(img_name)[0] + '.jpg')) + osp.join( + out_dir, 'images', 'training', + osp.splitext(img_name)[0].replace('_training', '') + + '.jpg')) now_dir = osp.join(tmp_dir, 'training', '1st_manual') for img_name in os.listdir(now_dir): @@ -72,8 +74,9 @@ def main(): img = mmcv.imread(osp.join(now_dir, img_name)) mmcv.imwrite( img, - osp.join(out_dir, 'images', 'validation', - osp.splitext(img_name)[0] + '.jpg')) + osp.join( + out_dir, 'images', 'validation', + osp.splitext(img_name)[0].replace('_test', '') + '.jpg')) now_dir = osp.join(tmp_dir, 'test', '1st_manual') if osp.exists(now_dir):