最近在做OCR的一个项目，手上的训练图片不够用，于是想通过 CGAN 的方法生产一些训练图片，遇到的第一个问题是将已有图片保存成 tfrecord 文件。
代码如下：
# coding: utf-8 from __future__ import absolute_import from __future__ import division from __future__ import print_function import gzip import os import sys import numpy as np from six.moves import urllib import tensorflow as tf from PIL import Image # The URLs where the MNIST data can be downloaded. # _DATA_URL = 'http://yann.lecun.com/exdb/mnist/' _TRAIN_DATA_FILENAME = './data/orginal_image/train/' # _TRAIN_LABELS_FILENAME = 'train-labels-idx1-ubyte.gz' _TEST_DATA_FILENAME = './data/orginal_image/test/' # _TEST_LABELS_FILENAME = 't10k-labels-idx1-ubyte.gz' _IMAGE_HEIGHT = 56 _IMAGE_WIDTH = 256 _NUM_CHANNELS = 3 # # The names of the classes. # _CLASS_NAMES = [ # 'zero', # 'one', # 'two', # 'three', # 'four', # 'five', # 'size', # 'seven', # 'eight', # 'nine', # ] def get_filelist(filename): return [os.path.join(filename,f) for f in os.listdir(filename)] def _extract_images(filename): """Extract the images into a numpy array. Args: filename: The path to an MNIST images file. num_images: The number of images in the file. Returns: A numpy array of shape [number_of_images, height, width, channels]. """ print('Extracting images from: ', filename) image_list = get_filelist(filename) data = [] num_images = len(image_list) for image in image_list: if image.endswith('DS_Store') or image.endswith('txt'): num_images = num_images-1 continue img = Image.open(image) img_resize = np.array(img.resize((_IMAGE_WIDTH,_IMAGE_HEIGHT))) # 注意顺序，一开始就把宽高顺序写反了，导致图片一致显示有误。 img_resize = img_resize.astype(np.uint8) data.append(img_resize) # data = np.array(data).astype(np.uint8) # data = data.reshape(num_images, _IMAGE_HEIGHT,_IMAGE_WIDTH,3) return data,num_images def _extract_labels(filename): """Extract the labels into a vector of int64 label IDs. Args: filename: The path to an MNIST labels file. num_labels: The number of labels in the file. Returns: A numpy array of shape [number_of_labels] """ print('Extracting labels from: ', filename) image_list = get_filelist(filename) labels = [] num_labels = len(image_list) for image in image_list: if image.endswith('DS_Store') or image.endswith('txt'): num_labels = num_labels-1 continue label = image.split('/')[-1].split('_')[0] labels.append(label) # print(label) return labels,num_labels def int64_feature(values): """Returns a TF-Feature of int64s. Args: values: A scalar or list of values. Returns: A TF-Feature. """ if not isinstance(values, (tuple, list)): values = [values] return tf.train.Feature(int64_list=tf.train.Int64List(value=values)) def bytes_feature(values): """Returns a TF-Feature of bytes. Args: values: A string. Returns: A TF-Feature. """ return tf.train.Feature(bytes_list=tf.train.BytesList(value=[values])) def image_to_tfexample(image_data, image_format, height, width, 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': bytes_feature(class_id), # 有更改 'image/height': int64_feature(height), 'image/width': int64_feature(width), })) def _add_to_tfrecord(data_filename, tfrecord_writer): """Loads data from the binary MNIST files and writes files to a TFRecord. Args: data_filename: The filename of the MNIST images. labels_filename: The filename of the MNIST labels. num_images: The number of images in the dataset. tfrecord_writer: The TFRecord writer to use for writing. """ images, num_images= _extract_images(data_filename) labels, num_labels = _extract_labels(data_filename) shape = (_IMAGE_HEIGHT,_IMAGE_WIDTH, _NUM_CHANNELS) with tf.Graph().as_default(): image = tf.placeholder(dtype=tf.uint8, shape=shape) encoded_jpeg = tf.image.encode_jpeg(image) # 这里要注意转码格式，因为我的代码是基于tf官方教程改的，是针对mnist图片，png格式的，我自己的训练图片是jpg的，转码不对也会出现问题。 with tf.Session('') as sess: for j in range(num_images): sys.stdout.write('\r>> Converting image %d/%d' % (j + 1, num_images)) sys.stdout.flush() png_string = sess.run(encoded_jpeg, feed_dict={image: images[j]}) example = image_to_tfexample( png_string, 'jpg'.encode(), _IMAGE_HEIGHT,_IMAGE_WIDTH, labels[j]) #这里同上，也要注意转码格式 tfrecord_writer.write(example.SerializeToString()) def _get_output_filename(dataset_dir, split_name): """Creates the output filename. Args: dataset_dir: The directory where the temporary files are stored. split_name: The name of the train/test split. Returns: An absolute file path. """ return '%s/invoice_%s.tfrecord' % (dataset_dir, split_name) def run(dataset_dir): """Runs the download and conversion operation. Args: dataset_dir: The dataset directory where the dataset is stored. """ if not tf.gfile.Exists(dataset_dir): tf.gfile.MakeDirs(dataset_dir) training_filename = _get_output_filename(dataset_dir, 'train') testing_filename = _get_output_filename(dataset_dir, 'test') print(training_filename) print(testing_filename) if tf.gfile.Exists(training_filename) and tf.gfile.Exists(testing_filename): print('Dataset files already exist. Exiting without re-creating them.') return # _download_dataset(dataset_dir) # First, process the training data: with tf.python_io.TFRecordWriter(training_filename) as tfrecord_writer: data_filename = _TRAIN_DATA_FILENAME _add_to_tfrecord(data_filename, tfrecord_writer) # Next, process the testing data: with tf.python_io.TFRecordWriter(testing_filename) as tfrecord_writer: data_filename = _TEST_DATA_FILENAME _add_to_tfrecord(data_filename, tfrecord_writer) # Finally, write the labels file: # labels_to_class_names = dict(zip(range(len(_CLASS_NAMES)), _CLASS_NAMES)) # dataset_utils.write_label_file(labels_to_class_names, dataset_dir) print('\nFinished converting the invoice dataset!') if __name__ == '__main__': dataset_dir = './data/tfrecord' run(dataset_dir) pass