将SynthText数据转成tfrecord

主要涉及两方面：

1. mat文件的读取
2. 构建自己的tfrecord数据集

python读取.mat文件

import scipy.io as sio
mat = sio.loadmat(matdir)

SynthText数据集

这一节讲gt.mat中的数据，可略过不读
下载地址：http://www.robots.ox.ac.uk/~vgg/data/scenetext/
包含858,750张图片
gt.mat中包含imnames,txt,globals,charBB,hearder,version,wordBB等
其中：
mat['imnames'][0] 放图片相对地址
mat['wordBB'][0] 放bbox的位置信息，张量的维度是24图片中包含的word数量，实际操作中一定要小心超出它的值图片大小范围
mat['txt'][0] 放每张图片中包含的文本字符串。注意，它将在相同区域呈现相同字体，颜色，变形等的组合在一起;因此可能与wordBB中对应图片中word数量不一致。
比如：

>>mat['imnames'][0][0]
array(['8/ballet_106_0.jpg'],
      dtype='<U18')

对应下图

ballet_106_0.jpg

>>mat['txt'][0][0]
array(['Lines:\nI lost\nKevin ', 'will                ',
       'line\nand            ', 'and\nthe             ',
       '(and                ', 'the\nout             ',
       'you                 ', "don't\n pkg          "],
      dtype='<U20')
>>mat['wordBB'][0][0].shape
(2,4,15)

因此我们对mat['txt']中的数据要经过strip()去掉空格，re.split()分割后在进行转tfrecord操作,以对mat['txt'][0][0]的处理为例

for val in mat['txt'][0][0]:
     v = [x.encode('ascii') for x in re.split("[ \n]", val.strip()) if x]
     str.extend(v)

代码

import numpy as np
import scipy.io as sio
import os
import re
import Image
import tensorflow as tf
import sys

def arr2list(x):
     ""'confirm every member is in [0.0,1.0]
        convert np.array to a list
    """
    x[x>1] = 1
    x[x<0] = 0
    return list(x)

def int64_feature(value):
    """Wrapper for inserting int64 features into Example proto.
    """
    if not isinstance(value, list):
        value = [value]
    return tf.train.Feature(int64_list=tf.train.Int64List(value=value))


def float_feature(value):
    """Wrapper for inserting float features into Example proto.
    """
    if not isinstance(value, list):
        value = [value]
    return tf.train.Feature(float_list=tf.train.FloatList(value=value))


def bytes_feature(value):
    """Wrapper for inserting bytes features into Example proto.
    """
    if not isinstance(value, list):
        value = [value]
    return tf.train.Feature(bytes_list=tf.train.BytesList(value=value))


def _convert_to_example(image_name, labels, labels_text, bboxes, shape,
                        difficult=0, truncated=0):
    """Build an Example proto for an image example.

    Args:
      image_data: string, JPEG encoding of RGB image;
      labels: list of integers, identifier for the ground truth;
      labels_text: list of text-string;
      bboxes: list of bounding boxes; each box is a list of integers;
          specifying [xmin, ymin, xmax, ymax]. All boxes are assumed to belong
          to the same label as the image label.
      shape: 3 integers, image shapes in pixels.
    Returns:
      Example proto
    """
    image_data = tf.gfile.FastGFile(image_name, 'rb').read()
    xmin = bboxes[0]
    ymin = bboxes[1]
    xmax = bboxes[2]
    ymax = bboxes[3]
    image_format = b'JPEG'
    example = tf.train.Example(features=tf.train.Features(feature={
        'image/height': int64_feature(shape[1]),
        'image/width': int64_feature(shape[0]),
        'image/channels': int64_feature(shape[2]),
        'image/shape': int64_feature(shape),
        'image/object/bbox/xmin': float_feature(xmin),
        'image/object/bbox/xmax': float_feature(xmax),
        'image/object/bbox/ymin': float_feature(ymin),
        'image/object/bbox/ymax': float_feature(ymax),
        'image/object/bbox/label': int64_feature(labels),
        'image/object/bbox/label_text': bytes_feature(labels_text),
        'image/object/bbox/difficult': int64_feature(difficult),
        'image/object/bbox/truncated': int64_feature(truncated),
        'image/format': bytes_feature(image_format),
        'image/encoded': bytes_feature(image_data)}))
    return example


mat_dir = 'gt.mat'
txt_dir = "info.txt"

# get gt.mat
mat = sio.loadmat(mat_dir)
print('load gt.mat')
input("continue")

# get imformation
imnames = mat['imnames'][0]
txt = mat['txt'][0]
wordBB = mat['wordBB'][0]

# set TFrecord dir set
tf_dir_set = set()
tf_writer = None
tf_dirs = "ttfrecords"

total_count = 0


i = 0
need = set()
while i < imnames.size:
    try:
        image_dir = imnames[i][0]
        tfrecord_name = os.path.split(image_dir)[0]
    
        if not tfrecord_name in tf_dir_set:
            # 新开一个tfwriter
            if tf_writer is not None:
                tf_writer.close()
            tf_name = ("%s/synthText_%s.tfrecords" % (tf_dirs, tfrecord_name))
            tf_writer = tf.python_io.TFRecordWriter(tf_name)
            tf_dir_set.add(tfrecord_name)

        img_size = Image.open(image_dir).size
        shape = [img_size[0], img_size[1], 3]
        if len(wordBB[i][0].shape) >1:
            minx = np.amin(wordBB[i][0], axis=0) / img_size[0]
            miny = np.amin(wordBB[i][1], axis=0) / img_size[1]
            maxx = np.amax(wordBB[i][0], axis=0) / img_size[0]
            maxy = np.amax(wordBB[i][1], axis=0) / img_size[1]
        else:
            minx = [np.amin(wordBB[i][0]) / img_size[0]]
            miny = [np.amin(wordBB[i][1]) / img_size[1]]
            maxx = [np.amax(wordBB[i][0]) / img_size[0]]
            maxy = [np.amax(wordBB[i][1]) / img_size[1]]
        #检查是否有>1的情况，并转为list
        minx = arr2list(minx)
        miny = arr2list(miny)
        maxy = arr2list(maxy)
        maxx = arr2list(maxx)
        bboxes = [minx, miny, maxx, maxy]

        str = []
        for val in txt[i]:
            v = [x.encode('ascii') for x in re.split("[ \n]", val.strip()) if x]
            str.extend(v)

        labels = [1] * len(str)

        example = _convert_to_example(image_dir, labels, str, bboxes, shape)
        tf_writer.write(example.SerializeToString())
        sys.stdout.write('\r>> Converting image %d/%d' % (i + 1, imnames.size))
        sys.stdout.flush()
        i = i + 1
        total_count += 1
    except Exception as e:
        print(e)
        choose = input("continue?Y/N")
        if choose == "Y":
            i = i+1
        else:
            sys.exit()
print("Converting image competely! totally %d records" % (total_count))