mask rcnn训练自己的数据集_深度学习目标检测工具箱mmdetection,训练自己的数据...

97dc9ebb65fd4b1db54cd572a4e4b7a0.png

一、简介

商汤科技(2018 COCO 目标检测挑战赛冠军)和香港中文大学最近开源了一个基于Pytorch实现的深度学习目标检测工具箱mmdetection,支持Faster-RCNN,Mask-RCNN,Fast-RCNN等主流的目标检测框架,后续会加入Cascade-RCNN以及其他一系列目标检测框架。

相比于Facebook开源的Detectron框架,作者声称mmdetection有三点优势:performance稍高、训练速度稍快、所需显存稍小。

我很早就听说了这个工具箱,但是一直没有开源。现在总算是开源了,于是就迫不及待地测试一下这个框架的效果。下面将记录一下我的测试过程,已经期间所遇到的一些坑。

2019.05.26更新了网盘里的demo.py文件,现在的demo应该可以匹配官方的最新代码了。

2019.02.22上传多个mmdetection模型,检测器包括Faster-RCNN、Mask-RCNN、RetinaNet,backbone包括Resnet-50、Resnet-101、ResNext-101,网盘链接:mmdetection(密码:dpyl)

二、安装教程

本人的系统环境:

  • Ubuntu 16.04
  • Cuda 9.0 + Cudnn 7.0.5
  • Python 3.6 (mmdetection要求Python版本需要3.4+)
  • Anaconda 3 (可选)

这里推荐大家使用Anaconda,可以比较方便的创建Python虚拟环境,避免不同的Python库之间产生冲突。在安装mmdetection之前,需要安装以下几个依赖库:

  • PyTorch 1.1 (Pytorch 0.4.1的版本需要切换branch,在clone了mmdetection的git之后需要git checkout pytorch-0.4.1)
  • Cython
  • mmcv

下面是我的安装和测试步骤,以Anaconda 3为例。

1. 使用conda创建Python虚拟环境(可选)

conda create -n mmdetection python=3.6
source activate mmdetection

这样就创建了名为mmdetection的Python3.6环境,并且在terminal中激活。如果不需要虚拟环境,则将下文的conda install改为pip install

2. 安装PyTorch 1.1

conda install pytorch=1.1 -c pytorch

安装好以后,进入Python环境,输入以下代码测试是否安装成功,不报错则说明安装成功

import torch

3. 安装Cython

conda install  cython

4. 安装mmcv

git clone https://github.com/open-mmlab/mmcv.git
cd mmcv
pip install .

注意pip install 后面有个点,别漏了

5. 安装mmdetection

git clone https://github.com/open-mmlab/mmdetection.git
cd mmdetection
./compile.sh
python setup.py install

到此,我们就完成了mmdetection及其依赖库的安装

6. 测试Demo 将下方的代码写入py文件,并存放到mmdetection文件夹目录下,然后运行。该代码的功能是检测图片中的目标,测试模型是官方给出的Faster-RCNN-fpn-resnet50的模型,运行代码会自动下载模型。由于模型是存储在亚马逊云服务器上,速度可能会稍慢,如果下载失败可以通过我的网盘链接mmdetection(密码:dpyl)进行下载,存放到mmdetection文件夹目录下,然后修改下方代码的相关部分

from 

### 7. 准备自己的数据 mmdetection支持coco格式和voc格式的数据集,下面将分别介绍这两种数据集的使用方式 - coco数据集 官方推荐coco数据集按照以下的目录形式存储,以coco2017数据集为例

mmdetection
├── mmdet
├── tools
├── configs
├── data
│   ├── coco
│   │   ├── annotations
│   │   ├── train2017
│   │   ├── val2017
│   │   ├── test2017

推荐以软连接的方式创建data文件夹,下面是创建软连接的步骤

cd mmdetection
mkdir data
ln -s $COCO_ROOT data

其中,$COCO_ROOT需改为你的coco数据集根目录 - voc数据集 与coco数据集类似,将voc数据集按照以下的目录形式存储,以VOC2007为例

mmdetection
├── mmdet
├── tools
├── configs
├── data
│   ├── VOCdevkit
│   │   ├── VOC2007
│   │   │   ├── Annotations
│   │   │   ├── JPEGImages
│   │   │   ├── ImageSets
│   │   │   │   ├── Main
│   │   │   │   │   ├── test.txt
│   │   │   │   │   ├── trainval.txt

同样推荐以软连接的方式创建

cd mmdetection
mkdir data
ln -s $VOC2007_ROOT data/VOCdevkit

其中,$VOC2007_ROOT需改为你的VOC2007数据集根目录

如果需要标注自己的数据,推荐使用LabelImg工具标注,然后需要修改mmdet/datasets/voc.py文件中的CLASSES为你自己的类别

### 8. 训练

官方推荐使用分布式的训练方式,这样速度更快,如果是coco训练集,修改CONFIG_FILE中的pretrained参数,改为你的模型路径,然后运行下方代码

./tools/dist_train.sh <CONFIG_FILE> <GPU_NUM> [optional arguments]

mmdetection官方代码已更新,目前已支持voc格式的数据集,不再需要自己修改

如果不想采用分布式的训练方式,或者你只有一块显卡,则运行下方的代码

python tools/train.py <CONFIG_FILE> --gpus <GPU_NUM> --work_dir <WORK_DIR>

至此,如果一切顺利的话,你的模型应该就开始训练了

  • 0
    点赞
  • 5
    收藏
    觉得还不错? 一键收藏
  • 0
    评论
Mask R-CNN 是一种基于 Faster R-CNN 的目标检测算法,可以检测出物体的位置,还可以精确地分割出物体的 mask。下面是一个 Mask R-CNN 训练自己数据集的基本步骤: 1. 准备数据集 首先,您需要准备自己的数据集,包括图片和标注。图片可以是任何大小,但最好是统一大小,以便于训练。标注可以使用标准的 COCO 标注格式,包括标注文件和类别文件。 2. 定义类别 您需要定义自己数据集的类别,例如人、车、动物等等。您可以创建一个名为 `CustomConfig` 的类,继承自 `Config` 类,并在其中设置 `NUM_CLASSES` 参数为您的类别数量加 1(因为需要加入一个背景类别): ```python class CustomConfig(Config): # Give the configuration a recognizable name NAME = "custom_dataset" # Train on 1 GPU and 1 image per GPU GPU_COUNT = 1 IMAGES_PER_GPU = 1 # Number of classes (including background) NUM_CLASSES = 1 + 4 # Background + 4 classes ``` 3. 定义数据集加载器 您需要定义一个数据集加载器,将准备好的数据集导入到模型中。您可以创建一个名为 `CustomDataset` 的类,继承自 `utils.Dataset` 类,并在其中实现 `load_dataset()`、`load_image()`、`load_mask()`、`image_reference()` 等方法,具体实现可以参考 Mask R-CNN 官方代码。 ```python class CustomDataset(utils.Dataset): def load_dataset(self, dataset_dir, subset): self.add_class("custom_dataset", 1, "class1") self.add_class("custom_dataset", 2, "class2") self.add_class("custom_dataset", 3, "class3") self.add_class("custom_dataset", 4, "class4") # Load annotations annotations = json.load(open(os.path.join(dataset_dir, "annotations.json"))) annotations = annotations["annotations"] # Add images and annotations to dataset for a in annotations: image_id = a["image_id"] image_path = os.path.join(dataset_dir, "images", str(image_id) + ".jpg") if not os.path.exists(image_path): continue if a["iscrowd"]: continue if a["category_id"] not in [1, 2, 3, 4]: continue self.add_image( "custom_dataset", image_id=image_id, path=image_path, width=a["width"], height=a["height"], annotations=a["bbox"] ) def load_mask(self, image_id): # Load annotations for image annotations = self.image_info[image_id]["annotations"] # Create one mask per instance masks = np.zeros([self.image_info[image_id]["height"], self.image_info[image_id]["width"], len(annotations)], dtype=np.uint8) # Load masks for i, a in enumerate(annotations): x1, y1, w, h = a x2 = x1 + w y2 = y1 + h masks[y1:y2, x1:x2, i] = 1 # Return masks and class IDs return masks, np.ones([len(annotations)], dtype=np.int32) def image_reference(self, image_id): info = self.image_info[image_id] return info["path"] ``` 4. 训练模型 在训练之前,您需要将预训练 COCO 权重加载到模型中: ```python model.load_weights(COCO_MODEL_PATH, by_name=True, exclude=["mrcnn_class_logits", "mrcnn_bbox_fc", "mrcnn_bbox", "mrcnn_mask"]) ``` 然后,您可以使用 `train()` 方法训练模型。在训练之前,您需要创建一个名为 `CustomConfig` 的配置对象,并设置好超参数和文件路径: ```python config = CustomConfig() config.display() model = modellib.MaskRCNN(mode="training", config=config, model_dir=MODEL_DIR) # Train the head branches model.train(dataset_train, dataset_val, learning_rate=config.LEARNING_RATE, epochs=30, layers='heads') ``` 5. 测试模型 在测试模型之前,您需要将模型切换到 inference 模式: ```python model = modellib.MaskRCNN(mode="inference", config=config, model_dir=MODEL_DIR) ``` 然后,您可以使用 `detect()` 方法对图片进行检测和分割: ```python results = model.detect([image], verbose=1) r = results[0] visualize.display_instances(image, r['rois'], r['masks'], r['class_ids'], class_names, r['scores']) ``` 以上就是使用 Mask R-CNN 训练自己数据集的基本步骤。具体实现可以参考 Mask R-CNN 官方代码。

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值