tensorflow 牌类分类,代码记录
# -*- coding: UTF-8 -*-
import os
import glob
import time
import numpy as np
import tensorflow as tf
import time
import cv2
from skimage import io, transform
from tensorflow.python.framework import graph_util
strIndexDict = {'1hongtao2': 0, '1hongtao3': 1, '1hongtao4': 2, '1hongtao5': 3, '1hongtao6': 4, '1hongtao7': 5, '1hongtao8': 6, '1hongtao9': 7, '1hongtao10': 8, '1hongtaoJ': 9, '1hongtaoQ': 10, '1hongtaoK': 11, '1hongtaoA': 12,
'2heitao2': 13, '2heitao3': 14, '2heitao4': 15, '2heitao5': 16, '2heitao6': 17, '2heitao7': 18, '2heitao8': 19, '2heitao9': 20, '2heitao10': 21, '2heitaoJ': 22, '2heitaoQ': 23, '2heitaoK': 24, '2heitaoA': 25,
'3fangkuai2':26, '3fangkuai3':27, '3fangkuai4':28, '3fangkuai5':29, '3fangkuai6':30, '3fangkuai7':31, '3fangkuai8':32, '3fangkuai9':33, '3fangkuai10':34, '3fangkuaiJ':35, '3fangkuaiQ':36, '3fangkuaiK':37, '3fangkuaiA':38,
'4yinghua2':39, '4yinghua3':40, '4yinghua4':41, '4yinghua5':42, '4yinghua6':43, '4yinghua7':44, '4yinghua8':45, '4yinghua9':46, '4yinghua10':47, '4yinghuaJ':48, '4yinghuaQ':49, '4yinghuaK':50, '4yinghuaA':51,
'5dawang': 52, '5xiaowang': 53}
np.set_printoptions(threshold=np.inf)
def convert_rgbtogray(img):
width, height, channel = img.shape()
grayImg = np.zeros((width, height))
for row in range(height):
for col in range(width):
R = img[col, row, 2]
G = img[col, row, 1]
B = img[col, row, 0]
grayImg[col, row] = R*0.299 + G*0.587 + B*0.114
return grayImg
# 读取图片
def read_img(path, w, h):
cate = [path + x for x in os.listdir(path) if os.path.isdir(path + x)]
# print(cate)
imgs = []
labels = []
print('Start read the image ...')
for index, folder in enumerate(cate):
# print(index, folder)
indexStr = folder[folder.rfind('/')+1:]
idx = strIndexDict[indexStr]
for im in glob.glob(folder + '/*.jpg'):
# print('Reading The Image: %s' % im)
img = cv2.imread(im, cv2.IMREAD_GRAYSCALE)
img = img * (1.0 / 255);
imgs.append(img)
labels.append(idx)
print('Finished ...')
return np.asarray(imgs, np.float32), np.asarray(labels, np.float32)
def read_img_data(filePathName):
imgs = []
labels = []
file = open(filePathName)
while 1:
linedata = file.readline()
if linedata == '':
break
dList = linedata.split(",")
labels.append(int(dList[0]))
img = []
for i in range(len(dList)):
if i == 0 or i> 1120:
continue
d = dList[i]
img.append(float(d))
#print(i, d)
imgs.append(img)
return np.asarray(imgs, np.float32), np.asarray(labels, np.int32)
def get_img_paths(path):
cate = [path + x for x in os.listdir(path) if os.path.isdir(path + x)]
# print(cate)
imgpaths = []
labels = []
print('Start read the image ...')
for index, folder in enumerate(cate):
# print(index, folder)
indexStr = folder[folder.rfind('/') + 1:]
idx = strIndexDict[indexStr]
for im in glob.glob(folder + '/*.jpg'):
imgpaths.append(im)
labels.append(idx)
print('Finished ...')
return imgpaths, labels
# 打乱顺序
def shuffleOrder(data, label):
num_example = data.shape[0]
arr = np.arange(num_example)
np.random.shuffle(arr)
data = data[arr]
label = label[arr]
return data, label
# 将所有数据分为训练集和验证集
def segmentation(data, label, ratio=0.8):
num_example = data.shape[0]
s = np.int(num_example * ratio)
x_train = data[:s]
y_train = label[:s]
x_val = data[s:]
y_val = label[s:]
return x_train, y_train, x_val, y_val
# 定义一个函数,按批次取数据
def minibatches(inputs=None, targets=None, batch_size=None, shuffle=False):
assert len(inputs) == len(targets)
if shuffle:
indices = np.arange(len(inputs))
np.random.shuffle(indices)
for start_idx in range(0, len(inputs) - batch_size + 1, batch_size):
if shuffle:
excerpt = indices[start_idx:start_idx + batch_size]
else:
excerpt = slice(start_idx, start_idx + batch_size)
# 生成label矩阵
targets2 = targets[excerpt]
# reshape inputs
inputs2 = inputs[excerpt].reshape(batch_size, 20*56*1)
yield inputs2, targets2
def reshapebatchs(inputs, labels):
# 重新生成形状
inputs_ = inputs.reshape((-1, 20*56*1))
# 生成label矩阵
labels_ = []
for labels_i in range(len(labels)):
labels_line = []
for idx in range(54):
if labels[labels_i] == idx:
labels_line.append(1)
else:
labels_line.append(0)
labels_.append(labels_line)
return inputs_, np.asarray(labels_, np.int32)
# weight initialization
def weight_variable(shape):
initial = tf.truncated_normal(shape, stddev=0.1, seed=1)
return tf.Variable(initial)
def bias_variable(shape):
initial = tf.constant(0.1, shape = shape)
return tf.Variable(initial)
# convolution
def conv2d(x, W):
return tf.nn.conv2d(x, W, strides=[1, 1, 1, 1], padding='SAME')
# pooling
def max_pool_2x2(x):
return tf.nn.max_pool(x, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='SAME')
def loss(logits,label_batches):
cross_entropy = tf.nn.sparse_softmax_cross_entropy_with_logits(logits=logits,labels=label_batches)
cost = tf.reduce_mean(cross_entropy)
return cost
def startTrain():
# start tensorflow interactiveSession
sess = tf.InteractiveSession()
# Create the model
# placeholder
x = tf.placeholder("float", [None, 1120], name='input')
y_ = tf.placeholder("float", [None, 54])
# first convolutinal layer
w_conv1 = weight_variable([5, 5, 1, 32])
b_conv1 = bias_variable([32])
x_image = tf.reshape(x, [-1, 56, 20, 1])
h_conv1 = tf.nn.relu(conv2d(x_image, w_conv1) + b_conv1)
h_pool1 = max_pool_2x2(h_conv1)
# second convolutional layer
w_conv2 = weight_variable([5, 5, 32, 64])
b_conv2 = bias_variable([64])
h_conv2 = tf.nn.relu(conv2d(h_pool1, w_conv2) + b_conv2)
h_pool2 = max_pool_2x2(h_conv2)
# densely connected layer
w_fc1 = weight_variable([14 * 5 * 64, 1024])
b_fc1 = bias_variable([1024])
h_pool2_flat = tf.reshape(h_pool2, [-1, 14 * 5 * 64])
h_fc1 = tf.nn.relu(tf.matmul(h_pool2_flat, w_fc1) + b_fc1)
# dropout
keep_prob = tf.placeholder("float", name="keep_prob")
h_fc1_drop = tf.nn.dropout(h_fc1, keep_prob)
# readout layer
w_fc2 = weight_variable([1024, 54])
b_fc2 = bias_variable([54])
y_conv = tf.nn.softmax(tf.matmul(h_fc1_drop, w_fc2) + b_fc2, name="softmax")
# train and evaluate the model
# loss = tf.losses.softmax_cross_entropy(y_, y_conv)
cross_entropy = -tf.reduce_sum(y_ * tf.log(y_conv))
train_step = tf.train.GradientDescentOptimizer(1e-4).minimize(cross_entropy)
argmaxy = tf.argmax(y_conv, 1, name="output")
argmaxy_ = tf.argmax(y_, 1)
correct_prediction = tf.equal(argmaxy, argmaxy_)
accuracy = tf.reduce_mean(tf.cast(correct_prediction, "float"))
# defind saver
saver = tf.train.Saver(max_to_keep=100)
# get data
imgpath = './charSamples/'
w = 20
h = 56
c = 1
ratio = 0.8 # 选取训练集的比例
data, label = read_img(path=imgpath, w=w, h=h)
#data, label = read_img_data(imgpath)
data, label = shuffleOrder(data=data, label=label)
x_train, y_train, x_val, y_val = segmentation(data=data, label=label, ratio=ratio)
n_epoch = 200
batch_size = 128
sess.run(tf.initialize_all_variables())
# 加载上次的数据
latestStep = 0
latestPath = tf.train.latest_checkpoint('model/')
if latestPath != None:
latestStep = int(latestPath[latestPath.rfind("-") + 1:])
latestMeta = latestPath + ".meta"
print(" latestPath %s, latestStep %d, latestMeta %s" % (latestPath, latestStep, latestMeta))
saver_r = tf.train.import_meta_graph(latestMeta)
saver_r.restore(sess, latestPath)
for i in range(n_epoch):
i = i + 1 + latestStep
t_start = time.clock()
for batch0_, batch1_ in minibatches(x_train, y_train, batch_size, shuffle=True):
batch0, batch1 = reshapebatchs(batch0_, batch1_)
train_step.run(feed_dict={x: batch0, y_: batch1, keep_prob: 0.5})
train_accuracy = accuracy.eval(feed_dict={x: batch0, y_: batch1, keep_prob: 1.0})
losses = cross_entropy.eval(feed_dict={x: batch0, y_: batch1, keep_prob: 1.0})
t_spent = time.clock() - t_start
print("step %d, train accuracy %g, loss: %g, spent %f" % (i, train_accuracy, losses, t_spent))
if i % 10 == 0:
saver.save(sess, "model/poker", global_step=i)
print("save model")
x_val_, y_val_ = reshapebatchs(x_val, y_val)
print("test accuracy %g" % accuracy.eval(feed_dict={x: x_val_, y_: y_val_, keep_prob: 1.0}))
constant_graph = graph_util.convert_variables_to_constants(sess, sess.graph_def, ["softmax"])
with tf.gfile.FastGFile("model/poker.pb", mode='wb') as f:
f.write(constant_graph.SerializeToString()
def predict2(img):
img.resize((1, 1120))
img = img * (1.0 / 255)
sess = tf.Session()
latestPath = tf.train.latest_checkpoint('model/')
if latestPath != None:
latestMeta = latestPath + ".meta"
print(" latestPath %s, latestMeta %s" % (latestPath, latestMeta))
saver_r = tf.train.import_meta_graph(latestMeta)
saver_r.restore(sess, latestPath)
# 获取placeholder变量
x = sess.graph.get_tensor_by_name("input:0")
print(x)
keep_prob = sess.graph.get_tensor_by_name("keep_prob:0")
print(keep_prob)
out_softmax = sess.graph.get_tensor_by_name("softmax:0")
print(out_softmax)
out_label = sess.graph.get_tensor_by_name("output:0")
print(out_label)
img_out_softmax = sess.run(out_softmax, feed_dict={x: img, keep_prob:1.0})
print("img_out_softmax:", img_out_softmax)
prediction_labels = np.argmax(img_out_softmax, axis=1)
print("label:", prediction_labels)