Created
February 11, 2017 16:09
-
-
Save AtiqueUrRehman/a07fca7fe7e8beeeccdd3980a3b0c48b to your computer and use it in GitHub Desktop.
An example of converting images to tfrecords, in this case we have an image-to-image mapping, so we have some input images and corresponding label images. Correspoing reading code is in comments below.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
from __future__ import absolute_import | |
from __future__ import division | |
from __future__ import print_function | |
from datetime import datetime | |
import os | |
import random | |
import sys | |
import threading | |
import math | |
import numpy as np | |
import tensorflow as tf | |
def _int64_feature(value): | |
"""Wrapper for inserting int64 features into Example proto.""" | |
if not isinstance(value, list): | |
value = [value] | |
return tf.train.Feature(int64_list=tf.train.Int64List(value=value)) | |
def _bytes_feature(value): | |
"""Wrapper for inserting bytes features into Example proto.""" | |
return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value])) | |
def _convert_to_example(orig_filename,label_filename, orig_image_buffer, label_image_buffer, | |
orig_height, orig_width, label_height, label_width): | |
"""Build an Example proto for an example. | |
Args: | |
orig_filename: string, path to an image file, e.g., '/path/to/example.JPG' | |
label_filename: string, path to an image file, e.g., '/path/to/example.JPG' | |
orig_image_buffer: string, JPEG encoding of RGB image | |
label_image_buffer: string, JPEG encoding of RGB image | |
orig_height: integer, image height in pixels | |
orig_width: integer, image width in pixels | |
label_height: integer, image height in pixels | |
label_width: integer, image width in pixels | |
Returns: | |
Example proto | |
""" | |
colorspace = 'RGB' | |
channels = 3 | |
image_format = 'JPEG' | |
example = tf.train.Example(features=tf.train.Features(feature={ | |
'orig/image/height': _int64_feature(orig_height), | |
'orig/image/width': _int64_feature(orig_width), | |
'orig/image/colorspace': _bytes_feature(tf.compat.as_bytes(colorspace)), | |
'orig/image/channels': _int64_feature(channels), | |
'orig/image/format': _bytes_feature(tf.compat.as_bytes(image_format)), | |
'orig/image/filename': _bytes_feature(tf.compat.as_bytes(os.path.basename(orig_filename))), | |
'orig/image/encoded': _bytes_feature(tf.compat.as_bytes(orig_image_buffer)), | |
'label/image/height': _int64_feature(label_height), | |
'label/image/width': _int64_feature(label_width), | |
'label/image/colorspace': _bytes_feature(tf.compat.as_bytes(colorspace)), | |
'label/image/channels': _int64_feature(channels), | |
'label/image/format': _bytes_feature(tf.compat.as_bytes(image_format)), | |
'label/image/filename': _bytes_feature(tf.compat.as_bytes(os.path.basename(label_filename))), | |
'label/image/encoded': _bytes_feature(tf.compat.as_bytes(label_image_buffer)) })) | |
return example | |
class ImageCoder(object): | |
"""Helper class that provides TensorFlow image coding utilities.""" | |
def __init__(self): | |
# Create a single Session to run all image coding calls. | |
self._sess = tf.Session() | |
# Initializes function that converts PNG to JPEG data. | |
self._png_data = tf.placeholder(dtype=tf.string) | |
image = tf.image.decode_png(self._png_data, channels=3) | |
self._png_to_jpeg = tf.image.encode_jpeg(image, format='rgb', quality=100) | |
# Initializes function that decodes RGB JPEG data. | |
self._decode_jpeg_data = tf.placeholder(dtype=tf.string) | |
self._decode_jpeg = tf.image.decode_jpeg(self._decode_jpeg_data, channels=3) | |
def png_to_jpeg(self, image_data): | |
return self._sess.run(self._png_to_jpeg, | |
feed_dict={self._png_data: image_data}) | |
def decode_jpeg(self, image_data): | |
image = self._sess.run(self._decode_jpeg, | |
feed_dict={self._decode_jpeg_data: image_data}) | |
assert len(image.shape) == 3 | |
assert image.shape[2] == 3 | |
return image | |
def _is_png(filename): | |
"""Determine if a file contains a PNG format image. | |
Args: | |
filename: string, path of the image file. | |
Returns: | |
boolean indicating if the image is a PNG. | |
""" | |
return '.png' in filename | |
def _process_image(filename, coder): | |
"""Process a single image file. | |
Args: | |
filename: string, path to an image file e.g., '/path/to/example.JPG'. | |
coder: instance of ImageCoder to provide TensorFlow image coding utils. | |
Returns: | |
image_buffer: string, JPEG encoding of RGB image. | |
height: integer, image height in pixels. | |
width: integer, image width in pixels. | |
""" | |
# Read the image file. | |
with tf.gfile.FastGFile(filename, 'r') as f: | |
image_data = f.read() | |
# Convert any PNG to JPEG's for consistency. | |
if _is_png(filename): | |
print('Converting PNG to JPEG for %s' % filename) | |
image_data = coder.png_to_jpeg(image_data) | |
# Decode the RGB JPEG. | |
image = coder.decode_jpeg(image_data) | |
# Check that image converted to RGB | |
assert len(image.shape) == 3 | |
height = image.shape[0] | |
width = image.shape[1] | |
assert image.shape[2] == 3 | |
return image_data, height, width | |
def _process_image_files_batch(coder, name, orig_filenames, label_filenames, output_directory, shards_size): | |
"""Processes and saves list of images as TFRecord in 1 thread. | |
Args: | |
coder: instance of ImageCoder to provide TensorFlow image coding utils. | |
name: string, unique identifier specifying the data set | |
orig_filenames: list of strings; each string is a path to an image file | |
label_filenames: list of strings; each string is a path to an image file | |
outpul_directory : Directory for output files | |
shards_size: integer size of shards for this data set. | |
""" | |
if (shards_size != -1): | |
total_shards = int(math.ceil(len(orig_filenames) / shards_size)) | |
else: | |
total_shards = 1 | |
shards_size = len(orig_filenames) | |
print ("Total Partitions %d, partition size %d "% (total_shards, shards_size)) | |
for shard in range(total_shards): | |
# Generate a sharded version of the file name, e.g. 'train-00002-of-00010' | |
output_filename = '%s-%.5d-of-%.5d' % (name, shard, total_shards) | |
output_file = os.path.join(output_directory, output_filename) | |
writer = tf.python_io.TFRecordWriter(output_file) | |
shard_counter = 0 | |
files_in_shard = np.arange(shard*shards_size, min((shard+1)*shards_size, len(orig_filenames) ), dtype=int) | |
for i in files_in_shard: | |
orig = orig_filenames[i] | |
label = label_filenames[i] | |
orig_image_buffer, orig_height, orig_width = _process_image(orig, coder) | |
label_image_buffer, label_height, label_width = _process_image(label, coder) | |
example = _convert_to_example(orig,label, orig_image_buffer, label_image_buffer, | |
orig_height, orig_width, label_height, label_width) | |
writer.write(example.SerializeToString()) | |
print ("Processed files %d of %d" % (shard*shards_size, len(orig_filenames))) | |
writer.close() | |
sys.stdout.flush() | |
shard_counter = 0 | |
sys.stdout.flush() | |
def main(orignal_image_folder, label_image_folder, output_directory, shards_size = -1): | |
orig_img_paths = [os.path.join(orignal_image_folder,im) for im in os.listdir(orignal_image_folder) if os.path.isfile (os.path.join(orignal_image_folder,im))] | |
label_img_paths = [os.path.join(label_image_folder,im) for im in os.listdir(label_image_folder) if os.path.isfile(os.path.join(label_image_folder,im))] | |
coder = ImageCoder() | |
_process_image_files_batch(coder, "data", orig_img_paths, label_img_paths, output_directory, shards_size) | |
if __name__ == '__main__': | |
if len(sys.argv) < 4 or len(sys.argv) > 5: | |
print ("Usage imagesToTfrecords <input_images_folder> <label_images_folder> <output_folder> (optional) <output partition size>") | |
elif len(sys.argv) == 4: | |
main(sys.argv[1], sys.argv[2], sys.argv[3]) | |
else: | |
main(sys.argv[1], sys.argv[2], sys.argv[3], int(sys.argv[4])) | |
# For reading files | |
# import tensorflow as tf | |
# import matplotlib.pyplot as plt | |
# filename = "../Data/tfrecords/cool-00000-of-00004" | |
# sess = tf.Session() | |
# for serialized_example in tf.python_io.tf_record_iterator(filename): | |
# example = tf.train.Example() | |
# example.ParseFromString(serialized_example) | |
# # traverse the Example format to get data | |
# img = example.features.feature['origimage/encoded'] | |
# # get the data out of tf record | |
# orignal_image_height = example.features.feature['orig/image/height'] | |
# orignal_image_width = example.features.feature['orig/image/width'] | |
# orignal_image_colors = example.features.feature['orig/image/colorspace'] | |
# orignal_image_channels = example.features.feature['orig/image/channels'] | |
# orignal_image_format = example.features.feature['orig/image/format'] | |
# orignal_image_filename = example.features.feature['orig/image/filename'] | |
# orignal_image_data = example.features.feature['orig/image/encoded'] | |
# noisy_image_height = example.features.feature['label/image/height'] | |
# noisy_image_width = example.features.feature['label/image/width'] | |
# noisy_image_colors = example.features.feature['label/image/colorspace'] | |
# noisy_image_channels = example.features.feature['label/image/channels'] | |
# noisy_image_format = example.features.feature['label/image/format'] | |
# noisy_image_filename = example.features.feature['label/image/filename'] | |
# noisy_image_data = example.features.feature['label/image/encoded'] | |
# orignal_image = sess.run(tf.image.decode_jpeg(orignal_image_data.bytes_list.value[0], channels=3)) | |
# noisy_image = sess.run(tf.image.decode_jpeg(noisy_image_data.bytes_list.value[0], channels=3)) | |
# plt.subplot(121) | |
# plt.title("Image Name : " + str(orignal_image_filename.bytes_list.value[0]) + "\n" + | |
# "Image Height : " + str(orignal_image_height.int64_list.value[0]) + "\n" + | |
# "Image Weight : " + str(orignal_image_width.int64_list.value[0]) + "\n" + | |
# "Image ColourSpace : " + str(orignal_image_colors.bytes_list.value[0]) + "\n" + | |
# "Image Channels : " + str(orignal_image_channels.int64_list.value[0]) + "\n" + | |
# "Image format : " + str(orignal_image_format.bytes_list.value[0]) + "\n") | |
# plt.imshow(orignal_image) | |
# plt.subplot(122) | |
# plt.title("Image Name : " + str(noisy_image_filename.bytes_list.value[0]) + "\n" + | |
# "Image Height : " + str(noisy_image_height.int64_list.value[0]) + "\n" + | |
# "Image Weight : " + str(noisy_image_width.int64_list.value[0]) + "\n" + | |
# "Image ColourSpace : " + str(noisy_image_colors.bytes_list.value[0]) + "\n" + | |
# "Image Channels : " + str(noisy_image_channels.int64_list.value[0]) + "\n" + | |
# "Image format : " + str(noisy_image_format.bytes_list.value[0]) + "\n") | |
# plt.imshow(noisy_image) | |
# plt.show() | |
# break |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment