天天看点

Mask-RCNN模型的实现自定义对象(无人机)检测

关注可学习更多的OpenCV教程。

问题或建议,请留言;

如果你觉得文章对你有帮助,欢迎转发支持

软件依赖与版本信息:

- tensorflow 1.11
- tensorflow object detection API
- opencv4.1.0
- python 3.6.5
- VS2015

数据标注工具:
https://github.com/abreheret/PixelAnnotationTool           

复制

Part.1-数据生成

打开标注工具PixelAnnotation 选择好dataset路径之后,顺序开始标注数据即可!

Mask-RCNN模型的实现自定义对象(无人机)检测

这个工具对每张原始图像会生成三张图像,分别是

IMAGENAME_color_mask.png

IMAGENAME_mask.png

IMAGENAME_watershed_mask.png

这里总计有546张图像,测试数据是一段每秒25帧的视频文件,最终标注完成之后,我们需要的是IMAGENAME_color_mask.png文件,

需要把它重命名为:

IMAGENAME.png

原图与标注之后的mask图像

Mask-RCNN模型的实现自定义对象(无人机)检测

最终得到的文件结构显示如下

dataset

|---Annotations

|---JPEGImages

|---train_data_dir

|---label_map.pbtxt

解释如下:

dataset 文件夹
Annotations 文件夹,存放的是标注之后mask数据
JPEGImages 文件夹,存放的是原始的JPG格式图像
train_data_dir 运行脚本创建tf record文件所在的目录
label_map.pbtxt 数据的分类JSON描述,这里针对固定翼无人机一个分类           

复制

创建tf record数据,

首先需要下载脚本文件create_mask_rcnn_tf_record.py,下载地址为:

https://github.com/vijendra1125/Custom-Mask-RCNN-using-Tensorfow-Object-detection-API           

复制

放置到tensorflow object detection API框架中的

research\object_detection\dataset_tools

目录下,然后打开修改代码行第57行,标注的每个对象mask都必须索引正确。

然后运行如下

Mask-RCNN模型的实现自定义对象(无人机)检测

就会得到生成的tf record文件。

上述数据标注与生成过程,只适合单对象的标注,如果有多个对象,请先通过labelImg标注好box框,生成XML标注文件,然后再通过PixelAnnotation生成mask数据。

Part.2-迁移学习

使用迁移学习训练Mask-RCNN实现自定义对象检测,首先需要一个预训练模型,这里使用的的预训练模型为:

mask_rcnn_inception_v2_coco

下载地址为:

http://download.tensorflow.org/models/object_detection/mask_rcnn_inception_v2_coco_2018_01_28.tar.gz           

复制

修改pipeline_config文件

从object_detection\samples\configs文件夹下找到对应的config文件

- mask_rcnn_inception_v2_coco.config           

复制

把所有PATH_TO_BE_CONFIGURED地方修改指向真实路径,然后根据需要调整如下几个参数。

keep_aspect_ratio_resizer

-min_dimension

-max_dimension

都修改为 600

-num_classes字段值修改为1,

表示我们只有一个类别

为了防止训练时候资源耗尽导致内存溢出错误,在train_config块中添加以下属性

batch_queue_capacity: 150

num_batch_queue_threads: 8

prefetch_queue_capacity: 10

gradient_clipping_by_norm: 10.0

保存好config文件之后,运行如下命令行开始启动训练

Mask-RCNN模型的实现自定义对象(无人机)检测

通过tensorboard查看训练

Mask-RCNN模型的实现自定义对象(无人机)检测
Mask-RCNN模型的实现自定义对象(无人机)检测

训练结果之后,运行如下脚本导出pb文件

Mask-RCNN模型的实现自定义对象(无人机)检测

Part.3-代码演示

使用导出pb文件,测试结果如下

Mask-RCNN模型的实现自定义对象(无人机)检测

相关演示代码如下:

import tensorflow as tf
import cv2 as cv

# Read the graph.
model_dir = 'D:/tensorflow/uav_train/export/frozen_inference_graph.pb'
with tf.gfile.FastGFile(model_dir, 'rb') as f:
    graph_def = tf.GraphDef()
    graph_def.ParseFromString(f.read())

cap = cv.VideoCapture("D:/images/video/uav_clip.avi")
vout = cv.VideoWriter("D:/mask_rcnn_demo.mp4", cv.VideoWriter_fourcc('D', 'I', 'V', 'X'), 25,
                     (1280, 720), True)
with tf.Session() as sess:
    # Restore session
    sess.graph.as_default()
    tf.import_graph_def(graph_def, name='')

    # process video stream
    count = 0
    while True:
        ret, img = cap.read()
        if ret is not True:
            break
        rows = img.shape[0]
        cols = img.shape[1]
        inp = cv.resize(img, (600, 600))
        inp = inp[:, :, [2, 1, 0]]  # BGR2RGB

        # Run the model
        out = sess.run([sess.graph.get_tensor_by_name('num_detections:0'),
                        sess.graph.get_tensor_by_name('detection_scores:0'),
                        sess.graph.get_tensor_by_name('detection_boxes:0'),
                        sess.graph.get_tensor_by_name('detection_classes:0')],
                       feed_dict={'image_tensor:0': inp.reshape(1, inp.shape[0], inp.shape[1], 3)})

        # Visualize detected bounding boxes.
        num_detections = int(out[0][0])
        for i in range(num_detections):
            classId = int(out[3][0][i])
            score = float(out[1][0][i])
            bbox = [float(v) for v in out[2][0][i]]
            if score > 0.5:
                x = bbox[1] * cols
                y = bbox[0] * rows
                right = bbox[3] * cols
                bottom = bbox[2] * rows
                cv.rectangle(img, (int(x), int(y)), (int(right), int(bottom)), (125, 255, 51), thickness=2)
        count += 1
        cv.imshow('mask_rcnn_demo', img)
        if 2500 < count < 3500:
            vout.write(img)
        c = cv.waitKey(1)
        if c == 27:
            break

# release resource
vout.release()
cap.release()           

复制

参考资料

https://github.com/vijendra1125/Custom-Mask-RCNN-using-Tensorfow-Object-detection-APIhttps://github.com/tensorflow/models/tree/master/research/object_detection