Labelme标注数据，使用Mask RCNN训练自己的数据集
发布时间
阅读量：
阅读量
代码：https://github.com/matterport/Mask_RCNN
训练环境：Ubuntu系统
一、数据准备
准备好数据图像，使用labelme进行多边形标注
1. 安装labelme
必须要有深度学习环境（安装annaconda）
打开cmd激活base环境（activate base）
使用命令 pip install labelme安装labelme工具2. 安装完成后在命令行输入labelme回车即可打开labelme工具
通过双击启动目标图像所在的文件夹，并执行opendir操作后，请您依次右键点击并选择创建多边形
标注完成后一张图像会生成一个json文件
3. 将json文件生成文件夹，文件夹里面包含以下几个文件
当前必须依赖于labelme 3.16.5版本（通过命令pip install labelme==3.16.5进行安装），并构建相应的代码文件夹以便运行。否则将导致系统报错提示：模块labelme.utils中没有属性draw_label
import argparse
import base64
import json
import os
import os.path as osp
import warnings
import PIL.Image
import yaml
from labelme import utils
###############################################增加的语句,改下路径即可##############################
import glob
json_list = glob.glob(os.path.join(r'./00/','*.json'))
############################################### end ##################################
def main():
# warnings.warn("This script is aimed to demonstrate how to convert the\n"
# "JSON file to a single image dataset, and not to handle\n"
# "multiple JSON files to generate a real-use dataset.")
parser = argparse.ArgumentParser()
############################################### 删除的语句 ##################################
# parser.add_argument('json_file')
# json_file = args.json_file
############################################### end ##################################
parser.add_argument('-o', '--out', default=None)
args = parser.parse_args()
###############################################增加的语句##################################
for json_file in json_list:
############################################### end ##################################
if args.out is None:
out_dir = osp.basename(json_file).replace('.', '_')
out_dir = osp.join(osp.dirname(json_file), out_dir)
else:
out_dir = args.out
if not osp.exists(out_dir):
os.mkdir(out_dir)
data = json.load(open(json_file))
if data['imageData']:
imageData = data['imageData']
else:
imagePath = os.path.join(os.path.dirname(json_file), data['imagePath'])
with open(imagePath, 'rb') as f:
imageData = f.read()
imageData = base64.b64encode(imageData).decode('utf-8')
img = utils.img_b64_to_arr(imageData)
label_name_to_value = {'_background_': 0}
for shape in sorted(data['shapes'], key=lambda x: x['label']):
label_name = shape['label']
if label_name in label_name_to_value:
label_value = label_name_to_value[label_name]
else:
label_value = len(label_name_to_value)
label_name_to_value[label_name] = label_value
lbl = utils.shapes_to_label(img.shape, data['shapes'], label_name_to_value)
label_names = [None] * (max(label_name_to_value.values()) + 1)
for name, value in label_name_to_value.items():
label_names[value] = name
lbl_viz = utils.draw_label(lbl, img, label_names)
PIL.Image.fromarray(img).save(osp.join(out_dir, 'img.png'))
utils.lblsave(osp.join(out_dir, 'label.png'), lbl)
PIL.Image.fromarray(lbl_viz).save(osp.join(out_dir, 'label_viz.png'))
with open(osp.join(out_dir, 'label_names.txt'), 'w') as f:
for lbl_name in label_names:
f.write(lbl_name + '\n')
warnings.warn('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)
print('Saved to: %s' % out_dir)
if __name__ == '__main__':
main()设置json_list变量用于获取所有*.json文件路径
在CMD窗口启动LabelMe并将其环境设置为3.16.5版本
切换至json2img.py所在的文件夹
执行Python脚本 json2img.py即可自动生成所需文件夹
该文件夹内将包含上述提到的五个文件
在此处创建一个新的文件夹并命名为labelme_json之后将所有相关文件放置到同一个目录中然后在此处创建一个新的文件夹并命名为labelme_json之后将所有相关文件放置到同一个目录中
4. 移动文件
将所有json文件复制至对应的json文件夹中，
将所有图像复制至pic文件夹中，
剩余cv_mask需通过脚本代码完成复制并重新命名为相应名称（请查阅相关资料）以供后续处理。
其中label.png应为8-bit图，
可通过属性检查确认其数据类型。
如果该图片是16-bit则需转换为8-bit。
# -*- coding: utf-8 -*-
"""
Created on Mon Jul 27 13:47:29 2020
@author: 45556
"""
import os
import shutil
import sys
import re
spath=r'./train/labelme_json/'
tpath=r'./train/cv_mask/'
sfolder=os.listdir(spath)
i=1
for tempfolder in sfolder:
temppath=os.path.join(spath,tempfolder)
flist=os.listdir(temppath)
if not 'label.png' in flist:
print('not find label.png in {}'.format(temppath))
continue
# 开始复制文件，注意这里不改变名字
# step 1 创建新的文件夹
shutil.copy(os.path.join(temppath,'label.png'),os.path.join(tpath,'{}.png'.format(tempfolder.replace("_json",""))))到这里所有的数据都准备好了，按照同样的步骤可以准备好测试集。
二、训练
1. train.py
打开代码项目，在主目录下新建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
from mrcnn 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(r"E:\Github\Mask_RCNN-master")
# 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 = 1 + 2 # background + 1 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['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 = "/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, "b_sidewalk")
self.add_class("shapes", 2, "z_crossing")
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")
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("b_sidewalk") != -1:
# print "car"
labels_form.append("b_sidewalk")
elif labels[i].find("z_crossing") != -1:
# print "leg"
labels_form.append("z_crossing")
# elif labels[i].find("well") != -1:
# # print "well"
# labels_form.append("well")
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 = "datasets/"
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')
# 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")将mask_rcnn_coco.h5权重文件下载至指定目录
对第26行进行主目录设置
在第53行中, 类别总数等于背景数量与各类别数量之和
第122至123行分别定义了各个类别的名称
对于第161至166行而言, 在类别数量较多的情况下建议采用条件判断结构
在第186行为数据预处理阶段设置了训练时长参数
在第240行为数据预处理阶段设置了训练时长参数
在第249行为模型优化阶段设置了精细调整参数
2. config.py
在mrcnn/config.py文件中进行如下修改：第55行设置BACKBONE为‘resnet50’；同时支持使用ResNet-101架构。第73行的NUM_CLASSES设置目前不支持包含背景类别。对于其他参数，则建议进行相应的调整。
在Ubuntu系统中训练的环境配置如下：
# Name Version Build Channel
_libgcc_mutex 0.1 main
_tflow_select 2.1.0 gpu
absl-py 0.9.0 py36_0
astor 0.8.1 py36_0
backcall 0.2.0 pypi_0 pypi
blas 1.0 mkl
c-ares 1.15.0 h7b6447c_1001
ca-certificates 2020.6.24 0
certifi 2020.6.20 py36_0
cudatoolkit 10.1.243 h6bb024c_0
cudnn 7.6.5 cuda10.1_0
cupti 10.1.168 0
cycler 0.10.0 pypi_0 pypi
decorator 4.4.2 pypi_0 pypi
freetype 2.10.2 h5ab3b9f_0
gast 0.4.0 py_0
google-pasta 0.2.0 py_0
grpcio 1.27.2 py36hf8bcb03_0
h5py 2.10.0 py36hd6299e0_1
hdf5 1.10.6 hb1b8bf9_0
imageio 2.9.0 pypi_0 pypi
importlib-metadata 1.7.0 py36_0
intel-openmp 2020.1 217
ipython 7.16.1 pypi_0 pypi
ipython-genutils 0.2.0 pypi_0 pypi
jedi 0.17.2 pypi_0 pypi
jpeg 9b h024ee3a_2
keras 2.0.8 pypi_0 pypi
keras-applications 1.0.8 py_1
keras-preprocessing 1.1.0 py_1
kiwisolver 1.2.0 pypi_0 pypi
lcms2 2.11 h396b838_0
ld_impl_linux-64 2.33.1 h53a641e_7
libedit 3.1.20191231 h14c3975_1
libffi 3.3 he6710b0_2
libgcc-ng 9.1.0 hdf63c60_0
libgfortran-ng 7.3.0 hdf63c60_0
libpng 1.6.37 hbc83047_0
libprotobuf 3.12.4 hd408876_0
libstdcxx-ng 9.1.0 hdf63c60_0
libtiff 4.1.0 h2733197_1
lz4-c 1.9.2 he6710b0_1
markdown 3.2.2 py36_0
matplotlib 3.3.1 pypi_0 pypi
mkl 2020.1 217
mkl-service 2.3.0 py36he904b0f_0
mkl_fft 1.1.0 py36h23d657b_0
mkl_random 1.1.1 py36h0573a6f_0
ncurses 6.2 he6710b0_1
networkx 2.4 pypi_0 pypi
ninja 1.10.0 py36hfd86e86_0
numpy 1.19.1 py36hbc911f0_0
numpy-base 1.19.1 py36hfa32c7d_0
olefile 0.46 py36_0
opencv-python 4.4.0.42 pypi_0 pypi
openssl 1.1.1g h7b6447c_0
parso 0.7.1 pypi_0 pypi
pexpect 4.8.0 pypi_0 pypi
pickleshare 0.7.5 pypi_0 pypi
pillow 7.2.0 py36hb39fc2d_0
pip 20.2.2 py36_0
prompt-toolkit 3.0.6 pypi_0 pypi
protobuf 3.12.4 py36he6710b0_0
ptyprocess 0.6.0 pypi_0 pypi
pygments 2.6.1 pypi_0 pypi
pyparsing 2.4.7 pypi_0 pypi
python 3.6.10 h7579374_2
python-dateutil 2.8.1 pypi_0 pypi
pytorch 1.6.0 py3.6_cuda10.1.243_cudnn7.6.3_0 pytorch
pywavelets 1.1.1 pypi_0 pypi
pyyaml 5.3.1 pypi_0 pypi
readline 8.0 h7b6447c_0
scikit-image 0.17.2 pypi_0 pypi
scipy 1.5.0 py36h0b6359f_0
setuptools 49.6.0 py36_0
six 1.15.0 py_0
sqlite 3.32.3 h62c20be_0
tensorboard 1.14.0 py36hf484d3e_0
tensorflow 1.14.0 gpu_py36h3fb9ad6_0
tensorflow-base 1.14.0 gpu_py36he45bfe2_0
tensorflow-estimator 1.14.0 py_0
tensorflow-gpu 1.14.0 h0d30ee6_0
termcolor 1.1.0 py36_1
tifffile 2020.8.13 pypi_0 pypi
tk 8.6.10 hbc83047_0
torchvision 0.7.0 py36_cu101 pytorch
traitlets 4.3.3 pypi_0 pypi
wcwidth 0.2.5 pypi_0 pypi
werkzeug 1.0.1 py_0
wheel 0.34.2 py36_0
wrapt 1.12.1 py36h7b6447c_1
xz 5.2.5 h7b6447c_0
zipp 3.1.0 py_0
zlib 1.2.11 h7b6447c_3
zstd 1.4.5 h9ceee32_0 选择合适的CUDA 版本配置好环境即可。
3. 训练
在终端中激活安装完成后，在该环境中切换至train.py文件所在目录。通过执行python train.py指令即可启动训练流程。接下来系统首先会读取并解析必要的数据。
训练中：
生成的模型将在每隔一个 epoch 位于 log 目录下的位置生成一个新的模型文件，并尽量确保盘内存储空间得到充分释放。在进行训练时，请进入 log 目录以监控训练过程中的损失指标。
三、评估
在训练完成后生成的模型必须进行评估的具体操作流程如下：
# -*- 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
import tensorflow as tf
from keras import backend as K
config = tf.ConfigProto()
config.gpu_options.allow_growth=True
sess = tf.Session(config=config)
K.set_session(sess)
# Root directory of the project
ROOT_DIR = os.path.abspath(r"E:\Github\Mask_RCNN-master")
# ROOT_DIR = os.getcwd()
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 = 5
# 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 = 256
IMAGE_MAX_DIM = 512
# 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 = 50
# Use a small epoch since the data is simple
STEPS_PER_EPOCH = 50
# use small validation steps since the epoch is small
VALIDATION_STEPS = 20
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(),Loader=yaml.FullLoader)
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
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, "b_sidewalk") # 脆性区域
self.add_class("shapes", 2, "z_crossing")
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_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("b_sidewalk") != -1:
# print "box"
labels_form.append("b_sidewalk")
elif labels[i].find("z_crossing")!=-1:
#print "column"
labels_form.append("z_crossing")
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
def list2array(list):
b = np.array(list[0])
for i in range(1, len(list)):
b = np.append(b, list[i], axis=0)
return b
def text_save(filename, data): # filename为写入CSV文件的路径，data为要写入数据列表.
file = open(filename, 'a')
for i in range(len(data)):
s = str(data[i]).replace('[', '').replace(']', '') # 去除[],这两行按数据不同，可以选择
s = s.replace("'", '').replace(',', '') + '\n' # 去除单引号，逗号，每行末尾追加换行符
file.write(s)
file.close()
print("保存txt文件成功")
# 测试集设置
dataset_root_path = "test-dataset/"
img_floder = dataset_root_path + "pic"
mask_floder = dataset_root_path + "cv2_mask"
imglist = os.listdir(img_floder)
count = len(imglist)
# 准备test数据集
dataset_test = DrugDataset()
dataset_test.load_shapes(count, img_floder, mask_floder, imglist, dataset_root_path)
dataset_test.prepare()
# mAP
# Compute VOC-Style mAP @ IoU=0.5
# Running on 10 images. Increase for better accuracy.
class InferenceConfig(ShapesConfig):
GPU_COUNT = 1
IMAGES_PER_GPU = 1
inference_config = InferenceConfig()
# Recreate the model in inference mode
model = modellib.MaskRCNN(mode="inference",
config=inference_config,
model_dir=MODEL_DIR)
model_path = os.path.join(MODEL_DIR, "res-50-399.h5") # 修改成自己训练好的模型
# Load trained weights
print("Loading weights from ", model_path)
model.load_weights(model_path, by_name=True)
img_list = np.random.choice(dataset_test.image_ids, 85)
APs = []
count1 = 0
# 遍历测试集
for image_id in img_list:
# 加载测试集的ground truth
image, image_meta, gt_class_id, gt_bbox, gt_mask = \
modellib.load_image_gt(dataset_test, inference_config,
image_id, use_mini_mask=False)
# 将所有ground truth载入并保存
if count1 == 0:
save_box, save_class, save_mask = gt_bbox, gt_class_id, gt_mask
else:
save_box = np.concatenate((save_box, gt_bbox), axis=0)
save_class = np.concatenate((save_class, gt_class_id), axis=0)
save_mask = np.concatenate((save_mask, gt_mask), axis=2)
molded_images = np.expand_dims(modellib.mold_image(image, inference_config), 0)
# 启动检测
results = model.detect([image], verbose=0)
r = results[0]
# 将所有检测结果保存
if count1 == 0:
save_roi, save_id, save_score, save_m = r["rois"], r["class_ids"], r["scores"], r['masks']
else:
save_roi = np.concatenate((save_roi, r["rois"]), axis=0)
save_id = np.concatenate((save_id, r["class_ids"]), axis=0)
save_score = np.concatenate((save_score, r["scores"]), axis=0)
save_m = np.concatenate((save_m, r['masks']), axis=2)
count1 += 1
# 计算AP, precision, recall
AP, precisions, recalls, overlaps = \
utils.compute_ap(save_box, save_class, save_mask,
save_roi, save_id, save_score, save_m)
print("AP: ", AP)
print("mAP: ", np.mean(AP))
# 绘制PR曲线
plt.plot(recalls, precisions, 'b', label='PR')
plt.title('precision-recall curve')
plt.xlabel('Recall')
plt.ylabel('Precision')
plt.legend()
plt.show()
# 保存precision, recall信息用于后续绘制图像
text_save('Kpreci.txt', precisions)
text_save('Krecall.txt', recalls)需要对第29行的根目录进行修改
第56行为分类数目
从第154到第300行为别名数量
从第300到307行为别名数量
在第308行为别名列表中添加新项
在测试集中增加一条新条目以补充现有数据
在训练集中增加一条新条目以补充现有数据
在验证集中增加一条新条目以补充现有数据
在验证集中新增一个字段用于存储验证结果信息
运行后输出P-R曲线图，PR坐标对及AP、mAP值
四、检测demo
在主目录下建立demo.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
import tensorflow as tf
from datetime import datetime
config = tf.ConfigProto(gpu_options=tf.GPUOptions(allow_growth=True))
sess = tf.Session(config=config)
# Root directory of the project
ROOT_DIR = os.path.abspath(r"E:\Github\Mask_RCNN-master")
# 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 samples.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 ,"****.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 = 1 + 1 # background + your classes num
# 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
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()
config.display()
# 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', '*****']
# Load a random image from the images folder
# #视频检测：
# video_path = './video/2.mp4'
# cap = cv2.VideoCapture(video_path)
#
# while(1):
# # get a frame
# ret, frame = cap.read()
# # show a frame
# start =time.clock()
# results = model.detect([frame], verbose=1)
# r = results[0]
# #cv2.imshow("capture", frame)
# visualize.display_instances(frame, r['rois'], r['masks'], r['class_ids'],
# class_names, r['scores'])
# end = time.clock()
# print(end-start)
# if cv2.waitKey(1) & 0xFF == ord('q'):
# break
#
# cap.release()
# cv2.destroyAllWindows()
file_names = next(os.walk(IMAGE_DIR))[2]
#image = skimage.io.imread(os.path.join(IMAGE_DIR, random.choice(file_names)))#多图像检测images文件夹下
image= skimage.io.imread(r"E:\Github\Mask_RCNN-master\images\28.jpg")
# Run detection
results = model.detect([image], verbose=1)
# print(end-start)
# Visualize results
r = results[0]
visualize.display_instances(image, r['rois'], r['masks'], r['class_ids'],
class_names, r['scores'])更新第十九行目录位置
整个过程都可以运行，如有问题可留言
如有侵权请联系删除谢谢！转载请注明出处。
全部评论 (0)
还没有任何评论哟~