AI challenger 场景分类（2） 读取tfrecord文件

tf.train.batch

# -*- coding: utf-8 -*-
"""
Created on Thu Sep 14 18:02:41 2017

@author: wayne


cifar10 官方样例   https://github.com/tensorflow/models/blob/master/tutorials/image/cifar10/cifar10_input.py
tensorflow cifar_10 代码阅读与理解 http://blog.csdn.net/leastsq/article/details/54374909
TensorFlow高效读取数据的方法  http://blog.csdn.net/u012759136/article/details/52232266
存为多个TFRecord文件 http://blog.csdn.net/xierhacker/article/details/72357651

"""

from PIL import Image
import numpy as np
import matplotlib.pyplot as plt
import tensorflow as tf

def read_and_decode(tfrecords_file, batch_size):  
    '''''read and decode tfrecord file, generate (image, label) batches 
    Args: 
        tfrecords_file: the directory of tfrecord file 
        batch_size: number of images in each batch 
    Returns: 
        image: 4D tensor - [batch_size, width, height, channel] 
        label: 1D tensor - [batch_size] 
    '''  
    # make an input queue from the tfrecord file  

    filename_queue = tf.train.string_input_producer([tfrecord_file])  
    reader = tf.TFRecordReader()  
    _, serialized_example = reader.read(filename_queue)  

    img_features = tf.parse_single_example(  
                                        serialized_example,  
                                        features={  
                                               'label': tf.FixedLenFeature([], tf.int64),  
                                               'h': tf.FixedLenFeature([], tf.int64),
                                               'w': tf.FixedLenFeature([], tf.int64),
                                               'c': tf.FixedLenFeature([], tf.int64),
                                               'image': tf.FixedLenFeature([], tf.string),  
                                               })  

    h = tf.cast(img_features['h'], tf.int32)
    w = tf.cast(img_features['w'], tf.int32)
    c = tf.cast(img_features['c'], tf.int32)

    image = tf.decode_raw(img_features['image'], tf.uint8)  
    image = tf.reshape(image, [h, w, c])

    label = tf.cast(img_features['label'],tf.int32) 
    label = tf.reshape(label, [1])

    ##########################################################  
    # you can put data augmentation here   
#    distorted_image = tf.random_crop(images, [530, 530, img_channel])
#    distorted_image = tf.image.random_flip_left_right(distorted_image)
#    distorted_image = tf.image.random_brightness(distorted_image, max_delta=63)
#    distorted_image = tf.image.random_contrast(distorted_image, lower=0.2, upper=1.8)
#    distorted_image = tf.image.resize_images(distorted_image, (imagesize,imagesize))
#    float_image = tf.image.per_image_standardization(distorted_image)

    image = tf.image.resize_images(image, (224,224))
    image = tf.reshape(image, [224, 224, 3])
    #image, label = tf.train.batch([image, label],  batch_size= batch_size)  

    image_batch, label_batch = tf.train.batch([image, label],  
                                                batch_size= batch_size,  
                                                num_threads= 64,    # 注意多线程有可能改变图片顺序
                                                capacity = 2000)  
    return image_batch, tf.reshape(label_batch, [batch_size])  

def read_tfrecord2(tfrecord_file, batch_size):
    train_batch, train_label_batch = read_and_decode(tfrecord_file, batch_size)

    with tf.Session() as sess:
        coord = tf.train.Coordinator()
        threads = tf.train.start_queue_runners(coord=coord)
        train_batch, train_label_batch = sess.run([train_batch, train_label_batch])
        coord.request_stop()
        coord.join(threads)
    return train_batch, train_label_batch


tfrecord_file = 'ai_challenger_scene_train_20170904/train.tfrecord'
batch_size = 2
train_batch, train_label_batch = read_tfrecord2(tfrecord_file, batch_size)

print(train_batch.shape)
print(train_label_batch)


im = Image.fromarray(np.uint8(train_batch[0,:,:,:])) #参考numpy和图片的互转：http://blog.csdn.net/zywvvd/article/details/72810360
im.show()
im = Image.fromarray(np.uint8(train_batch[1,:,:,:])) #参考numpy和图片的互转：http://blog.csdn.net/zywvvd/article/details/72810360
im.show()

#[16 20 63 10  5]