利用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'])
相关推荐
落樱弥城6 分钟前
角点特征:从传统算法到深度学习算法演进
人工智能·深度学习·算法
StarRocks_labs17 分钟前
StarRocks MCP Server 开源发布:为 AI 应用提供强大分析中枢
数据库·starrocks·人工智能·开源·olap·mcp
珂朵莉MM1 小时前
2024 睿抗机器人开发者大赛CAIP-编程技能赛-专科组(国赛)解题报告 | 珂学家
开发语言·人工智能·算法·leetcode·职场和发展·深度优先·图论
ㄣ知冷煖★1 小时前
【论文阅读】A Survey on Multimodal Large Language Models
人工智能·语言模型·自然语言处理
视觉语言导航2 小时前
兼顾长、短视频任务的无人机具身理解!AirVista-II:面向动态场景语义理解的无人机具身智能体系统
人工智能·无人机·具身智能
墨绿色的摆渡人2 小时前
pytorch小记(二十二):全面解读 PyTorch 的 `torch.cumprod`——累积乘积详解与实战示例
人工智能·pytorch·python
moonsims2 小时前
低空态势感知:基于AI的DAA技术是低空飞行的重要安全保障-机载端&地面端
人工智能·安全
若叶时代2 小时前
数据分析_Python
人工智能·python·数据分析
虾球xz2 小时前
游戏引擎学习第286天:开始解耦实体行为
c++·人工智能·学习·游戏引擎
武子康2 小时前
大语言模型 11 - 从0开始训练GPT 0.25B参数量 MiniMind2 准备数据与训练模型 DPO直接偏好优化
人工智能·gpt·ai·语言模型·自然语言处理