利用maskrcnn来实现目标检测与追踪

首先下载源代码仓库,链接地址如下:

maskrcnn

能够实现的效果如图所示:

该存储库包括:

  • 基于FPN和ResNet101构建的Mask R-CNN的源代码。
  • MS COCO 的训练代码
  • MS COCO 的预训练砝码
  • Jupyter 笔记本,用于可视化每一步的检测管道
  • 用于多 GPU 训练的并行模型类
  • 对 MS COCO 指标 (AP) 的评估
  • 在自己的数据集上进行训练的示例

下载代码仓库,进行解压后的目录如下:

可以使用下面:

pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple

也可以使用

python setup.py install

来安装相关的依赖包,安装完成后,还需要下载模型文件,

下载链接地址如下:

mask_rcnn_balloon.h5

测试代码如下所示:

python 复制代码
import os
import sys
import random
import math
import numpy as np
import skimage.io
import matplotlib
import matplotlib.pyplot as plt

# Root directory of the project
ROOT_DIR = os.path.abspath("../")

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

%matplotlib inline 

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

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

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

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

# 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'])
相关推荐
Luis Li 的猫猫2 小时前
深度学习中的知识蒸馏
人工智能·经验分享·深度学习·学习·算法
木觞清4 小时前
PyTorch与TensorFlow的对比:哪个框架更适合你的项目?
人工智能·pytorch·tensorflow
wyg_0311137 小时前
用deepseek学大模型04-模型可视化与数据可视化
人工智能·机器学习·信息可视化
陈敬雷-充电了么-CEO兼CTO8 小时前
DeepSeek核心算法解析:如何打造比肩ChatGPT的国产大模型
人工智能·神经网络·自然语言处理·chatgpt·大模型·aigc·deepseek
南风过闲庭9 小时前
人工智能泡沫效应
大数据·人工智能·科技·搜索引擎·百度·ai
我是一个对称矩阵9 小时前
YOLOv5-Seg 深度解析:与 YOLOv5 检测模型的区别
人工智能·yolo·目标跟踪
AomanHao9 小时前
图像质量评价指标-UCIQE-UIQM
图像处理·人工智能·计算机视觉·评价指标
MYT_flyflyfly9 小时前
计算机视觉-尺度不变区域
人工智能·计算机视觉
何小Ai同学9 小时前
Deepseek赚钱密码:小场景闭环如何让你快速盈利?
人工智能·架构·deepseek
AI服务老曹9 小时前
通过感知、分析、预测、控制,最大限度发挥效率的智慧油站开源了
人工智能·开源·自动化·音视频