TensorFlow2学习22、手工标注数据并生成TFRecords Format格式

一、说明

本文实验环境：

• TF2.0
• ubuntu16.04
• python3.6
• LabelImg
• pandas

二、准备图片

为了简单起见，我下载的Kaggle猫狗大战图片。
将图片分成2个文件夹存放，分别为images/test,images/train
我这里只放了少量图片，实际使用中需要准备大量的图片。

三、软件设置

1. 下载labelImage软件

https://github.com/tzutalin/labelImg/releases

四、软件使用

打开labelImg，选择OpenDir，指向自己的Images文件夹

选择Create RectBox手工标注

• 操作时，LabelImg可能只能画正方形。可以按alt键切换成长方形。但标注一个后，后面的就可以自动支持长方形标注了。
• 标注后，可以右键选择Edit RectBox修改
• 使用alt可以标注多个目标
• 标注后，在弹出的标签里输入或选择标签

标注好一个后，点击Save-Save，FileName后面手工加上.xml扩展名（可能这个版本的bug，文件名里带.的保存时会被移除）

常用快捷键

• ctrl + s 保存
• ctrl + d 复制标签和矩形框
• 空格 标记当前图片已标记
• w 创建一个矩形
• d 下一张图片
• a 上一张图片

标注生成示例：

五、把生成的xml汇总到.csv中

使用下面的脚本

import os
import glob
import pandas as pd
import xml.etree.ElementTree as ET

path = '/home/admin1/testpython/models/research/object_detection/images/test'
target = 'test.csv'
os.chdir(path)


def xml_to_csv(path):
    xml_list = []
    for xml_file in glob.glob(path + '/*.xml'):
        tree = ET.parse(xml_file)
        root = tree.getroot()
        for member in root.findall('object'):
            value = (root.find('filename').text,
                     int(root.find('size')[0].text),
                     int(root.find('size')[1].text),
                     member[0].text,
                     int(member[4][0].text),
                     int(member[4][1].text),
                     int(member[4][2].text),
                     int(member[4][3].text)
                     )
            xml_list.append(value)
    column_name = ['filename', 'width', 'height', 'class', 'xmin', 'ymin', 'xmax', 'ymax']
    xml_df = pd.DataFrame(xml_list, columns=column_name)
    return xml_df


def main():
    image_path = path
    xml_df = xml_to_csv(image_path)
    xml_df.to_csv(target, index=None)
    print('Successfully converted xml to csv.')


main()

生成示例：

filename,width,height,class,xmin,ymin,xmax,ymax
dog.12450.jpg,276,397,dog,50,34,242,392
dog.12451.jpg,373,500,dog,14,10,372,488
cat.12409.jpg,500,424,cat,18,30,426,418
cat.12412.jpg,499,334,cat,135,8,353,332
dog.12456.jpg,499,435,dog,36,19,456,435
cat.12415.jpg,500,419,cat,44,23,500,409
cat.12413.jpg,499,375,cat,144,140,273,287
cat.12413.jpg,499,375,cat,237,83,329,147
dog.12459.jpg,249,199,dog,31,1,232,199
cat.12407.jpg,362,499,cat,5,11,358,488
cat.12408.jpg,500,438,cat,12,23,473,401
cat.12414.jpg,500,374,cat,19,55,481,265
cat.12411.jpg,398,299,cat,35,2,390,269
cat.12416.jpg,387,500,cat,14,19,343,478
cat.12410.jpg,341,499,cat,24,2,292,482
dog.12458.jpg,185,339,dog,22,1,180,330
dog.12455.jpg,469,453,dog,57,28,423,451
dog.12457.jpg,500,442,dog,49,7,471,431
dog.12453.jpg,432,419,dog,73,44,361,417
dog.12452.jpg,500,367,dog,19,14,470,346
dog.12454.jpg,304,270,dog,105,5,213,144
dog.12454.jpg,304,270,dog,51,101,266,258

六、从csv创建tfrecord格式

脚本：

# -*- coding: utf-8 -*-


"""
Usage:
  # From tensorflow/models/
  # Create train data:
  python3 generate_tfrecord.py --csv_input=test/test.csv  --output_path=train.record
  # Create test data:
  python3 generate_tfrecord.py --csv_input=train/train.csv  --output_path=test.record
"""


import os
import io
import pandas as pd
import tensorflow as tf

from PIL import Image
from object_detection.utils import dataset_util
from collections import namedtuple, OrderedDict

os.chdir('/home/admin1/testpython/models/research/object_detection/images')

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


def class_text_to_int(row_label):
    if row_label == 'dog':     # 改为自己的分类
        return 1
    elif row_label=='cat':
        return 0
    else:
        return None


def split(df, group):
    data = namedtuple('data', ['filename', 'object'])
    gb = df.groupby(group)
    return [data(filename, gb.get_group(x)) for filename, x in zip(gb.groups.keys(), gb.groups)]


def create_tf_example(group, path):
    with tf.gfile.GFile(os.path.join(path, '{}'.format(group.filename)), 'rb') as fid:
        encoded_jpg = fid.read()

    encoded_jpg_io = io.BytesIO(encoded_jpg)
    image = Image.open(encoded_jpg_io)
    width, height = image.size
    filename = group.filename.encode('utf8')
    image_format = b'jpg'
    xmins = []
    xmaxs = []
    ymins = []
    ymaxs = []
    classes_text = []
    classes = []

    for index, row in group.object.iterrows():
        xmins.append(row['xmin'] / width)
        xmaxs.append(row['xmax'] / width)
        ymins.append(row['ymin'] / height)
        ymaxs.append(row['ymax'] / height)
        classes_text.append(row['class'].encode('utf8'))
        classes.append(class_text_to_int(row['class']))
    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_jpg),
        '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)
    path = os.path.join(os.getcwd(), 'test')         #  可以打印出path看看值对不对
    examples = pd.read_csv(FLAGS.csv_input)

    grouped = split(examples, 'filename')
    for group in grouped:

        tf_example = create_tf_example(group, path)
        print(tf_example.SerializeToString())
        writer.write(tf_example.SerializeToString())

    writer.close()
    output_path = os.path.join(os.getcwd(), FLAGS.output_path)
    print('Successfully created the TFRecords: {}'.format(output_path))


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

要点：

• 32行 class_text_to_int要将自己分类都用数值转换，需要手工调整程序

运行结果：