MaskRCNN训练自己的数据集

一. 数据标注

1. 安装labelme软件

conda install labelme

2. 运行labelme

# 命令行中直接输入
labelme

3. 标注

二、训练数据处理

1. 在根目录下创建datasets/demo1文件夹，创建如下几个文件夹

2. 将标注好的.json文件放在json文件夹中

3. 原图放在pic文件夹中

4. 批量处理JSON文件

现在的.json包含了我们标注的信息，但是还不是可以让代码直接读取的数据格式，对于不同的深度学习代码，数据存放的格式要根据不同代码的写法来定，下面以上述的mask-rcnn为例，将数据修改为我们需要的格式。

在json文件夹中创建test.bat脚本

@echo off
for %%i in (*.json) do labelme_json_to_dataset "%%i"
pause

在json文件夹中运行（要激活maskrcnn虚拟环境）

批处理生成的文件夹

5. 将批处理生成的文件夹复制到labelme_json文件夹中

6. 将labelme_json文件夹中的label.png复制到cv2_mask中

单个复制太麻烦，我们写脚本批处理， 创建copy.py

import os 
import shutil

for dir_name in os.listdir('./labelme_json'):
    pic_name = dir_name[:-5] + '.png'
    from_dir = './labelme_json/'+dir_name+'./label.png'
    to_dir = './cv2_mask/'+pic_name
    shutil.copyfile(from_dir, to_dir)
        
    print (from_dir)
    print (to_dir)

三、训练数据

根目录下创建train_test.py

# -*- coding: utf-8 -*-
'''
Author: qingqingdan 1306047688@qq.com
Date: 2024-11-22 17:48:33
LastEditTime: 2024-11-29 12:16:01
Description: labelme标注图片 训练自己的数据集
'''

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("../")
# 训练后的模型
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)


#基础设置
dataset_root_path="datasets/demo2/"
img_floder = dataset_root_path + "pic"
mask_floder = dataset_root_path + "cv2_mask"
imglist = os.listdir(img_floder)
count = len(imglist)



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 = 1 + 2  # background + 2 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 = 100

    # 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