Training works for SegFormer using MSD dataset works

.gitignore

@@ -150,3 +150,5 @@ logs/*.log
 /bck/
 /data/ADE/images/validation/
 /pretrained/ImageNet-1K/mit_b5.pth
+/data/MSD/
+/pretrained/ImageNet-1K/mit_b0.pth

SegFormer/local_configs/_base_/datasets/msd.py

@@ -0,0 +1,54 @@
+# dataset settings
+dataset_type = 'MSDDataset'
+data_root = 'data/MSD/Task09_Spleen_RGB_2D_512'
+img_norm_cfg = dict(
+    mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True)
+crop_size = (512, 512)
+train_pipeline = [
+    dict(type='LoadImageFromFile'),
+    dict(type='LoadAnnotations'),
+    dict(type='Resize', img_scale=(512, 512), ratio_range=(0.5, 2.0)),
+    dict(type='RandomCrop', crop_size=crop_size, cat_max_ratio=0.75),
+    dict(type='RandomFlip', prob=0.5),
+    dict(type='PhotoMetricDistortion'),
+    dict(type='Normalize', **img_norm_cfg),
+    dict(type='Pad', size=crop_size, pad_val=0, seg_pad_val=255),
+    dict(type='DefaultFormatBundle'),
+    dict(type='Collect', keys=['img', 'gt_semantic_seg']),
+]
+test_pipeline = [
+    dict(type='LoadImageFromFile'),
+    dict(
+        type='MultiScaleFlipAug',
+        img_scale=(512, 512),
+        # img_ratios=[0.5, 0.75, 1.0, 1.25, 1.5, 1.75],
+        flip=False,
+        transforms=[
+            dict(type='Resize', keep_ratio=True),
+            dict(type='RandomFlip'),
+            dict(type='Normalize', **img_norm_cfg),
+            dict(type='ImageToTensor', keys=['img']),
+            dict(type='Collect', keys=['img']),
+        ])
+]
+data = dict(
+    samples_per_gpu=2,
+    workers_per_gpu=2,
+    train=dict(
+        type=dataset_type,
+        data_root=data_root,
+        img_dir='images/training',
+        ann_dir='annotations/training',
+        pipeline=train_pipeline),
+    val=dict(
+        type=dataset_type,
+        data_root=data_root,
+        img_dir='images/validation',
+        ann_dir='annotations/validation',
+        pipeline=test_pipeline),
+    test=dict(
+        type=dataset_type,
+        data_root=data_root,
+        img_dir='images/validation',
+        ann_dir='annotations/validation',
+        pipeline=test_pipeline))

SegFormer/local_configs/_base_/schedules/schedule_20k.py

@@ -4,6 +4,6 @@
 # learning policy
 lr_config = dict(policy='poly', power=0.9, min_lr=1e-4, by_epoch=False)
 # runtime settings
-runner = dict(type='IterBasedRunner', max_iters=20000)
+runner = dict(type='IterBasedRunner', max_iters=2000)
 checkpoint_config = dict(by_epoch=False, interval=2000)
 evaluation = dict(interval=2000, metric='mIoU')

SegFormer/local_configs/segformer/MSD/segformer.b1.512x512.MSD.40k.py → SegFormer/local_configs/segformer/MSD/segformer.512x512.msd.40k.py

@@ -1,22 +1,22 @@
 _base_ = [
     '../../_base_/models/segformer.py',
-    '../../_base_/datasets/ade20k_repeat.py',
+    '../../_base_/datasets/msd.py',
     '../../_base_/default_runtime.py',
-    '../../_base_/schedules/schedule_40k_adamw.py'
+    '../../_base_/schedules/schedule_20k.py'
 ]
 
 # model settings
 norm_cfg = dict(type='SyncBN', requires_grad=True)
 find_unused_parameters = True
 model = dict(
     type='EncoderDecoder',
-    pretrained='../../pretrained/ImageNet-1K/mit_b5.pth',
+    pretrained='../../pretrained/ImageNet-1K/mit_b0.pth',
     backbone=dict(
-        type='mit_b5',
+        type='mit_b0',
         style='pytorch'),
     decode_head=dict(
         type='SegFormerHead',
-        in_channels=[64, 128, 320, 512],
+        in_channels=[32, 64, 160, 256],
         in_index=[0, 1, 2, 3],
         feature_strides=[4, 8, 16, 32],
         channels=128,
@@ -43,6 +43,5 @@
                  warmup_ratio=1e-6,
                  power=1.0, min_lr=0.0, by_epoch=False)
 
-
 data = dict(samples_per_gpu=2)
 evaluation = dict(interval=16000, metric='mIoU')

SegFormer/mmseg/datasets/__init__.py

@@ -1,4 +1,5 @@
 from .ade import ADE20KDataset
+from .msd import MSDDataset
 from .builder import DATASETS, PIPELINES, build_dataloader, build_dataset
 from .chase_db1 import ChaseDB1Dataset
 from .cityscapes import CityscapesDataset
@@ -16,5 +17,6 @@
     'CustomDataset', 'build_dataloader', 'ConcatDataset', 'RepeatDataset',
     'DATASETS', 'build_dataset', 'PIPELINES', 'CityscapesDataset',
     'PascalVOCDataset', 'ADE20KDataset', 'PascalContextDataset',
-    'ChaseDB1Dataset', 'DRIVEDataset', 'HRFDataset', 'STAREDataset', 'MapillaryDataset', 'CocoStuff'
+    'ChaseDB1Dataset', 'DRIVEDataset', 'HRFDataset', 'STAREDataset', 'MapillaryDataset', 'CocoStuff',
+    'MSDDataset'
 ]

SegFormer/mmseg/datasets/builder.py

@@ -60,7 +60,8 @@ def _concat_dataset(cfg, default_args=None):
 
 def build_dataset(cfg, default_args=None):
     """Build datasets."""
-    cfg.data_root = '../../data/ADEChallengeData2016'
+    cfg.data_root = '../../data/MSD/Task09_Spleen_RGB_2D_512'
+    print("groun")
 
     from .dataset_wrappers import ConcatDataset, RepeatDataset
     if isinstance(cfg, (list, tuple)):

SegFormer/mmseg/datasets/msd.py

@@ -0,0 +1,24 @@
+from .builder import DATASETS
+from .custom import CustomDataset
+
+
+@DATASETS.register_module()
+class MSDDataset(CustomDataset):
+    """ADE20K dataset.
+
+    In segmentation map annotation for ADE20K, 0 stands for background, which
+    is not included in 150 categories. ``reduce_zero_label`` is fixed to True.
+    The ``img_suffix`` is fixed to '.jpg' and ``seg_map_suffix`` is fixed to
+    '.png'.
+    """
+    CLASSES = (
+        "no Ailment", "silment")
+
+    PALETTE = [[120, 120, 120], [92, 0, 255]]
+
+    def __init__(self, **kwargs):
+        super(MSDDataset, self).__init__(
+            img_suffix='.png',
+            seg_map_suffix='.png',
+            reduce_zero_label=True,
+            **kwargs)

requirements.txt

@@ -14,4 +14,5 @@ scikit-image
 attr
 ipython
 timm==0.3.2
-opencv-python==4.5.1.48
+opencv-python==4.5.1.48
+imageio