1. 程式人生 > >【CV】如何使用Tensorflow提供的Object Detection API --2--資料轉換為TFRecord格式

【CV】如何使用Tensorflow提供的Object Detection API --2--資料轉換為TFRecord格式

本篇主要講的是如何將現存的資料變成Tensorflow記錄格式,然後我們就可以用這些資料來進行微調模型,以解決我們關心的問題了。

什麼是TFRecord格式

一般使用TF讀取資料有四種方式:

  • 預先把所有資料載入進記憶體
  • 在每輪訓練中使用原生Python程式碼讀取一部分資料,然後使用feed_dict輸入到計算圖
  • 利用Threading和Queues從TFRecord中分批次讀取資料
  • 使用Dataset API

後面兩種都用的是TFRecord格式。

Tensorflow物體檢測API需要訓練資料格式是TFRecord格式。

TFRecord

格式會將所有的標記,比如邊界框以及原始影象打包成一個檔案,建立TFRecord格式的檔案會有一些麻煩,但是一旦建立好以後,使用起來非常方便。

建立單個TFRecord檔案

TF給了我們一個樣例程式:

def create_tf_example(label_and_data_info):
	height = None
	width = None
	filename = None
	encoded_image_data = None # encoded image bytes
	image_format = None # b'jpeg' or b'png'
	# 儲存標準化以後的資料,一個框一個值
xmins = [] xmaxs = [] ymins = [] ymaxs = [] classes_text = [] # 邊界框的類別名 classes = [] # 邊界框的數字編號 # TODO tf_label_and_data = tf.train.Example(features=tf.train.Features(feature={ 'image/height':dataset_util_int64_feature(height), 'image/width': dataset_util_int64_feature(width), 'image/filename'
: dataset_util.bytes_feature(filename), 'image/source_id': dataset_util.bytes_feature(filename), 'image/encoded': dataset_util.bytes_feature(encoded_image_data), 'image/format': dataset_util.bytes_feature(image_format), 'image/object/bbox/xmin': dataset_util.float_list_feature(xmins), 'image/object/bbox/xmax': dataset_util.float_list_feature(xmaxs), 'image/object/bbox/ymin': dataset_util.float_list_feature(ymins), 'image/object/bbox/ymax': dataset_util.float_list_feature(ymaxs), 'image/object/class/text': dataset_util.bytes_list_feature(classes_text), 'image/object/class/label': dataset_util.int64_list_feature(classes), })) return tf_label_and_data

這個函式就是從單張圖片中提取資訊,注意到在方框資訊和類別資訊之外,我們還需要提供編碼的圖片資料,這個可以使用tf.gfile.GFile()函式來完成。

建立完整的TFRecord檔案

import tensorflow as tf
from object_detection.utils import dataset_util

flags = tf.app.flags
# 設定一個輸入命令列引數
flags.DEFINE_string('output_path', '', 'Path to output TFRecord')
FLAGS = flags.FLAGS

def create_tf_example(data_and_label_info):
	...
	return tf_data_and_label
def main():
	writer = tf.python_io.TFRecordWriter(FLAGS.output_path)
	# TODO
	file_loc = None
	all_data_and_label_info = LOAD(file_loc)
	
	for data_and_label_info in all_data_and_label_info:
		tf_example = create_tf_example(data_and_label_info)
		writer.write(tf_example_SerializeToString())
	writer.close()
if __name__ == '__main__':
	tf.app.run()

更具體的,針對Bosch Small Traffic Lights Dataset資料集,看這個老哥的實現,程式碼我拷貝如下:

import tensorflow as tf
import yaml
import os
from object_detection.utils import dataset_util


flags = tf.app.flags
flags.DEFINE_string('output_path', '', 'Path to output TFRecord')
FLAGS = flags.FLAGS

LABEL_DICT =  {
    "Green" : 1,
    "Red" : 2,
    "GreenLeft" : 3,
    "GreenRight" : 4,
    "RedLeft" : 5,
    "RedRight" : 6,
    "Yellow" : 7,
    "off" : 8,
    "RedStraight" : 9,
    "GreenStraight" : 10,
    "GreenStraightLeft" : 11,
    "GreenStraightRight" : 12,
    "RedStraightLeft" : 13,
    "RedStraightRight" : 14
    }

def create_tf_example(example):
    
    # Bosch
    height = 720 # Image height
    width = 1280 # Image width

    filename = example['path'] # Filename of the image. Empty if image is not from file
    filename = filename.encode()

    with tf.gfile.GFile(example['path'], 'rb') as fid:
        encoded_image = fid.read()

    image_format = 'png'.encode() 

    xmins = [] # List of normalized left x coordinates in bounding box (1 per box)
    xmaxs = [] # List of normalized right x coordinates in bounding box
                # (1 per box)
    ymins = [] # List of normalized top y coordinates in bounding box (1 per box)
    ymaxs = [] # List of normalized bottom y coordinates in bounding box
                # (1 per box)
    classes_text = [] # List of string class name of bounding box (1 per box)
    classes = [] # List of integer class id of bounding box (1 per box)

    for box in example['boxes']:
        #if box['occluded'] is False:
        #print("adding box")
        xmins.append(float(box['x_min'] / width))
        xmaxs.append(float(box['x_max'] / width))
        ymins.append(float(box['y_min'] / height))
        ymaxs.append(float(box['y_max'] / height))
        classes_text.append(box['label'].encode())
        classes.append(int(LABEL_DICT[box['label']]))


    tf_example = tf.train.Example(features=tf.train.Features(feature={
        'image/height': dataset_util.int64_feature(height),
        'image/width': dataset_util.int64_feature(width),
        'image/filename': dataset_util.bytes_feature(filename),
        'image/source_id': dataset_util.bytes_feature(filename),
        'image/encoded': dataset_util.bytes_feature(encoded_image),
        'image/format': dataset_util.bytes_feature(image_format),
        'image/object/bbox/xmin': dataset_util.float_list_feature(xmins),
        'image/object/bbox/xmax': dataset_util.float_list_feature(xmaxs),
        'image/object/bbox/ymin': dataset_util.float_list_feature(ymins),
        'image/object/bbox/ymax': dataset_util.float_list_feature(ymaxs),
        'image/object/class/text': dataset_util.bytes_list_feature(classes_text),
        'image/object/class/label': dataset_util.int64_list_feature(classes),
    }))

    return tf_example


def main(_):
    
    writer = tf.python_io.TFRecordWriter(FLAGS.output_path)
    
    # BOSCH
    INPUT_YAML = "data/test-bosch/dataset_test_rgb/test.yaml"
    examples = yaml.load(open(INPUT_YAML, 'rb').read())

    #examples = examples[:10]  # for testing
    len_examples = len(examples)
    print("Loaded ", len(examples), "examples")

    for i in range(len(examples)):
        examples[i]['path'] = os.path.abspath(os.path.join(os.path.dirname(INPUT_YAML), examples[i]['path']))
    
    counter = 0
    for example in examples:
        tf_example = create_tf_example(example)
        writer.write(tf_example.SerializeToString())

        if counter % 10 == 0:
            print("Percent done", (counter/len_examples)*100)
        counter += 1

    writer.close()

if __name__ == '__main__':
    tf.app.run()

其中,yaml檔案樣式是:

- boxes:
  - {label: Green, occluded: false, x_max: 582.3417892052, x_min: 573.3726437481,
    y_max: 276.6271175345, y_min: 256.3114627642}
  - {label: Green, occluded: false, x_max: 517.6267821724, x_min: 510.0276868266,
    y_max: 273.164089267, y_min: 256.4279864221}
  path: ./rgb/train/2015-10-05-16-02-30_bag/720654.png
- boxes: []
  path: ./rgb/train/2015-10-05-16-02-30_bag/720932.png

寫完執行的方式是:

python tf_record.py --output_path training.record

在上面程式碼中出現的命令列引數使用,下面展開細說一下。

tf.app.flags使用

  • tf.app.flags.DEFINE_string函式
  • tf.app.flags.FLAGS變數

flags其實就是對ArgumentParser的一層淺淺的封裝。

使用tf.app.flags.DEFINE_xxx()就是用來新增命令列的可選引數,使用tf.app.flags.FLAGS就是用來從對應的命令列引數取出引數。

其中,xxx可以是:

  • string
  • float
  • int
  • bool

等等。

DEFINE_XXX(flag_name, default_value, doctstring)

這個函式呼叫的是_define_helper函式

def _define_helper(flag_name, default_value, docstring, flagtype):
	_global_parser.add_argument('--' + flag_name,
								default=default_value,
								help=docstring, 
								type=flagtype)

使用示例

import tensorflow as tf
FLAGS = tf.app.flags.FLAGS # 用於提取命令列引數
tf.app.flags.DEFINE_float('FLOAT', 0.01, 'Input a float number')
tf.app.flags.DEFINE_integer('INT',100, 'Input an integer number')
tf.app.flags.DEFINE_boolean('BOOL', True, 'Input a boolean number')

# 直接通過FLAGS物件提取命令列引數
print(FLAGS.FLOAT)
print(FLAGS.INT)
print(FLAGS.BOOL)

END.

參考:

https://blog.csdn.net/m0_37041325/article/details/77448971
https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/platform/flags.py