为图像分割任务创建HDF5格式

时间:2018-04-29 16:20:22

标签: python-2.7 caffe hdf5 pycaffe

我开始编写用于为图像分割任务创建HDF5的python代码。我使用了this link中的代码和Shai提供的链接。我的图片是一个频道,格式为.mat。我写了以下代码,我只想与专家核实这段代码是否正确。专家可以看看吗?感谢

import os, h5py
import caffe
import numpy as np
import scipy
import scipy.io as sio
from array import array
import cv2
import matplotlib.pyplot as plt

caffe_root='/home/ss/caffe/'
import sys
sys.path.insert(0,caffe_root+'python')

def img_to_hdf5(paths_src_file,paths_lbl_file,path_dst,msg):
    """
    paths_src_file : path to the image paths in a txt file
    paths_lbl_file : path to the image paths in a txt file
    path_dst = path to the hdf5 file  
    """

    print(msg)
    arrays = {}
    SIZE=256   #fixed size of all images

    #read the lines of img and lbl path from text file and save into paths_src and paths_lbl
    paths_src = []
    with open(paths_src_file) as f:
        for line in f.readlines():
            line = line.strip('\n')
            paths_src.append(line)

    paths_lbl=[]
    with open(paths_lbl_file) as f:
        for line in f.readlines():
            line=line.strip('\n')
            paths_lbl.append(line)

    data = np.zeros( (len(paths_src), 1, SIZE, SIZE), dtype='f4' )  #1 channel grayscale image
    label = np.zeros( (len(paths_lbl), 1, SIZE, SIZE), dtype='f4' ) #1 channel label image
    for in_idx, in_ in enumerate(paths_src):
        print in_idx,in_
        f=h5py.File(in_,'r')
        mat=f['image'].value
        im=np.array(mat,dtype=np.float32)
        #im = cv2.cvtColor(im,cv2.COLOR_GRAY2RGB)
        #im = im[:,:,::-1]  #switch from RGB to BGR
        im = im.reshape(im.shape[0], im.shape[1], 1)
        im = im.transpose((2,0,1)) # convert to CxHxW
        data[in_idx]=im

    for in_idx, in_ in enumerate(paths_lbl):
        print in_idx,in_
        f=h5py.File(in_,'r')
        mat=f['image'].value
        im=np.array(mat,dtype=np.float32)
        #im = cv2.cvtColor(im,cv2.COLOR_GRAY2RGB)
        #im = im[:,:,::-1]  #switch from RGB to BGR

        im = im.reshape(im.shape[0], im.shape[1], 1)
        im = im.transpose((2,0,1)) # convert to CxHxW
        label[in_idx]=im
    h5_train = os.path.join(path_dst, 'train_data.h5')
    with h5py.File(h5_train,'w') as H:
        H.create_dataset( 'data', data=data ) # note the name X given to the dataset!
        H.create_dataset( 'label', data=label ) # note the name y given to the dataset!
    text_train = os.path.join(path_dst, 'train-path.txt')
    with open(text_train,'w') as L:
        L.write(h5_train) # list all h5 files you are going to use


train_img_paths = './train_img.txt'  #text file of paths to images
train_label_paths = './train_label.txt'  #text file of paths to label images (ground truth)
train_img_hdf5 = '/home/ss/workspace/create_hdf5/' # Your path to h5 file                       

st='Creating Training Data HDF5 File .....'
img_to_hdf5(train_img_paths, train_label_paths,train_img_hdf5,st)
print('DONE...')

0 个答案:

没有答案