data_utils.py

import sys
import cPickle as pickle
import numpy as np
import os
from scipy.misc import imread

import numpy as np
import lmdb
import caffe

def load_CIFAR_batch(filename, pad=True):
  """ load single batch of cifar """
  with open(filename, 'rb') as f:
    datadict = pickle.load(f)
    X = datadict['data']
    Y = datadict['labels']
    X = X.reshape(10000, 3, 32, 32).astype(np.uint8)
    padded = np.zeros((10000, 3, 40, 40), dtype=np.uint8)
    padded[:,:,:,:] = 128
    padded[:,:,4:-4, 4:-4] = X
    Y = np.array(Y, dtype=np.int64) 
    if not pad:
      return X, Y
    return padded, Y

def load_CIFAR10(ROOT):
  """ load all of cifar """
  xs = []
  ys = []
  for b in range(1,6):
    f = os.path.join(ROOT, 'data_batch_%d' % (b, ))
    X, Y = load_CIFAR_batch(f)
    xs.append(X)
    ys.append(Y)    
  Xtr = np.concatenate(xs)
  Ytr = np.concatenate(ys)
  idx = np.arange(len(Ytr))
  np.random.shuffle(idx)
  print 'shuffle training data', len(idx)
  Xtr = Xtr[idx]
  Ytr = Ytr[idx]
  print idx
  print 'tr label',Ytr.min(), Ytr.max()
  del X, Y
  Xte, Yte = load_CIFAR_batch(os.path.join(ROOT, 'test_batch'), pad=False)
  print 'te label',Yte.min(), Yte.max()
  print Xtr.shape
  print Ytr.shape
  print Xte.shape
  print Yte.shape
  return Xtr, Ytr, Xte, Yte

def py2lmdb(X, y, save_path):
  # Let's pretend this is interesting data

  assert X.dtype == np.uint8
  N = X.shape[0]
  assert N == y.shape[0], str(N) + ' ' + str(y.shape)
  
  
  # We need to prepare the database for the size. We'll set it 10 times
  # greater than what we theoretically need. There is little drawback to
  # setting this too big. If you still run into problem after raising
  # this, you might want to try saving fewer entries in a single
  # transaction.
  map_size = X.nbytes * 10
  
  env = lmdb.open(save_path, map_size=map_size)
  
  with env.begin(write=True) as txn:
      # txn is a Transaction object
      for i in range(N):
          datum = caffe.proto.caffe_pb2.Datum()
          datum.channels = X.shape[1]
          datum.height = X.shape[2]
          datum.width = X.shape[3]
          datum.data = X[i].tobytes()  # or .tostring() if numpy < 1.9
          datum.label = int(y[i])
          str_id = '{:08}'.format(i)
  
          # The encode is only essential in Python 3
          txn.put(str_id.encode('ascii'), datum.SerializeToString())


if __name__ == '__main__':
  root = sys.argv[1]
  Xtr, Ytr, Xte, Yte = load_CIFAR10(root)
  paths = [ os.path.join(root, i) for i in ['train', 'test']]
  py2lmdb(Xtr, Ytr, paths[0])
  py2lmdb(Xte, Yte, paths[1])
  for i in paths:
    print 'saved to', i