import os
import tensorflow as tf
from PIL import Image
import numpy as np

data_dir = 'train/'
classes={'不导电','其他','凸粉','擦花','无瑕疵样本','桔皮','横条压凹','涂层开裂','漏底','碰伤','脏点','起坑'}
writer= tf.python_io.TFRecordWriter("train.tfrecords")

for index,name in enumerate(classes):
    class_path = data_dir + name
    for img_name in os.listdir(class_path):
        img_path = os.path.join(class_path, img_name)
        img = Image.open(img_path)
        img = img.resize((224,224))
        img_raw = img.tobytes()
        example = tf.train.Example(features=tf.train.Features(feature={
            "label": tf.train.Feature(int64_list=tf.train.Int64List(value=[index])),
            'img_raw': tf.train.Feature(bytes_list=tf.train.BytesList(value=[img_raw]))
        }))
        writer.write(example.SerializeToString())

writer.close()


'''

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]))

for index,name in enumerate(classes):
    class_path = data_dir + name
    for img_name in os.listdir(class_path):
        img_path = os.path.join(class_path, img_name)
        img = Image.open(img_path)
        img = img.resize((224,224))
        img_raw = img.tobytes()
        features=tf.train.Features(feature={
            "label": _int64_feature(index),
            'img_raw': _bytes_feature(img_raw)
        })
        example = tf.train.Example(features=feature)

        writer.write(example.SerializeToString())

writer.close()
'''