TensorFlow提供了TFRecords的格式來統一存儲數據,理論上,TFRecords可以存儲任何形式的數據。
TFRecords文件中的數據都是通過tf.train.Example Protocol Buffer的格式存儲的。以下的代碼給出了tf.train.Example的定義。
message Example { Features features = 1; }; message Features { map<string, Feature> feature = 1; }; message Feature { oneof kind { BytesList bytes_list = 1; FloatList float_list = 2; Int64List int64_list = 3; } }; 下面將介紹如何生成和讀取tfrecords文件:
首先介紹tfrecords文件的生成,直接上代碼:
from random import shuffle import numpy as np import glob import tensorflow as tf import cv2 import sys import os # 因為我裝的是CPU版本的,運行起來會有'warning',解決方法入下,眼不見為凈~ os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2' shuffle_data = True image_path = '/path/to/image/*.jpg' # 取得該路徑下所有圖片的路徑,type(addrs)= list addrs = glob.glob(image_path) # 標簽數據的獲得具體情況具體分析,type(labels)= list labels = ... # 這里是打亂數據的順序 if shuffle_data: c = list(zip(addrs, labels)) shuffle(c) addrs, labels = zip(*c) # 按需分割數據集 train_addrs = addrs[0:int(0.7*len(addrs))] train_labels = labels[0:int(0.7*len(labels))] val_addrs = addrs[int(0.7*len(addrs)):int(0.9*len(addrs))] val_labels = labels[int(0.7*len(labels)):int(0.9*len(labels))] test_addrs = addrs[int(0.9*len(addrs)):] test_labels = labels[int(0.9*len(labels)):] # 上面不是獲得了image的地址么,下面這個函數就是根據地址獲取圖片 def load_image(addr): # A function to Load image img = cv2.imread(addr) img = cv2.resize(img, (224, 224), interpolation=cv2.INTER_CUBIC) img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) # 這里/255是為了將像素值歸一化到[0,1] img = img / 255. img = img.astype(np.float32) return img # 將數據轉化成對應的屬性 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])) def _float_feature(value): return tf.train.Feature(float_list=tf.train.FloatList(value=[value])) # 下面這段就開始把數據寫入TFRecods文件 train_filename = '/path/to/train.tfrecords' # 輸出文件地址 # 創建一個writer來寫 TFRecords 文件 writer = tf.python_io.TFRecordWriter(train_filename) for i in range(len(train_addrs)): # 這是寫入操作可視化處理 if not i % 1000: print('Train data: {}/{}'.format(i, len(train_addrs))) sys.stdout.flush() # 加載圖片 img = load_image(train_addrs[i]) label = train_labels[i] # 創建一個屬性(feature) feature = {'train/label': _int64_feature(label), 'train/image': _bytes_feature(tf.compat.as_bytes(img.tostring()))} # 創建一個 example protocol buffer example = tf.train.Example(features=tf.train.Features(feature=feature)) # 將上面的example protocol buffer寫入文件 writer.write(example.SerializeToString()) writer.close() sys.stdout.flush()
新聞熱點
疑難解答