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

EdenGabriel

已于 2022-04-06 17:42:36 修改

阅读量9.2k

点赞数 4

分类专栏：图像分割文章标签：深度学习 pytorch

于 2020-05-15 22:27:43 首次发布

本文链接：https://blog.csdn.net/qq_38587510/article/details/106140983

版权

图像分割专栏收录该内容

5 篇文章 3 订阅

订阅专栏

写在前面：由于课程试验要求，需要基于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文件夹，新建 $\color{red}{datasets}$ 、 $\color{red}{myconfig}$ 文件夹，
在datasets文件夹下继续新建一个名为coco的文件夹，在coco文件夹下新建三个文件夹，名称分别为 $\color{red}{annotations}$ 、 $\color{red}{train2020}$ 、 $\color{red}{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文件，将训练样本和验证样本的路径修改为自己的数据路径：在这里插入图片描述
接下来修改.yaml模型配置参数，主要修改NUM_CLASSES，DATASETS中的train和test ，模型保存文件OUTPUT_DIR以及PATHS_CATALOG
我的训练类别只有一类，所以NUM_CLASSES=2

DATASETS如上，这里要注意的时，即使你的训练或测试只有一个数据样本，后面的逗号也不可以去掉。 在这里插入图片描述
最后加上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的方法。

4
点赞
踩
81

收藏

觉得还不错? 一键收藏
11
评论
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版本的。因为在试验过程中先踩了这部分的坑，.
复制链接

扫一扫