Define blood vessel dataset and fix filename bug (#203)

* fix filename bug

* define blood vessel dataset

* redo debug

* fix small bug

* rename dataset

Co-authored-by: yamengxi <yamengxi@sensetime.com>

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'
 ]

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)

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)

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)

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)

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):