Mask RCNN训练自己的数据集——程序调试记录

这篇主要记录用Mask RCNN训练自己的数据集,防止以后忘记。
主要参考:

  1. https://blog.csdn.net/xxxy502/article/details/97754384
  2. https://github.com/matterport/Mask_RCNN (开放源代码,需要下载进行调试)
  3. https://pan.baidu.com/s/1VBIFgpX95FXi6-5u2SB2Cw ( 密码:pn2f,需下载预训练权重文件mask_rcnn_coco.h5)

本文主要参考了博客https://blog.csdn.net/weixin_44826203。非常感谢该博主提供的训练脚本以及一些非常实用的代码,让我少走了很多弯路。如果本文内容有所侵权,请提醒我删除,谢谢。

0 环境介绍

环境windows10、pycharm
numpy == 1.19.5
scipy
Pillow
cython
matplotlib
scikit-image
tensorflow ==1.5.0
keras == 2.1.6
opencv-python
h5py == 2.10.0
imgaug
IPython[all]

这里引用的是Mask RCNN 里面的requirements.txt文件,上述是自己调通的版本。

1 准备工作

1.1 下载Mask RCNN源代码,配置环境,创建所需的文件及目录。

  • 在Mask_RCNN-master/samples 的文件夹下,自己建一个数据库目录命名为my_data。
  • 在Mask_RCNN-master 文件夹下建立train.py、predict.py
    如下图所示
    在这里插入图片描述

1.2 建立自己的数据集(方法见 Sec. 2)

建立的数据库(my_data)应有如下的内容
在这里插入图片描述

2 数据集的制作

2.1 使用labelme对数据进行标注

首先,在对数据进行标注之前,建议先把所有数据转为统一的尺寸,并且重命名使得更加规整。
可以参考下面这段代码(但是我这个代码执行完之后,会多出一倍的图片,比如原来是5张,他会变成10张,我不知道是为什么,但不影响使用,把后面的多出来删掉就好了,按名称排序前5张和后5张是完全相同的)

import os
import os.path
from PIL import Image
def ResizeImage(f_path, out_dir, width, height, type):
    i = 0
    for root, dirs, files in os.walk(f_path):
        # 使用os.walk()方法遍历输出一个文件夹下的所有文件名
        # os.path.join()方法拼接文件名返回所有文件的路径
        for name in files:
            file = os.path.join(root, name)
            img = Image.open(file)
            out = img.resize((width, height), Image.ANTIALIAS)

            i += 1
            f_name = str(i) + '.jpeg'
            out_name = os.path.join(out_dir, f_name)
            out.save(out_name, type)

f_path = r'************'          # 图片所在路径
out_dir = r'***********/out_dir'  # 图片输出路径
width = 384                       # 定义输出图片的宽和高(最好为8的倍数)
height = 512
type = 'jpeg'                     # 定义输出图片类型

ResizeImage(f_path, out_dir, width, height, type)

接下来,使用labelme工具对图片进行标注,如果你没有用过labelme工具,首先一键安装它。

conda install pyqt
pip install labelme -i https://pypi.tuna.tsinghua.edu.cn/simple

安装很简单,遇到问题百度很容易解决。
然后在cmd中输入指令“labelme”则自动弹出窗口。
在这里插入图片描述
左上角open_dir 选择刚才输出的out_dir ,选择左边create那个,就可以创建区域了,然后输入标签的名字(这个名字要记好,因为后面还会再用到),比如我这里创建一个类别叫patrick_star。当然如果有更多的类,你也可以创建第二类第三类。

创建好了之后,会生成一个.json 文件,里边记录这各个标记点的位置信息和类别,然后需要把这个json文件转换成png格式的mask文件。

2.2 将json文件转换为模型需要的mask文件

使用的是labelme的labelme_json_to_dataset函数。
可以写一个简单的脚本让它批量执行。在json文件的同一目录下创建一个json2dataset.py文件,里面是这些内容:

import os
files=os.listdir('./')
files.remove('json2dataset.py')   # 在目录信息中删除这个py文件本身
for i in range(len(files)):
    os.system('labelme_json_to_dataset '+files[i])

然后在当前目录开启命令行,用python执行它就好了。
转换完之后,你看到的应该是这个样子:
在这里插入图片描述

每个文件夹里边应该包含这些文件:
在这里插入图片描述
注意这个1.png最初你看到的应该是label.png,但是接下来在2.3中我们要把它转换成对应文件夹的名字。
如果你生成的文件夹里没有.yaml文件,不要着急,跟着下面的操作来。
首先在你的电脑上找到json_to_dataset.py这个脚本,打开之后做如下修改:

# 最前面加入导包
import yaml

# 中间是代码的主体部分就不贴了
# 然后在最下面main函数之前加上这一部分:
    logger.warning('info.yaml is being replaced by label_names.txt')
    info = dict(label_names=label_names)
    with open(osp.join(out_dir, 'info.yaml'), 'w') as f:
        yaml.safe_dump(info, f, default_flow_style=False)
    logger.info('Saved to: {}'.format(out_dir))

# 下面就是main函数了
if __name__ == '__main__':
    main()

再重新执行一次json_to_dataset,然后就生成yaml了。

2.3 将数据整理成模型认可的形式

为了适应模型内部默认的路径格式,需要对label.png进行简单的重命名(否则就要去代码里边改,比较麻烦)。比如你的json文件夹叫1_json,那这个png的图就应该改成1.png,为了大家批量操作,我写了个简单的脚本。

# 把label.png改名为1.png
import os
for root, dirs, names in os.walk(r'******out_dir'):   # 改成你自己的json文件夹所在的目录
    for dr in dirs:
        file_dir = os.path.join(root, dr)
        # print(dr)
        file = os.path.join(file_dir, 'label.png')
        # print(file)
        new_name = dr.split('_')[0] + '.png'
        new_file_name = os.path.join(file_dir, new_name)
        os.rename(file, new_file_name)

找到之前的文件夹my_data,在里边创建四个文件夹,并分别存放以下内容:

文件夹内容
cv2_maskjson_to_dataset生成文件夹中的png格式label文件
labelme_jsonjson_to_dataset生成的文件夹
jsonlabelme生成的json文件
pic尺寸标准化之后的原图

其他三个文件夹都比较容易准备好,cv2_mask这个挨个复制比较麻烦,可以用下面的方法批量复制:

import os
from shutil import copyfile
for root, dirs, names in os.walk(r'******'):   # 改成你自己的json文件夹所在的目录
    for dr in dirs:
        file_dir = os.path.join(root, dr)
        print(dr)
        file = os.path.join(file_dir,'label.png')
        print(file)
        new_name = dr.split('_')[0] + '.png'
        new_file_name = os.path.join(file_dir, new_name)
        print(new_file_name)
        
        tar_root = r'******my_data/cv2_mask'      # 目标路径
        tar_file = os.path.join(tar_root, new_name)
        copyfile(new_file_name, tar_file)

至此,准备工作已经完成了。

3 训练过程

3.1 下载预训练权重文件(参考链接3)

这个权重文件mask_rcnn_coco.h5下载参考链接,需要的同学可以自行下载。此权重是在coco数据集上训练,迁移过来的。
下载好了之后放在项目根目录Mask_RCNN-master中就可以了。

3.2 训练代码

前期在Mask_RCNN-master目录下创建了一个名为train.py的文件,加入以下内容:
(首先声明一下,这段代码是从别人那里参考来的,侵权请提醒删除)

# -*- coding: utf-8 -*-
 
import os
import sys
import random
import math
import re
import time
import numpy as np
import cv2
# import matplotlib
# import matplotlib.pyplot as plt
import tensorflow as tf
from mrcnn.config import Config
# import utils
from mrcnn import model as modellib, utils
from mrcnn import visualize
import yaml
from mrcnn.model import log
from PIL import Image
 
# os.environ["CUDA_VISIBLE_DEVICES"] = "0"
# Root directory of the project
ROOT_DIR = os.getcwd()
 
# ROOT_DIR = os.path.abspath("../")
# Directory to save logs and trained model
MODEL_DIR = os.path.join(ROOT_DIR, "logs")
 
iter_num = 0
 
# Local path to trained weights file
COCO_MODEL_PATH = os.path.join(ROOT_DIR, "mask_rcnn_coco.h5")
# Download COCO trained weights from Releases if needed
# if not os.path.exists(COCO_MODEL_PATH):
#     utils.download_trained_weights(COCO_MODEL_PATH)
 
 
class ShapesConfig(Config):
    """Configuration for training on the toy shapes dataset.
    Derives from the base Config class and overrides values specific
    to the toy shapes dataset.
    """
    # Give the configuration a recognizable name
    NAME = "shapes"
 
    # Train on 1 GPU and 8 images per GPU. We can put multiple images on each
    # GPU because the images are small. Batch size is 8 (GPUs * images/GPU).
    GPU_COUNT = 1
    IMAGES_PER_GPU = 1
 
    # Number of classes (including background)
    NUM_CLASSES = 4 + 1  # background + 1 shapes  注意这里我是4类,所以是4+1
 
    # Use small images for faster training. Set the limits of the small side
    # the large side, and that determines the image shape.
    IMAGE_MIN_DIM = 256
    IMAGE_MAX_DIM = 1024
 
    # Use smaller anchors because our image and objects are small
    # RPN_ANCHOR_SCALES = (8 * 6, 16 * 6, 32 * 6, 64 * 6, 128 * 6)  # anchor side in pixels
    RPN_ANCHOR_SCALES = (16 * 6, 32 * 6, 64 * 6, 128 * 6, 256 * 6)    # 我的图片中目标比较大,所以我把anchor的尺寸也设置的大了一点
 
    # Reduce training ROIs per image because the images are small and have
    # few objects. Aim to allow ROI sampling to pick 33% positive ROIs.
    TRAIN_ROIS_PER_IMAGE = 100
 
    # Use a small epoch since the data is simple
    STEPS_PER_EPOCH = 50     #  每个epoch中迭代的step,最好不要改动
 
    # use small validation steps since the epoch is small
    VALIDATION_STEPS = 50
 
 
config = ShapesConfig()
config.display()
 
 
class DrugDataset(utils.Dataset):
    # 得到该图中有多少个实例(物体)
    def get_obj_index(self, image):
        n = np.max(image)
        return n
 
    # 解析labelme中得到的yaml文件,从而得到mask每一层对应的实例标签
    def from_yaml_get_class(self, image_id):
        info = self.image_info[image_id]
        with open(info['yaml_path']) as f:
            temp = yaml.load(f.read())
            labels = temp['label_names']
            del labels[0]
        return labels
 
    # 重新写draw_mask
    def draw_mask(self, num_obj, mask, image, image_id):
        # print("draw_mask-->",image_id)
        # print("self.image_info",self.image_info)
        info = self.image_info[image_id]
        # print("info-->",info)
        # print("info[width]----->",info['width'],"-info[height]--->",info['height'])
        for index in range(num_obj):
            for i in range(info['width']):
                for j in range(info['height']):
                    # print("image_id-->",image_id,"-i--->",i,"-j--->",j)
                    # print("info[width]----->",info['width'],"-info[height]--->",info['height'])
                    at_pixel = image.getpixel((i, j))
                    if at_pixel == index + 1:
                        mask[j, i, index] = 1
        return mask
 
    # 重新写load_shapes,里面包含自己的自己的类别
    # 并在self.image_info信息中添加了path、mask_path 、yaml_path
    # yaml_pathdataset_root_path = "/dateset/"
    # img_floder = dataset_root_path + "rgb"
    # mask_floder = dataset_root_path + "mask"
    # dataset_root_path = "/tongue_dateset/"
    def load_shapes(self, count, img_floder, mask_floder, imglist, dataset_root_path):
        """Generate the requested number of synthetic images.
        count: number of images to generate.
        height, width: the size of the generated images.
        """
        # Add classes
        self.add_class("shapes", 1, "leibie1")
        self.add_class("shapes", 2, "leibie2")
        self.add_class("shapes", 3, "leibie3")
        self.add_class("shapes", 4, "leibie4")
 
        for i in range(count):
            # 获取图片宽和高
            print(i)
            filestr = imglist[i].split(".")[0]
            # print(imglist[i],"-->",cv_img.shape[1],"--->",cv_img.shape[0])
            # print("id-->", i, " imglist[", i, "]-->", imglist[i],"filestr-->",filestr)
            # filestr = filestr.split("_")[1]
            mask_path = mask_floder + "/" + filestr + ".png"
            yaml_path = dataset_root_path + "labelme_json/" + filestr + "_json/info.yaml"
            print(dataset_root_path + "labelme_json/" + filestr + "_json/img.png")
            cv_img = cv2.imread(dataset_root_path + "labelme_json/" + filestr + "_json/img.png")
            print(type(cv_img))
 
            self.add_image("shapes", image_id=i, path=img_floder + "/" + imglist[i],
                           width=cv_img.shape[1], height=cv_img.shape[0], mask_path=mask_path, yaml_path=yaml_path)
 
    # 重写load_mask
    def load_mask(self, image_id):
        """Generate instance masks for shapes of the given image ID.
        """
        global iter_num
        print("image_id", image_id)
        info = self.image_info[image_id]
        count = 1  # number of object
        img = Image.open(info['mask_path'])
        num_obj = self.get_obj_index(img)
        mask = np.zeros([info['height'], info['width'], num_obj], dtype=np.uint8)
        mask = self.draw_mask(num_obj, mask, img, image_id)
        occlusion = np.logical_not(mask[:, :, -1]).astype(np.uint8)
        for i in range(count - 2, -1, -1):
            mask[:, :, i] = mask[:, :, i] * occlusion
 
            occlusion = np.logical_and(occlusion, np.logical_not(mask[:, :, i]))
        labels = []
        labels = self.from_yaml_get_class(image_id)
        labels_form = []
        for i in range(len(labels)):
            if labels[i].find("leibie1") != -1:
                labels_form.append("leibie1")
            elif labels[i].find("leibie2") != -1:
                labels_form.append("leibie2")
            elif labels[i].find("leibie3") != -1:
                labels_form.append("leibie3")
            elif labels[i].find("leibie4") != -1:
                labels_form.append("leibie4")

        class_ids = np.array([self.class_names.index(s) for s in labels_form])
        return mask, class_ids.astype(np.int32)
 
'''
def get_ax(rows=1, cols=1, size=8):
    """Return a Matplotlib Axes array to be used in
    all visualizations in the notebook. Provide a
    central point to control graph sizes.
    Change the default size attribute to control the size
    of rendered images
    """
    _, ax = plt.subplots(rows, cols, figsize=(size * cols, size * rows))
    return ax
'''
 
# 基础设置
dataset_root_path = "samples/my_data/"    # 你的数据的路径
img_floder = dataset_root_path + "pic"
mask_floder = dataset_root_path + "cv2_mask"
# yaml_floder = dataset_root_path
imglist = os.listdir(img_floder)
count = len(imglist)
 
# train与val数据集准备
dataset_train = DrugDataset()
dataset_train.load_shapes(count, img_floder, mask_floder, imglist, dataset_root_path)
dataset_train.prepare()
 
# print("dataset_train-->",dataset_train._image_ids)
 
dataset_val = DrugDataset()
dataset_val.load_shapes(count, img_floder, mask_floder, imglist, dataset_root_path)
dataset_val.prepare()
 
# print("dataset_val-->",dataset_val._image_ids)
 
# Load and display random samples
# image_ids = np.random.choice(dataset_train.image_ids, 4)
# for image_id in image_ids:
#    image = dataset_train.load_image(image_id)
#    mask, class_ids = dataset_train.load_mask(image_id)
#    visualize.display_top_masks(image, mask, class_ids, dataset_train.class_names)
 
# Create model in training mode
model = modellib.MaskRCNN(mode="training", config=config,
                          model_dir=MODEL_DIR)
 
# Which weights to start with?
init_with = "coco"  # imagenet, coco, or last
 
if init_with == "imagenet":
    model.load_weights(model.get_imagenet_weights(), by_name=True)
elif init_with == "coco":
    # Load weights trained on MS COCO, but skip layers that
    # are different due to the different number of classes
    # See README for instructions to download the COCO weights
    # print(COCO_MODEL_PATH)
    model.load_weights(COCO_MODEL_PATH, by_name=True,
                       exclude=["mrcnn_class_logits", "mrcnn_bbox_fc",
                                "mrcnn_bbox", "mrcnn_mask"])
elif init_with == "last":
    # Load the last model you trained and continue training
    model.load_weights(model.find_last()[1], by_name=True)
 
# Train the head branches
# Passing layers="heads" freezes all layers except the head
# layers. You can also pass a regular expression to select
# which layers to train by name pattern.
model.train(dataset_train, dataset_val,
            learning_rate=config.LEARNING_RATE,
            epochs=10,
            layers='heads')         # 固定其他层,只训练head,epoch为10
 
# Fine tune all layers
# Passing layers="all" trains all layers. You can also
# pass a regular expression to select which layers to
# train by name pattern.
model.train(dataset_train, dataset_val,
            learning_rate=config.LEARNING_RATE / 10,
            epochs=10,
            layers="all")           # 微调所有层的参数,epoch为10

3.3 修改自己的参数

在上面的代码中主要有这几个地方需要修改:
1、第55行左右的NUM_CLASSES,类别总数,要包含背景数。比如你只有一类patrick_star,那这里就是1 + 1 ,如果你有两类,那这里就是2 + 1。
2、第124行左右添加自己的类别

# Add classes
self.add_class("shapes", 1, "leibie1")

如果你只有一类patrick_star,那就把这里的leibie1换成patrick_star,注意这里的名字一定要与你在labelme中标注的类别名称保持一致,否则会报错。
3、166行左右修改类别,和第2点一样,leibie1换成相应的名字。

        for i in range(len(labels)):
            if labels[i].find("leibie1") != -1:
                labels_form.append("leibie1")

4、修改数据集所在路径dataset_root_path = “samples/my_data/”
5、其他设置
代码默认的将batch_size设置为8,效果还不错,没有必要修改它(除非你的数据集特别大/小)
epoch的设置在代码最后面。
step的设置在代码前段,默认是50,改的话最好是50的倍数,因为模型的日志保存是每100步保存一次,我之前把它设置成30结果tensorboard什么也没存下来(也有可能是我的问题,没仔细看),当然你也可以修改模型保存的参数。
anchor的尺寸,在step设置的前面。

然后python执行它就可以开始训练了,应该会提示一堆缺少model,根据提示安装它就可以了。
如果报了错误IndexError: boolean index did not match indexed array along dimension 0; dim
dimension相关的错误,要么是你的数据没有没有处理好(比如名称、路径不对,labelme中的类别名和代码中的类别名不匹配等等),要么是代码里的类别数量没设置对,或者某些拼写错误,总之仔细检查一遍应该可以解决。

开始正确训练后会看见以下界面:
在这里插入图片描述
有两个警告分别是

  1. YAMLLoadWarning: calling yaml.load() without Loader=… is deprecated, as the default Loader is unsafe. Please read https://msg.pyyaml.org/load for full details. temp = yaml.load(f.read())
    需要将第91行代码换成
    在这里插入图片描述
 temp =yaml.load(f.read(), Loader=yaml.FullLoader)

2.FutureWarning: Input image dtype is bool. Interpolation is not defined with bool data type. Please set order to 0 or explicitely cast input image to another data type. Starting from version 0.19 a ValueError will be raised instead of this warning. order
需要将 scikit-image-0.17.2 换成 scikit-image-0.16.2 (前面对其版本无要求,但别的版本会报错)
在这里插入图片描述
再次训练即可。上述两警告对结果暂时还没有影响,改不改全看个人。
在这里插入图片描述

3.4 训练结束的状态

训练结束后,根目录下会多出一个logs文件夹。
在这里插入图片描述
里面装的是记录训练的模型为.h5文件,名字由日期加时间构成。后面预测阶段会用到训练好的模型。以第一个为例:里面记录了迭代10次的训练模型
在这里插入图片描述

4 测试效果

4.1 测试结果代码

# -*- coding: utf-8 -*-
import os
import sys
import random
import math
import numpy as np
import skimage.io
import matplotlib
import matplotlib.pyplot as plt
import cv2
import time
from mrcnn.config import Config
from datetime import datetime
# Root directory of the project
ROOT_DIR = os.getcwd()
 
# Import Mask RCNN
sys.path.append(ROOT_DIR)  # To find local version of the library
from mrcnn import utils
import mrcnn.model as modellib
from mrcnn import visualize
# Import COCO config
# sys.path.append(os.path.join(ROOT_DIR, "samples/coco/"))  # To find local version
# from samples.coco import coco
 
 
# Directory to save logs and trained model
MODEL_DIR = os.path.join(ROOT_DIR, "logs")
 
# Local path to trained weights file
COCO_MODEL_PATH = "*******/Mask_RCNN-master/logs/shapes*****/mask_rcnn_shapes_00**.h5"   #  模型保存目录
# Download COCO trained weights from Releases if needed
if not os.path.exists(COCO_MODEL_PATH):
    utils.download_trained_weights(COCO_MODEL_PATH)
    print("cuiwei***********************")
 
# Directory of images to run detection on
IMAGE_DIR = os.path.join(ROOT_DIR, "images")
 
class ShapesConfig(Config):
    """Configuration for training on the toy shapes dataset.
    Derives from the base Config class and overrides values specific
    to the toy shapes dataset.
    """
    # Give the configuration a recognizable name
    NAME = "shapes"
 
    # Train on 1 GPU and 8 images per GPU. We can put multiple images on each
    # GPU because the images are small. Batch size is 8 (GPUs * images/GPU).
    GPU_COUNT = 1
    IMAGES_PER_GPU = 1
 
    # Number of classes (including background)
    NUM_CLASSES = 4 + 1  # background + 3 shapes
 
    # Use small images for faster training. Set the limits of the small side
    # the large side, and that determines the image shape.
    IMAGE_MIN_DIM = 320
    IMAGE_MAX_DIM = 384
 
    # Use smaller anchors because our image and objects are small
    RPN_ANCHOR_SCALES = (8 * 6, 16 * 6, 32 * 6, 64 * 6, 128 * 6)  # anchor side in pixels
 
    # Reduce training ROIs per image because the images are small and have
    # few objects. Aim to allow ROI sampling to pick 33% positive ROIs.
    TRAIN_ROIS_PER_IMAGE =100
 
    # Use a small epoch since the data is simple
    STEPS_PER_EPOCH = 50
 
    # use small validation steps since the epoch is small
    VALIDATION_STEPS = 50
 
#import train_tongue
#class InferenceConfig(coco.CocoConfig):
class InferenceConfig(ShapesConfig):
    # Set batch size to 1 since we'll be running inference on
    # one image at a time. Batch size = GPU_COUNT * IMAGES_PER_GPU
    GPU_COUNT = 1
    IMAGES_PER_GPU = 1
 
config = InferenceConfig()
 
model = modellib.MaskRCNN(mode="inference", model_dir=MODEL_DIR, config=config)
 
 
# Create model object in inference mode.
model = modellib.MaskRCNN(mode="inference", model_dir=MODEL_DIR, config=config)
 
# Load weights trained on MS-COCO
model.load_weights(COCO_MODEL_PATH, by_name=True)
 
# COCO Class names
# Index of the class in the list is its ID. For example, to get ID of
# the teddy bear class, use: class_names.index('teddy bear')
class_names = ['BG', 'leibie1'']   # 注意修改类别名称
# Load a random image from the images folder
file_names = next(os.walk(IMAGE_DIR))[2]
image = skimage.io.imread("img.png")      # 你想要测试的图片
 
a=datetime.now()
# Run detection
results = model.detect([image], verbose=1)
b=datetime.now()
# Visualize results
print("shijian",(b-a).seconds)
r = results[0]
visualize.display_instances(image, r['rois'], r['masks'], r['class_ids'],
                            class_names, r['scores'])


4.2 修改参数

  1. 第58行NUM_CLASSES,类别总数,要和上面训练过程保持一致,
  2. 第34行需要预测模型的地址,应当是训练所保留的模型目录在logs/shape*****/***.h5
  3. 第101行修改类别名称BG为背景不能改,剩下的将自己类别名称依次写入即可
  4. 第104行修改测试的图像所在位置

4.3 测试结果

会出现
在这里插入图片描述

然后过一会儿会弹出图片。

如果卡在了Downloading pretrained model to *****,说明你的模型路径没设置好,仔细检查。
至此应该是运行完成了,剩下的调参工作因人而异,仔细看代码吧。

5 常见问题

  1. 刚运行时可能有个有关 keepdims 的报错,将其换成 keep_dims即可
    若是还有问题评论区留言
  • 3
    点赞
  • 24
    收藏
    觉得还不错? 一键收藏
  • 打赏
    打赏
  • 5
    评论
Mask R-CNN 是一种基于 Faster R-CNN 的目标检测算法,可以检测出物体的位置,还可以精确地分割出物体的 mask。下面是一个 Mask R-CNN 训练自己数据集的基本步骤: 1. 准备数据集 首先,您需要准备自己的数据集,包括图片和标注。图片可以是任何大小,但最好是统一大小,以便于训练。标注可以使用标准的 COCO 标注格式,包括标注文件和类别文件。 2. 定义类别 您需要定义自己数据集的类别,例如人、车、动物等等。您可以创建一个名为 `CustomConfig` 的类,继承自 `Config` 类,并在其中设置 `NUM_CLASSES` 参数为您的类别数量加 1(因为需要加入一个背景类别): ```python class CustomConfig(Config): # Give the configuration a recognizable name NAME = "custom_dataset" # Train on 1 GPU and 1 image per GPU GPU_COUNT = 1 IMAGES_PER_GPU = 1 # Number of classes (including background) NUM_CLASSES = 1 + 4 # Background + 4 classes ``` 3. 定义数据集加载器 您需要定义一个数据集加载器,将准备好的数据集导入到模型中。您可以创建一个名为 `CustomDataset` 的类,继承自 `utils.Dataset` 类,并在其中实现 `load_dataset()`、`load_image()`、`load_mask()`、`image_reference()` 等方法,具体实现可以参考 Mask R-CNN 官方代码。 ```python class CustomDataset(utils.Dataset): def load_dataset(self, dataset_dir, subset): self.add_class("custom_dataset", 1, "class1") self.add_class("custom_dataset", 2, "class2") self.add_class("custom_dataset", 3, "class3") self.add_class("custom_dataset", 4, "class4") # Load annotations annotations = json.load(open(os.path.join(dataset_dir, "annotations.json"))) annotations = annotations["annotations"] # Add images and annotations to dataset for a in annotations: image_id = a["image_id"] image_path = os.path.join(dataset_dir, "images", str(image_id) + ".jpg") if not os.path.exists(image_path): continue if a["iscrowd"]: continue if a["category_id"] not in [1, 2, 3, 4]: continue self.add_image( "custom_dataset", image_id=image_id, path=image_path, width=a["width"], height=a["height"], annotations=a["bbox"] ) def load_mask(self, image_id): # Load annotations for image annotations = self.image_info[image_id]["annotations"] # Create one mask per instance masks = np.zeros([self.image_info[image_id]["height"], self.image_info[image_id]["width"], len(annotations)], dtype=np.uint8) # Load masks for i, a in enumerate(annotations): x1, y1, w, h = a x2 = x1 + w y2 = y1 + h masks[y1:y2, x1:x2, i] = 1 # Return masks and class IDs return masks, np.ones([len(annotations)], dtype=np.int32) def image_reference(self, image_id): info = self.image_info[image_id] return info["path"] ``` 4. 训练模型 在训练之前,您需要将预训练 COCO 权重加载到模型中: ```python model.load_weights(COCO_MODEL_PATH, by_name=True, exclude=["mrcnn_class_logits", "mrcnn_bbox_fc", "mrcnn_bbox", "mrcnn_mask"]) ``` 然后,您可以使用 `train()` 方法训练模型。在训练之前,您需要创建一个名为 `CustomConfig` 的配置对象,并设置好超参数和文件路径: ```python config = CustomConfig() config.display() model = modellib.MaskRCNN(mode="training", config=config, model_dir=MODEL_DIR) # Train the head branches model.train(dataset_train, dataset_val, learning_rate=config.LEARNING_RATE, epochs=30, layers='heads') ``` 5. 测试模型 在测试模型之前,您需要将模型切换到 inference 模式: ```python model = modellib.MaskRCNN(mode="inference", config=config, model_dir=MODEL_DIR) ``` 然后,您可以使用 `detect()` 方法对图片进行检测和分割: ```python results = model.detect([image], verbose=1) r = results[0] visualize.display_instances(image, r['rois'], r['masks'], r['class_ids'], class_names, r['scores']) ``` 以上就是使用 Mask R-CNN 训练自己数据集的基本步骤。具体实现可以参考 Mask R-CNN 官方代码。
评论 5
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包

打赏作者

张十三、

你的鼓励将是我创作的最大动力

¥1 ¥2 ¥4 ¥6 ¥10 ¥20
扫码支付:¥1
获取中
扫码支付

您的余额不足,请更换扫码支付或充值

打赏作者

实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值