参考代码来源于 https://github.com/dbolya/yolact

实例分割 yolact

最近进入新项目，又开始忙了，赶着项目没什么时间写博客。
最近项目是分割类的，主要训练慢，而且成效也慢，加上小目标检测效果不好，一直忙着在想技巧，分析数据分析问题。利用周末记录一下用到的实例分割模型。

---

更新：（这上面看一些技巧挺有用的，我也用了一些，对模型多少会有些提升。https://www.aiuai.cn/aifarm1370.html

1. 环境配置

首先将项目安装下来，或者git clone到自己本地。

因为我这边跑的是yolact++（在data/config.py内设置），所以按官方要求是要安装DCNv2

即执行下面代码。

cd external/DCNv2
python setup.py build develop

这里因为我一开始用的pytorch1.6版本，一直安装不下，也可能是对应的cuda版本问题。后来把版本降低后就可以正常执行安装了。版本：torch 1.4.0； torchvision 0.5.0 所以这里最好自己用conda创建一个专门用来跑yolact的虚拟环境。其他环境配置可以参照官方说明environment.yml。

2. 实例分割数据集准备

yolact官方代码训练用的coco数据集，所以如果要训练自己的数据集，就要先把数据集转换成coco的格式。我这边是用labelme标注数据集的，所以在网上找了个labelme2coco的转换脚本。

（1）labelme2coco

我这里主要数据标注的时候有多种类别，但是最终用来训练的只有两种类别，所以用label_1和lable_2将不同的类别归为一类。

import argparse
import json
import matplotlib.pyplot as plt
import skimage.io as io
import cv2
from labelme import utils
import numpy as np
import glob
import PIL.Image
import sys, os, time
from shapely.geometry import Polygon
labels_1 = ['a','b']
labels_2 = ['c','d']
class labelme2coco(object):
    def __init__(self,labelme_json=[], path="", save_json_path=r''):
        '''
        :param labelme_json: 所有labelme的json文件路径组成的列表
        :param save_json_path: json保存位置
        '''
        self.labelme_json=labelme_json
        self.save_json_path=save_json_path
        self.info = {
    }
        self.licenses = []
        self.images=[]
        self.categories=[]
        self.annotations=[]
        # self.data_coco = {}
        self.label=[]
        self.annID=1
        self.height=0
        self.width=0
        self.path = path
        self.save_json()

    def data_transfer(self):
        self.info = self.info_dict()
        self.licenses = self.licenses_dict()
        count=0
        for num,json_file in enumerate(self.labelme_json):
            # print(num)

            with open(json_file,'r') as fp:
                print('[{}/{} count {} {}]'.format(num,len(self.labelme_json),count,json_file))
                try:
                    data = json.load(fp)  # 加载json文件
                    self.images.append(self.image(data,count))
                    for shapes in data['shapes']:
                        #label=shapes['label'].split('_')
                        if shapes['label'] in labels_1:
                            label = '1'
                        elif shapes['label'] in labels_2:
                            label = '2'
                        else:
                            print(shapes['label'],'error')
                            continue
                        # label=shapes['label']
                        #print(shapes['label'])
                        if label not in self.label:
                            self.categories.append(self.categorie(label))
                            self.label.append(label)
                        points=shapes['points']
                        # points = np.array(points,'int').tolist()
                        if shapes['shape_type'] == 'rectangle':
                            x0,y0,x1,y1 = points[0][0],points[0][1],points[1][0],points[1][1]
                            points = [[x0, y0], [x1, y0], [x1, y1], [x0, y1]]
                        if shapes['shape_type'] == 'polygon':
                            if len(points) <= 2:
                                continue
                        self.annotations.append(self.annotation(shapes['shape_type'],points,label,count))
                        self.annID+=1
                    count+=1
                except:
                    print(data['imagePath'],'is error')
        print(self.label)

    def info_dict(self):
        info = {
    }
        info['description'] = None
        info['url'] = None
        info['contributor'] = None
        info['day'] = time.strftime('%Y.%m.%d',time.localtime(time.time()))
        info['data_created'] = time.strftime('%Y-%m-%d %H:%M:%S',time.localtime(time.time()))
        return info

    def licenses_dict(self):
        license = []
        license = [{
    'url':None,'id':0,'name':None}]
        return license

    def image(self,data,num):
        image={
    }
        img = cv2.imdecode(np.fromfile(os.path.join(self.path, data['imagePath']), dtype=np.uint8), 0)
        height, width = img.shape[:2]
        image['license'] = 0
        image['url'] = None
        image['file_name'] = data['imagePath'].split('/')[-1]
        image['height'] = height
        image['width'] = width
        image['date_captured'] = None
        image['id'] = num
        self.height=height
        self.width=width
        return image

    def categorie(self,label):
        categorie={
    }
        categorie['supercategory'] = None
        categorie['id']=len(self.label)+1 # 0 默认为背景
        categorie['name'] = label
        return categorie

    def annotation(self,type,points,label,num):
        annotation={
    }
        annotation['id'] = self.annID
        annotation['image_id'] = num
        annotation['category_id'] = self.getcatid(label)
        annotation['segmentation']=[list(map(int, list(np.asarray(points).flatten())))]
        poly = Polygon(points)
        area_ = round(poly.area,6)
        annotation['area'] = area_
        annotation['bbox'] = list(map(float, self.getbbox(points)))
        annotation['iscrowd'] = 0
        # annotation['bbox'] = str(self.getbbox(points)) # 使用list保存json文件时报错（不知道为什么）
        # list(map(int,a[1:-1].split(','))) a=annotation['bbox'] 使用该方式转成list
        return annotation

    def getcatid(self,label):
        for categorie in self.categories:
            if label==categorie['name']:
                return categorie['id']
        return -1

    def getbbox(self,points):
        # img = np.zeros([self.height,self.width],np.uint8)
        # cv2.polylines(img, [np.asarray(points)], True, 1, lineType=cv2.LINE_AA)  # 画边界线
        # cv2.fillPoly(img, [np.asarray(points)], 1)  # 画多边形 内部像素值为1
        polygons = points
        mask = self.polygons_to_mask([self.height,self.width], polygons)
        return self.mask2box(mask)

    def mask2box(self, mask):
        '''从mask反算出其边框
        mask：[h,w]  0、1组成的图片
        1对应对象，只需计算1对应的行列号（左上角行列号，右下角行列号，就可以算出其边框）
        '''
        # np.where(mask==1)
        index = np.argwhere(mask == 1)
        rows = index[:, 0]
        clos = index[:, 1]
        # 解析左上角行列号
        left_top_r = np.min(rows)  # y
        left_top_c = np.min(clos)  # x
        # 解析右下角行列号
        right_bottom_r = np.max(rows)
        right_bottom_c = np.max(clos)
        # return [(left_top_r,left_top_c),(right_bottom_r,right_bottom_c)]
        # return [(left_top_c, left_top_r), (right_bottom_c, right_bottom_r)]
        # return [left_top_c, left_top_r, right_bottom_c, right_bottom_r]  # [x1,y1,x2,y2]
        return [left_top_c, left_top_r, right_bottom_c-left_top_c, right_bottom_r-left_top_r]  # [x1,y1,w,h] 对应COCO的bbox格式

    def polygons_to_mask(self,img_shape, polygons):
        mask = np.zeros(img_shape, dtype=np.uint8)
        mask = PIL.Image.fromarray(mask)
        xy = list(map(tuple, polygons))
        PIL.ImageDraw.Draw(mask).polygon(xy=xy, outline=1, fill=1)
        mask = np.array(mask, dtype=bool)
        return mask

    def data2coco(self):
        data_coco={
    }
        data_coco['info']=self.info
        data_coco['licenses']=self.licenses
        data_coco['images']=self.images
        data_coco['categories']=self.categories
        data_coco['annotations']=self.annotations
        return data_coco

    def save_json(self):
        self.data_transfer()
        self.data_coco = self.data2coco()
        # 保存json文件
        json.dump(self.data_coco, open(self.save_json_path, 'w'), indent=4)  # indent=4 更加美观显示

if __name__ == '__main__':
    path = r''# 存放labelme标注好的数据（图片以及json标签）
    print(os.path.join(path, '*.json'))
    labelme_json=glob.glob(os.path.join(path, '*.json'))
    print("json data numbers: ", len(labelme_json))
    save_path = 'save.json'# coco的格式主要就是将所有数据存放在一个json文件内。
    if not os.path.exists(save_path):
        os.mkdir(save_path)
    labelme2coco((labelme_json), path, save_path)

（2）visual_coco

转换好coco 标签格式后，最好可视化检查一下标签是否正常。
这里可以通过pycocotools写一个可视化脚本。

import os
from pycocotools.coco import COCO

from skimage import io
from matplotlib import pyplot as plt

img_path = r'dataset\train_dataset'#图片路径
path_save = r'hardSamples/mult_lossSORT_biaozhu'# 输出可视化图片
json_file = r'dataset/train.json' #转换好coco格式的json文件
os.makedirs(path_save, exist_ok=True)

coco = COCO(json_file)
print(coco)

catIds = coco.getCatIds()
imgIds = coco.getImgIds()

for i in range(len(imgIds)):
    # print('11')
    img = coco.loadImgs(imgIds[i])[0]
    img_name = img['file_name']
    print('[{}/{}:{}]'.format(i+1,len(imgIds),img_name.replace('/','_')))

    img_fullpath = os.path.join(img_path,img_name)
    I = io.imread(img_fullpath)
    plt.axis('off')
    plt.imshow(I) #绘制图像，显示交给plt.show()处理
    annIds = coco.getAnnIds(imgIds=img['id'], iscrowd=None)
    anns = coco.loadAnns(annIds)
    coco.showAnns(anns)
    # plt.show() #显示图像

    img_save_fullpath = os.path.join(path_save,img_name.replace('/','_'))
    plt.savefig(img_save_fullpath)
    plt.cla()

3. 修改代码配置文件

首先是修改data/config.py ，自己写一个my_custom_dataset如下，里面放入你做好的数据集和标签。对应的yolact_base_config里面改dataset': my_custom_dataset,里面还有其他配置可以改的，
如网络：我这里使用的yolact_plus，输入尺寸用的550。里面还可以设置损失函数之类的，图形增强的一些配置，可以自己去看。因为这里面几个配置文件copy来copy去的，容易把自己绕晕。可以自己根据配置写一个Config。然后修改train.py里面的config名字就行了。

my_custom_dataset = Config({
    
    'name': 'my_custom_dataset',

    'train_images': '/data/train/',
    'train_info': '/data/tran_data.json',
    'valid_images': '/data/val/',
    'valid_info': '/data/val_data.json',

    'has_gt': True,
    'class_names': ('1', '2') # 类别
})

# ----------------------- YOLACT v1.0 CONFIGS ----------------------- #

yolact_base_config = coco_base_config.copy({
    
    'name': 'yolact_base',

    # Dataset stuff
    'dataset': my_custom_dataset,
    'num_classes': len(my_custom_dataset.class_names) + 1,

    # Image Size
    'max_size': 550,

    # Training params
    'lr_steps': (280000, 600000, 700000, 750000),
    'max_iter': 800000,

    # Backbone Settings
    'backbone': resnet101_backbone.copy({
    
        'selected_layers': list(range(1, 4)),
        'use_pixel_scales': True,
        'preapply_sqrt': False,
        'use_square_anchors': True,  # This is for backward compatability with a bug

        'pred_aspect_ratios': [[[1, 1 / 2, 2]]] * 5,
        'pred_scales': [[24], [48], [96], [192], [384]],
    }),

    # FPN Settings
    'fpn': fpn_base.copy({
    
        'use_conv_downsample': True,
        'num_downsample': 2,
    }),

    # Mask Settings
    'mask_type': mask_type.lincomb,
    'mask_alpha': 6.125,
    'mask_proto_src': 0,
    'mask_proto_net': [(256, 3, {
    'padding': 1})] * 3 + [(None, -2, {
    }), (256, 3, {
    'padding': 1})] + [(32, 1, {
    })],
    'mask_proto_normalize_emulate_roi_pooling': True,

    # Other stuff
    'share_prediction_module': True,
    'extra_head_net': [(256, 3, {
    'padding': 1})],

    'positive_iou_threshold': 0.5,
    'negative_iou_threshold': 0.4,

    'crowd_iou_threshold': 0.7,

    'use_semantic_segmentation_loss': True,
})

4. 训练

修改一下train.py里面的配置。

batch_size： 模型有点大，我用1080Ti 11G显存，一块最多只能训练14个batch_size（输入尺寸550）。
batch_alloc： 你用多块GPU的话，这里可以对应写每块GPU要多少个batchsize。（如上面batchsize为28，这里就是14,14）
config: 对应刚刚配置数据的代码。
save_folder： 保存权重的地址。
init_weights： 第一次训练使用的权重。（如果save_folder中没有权重的话）
resume： 这个为latest的时候会优先使用save_folder中的最后一个权重进行训练。

训练的优化器，这里默认使用SGD。这里要到train.py中optimizer = optim.SGD(net.parameters(), lr=args.lr, momentum=args.momentum, weight_decay=args.decay)这一块去改。我也不太清楚是不是分割模型用SGD比较好，后面我也有改成Adam优化器，损失会降的比较快。

这里如果要指定那几块GPU进行训练，需要在train.py最上端加入这一串代码。

import os
os.environ['CUDA_VISIBLE_DEVICES'] = '6,7'  # 

5. 预测

训练完模型要跑预测只需要按照eval.py中的配置进行就行了。

6. 总结

实例分割模型训练起来比较慢，而且模型训练起来也比较难。我正在慢慢摸索，也尝试了很多技巧。（最近也还在忙这个项目，后面又时间再把使用过的技巧，还有一些坑总结一下。）