如何制作将分割数据集封装成tfrecords

import os
from PIL import Image
import tensorflow as tf

IMAGE_SIZE = 224

#数据集的地址,最后的反斜杠别漏掉了
train_label_path=r'/home/user/python/tensorflow_1/FCN.tensorflow-master/Data_zoo/MIT_SceneParsing/ADEChallengeData2016/annotations/validation/'
train_path = r'/home/user/python/tensorflow_1/FCN.tensorflow-master/Data_zoo/MIT_SceneParsing/ADEChallengeData2016/images/validation/'

writer = tf.python_io.TFRecordWriter('fenge_test.tfrecords') #输出成tfrecord文件

#编码方式
def _bytes_feature(value):
    return tf.train.Feature(bytes_list = tf.train.BytesList(value = [value]))

train_label = os.listdir(train_label_path)
train= os.listdir(train_path)
# for img_name in os.listdir(data_path):
#     print(img_name)
#确保名字统一
train_label.sort()
train.sort()

num=1
for ln,lm  in zip(train, train_label):
    img_path = train_path + ln
    label_path = train_label_path+lm
    img = Image.open(img_path)
    lab = Image.open(label_path)
    try:
        #cheke image whether 3input
        r, g, b = img.split()
        #用来验证是否一致
        # aa=os.path.splitext(ln.split('.')[0])[0]
        # bb=os.path.splitext(lm.split('.')[0])[0]
        # print(num,aa==bb)
        # num += 1
        img = img.resize((224, 224))
        lab = lab.resize((224, 224))
        img_raw = img.tobytes()
        lab_raw = lab.tobytes()
        example = tf.train.Example(features=tf.train.Features(feature={"label": _bytes_feature(lab_raw),
                                                                       "img_raw": _bytes_feature(img_raw)
                                                                       }))
        writer.write(example.SerializeToString())
    except ValueError:
        continue
writer.close()
print("finish to write data to tfrecord file!")

 

你可能感兴趣的:(python,tensorflow)