tiff_to_lmdb.py

import numpy as np
import deepdish as dd
import lmdb
import caffe
import glob
import random
import tifffile

#creates a numpy array data 
fdata = glob.glob("/data/*.tif")
n_data = len(fdata)
final_data = np.zeros((n_data,256,256,5))

for k in range (len(fdata)):
    fcurrent = fdata[k]
    img = tifffile.imread(fcurrent)
    img = img.reshape(1,256,256,5)   
    final_data[k] = img
    
#input format image,channel,width,height
final_data = np.transpose(final_data,(0,3,1,2))
data = final_data
print 'data shape : ',data.shape

#creates a numpy array of label Y

flabel = glob.glob("/label/*.tif")
n_label = len(flabel)
final_label = np.zeros((n_label,256,256,1))

for k in range (len(flabel)):
    fcurrent = flabel[k]
    img = tifffile.imread(fcurrent)
    img = img.reshape(1,256,256,1)   
    final_label[k] = img

#input format image,channel,width,height
final_label = np.transpose(final_label,(0,3,1,2))    
Y =  final_label
print 'label shape : ',Y.shape


#shuffle
N = range(len(final_data))
random.shuffle(N)

map_size = Y.nbytes*2
env = lmdb.open('labels', map_size=map_size)
for z in N:
    i = N[z]
    im_dat = caffe.io.array_to_datum(np.array(Y[i]).astype(float))
    str_id = '{:0>10d}'.format(i)
    with env.begin(write=True) as txn:
        txn.put(str_id, im_dat.SerializeToString())


print 'Creating image dataset.'
X = np.array(data)
X = np.array(X,dtype=np.float32)
map_size = X.nbytes*2
env = lmdb.open('data_lmdb', map_size=map_size)
for z in N:
    i = N[z]
    im_dat = caffe.io.array_to_datum(np.array(X[i]).astype(float))
    str_id = '{:0>10d}'.format(i)
    with env.begin(write=True) as txn:
        txn.put(str_id, im_dat.SerializeToString())