使用Mask_RCNN进行预测

今天我们来说一下,如何使用自己训练出来的Mask_RCNN模型,或是官方自己的模型权重来进行预测:

该Mask_RCNN版本基于:Python3,Keras,TensorFlow,我使用的具体版本为:

  • Python 3.6.3
  • TensorFlow 1.7
  • Keras 2.1.5

tensorflow安装:

https://blog.csdn.net/qq_15969343/article/details/79971469

Mask-RCNN :

https://github.com/matterport/Mask_RCNN


1.加载模块以及载入模型权重:

import os
import sys
import random
import math
import numpy as np
import skimage.io
import matplotlib
import matplotlib.pyplot as plt

import coco
import utils
import model as modellib
import visualize
from config import Config

# Root directory of the project
ROOT_DIR = os.getcwd()

# Directory to save logs and trained model
MODEL_DIR = os.path.join(ROOT_DIR, "logs")

# Local path to trained weights file
MODEL_WEIGHT_PATH = os.path.join(ROOT_DIR, "mask_rcnn_val_all.h5")  #这里输入模型权重的路径
# Download COCO trained weights from Releases if needed
if not os.path.exists(MODEL_WEIGHT_PATH):
    utils.download_trained_weights(MODEL_WEIGHT_PATH)

# Directory of images to run detection on
IMAGE_DIR = os.path.join(ROOT_DIR, "samples\\balloon\\datasets\\balloon\\val2")   #这是输入你要预测图片的路径


2.根据自己训练时候的配置,从Config类中继承创建一个新类:

例如,我要检测的物体有5种:'person', 'book', 'chair', 'bottle', 'sports bag'再加上背景则一共是6种:

class somethingConfig(Config):
    """Configuration for training on MS COCO.
    Derives from the base Config class and overrides values specific
    to the COCO dataset.
    """
    # Give the configuration a recognizable name
    NAME = "something"

    # We use a GPU with 12GB memory, which can fit two images.
    # Adjust down if you use a smaller GPU.
    IMAGES_PER_GPU = 2

    # Uncomment to train on 8 GPUs (default is 1)
    # GPU_COUNT = 8

    # Number of classes (including background)
    NUM_CLASSES = 1 + 5  # COCO has 80 classes


3.根据上面的类,新建一个专门用于预测的类:

class InferenceConfig(somethingConfig):
    # 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()


4.载入图像,进行预测:

# 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(MODEL_WEIGHT_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', 'person', 'book', 'chair', 'bottle', 'sports bag']

# Load a random image from the images folder
file_names = next(os.walk(IMAGE_DIR))[2]
for x in range(len(file_names)):
    image = skimage.io.imread(os.path.join(IMAGE_DIR, file_names[x]))

# Run detection
    results = model.detect([image], verbose=1)

# Visualize results
    r = results[0]
    visualize.display_instances(image, r['rois'], r['masks'], r['class_ids'],
                            class_names, r['scores'])


5.预测结果:


猜你喜欢

转载自blog.csdn.net/qq_15969343/article/details/80388311
今日推荐