【Pytorch】pytorch版本deeplabv3p训练语义分割数据集

开发环境

  • Pytorch
  • Python
  • GPU / CPU
  • 代码框架: pytorch-deeplab-xception

一、数据集准备

  1. 制作语义分割数据集可用labelme

    安装 >> pip install labelme 启动 >> labelme
    注意: 通过labelme制作的语义分割数据的json文件中mask是连串的坐标点,需要使用该坐标转换成mask.png
    若语义分割的类别总共有3类,mask.png的像素组成应该:
    第一类的像素值为(0, 0, 0), 第二类的像素值为(1, 1, 1), 第三类的像素值为(2, 2, 2),以此类推

  2. 代码框架的数据集分布

    1. 可以现在根目录创建data文件夹用于存放数据集,例如需要分割表计任务,创建后文件夹分布为 data/meter

    2. data/meter目录下创建三个文件夹:

      • ImageSets: 文件夹内有一个文件夹Segmentation,Segmentation里面存放的是训练集、验证集、测试集的信息如: train.txt、val.txt
      • JPEGImages: 存放原图文件如: *.jpg
      • SegmentationClass: 存放原图对应的mask图片文件如: *.png

      最后需要创建的文件夹有:
      data/meter/ImageSets
      data/meter/ImageSets/Segmentation
      data/meter/JPEGImages
      data/meter/SegmentationClass

二、修改源码

请仔细一步一步跟着做,如果后面训练的时候报错了,请返回重新检查。

  1. 在根目录mypath.py文件,定义自己的数据集叫meter, 该名字会在最后train.py文件中调用命令行训练中会用到。代码中更改:添加一个判断分支为数据集名字如meter, 然后返回数据集的绝对路径,而这个绝对路径就是在前面项目的根目录创建的data文件夹中meter数据集。
    【Pytorch】pytorch版本deeplabv3p训练语义分割数据集_第1张图片
  2. 然后在dataloaders/datasets路径下创建自己的数据集文件meter.py,我是按照VOC数据集格式,所以可以直接将pascal.py内容复制过来修改。
    【Pytorch】pytorch版本deeplabv3p训练语义分割数据集_第2张图片
    如类别数,表盘分割任务只需要指针、刻度和背景,所以class为3
  3. 修改dataloaders/utils.py,在文件中创建一个get_meter_labels()函数,根据自己数据集的类别数和mask使用的颜色进行填写,如表盘分割任务有3个类别,则填写3个颜色通道的数据,用于预测显示mask的颜色:
    【Pytorch】pytorch版本deeplabv3p训练语义分割数据集_第3张图片
  4. 修改dataloaders/utils.pydecode_segmap函数,添加判断分支和加入类别数:
    【Pytorch】pytorch版本deeplabv3p训练语义分割数据集_第4张图片
  5. 修改dataloaders/init.py文件
    【Pytorch】pytorch版本deeplabv3p训练语义分割数据集_第5张图片
    【Pytorch】pytorch版本deeplabv3p训练语义分割数据集_第6张图片
  6. 修改train.py,在大概187行添加上自己的数据集名字
    【Pytorch】pytorch版本deeplabv3p训练语义分割数据集_第7张图片

三、训练模型

作者提供了四种的backbone,分别是:resnet、xception、drn、mobilenet,并且提供了预训练权重,在训练开始时会自动下载,可以根据自己电脑的配置进行选择,如果显存不足可以使用Mobilenet或Resnet作为backbone

python train.py --dataset meter --backbone resnet --lr 0.007 --workers 1 --epochs 50 --batch-size 32 --gpu-ids 0 --checkname meter-deeplab-resnet 

更多的参数传参在train.py里面查看,如果是多GPU的话可以指定GPU信息,具体可以参考程序里面的train_sh文件
【Pytorch】pytorch版本deeplabv3p训练语义分割数据集_第8张图片
【Pytorch】pytorch版本deeplabv3p训练语义分割数据集_第9张图片

四、模型测试

新建detect.py并写入

import argparse
import os
import numpy as np
import time

from modeling.deeplab import *
from dataloaders import custom_transforms as tr
from PIL import Image
from torchvision import transforms
from dataloaders.utils import  *
from torchvision.utils import make_grid, save_image

def main():

    parser = argparse.ArgumentParser(description="PyTorch DeeplabV3Plus Training")
    parser.add_argument('--in-path', type=str, required=True, help='image to test')
    # parser.add_argument('--out-path', type=str, required=True, help='mask image to save')
    parser.add_argument('--backbone', type=str, default='resnet',
                        choices=['resnet', 'xception', 'drn', 'mobilenet'],
                        help='backbone name (default: resnet)')
    parser.add_argument('--ckpt', type=str, default='deeplab-resnet.pth',
                        help='saved model')
    parser.add_argument('--out-stride', type=int, default=16,
                        help='network output stride (default: 8)')
    parser.add_argument('--no-cuda', action='store_true', default=False,
                        help='disables CUDA training')
    parser.add_argument('--gpu-ids', type=str, default='0',
                        help='use which gpu to train, must be a \
                        comma-separated list of integers only (default=0)')
    parser.add_argument('--dataset', type=str, default='pascal',
                        choices=['pascal', 'coco', 'cityscapes','invoice'],
                        help='dataset name (default: pascal)')
    parser.add_argument('--crop-size', type=int, default=513,
                        help='crop image size')
    parser.add_argument('--num_classes', type=int, default=2,
                        help='crop image size')
    parser.add_argument('--sync-bn', type=bool, default=None,
                        help='whether to use sync bn (default: auto)')
    parser.add_argument('--freeze-bn', type=bool, default=False,
                        help='whether to freeze bn parameters (default: False)')

    args = parser.parse_args()
    args.cuda = not args.no_cuda and torch.cuda.is_available()
    if args.cuda:
        try:
            args.gpu_ids = [int(s) for s in args.gpu_ids.split(',')]
        except ValueError:
            raise ValueError('Argument --gpu_ids must be a comma-separated list of integers only')

    if args.sync_bn is None:
        if args.cuda and len(args.gpu_ids) > 1:
            args.sync_bn = True
        else:
            args.sync_bn = False
    model_s_time = time.time()
    model = DeepLab(num_classes=args.num_classes,
                    backbone=args.backbone,
                    output_stride=args.out_stride,
                    sync_bn=args.sync_bn,
                    freeze_bn=args.freeze_bn)

    ckpt = torch.load(args.ckpt, map_location='cpu')
    model.load_state_dict(ckpt['state_dict'])
    model = model.cuda()
    model_u_time = time.time()
    model_load_time = model_u_time-model_s_time
    print("model load time is {}".format(model_load_time))

    composed_transforms = transforms.Compose([
        tr.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
        tr.ToTensor()])
    for name in os.listdir(args.in_path):
        s_time = time.time()
        image = Image.open(args.in_path+"/"+name).convert('RGB')

        # image = Image.open(args.in_path).convert('RGB')
        target = Image.open(args.in_path+"/"+name).convert('L')
        sample = {'image': image, 'label': target}
        tensor_in = composed_transforms(sample)['image'].unsqueeze(0)

        model.eval()
        if args.cuda:
            tensor_in = tensor_in.cuda()
        with torch.no_grad():
            output = model(tensor_in)

        grid_image = make_grid(decode_seg_map_sequence(torch.max(output[:3], 1)[1].detach().cpu().numpy()),
                                3, normalize=False, range=(0, 255))
        save_image(grid_image,args.in_path+"/"+"{}_mask.png".format(name[0:-4]))
        u_time = time.time()
        img_time = u_time-s_time
        print("image:{} time: {} ".format(name,img_time))
        # save_image(grid_image, args.out_path)
        # print("type(grid) is: ", type(grid_image))
        # print("grid_image.shape is: ", grid_image.shape)
    print("image save in in_path.")
if __name__ == "__main__":
   main()

# python demo.py --in-path your_file --out-path your_dst_file

python detect.py --in-path /home/dataset/...... --ckpt run/meter/deeplab-mobilenet/model_best.pth.tar --backbone resnet --dataset meter --num_classes 3
  • –in-path: 测试图片路径目录
  • –ckpt: 模型路径,如刚训练完生成的 run/invoice/deeplab-mobilnet/model_best.pth.tar
  • num_classes: 分割对象的类别总数

结果显示,如下图模型已经成功完成分割,可能分割的性能可能还有待改善(但模型已经跑通了):
【Pytorch】pytorch版本deeplabv3p训练语义分割数据集_第10张图片
【Pytorch】pytorch版本deeplabv3p训练语义分割数据集_第11张图片

你可能感兴趣的:(深度学习)