PyTorch-Mask-RCNN

这是一个基于 PyTorch 的 Mask R-CNN 实现，其大部分代码源自 Matterport 的 Mask_RCNN。Matterport 的仓库是基于 Keras 和 TensorFlow 的实现。以下 README 中的部分内容摘自 Matterport 的 README。关于本仓库的依赖要求、在 MS COCO 数据集上的训练以及检测结果的详细信息，请参阅文档末尾。

Mask R-CNN 模型为图像中每个对象实例生成边界框和分割掩码。它基于特征金字塔网络（FPN）和 ResNet101 主干网络。

接下来的四张图片展示了检测流程中的不同阶段：

1. 锚框排序与过滤

区域建议网络会提出可能属于目标对象的候选边界框。图中展示了正锚框、负锚框以及锚框的精炼过程。

2. 边界框精炼

这是第二阶段中最终检测框（虚线）及其精炼结果（实线）的示例。

3. 掩码生成

生成的掩码示例。这些掩码随后会被缩放并放置到图像的正确位置上。

4. 将各个部分组合成最终结果

需求

• Python 3
• PyTorch 0.3
• matplotlib, scipy, skimage, h5py

安装

1. 克隆本仓库。

    git clone https://github.com/multimodallearning/pytorch-mask-rcnn.git
   

2. 我们使用了另外两个仓库中的函数，需要以正确的 --arch 选项进行编译以支持 CUDA。这两个函数分别是 ruotianluo 的 pytorch-faster-rcnn 仓库中的非极大值抑制函数，以及 longcw 的 RoiAlign 函数。

   显卡	arch
   TitanX	sm_52
   GTX 960M	sm_50
   GTX 1070	sm_61
   GTX 1080 (Ti)	sm_61

    cd nms/src/cuda/
    nvcc -c -o nms_kernel.cu.o nms_kernel.cu -x cu -Xcompiler -fPIC -arch=[arch]
    cd ../../
    python build.py
    cd ../
   
    cd roialign/roi_align/src/cuda/
    nvcc -c -o crop_and_resize_kernel.cu.o crop_and_resize_kernel.cu -x cu -Xcompiler -fPIC -arch=[arch]
    cd ../../
    python build.py
    cd ../../
   

3. 由于我们使用 COCO 数据集，请安装 Python COCO API 并创建一个符号链接。

    ln -s /path/to/coco/cocoapi/PythonAPI/pycocotools/ pycocotools
   

4. 从 Google Drive 下载 COCO 和 ImageNet 上的预训练模型。

演示

要测试您的安装是否成功，只需运行演示脚本：

python demo.py

该脚本可以在 CPU 或 GPU 上运行，输出结果应如下所示：

在 COCO 数据集上训练

训练和评估代码位于 coco.py 文件中。您可以通过命令行运行如下命令：

# 从 COCO 预训练权重开始训练新模型
python coco.py train --dataset=/path/to/coco/ --model=coco

# 从 ImageNet 权重开始训练新模型
python coco.py train --dataset=/path/to/coco/ --model=imagenet

# 继续训练之前已经训练过的模型
python coco.py train --dataset=/path/to/coco/ --model=/path/to/weights.h5

# 继续训练最近一次保存的模型。此命令会自动在模型目录中查找最新的权重文件。
python coco.py train --dataset=/path/to/coco/ --model=last

如果您尚未下载 COCO 数据集，请在运行命令时设置下载选项，例如：

# 从 COCO 预训练权重开始训练新模型
python coco.py train --dataset=/path/to/coco/ --model=coco --download=true

您还可以通过以下命令运行 COCO 评估代码：

# 对最近一次训练的模型进行 COCO 评估
python coco.py evaluate --dataset=/path/to/coco/ --model=last

训练计划、学习率及其他参数可在 coco.py 文件中进行设置。

结果

基于默认配置以及使用预训练 ImageNet 权重初始化主干网络的训练，报告了 COCO 数据集上的边界框和分割结果。使用的指标是 IoU=0.50:0.95 下的 AP 值。

	从零开始	从 Keras 转换而来	Matterport 的 Mask_RCNN	Mask R-CNN 论文
bbox	待定	0.347	0.347	0.382
segm	待定	0.296	0.296	0.354

pytorch-mask-rcnn 快速上手指南

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统: Linux (推荐)
• Python: 3.x
• 深度学习框架: PyTorch 0.3
• 系统依赖: CUDA (需根据显卡型号选择对应的 arch 编译选项)
• Python 库依赖:
  • matplotlib
  • scipy
  • skimage
  • h5py
  • pycocotools (COCO API)

注意：本项目基于较早期的 PyTorch 版本 (0.3)，若使用新版 PyTorch 可能需要修改代码以适配 API 变化。

安装步骤

1. 克隆项目代码

git clone https://github.com/multimodallearning/pytorch-mask-rcnn.git
cd pytorch-mask-rcnn

2. 编译 CUDA 扩展

本项目依赖两个需要编译的 CUDA 模块：非极大值抑制 (NMS) 和 RoIAlign。 重要：请根据您的 GPU 型号替换下方命令中的 [arch] 参数：

GPU 型号	arch 参数
TitanX	sm_52
GTX 960M	sm_50
GTX 1070 / 1080 (Ti)	sm_61

编译 NMS 模块：

cd nms/src/cuda/
# 请将 [arch] 替换为您的实际架构，例如 sm_61
nvcc -c -o nms_kernel.cu.o nms_kernel.cu -x cu -Xcompiler -fPIC -arch=[arch]
cd ../../
python build.py
cd ../

编译 RoIAlign 模块：

cd roialign/roi_align/src/cuda/
# 请将 [arch] 替换为您的实际架构，例如 sm_61
nvcc -c -o crop_and_resize_kernel.cu.o crop_and_resize_kernel.cu -x cu -Xcompiler -fPIC -arch=[arch]
cd ../../
python build.py
cd ../../

3. 配置 COCO API

安装 Python COCO API 并创建软链接（假设您已下载 cocoapi 到本地）：

# 请将 /path/to/coco/cocoapi/PythonAPI/pycocotools/ 替换为您的实际路径
ln -s /path/to/coco/cocoapi/PythonAPI/pycocotools/ pycocotools

4. 下载预训练模型

从 Google Drive 下载在 COCO 和 ImageNet 上预训练的权重文件，并将其放置在项目适当位置（具体加载路径需在代码中确认或参考 demo 脚本）：

• 下载地址：Google Drive Link

基本使用

运行演示脚本

安装完成后，运行以下命令测试环境是否配置成功。该脚本支持 CPU 或 GPU 运行，将输出带有实例分割结果的图像。

python demo.py

训练与评估 (基于 COCO 数据集)

所有训练和评估逻辑均在 coco.py 中。以下是常用命令示例：

开始新训练 (使用 COCO 预训练权重):

python coco.py train --dataset=/path/to/coco/ --model=coco

开始新训练 (使用 ImageNet 预训练权重):

python coco.py train --dataset=/path/to/coco/ --model=imagenet

断点续训:

# 指定权重文件
python coco.py train --dataset=/path/to/coco/ --model=/path/to/weights.h5

# 或自动查找上次训练的权重
python coco.py train --dataset=/path/to/coco/ --model=last

自动下载数据集并训练: 如果尚未下载 COCO 数据集，可添加 --download=true 参数：

python coco.py train --dataset=/path/to/coco/ --model=coco --download=true

执行评估:

python coco.py evaluate --dataset=/path/to/coco/ --model=last

提示: 学习率、训练轮次等超参数可直接在 coco.py 文件中进行调整。