1.1 mmdetction 安装
1.1.1 系统环境需求
参考 mmdetection 官方文档 https://mmdetection.readthedocs.io/en/latest/INSTALL.html,系统环境需求如下:
- Linux (Windows is not officially supported)
- Python 3.5+
- PyTorch 1.1 or higher
- CUDA 9.0 or higher
- NCCL 2
- GCC 4.9 or higher
- mmcv
我的系统环境:
- CentOS 7.2
- Python 3.7
- PyTorch 1.1.0
- CUDA 10.0
- NCCL 2
- GCC 7.4.0
- mmcv
安装 mmcv 时的依赖项如下:
addict
numpy
pyyaml
six
其中,我的环境中缺少 addict 和 pyyaml,从 https://pypi.org/ 中下载源码离线安装。
附:查看深度学习软件/库/工具的命令速查表:
1.1.2 安装 mmdetection
官方文档的安装说明如下,适合网络环境好的条件下进行在线安装,
# Clone the mmdetection repository
git clone https://github.com/open-mmlab/mmdetection.git
cd mmdetection
# Install build requirements and then install mmdetection.
pip install -r requirements/build.txt
pip install "git+https://github.com/cocodataset/cocoapi.git#subdirectory=PythonAPI"
pip install -v -e . # or "python setup.py develop"
这里采用离线安装方式,
(1)检查 mmdetection 的依赖项,满足要求。
cat requirements/build.txt
# These must be installed before building mmdetection
numpy
torch>=1.1
(2)安装 cocoapi,
git clone https://github.com/cocodataset/cocoapi.git
cd cocoapi/PythonAPI
python setup.py install
(3)安装 mmdetection,依赖项检查结果,terminaltables 没有,采用源码离线安装。
安装完成以后,打开 Python 进行验证,
from mmdet.apis import init_detector, inference_detector, show_result
1.2 训练自定义数据集 CatDog
1.2.1 准备数据集
创建 data 文件夹,软链接到数据集根目录,其中标记好的数据集采用 coco 数据格式,
cd mmdetection
mkdir data
export COCO_ROOT=/data1/Projects/datasets/coco
ln -s $COCO_ROOT data
在 mmdet/datasets/cat_dog.py
:
from .coco import CocoDataset
from .registry import DATASETS
@DATASETS.register_module
class CatDog(CocoDataset):
CLASSES = ('dog', 'cat')
在 mmdet/datasets/__init__.py
:
from .cat_dog import CatDog
1.2.2 修改 faster_rcnn 模型配置
下载 resnet50 的预训练模型,放入 $TORCH_HOME
,
export TORCH_HOME=/data1/Projects/pretrained_models
echo $TORCH_HOME
mkdir -p /data1/Projects/pretrained_models/checkpoints/
mv resnet50-19c8e357.pth /data1/Projects/pretrained_models/checkpoints/
用 faster_rcnn 做为模型进行目标检测,拷贝一份 configs/faster_rcnn_r50_fpn_1x.py
为 configs/cat_dog_faster_rcnn_r50_fpn_1x.py
,在 cat_dog_faster_rcnn_r50_fpn_1x.py
,加入预训练的 resnet50 加载路径,
# model settings
import os
os.environ['TORCH_HOME'] = '/data1/Projects/pretrained_models'
在 config 文件 configs/cat_dog_faster_rcnn_r50_fpn_1x.py
使用 CatDog 数据集,分类数 3 包括狗,猫和背景,
num_classes=3,
CatDog 数据集为类 coco 数据集,针对 coco 数据集修改:
- 定义数据种类,修改
mmdetection/mmdet/datasets/coco.py
,把 CLASSES 的 tuple 改为自己数据集对应的种类。
CLASSES = ('dog', 'cat')
- 在
mmdetection/mmdet/core/evaluation/class_names.py
修改 coco_classes 数据集类别,这个关系到后面 test 的时候结果图中显示的类别名称。
def coco_classes():
return [ 'dog', 'cat']
1.2.3 训练模型
使用单个 GPU 进行训练,指定 --work_dir
保存模型结果,
python tools/train.py configs/cat_dog_faster_rcnn_r50_fpn_1x.py \
--gpus 1 \
--work_dir './work_dirs/cat_dog_faster_rcnn_r50_fpn_1x'
1.2.4 测试图片
1.2.4.1 测试单张图片
import numpy as np
from mmdet.apis import init_detector, inference_detector
import mmcv
import cv2
threshold = 0.9 # confidence score
config_file = './configs/cat_dog_faster_rcnn_r50_fpn_1x.py'
checkpoint_file = './work_dirs/cat_dog_faster_rcnn_r50_fpn_1x/latest.pth'
# 通过配置文件(config file)和模型文件(checkpoint file)构建检测模型
model = init_detector(config_file, checkpoint_file, device='cuda:0')
# 测试单张图片并展示结果
img_path = '/data1/Projects/datasets/cat_dog_single/cat.12176.jpg'
result = inference_detector(model, img_path)
bboxes = np.vstack(result)
# print(bboxes.shape)
labels = [
np.full(bbox.shape[0], i, dtype=np.int32)
for i, bbox in enumerate(result)
]
labels = np.concatenate(labels)
# print(labels)
img = cv2.imread(img_path)
scores = bboxes[:, -1]
inds = scores > threshold
bboxes = bboxes[inds, :]
# print(bboxes.shape)
labels = labels[inds]
# print(labels)
class_names = model.CLASSES
cat_dog_dict = {}
for label in labels:
cat_dog_dict[class_names[label]] = cat_dog_dict.get(class_names[label], 0) + 1
print(cat_dog_dict)
for k, v in cat_dog_dict.items():
print('{0} 有 {1} 只 {2}'.format(img_path, v, k))
for bbox in bboxes:
left_top = (bbox[0], bbox[1])
right_bottom = (bbox[2], bbox[3])
cv2.rectangle(img, left_top, right_bottom, color=(0, 255, 0))
cv2.imwrite('/data1/Projects/datasets/cat_dog_single/res_cat.12176.jpg', img)
1.2.4.2 测试多张图片
from mmdet.apis import init_detector, inference_detector, show_result
import mmcv
import numpy as np
import glob
import os
config_file = 'configs/cat_dog_faster_rcnn_r50_fpn_1x.py'
checkpoint_file = 'work_dirs/cat_dog_faster_rcnn_r50_fpn_1x/latest.pth'
score_thr = 0.9
# 通过配置文件(config file)和模型文件(checkpoint file)构建检测模型
model = init_detector(config_file, checkpoint_file, device='cuda:0')
# 测试单张图片
# img = '/data1/Projects/datasets/cat_dog_single/cat.12176.jpg'
# result = inference_detector(model, img)
# show_result(img, result, model.CLASSES, score_thr=score_thr,
# out_file='/data1/Projects/datasets/cat_dog_single/res_cat.12176.jpg')
# 测试多张图片
# imgs = glob.glob('/data1/Projects/datasets/coco/val2017/*.jpg')
imgs = glob.glob('/data1/Projects/datasets/test/*.jpg')
for i, img in enumerate(imgs):
# 画 bounding boxes 到图片上
# print(i, imgs[i])
result = inference_detector(model, img)
file_name = imgs[i].split('/')[-1]
out_file = os.path.join('/data1/Projects/datasets/test_det', file_name)
show_result(img, result, model.CLASSES, score_thr=score_thr, out_file=out_file)
# 输出图片的猫和狗的数量
img = mmcv.imread(img)
img = img.copy()
bbox_result = result
bboxes = np.vstack(bbox_result)
labels = [
np.full(bbox.shape[0], i, dtype=np.int32)
for i, bbox in enumerate(bbox_result)
]
labels = np.concatenate(labels)
# 根据阈值调整输出的 bboxes 和 labels
scores = bboxes[:, -1]
inds = scores > score_thr
bboxes = bboxes[inds, :]
labels = labels[inds]
class_names = model.CLASSES
cat_dog_dict = {}
for label in labels:
cat_dog_dict[class_names[label]] = cat_dog_dict.get(class_names[label], 0) + 1
for k, v in cat_dog_dict.items():
print('{0} 有 {1} 只 {2}'.format(imgs[i].split('/')[-1], v, k))
print('--------------------')
测试多张图片,输出结果如下:
微信公众号「padluo」,分享数据科学家的自我修养,既然遇见,不如一起成长。