Last active
July 15, 2019 19:42
-
-
Save MathiasGruber/8debc802464a48efc8c22a2064e0bf78 to your computer and use it in GitHub Desktop.
Convert image data to tf.records
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
# Copyright 2016 Google Inc. All Rights Reserved. | |
# | |
# Licensed under the Apache License, Version 2.0 (the "License"); | |
# you may not use this file except in compliance with the License. | |
# You may obtain a copy of the License at | |
# | |
# http://www.apache.org/licenses/LICENSE-2.0 | |
# | |
# Unless required by applicable law or agreed to in writing, software | |
# distributed under the License is distributed on an "AS IS" BASIS, | |
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
# See the License for the specific language governing permissions and | |
# limitations under the License. | |
# ============================================================================== | |
""" | |
Adapted from: https://github.com/kmonachopoulos/ImageNet-to-TFrecord | |
Converts image data to TFRecords file format with Example protos. | |
The raw data set is expected to reside in JPEG files located in the | |
following directory structure. | |
data_dir/classname1/filename1.JPEG | |
data_dir/classname2/filename2.JPEG | |
... | |
Example commands used to run this script: | |
============================= | |
python build_tfrecords.py -data_directory ./COCO/train/ -output_directory ./COCO/tf_records_train/ -data_shards 80 -num_threads 8 | |
python build_tfrecords.py -data_directory ./COCO/val/ -output_directory ./COCO/tf_records_val/ -data_shards 80 -num_threads 8 | |
python build_tfrecords.py -data_directory ./wikiart/wikiart/ -output_directory ./wikiart/tfrecords/ -data_shards 80 -num_threads 8 | |
python build_tfrecords.py -data_directory ./ILSVRC/Data/CLS-LOC/train -output_directory ./ILSVRC/Data/CLS-LOC/tfrecords_train -data_shards 1024 -num_threads 8 | |
python build_tfrecords.py -data_directory ./ILSVRC/Data/CLS-LOC/val -output_directory ./ILSVRC/Data/CLS-LOC/tfrecords_val -data_shards 1 -num_threads 8 | |
python build_tfrecords.py -data_directory ./ILSVRC/Data/CLS-LOC/test -output_directory ./ILSVRC/Data/CLS-LOC/tfrecords_test -data_shards 1 -num_threads 8 | |
Each record within the TFRecord file is a serialized Example proto. | |
The Example proto contains the following fields: | |
image/encoded: string containing JPEG encoded image in RGB colorspace | |
image/height: integer, image height in pixels | |
image/width: integer, image width in pixels | |
image/colorspace: string, specifying the colorspace, always 'RGB' | |
image/channels: integer, specifying the number of channels, always 3 | |
image/format: string, specifying the format, always'JPEG' | |
image/filename: string containing the basename of the image file | |
e.g. 'n01440764_10026.JPEG' or 'ILSVRC2012_val_00000293.JPEG' | |
image/class/label: integer specifying the index in a classification layer. | |
The label ranges from [1, 1000] where 0 is not used. | |
image/class/synset: string specifying the unique ID of the label, | |
e.g. 'n01440764' | |
image/class/text: string specifying the human-readable version of the label | |
e.g. 'red fox, Vulpes vulpes' | |
Note that the length of xmin is identical to the length of xmax, ymin and ymax | |
for each example. | |
""" | |
from __future__ import absolute_import | |
from __future__ import division | |
from __future__ import print_function | |
from datetime import datetime | |
import os | |
from PIL import Image | |
import gc | |
import random | |
import sys | |
import threading | |
import numpy as np | |
from six.moves import xrange # pylint: disable=redefined-builtin | |
import tensorflow as tf | |
tf.app.flags.DEFINE_string('data_directory', None, | |
'Data directory') | |
tf.app.flags.DEFINE_string('output_directory', None, | |
'Output data directory') | |
tf.app.flags.DEFINE_integer('data_shards', 1024, | |
'Number of shards in training TFRecord files.') | |
tf.app.flags.DEFINE_integer('num_threads', 1, | |
'Number of threads to preprocess the images.') | |
FLAGS = tf.app.flags.FLAGS | |
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 _float_feature(value): | |
"""Wrapper for inserting float features into Example proto.""" | |
if not isinstance(value, list): | |
value = [value] | |
return tf.train.Feature(float_list=tf.train.FloatList(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(filename, image_buffer, label, synset, human, height, width): | |
"""Build an Example proto for an example. | |
Args: | |
filename: string, path to an image file, e.g., '/path/to/example.JPG' | |
image_buffer: string, JPEG encoding of RGB image | |
label: integer, identifier for the ground truth for the network | |
synset: string, unique WordNet ID specifying the label, e.g., 'n02323233' | |
human: string, human-readable label, e.g., 'red fox, Vulpes vulpes' | |
height: integer, image height in pixels | |
width: integer, image width in pixels | |
Returns: | |
Example proto | |
""" | |
colorspace = b'RGB' | |
channels = 3 | |
image_format = b'JPEG' | |
example = tf.train.Example(features=tf.train.Features(feature={ | |
'image/height': _int64_feature(height), | |
'image/width': _int64_feature(width), | |
'image/colorspace': _bytes_feature(colorspace), | |
'image/channels': _int64_feature(channels), | |
'image/class/label': _int64_feature(label), | |
'image/class/synset': _bytes_feature(bytes(synset,'utf-8')), | |
'image/class/text': _bytes_feature(bytes(human,'utf-8')), | |
'image/format': _bytes_feature(image_format), | |
'image/filename': _bytes_feature(bytes(os.path.basename(filename),'utf-8')), | |
'image/encoded': _bytes_feature(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 converts CMYK JPEG data to RGB JPEG data. | |
self._cmyk_data = tf.placeholder(dtype=tf.string) | |
image = tf.image.decode_jpeg(self._cmyk_data, channels=0) | |
self._cmyk_to_rgb = 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 cmyk_to_rgb(self, image_data): | |
return self._sess.run(self._cmyk_to_rgb, | |
feed_dict={self._cmyk_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 _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. | |
image_data = tf.gfile.GFile(filename, 'rb').read() | |
# 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, thread_index, ranges, name, filenames, | |
synsets, labels, humans, num_shards): | |
"""Processes and saves list of images as TFRecord in 1 thread. | |
Args: | |
coder: instance of ImageCoder to provide TensorFlow image coding utils. | |
thread_index: integer, unique batch to run index is within [0, len(ranges)). | |
ranges: list of pairs of integers specifying ranges of each batches to | |
analyze in parallel. | |
name: string, unique identifier specifying the data set | |
filenames: list of strings; each string is a path to an image file | |
synsets: list of strings; each string is a unique WordNet ID | |
labels: list of integer; each integer identifies the ground truth | |
humans: list of strings; each string is a human-readable label | |
num_shards: integer number of shards for this data set. | |
""" | |
# Each thread produces N shards where N = int(num_shards / num_threads). | |
# For instance, if num_shards = 128, and the num_threads = 2, then the first | |
# thread would produce shards [0, 64). | |
num_threads = len(ranges) | |
assert not num_shards % num_threads | |
num_shards_per_batch = int(num_shards / num_threads) | |
shard_ranges = np.linspace(ranges[thread_index][0], | |
ranges[thread_index][1], | |
num_shards_per_batch + 1).astype(int) | |
num_files_in_thread = ranges[thread_index][1] - ranges[thread_index][0] | |
counter = 0 | |
for s in xrange(num_shards_per_batch): | |
# Generate a sharded version of the file name, e.g. 'train-00002-of-00010' | |
shard = thread_index * num_shards_per_batch + s | |
output_filename = '%s-%.5d-of-%.5d' % (name, shard, num_shards) | |
output_file = os.path.join(FLAGS.output_directory, output_filename) | |
writer = tf.python_io.TFRecordWriter(output_file) | |
shard_counter = 0 | |
files_in_shard = np.arange(shard_ranges[s], shard_ranges[s + 1], dtype=int) # HERE | |
for i in files_in_shard: | |
filename = filenames[i] | |
label = labels[i] | |
synset = synsets[i] | |
human = humans[i] | |
try: | |
image_buffer, height, width = _process_image(filename, coder) | |
example = _convert_to_example(filename, image_buffer, label, synset, human, height, width) | |
writer.write(example.SerializeToString()) | |
shard_counter += 1 | |
counter += 1 | |
if not counter % 1000: | |
print('%s [thread %d]: Processed %d of %d images in thread batch.' % | |
(datetime.now(), thread_index, counter, num_files_in_thread)) | |
sys.stdout.flush() | |
except Exception as e: | |
print("Error: ", filename, e) | |
writer.close() | |
print('%s [thread %d]: Wrote %d images to %s' % | |
(datetime.now(), thread_index, shard_counter, output_file)) | |
sys.stdout.flush() | |
shard_counter = 0 | |
print('%s [thread %d]: Wrote %d images to %d shards.' % | |
(datetime.now(), thread_index, counter, num_files_in_thread)) | |
sys.stdout.flush() | |
def _process_image_files(name, filenames, synsets, labels, num_shards): | |
"""Process and save list of images as TFRecord of Example protos. | |
Args: | |
name: string, unique identifier specifying the data set | |
filenames: list of strings; each string is a path to an image file | |
synsets: list of strings; each string is a unique WordNet ID | |
labels: list of integer; each integer identifies the ground truth | |
humans: list of strings; each string is a human-readable label | |
num_shards: integer number of shards for this data set. | |
""" | |
assert len(filenames) == len(synsets) | |
assert len(filenames) == len(labels) | |
humans = [str(h) for h in labels] | |
# Break all images into batches with a [ranges[i][0], ranges[i][1]]. | |
spacing = np.linspace(0, len(filenames), FLAGS.num_threads + 1).astype(np.int) | |
ranges = [] | |
threads = [] | |
for i in xrange(len(spacing) - 1): | |
ranges.append([spacing[i], spacing[i+1]]) | |
# Launch a thread for each batch. | |
print('Launching %d threads for spacings: %s' % (FLAGS.num_threads, ranges)) | |
sys.stdout.flush() | |
# Create a mechanism for monitoring when all threads are finished. | |
coord = tf.train.Coordinator() | |
# Create a generic TensorFlow-based utility for converting all image codings. | |
coder = ImageCoder() | |
threads = [] | |
for thread_index in xrange(len(ranges)): | |
args = (coder, thread_index, ranges, name, filenames, | |
synsets, labels, humans, num_shards) | |
t = threading.Thread(target=_process_image_files_batch, args=args) | |
t.start() | |
threads.append(t) | |
# Wait for all the threads to terminate. | |
coord.join(threads) | |
print('%s: Finished writing all %d images in data set.' % | |
(datetime.now(), len(filenames))) | |
sys.stdout.flush() | |
def _find_image_files(data_dir): | |
"""Build a list of all images files and labels in the data set. | |
Args: | |
data_dir: string, path to the root directory of images. | |
Assumes that the ImageNet data set resides in JPEG files located in | |
the following directory structure. | |
data_dir/n01440764/ILSVRC2012_val_00000293.JPEG | |
data_dir/n01440764/ILSVRC2012_val_00000543.JPEG | |
where 'n01440764' is the unique synset label associated with these images. | |
Returns: | |
filenames: list of strings; each string is a path to an image file. | |
synsets: list of strings; each string is a unique WordNet ID. | |
labels: list of integer; each integer identifies the ground truth. | |
""" | |
print('Determining list of input files and labels from %s.' % data_dir) | |
challenge_synsets = [ name for name in os.listdir(data_dir) if os.path.isdir(os.path.join(data_dir, name)) ] | |
print("Identified categories: ", challenge_synsets) | |
labels = [] | |
filenames = [] | |
synsets = [] | |
# Leave label index 0 empty as a background class. | |
label_index = 1 | |
# Construct the list of JPEG files and labels. | |
for synset in challenge_synsets: | |
jpeg_file_path = '%s/%s/*.*' % (data_dir, synset) | |
matching_files = tf.gfile.Glob(jpeg_file_path) | |
labels.extend([label_index] * len(matching_files)) | |
synsets.extend([synset] * len(matching_files)) | |
filenames.extend(matching_files) | |
if not label_index % 100: | |
print('Finished finding files in %d of %d classes.' % ( | |
label_index, len(challenge_synsets))) | |
label_index += 1 | |
# Shuffle the ordering of all image files in order to guarantee | |
# random ordering of the images with respect to label in the | |
# saved TFRecord files. Make the randomization repeatable. | |
shuffled_index = range(len(filenames)) | |
random.seed(12345) | |
random.shuffle(list(range(len(shuffled_index)))) | |
filenames = [filenames[i] for i in shuffled_index] | |
synsets = [synsets[i] for i in shuffled_index] | |
labels = [labels[i] for i in shuffled_index] | |
print('Found %d files across %d labels inside %s.' % | |
(len(filenames), len(challenge_synsets), data_dir)) | |
return filenames, synsets, labels | |
def _find_human_readable_labels(synsets, synset_to_human): | |
"""Build a list of human-readable labels. | |
Args: | |
synsets: list of strings; each string is a unique WordNet ID. | |
synset_to_human: dict of synset to human labels, e.g., | |
'n02119022' --> 'red fox, Vulpes vulpes' | |
Returns: | |
List of human-readable strings corresponding to each synset. | |
""" | |
humans = [] | |
for s in synsets: | |
assert s in synset_to_human, ('Failed to find: %s' % s) | |
humans.append(synset_to_human[s]) | |
return humans | |
def _process_dataset(name, directory, num_shards): | |
"""Process a complete data set and save it as a TFRecord. | |
Args: | |
name: string, unique identifier specifying the data set. | |
directory: string, root path to the data set. | |
num_shards: integer number of shards for this data set. | |
synset_to_human: dict of synset to human labels, e.g., | |
'n02119022' --> 'red fox, Vulpes vulpes' | |
""" | |
filenames, synsets, labels = _find_image_files(directory) | |
_process_image_files(name, filenames, synsets, labels, num_shards) | |
def _verify_images(datadir): | |
"""Verifies the integrity of all the images, and flags if some cannot be decoded""" | |
import os | |
import gc | |
from PIL import Image | |
import warnings | |
warnings.filterwarnings("error") | |
print(" -- Verifying file integrity --") | |
i = 0 | |
for subdir, dirs, files in os.walk(rootdir): | |
for file in files: | |
i += 1 | |
print("Total images: ", i) | |
n = 0 | |
fucked = [] | |
for subdir, dirs, files in os.walk(rootdir): | |
for file in files: | |
fn = os.path.join(subdir, file) | |
n += 1 | |
try: | |
with Image.open(fn) as img: | |
img.verify() | |
gc.collect() | |
if n % 1000 == 0: | |
print('Progress: {} out of {}'.format(n, i)) | |
except Exception as e: | |
print("Found issue with: ", fn, e) | |
fucked.append(fn) | |
for fuck in fucked: | |
if os.path.exists(fuck): | |
os.remove(fuck) | |
else: | |
print("already removed", fuck) | |
def main(unused_argv): | |
assert not FLAGS.data_shards % FLAGS.num_threads, ( | |
'Please make the FLAGS.num_threads commensurate with FLAGS.train_shards') | |
print('Saving results to %s' % FLAGS.output_directory) | |
if(FLAGS.data_directory != None): | |
_process_dataset('shard', FLAGS.data_directory, FLAGS.data_shards) | |
if __name__ == '__main__': | |
tf.app.run() |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment