支持国产——使用mmdetection进行目标检测并保存推理结果图片

⭐️ 概述

MMDetection是OpenMMLab算法体系的一部分，OpenMMLab是一个国产的开源计算机视觉算法体系。

今天小编初次体验了MMDetection，怀着激动的心情，将这次体验记录下来。

⭐️ 准备工作

笔者的cuda版本是11.8

python环境中 pytorch 相关版本是：

torch                        1.13.1
torchaudio                   0.13.1
torchvision                  0.14.1

安装 mmcv

pip install -U openmim
mim install mmcv==2.1.0

安装 mmdetection

新建一个目录，在该目录下执行：

git clone git@github.com:open-mmlab/mmdetection.git
cd mmdetection
pip install -r requirements.txt
python setup.py develop

注意，执行完上面的命令后，本步骤所建目录的子目录 mmdetection，即为 mmdetection 安装目录。

目前mmdetection的版本为 3.3.0。

⭐️ 开始使用

新建项目目录，该目录独立于 mmdetection 安装目录

项目目录结构如下：

其中demo.py 的内容如下

import os
import mmcv
from mmdet.apis import init_detector, inference_detector
from mmdet.visualization import DetLocalVisualizer


# mmdetection 安装目录，该值取决于你的安装环境
base_dir = 'xxx'
# 输入图片
in_img = 'demo1.jpg'
# in_img = 'demo2.jpg'
# 结果图片
out_img = 'out1.jpg'
# out_img = 'out2.jpg'

if __name__ == '__main__':
    # 模型配置文件
    config_file = os.path.join(base_dir, 'configs/faster_rcnn/faster-rcnn_r50_fpn_1x_coco.py')
    # 模型 checkpoint 文件
    # 下载地址: https://download.openmmlab.com/mmdetection/v2.0/faster_rcnn/faster_rcnn_r50_fpn_1x_coco/faster_rcnn_r50_fpn_1x_coco_20200130-047c8118.pth
    checkpoint_file = r'checkpoints/faster_rcnn_r50_fpn_1x_coco_20200130-047c8118.pth'
    # 构建模型
    model = init_detector(config_file, checkpoint_file, device='cuda:0')
    # 测试单张图片
    img = mmcv.imread(in_img)
    # 利用模型对图片进行推理
    result = inference_detector(model, img)
    # 打印推理结果
    print(result)
    # 可视化推理结果并保存图片
    visualizer = DetLocalVisualizer()
    visualizer.add_datasample('detection', img, result, draw_pred=True, show=False, out_file=out_img)

base_dir一定要换成 mmdetection 安装目录，

我们准备利用Faster Rcnn 对两张图片进行推理，图片分别是

demo1.jpg

demo2.jpg

推理结果如下

效果感觉刚刚的