cnn汉字识别 tensorflow demo

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147

148

149

150

151

152

153

154

155

156

157

158

159

160

161

162

163

164

165

166

167

168

169

170

171

172

173

174

175

# -*- coding: utf-8 -*-
 
 
import tensorflow as tf
import os
import random
import tensorflow.contrib.slim as slim
import time
import numpy as np
import pickle
from PIL import Image
 
 
mode = "inference"
char_size = 3755
epochs = 5
batch_size = 128
checkpoint_dir = '/aiml/code/'
#train_data_dir = 'D:/Yang/softwares/Spider_ws/WordRecognition/data/train/'
#test_data_dir = 'D:/Yang/softwares/Spider_ws/WordRecognition/data/test/'
 
class DataIterator:
    def __init__(self, data_dir):
        self.image_names = []
        for root, sub_folder, file_list in os.walk(data_dir):
            self.image_names += [os.path.join(root, file_path) for file_path in file_list]
        random.shuffle(self.image_names)
        self.labels = [int(file_name[len(data_dir):].split(os.sep)[0]) for file_name in self.image_names]
 
    @property
    def size(self):
        return len(self.labels)
 
    def input_pipeline(self, batch_size, num_epochs=None):
        images_tensor = tf.convert_to_tensor(self.image_names, dtype=tf.string)
        labels_tensor = tf.convert_to_tensor(self.labels, dtype=tf.int64)
        input_queue = tf.train.slice_input_producer([images_tensor, labels_tensor], num_epochs=num_epochs)
 
        labels = input_queue[1]
        images_content = tf.read_file(input_queue[0])
        images = tf.image.convert_image_dtype(tf.image.decode_png(images_content, channels=1), tf.float32)
        new_size = tf.constant([64, 64], dtype=tf.int32)
        images = tf.image.resize_images(images, new_size)
        image_batch, label_batch = tf.train.shuffle_batch([images, labels], batch_size=batch_size, capacity=50000,
                                                          min_after_dequeue=10000)
        return image_batch, label_batch
 
 
def build_graph(top_k):
    # with tf.device('/cpu:0'):
    images = tf.placeholder(dtype=tf.float32, shape=[None, 64, 64, 1], name='input_image')
    labels = tf.placeholder(dtype=tf.int64, shape=[None], name='label_batch')
 
    conv_1 = slim.conv2d(images, 64, [3, 3], 1, padding='SAME', scope='conv1')
    max_pool_1 = slim.max_pool2d(conv_1, [2, 2], [2, 2], padding='SAME')
    conv_2 = slim.conv2d(max_pool_1, 128, [3, 3], padding='SAME', scope='conv2')
    max_pool_2 = slim.max_pool2d(conv_2, [2, 2], [2, 2], padding='SAME')
    conv_3 = slim.conv2d(max_pool_2, 256, [3, 3], padding='SAME', scope='conv3')
    max_pool_3 = slim.max_pool2d(conv_3, [2, 2], [2, 2], padding='SAME')
     
    flatten = slim.flatten(max_pool_3)
    fc1 = slim.fully_connected(flatten, 1024, activation_fn=tf.nn.tanh, scope='fc1')
    logits = slim.fully_connected(fc1, char_size, activation_fn=None, scope='output_logit')
     
    loss = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(logits=logits, labels=labels))
    accuracy = tf.reduce_mean(tf.cast(tf.equal(tf.argmax(logits, 1), labels), tf.float32))
 
    global_step = tf.get_variable("step", [], initializer=tf.constant_initializer(0.0), trainable=False)
    rate = tf.train.exponential_decay(2e-4, global_step, decay_steps=2000, decay_rate=0.97, staircase=True)
    train_op = tf.train.AdamOptimizer(learning_rate=rate).minimize(loss, global_step = global_step)
    probabilities = tf.nn.softmax(logits)
    pred = tf.identity(probabilities, name = 'prediction')
     
    return {'images': images,
            'labels': labels,
            'global_step': global_step,
            'train_op': train_op,
            'loss': loss,
            'accuracy': accuracy}
 
def train():
    train_feeder = DataIterator(data_dir=train_data_dir)
    test_feeder = DataIterator(data_dir=test_data_dir)
    with tf.Session() as sess:
        train_images, train_labels = train_feeder.input_pipeline(batch_size)
        test_images, test_labels = test_feeder.input_pipeline(batch_size)
        graph = build_graph(top_k=1)
        sess.run(tf.global_variables_initializer())
        coord = tf.train.Coordinator()
        threads = tf.train.start_queue_runners(sess=sess, coord=coord)
        saver = tf.train.Saver()
 
        print (':::Training Start:::')
        try:
            while not coord.should_stop():
                start_time = time.time()
                train_images_batch, train_labels_batch = sess.run([train_images, train_labels])
                feed_dict = {graph['images']: train_images_batch,
                             graph['labels']: train_labels_batch}
                _, loss_val, step = sess.run(
                    [graph['train_op'], graph['loss'], graph['global_step']],
                    feed_dict=feed_dict)
                end_time = time.time()
                if step % 10 == 1:
                    print ("the step {0} takes {1} loss {2}".format(step, end_time - start_time, loss_val))
                if step > 200000:
                    break
                if step % 50 == 1:
                    test_images_batch, test_labels_batch = sess.run([test_images, test_labels])
                    feed_dict = {graph['images']: test_images_batch,
                                 graph['labels']: test_labels_batch}
                    accuracy_test = sess.run(
                        graph['accuracy'],
                        feed_dict=feed_dict)
                    print ('===============Eval a batch=======================')
                    print ('the step {0} test accuracy: {1}'.format(step, accuracy_test))
                    print ('===============Eval a batch=======================')
                if step % 200 == 1:
                    print ('Save the ckpt of {0}'.format(step))
                    saver.save(sess, os.path.join(checkpoint_dir, 'my-model'),
                               global_step=graph['global_step'])
        except tf.errors.OutOfRangeError:
            print ('==================Train Finished================')
            saver.save(sess, os.path.join(checkpoint_dir, 'my-model'), global_step=graph['global_step'])
        finally:
            coord.request_stop()
        coord.join(threads)
 
def new_inference(predict_dir):
    saver = tf.train.import_meta_graph( checkpoint_dir + "my-model-164152.meta", clear_devices=True)
    image_list = []
    new_file_list = []
    for root, _, file_list in os.walk(predict_dir):
        new_file_list += [file for file in file_list if ".nfs" not in file]
        new_file_list.sort(key= lambda x:int(x[:-4]))
        for file in new_file_list:
#            print (new_file_list)
            image = os.path.join(root, file)         
            temp_image = Image.open(image).convert('L')
            temp_image = temp_image.resize((64, 64), Image.ANTIALIAS)           
            temp_image = np.asarray(temp_image) / 255.0
            image_list.append(temp_image)
    image_list = np.asarray(image_list)       
    temp_image = image_list.reshape([len(new_file_list), 64, 64, 1])
    with tf.Session() as sess: 
        saver.restore(sess, checkpoint_dir + "my-model-164152") #读入模型参数
        graph = tf.get_default_graph()
        op = graph.get_tensor_by_name("prediction:0")
        input_tensor = graph.get_tensor_by_name('input_image:0')
        probs = sess.run(op,feed_dict = {input_tensor:temp_image})
        result = []
        for word in probs:
            result.append(np.argsort(-word)[:3])
        return result
         
def main():
 
    if mode == "train":
        train()
    if mode == 'inference':
        word_dict = pickle.load(open("/aiml/code/word_dict", "rb"))
        image_path = '/aiml/data/'
        index = new_inference(image_path)
        file = open("/aiml/result/result.txt", "w")        
#        print ("预测文字为: ")
        pred_list = []
        for i in index:
#            print ("最大几率三个:")
#            print (word_dict[str(i[0])],word_dict[str(i[1])],word_dict[str(i[2])])                
            pred_list.append(word_dict[str(i[0])])
            file.write(word_dict[str(i[0])])
         
if __name__ == "__main__":
#    tf.app.run()
    main()