MM Segmentation 训练自己的数据集

在/mmsegmentation-master/mmseg/datasets路径下新建一个mydataset.py文件

# Copyright (c) OpenMMLab. All rights reserved.

from .builder import DATASETS
from .custom import CustomDataset

@DATASETS.register_module()
class MyDataset(CustomDataset):
    """My Dataset.

    """
    CLASSES = ('background', 'alpha')  # 需要修改的地方
    PALETTE = [[0, 0, 0], [255, 255, 255]]  # 需要修改的地方

    def __init__(self, **kwargs):
        super(MyDatasetDataset, self).__init__(
            img_suffix='.png',  # 需要修改的地方
            seg_map_suffix='.png', #需要修改的地方
            reduce_zero_label=False,
            **kwargs)
        assert self.file_client.exists(self.img_dir)

同时在/mmsegmentation-master/mmseg/datasets中的____init____.py添加刚刚新建的数据集

# Copyright (c) OpenMMLab. All rights reserved.
from .ade import ADE20KDataset
from .builder import DATASETS, PIPELINES, build_dataloader, build_dataset
from .chase_db1 import ChaseDB1Dataset
from .cityscapes import CityscapesDataset
from .coco_stuff import COCOStuffDataset
from .custom import CustomDataset
from .dark_zurich import DarkZurichDataset
from .dataset_wrappers import (ConcatDataset, MultiImageMixDataset,
                               RepeatDataset)
from .drive import DRIVEDataset
from .hrf import HRFDataset
from .isaid import iSAIDDataset
from .isprs import ISPRSDataset
from .loveda import LoveDADataset
from .night_driving import NightDrivingDataset