寫在之前：斷斷續續在網上大浪淘沙，個人感覺這段代碼不錯，之前在windows運行出現了一些bug，覺得煩，在Ununtu上稍微改了一下就ok了。有需要的直接拿走~ #Author:Anthony #導入相應的模塊 import tensorflow as tf import os import random import math import sys #劃分驗證集訓練集 _NUM_TEST = 40 #random seed _RANDOM_SEED = 0 #數據塊 _NUM_SHARDS = 2 #數據集路徑 DATASET_DIR = '/home/anthony/文檔/數據集_帶標簽/SHIYAN_SAMEZIZE' #標簽文件 LABELS_FILENAME = '/home/anthony/文檔/數據集_帶標簽/SHIYAN_SAMEZIZE_labels.txt' #定義tfrecord 的路徑和名稱 def _get_dataset_filename(dataset_dir,split_name,shard_id): output_filename = 'image_%s_%05d-of-%05d.tfrecord' % (split_name,shard_id,_NUM_SHARDS) return os.path.join(dataset_dir,output_filename) #判斷tfrecord文件是否存在 def _dataset_exists(dataset_dir): for split_name in ['train','test']: for shard_id in range(_NUM_SHARDS): #定義tfrecord的路徑名字 output_filename = _get_dataset_filename(dataset_dir,split_name,shard_id) if not tf.gfile.Exists(output_filename): return False return True #獲取圖片以及分類 def _get_filenames_and_classes(dataset_dir): #數據目錄 directories = [] #分類名稱 class_names = [] for filename in os.listdir(dataset_dir): #合并文件路徑 path = os.path.join(dataset_dir,filename) #判斷路徑是否是目錄 if os.path.isdir(path): #加入數據目錄 directories.append(path) #加入類別名稱 class_names.append(filename) photo_filenames = [] #循環分類的文件夾 for directory in directories: for filename in os.listdir(directory): path = os.path.join(directory,filename) #將圖片加入圖片列表中 photo_filenames.append(path) #返回結果 return photo_filenames ,class_names def int64_feature(values): if not isinstance(values,(tuple,list)): values = [values] return tf.train.Feature(int64_list=tf.train.Int64List(value=values)) def bytes_feature(values): return tf.train.Feature(bytes_list=tf.train.BytesList(value=[values])) #圖片轉換城tfexample函數 def image_to_tfexample(image_data,image_format,class_id): return tf.train.Example(features=tf.train.Features(feature={ 'image/encoded': bytes_feature(image_data), 'image/format': bytes_feature(image_format), 'image/class/label': int64_feature(class_id) })) def write_label_file(labels_to_class_names,dataset_dir,filename=LABELS_FILENAME): label_filename = os.path.join(dataset_dir,filename) with tf.gfile.Open(label_filename,'w') as f: for label in labels_to_class_names: class_name = labels_to_class_names[label] f.write('%d:%s\n' % (label, class_name)) #數據轉換城tfrecorad格式 def _convert_dataset(split_name,filenames,class_names_to_ids,dataset_dir): assert split_name in ['train','test'] #計算每個數據塊的大小 num_per_shard = int(len(filenames) / _NUM_SHARDS) with tf.Graph().as_default(): with tf.Session() as sess: for shard_id in range(_NUM_SHARDS): #定義tfrecord的路徑名字 output_filename = _get_dataset_filename(dataset_dir,split_name,shard_id) with tf.python_io.TFRecordWriter(output_filename) as tfrecord_writer: #每個數據塊開始的位置 start_ndx = shard_id * num_per_shard #每個數據塊結束的位置 end_ndx = min((shard_id+1) * num_per_shard,len(filenames)) for i in range(start_ndx,end_ndx): try: sys.stdout.write('\r>> Converting image %d/%d shard %d '% (i+1,len(filenames),shard_id)) sys.stdout.flush() #讀取圖片 image_data = tf.gfile.FastGFile(filenames[i],'rb').read() #獲取圖片的類別名稱 #basename獲取圖片路徑最后一個字符串 #dirname是除了basename之外的前面的字符串路徑r class_name = os.path.basename(os.path.dirname(filenames[i])) #獲取圖片的id class_id = class_names_to_ids[class_name] #生成tfrecord文件 example = image_to_tfexample(image_data,b'jpg',class_id) #寫入數據 tfrecord_writer.write(example.SerializeToString()) except IOError as e: print ('could not read:',filenames[1]) print ('error:' , e) print ('skip it \n') sys.stdout.write('\n') sys.stdout.flush() if __name__ == '__main__': #判斷tfrecord文件是否存在 if _dataset_exists(DATASET_DIR): print ('tfrecord exists') else: #獲取圖片以及分類 photo_filenames,class_names = _get_filenames_and_classes(DATASET_DIR) #將分類的list轉換成dictionary{‘house':3,'flowers:2'} class_names_to_ids = dict(zip(class_names,range(len(class_names)))) #切分數據為測試訓練集 random.seed(_RANDOM_SEED) random.shuffle(photo_filenames) training_filenames = photo_filenames[_NUM_TEST:] testing_filenames = photo_filenames[:_NUM_TEST] #數據轉換 _convert_dataset('train',training_filenames,class_names_to_ids,DATASET_DIR) _convert_dataset('test',testing_filenames,class_names_to_ids,DATASET_DIR) #輸出lables文件 #與前面的 class_names_to_ids中的元素位置相反{1:'people,2:'flowers'} labels_to_class_names = dict(zip(range(len(class_names)),class_names)) write_label_file(labels_to_class_names,DATASET_DIR)

總結

以上是生活随笔為你收集整理的TFRecord数据集制作的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。