pytorch 版本deeplabv3+ 训练自己的语义分割数据集

一、数据集制备

用于海岸线项目的语义分割数据集制备与VOC格式相同（关于VOC格式说明可参加链接1）。首先需要在dataset目录下建立如下图所示的自己的数据集目录，比如我的是海岸线分割数据集，那么在dataset下就建立起sealand文件夹。然后在sealand文件夹下建立起ImageSets、JPEGImages、SegmentationClass三个文件。其中JPEGImages，存放所有的原始图片;SegmentationClass存放所有的标注文件mask图片;ImageSets/Segmentation/下存放train.txt test.txt  val.txt 。算法会自动根据train.txt的路径信息自动去从标注文件和原始文件中读取训练集测试集验证集等。数据的分堆方式可以参考如下代码（maketxt.py）。

注意：label的标签一定要是单通道灰度图，且里面的类别标签不能是255（因为有些卷积在给特征图补额外像素的时候会将对应的值补255）。比如对海岸线分割的二分类任务，海的区域是1，陆地为0。云分割任务，薄云、厚云为1，2，陆地以及其它背景区域则为0。

from sklearn.model_selection import train_test_split
import os
“”“”
train_test_split的用法
train_size：训练集大小

　　float：0-1之间，表示训练集所占的比例

　　int：直接指定训练集的数量

　　None：自动为测试集的补集，也就是原始数据集减去测试集

test_size：测试集大小，默认值是0.25

　　float：0-1之间，表示测试集所占的比例

　　int：直接指定测试集的数量

　　None：自动为训练集的补集，也就是原始数据集减去训练集

random_state：可以理解为随机数种子，主要是为了复现结果而设置

shuffle：表示是否打乱数据位置，True或者False，默认是True

stratify：表示是否按照样本比例（不同类别的比例）来划分数据集，例如原始数据集 类A:类B = 75%:25%，那么划分的测试集和训练集中的A:B的比例都会是75%:25%；可用于样本类别差异很大的情况，一般使用为：stratify=y，即用数据集的标签y来进行划分。
“”“

imagedir = '/home/jd/projects/pytorch-deeplab-xception/dataloaders/datasets/wind/JPEGImages/'

outdir = '/home/jd/projects/pytorch-deeplab-xception/dataloaders/datasets/wind/ImageSets/'

images = []
for file in os.listdir(imagedir):
        filename = file.split('.')[0]
        images.append(filename)
#训练集测试集验证集比例为：4：2：2
train, test = train_test_split(images, train_size=0.5, random_state=0)
val, test = train_test_split(test, train_size=0.5, random_state=0)

with open(outdir+"train.txt", 'w') as f:
	f.write('\n'.join(train))
with open(outdir+"val.txt", 'w') as f:
	f.write('\n'.join(val))
with open(outdir+"test.txt", 'w') as f:
	f.write('\n'.join(test))

                                                        

                                                               图1.自己数据集的制备目录

二、修改参数文件

1.mypath.py 文件的修改

    需要在mypath.py文件中加入自己的数据集此处我加入的是sealand

   

                                                                                 图2.mypath.py文件修改

2.sealand.py 文件的添加

      在dataloaders/datasets路径下创建自己的数据集文件，在海岸线分割任务中创建的是sealand.py。因为我们的数据集格式与VOC一致，所以直接将pascal.py内容复制过来修改，这里将类别数和数据集改成自己的数据集，类别记得加上背景类（比如有20类目标，加上背景类，NUM_CLASSES=21）。我这里一共是2类。

3.修改dataloaders/utils.py文件

修改utils.py文件中的decode_segmap函数，添加自己的数据集和类别信息；同时增加def get_sealand_labels()函数，用于给mask图着色

def decode_segmap(label_mask, dataset, plot=False):
    """Decode segmentation class labels into a color image
       用于把分类的mask图转化成彩色图像
    """
    if dataset == 'pascal' or dataset == 'coco':
        n_classes = 21
        label_colours = get_pascal_labels()
    elif dataset == 'cityscapes':
        n_classes = 19
        label_colours = get_cityscapes_labels()
    elif dataset == 'sealand':    #####海岸线分割
        n_classes = 2
        label_colours = get_sealand_labels()
    elif dataset == 'wind':
        n_classes = 2
        label_colours = get_wind_labels()
    else:
        raise NotImplementedError

#根据网络输出的mask图，给不同的类别染不同的颜色，[0,0,0]为黑色背景
def get_sealand_labels():
    return np.asarray([[0,0,0],[128,0,0]])

4.修改dataloaders/__init__.py文件

在__init__.py 中添加sealand数据集的描述部分，如下所示

    elif args.dataset == 'sealand':
        train_set = sealand.VOCSegmentation(args, split='train')
        val_set = sealand.VOCSegmentation(args, split='val')
        """
        if args.use_sbd:
            sbd_train = sbd.SBDSegmentation(args, split=['train', 'val'])
            train_set = combine_dbs.CombineDBs([train_set, sbd_train], excluded=[val_set])
        """ 
        num_class = train_set.NUM_CLASSES
        train_loader = DataLoader(train_set, batch_size=args.batch_size, shuffle=True, **kwargs)
        val_loader = DataLoader(val_set, batch_size=args.batch_size, shuffle=False, **kwargs)
        test_loader = None

        return train_loader, val_loader, test_loader, num_class

三、网络训练

python train.py --backbone mobilenet  --dataset sealand --lr 0.007 --workers 1 --epochs 50 --batch-size 8 --gpu-ids 0 --checkname deeplab-mobilenet 

四、模型测试

   输出的二值图在output_new文件夹里

python demo.py --in-path test/ --ckpt run/sealand/deeplab-mobilenet/model_best.pth.tar --backbone mobilenet

链接

【1】https://blog.csdn.net/qq_41997920/article/details/104914129 VOC数据集详解