【OpenMMLab 实战营打卡-第5课】MMDetection代码教学

一、创建环境

1. 加载 anaconda, 创建⼀个opennmmlab_mmdetection的conda环境。

2. 安装PyTorch & MMCV

由于北京超算GPU资源为RTX3090, 因此加载的cuda版本需要大于等于11.1。使用 pip 安装的torch 不包括 cuda，所以需要使⽤ module 加载 cuda/11.1 模块。本教程使用的PyTorch版本1.10.0+cu111。

3. 安装mmdetection

采用下载编译的方式安装mmdetection模块，要求编译器gcc版本大于5。由于超算上原本编译器版本低于5，因此需要module命令加载一个高版本的gcc，本教程指定gcc版本为7.3。

二、创建数据集

balloon是带有mask的气球数据集，其中训练集包含61张图片，验证集包含13张图片。

下载链接：https://github.com/matterport/Mask_RCNN/releases/download/v2.1/balloon_dataset.zip

# 切换到mmdetection目录下
cd mmdetection

# 创建ballon2coco.py
vim ballon2coco.py

# 创建data目录，将本地数据集上传到服务器
mkdir data && cd data
unzip balloon_dataset.zip

# 转换为COCO数据集
python ballon2coco.py

ballon2coco.py源码

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))

        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)


if __name__ == '__main__':
    convert_balloon_to_coco('data/balloon/train/via_region_data.json',
                            'data/balloon/train.json', 'data/balloon/train/')

    convert_balloon_to_coco('data/balloon/val/via_region_data.json',
                            'data/balloon/val.json', 'data/balloon/val/')
                            

【注】执行上述命令后balloon目录下会多了train.json和val.json文件，这里分别将train.json和val.json放到train和val目录下。

三、构建配置文件

从 Model Zoo 中找到 mask rcnn 模型并找到 configs/mask_rcnn/ 中对应的模型配置文件。将模型下载到环境中，通常放在放置在 checkpoints 文件夹中。

# 创建模型配置文件
cd mmdetection/configs
mkdir balloon

# 创建配置文件
vim mask_rcnn_r50_caffe_fpn_mstrain-poly_1x_balloon.py

# 下载模型
cd ..
mkdir checkpoints && cd checkpoints
wget https://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

四、模型微调

通过命令行工具或者 Python API 完成模型微调，并在验证集上完成测试和评价。可以调整参数，或使用不同的预训练模型来获得更高的评分。

# 创建任务训练脚本
cd mmdetection
vim run.sh

# 提交任务
sbatch --gpus=1 run.sh

run.sh

四、Color splash特效制作

在获取到图像的mask之后，请同学们将图像转为灰度图像，并在mask区域将原图像的像素值拷贝到灰度图像上即可完成特效制作。测试样例视频：test_video.mp4

# 创建视频处理python脚本
cd mmdetection/demo
vim video_balloon.py

# 创建视频处理GPU任务脚本
cd ..
vim video_process.sh

# 提交任务
sbatch --gpus=1 video_process.sh

video_balloon.py

video_process.sh

效果展示

以上就是第二次作业的全部过程，感谢您的阅读~