AI_Scene classification 数据制作成TFrecord 格式

# encoding: utf-8
'''
@author: tourior
@software: python3.6
@file: Tfrecord.py
@time: 2017/10/10 10:35
@desc:
'''

import tensorflow as tf
import numpy as np
import cv2
import json
import glob
import matplotlib.pyplot as plt
import sys


#读json数据
def read_json(json_path):
    data={}
    with open(json_path , 'r',encoding ='utf-8' ) as fp:
        data = json.load(fp)
    return data

#生成整数型的属性。
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]))

#读入图片并且resize成224,224,3 , 然后将图片转换成字符串返回
def read_image(img_path):
    img = plt.imread(img_path) # read pic
    img = cv2.resize(img,(224,224),interpolation=cv2.INTER_CUBIC)#resize 224,224,3
    return img.tostring()    #转成成字符串返回


def view_bar(num, total):
    rate = float(num) / float(total)
    rate_num = int(rate * 100)
    r = '\r[%s%s]%d%%,%d' % ("="*rate_num, " "*(100-rate_num), rate_num, num )
    sys.stdout.write(r)
    sys.stdout.flush()


# 保存进tfrecord
def json_data2tfrecord(json_data, base_path, tfrecord_name):
    total = len(json_data)
    count = 0
    writer = tf.python_io.TFRecordWriter(str(tfrecord_name) + '.tfrecords')
    for data in json_data:
        img_path = base_path + data['image_id']
        raw_image = read_image(img_path)
        label = int(data['label_id'])
        example = tf.train.Example(features=tf.train.Features(feature={
            'image_raw': _bytes_feature(raw_image),
            'label': _int64_feature(label)
        }))
        writer.write(example.SerializeToString())

        view_bar(count, total)
        count = count + 1

    writer.close()
    print('finish')


if __name__=="__main__":
    base_path = 'G:/AI_Scene classification/image_data/validation/scene_validation_images_20170908/'
    data = read_json('G:/AI_Scene classification/image_data/validation/scene_validation_annotations_20170908.json')
    json_data2tfrecord(data , base_path , 'ai_challenger_scene_validation')