【Python 代码】生成hdf5文件

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!'