????????前面的文章 TensorFlow 訓練自己的目標檢測器 寫作的時候,TensorFlow models 項目下的目標檢測專題 object_detection 還沒有給出用于實例分割的預訓練模型,但其實這個專題中的 Faster R-CNN 模型是按照 Mask R-CNN 來寫的,只要用戶在訓練時傳入了 mask,則模型也會預測 mask,這可以從該專題下的文件
object_detection/meta_architectures/faster_rcnn_meta_arch.py
看出來。
????????現(xiàn)在,TensorFlow object_detection API 官方已經(jīng)放出 Mask R-CNN 預訓練模型,而且對目標檢測的源代碼也做了一些改動(主要是引入了 TensorFlow 的兩個高級 API 模塊:tf.estimator 以及 tf.keras),為了與官方的迭代同步,以及作為文章 TensorFlow 訓練自己的目標檢測器 的補充,特別記錄一下利用 TensorFlow/Object Detection API 來訓練 Mask R-CNN 模型的過程。
????????首先,來體驗一下官方公布的預訓練模型的分割效果(此處引用的預訓練模型是 mask_rcnn_resnet101_atrous_coco):

顯然,效果還是不錯的。
????????關于 TensorFlow/Object Detection API 的安裝以及其它信息請參考前面文章 TensorFlow 訓練自己的目標檢測器。如果你對此不陌生(且成功更新了該 API)的話,那么話不多說,馬上進入正題。
????????所有代碼見 github。如果不想自己標注訓練數(shù)據(jù),請使用下一篇文章生成的數(shù)據(jù) train.record 和 val.record 以及對應的 shape_label_map.pbtxt 來訓練 Mask R-CNN 模型。
一、數(shù)據(jù)準備
????????因為只是詳實的記錄一下訓練過程,所以數(shù)據(jù)量不需要太多,我們以數(shù)據(jù)集 Oxford-IIIT Pet 中的 阿比西尼亞貓(Abyssinian) 為例來說明。數(shù)據(jù)集 Oxford-IIIT Pet 可以從 這里 下載,數(shù)據(jù)量不大,只有 800M 不到。其中,阿比西尼亞貓的圖像只有 232 張,這種貓的長相如下:

要訓練 Mask R-CNN 實例分割模型,我們首先要準備圖像的掩模(mask),使用標注工具 labelme(支持 Windows 和 Ubuntu,使用 (sudo) pip install labelme 安裝,需要安裝依賴項:(sudo) pip install pyqt5)來完成這一步。安裝完 labelme 之后,在命令行執(zhí)行 labelme 會彈出一個標注窗口:

跟 labelImg 幾乎一樣。從 Open 打開一張圖像 Abyssinian_65.jpg,之后使用 Create Polygons 描出目標所在的近似多邊形區(qū)域:

點擊 Save 之后選擇路徑保存為一個 json 文件:Abyssinian_65.json。
????????【將命令行來到 Abyssinian_65.json 文件所在的文件夾,執(zhí)行
labelme_json_to_dataset Abyssinian_65.json
會在當前目錄下生成一個名叫 Abyssinian_65_json 的文件夾,里面包含如下文件:

其中的 label.png 圖像:

正是在公開數(shù)據(jù)集經(jīng)常見到的實例分割掩模?!俊?strong>這一段只用來描述 labelme 的完整功能,實際上本文不需要執(zhí)行這個過程!
????????但是 labelme 有一個很大的缺陷,即它只能標注首尾相連的多邊形,如果一個目標實例包含一個洞(如第二幅圖像 Abyssinian_65.jpg 的貓的兩腿之間的空隙),那么這個洞也會算成這個目標實例的一部分,而這顯然是不正確的。為了避免這個缺陷,在標注目標實例時,可以增加一個額外的類 hole(如上圖的 綠色 部分),實際使用時只要把 hole 部分去掉即可,如:

????????TensorFlow 訓練時要求 mask 是跟原圖像一樣大小的二值(0-1)png 圖像(如上圖),而且數(shù)據(jù)輸入格式必須為 tfrecord 文件,所以還需要寫一個數(shù)據(jù)格式轉化的輔助 python 文件,該文件可以參考 TensorFlow 目標檢測官方的文件 create_coco_tf_record.py 來寫。
????????在寫之前,強調(diào)說明一下數(shù)據(jù)輸入的格式:對每張圖像中的每個目標,該目標的 mask 是一張與原圖一樣大小的 0-1 二值圖像,該目標所在區(qū)域的值為 1,其他區(qū)域全為 0(見 TensorFlow/object_detection 官方說明:Run an Instance Segmentation Model/PNG Instance Segmentation Masks)。也就是說,同一張圖像中的所有目標的 mask 都需要從單個標注文件中分割出來。這可以使用 OpenCV 的 cv2.fillPoly 函數(shù)來實現(xiàn),該函數(shù)將指定多邊形區(qū)域內(nèi)部的值都填充為用戶設定的值。
????????假設已經(jīng)準備好了 mask 標注數(shù)據(jù),因為包圍每個目標的 mask 的最小矩形就是該目標的 boundingbox,所以目標檢測的標注數(shù)據(jù)也就同時有了。接下來,只需要將這些標注數(shù)據(jù)(原始圖像,以及 labelme 標注生成的 json 文件)轉換成 TFRecord 文件即可,使用如下代碼完成這一步操作(命名為 create_tf_record.py,見 github):
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Sun Aug 26 10:57:09 2018
@author: shirhe-lyh
"""
"""Convert raw dataset to TFRecord for object_detection.
Please note that this tool only applies to labelme's annotations(json file).
Example usage:
python3 create_tf_record.py \
--images_dir=your absolute path to read images.
--annotations_json_dir=your path to annotaion json files.
--label_map_path=your path to label_map.pbtxt
--output_path=your path to write .record.
"""
import cv2
import glob
import hashlib
import io
import json
import numpy as np
import os
import PIL.Image
import tensorflow as tf
import read_pbtxt_file
flags = tf.app.flags
flags.DEFINE_string('images_dir', None, 'Path to images directory.')
flags.DEFINE_string('annotations_json_dir', 'datasets/annotations',
'Path to annotations directory.')
flags.DEFINE_string('label_map_path', None, 'Path to label map proto.')
flags.DEFINE_string('output_path', None, 'Path to the output tfrecord.')
FLAGS = flags.FLAGS
def int64_feature(value):
return tf.train.Feature(int64_list=tf.train.Int64List(value=[value]))
def int64_list_feature(value):
return tf.train.Feature(int64_list=tf.train.Int64List(value=value))
def bytes_feature(value):
return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))
def bytes_list_feature(value):
return tf.train.Feature(bytes_list=tf.train.BytesList(value=value))
def float_list_feature(value):
return tf.train.Feature(float_list=tf.train.FloatList(value=value))
def create_tf_example(annotation_dict, label_map_dict=None):
"""Converts image and annotations to a tf.Example proto.
Args:
annotation_dict: A dictionary containing the following keys:
['height', 'width', 'filename', 'sha256_key', 'encoded_jpg',
'format', 'xmins', 'xmaxs', 'ymins', 'ymaxs', 'masks',
'class_names'].
label_map_dict: A dictionary maping class_names to indices.
Returns:
example: The converted tf.Example.
Raises:
ValueError: If label_map_dict is None or is not containing a class_name.
"""
if annotation_dict is None:
return None
if label_map_dict is None:
raise ValueError('`label_map_dict` is None')
height = annotation_dict.get('height', None)
width = annotation_dict.get('width', None)
filename = annotation_dict.get('filename', None)
sha256_key = annotation_dict.get('sha256_key', None)
encoded_jpg = annotation_dict.get('encoded_jpg', None)
image_format = annotation_dict.get('format', None)
xmins = annotation_dict.get('xmins', None)
xmaxs = annotation_dict.get('xmaxs', None)
ymins = annotation_dict.get('ymins', None)
ymaxs = annotation_dict.get('ymaxs', None)
masks = annotation_dict.get('masks', None)
class_names = annotation_dict.get('class_names', None)
labels = []
for class_name in class_names:
label = label_map_dict.get(class_name, None)
if label is None:
raise ValueError('`label_map_dict` is not containing {}.'.format(
class_name))
labels.append(label)
encoded_masks = []
for mask in masks:
pil_image = PIL.Image.fromarray(mask.astype(np.uint8))
output_io = io.BytesIO()
pil_image.save(output_io, format='PNG')
encoded_masks.append(output_io.getvalue())
feature_dict = {
'image/height': int64_feature(height),
'image/width': int64_feature(width),
'image/filename': bytes_feature(filename.encode('utf8')),
'image/source_id': bytes_feature(filename.encode('utf8')),
'image/key/sha256': bytes_feature(sha256_key.encode('utf8')),
'image/encoded': bytes_feature(encoded_jpg),
'image/format': bytes_feature(image_format.encode('utf8')),
'image/object/bbox/xmin': float_list_feature(xmins),
'image/object/bbox/xmax': float_list_feature(xmaxs),
'image/object/bbox/ymin': float_list_feature(ymins),
'image/object/bbox/ymax': float_list_feature(ymaxs),
'image/object/mask': bytes_list_feature(encoded_masks),
'image/object/class/label': int64_list_feature(labels)}
example = tf.train.Example(features=tf.train.Features(
feature=feature_dict))
return example
def _get_annotation_dict(images_dir, annotation_json_path):
"""Get boundingboxes and masks.
Args:
images_dir: Path to images directory.
annotation_json_path: Path to annotated json file corresponding to
the image. The json file annotated by labelme with keys:
['lineColor', 'imageData', 'fillColor', 'imagePath', 'shapes',
'flags'].
Returns:
annotation_dict: A dictionary containing the following keys:
['height', 'width', 'filename', 'sha256_key', 'encoded_jpg',
'format', 'xmins', 'xmaxs', 'ymins', 'ymaxs', 'masks',
'class_names'].
#
# Raises:
# ValueError: If images_dir or annotation_json_path is not exist.
"""
# if not os.path.exists(images_dir):
# raise ValueError('`images_dir` is not exist.')
#
# if not os.path.exists(annotation_json_path):
# raise ValueError('`annotation_json_path` is not exist.')
if (not os.path.exists(images_dir) or
not os.path.exists(annotation_json_path)):
return None
with open(annotation_json_path, 'r') as f:
json_text = json.load(f)
shapes = json_text.get('shapes', None)
if shapes is None:
return None
image_relative_path = json_text.get('imagePath', None)
if image_relative_path is None:
return None
image_name = image_relative_path.split('/')[-1]
image_path = os.path.join(images_dir, image_name)
image_format = image_name.split('.')[-1].replace('jpg', 'jpeg')
if not os.path.exists(image_path):
return None
with tf.gfile.GFile(image_path, 'rb') as fid:
encoded_jpg = fid.read()
image = cv2.imread(image_path)
height = image.shape[0]
width = image.shape[1]
key = hashlib.sha256(encoded_jpg).hexdigest()
xmins = []
xmaxs = []
ymins = []
ymaxs = []
masks = []
class_names = []
hole_polygons = []
for mark in shapes:
class_name = mark.get('label')
class_names.append(class_name)
polygon = mark.get('points')
polygon = np.array(polygon)
if class_name == 'hole':
hole_polygons.append(polygon)
else:
mask = np.zeros(image.shape[:2])
cv2.fillPoly(mask, [polygon], 1)
masks.append(mask)
# Boundingbox
x = polygon[:, 0]
y = polygon[:, 1]
xmin = np.min(x)
xmax = np.max(x)
ymin = np.min(y)
ymax = np.max(y)
xmins.append(float(xmin) / width)
xmaxs.append(float(xmax) / width)
ymins.append(float(ymin) / height)
ymaxs.append(float(ymax) / height)
# Remove holes in mask
for mask in masks:
mask = cv2.fillPoly(mask, hole_polygons, 0)
annotation_dict = {'height': height,
'width': width,
'filename': image_name,
'sha256_key': key,
'encoded_jpg': encoded_jpg,
'format': image_format,
'xmins': xmins,
'xmaxs': xmaxs,
'ymins': ymins,
'ymaxs': ymaxs,
'masks': masks,
'class_names': class_names}
return annotation_dict
def main(_):
if not os.path.exists(FLAGS.images_dir):
raise ValueError('`images_dir` is not exist.')
if not os.path.exists(FLAGS.annotations_json_dir):
raise ValueError('`annotations_json_dir` is not exist.')
if not os.path.exists(FLAGS.label_map_path):
raise ValueError('`label_map_path` is not exist.')
label_map = read_pbtxt_file.get_label_map_dict(FLAGS.label_map_path)
writer = tf.python_io.TFRecordWriter(FLAGS.output_path)
num_annotations_skiped = 0
annotations_json_path = os.path.join(FLAGS.annotations_json_dir, '*.json')
for i, annotation_file in enumerate(glob.glob(annotations_json_path)):
if i % 100 == 0:
print('On image %d', i)
annotation_dict = _get_annotation_dict(
FLAGS.images_dir, annotation_file)
if annotation_dict is None:
num_annotations_skiped += 1
continue
tf_example = create_tf_example(annotation_dict, label_map)
writer.write(tf_example.SerializeToString())
print('Successfully created TFRecord to {}.'.format(FLAGS.output_path))
if __name__ == '__main__':
tf.app.run()
假設你的所有原始圖像的路徑為 path_to_images_dir,使用 labelme 標注產(chǎn)生的所有用于 訓練 的 json 文件的路徑為 path_to_train_annotations_json_dir,所有用于 驗證 的 json 文件的路徑為 path_to_val_annotaions_json_dir,在終端先后執(zhí)行如下指令:
$ python3 create_tf_record.py \
--images_dir=path_to_images_dir \
--annotations_json_dir=path_to_train_annotations_json_dir \
--label_map_path=path_to_label_map.pbtxt \
--output_path=path_to_train.record
$ python3 create_tf_record.py \
--images_dir=path_to_images_dir \
--annotations_json_dir=path_to_val_annotations_json_dir \
--label_map_path=path_to_label_map.pbtxt \
--output_path=path_to_val.record
其中,以上所有路徑都支持相對路徑。output_path 為輸出的 train.record 以及 val.record 的路徑,label_map_path 是所有需要檢測的類名及類標號的配置文件,該文件的后綴名為 .pbtxt,寫法很簡單,假如你要檢測 ’person' , 'car' ,'bicycle' 等類目標,則寫入如下內(nèi)容:
item {
????????id: 1
????????name: 'person'
}item {
????????id: 2
????????name: 'car'
}item {
????????id: 3
????????name: 'bicycle'
}...
這里我們只檢測 阿比西尼亞貓(Abyssinian)一個類,所以只需要寫入:
item {
????????id: 1
????????name: 'Abyssinian'
}
即可,命名為 Abyssinian_label_map.pbtxt。
???????通過以上源代碼的部分內(nèi)容:
feature_dict = {
'image/height': int64_feature(height),
'image/width': int64_feature(width),
'image/filename': bytes_feature(filename.encode('utf8')),
'image/source_id': bytes_feature(filename.encode('utf8')),
'image/key/sha256': bytes_feature(sha256_key.encode('utf8')),
'image/encoded': bytes_feature(encoded_jpg),
'image/format': bytes_feature(image_format.encode('utf8')),
'image/object/bbox/xmin': float_list_feature(xmins),
'image/object/bbox/xmax': float_list_feature(xmaxs),
'image/object/bbox/ymin': float_list_feature(ymins),
'image/object/bbox/ymax': float_list_feature(ymaxs),
'image/object/mask': bytes_list_feature(encoded_masks),
'image/object/class/label': int64_list_feature(labels)}
我們知道,寫入 tfrecord 的文件內(nèi)容包括:原始圖像的寬高,圖像保存名,圖像本身,圖像格式,目標邊框(boundingbox),掩模(mask),以及類標號(label)等。而且還需要注意的是:boungingbox 必須是正規(guī)化坐標(除以圖像寬或高,0-1 之間取值)。到此,訓練數(shù)據(jù)準備完畢,進入訓練時間。
二、訓練 Mask R-CNN 模型
???????訓練過程完全類似文章 TensorFlow 訓練自己的目標檢測器,只需要下載一個 Mask R-CNN 的預訓練模型,以及配置一下訓練參數(shù)即可。預訓練模型下載請前往鏈接 Mask R-CNN 預訓練模型,使用預訓練模型(以及下面精調(diào)后的模型)可以參考 官方示例代碼,注意要將 sess 當作函數(shù) run_inference_for_single_image 的參數(shù)傳入,否則預測每幅圖像都要重新生成會話會消耗大量時間,此時原來的 for 循環(huán)應該這么寫:
def run_inference_for_single_image(image, graph, sess):
with graph.as_default():
# Get handles to input and output tensors
ops = tf.get_default_graph().get_operations()
...
return return output_dict
with tf.Session(graph=detection_graph) as sess:
for image_path in TEST_IMAGE_PATHS:
...
output_dict = run_inference_for_single_image(image_np, detection_graph, sess)
???????接下來是使用預訓練模型和第一步標注的數(shù)據(jù)來微調(diào)模型了,假如下載的預訓練模型是 mask_rcnn_inception_v2_coco,那么復制
TensorFlow models/research/object_detection/samples/configs/mask_rcnn_inception_v2_coco.config
文件到你的訓練項目下,將其中的 num_classes : 90 改為你要檢測的目標總類目數(shù),比如,因為這里我們只檢測 阿比西尼亞貓 一個類,所以改為 1。另外,還需要將該文件中 5 個 PATH_TO_BE_CONFIGURED 改為相應文件的路徑,詳情參考文章 TensorFlow 訓練自己的目標檢測器。
???????所有配置都完成后,在 models/research/object_detection 目錄的終端下執(zhí)行:
$ python3 model_main.py \
--model_dir=path/to/save/directory \
--pipeline_config_path=path/to/mask_rcnn_inception_v2_xxx.config
開始訓練。訓練成功開始后,新開一個終端,可以使用 tensorboard 在瀏覽器上實時監(jiān)督訓練過程,如果想要提前終止訓練,請用 Ctrl + C 中斷。訓練結束后,模型轉換以及使用請參考文章 TensorFlow 訓練自己的目標檢測器 或者 官方文檔。如果執(zhí)行以上訓練指令時,有 TensorFlow 本身代碼報錯,請使用:
$ sudo pip/pip3 install --upgrade tensorflow-gpu
升級 TensorFlow 版本。除此之外,如果說缺乏 pycocotools 模塊,可以使用
$ sudo pip/pip3 install Cython pycocotools
安裝。
說明:
???????1.有關本文章代碼以及數(shù)據(jù)都在 github,文件夾 datasets/images 下有 232 張 阿比西尼亞貓 的圖像,文件夾 datasets/annotations 下有其中 10 張圖像的 mask 標注數(shù)據(jù),可以將該文件夾一分為二,比如挑選其中 8 張用于生成 train.record,其它 2 張用于生成 val.record。然后,修改文件夾 training 下的配置文件 mask_rcnn_inception_v2_coco.config (假如對應的預訓練模型是 mask_rcnn_inception_v2_coco。其它預訓練模型對應的配置文件請到文件夾 models/research/object_detection/samples/configs 內(nèi)復制),之后就可以開始訓練了。因為只是完整演示整個訓練過程,所以數(shù)據(jù)多少無所謂。
???????2.為了防止訓練過程中出現(xiàn)大概某某 groundtruth 已經(jīng)加入將忽略的問題,請確保訓練集和驗證集圖像的名字不要重名。
【TensorFlow bug】:訓練過程如果報如下錯誤:TypeError: can't pickle dict_values objects,則將 models/research/object_detection/model_lib.py 中第 418 行的 category_index.values() 改成 list(category_index.values()) 即可。