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 shutilfor 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_nameshutil.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 specificto the toy shapes dataset."""# Give the configuration a recognizable nameNAME = "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 = 1IMAGES_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 = 320IMAGE_MAX_DIM = 384# Use smaller anchors because our image and objects are smallRPN_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 simpleSTEPS_PER_EPOCH = 100# use small validation steps since the epoch is smallVALIDATION_STEPS = 50config = 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(), Loader=yaml.FullLoader)labels = temp['label_names']del labels[0]return labels# 重新写draw_maskdef draw_mask(self, num_obj, mask, image,image_id):info = self.image_info[image_id]for index in range(num_obj):for i in range(info['width']):for j in range(info['height']):at_pixel = image.getpixel((i, j))if at_pixel == index + 1:mask[j, i, index] = 1return mask# 重新写load_shapes，里面包含自己的自己的类别# 并在self.image_info信息中添加了path、mask_path 、yaml_path# yaml_pathdataset_root_path = "/tongue_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, "tank") # 黑色素瘤#self.add_class("shapes", 1, "danhe") self.add_class("shapes", 2, "linba")  for i in range(count):# 获取图片宽和高filestr = imglist[i].split(".")[0]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")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_maskdef load_mask(self, image_id):"""Generate instance masks for shapes of the given image ID."""global iter_numprint("image_id",image_id)info = self.image_info[image_id]count = 1  # number of objectimg = 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] * occlusionocclusion = 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("danhe") != -1:labels_form.append("danhe")if labels[i].find("linba") != -1:labels_form.append("linba")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 inall visualizations in the notebook. Provide acentral point to control graph sizes.Change the default size attribute to control the sizeof rendered images"""_, ax = plt.subplots(rows, cols, figsize=(size * cols, size * rows))return ax#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(7, 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 lastif 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 weightsmodel.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 trainingmodel.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')# 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")

代码有3处需要修改的地方

四、预测数据

1. 训练完成后，会在根目录下生成logs文件夹，选择最后生成的训练模型

2. 根目录下创建预测脚本 model_test.py

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 import colorsys
from skimage.measure import find_contoursfrom PIL import Image# 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")
MODEL_WEIGHT = './logs/shapes20241122T2251/mask_rcnn_shapes_0010.h5'# 预测结果保存地址
OUTPUT_DIR = "./output_results"
if not os.path.exists(OUTPUT_DIR):os.makedirs(OUTPUT_DIR)# 获取待预测文件名称，用于保存同名文件
def get_last_part_of_string(path):return os.path.basename(path)# 需要预测的数据地址
IMAGE_DIR = os.path.join(ROOT_DIR, "datasets/demo2/pic")class ShapesConfig(Config):"""Configuration for training on the toy shapes dataset.Derives from the base Config class and overrides values specificto the toy shapes dataset."""# Give the configuration a recognizable nameNAME = "shapes"# os.path.exists()# 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 = 1IMAGES_PER_GPU = 1# Number of classes (including background)NUM_CLASSES = 1 + 2  # 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 = 320IMAGE_MAX_DIM = 384# Use smaller anchors because our image and objects are smallRPN_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 simpleSTEPS_PER_EPOCH = 100# use small validation steps since the epoch is smallVALIDATION_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_GPUGPU_COUNT = 1IMAGES_PER_GPU = 1# 以下def函数从 mrcnn/visualize.py 中复制过来def apply_mask(image, mask, color, alpha=0.5):"""打上mask图标"""for c in range(3):image[:, :, c] = np.where(mask == 1,image[:, :, c] *(1 - alpha) + alpha * color[c] * 255,image[:, :, c])return imagedef random_colors(N, bright=True):"""生成随机颜色"""brightness = 1.0 if bright else 0.7hsv = [(i / N, 1, brightness) for i in range(N)]colors = list(map(lambda c: colorsys.hsv_to_rgb(*c), hsv))return colors# 绘制结果
def display_instances(image, boxes, masks, class_ids, class_names, scores=None, filename="", title="",figsize=(16, 16),show_mask=True, show_bbox=True,colors=None, captions=None):# instance的数量N = boxes.shape[0]if not N:print("\n*** No instances to display *** \n")else:assert boxes.shape[0] == masks.shape[-1] == class_ids.shape[0]colors = colors or random_colors(N)# 当masked_image为原图时是在原图上绘制# 如果不想在原图上绘制，可以把masked_image设置成等大小的全0矩阵masked_image = np.array(image, np.uint8)for i in range(N):color = colors[i]# 该部分用于显示bboxif not np.any(boxes[i]):continuey1, x1, y2, x2 = boxes[i]# 绘制方形边界框if show_bbox:# 打印边界框坐标# print(f"Box {i + 1}: (x1, y1, x2, y2) = ({x1}, {y1}, {x2}, {y2})")cropped_image_rgb = image[y1:y2, x1:x2]image_bgr = cv2.cvtColor(cropped_image_rgb, cv2.COLOR_RGB2BGR)# cv2.rectangle(masked_image, (x1, y1), (x2, y2), (color[0] * 255, color[1] * 255, color[2] * 255), 2)# 该部分用于显示文字与置信度if not captions:class_id = class_ids[i]score = scores[i] if scores is not None else Nonelabel = class_names[class_id]caption = "{} {:.3f}".format(label, score) if score else labelelse:caption = captions[i]# 绘制文字（类别、分数）font = cv2.FONT_HERSHEY_SIMPLEXcv2.putText(masked_image, caption, (x1, y1 + 8), font, 0.5, (255, 255, 255), 1)# 绘制语义分割（遮挡物体面积部分）mask = masks[:, :, i]if show_mask:masked_image = apply_mask(masked_image, mask, color)# 画出语义分割的范围padded_mask = np.zeros((mask.shape[0] + 2, mask.shape[1] + 2), dtype=np.uint8)padded_mask[1:-1, 1:-1] = maskcontours = find_contours(padded_mask, 0.5)# 绘制边缘轮廓for verts in contours:verts = np.fliplr(verts) - 1cv2.polylines(masked_image, [np.array([verts], np.int)], 1,(color[0] * 255, color[1] * 255, color[2] * 255), 1)# 将绘制好的图片保存在制定文件夹中save_path = os.path.join(OUTPUT_DIR, filename)cv2.imwrite(save_path, masked_image)img = Image.fromarray(np.uint8(masked_image))return img# 加载配置
config = InferenceConfig()# 加载模型
model = modellib.MaskRCNN(mode="inference", model_dir=MODEL_DIR, config=config)#  加载权重
model.load_weights(MODEL_WEIGHT, 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', 'danhe', 'linba']
# Load a random image from the images folder############################## 单张预测图片 ################################
# file_names = next(os.walk(IMAGE_DIR))[2]
# # 随机读取一张图片，进行预测
# image = skimage.io.imread(os.path.join(IMAGE_DIR, random.choice(file_names)))
# # 当前图片名称
# filename = random.choice(file_names)# a=datetime.now() 
# # Run detection
# results = model.detect([image], verbose=1)
# b=datetime.now() 
# # Visualize results
# print("time:",(b-a).seconds)
# r = results[0]
# display_instances(image, r['rois'], r['masks'], r['class_ids'], class_names, r['scores'], filename)# # 获取物体数量
# num_objects = 0
# for roi, class_id, score, mask in zip(r['rois'],  r['class_ids'], r['scores'], r['masks']):
#     if score > 0.5:  # 设置置信度阈值
#         num_objects += 1# print('Number of objects detected:', num_objects)############################## 批量预测图片 ################################
count = os.listdir(IMAGE_DIR)for i in range(0,len(count)):path = os.path.join(IMAGE_DIR, count[i])if os.path.isfile(path):file_names = next(os.walk(IMAGE_DIR))[2]image = skimage.io.imread(os.path.join(IMAGE_DIR, count[i]))# Run detectionresults = model.detect([image], verbose=1)r = results[0]print('FileNmame: ',count[i])display_instances(image, r['rois'], r['masks'], r['class_ids'], class_names, r['scores'], count[i])