mask-rcnn训练完自己的数据集之后的测试demo

最新推荐文章于 2025-05-26 17:07:00 发布

ZealCV

最新推荐文章于 2025-05-26 17:07:00 发布

阅读量4.3k

点赞数

CC 4.0 BY-SA版权

分类专栏：深度学习文章标签：深度学习图像语义分割

本文链接：https://blog.youkuaiyun.com/u011574296/article/details/79903297

深度学习专栏收录该内容

19 篇文章

订阅专栏

本文介绍如何使用 Mask R-CNN 模型进行实物检测。通过加载预训练权重，随机选择一张图片进行实例分割，并展示检测结果。文中涉及模型配置、数据路径设置、图像读取及显示等关键步骤。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

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

class_names = ['BG', 'your class1', 'your class2', 'your class3']

class InferenceConfig(coco.CocoConfig):
    # 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

    NUM_CLASSES = 1 + 3  # background + 3 class

config = InferenceConfig()
config.display()

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

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

# Directory of images to run detection on
IMAGE_DIR = os.path.join(ROOT_DIR, "images")

# 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)))

# Create model object in inference mode.
model = modellib.MaskRCNN(mode="inference", config=config , model_dir=MODEL_DIR)

# Local path to trained weights file
COCO_MODEL_PATH = os.path.join(ROOT_DIR, "mask_rcnn_shapes_0002.h5")
# Load weights trained on MS-COCO
model.load_weights(COCO_MODEL_PATH, by_name=True)


# 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'])