项目地址:https://github.com/matterport/Mask_RCNN
文件目录
预测
-
mrcnn:这是保存项目 Python 代码的核心目录。
__init__.py
:将mrcnn文件夹标记为 Python 库。
model.py
:具有用于构建层和模型的函数和类。
config.py
:保存一个名为Config
的类,该类包含有关模型的一些配置参数。
utils.py
:包括一些辅助函数和类。
visualize.py
:可视化模型的结果。
parallel_model.py
:用于支持多个 GPU
-
samples:Jupyter 笔记本提供了一些使用项目的示例。
-
images:测试图像的集合。
-
asserts:一些带注释的图片
逻辑
要构建 Mask R-CNN 模型,必须指定几个参数。这些参数控制非最大抑制 (NMS)、交集并集 (IoU)、图像大小、每个图像的 ROI 数量、ROI 池层等
-
准备模型配置参数
-
构建掩码 R-CNN 模型架构
-
加载模型权重
-
读取输入图像
-
检测对象
-
可视化结果
-
预测代码
-
import mrcnn import mrcnn.config #配置文件 import mrcnn.model import mrcnn.visualize import cv2 import os CLASS_NAMES = ['BG', 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light', 'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard', 'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch', 'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone', 'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush'] class SimpleConfig(mrcnn.config.Config): NAME = "coco_inference" GPU_COUNT = 1 IMAGES_PER_GPU = 1 # BATCH_SIZE = IMAGES_PER_GPU * GPU_COUNT 批量大小 NUM_CLASSES = len(CLASS_NAMES) #背景是一类 model = mrcnn.model.MaskRCNN(mode="inference", config=SimpleConfig(), model_dir=os.getcwd()) # 用于保存训练日志和训练权重的目录 # model.keras_model.summary() 可以打印模型的摘要参数量等 model.load_weights(filepath="mask_rcnn_coco.h5", by_name=True) # 如果为 True,则根据其名称为每个图层分配权重 image = cv2.imread("sample2.jpg") image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) r = model.detect([image], verbose=0) # 设置为 1 时,将打印一些日志消息 r = r[0] # 对于每个输入图像,该方法返回一个字典,其中包含有关检测到的对象的信息。要返回有关馈送到模型的第一张图像的信息,则将索引与变量一起使用 # print(r.keys()) # 可视化 mrcnn.visualize.display_instances(image=image, boxes=r['rois'], masks=r['masks'], class_ids=r['class_ids'], class_names=CLASS_NAMES, scores=r['scores'])
训练
- 数据集 mrcnn.utils.Dataset
images:数据集中的图像
annots:每个图像的注释作为单独的 XML 文件
add_class()
:添加新类。add_image()
:将新图像添加到数据集。image_reference()
:检索图像的引用(例如路径或链接)。prepare()
:将所有类和图像添加到数据集后,此方法将准备数据集以供使用。source_image_link()
:返回图像的路径或链接。load_image()
:读取并返回图像。load_mask()
:加载图像中对象的蒙版。
在前面列出的所有方法中,必须重写该方法。原因是检索对象的蒙版因注释文件格式而异,因此没有加载蒙版的单一方法。因此,加载掩码是开发人员必须完成的任务。load_mask()
-
训练代码
import os
import xml.etree
from numpy import zeros, asarray
import mrcnn.utils
import mrcnn.config
import mrcnn.model
class KangarooDataset(mrcnn.utils.Dataset):
def load_dataset(self, dataset_dir, is_train=True):
self.add_class("dataset", 1, "kangaroo")
images_dir = dataset_dir + '/images/'
annotations_dir = dataset_dir + '/annots/'
for filename in os.listdir(images_dir):
image_id = filename[:-4]
if image_id in ['00090']:
continue
if is_train and int(image_id) >= 150:
continue
if not is_train and int(image_id) < 150:
continue
img_path = images_dir + filename
ann_path = annotations_dir + image_id + '.xml'
self.add_image('dataset', image_id=image_id, path=img_path, annotation=ann_path)
def extract_boxes(self, filename): # 返回每个边界框的坐标以及每个图像的宽度和高度
tree = xml.etree.ElementTree.parse(filename)
root = tree.getroot()
boxes = list()
for box in root.findall('.//bndbox'):
xmin = int(box.find('xmin').text)
ymin = int(box.find('ymin').text)
xmax = int(box.find('xmax').text)
ymax = int(box.find('ymax').text)
coors = [xmin, ymin, xmax, ymax]
boxes.append(coors)
width = int(root.find('.//size/width').text)
height = int(root.find('.//size/height').text)
return boxes, width, height
def load_mask(self, image_id):
info = self.image_info[image_id]
path = info['annotation']
boxes, w, h = self.extract_boxes(path)
masks = zeros([h, w, len(boxes)], dtype='uint8')
class_ids = list()
for i in range(len(boxes)):
box = boxes[i]
row_s, row_e = box[1], box[3]
col_s, col_e = box[0], box[2]
masks[row_s:row_e, col_s:col_e, i] = 1
class_ids.append(self.class_names.index('kangaroo'))
return masks, asarray(class_ids, dtype='int32')
class KangarooConfig(mrcnn.config.Config):
NAME = "kangaroo_cfg"
GPU_COUNT = 1
IMAGES_PER_GPU = 1
NUM_CLASSES = 2
STEPS_PER_EPOCH = 131
train_set = KangarooDataset()
train_set.load_dataset(dataset_dir='kangaroo', is_train=True)
train_set.prepare()
valid_dataset = KangarooDataset()
valid_dataset.load_dataset(dataset_dir='kangaroo', is_train=False)
valid_dataset.prepare()
kangaroo_config = KangarooConfig()
model = mrcnn.model.MaskRCNN(mode='training',
model_dir='./',
config=kangaroo_config)
model.load_weights(filepath='mask_rcnn_coco.h5',
by_name=True,
exclude=["mrcnn_class_logits", "mrcnn_bbox_fc", "mrcnn_bbox", "mrcnn_mask"])
# 不加载其权重的图层的名称。这些是架构顶部的层,它们根据问题类型(例如类的数量)而变化,负责生成类概率、边界框和掩码的层
model.train(train_dataset=train_set,
val_dataset=valid_dataset,
learning_rate=kangaroo_config.LEARNING_RATE,
epochs=1,
layers='heads')
model_path = 'Kangaro_mask_rcnn.h5'
model.keras_model.save_weights(model_path)