OpenMMLab AI实战营第五课笔记

在气球数据集上训练检测模型

训练新模型通常有三个步骤:

支持新数据集
修改配置文件
训练模型

而MMDetection中有三种方法来支持新数据集:

将数据集整理为COCO格式
将数据集整理为中间格式
直接实现新数据集的支持

我们将基于数据集balloon dataset，通过三种方法来完成本节教程。这是一个带有mask的数据集，其中训练集包含61张图片，验证集包含13张图片。其他方法以及进阶用法可以参考doc。

进入 mmdetection 主目录

In [ ]:

import os
os.chdir('mmdetection')

下载并观察气球数据集

In [ ]:

#下载并解压数据集
!wget https://github.com/matterport/Mask_RCNN/releases/download/v2.1/balloon_dataset.zip
!unzip balloon_dataset.zip > /dev/null

In [ ]:

!tree balloon

In [ ]:

# 简单查看数据集的图片
from PIL import Image
Image.open('balloon/val/14898532020_ba6199dd22_k.jpg')

这里我们以单张图像为例查看balloon数据集的标签json格式，可以看到标签包含基本的图像名，大小信息以及mask所在的区域。

In [ ]:

# 查看单张图片的标签格式
import json

with open('balloon/train/via_region_data.json') as f:
    data = json.load(f)
anno1 = list(data.keys())[0]
print(anno1)
print(json.dumps(data[anno1], indent=2))

下载训练好的mask rcnn模型

In [ ]:

# 下载训练好的 Mask R CNN 模型到checkpoints文件夹
!wget http://download.openmmlab.com/mmdetection/v2.0/mask_rcnn/mask_rcnn_r50_caffe_fpn_mstrain-poly_3x_coco/mask_rcnn_r50_caffe_fpn_mstrain-poly_3x_coco_bbox_mAP-0.408__segm_mAP-0.37_20200504_163245-42aa3d00.pth -P checkpoints

方法一：将数据集整理为COCO格式

这里我们需要了解COCO数据集的格式，并将balloon的标注文件转换为coco格式。

查看COCO数据集格式

对于实例分割来说，比较重要的COCO格式如下所示。

更详细完整的格式说明可以参考这里

{
    "images": [image],
    "annotations": [annotation],
    "categories": [category]
}

# 图像信息
image = {
    "id": int, #图像id
    "width": int, #宽
    "height": int, #高
    "file_name": str, #图像文件名
}

#标注信息
annotation = {
    "id": int, #标注id   
    "image_id": int, #对应的图像id
    "category_id": int, #类别id
    "segmentation": RLE or [polygon], #分割mask信息
    "area": float, #mask面积
    "bbox": [x,y,width,height], #边界框信息
    "iscrowd": 0 or 1, #是否为密集图片
}

#类别信息
categories = [{
    "id": int, #类别id
    "name": str, #类别名
    "supercategory": str, #母类名称
}]

转换标注文件格式

编写功能函数convert_balloon_to_coco完成格式转换。

In [ ]:

import os.path as osp
import mmcv

def convert_balloon_to_coco(ann_file, out_file, image_prefix):
    data_infos = mmcv.load(ann_file)

    annotations = []
    images = []
    obj_count = 0
    for idx, v in enumerate(mmcv.track_iter_progress(data_infos.values())):
        filename = v['filename']
        img_path = osp.join(image_prefix, filename)
        height, width = mmcv.imread(img_path).shape[:2]

        images.append(dict(
            id=idx,
            file_name=filename,
            height=height,
            width=width))

        bboxes = []
        labels = []
        masks = []
        for _, obj in v['regions'].items():
            assert not obj['region_attributes']
            obj = obj['shape_attributes']
            px = obj['all_points_x']
            py = obj['all_points_y']
            poly = [(x + 0.5, y + 0.5) for x, y in zip(px, py)]
            poly = [p for x in poly for p in x]

            x_min, y_min, x_max, y_max = (
                min(px), min(py), max(px), max(py))

            data_anno = dict(
                image_id=idx,
                id=obj_count,
                category_id=0,
                bbox=[x_min, y_min, x_max - x_min, y_max - y_min],
                area=(x_max - x_min) * (y_max - y_min),
                segmentation=[poly],
                iscrowd=0)
            annotations.append(data_anno)
            obj_count += 1

    coco_format_json = dict(
        images=images,
        annotations=annotations,
        categories=[{'id':0, 'name': 'balloon'}])
    mmcv.dump(coco_format_json, out_file)
convert_balloon_to_coco('balloon/train/via_region_data.json', 'balloon/train/coco.json', 'balloon/train')
convert_balloon_to_coco('balloon/val/via_region_data.json', 'balloon/val/coco.json', 'balloon/val')

In [ ]:

# import json

# with open('balloon/val/coco.json') as f:
#     data = json.load(f)

# with open('coco.json', 'w') as f:
#     json.dump(data, f, indent=4)

修改 config 配置文件并保存

In [ ]:

from mmcv import Config
from mmdet.apis import set_random_seed

# 获取基本配置文件参数
cfg = Config.fromfile('./configs/mask_rcnn/mask_rcnn_r50_caffe_fpn_mstrain-poly_1x_coco.py')

# 修改数据集类型以及文件路径
cfg.dataset_type = 'COCODataset'
cfg.data_root = 'balloon/'
cfg.classes = ('balloon',)

cfg.data.train.data_root = 'balloon/'
cfg.data.train.ann_file = 'train/coco.json'
cfg.data.train.img_prefix = 'train'
cfg.data.train.classes = cfg.classes

cfg.data.val.data_root = 'balloon/'
cfg.data.val.ann_file = 'val/coco.json'
cfg.data.val.img_prefix = 'val'
cfg.data.val.classes = cfg.classes

cfg.data.test.data_root = 'balloon/'
cfg.data.test.ann_file = 'val/coco.json'
cfg.data.test.img_prefix = 'val'
cfg.data.test.classes = cfg.classes

# 修改bbox_head和mask_head中的类别数
cfg.model.roi_head.bbox_head.num_classes = 1
cfg.model.roi_head.mask_head.num_classes = 1
# 使用预训练好的faster rcnn模型用于finetuning
cfg.load_from = 'checkpoints/mask_rcnn_r50_caffe_fpn_mstrain-poly_3x_coco_bbox_mAP-0.408__segm_mAP-0.37_20200504_163245-42aa3d00.pth'
# 设置工作目录用于存放log和临时文件
cfg.work_dir = 'work_dirs/balloon'

# 原本的学习率是在8卡基础上训练设置的，现在单卡需要除以8
cfg.optimizer.lr = 0.02 / 8
cfg.lr_config.warmup = None
cfg.log_config.interval = 10

# # 设置evaluation间隔减少运行时间
# cfg.evaluation.interval = 12
# # 设置存档点间隔减少存储空间的消耗
# cfg.checkpoint_config.interval = 12

# 设置运行epoch数
cfg.runner.max_epochs = 1

# 固定随机种子使得结果可复现
cfg.seed = 0
set_random_seed(0, deterministic=False)
cfg.gpu_ids = range(1)

cfg.device = 'cuda'

# 打印所有的配置参数
# print(f'Config:\n{cfg.pretty_text}')

# 保存配置文件
mmcv.mkdir_or_exist('work_dirs/balloon')
cfg.dump('work_dirs/balloon/cocoformat.py')

训练新模型

In [ ]:

from mmdet.datasets import build_dataset
from mmdet.models import build_detector
from mmdet.apis import train_detector

# 构建数据集
datasets = [build_dataset(cfg.data.train)]

# 构建检测模型
model = build_detector(cfg.model)
# 添加类别文字属性提高可视化效果
model.CLASSES = datasets[0].CLASSES

# 创建工作目录并训练模型
mmcv.mkdir_or_exist(osp.abspath(cfg.work_dir))

In [ ]:

# 训练模型
train_detector(model, datasets, cfg, distributed=False, validate=True)

在测试集评价模型

训练完模型后，需要对模型进行评价分析，这里可以用mmdet提供的test.py来直接运行。更多评价方法可以查看这里

In [ ]:

!python tools/test.py work_dirs/balloon/cocoformat.py work_dirs/balloon/latest.pth --eval bbox segm

预测新图像

In [ ]:

from mmdet.apis import inference_detector, show_result_pyplot

img = mmcv.imread('balloon/val/14898532020_ba6199dd22_k.jpg')

model.cfg = cfg
result = inference_detector(model, img)
show_result_pyplot(model, img, result)

In [ ]:

result[0][0].shape

In [ ]:

len(result[1][0])

In [ ]:

result[1][0][0].shape

方法二：将数据集转换为中间格式

除了coco格式，mmdet还支持中间格式的定义。

数据集中间格式

[
    {
        'filename': 'a.jpg',
        'width': 1280,
        'height': 720,
        'ann': {
            'bboxes': <np.ndarray> (n, 4),
            'labels': <np.ndarray> (n, ),
            'bboxes_ignore': <np.ndarray> (k, 4), (optional field)
            'labels_ignore': <np.ndarray> (k, 4) (optional field),
            'masks': [poly]
        }
    },
    ...
]

我们需要根据上述中间类型来编写转换函数。

转换标注文件格式

编写功能函数convert_balloon_to_middle完成格式转换。

In [ ]:

import os.path as osp
import numpy as np

def convert_balloon_to_middle(ann_file, out_file, image_prefix):
    data_infos = mmcv.load(ann_file)

    middle_format = []
    obj_count = 0
    for idx, v in enumerate(mmcv.track_iter_progress(data_infos.values())):
        filename = v['filename']
        img_path = osp.join(image_prefix, filename)
        height, width = mmcv.imread(img_path).shape[:2]

        data_info = dict(filename=filename, width=width, height=height)

        gt_bboxes = []
        gt_labels = []
        gt_masks_ann = []

        for _, obj in v['regions'].items():
            assert not obj['region_attributes']
            obj = obj['shape_attributes']
            px = obj['all_points_x']
            py = obj['all_points_y']
            poly = [(x + 0.5, y + 0.5) for x, y in zip(px, py)]
            poly = [p for x in poly for p in x]

            x_min, y_min, x_max, y_max = (
                min(px), min(py), max(px), max(py))

            bbox=[x_min, y_min, x_max, y_max]
            label = 0
            segmentation=[poly]
            gt_bboxes.append(bbox)
            gt_labels.append(label)
            gt_masks_ann.append(segmentation)

        data_anno = dict(
            bboxes=np.array(gt_bboxes, dtype=np.float32).reshape(-1, 4),
            labels=np.array(gt_labels, dtype=np.long),
            # bboxes_ignore = np.zeros((0, 4), dtype=np.float32),
            # masks=gt_masks_ann
            )         
     
        data_info.update(ann=data_anno)
        middle_format.append(data_info)

    mmcv.dump(middle_format, out_file)
convert_balloon_to_middle('balloon/train/via_region_data.json', 'balloon/train/middle.pkl', 'balloon/train')
convert_balloon_to_middle('balloon/val/via_region_data.json', 'balloon/val/middle.pkl', 'balloon/val')

In [ ]:

# 查看单张图片的标签格式
import json

data = mmcv.load('balloon/train/middle.pkl')
print(data[3])

修改 config 配置并保存

注意这里记得修改各个部分的数据集类型为CustomDataset、对应标注文件为middle.pkl、工作目录和配置文件中的部分为middleformat。

In [ ]:

from mmcv import Config
from mmdet.apis import set_random_seed

# 获取基本配置文件参数
# cfg = Config.fromfile('./configs/mask_rcnn/mask_rcnn_r50_caffe_fpn_mstrain-poly_1x_coco.py')
cfg = Config.fromfile('./configs/faster_rcnn/faster_rcnn_r50_caffe_fpn_mstrain_1x_coco.py')

# 修改数据集类型以及文件路径
cfg.dataset_type = 'CustomDataset'
cfg.data_root = 'balloon/'
cfg.classes = ('balloon',)

cfg.data.train.type = 'CustomDataset'
cfg.data.train.data_root = 'balloon/'
cfg.data.train.ann_file = 'train/middle.pkl'
cfg.data.train.img_prefix = 'train'
cfg.data.train.classes = cfg.classes

cfg.data.val.type = 'CustomDataset'
cfg.data.val.data_root = 'balloon/'
cfg.data.val.ann_file = 'val/middle.pkl'
cfg.data.val.img_prefix = 'val'
cfg.data.val.classes = cfg.classes

cfg.data.test.type = 'CustomDataset'
cfg.data.test.data_root = 'balloon/'
cfg.data.test.ann_file = 'val/middle.pkl'
cfg.data.test.img_prefix = 'val'
cfg.data.test.classes = cfg.classes

# 修改bbox_head和mask_head中的类别数
cfg.model.roi_head.bbox_head.num_classes = 1
# cfg.model.roi_head.mask_head.num_classes = 1
# 使用预训练好的faster rcnn模型用于finetuning
cfg.load_from = 'checkpoints/mask_rcnn_r50_caffe_fpn_mstrain-poly_3x_coco_bbox_mAP-0.408__segm_mAP-0.37_20200504_163245-42aa3d00.pth'
# 设置工作目录用于存放log和临时文件
cfg.work_dir = 'work_dirs/balloon/middleformat'

# 原本的学习率是在8卡基础上训练设置的，现在单卡需要除以8
cfg.optimizer.lr = 0.02 / 8
cfg.lr_config.warmup = None
cfg.log_config.interval = 10

# 由于是自定义数据集，需要修改评价方法
cfg.evaluation.metric = 'mAP'
# # 设置evaluation间隔减少运行时间
# cfg.evaluation.interval = 12
# # 设置存档点间隔减少存储空间的消耗
# cfg.checkpoint_config.interval = 12

# 设置运行epoch数
cfg.runner.max_epochs = 1

# 固定随机种子使得结果可复现
cfg.seed = 0
set_random_seed(0, deterministic=False)
cfg.gpu_ids = range(1)

cfg.device = 'cuda'

# 打印所有的配置参数
# print(f'Config:\n{cfg.pretty_text}')

# 保存配置文件
mmcv.mkdir_or_exist('work_dirs/balloon/middleformat')
cfg.dump('work_dirs/balloon/middleformat/middleformat.py')

训练新模型

这里需要根据配置文件构建数据集，检测模型并完成训练。

In [ ]:

from mmdet.datasets import build_dataset
from mmdet.models import build_detector
from mmdet.apis import train_detector

# 构建数据集
datasets = [build_dataset(cfg.data.train)]

# 构建检测模型
model = build_detector(cfg.model)
# 添加类别文字属性提高可视化效果
model.CLASSES = datasets[0].CLASSES

# 创建工作目录并训练模型
mmcv.mkdir_or_exist(osp.abspath(cfg.work_dir))

In [ ]:

# 训练模型
train_detector(model, datasets, cfg, distributed=False, validate=True)

测试集上评价模型

In [ ]:

!python tools/test.py work_dirs/balloon/middleformat/middleformat.py work_dirs/balloon/middleformat/latest.pth --eval mAP

预测新图像

In [ ]:

from mmdet.apis import inference_detector, show_result_pyplot

img = mmcv.imread('balloon/val/14898532020_ba6199dd22_k.jpg')

model.cfg = cfg
result = inference_detector(model, img)

show_result_pyplot(model, img, result)

方法三：直接实现新数据集的支持

在mmdet/datasets新建BalloonDataset.py文件，填充内容如下：

并在mmdet/datasets/__init__.py中，增加from .BalloonDataset import BalloonDataset

In [ ]:

# mmdet/datasets/BalloonDataset.py

import copy
import os.path as osp

import mmcv
import numpy as np

from mmdet.datasets.builder import DATASETS
from mmdet.datasets.custom import CustomDataset

@DATASETS.register_module()
class BalloonDataset(CustomDataset):

    CLASSES = ('balloon',)

    def load_annotations(self, ann_file):
        data_infos = mmcv.load(self.ann_file)

        middle_format = []
        obj_count = 0
        for idx, v in enumerate(mmcv.track_iter_progress(data_infos.values())):
            filename = v['filename']
            img_path = osp.join(self.img_prefix, filename)
            height, width = mmcv.imread(img_path).shape[:2]

            data_info = dict(filename=filename, width=width, height=height)

            gt_bboxes = []
            gt_labels = []
            gt_masks_ann = []

            for _, obj in v['regions'].items():
                assert not obj['region_attributes']
                obj = obj['shape_attributes']
                px = obj['all_points_x']
                py = obj['all_points_y']
                poly = [(x + 0.5, y + 0.5) for x, y in zip(px, py)]
                poly = [p for x in poly for p in x]

                x_min, y_min, x_max, y_max = (
                    min(px), min(py), max(px), max(py))

                bbox=[x_min, y_min, x_max, y_max]
                label = 0
                segmentation=[poly]
                gt_bboxes.append(bbox)
                gt_labels.append(label)
                gt_masks_ann.append(segmentation)

            data_anno = dict(
                bboxes=np.array(gt_bboxes, dtype=np.float32).reshape(-1, 4),
                labels=np.array(gt_labels, dtype=np.long),
                # bboxes_ignore = np.zeros((0, 4), dtype=np.float32),
                # masks=gt_masks_ann
                )         
        
            data_info.update(ann=data_anno)
            middle_format.append(data_info)

        return middle_format

修改config配置文件并保存

这里同样注意修改数据集类型，标注文件名以及工作目录和配置文件存储路径等

In [ ]:

from mmcv import Config
from mmdet.apis import set_random_seed

# 获取基本配置文件参数
# cfg = Config.fromfile('./configs/mask_rcnn/mask_rcnn_r50_caffe_fpn_mstrain-poly_1x_coco.py')
cfg = Config.fromfile('./configs/faster_rcnn/faster_rcnn_r50_caffe_fpn_mstrain_1x_coco.py')

# 修改数据集类型以及文件路径
cfg.dataset_type = 'BalloonDataset'
cfg.data_root = 'balloon/'

cfg.data.train.type = 'BalloonDataset'
cfg.data.train.data_root = 'balloon/'
cfg.data.train.ann_file = 'train/via_region_data.json'
cfg.data.train.img_prefix = 'train'

cfg.data.val.type = 'BalloonDataset'
cfg.data.val.data_root = 'balloon/'
cfg.data.val.ann_file = 'val/via_region_data.json'
cfg.data.val.img_prefix = 'val'

cfg.data.test.type = 'BalloonDataset'
cfg.data.test.data_root = 'balloon/'
cfg.data.test.ann_file = 'val/via_region_data.json'
cfg.data.test.img_prefix = 'val'

# 修改bbox_head和mask_head中的类别数
cfg.model.roi_head.bbox_head.num_classes = 1
# cfg.model.roi_head.mask_head.num_classes = 1
# 使用预训练好的faster rcnn模型用于finetuning
cfg.load_from = 'checkpoints/mask_rcnn_r50_caffe_fpn_mstrain-poly_3x_coco_bbox_mAP-0.408__segm_mAP-0.37_20200504_163245-42aa3d00.pth'
# 设置工作目录用于存放log和临时文件
cfg.work_dir = 'work_dirs/balloon/newdataset'

# 原本的学习率是在8卡基础上训练设置的，现在单卡需要除以8
cfg.optimizer.lr = 0.02 / 8
cfg.lr_config.warmup = None
cfg.log_config.interval = 10

# 由于是自定义数据集，需要修改评价方法
cfg.evaluation.metric = 'mAP'
# # 设置evaluation间隔减少运行时间
# cfg.evaluation.interval = 12
# # 设置存档点间隔减少存储空间的消耗
# cfg.checkpoint_config.interval = 12

# 设置运行epoch数
cfg.runner.max_epochs = 1

# 固定随机种子使得结果可复现
cfg.seed = 0
set_random_seed(0, deterministic=False)
cfg.gpu_ids = range(1)

cfg.device = 'cuda'

# 打印所有的配置参数
# print(f'Config:\n{cfg.pretty_text}')

# 保存配置文件
mmcv.mkdir_or_exist('work_dirs/balloon/newdataset')
cfg.dump('work_dirs/balloon/newdataset/newdataset.py')

训练新模型

In [ ]:

from mmdet.datasets import build_dataset
from mmdet.models import build_detector
from mmdet.apis import train_detector

# 构建数据集
datasets = [build_dataset(cfg.data.train)]

# 构建检测模型
model = build_detector(cfg.model)
# 添加类别文字属性提高可视化效果
model.CLASSES = datasets[0].CLASSES

# 创建工作目录并训练模型
mmcv.mkdir_or_exist(osp.abspath(cfg.work_dir))

In [ ]:

# 训练模型
train_detector(model, datasets, cfg, distributed=False, validate=True)

测试集上评价模型

In [ ]:

!python tools/test.py work_dirs/ballon/newdataset/newdataset.py work_dirs/ballon/newdataset/latest.pth --eval mAP

预测新图像

In [ ]:

from mmdet.apis import inference_detector, show_result_pyplot

img = mmcv.imread('balloon/val/14898532020_ba6199dd22_k.jpg')

model.cfg = cfg
result = inference_detector(model, img)

show_result_pyplot(model, img, result)

posted @ 2023-02-10 22:57 Xu_Lin 阅读(167) 评论(0) 编辑收藏举报

会员力量，点亮园子希望

刷新页面返回顶部

Xu_Lin

Do more; Learn more; Be more

OpenMMLab AI实战营第五课笔记

OpenMMLab AI实战营第五课笔记

在气球数据集上训练检测模型

进入 mmdetection 主目录

下载并观察气球数据集

下载训练好的mask rcnn模型

方法一：将数据集整理为COCO格式

查看COCO数据集格式

转换标注文件格式

修改 config 配置文件并保存

训练新模型

在测试集评价模型

预测新图像

方法二：将数据集转换为中间格式

数据集中间格式

转换标注文件格式

修改 config 配置并保存

训练新模型

测试集上评价模型

预测新图像

方法三：直接实现新数据集的支持

修改config配置文件并保存

训练新模型

测试集上评价模型

预测新图像

公告

Xu_Lin

Do more; Learn more; Be more

OpenMMLab AI实战营 第五课笔记

OpenMMLab AI实战营 第五课笔记

在气球数据集上训练检测模型

进入 mmdetection 主目录

下载并观察气球数据集

下载训练好的mask rcnn模型

方法一：将数据集整理为COCO格式

查看COCO数据集格式

转换标注文件格式

修改 config 配置文件并保存

训练新模型

在测试集评价模型

预测新图像

方法二：将数据集转换为中间格式

数据集中间格式

转换标注文件格式

修改 config 配置并保存

训练新模型

测试集上评价模型

预测新图像

方法三：直接实现新数据集的支持

修改config配置文件并保存

训练新模型

测试集上评价模型

预测新图像

公告

OpenMMLab AI实战营第五课笔记

OpenMMLab AI实战营第五课笔记