模型的检测过程:
(1)自定义配置信息
(2)准备测试集test
(3)创建inference模式下的模型
(4)调用model.detect进行检测
(5)可视化,或者计算评价指标(AP等)
代码如下:(与训练过程的代码重复的部分,不再注释)
# -*- coding: utf-8 -*-
"""
Created on Tue Dec 17 15:01:49 2019
@author: Lenovo
"""
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
from PIL import Image
import yaml
import skimage.io
ROOT_DIR = os.path.abspath("../../")
sys.path.append(ROOT_DIR)
from mrcnn.config import Config
from mrcnn import utils
import mrcnn.model as modellib
from mrcnn import visualize
from mrcnn.model import log
################################################################################
MODEL_DIR = os.path.join(ROOT_DIR, "logs")
COCO_MODEL_PATH = os.path.join(ROOT_DIR, "mask_rcnn_coco.h5")
if not os.path.exists(COCO_MODEL_PATH):
utils.download_trained_weights(COCO_MODEL_PATH)
################################################################################
class HooksConfig(Config):
"""
在shapes数据集上进行配置修改,即创建继承于基类Config的子类并且修改对应的属性值即可
"""
NAME = "hooks"
# 当使用GPU同时测试多张图片时,输入的图片数量必须=BatchSize=GPU_COUNT*IMAGES_PER_GPU
GPU_COUNT = 1
IMAGES_PER_GPU = 2
NUM_CLASSES = 1 + 1 # background + hooks
IMAGE_MIN_DIM = 128*8
IMAGE_MAX_DIM = 128*8
RPN_ANCHOR_SCALES = (8*8, 16*8, 32*8, 64*8, 128*8)
TRAIN_ROIS_PER_IMAGE = 32
STEPS_PER_EPOCH = 100
VALIDATION_STEPS = 5
config = HooksConfig()
config.display()
################################################################################
def get_ax(rows=1, cols=1, size=8):
_, ax = plt.subplots(rows, cols, figsize=(size*cols, size*rows))
return ax
################################################################################
class HookDataset(utils.Dataset):
def get_obj_index(self, image):
n = np.max(image)
return n
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']
#labels = list(labels.keys())
del labels[0]
return labels
def draw_mask(self, num_obj, mask, image):
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] =1
return mask
def load_shapes(self, count, height, width, img_floder, mask_floder