mask rcnn训练自己的数据集

最新推荐文章于 2024-09-14 21:01:48 发布

置顶

Tom Hardy

最新推荐文章于 2024-09-14 21:01:48 发布

阅读量10w+

点赞数 94

分类专栏：深度学习

本文链接：https://blog.csdn.net/qq_29462849/article/details/81037343

版权

本文介绍如何基于Mask R-CNN训练自己的数据集，包括数据准备、源代码使用、训练参数设置和训练过程中遇到的问题。参考了开源项目、图像标记工具LabelMe，以及训练数据转换工具。训练数据集包含图像、JSON文件和掩码标签，训练代码需配合pycocotools库。文章还提供了训练和测试模型的代码，以及模型的保存和使用方式。

摘要由CSDN通过智能技术生成

前言

最近迷上了mask rcnn，也是由于自己工作需要吧，特意研究了其源代码，并基于自己的数据进行训练~
本博客参考https://blog.csdn.net/disiwei1012/article/details/79928679#commentsedit

实验目的

这里写图片描述

哎~说多了都是泪，谁让我是工科生呢？只能检测工件了。。。做不了高大上的东西了，哈哈

主要参考及工具

基于Mask RCNN开源项目：https://github.com/matterport/Mask_RCNN
图片标记工具基于开源项目：https://github.com/wkentaro/labelme
训练工具：win10+GTX1060+cuda9.1+cudnn7+tensorflow-gpu-1.6.0+keras-2.1.6，140幅图像，一共3类，1小时左右

有关labelme的使用可以参考：https://blog.csdn.net/shwan_ma/article/details/77823281

有关mask-rcnn和Faster RCNN算法可以参考：
https://blog.csdn.net/linolzhang/article/details/71774168
https://blog.csdn.net/lk123400/article/details/54343550/

准备训练数据集

这是我简历的四个文件夹，下面一一道来~
这里写图片描述
1.pic

这是训练的图像，一共700幅

2.json
这里写图片描述
这是通过labelme处理训练图像后生成的文件

3.labelme_json
这里写图片描述

这个是处理.json文件后产生的数据，使用方法为labelme_json_to_dataset+空格+文件名称.json，这个前提是labelme要准确安装并激活。但是这样会产生一个问题，对多幅图像这样处理，太麻烦，在这里提供一个工具，可以直接在.json文件目录下转换所有的json文件，链接：json数据转换工具。

4.cv2_mask文件

由于labelme生成的掩码标签 label.png为16位存储，opencv默认读取8位，需要将16位转8位，可通过C++程序转化，代码请参考这篇博文：http://blog.csdn.net/l297969586/article/details/79154150
这里写图片描述
一团黑，不过不要怕，正常的~

源代码

运行该代码，需要安装pycocotools，在windows下安装该工具非常烦，有的可以轻松的安装成功，有的重装系统也很难成功，哎，都是坑~~关于Windows下安装pycocotools请参考：https://blog.csdn.net/chixia1785/article/details/80040172，https://blog.csdn.net/gxiaoyaya/article/details/78363391

测试的源代码

Github上开源的代码，是基于ipynb的，我直接把它转换成.py文件，首先做个测试，基于coco数据集上训练好的模型，可以调用摄像头~~~

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
# Root directory of the project
ROOT_DIR = os.path.abspath("../")

# 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
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 = os.path.join(MODEL_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)
    print("cuiwei***********************")

# Directory of images to run detection on
IMAGE_DIR = os.path.join(ROOT_DIR, "images")

class InferenceConfig(coco.CocoConfig):
    # 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()
config.display()


# Create model object in inference mode.
model = modellib.MaskRCNN(mode="inference", model_dir=MODEL_DIR, config=config)

# Load w