利用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'])
相关推荐
IT_陈寒3 小时前
闭包陷阱让我加了两天班,JavaScript你真行
前端·人工智能·后端
wumingxiaoyao3 小时前
从 0 开始学 AI:第 2 课,AI、机器学习、深度学习和大模型是什么关系?
人工智能·深度学习·机器学习·ai·大模型·llm
2601_962683893 小时前
治理遗留系统中的“生肉 SQL”:一次用多模型协作优化慢查询的实战复盘
数据库·人工智能·sql
踩着两条虫3 小时前
可视化 vs 终端 vs 云端:VTJ.PRO、Claude Code、Codex 三强横评
前端·vue.js·人工智能·低代码·架构
字节跳动视频云技术团队3 小时前
Agentic 范式下的视频画质优化:火山引擎的新路径
人工智能·音视频开发
欢子4 小时前
Android 端接入AI
人工智能
酱学编程4 小时前
【从零到一实现一个 AI Agent 框架 · 第四篇】04. 任务规划:拆解复杂目标 -
服务器·网络·数据库·人工智能
手写码匠4 小时前
手写 AI 上下文压缩系统:从零实现 Prompt 压缩与选择性上下文管理
人工智能·深度学习·算法·aigc
机器之心4 小时前
世界模型评测的最大盲区,被这个新基准捅破了
人工智能·openai
机器之心4 小时前
Claude、GLM、GPT谁才是真正的AI软件工程师?首个持续更新Visual Spec-to-App Benchmark发布
人工智能·openai