2.1【Pytorch版(torch-cuda版 ) Mask-RCNN 训练自己的数据集】

写在前面:由于课程试验要求,需要基于pytorch实现maskrcnn,so最近又跑了一下pytorch版的maskrcnn,官方已经给出了详细的教程,虽然说支持cpu推理,但是不支持cpu训练啊,奈何手头上只有一个cpu本,也没有nvida显卡,只有intel的集显,so整理一波本次训练maskrcnn的过程。

这篇blog是按照常规的torch版的maskrcnn的训练教程写的,理论上,所依赖的是torch-cuda版本,后续还有一篇是基于torch-cpu版本的。因为在试验过程中先踩了这部分的坑,所以写在前面。(再次提醒,如果电脑和本人一样只有intel集显的话,下面的一系列步骤对成功在电脑上跑通maskrcnn没有任何帮助,本人的下一篇blog会提到如何不用cuda在cpu本下跑maskrcnn。这篇blog仅仅是为了方便疫情过后可以回到实验室继续搞,提前做好记录。)

环境
Ubuntu16.04
torch == 1.5.0+cu92
torchvision == 0.6.0+cu92

这里要注意,torch版本>=0.3.0即可,使用的torch和torchvision是匹配的,且都是cuda版本,如何选择这两个的匹配版本,请看:https://pytorch.org/
具体操作流程见我的另一篇blog,里面有提到:
【YoloV3–pytorch】Part One:基于Pytorch的YoloV3训练自己的数据集----准备数据集、配置文件并下载预训练权重文件

一、下载官方demo并配置环境

先去官网将代码下载下来 https://github.com/facebookresearch/maskrcnn-benchmark

cd maskrcnn-benchmark-master
sudo python3 setup.py build develop

笔者现有打包好的pycharm虚拟环境中包含了大部分需要配置安装的模块,so最初笔者是缺什么模块补什么模块,折腾了一番还是有问题,最终屈服了,直接运行上述命令,bingo。

安装好setup中要求的环境之后,可以根据官网的教程进行测试。

cd demo
# by default, it runs on the GPU
# for best results, use min-image-size 800
python webcam.py --min-image-size 800
# can also run it on the CPU
python webcam.py --min-image-size 300 MODEL.DEVICE cpu
# or change the model that you want to use
python webcam.py --config-file ../configs/caffe2/e2e_mask_rcnn_R_101_FPN_1x_caffe2.yaml --min-image-size 300 MODEL.DEVICE cpu
# in order to see the probability heatmaps, pass --show-mask-heatmaps
python webcam.py --min-image-size 300 --show-mask-heatmaps MODEL.DEVICE cpu
# for the keypoint demo
python webcam.py --config-file ../configs/caffe2/e2e_keypoint_rcnn_R_50_FPN_1x_caffe2.yaml --min-image-size 300 MODEL.DEVICE cpu

简单的用前两种测试一下就可以了,在测试过程中可能会提示你缺少某些模块,对应安装就可以了。(如何解决python安装模块慢的问题,可以参看我的另一篇blog,解决Ubuntu下python安装模块速度慢的问题,不得不说,工具选得好,效率刷刷得往上彪)

至此,倘若还缺少apex库,建议通过github下载,然后再进行安装,因为本人直接通过pip安装后,后续运行demo的时候还会报错,保险起见,采用前者。
https://github.com/NVIDIA/apex

cd apex-master
sudo python3 setup.py install

解决方式参考blog

二、数据准备及模型配置

Step1:转换数据格式

笔者是采用labelme来标注的,需要对格式进行转换,转为coco格式,转换代码如下:

# -*- coding:utf-8 -*-
# !/usr/bin/env python

import argparse
import json
import matplotlib.pyplot as plt
from labelme import utils
import skimage.io as io
import sys
sys.path.remove('/opt/ros/kinetic/lib/python2.7/dist-packages')

import cv2
import numpy as np
import glob
import PIL.Image
from shapely.geometry import Polygon#https://shapely.readthedocs.io/en/latest/manual.html#geometric-objects

class labelme2coco(object):
    def __init__(self,labelme_json=[],save_json_path='/home/eden/Downloads/maskrcnn-benchmark/datasets/val2020.json'):
        '''
        :param labelme_json: 所有labelme的json文件路径组成的列表
        :param save_json_path: json保存位置
        '''
        self.labelme_json=labelme_json#所有的json文件
        self.save_json_path=save_json_path#输出文件new.json的路径
        self.images=[]
        self.categories=[]
        self.annotations=[]
        # self.data_coco = {}
        self.label=[]
        self.annID=1
        self.height=0
        self.width=0

        self.save_json()

#我的数据用label标注的名称格式是:000001.jpg、000002.jpg....
    def data_transfer(self):
        for num,json_file in enumerate(self.labelme_json):
            with open(json_file,'r') as fp:
                data = json.load(fp)  # 加载json文件
                self.images.append(self.image(data,num))
                for shapes in data['shapes']:
                    label=shapes['label'].split('_')
                    if label[0] not in self.label:
                        self.categories.append(self.categorie(label))
                        self.label.append(label[0])
                    points=shapes['points']
                    self.annotations.append(self.annotation(points,label,num))
                    self.annID+=1
    
    def image(self,data,num):
        image={}
        img = utils.img_b64_to_arr(data['imageData'])  # 解析原图片数据
        # img=io.imread(data['imagePath']) # 通过图片路径打开图片
        # img = cv2.imread(data['imagePath'], 0)
        height, width = img.shape[:2]
        img = None
        image['height']=height
        image['width'] = width
        image['id']=num+1
        image['file_name'] = data['imagePath'].split('/')[-1]

        self.height=height
        self.width=width

        return image

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

    def annotation(self,points,label,num):
        annotation={}
        annotation['segmentation']=[list(np.asarray(points).flatten())]
        poly = Polygon(points)
        annotation['iscrowd'] = 0
        annotation['image_id'] = num+1
        area_ = round(poly.area,6)
        annotation['area'] = area_

        # annotation['bbox'] = str(self.getbbox(points)) # 使用list保存json文件时报错(不知道为什么)
        # list(map(int,a[1:-1].split(','))) a=annotation['bbox'] 使用该方式转成list
        annotation['bbox'] = list(map(float,self.getbbox(points)))

        annotation['category_id'] = self.getcatid(label)
        annotation['id'] = self.annID
        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_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['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 更加美观显示

labelme_json=[]
for i in range(1,71):#遍历所有的labelme 标准的json,载入并转换

    ind='/home/eden/Mask_RCNN-master/data/train_data/json/'+str(i).zfill(6)+'.json'

    labelme_json.append(ind)

lab = labelme2coco(labelme_json)#调用labelme2coco类
print('Saved to :',lab.save_json_path)#保存

我的数据的名称格式是:000001.jpg、000002.jpg…。一般的数据格式是 父类_子类_***.jpg,倘若你的数据格式是这样的,主要修改这个函数的代码。

def data_transfer(self):

主要用来修改label的截断以符合自己的数据格式,同时还要注意对应修改该函数下的label

def categorie(self,label):

转换为之后会在当前目录下生成一个新的json文件,该json文件满足coco格式(:上述代码生成的coco json文件并不包含area参数,实验说明即使json文件中没有area参数也不影响后续正常训练。)

关于coco数据格式的一些资料:
参考blog1:该blog有包含area的demo
对coco数据集的详细解读:
参考blog2
参考blog3
参考blog4
Step2:配置数据格式

接下来进入到maskrcnn-benchmark-master下的maskrcnn_benchmark文件夹,新建 d a t a s e t s \color{red}{datasets} datasets m y c o n f i g \color{red}{myconfig} myconfig文件夹,
在datasets文件夹下继续新建一个名为coco的文件夹,在coco文件夹下新建三个文件夹,名称分别为 a n n o t a t i o n s \color{red}{annotations} annotations t r a i n 2020 \color{red}{train2020} train2020 v a l 2020 \color{red}{val2020} val2020
(其实没有必要严格的按照这个格式来,实际上只是为了和目前的一些数据格式保持一致,只要后面对数据集路径的修改都严格的改成自己的路径就可以了。)
annotations文件夹下存放两个json文件,分别为train2020.json和val2020.json(这两个文件就是经过上述格式转换demo得到的,其中所对应的图片是train2020和val2020文件夹下的图片)
图片数据整理好了之后,接下里就是配置模型文件了。

插入一个常识:有时在训练时为了方便,直接将训练集和验证集使用同一个数据集,这样做并不可取。理由见该博文:
验证集和训练集为什么要不一致

Step3:配置模型文件

1. 在新建的myconfig文件下将./maskrcnn-benchmark-master/configs里想要训练的.yaml配置文件复制过去,把并且把
./maskrcnn-benchmark-master/maskrcnn_benchmark/
config文件夹下的__pycache__和paths_catalogs.py复制到到myconfig文件夹下,本文使用的是e2e_mask_rcnn_fbnet.yaml。
2. 修改maskrcnn_benchmark/utils下的checkpoint.py,注释65和68两行(self.optimizer.load_state…和self.scheduler.load_…)
3. 修改paths_catalogs.py文件,将训练样本和验证样本的路径修改为自己的数据路径:2.1【Pytorch版(torch-cuda版 ) Mask-RCNN 训练自己的数据集】_第1张图片
接下来修改.yaml模型配置参数,主要修改NUM_CLASSES,DATASETS中的train和test ,模型保存文件OUTPUT_DIR以及PATHS_CATALOG
我的训练类别只有一类,所以NUM_CLASSES=2
在这里插入图片描述
DATASETS如上,这里要注意的时,即使你的训练或测试只有一个数据样本,后面的逗号也不可以去掉。2.1【Pytorch版(torch-cuda版 ) Mask-RCNN 训练自己的数据集】_第2张图片
最后加上OUTPUT_DIR和PATHS_CATALOG。其他的关于batch size、学习率等等,依照个人实际情况进行配置即可。

三、 训练

戏剧性的转折来了,理论上来说到此为止,只要在文件跟目录下运行:

python3 tools/train_net.py --config-file=.yaml文件的路径

这个时候就已经开始训练了,但无奈的是,本人安装的是torch-cpu版本,所以会显示如下错误,在这里插入图片描述
这是因为apex依赖于cuda,本人电脑也没有nvida显卡,所以也没有安装cuda版本的torch,通过一番搜寻发现,有网友说即便如此,只要安装cuda版本的torch和torchvision也是可以正常运行的,由于时间关系,重新搭建一个虚拟环境再下载一大堆的安装包时间很长(主要是本人硬件垃圾,而且网速也差劲),这一方法就没有尝试,待疫情过后到实验室再尝试吧。如果你的电脑可安装cuda,那么可继续按照下面的来,如果像我一样,只有intel集显的话,请看下一篇blog中提到的一种在cpu下训练maskrcnn的方法。

本篇blog主要参考资料:

参考blog1
参考blog2
参考blog3
参考blog4

你可能感兴趣的:(图像分割)