Merge pull request #1 from rai-project/feature/tfrecord

implements tfrecord reader

Author: abdul dakkak

reader/.gitignore

@@ -1,2 +1,5 @@
 
 image.png
+validation.tfrecord
+*tar.gz
+cifar-10-batches-py

reader/_fixtures/.gitignore

@@ -0,0 +1 @@
+test.png

reader/_fixtures/cifar10_validation.tfrecord

@@ -0,0 +1,123 @@
+# Copyright 2017 The TensorFlow Authors. 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.
+# ==============================================================================
+"""Read CIFAR-10 data from pickled numpy arrays and writes TFRecords.
+
+Generates tf.train.Example protos and writes them to TFRecord files from the
+python version of the CIFAR-10 dataset downloaded from
+https://www.cs.toronto.edu/~kriz/cifar.html.
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import argparse
+import os
+import sys
+
+import tarfile
+from six.moves import cPickle as pickle
+from six.moves import xrange  # pylint: disable=redefined-builtin
+import tensorflow as tf
+
+CIFAR_FILENAME = 'cifar-10-python.tar.gz'
+CIFAR_DOWNLOAD_URL = 'https://www.cs.toronto.edu/~kriz/' + CIFAR_FILENAME
+CIFAR_LOCAL_FOLDER = 'cifar-10-batches-py'
+
+
+def download_and_extract(data_dir):
+  # download CIFAR-10 if not already downloaded.
+  tf.contrib.learn.datasets.base.maybe_download(CIFAR_FILENAME, data_dir,
+                                                CIFAR_DOWNLOAD_URL)
+  tarfile.open(os.path.join(data_dir, CIFAR_FILENAME),
+               'r:gz').extractall(data_dir)
+
+
+def _int64_feature(value):
+  if not isinstance(value, list):
+    value = [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 _get_file_names():
+  """Returns the file names expected to exist in the input_dir."""
+  file_names = {}
+  # file_names['train'] = ['data_batch_%d' % i for i in xrange(1, 5)]
+  file_names['validation'] = ['data_batch_5']
+  # file_names['eval'] = ['test_batch']
+  return file_names
+
+
+def read_pickle_from_file(filename):
+  with tf.gfile.Open(filename, 'rb') as f:
+    if sys.version_info >= (3, 0):
+      data_dict = pickle.load(f, encoding='bytes')
+    else:
+      data_dict = pickle.load(f)
+  return data_dict
+
+
+def convert_to_tfrecord(input_files, output_file):
+  """Converts a file to TFRecords."""
+  print('Generating %s' % output_file)
+  with tf.python_io.TFRecordWriter(output_file) as record_writer:
+    for input_file in input_files:
+      data_dict = read_pickle_from_file(input_file)
+      data = data_dict[b'data']
+      labels = data_dict[b'labels']
+      num_entries_in_batch = len(labels)
+      for i in range(num_entries_in_batch):
+        example = tf.train.Example(features=tf.train.Features(
+            feature={
+                'image/id': _int64_feature(i),
+                'image/encoded': _bytes_feature(data[i].tobytes()),
+                'image/format': _bytes_feature(b'cifar'),
+                'image/label': _int64_feature(labels[i]),
+                'image/height': _int64_feature(32),
+                'image/width': _int64_feature(32),
+            }))
+        record_writer.write(example.SerializeToString())
+
+
+def main(data_dir):
+  print('Download from {} and extract.'.format(CIFAR_DOWNLOAD_URL))
+  download_and_extract(data_dir)
+  file_names = _get_file_names()
+  input_dir = os.path.join(data_dir, CIFAR_LOCAL_FOLDER)
+  for mode, files in file_names.items():
+    input_files = [os.path.join(input_dir, f) for f in files]
+    output_file = os.path.join(data_dir, mode + '.tfrecords')
+    try:
+      os.remove(output_file)
+    except OSError:
+      pass
+    # Convert to tf.train.Example and write the to TFRecords.
+    convert_to_tfrecord(input_files, output_file)
+  print('Done!')
+
+
+if __name__ == '__main__':
+  parser = argparse.ArgumentParser()
+  parser.add_argument(
+      '--data-dir',
+      type=str,
+      default='',
+      help='Directory to download and extract CIFAR-10 to.')
+
+  args = parser.parse_args()
+  main(args.data_dir)

reader/_fixtures/generate_cifar10_tfrecords.py

@@ -22,12 +22,6 @@ type RecordIOReader struct {
 	r io.ReadCloser
 }
 
-type Record struct {
-	ID         uint64
-	LabelIndex float32
-	Image      *types.RGBImage
-}
-
 func NewRecordIOReader(path string) (*RecordIOReader, error) {
 	r, err := os.Open(path)
 	if err != nil {
@@ -38,7 +32,7 @@ func NewRecordIOReader(path string) (*RecordIOReader, error) {
 	}, nil
 }
 
-func (r *RecordIOReader) Next(ctx context.Context) (*Record, error) {
+func (r *RecordIOReader) Next(ctx context.Context) (*ImageRecord, error) {
 	f := r.r
 
 	var magic uint32
@@ -113,7 +107,7 @@ func (r *RecordIOReader) Next(ctx context.Context) (*Record, error) {
 		return nil, errors.Errorf("expecting an rgb image")
 	}
 
-	return &Record{
+	return &ImageRecord{
 		ID:         imageId1,
 		LabelIndex: label,
 		Image:      rgbImage,

reader/mxnet_recordio.go

@@ -0,0 +1,9 @@
+package reader
+
+import "github.com/rai-project/image/types"
+
+type ImageRecord struct {
+	ID         uint64
+	LabelIndex float32
+	Image      *types.RGBImage
+}

reader/record.go

@@ -0,0 +1,87 @@
+package reader
+
+import (
+	"bytes"
+	context "context"
+	goimage "image"
+	"io"
+	"os"
+	"strings"
+
+	"github.com/pkg/errors"
+	"github.com/rai-project/image"
+	"github.com/rai-project/image/types"
+	"github.com/ubccr/terf"
+)
+
+type TFRecordReader struct {
+	r io.ReadCloser
+	*terf.Reader
+}
+
+func NewTFRecordReader(path string) (*TFRecordReader, error) {
+	r, err := os.Open(path)
+	if err != nil {
+		return nil, errors.Wrapf(err, "cannot open %v", path)
+	}
+	return &TFRecordReader{
+		r:      r,
+		Reader: terf.NewReader(r),
+	}, nil
+}
+
+func (r *TFRecordReader) Next(ctx context.Context) (*ImageRecord, error) {
+	nxt, err := r.Reader.Next()
+	if err != nil {
+		return nil, err
+	}
+
+	imgRecord := new(terf.Image)
+	err = imgRecord.UnmarshalExample(nxt)
+	if err != nil {
+		return nil, errors.Wrap(err, "unable to unmarshal image")
+	}
+
+	if strings.ToLower(imgRecord.Format) == "cifar" {
+		img := types.NewRGBImage(goimage.Rect(0, 0, imgRecord.Width, imgRecord.Height))
+		imgPix := img.Pix
+		inputPix := imgRecord.Raw
+		channels := 3
+
+		for h := 0; h < imgRecord.Height; h++ {
+			for w := 0; w < imgRecord.Width; w++ {
+				for c := 0; c < channels; c++ {
+					imgPix[channels*(h*imgRecord.Width+w)+c] =
+						// inputPix format = The first 1024 entries contain the red channel values,
+						// the next 1024 the green, and the final 1024 the blue.
+						inputPix[c*imgRecord.Height*imgRecord.Width+h*imgRecord.Width+w]
+				}
+			}
+		}
+		return &ImageRecord{
+			ID:         uint64(imgRecord.ID),
+			LabelIndex: float32(imgRecord.LabelID),
+			Image:      img,
+		}, nil
+	}
+
+	img, err := image.Read(bytes.NewBuffer(imgRecord.Raw), image.Context(nil))
+	if err != nil {
+		return nil, err
+	}
+
+	rgbImage, ok := img.(*types.RGBImage)
+	if !ok {
+		return nil, errors.Errorf("expecting an rgb image")
+	}
+
+	return &ImageRecord{
+		ID:         uint64(imgRecord.ID),
+		LabelIndex: float32(imgRecord.LabelID),
+		Image:      rgbImage,
+	}, nil
+}
+
+func (r *TFRecordReader) Close() error {
+	return r.r.Close()
+}

reader/tfrecord.go

@@ -0,0 +1,43 @@
+package reader
+
+import (
+	context "context"
+	"image/png"
+	"os"
+	"path/filepath"
+	"testing"
+
+	"github.com/GeertJohan/go-sourcepath"
+	"github.com/stretchr/testify/assert"
+)
+
+var (
+	fixturesPath = filepath.Join(sourcepath.MustAbsoluteDir(), "_fixtures")
+)
+
+func TestTFRecord(t *testing.T) {
+	reader, err := NewTFRecordReader(filepath.Join(fixturesPath, "cifar10_validation.tfrecord"))
+	assert.NoError(t, err)
+	assert.NotEmpty(t, reader)
+
+	defer reader.Close()
+
+	rec, err := reader.Next(context.Background())
+	assert.NoError(t, err)
+	assert.NotEmpty(t, rec)
+
+	assert.Equal(t, rec.ID, uint64(0))
+
+	rec, err = reader.Next(context.Background())
+	assert.NoError(t, err)
+	assert.NotEmpty(t, rec)
+
+	assert.Equal(t, rec.ID, uint64(1))
+	assert.NotEmpty(t, rec.Image.Pix)
+
+	out, _ := os.Create("_fixtures/test.png")
+	defer out.Close()
+
+	err = png.Encode(out, rec.Image.ToRGBAImage())
+	assert.NoError(t, err)
+}

reader/tfrecord_test.go

@@ -7,6 +7,7 @@ import (
 	"strings"
 
 	context "context"
+
 	"github.com/Unknwon/com"
 	"github.com/pkg/errors"
 	"github.com/rai-project/config"
@@ -37,7 +38,7 @@ type ILSVRC2012ValidationRecordIO struct {
 }
 
 type iLSVRC2012ValidationRecordIOLabeledData struct {
-	*reader.Record
+	*reader.ImageRecord
 }
 
 type recordIoOffset struct {
@@ -189,7 +190,7 @@ func (d *ILSVRC2012ValidationRecordIO) Next(ctx context.Context) (dldataset.Labe
 	}
 
 	return &iLSVRC2012ValidationRecordIOLabeledData{
-		Record: rec,
+		ImageRecord: rec,
 	}, nil
 }