利用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'])
相关推荐
最贪吃的虎3 分钟前
MIT新论文:Hyperloop Transformers
人工智能·python·语言模型·langchain
懂懂tty7 分钟前
详解Cursor 中 Skill、Spec
人工智能·ai编程
碧海银沙音频科技研究院13 分钟前
BES2800BP_nuttx编译环境搭建方法
人工智能·深度学习·算法
墨染天姬18 分钟前
cursor的MCP怎么配置使用?
人工智能
colofullove18 分钟前
幂等设计与数据一致性:确保小说知识库的可靠更新
人工智能
庄小焱23 分钟前
【AI模型】——基于知识图谱的RAG
人工智能·大模型·知识图谱·rag·ai模型·ai系统
❆VE❆24 分钟前
python实战(一):对接AI大模型并应用
开发语言·人工智能·python·ai
格林威30 分钟前
堡盟Baumer VCX系列工业相机供电与触发:网口(GigE) vs USB3.0
开发语言·人工智能·数码相机·计算机视觉·视觉检测·工业相机·高速相机
三毛的二哥30 分钟前
BEV:感知抖动问题及解决办法
人工智能·算法·计算机视觉
光泽雨31 分钟前
VM图像处理(1、图像二值化和图像滤波,Sobel提取过程)
图像处理·人工智能