import random

from PIL import Image

import numpy as np

import os

import h5py

from PIL import Image

LIST_FILE = ['list_train.txt', 'list_test.txt']######################

HDF5_LIST = 'HDF5/list_hdf5.txt'##############

print '\nplease wait...'

#write

Phase='TRAIN'

slice_num=100 #every 100 img form a hdf5 file

if Phase=='TRAIN':

    image_dir=LIST_FILE[0]

elif Phase=='TEST':

    image_dir=LIST_FILE[1]

else:

    print 'error!'

files=[]

with open(image_dir) as f0:

    for line in f0.readlines():

        files.append(line[:-1])  

random.shuffle(files)#随机打乱文件顺序#############

sum_file=len(files) # sum of img

count_end=int(sum_file/slice_num) #num of hdf5 files:count_end+1

for count in range (count_end+1):

    files_part=[]

    if count==count_end:

        files_part=files[count_end*slice_num:]

    else:

        files_part=files[count*slice_num:(count+1)*slice_num]

    # data :eg 1channel 96*32

    datas = np.zeros((len(files_part),1, 512, 512)) ###输入tif尺寸

    # label eg 1*2

    labels = np.zeros((len(files_part),1,340, 340)) ####label尺寸

    for ii, _file in enumerate(files_part):

        train_img='./train/train_img_512/'+_file #######################

        train_label='./train/train_label_crop340/'+_file ################

        #print _file

        datas[ii, :, :] = np.array(Image.open(train_img)).astype(np.float32) / 256

        labels[ii, :, :] = \

            ((np.array(Image.open(train_label)).astype(np.float32) / 256)>0.5)\

                .astype(np.float32)####大于0.5输出1，否则输出0

    #New=Image.fromarray(labels[0][0])

    #New.show()

    hdf5filename=Phase+'_hdf5_'+str(count)+'.h5'

    with h5py.File('HDF5/'+hdf5filename, 'w') as f:

        f['data'] = datas

        f['label'] = labels

        f.close()

    #生成hdf5文件列表用于prototxt中

    with open(HDF5_LIST, 'a') as f:

        f.write('caffe-unet-src/cell_data/HDF5/'+hdf5filename + '\n')

        f.close()

    print 'hdf5 file : %d'%(count+1)

print '\ndone!'