【记录】mmsegmentation 训练自己的数据集

数据集标注工具选择

labeme 标注

感觉不能很好的贴合轮廓

精灵标注助手

利用涂抹工具

photoshop标注

利用魔棒或磁性套索工具，创建选区后，赋予不同的颜色，提高标注效率

标准格式转换

下面介绍两种，但在图像分割领域常用的格式为cityscapes数据集格式和VOC数据集格式
coco数据集格式在mmdetection中的实例分割中用得到

转化为COCO格式

label2coco.py

 python labelme2coco.py --input_dir images --output_dir coco --labels labels.txt

其中labels.txt中每一行为一个类别名称,如下所示：

__ignore__
_background_
xxx

# 命令行执行： python labelme2coco.py --input_dir images --output_dir coco --labels labels.txt
# 输出文件夹必须为空文件夹
 ###  python labelme2coco.py --input_dir images --output_dir coco --labels label.txt

import argparse
import collections
import datetime
import glob
import json
import os
import os.path as osp
import sys
import uuid
import imgviz
import numpy as np
import labelme
from sklearn.model_selection import train_test_split
 
try:
    import pycocotools.mask
except ImportError:
    print("Please install pycocotools:\n\n    pip install pycocotools\n")
    sys.exit(1)
 
 
def to_coco(args,label_files,train):
 
    # 创建 总标签data 
    now = datetime.datetime.now()
    data = dict(
        info=dict(
            description=None,
            url=None,
            version=None,
            year=now.year,
            contributor=None,
            date_created=now.strftime("%Y-%m-%d %H:%M:%S.%f"),
        ),
        licenses=[dict(url=None, id=0, name=None,)],
        images=[
            # license, url, file_name, height, width, date_captured, id
        ],
        type="instances",
        annotations=[
            # segmentation, area, iscrowd, image_id, bbox, category_id, id
        ],
        categories=[
            # supercategory, id, name
        ],
    )
 
    # 创建一个 {类名 : id} 的字典，并保存到 总标签data 字典中。
    class_name_to_id = {
    
    }
    for i, line in enumerate(open(args.labels).readlines()):
        class_id = i - 1  # starts with -1
        class_name = line.strip()   # strip() 方法用于移除字符串头尾指定的字符(默认为空格或换行符)或字符序列。
        if class_id == -1:
            assert class_name == "__ignore__"   # background:0, class1:1, ,,
            continue
        class_name_to_id[class_name] = class_id
        data["categories"].append(
            dict(supercategory=None, id=class_id, name=class_name,)
        )
 
 
    if train:
        out_ann_file = osp.join(args.output_dir, "annotations","instances_train2017.json")
    else:
        out_ann_file = osp.join(args.output_dir, "annotations","instances_val2017.json")
 
 
    for image_id, filename in enumerate(label_files):
 
        label_file = labelme.LabelFile(filename=filename)
        base = osp.splitext(osp.basename(filename))[0]      # 文件名不带后缀
        if train:
            out_img_file = osp.join(args.output_dir, "train2017", base + ".jpg")
        else:
            out_img_file = osp.join(args.output_dir, "val2017", base + ".jpg")
        
        print("| ",out_img_file)
 
        # ************************** 对图片的处理开始 *******************************************
        # 将标签文件对应的图片进行保存到对应的 文件夹。train保存到 train2017/ test保存到 val2017/
        img = labelme.utils.img_data_to_arr(label_file.imageData)   # .json文件中包含图像，用函数提出来
        imgviz.io.imsave(out_img_file, img)     # 将图像保存到输出路径
 
        # ************************** 对图片的处理结束 *******************************************
 
        # ************************** 对标签的处理开始 *******************************************
        data["images"].append(
            dict(
                license=0,
                url=None,
                file_name=osp.relpath(out_img_file, osp.dirname(out_ann_file)),
                #   out_img_file = "/coco/train2017/1.jpg"
                #   out_ann_file = "/coco/annotations/annotations_train2017.json"
                #   osp.dirname(out_ann_file) = "/coco/annotations"
                #   file_name = ..\train2017\1.jpg   out_ann_file文件所在目录下 找 out_img_file 的相对路径
                height=img.shape[0],
                width=img.shape[1],
                date_captured=None,
                id=image_id,
            )
        )
 
        masks = {
    
    }  # for area
        segmentations = collections.defaultdict(list)  # for segmentation
        for shape in label_file.shapes:
            points = shape["points"]
            label = shape["label"]
            group_id = shape.get("group_id")
            shape_type = shape.get("shape_type", "polygon")
            mask = labelme.utils.shape_to_mask(
                img.shape[:2], points, shape_type
            )
 
            if group_id is None:
                group_id = uuid.uuid1()
 
            instance = (label, group_id)
 
            if instance in masks:
                masks[instance] = masks[instance] | mask
            else:
                masks[instance] = mask
 
            if shape_type == "rectangle":
                (x1, y1), (x2, y2) = points
                x1, x2 = sorted([x1, x2])
                y1, y2 = sorted([y1, y2])
                points = [x1, y1, x2, y1, x2, y2, x1, y2]
            else:
                points = np.asarray(points).flatten().tolist()
 
            segmentations[instance].append(points)
        segmentations = dict(segmentations)
 
        for instance, mask in masks.items():
            cls_name, group_id = instance
            if cls_name not in class_name_to_id:
                continue
            cls_id = class_name_to_id[cls_name]
 
            mask = np.asfortranarray(mask.astype(np.uint8))
            mask = pycocotools.mask.encode(mask)
            area = float(pycocotools.mask.area(mask))
            bbox = pycocotools.mask.toBbox(mask).flatten().tolist()
 
            data["annotations"].append(
                dict(
                    id=len(data["annotations"]),
                    image_id=image_id,
                    category_id=cls_id,
                    segmentation=segmentations[instance],
                    area=area,
                    bbox=bbox,
                    iscrowd=0,
                )
            )
        # ************************** 对标签的处理结束 *******************************************
 
        # ************************** 可视化的处理开始 *******************************************
        if not args.noviz:
            labels, captions, masks = zip(
                *[
                    (class_name_to_id[cnm], cnm, msk)
                    for (cnm, gid), msk in masks.items()
                    if cnm in class_name_to_id
                ]
            )
            viz = imgviz.instances2rgb(
                image=img,
                labels=labels,
                masks=masks,
                captions=captions,
                font_size=15,
                line_width=2,
            )
            out_viz_file = osp.join(
                args.output_dir, "visualization", base + ".jpg"
            )
            imgviz.io.imsave(out_viz_file, viz)
        # ************************** 可视化的处理结束 *******************************************
 
    with open(out_ann_file, "w") as f:  # 将每个标签文件汇总成data后，保存总标签data文件
        json.dump(data, f)
 
 
# 主程序执行
def main():
    parser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    parser.add_argument("--input_dir", help="input annotated directory")
    parser.add_argument("--output_dir", help="output dataset directory")
    parser.add_argument("--labels", help="labels file", required=True)
    parser.add_argument("--noviz", help="no visualization", action="store_true")
    args = parser.parse_args()
 
    if osp.exists(args.output_dir):
        print("Output directory already exists:", args.output_dir)
        sys.exit(1)
    os.makedirs(args.output_dir)
    print("| Creating dataset dir:", args.output_dir)
    if not args.noviz:
        os.makedirs(osp.join(args.output_dir, "visualization"))
 
    # 创建保存的文件夹
    if not os.path.exists(osp.join(args.output_dir, "annotations")):
        os.makedirs(osp.join(args.output_dir, "annotations"))
    if not os.path.exists(osp.join(args.output_dir, "train2017")):
        os.makedirs(osp.join(args.output_dir, "train2017"))
    if not os.path.exists(osp.join(args.output_dir, "val2017")):
        os.makedirs(osp.join(args.output_dir, "val2017"))
 
    # 获取目录下所有的.jpg文件列表
    feature_files = glob.glob(osp.join(args.input_dir, "*.jpg"))
    print('| Image number: ', len(feature_files))
 
    # 获取目录下所有的joson文件列表
    label_files = glob.glob(osp.join(args.input_dir, "*.json"))
    print('| Json number: ', len(label_files))
    
 
    # feature_files:待划分的样本特征集合    label_files:待划分的样本标签集合    test_size:测试集所占比例 
    # x_train:划分出的训练集特征      x_test:划分出的测试集特征     y_train:划分出的训练集标签    y_test:划分出的测试集标签
    x_train, x_test, y_train, y_test = train_test_split(feature_files, label_files, test_size=0.2)
    print("| Train number:", len(y_train), '\t Value number:', len(y_test))
 
    # 把训练集标签转化为COCO的格式，并将标签对应的图片保存到目录 /train2017/
    print("—"*50) 
    print("| Train images:")
    to_coco(args,y_train,train=True)
    
    # 把测试集标签转化为COCO的格式，并将标签对应的图片保存到目录 /val2017/ 
    print("—"*50)
    print("| Test images:")
    to_coco(args,y_test,train=False)
    
 
if __name__ == "__main__":
    print("—"*50)
    main()
    print("—"*50)

转化为VOC格式

分为两种来源：labelme的json文件和ps等工具直接生成的mask图像

json文件

#!/usr/bin/env python

from __future__ import print_function

import argparse
import glob
import json
import os
import os.path as osp
import sys

import numpy as np
import PIL.Image

import labelme


def main():
    parser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    parser.add_argument('input_dir', help='input annotated directory')
    parser.add_argument('output_dir', help='output dataset directory')
    parser.add_argument('--labels', help='labels file', required=True)
    args = parser.parse_args()

    if osp.exists(args.output_dir):
        print('Output directory already exists:', args.output_dir)
        sys.exit(1)
    os.makedirs(args.output_dir)
    os.makedirs(osp.join(args.output_dir, 'JPEGImages'))
    os.makedirs(osp.join(args.output_dir, 'SegmentationClass'))
    os.makedirs(osp.join(args.output_dir, 'SegmentationClassPNG'))
    os.makedirs(osp.join(args.output_dir, 'SegmentationClassVisualization'))
    print('Creating dataset:', args.output_dir)

    class_names = []
    class_name_to_id = {
    
    }
    for i, line in enumerate(open(args.labels).readlines()):
        class_id = i - 1  # starts with -1
        class_name = line.strip()
        class_name_to_id[class_name] = class_id
        if class_id == -1:
            assert class_name == '__ignore__'
            continue
        elif class_id == 0:
            assert class_name == '_background_'
        class_names.append(class_name)
    class_names = tuple(class_names)
    print('class_names:', class_names)
    out_class_names_file = osp.join(args.output_dir, 'class_names.txt')
    with open(out_class_names_file, 'w') as f:
        f.writelines('\n'.join(class_names))
    print('Saved class_names:', out_class_names_file)

    colormap = labelme.utils.label_colormap(255)

    for label_file in glob.glob(osp.join(args.input_dir, '*.json')):
        print('Generating dataset from:', label_file)
        with open(label_file) as f:
            base = osp.splitext(osp.basename(label_file))[0]
            out_img_file = osp.join(
                args.output_dir, 'JPEGImages', base + '.jpg')
            out_lbl_file = osp.join(
                args.output_dir, 'SegmentationClass', base + '.npy')
            out_png_file = osp.join(
                args.output_dir, 'SegmentationClassPNG', base + '.png')
            out_viz_file = osp.join(
                args.output_dir,
                'SegmentationClassVisualization',
                base + '.jpg',
            )

            data = json.load(f)

            img_file = osp.join(osp.dirname(label_file), data['imagePath'])
            img = np.asarray(PIL.Image.open(img_file))
            PIL.Image.fromarray(img).save(out_img_file)

            lbl = labelme.utils.shapes_to_label(
                img_shape=img.shape,
                shapes=data['shapes'],
                label_name_to_value=class_name_to_id,
            )
            labelme.utils.lblsave(out_png_file, lbl)

            np.save(out_lbl_file, lbl)

            viz = labelme.utils.draw_label(
                lbl, img, class_names, colormap=colormap)
            PIL.Image.fromarray(viz).save(out_viz_file)


if __name__ == '__main__':
    main()

mask图像

以二值图像为例

from PIL import Image
import numpy as np
import os



def label_colormap(N=256):

    def bitget(byteval, idx):
        return ((byteval & (1 << idx)) != 0)

    cmap = np.zeros((N, 3))
    for i in range(0, N):
        id = i
        r, g, b = 0, 0, 0
        for j in range(0, 8):
            r = np.bitwise_or(r, (bitget(id, 0) << 7 - j))
            g = np.bitwise_or(g, (bitget(id, 1) << 7 - j))
            b = np.bitwise_or(b, (bitget(id, 2) << 7 - j))
            id = (id >> 3)
        cmap[i, 0] = r
        cmap[i, 1] = g
        cmap[i, 2] = b
    cmap = cmap.astype(np.float32) / 255
    return cmap

if __name__ == '__main__':
    work_dir = "./SegmentationClass1" # 图像所处文件夹
    file_names = os.listdir(work_dir)
    for file_name in file_names:
        
        file_path = os.path.join(work_dir,file_name)
        # file_path = "./SegmentationClass1/0001_mask.png"
        image = Image.open(file_path).convert('L')
        img = np.array(image)
        print(np.sum(img))
        img[img==255] = 1
        
        print(np.sum(img))
        # 重新保存
        # image = Image.fromarray(img,'L')
        image = Image.fromarray(img, mode='P')
        colormap = label_colormap(255)
        image.putpalette((colormap * 255).astype(np.uint8).flatten())
        new_name = file_name[:-4]
        new_name = new_name.strip("_mask") # 文件名处理成和图像一样的名字

        image.save(f'{
      
      new_name}.png')        

划分训练集和测试集

运行后会在* ./ImageSets/Segmentation/. *文件夹下生成train.txt test.txt

import os
import random

import numpy as np
from PIL import Image
from tqdm import tqdm

#-------------------------------------------------------#
#   想要增加测试集修改trainval_percent 
#   修改train_percent用于改变验证集的比例 9:1
#   
#   当前该库将测试集当作验证集使用，不单独划分测试集
#-------------------------------------------------------#
trainval_percent    = 0.9
train_percent       = 1
#-------------------------------------------------------#
#   指向VOC数据集所在的文件夹
#   默认指向根目录下的VOC数据集
#-------------------------------------------------------#
VOCdevkit_path      = './'

if __name__ == "__main__":
    random.seed(0)
    print("Generate txt in ImageSets.")
    segfilepath     = os.path.join(VOCdevkit_path, 'SegmentationClass')
    saveBasePath    = os.path.join(VOCdevkit_path, 'ImageSets/Segmentation')
    
    temp_seg = os.listdir(segfilepath)
    total_seg = []
    for seg in temp_seg:
        if seg.endswith(".png"):
            total_seg.append(seg)

    num     = len(total_seg)  
    list    = range(num)  
    tv      = int(num*trainval_percent)  
    tr      = int(tv*train_percent)  
    trainval= random.sample(list,tv)  
    train   = random.sample(trainval,tr)  
    
    print("train and val size",tv)
    print("traub suze",tr)
    ftrainval   = open(os.path.join(saveBasePath,'trainval.txt'), 'w')  
    ftest       = open(os.path.join(saveBasePath,'test.txt'), 'w')  
    ftrain      = open(os.path.join(saveBasePath,'train.txt'), 'w')  
    fval        = open(os.path.join(saveBasePath,'val.txt'), 'w')  
    
    for i in list:  
        name = total_seg[i][:-4]+'\n'  
        if i in trainval:  
            ftrainval.write(name)  
            if i in train:  
                ftrain.write(name)  
            else:  
                fval.write(name)  
        else:  
            ftest.write(name)  
    
    ftrainval.close()  
    ftrain.close()  
    fval.close()  
    ftest.close()
    print("Generate txt in ImageSets done.")

修改配置文件

假设使用数据集格式为VOC格式，网络为deeplabv3+,直接在命令行里输入

python tools/train.py configs/deeplabv3/deeplabv3-r50-d8512x51220k_voc12aug.py

会work_dirs文件夹下生成一个对应的文件夹，把里面的deeplabv3-r50-d8512x51220k_voc12aug.py文件复制出来修改

修改data_root为自己的路径

搜索num_classes,改为类别数+1

修改voc.py

在mmsegmentation/mmseg/datasets路径下，找到PascalVOCDataset类，将CLASSES 改为自己的类别名称，PALETTE改为对应的数量

修改class_names.py

在mmsegmentation/mmseg/core/evaluation下，找到voc_classes函数，return返回值改为自己的类别名称。

开始训练

python tools/train.py deeplabv3_r50-d8.py

推理

将测试结果保存到文件夹中

python tools/test.py deeplabv3plus.py  work_dirs/deeplabv3plus_r50-d8_512x512_20k_voc12aug/iter_20000.pth  --show-dir deeplabv3plus  --eval mIoU

检查数据集

python tools/browse_dataset.py deeplabv3_unet.py --show-origin --output-dir ./outputs --opacity 0.2

-show-origin 显示原始图，不加则会显示增强图
–ouput-dir 图像保存的路径
–opacity 不透明度

利用训练好的模型进行推理

import mmcv
import argparse
import cv2
import numpy as np
from mmseg.apis import inference_segmentor, init_segmentor
from pathlib import Path


def parse_args():
    parser = argparse.ArgumentParser(
        description='mmseg inference a model')
    parser.add_argument('-cfg','--config-file',type=str, default="../unet.py", help='test config file path')
    parser.add_argument('--checkpoint-file', type=str, default='../work_dirs/latest.pth',help='checkpoint file')
    parser.add_argument(
        '--image-dir', default="../test",
        help=('save dir'))
    parser.add_argument(
        '--output-dir', default="../results",help='save dirs')
    parser.add_argument('--show-origin',action='store_true', help='show origal imgae')
    args = parser.parse_args()
    return args
if __name__ == '__main__':
    args = parse_args()

  
    model = init_segmentor(args.config_file, args.checkpoint_file, device='cuda:0')
    image_dirs = Path(args.image_dir)
    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
   
    for img in image_dirs.rglob('*.png'):

        image_name = img.stem
       
        result = inference_segmentor(model, img)
        if hasattr(model, 'module'):
            model = model.module
        img = model.show_result(
            img, result, palette=None, show=False, opacity=0.2)
       
        if args.show_origin:
            origin_image = cv2.imread(str(img))

            img = np.hstack((origin_image, img))
        out_file = output_dir / f'{
      
      image_name}.png'
        cv2.imwrite(str(out_file),img)
  

如果显示原图与分割后的图

python inference.py --config-file ../unet.py checkpoint-file latest.pt  --image-dir ./test ----output-dir ./result --show_origin