import numpy as np
import tensorflow as tf
import time
import os
import cv2
from sklearn.utils import shuffle
# 图片存放位置
PATH_DES = [
r'data_tfrecords/integers_tfrecords/',
r'data_tfrecords/alphabets_tfrecords/',
r'data_tfrecords/Chinese_letters_tfrecords/'
]
PATH_RES = [r'data/integers/',
r'data/alphabets/',
r'data/Chinese_letters/']
PATH = list(zip(PATH_RES, PATH_DES))
# transformation between integer <-> string
# 用于车牌识别时--数字+字母+32省份
integers = {
'0': 0,
'1': 1,
'2': 2,
'3': 3,
'4': 4,
'5': 5,
'6': 6,
'7': 7,
'8': 8,
'9': 9
}
alphabets = {
'A': 10,
'B': 11,
'C': 12,
'D': 13,
'E': 14,
'F': 15,
'G': 16,
'H': 17,
'I': 18,
'J': 19,
'K': 20,
'L': 21,
'M': 22,
'N': 23,
'O': 24,
'P': 25,
'Q': 26,
'R': 27,
'S': 28,
'T': 29,
'U': 30,
'V': 31,
'W': 32,
'X': 33,
'Y': 34,
'Z': 35
}
provinces = {
'藏': 36,
'川': 37,
'鄂': 38,
'甘': 39,
'赣': 40,
'广': 41,
'桂': 42,
'贵': 43,
'黑': 44,
'沪': 45,
'吉': 46,
'冀': 47,
'津': 48,
'晋': 49,
'京': 50,
'辽': 51,
'鲁': 52,
'蒙': 53,
'闽': 54,
'宁': 55,
'青': 56,
'琼': 57,
'陕': 58,
'苏': 59,
'皖': 60,
'湘': 61,
'新': 62,
'渝': 63,
'豫': 64,
'粤': 65,
'云': 66,
'浙': 67
}
label_ref = [
integers,
alphabets,
provinces
]
# 图片信息
IMG_HEIGHT = 28
IMG_WIDTH = 16
IMG_CHANNELS = 1
# NUM_TRAIN = 7000
NUM_VALIDARION = [sum([len(os.listdir(r + i))
for i in os.listdir(r)]) // 5 for r in PATH_RES]
# 读取图片
def read_images(path_res, label_ref, num_validation):
imgs = []
labels = []
path_res_dirs = sorted(os.listdir(path_res))
for i in path_res_dirs:
paths_images = os.listdir(path_res + i) # 本想排序的, 但是字符串排序效果不尽人意.
t_lst = [''.join((path_res, i, '/', t)) for t in paths_images]
paths_images = t_lst.copy()
del t_lst
for j in range(len(paths_images)):
c = 0
img = cv2.imread(paths_images[j], 0)
img_blur = cv2.bilateralFilter(img, 3, 45, 45)
img_current = cv2.resize(img_blur, (28, 28))
ret, img_current_threshed = cv2.threshold(img_current,
127, 255,
cv2.THRESH_OTSU)
h, w = img_current_threshed.shape
t_c = np.array([[img_current_threshed[0][0],
img_current_threshed[0, w-1]],
[img_current_threshed[h-1, 0],
img_current_threshed[h-1, w-1]]])
c = sum([(t_c[0, 0]//255), (t_c[1, 1]//255),
(t_c[0, 1]//255), (t_c[1, 0]//255)])
if_reverse = sum([sum(img_current_threshed[0, :] // 255),
sum(img_current_threshed[:, w-1] // 255),
sum(img_current_threshed[h-1, :] // 255),
sum(img_current_threshed[:, 0] // 255)])\
/ ((h + w) * 2 + 4) > 0.5
# if c >= 1:
# img_current_threshed = 255 - img_current_threshed
if c > 2 or (c > 1 and if_reverse):
img_current_threshed = 255 - img_current_threshed
# img_current_threshed = img_current
label_current = paths_images[j].split("/")[-2]
# if i == '2':
# fig, ax = plt.subplots(1, 2, figsize=(16, 8))
# ax0, ax1 = ax.ravel()
# ax0.imshow(img_current, cmap="gray")
# ax1.imshow(img_current_threshed, cmap="gray")
# plt.title(c)
# # print([img_current_threshed[0][0],
# # img_current_threshed[0, w-1],
# # img_current_threshed[h-1, 0],
# # img_current_threshed[h-1, w-1]])
# plt.show()
imgs.append((img_current_threshed // 255).astype(np.uint8))
labels.append(np.uint8(label_ref[label_current]))
imgs = np.array(imgs)
imgs = imgs.reshape(imgs.shape[0], -1)
labels = np.array(labels)
labels = labels.reshape(labels.shape[0], -1)
data = np.hstack((labels, imgs))
data = shuffle(data)
test_labels = data[:num_validation, 0]
test_images = data[:num_validation, 1:]
train_labels = data[num_validation:, 0]
train_images = data[num_validation:, 1:]
return train_labels, train_images, test_labels, test_images
# 生成整数型的属性
def _int64_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 convert(images, labels, filename):
# 获取要转换为TFRecord文件的图片数目
num = images.shape[0]
print("num:", num)
print("images.shape:", images.shape)
# 输出TFRecord文件的文件名
print('Writting', filename)
# 创建一个writer来写TFRecord文件
writer = tf.python_io.TFRecordWriter(filename)
for i in range(num):
# 将图像矩阵转化为一个字符串
img_raw = images[i].tostring()
# 将一个样例转化为Example Protocol Buffer，并将所有需要的信息写入数据结构
example = tf.train.Example(features=tf.train.Features(feature={
'label': _int64_feature(int(labels[i])),
'image_raw': _bytes_feature(img_raw)}))
# 将example写入TFRecord文件
writer.write(example.SerializeToString())
writer.close()
print('Writting End')
def main():
start_time = time.time()
for i in range(len(PATH)):
print('reading images from {} begin'.format(PATH_RES[i]))
data = read_images(PATH_RES[i], label_ref[i], NUM_VALIDARION[i])
train_labels, train_images, test_labels, test_images = data
# Slice data here.
print('convert to tfrecords into {} begin'.format(PATH_DES[i]))
convert(train_images, train_labels, PATH_DES[i]+"train.tfrecords")
convert(test_images, test_labels, PATH_DES[i]+"test.tfrecords")
duration = time.time() - start_time
print('Converting end , total cost = %d sec' % duration)
if __name__ == '__main__':
main()
